1 files changed, 44 insertions, 12 deletions

diff --git a/infer.py b/infer.py
index f88245a..c4d1e0d 100644
--- a/infer.py
+++ b/infer.py
@@ -45,6 +45,7 @@ default_args = {
 
 
 default_cmds = {
+    "project": "",
     "scheduler": "dpmsm",
     "prompt": None,
     "negative_prompt": None,
@@ -104,6 +105,12 @@ def create_cmd_parser():
         description="Simple example of a training script."
     )
     parser.add_argument(
+        "--project",
+        type=str,
+        default=None,
+        help="The name of the current project.",
+    )
+    parser.add_argument(
         "--scheduler",
         type=str,
         choices=["plms", "ddim", "klms", "dpmsm", "dpmss", "euler_a", "kdpm2", "kdpm2_a"],
@@ -184,7 +191,16 @@ def save_args(basepath, args, extra={}):
         json.dump(info, f, indent=4)
 
 
-def create_pipeline(model, embeddings_dir, dtype):
+def load_embeddings(pipeline, embeddings_dir):
+    added_tokens = load_embeddings_from_dir(
+        pipeline.tokenizer,
+        pipeline.text_encoder.text_model.embeddings,
+        Path(embeddings_dir)
+    )
+    print(f"Added {len(added_tokens)} tokens from embeddings dir: {added_tokens}")
+
+
+def create_pipeline(model, dtype):
     print("Loading Stable Diffusion pipeline...")
 
     tokenizer = MultiCLIPTokenizer.from_pretrained(model, subfolder='tokenizer', torch_dtype=dtype)
@@ -193,10 +209,7 @@ def create_pipeline(model, embeddings_dir, dtype):
     unet = UNet2DConditionModel.from_pretrained(model, subfolder='unet', torch_dtype=dtype)
     scheduler = DDIMScheduler.from_pretrained(model, subfolder='scheduler', torch_dtype=dtype)
 
-    embeddings = patch_managed_embeddings(text_encoder)
-    added_tokens = load_embeddings_from_dir(tokenizer, embeddings, Path(embeddings_dir))
-
-    print(f"Added {len(added_tokens)} tokens from embeddings dir: {added_tokens}")
+    patch_managed_embeddings(text_encoder)
 
     pipeline = VlpnStableDiffusion(
         text_encoder=text_encoder,
@@ -220,7 +233,14 @@ def generate(output_dir, pipeline, args):
         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[0])[:100]}")
+    use_subdirs = len(args.prompt) != 1
+    if use_subdirs:
+        if len(args.project) != 0:
+            output_dir = output_dir.joinpath(f"{now}_{slugify(args.project)}")
+        else:
+            output_dir = output_dir.joinpath(now)
+    else:
+        output_dir = output_dir.joinpath(f"{now}_{slugify(args.prompt[0])[:100]}")
     output_dir.mkdir(parents=True, exist_ok=True)
 
     args.seed = args.seed or torch.random.seed()
@@ -257,7 +277,8 @@ def generate(output_dir, pipeline, args):
             dynamic_ncols=True
         )
 
-        generator = torch.Generator(device="cuda").manual_seed(args.seed + i)
+        seed = args.seed + i
+        generator = torch.Generator(device="cuda").manual_seed(seed)
         images = pipeline(
             prompt=args.prompt,
             negative_prompt=args.negative_prompt,
@@ -272,8 +293,13 @@ def generate(output_dir, pipeline, args):
         ).images
 
         for j, image in enumerate(images):
-            image.save(output_dir.joinpath(f"{args.seed + i}_{j}.png"))
-            image.save(output_dir.joinpath(f"{args.seed + i}_{j}.jpg"), quality=85)
+            image_dir = output_dir
+            if use_subdirs:
+                idx = j % len(args.prompt)
+                image_dir = image_dir.joinpath(slugify(args.prompt[idx])[: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)
 
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
@@ -283,10 +309,11 @@ class CmdParse(cmd.Cmd):
     prompt = 'dream> '
     commands = []
 
-    def __init__(self, output_dir, pipeline, parser):
+    def __init__(self, output_dir, ti_embeddings_dir, pipeline, parser):
         super().__init__()
 
         self.output_dir = output_dir
+        self.ti_embeddings_dir = ti_embeddings_dir
         self.pipeline = pipeline
         self.parser = parser
 
@@ -302,6 +329,10 @@ class CmdParse(cmd.Cmd):
         if elements[0] == 'q':
             return True
 
+        if elements[0] == 'reload_embeddings':
+            load_embeddings(self.pipeline, self.ti_embeddings_dir)
+            return
+
         try:
             args = run_parser(self.parser, default_cmds, elements)
 
@@ -337,9 +368,10 @@ def main():
     output_dir = Path(args.output_dir)
     dtype = {"fp32": torch.float32, "fp16": torch.float16, "bf16": torch.bfloat16}[args.precision]
 
-    pipeline = create_pipeline(args.model, args.ti_embeddings_dir, dtype)
+    pipeline = create_pipeline(args.model, dtype)
+    load_embeddings(pipeline, args.ti_embeddings_dir)
     cmd_parser = create_cmd_parser()
-    cmd_prompt = CmdParse(output_dir, pipeline, cmd_parser)
+    cmd_prompt = CmdParse(output_dir, args.ti_embeddings_dir, pipeline, cmd_parser)
     cmd_prompt.cmdloop()