1 files changed, 37 insertions, 2 deletions

diff --git a/infer.py b/infer.py
index ed86ab1..93848d7 100644
--- a/infer.py
+++ b/infer.py
@@ -26,6 +26,8 @@ from diffusers import (
     DEISMultistepScheduler,
     UniPCMultistepScheduler
 )
+from peft import LoraConfig, LoraModel, set_peft_model_state_dict
+from safetensors.torch import load_file
 from transformers import CLIPTextModel
 
 from data.keywords import str_to_keywords, keywords_to_str
@@ -43,7 +45,7 @@ default_args = {
     "model": "stabilityai/stable-diffusion-2-1",
     "precision": "fp32",
     "ti_embeddings_dir": "embeddings_ti",
-    "lora_embeddings_dir": "embeddings_lora",
+    "lora_embedding": None,
     "output_dir": "output/inference",
     "config": None,
 }
@@ -99,7 +101,7 @@ def create_args_parser():
         type=str,
     )
     parser.add_argument(
-        "--lora_embeddings_dir",
+        "--lora_embedding",
         type=str,
     )
     parser.add_argument(
@@ -236,6 +238,38 @@ def load_embeddings(pipeline, embeddings_dir):
     print(f"Added {len(added_tokens)} tokens from embeddings dir: {list(zip(added_tokens, added_ids))}")
 
 
+def load_lora(pipeline, path):
+    if path is None:
+        return
+
+    path = Path(path)
+
+    with open(path / "lora_config.json", "r") as f:
+        lora_config = json.load(f)
+
+    tensor_files = list(path.glob("*_end.safetensors"))
+
+    if len(tensor_files) == 0:
+        return
+
+    lora_checkpoint_sd = load_file(path / tensor_files[0])
+    unet_lora_ds = {k: v for k, v in lora_checkpoint_sd.items() if "text_encoder_" not in k}
+    text_encoder_lora_ds = {
+        k.replace("text_encoder_", ""): v for k, v in lora_checkpoint_sd.items() if "text_encoder_" in k
+    }
+
+    unet_config = LoraConfig(**lora_config["peft_config"])
+    pipeline.unet = LoraModel(unet_config, pipeline.unet)
+    set_peft_model_state_dict(pipeline.unet, unet_lora_ds)
+
+    if "text_encoder_peft_config" in lora_config:
+        text_encoder_config = LoraConfig(**lora_config["text_encoder_peft_config"])
+        pipeline.text_encoder = LoraModel(text_encoder_config, pipeline.text_encoder)
+        set_peft_model_state_dict(pipeline.text_encoder, text_encoder_lora_ds)
+
+    return
+
+
 def create_scheduler(config, scheduler: str, subscheduler: Optional[str] = None):
     if scheduler == "plms":
         return PNDMScheduler.from_config(config)
@@ -441,6 +475,7 @@ def main():
     pipeline = create_pipeline(args.model, dtype)
 
     load_embeddings(pipeline, args.ti_embeddings_dir)
+    load_lora(pipeline, args.lora_embedding)
     # pipeline.unet.load_attn_procs(args.lora_embeddings_dir)
 
     cmd_parser = create_cmd_parser()