From 6a49074dce78615bce54777fb2be3bfd0dd8f780 Mon Sep 17 00:00:00 2001 From: Volpeon Date: Fri, 14 Oct 2022 20:03:01 +0200 Subject: Removed aesthetic gradients; training improvements --- infer.py | 32 ++++---------------------------- 1 file changed, 4 insertions(+), 28 deletions(-) (limited to 'infer.py') 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, } @@ -77,10 +76,6 @@ def create_args_parser(): "--ti_embeddings_dir", 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() -- cgit v1.2.3-54-g00ecf