1 files changed, 5 insertions, 10 deletions

diff --git a/infer.py b/infer.py
index d744768..8e17c4e 100644
--- a/infer.py
+++ b/infer.py
@@ -19,9 +19,6 @@ from pipelines.stable_diffusion.vlpn_stable_diffusion import VlpnStableDiffusion
 torch.backends.cuda.matmul.allow_tf32 = True
 
 
-line_sep = " <OR> "
-
-
 default_args = {
     "model": None,
     "scheduler": "euler_a",
@@ -254,8 +251,11 @@ def create_pipeline(model, scheduler, ti_embeddings_dir, dtype):
 
 
 def generate(output_dir, pipeline, args):
+    if isinstance(args.prompt, str):
+        args.prompt = [args.prompt]
+
     now = datetime.datetime.now().strftime("%Y-%m-%dT%H-%M-%S")
-    output_dir = output_dir.joinpath(f"{now}_{slugify(args.prompt)[:100]}")
+    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()
@@ -276,14 +276,9 @@ def generate(output_dir, pipeline, args):
                 dynamic_ncols=True
             )
 
-            if isinstance(args.prompt, str):
-                args.prompt = [args.prompt]
-
-            prompt = [p.split(line_sep) for p in args.prompt] * args.batch_size
-
             generator = torch.Generator(device="cuda").manual_seed(seed + i)
             images = pipeline(
-                prompt=prompt,
+                prompt=args.prompt * (args.batch_size // len(args.prompt)),
                 height=args.height,
                 width=args.width,
                 negative_prompt=args.negative_prompt,