|
""" |
|
Generates a set of new images to test the quality of the model. |
|
""" |
|
from hashlib import md5 |
|
from pathlib import Path |
|
|
|
import torch |
|
from diffusers import StableDiffusionPipeline |
|
from more_itertools import chunked |
|
|
|
from realfake.utils import Args, inject_args |
|
|
|
|
|
class GenerateImagesParams(Args): |
|
prompts_file: Path |
|
output_dir: Path |
|
num_images: int = 100 |
|
batch_size: int = 16 |
|
model_id: str = "CompVis/stable-diffusion-v1-4" |
|
device: str = "cuda" |
|
|
|
|
|
@inject_args |
|
def main(args: GenerateImagesParams) -> None: |
|
pipe = StableDiffusionPipeline.from_pretrained(args.model_id, torch_dtype=torch.float16) |
|
pipe.to(args.device) |
|
prompts = [p for p in (p.strip() for p in args.prompts_file.read_text().split("\n")) if p] |
|
|
|
batch_size = args.batch_size |
|
|
|
if args.num_images == 1: |
|
args.output_dir.mkdir(exist_ok=True, parents=True) |
|
|
|
count = 0 |
|
for chunk in chunked(prompts, batch_size): |
|
output = pipe(chunk, num_images_per_prompt=1) |
|
for image in output.images: |
|
image.save(args.output_dir/f"{count:04d}.png") |
|
count += 1 |
|
|
|
with (args.output_dir/"prompts.txt").open("w") as fp: |
|
fp.write("\n".join(prompts)) |
|
|
|
else: |
|
for prompt in prompts: |
|
prompt_signature = md5(prompt.encode("utf-8")).hexdigest() |
|
prompt_dir = args.output_dir/prompt_signature |
|
prompt_dir.mkdir(exist_ok=True, parents=True) |
|
(prompt_dir/"prompt.txt").write_text(prompt) |
|
|
|
left_images = args.num_images |
|
count = 0 |
|
while left_images > 0: |
|
output = pipe(prompt, num_images_per_prompt=min(batch_size, left_images)) |
|
for image in output.images: |
|
image.save(prompt_dir/f"{count:04d}.png") |
|
count += 1 |
|
left_images -= batch_size |
|
|
|
|
|
if __name__ == "__main__": |
|
main() |
|
|
