Add optional TI functionality to Dreambooth

1 files changed, 3 insertions, 3 deletions

diff --git a/infer.py b/infer.py
index 8e17c4e..01010eb 100644
--- a/infer.py
+++ b/infer.py
@@ -258,7 +258,7 @@ def generate(output_dir, pipeline, args):
     output_dir = output_dir.joinpath(f"{now}_{slugify(args.prompt[0])[:100]}")
     output_dir.mkdir(parents=True, exist_ok=True)
 
-    seed = args.seed or torch.random.seed()
+    args.seed = args.seed or torch.random.seed()
 
     save_args(output_dir, args)
 
@@ -276,7 +276,7 @@ def generate(output_dir, pipeline, args):
                 dynamic_ncols=True
             )
 
-            generator = torch.Generator(device="cuda").manual_seed(seed + i)
+            generator = torch.Generator(device="cuda").manual_seed(args.seed + i)
             images = pipeline(
                 prompt=args.prompt * (args.batch_size // len(args.prompt)),
                 height=args.height,
@@ -290,7 +290,7 @@ def generate(output_dir, pipeline, args):
             ).images
 
             for j, image in enumerate(images):
-                image.save(output_dir.joinpath(f"{seed + i}_{j}.jpg"))
+                image.save(output_dir.joinpath(f"{args.seed + i}_{j}.jpg"))
 
     if torch.cuda.is_available():
         torch.cuda.empty_cache()