Update

1 files changed, 3 insertions, 4 deletions

diff --git a/infer.py b/infer.py
index c4d1e0d..b29b136 100644
--- a/infer.py
+++ b/infer.py
@@ -295,11 +295,10 @@ def generate(output_dir, pipeline, args):
         for j, image in enumerate(images):
             image_dir = output_dir
             if use_subdirs:
-                idx = j % len(args.prompt)
-                image_dir = image_dir.joinpath(slugify(args.prompt[idx])[:100])
+                image_dir = image_dir.joinpath(slugify(args.prompt[j % len(args.prompt)])[:100])
             image_dir.mkdir(parents=True, exist_ok=True)
-            image.save(image_dir.joinpath(f"{seed}_{j}.png"))
-            image.save(image_dir.joinpath(f"{seed}_{j}.jpg"), quality=85)
+            image.save(image_dir.joinpath(f"{seed}_{j // len(args.prompt)}.png"))
+            image.save(image_dir.joinpath(f"{seed}_{j // len(args.prompt)}.jpg"), quality=85)
 
     if torch.cuda.is_available():
         torch.cuda.empty_cache()