1 files changed, 52 insertions, 23 deletions

diff --git a/infer.py b/infer.py
index 63b16d8..650c119 100644
--- a/infer.py
+++ b/infer.py
@@ -23,7 +23,8 @@ default_args = {
     "model": None,
     "scheduler": "euler_a",
     "precision": "fp32",
-    "embeddings_dir": "embeddings",
+    "ti_embeddings_dir": "embeddings_ti",
+    "ag_embeddings_dir": "embeddings_ag",
     "output_dir": "output/inference",
     "config": None,
 }
@@ -73,7 +74,11 @@ def create_args_parser():
         choices=["fp32", "fp16", "bf16"],
     )
     parser.add_argument(
-        "--embeddings_dir",
+        "--ti_embeddings_dir",
+        type=str,
+    )
+    parser.add_argument(
+        "--ag_embeddings_dir",
         type=str,
     )
     parser.add_argument(
@@ -167,42 +172,63 @@ def save_args(basepath, args, extra={}):
         json.dump(info, f, indent=4)
 
 
-def load_embeddings(tokenizer, text_encoder, embeddings_dir):
+def load_embeddings_ti(tokenizer, text_encoder, embeddings_dir):
+    print(f"Loading Textual Inversion embeddings")
+
     embeddings_dir = Path(embeddings_dir)
     embeddings_dir.mkdir(parents=True, exist_ok=True)
 
     for file in embeddings_dir.iterdir():
-        placeholder_token = file.stem
+        if file.is_file():
+            placeholder_token = file.stem
 
-        num_added_tokens = tokenizer.add_tokens(placeholder_token)
-        if num_added_tokens == 0:
-            raise ValueError(
-                f"The tokenizer already contains the token {placeholder_token}. Please pass a different"
-                " `placeholder_token` that is not already in the tokenizer."
-            )
+            num_added_tokens = tokenizer.add_tokens(placeholder_token)
+            if num_added_tokens == 0:
+                raise ValueError(
+                    f"The tokenizer already contains the token {placeholder_token}. Please pass a different"
+                    " `placeholder_token` that is not already in the tokenizer."
+                )
 
     text_encoder.resize_token_embeddings(len(tokenizer))
 
     token_embeds = text_encoder.get_input_embeddings().weight.data
 
     for file in embeddings_dir.iterdir():
-        placeholder_token = file.stem
-        placeholder_token_id = tokenizer.convert_tokens_to_ids(placeholder_token)
+        if file.is_file():
+            placeholder_token = file.stem
+            placeholder_token_id = tokenizer.convert_tokens_to_ids(placeholder_token)
+
+            data = torch.load(file, map_location="cpu")
+
+            assert len(data.keys()) == 1, 'embedding file has multiple terms in it'
+
+            emb = next(iter(data.values()))
+            if len(emb.shape) == 1:
+                emb = emb.unsqueeze(0)
 
-        data = torch.load(file, map_location="cpu")
+            token_embeds[placeholder_token_id] = emb
 
-        assert len(data.keys()) == 1, 'embedding file has multiple terms in it'
+            print(f"Loaded {placeholder_token}")
 
-        emb = next(iter(data.values()))
-        if len(emb.shape) == 1:
-            emb = emb.unsqueeze(0)
 
-        token_embeds[placeholder_token_id] = emb
+def load_embeddings_ag(pipeline, embeddings_dir):
+    print(f"Loading Aesthetic Gradient embeddings")
 
-        print(f"Loaded embedding: {placeholder_token}")
+    embeddings_dir = Path(embeddings_dir)
+    embeddings_dir.mkdir(parents=True, exist_ok=True)
+
+    for file in embeddings_dir.iterdir():
+        if file.is_file():
+            placeholder_token = file.stem
 
+            data = torch.load(file, map_location="cpu")
 
-def create_pipeline(model, scheduler, embeddings_dir, dtype):
+            pipeline.add_aesthetic_gradient_embedding(placeholder_token, data)
+
+            print(f"Loaded {placeholder_token}")
+
+
+def create_pipeline(model, scheduler, ti_embeddings_dir, ag_embeddings_dir, dtype):
     print("Loading Stable Diffusion pipeline...")
 
     tokenizer = CLIPTokenizer.from_pretrained(model, subfolder='tokenizer', torch_dtype=dtype)
@@ -210,7 +236,7 @@ def create_pipeline(model, scheduler, embeddings_dir, dtype):
     vae = AutoencoderKL.from_pretrained(model, subfolder='vae', torch_dtype=dtype)
     unet = UNet2DConditionModel.from_pretrained(model, subfolder='unet', torch_dtype=dtype)
 
-    load_embeddings(tokenizer, text_encoder, embeddings_dir)
+    load_embeddings_ti(tokenizer, text_encoder, ti_embeddings_dir)
 
     if scheduler == "plms":
         scheduler = PNDMScheduler(
@@ -236,10 +262,13 @@ def create_pipeline(model, scheduler, embeddings_dir, dtype):
         tokenizer=tokenizer,
         scheduler=scheduler,
     )
+    pipeline.aesthetic_gradient_iters = 30
     pipeline.to("cuda")
 
     print("Pipeline loaded.")
 
+    load_embeddings_ag(pipeline, ag_embeddings_dir)
+
     return pipeline
 
 
@@ -259,7 +288,7 @@ def generate(output_dir, pipeline, args):
     else:
         init_image = None
 
-    with torch.autocast("cuda"), torch.inference_mode():
+    with torch.autocast("cuda"):
         for i in range(args.batch_num):
             pipeline.set_progress_bar_config(
                 desc=f"Batch {i + 1} of {args.batch_num}",
@@ -337,7 +366,7 @@ 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.scheduler, args.embeddings_dir, dtype)
+    pipeline = create_pipeline(args.model, args.scheduler, args.ti_embeddings_dir, args.ag_embeddings_dir, dtype)
     cmd_parser = create_cmd_parser()
     cmd_prompt = CmdParse(output_dir, pipeline, cmd_parser)
     cmd_prompt.cmdloop()