Removed aesthetic gradients; training improvements

1 files changed, 4 insertions, 28 deletions

diff --git a/infer.py b/infer.py
index 650c119..1a0baf5 100644
--- a/infer.py
+++ b/infer.py
@@ -24,7 +24,6 @@ default_args = {
     "scheduler": "euler_a",
     "precision": "fp32",
     "ti_embeddings_dir": "embeddings_ti",
-    "ag_embeddings_dir": "embeddings_ag",
     "output_dir": "output/inference",
     "config": None,
 }
@@ -78,10 +77,6 @@ def create_args_parser():
         type=str,
     )
     parser.add_argument(
-        "--ag_embeddings_dir",
-        type=str,
-    )
-    parser.add_argument(
         "--output_dir",
         type=str,
     )
@@ -211,24 +206,7 @@ def load_embeddings_ti(tokenizer, text_encoder, embeddings_dir):
             print(f"Loaded {placeholder_token}")
 
 
-def load_embeddings_ag(pipeline, embeddings_dir):
-    print(f"Loading Aesthetic Gradient embeddings")
-
-    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")
-
-            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):
+def create_pipeline(model, scheduler, ti_embeddings_dir, dtype):
     print("Loading Stable Diffusion pipeline...")
 
     tokenizer = CLIPTokenizer.from_pretrained(model, subfolder='tokenizer', torch_dtype=dtype)
@@ -262,13 +240,11 @@ def create_pipeline(model, scheduler, ti_embeddings_dir, ag_embeddings_dir, dtyp
         tokenizer=tokenizer,
         scheduler=scheduler,
     )
-    pipeline.aesthetic_gradient_iters = 30
+    pipeline.aesthetic_gradient_iters = 20
     pipeline.to("cuda")
 
     print("Pipeline loaded.")
 
-    load_embeddings_ag(pipeline, ag_embeddings_dir)
-
     return pipeline
 
 
@@ -288,7 +264,7 @@ def generate(output_dir, pipeline, args):
     else:
         init_image = None
 
-    with torch.autocast("cuda"):
+    with torch.autocast("cuda"), torch.inference_mode():
         for i in range(args.batch_num):
             pipeline.set_progress_bar_config(
                 desc=f"Batch {i + 1} of {args.batch_num}",
@@ -366,7 +342,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.ti_embeddings_dir, args.ag_embeddings_dir, dtype)
+    pipeline = create_pipeline(args.model, args.scheduler, args.ti_embeddings_dir, dtype)
     cmd_parser = create_cmd_parser()
     cmd_prompt = CmdParse(output_dir, pipeline, cmd_parser)
     cmd_prompt.cmdloop()