Generating one image per prompt

realfake/bin/generate_images.py

@@ -6,6 +6,7 @@ from pathlib import Path
 
 import torch
 from diffusers import StableDiffusionPipeline
+from more_itertools import chunked
 
 from realfake.utils import Args, inject_args
 
@@ -14,6 +15,7 @@ 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"
 
@@ -22,24 +24,38 @@ class GenerateImagesParams(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)
+    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)
+    batch_size = args.batch_size
 
-        left_images = args.num_images
+    if args.num_images == 1:
+        args.output_dir.mkdir(exist_ok=True, parents=True)
+        
         count = 0
-        while left_images > 0:
-            output = pipe(prompt, num_images_per_prompt=min(batch_size, left_images))            
+        for chunk in chunked(prompts, batch_size):
+            output = pipe(chunk, num_images_per_prompt=1)
             for image in output.images:
-                image.save(prompt_dir/f"{count}.png")
+                image.save(args.output_dir/f"{count:04d}.png")
                 count += 1
-            left_images -= batch_size
+
+        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__":