Generating fakes with SD

realfake/bin/generate_images.py

@@ -0,0 +1,44 @@
+"""
+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 realfake.utils import Args, inject_args
+
+
+class GenerateImagesParams(Args):
+    prompts_file: Path
+    output_dir: Path
+    num_images: int = 100
+    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 = 16
+    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
+        while left_images > 0:
+            output = pipe(prompt, num_images_per_prompt=min(batch_size, left_images))            
+            for i, image in enumerate(output.images):
+                image.save(prompt_dir/f"{i}.png")
+            left_images -= batch_size
+
+
+if __name__ == "__main__":
+    main()