Merge pull request #51 from ExponentialML/feat/lora-infer

Add LoRA Inference.

Author: ExponentialML

inference.py

@@ -3,7 +3,7 @@
 import warnings
 from pathlib import Path
 from uuid import uuid4
-
+from utils.lora import inject_inferable_lora
 import torch
 from diffusers import DPMSolverMultistepScheduler, TextToVideoSDPipeline
 from einops import rearrange
@@ -111,10 +111,12 @@ def inference(
     device="cuda",
     xformers=False,
     sdp=False,
+    lora_path='',
+    lora_rank=64
 ):
     with torch.autocast(device, dtype=torch.half):
         pipeline = initialize_pipeline(model, device, xformers, sdp)
-
+        inject_inferable_lora(pipeline, lora_path, r=lora_rank)
         prompt = [prompt] * batch_size
         negative_prompt = ([negative_prompt] * batch_size) if negative_prompt is not None else None
 
@@ -168,6 +170,8 @@ def inference(
     parser.add_argument("-d", "--device", type=str, default="cuda")
     parser.add_argument("-x", "--xformers", action="store_true")
     parser.add_argument("-S", "--sdp", action="store_true")
+    parser.add_argument("-lP", "--lora_path", type=str, default="")
+    parser.add_argument("-lR", "--lora_rank", type=int, default=64)
     parser.add_argument("-rw", "--remove-watermark", action="store_true")
     args = vars(parser.parse_args())
 

utils/lora.py

@@ -1,6 +1,7 @@
 import json
 import math
 from itertools import groupby
+import os
 from typing import Callable, Dict, List, Optional, Set, Tuple, Type, Union
 
 import numpy as np
@@ -474,6 +475,52 @@ def inject_trainable_lora_extended(
     return require_grad_params, names
 
 
+def inject_inferable_lora(
+        model, 
+        lora_path='', 
+        unet_replace_modules=["UNet3DConditionModel"], 
+        text_encoder_replace_modules=["CLIPEncoderLayer"],
+        is_extended=False, 
+        r=16
+    ):    
+    from transformers.models.clip import CLIPTextModel
+    from diffusers import UNet3DConditionModel
+
+    def is_text_model(f): return 'text_encoder' in f and isinstance(model.text_encoder, CLIPTextModel)
+    def is_unet(f): return 'unet' in f and model.unet.__class__.__name__ == "UNet3DConditionModel"
+
+    if os.path.exists(lora_path):
+        try:
+            for f in os.listdir(lora_path):
+                if f.endswith('.pt'):
+                    lora_file = os.path.join(lora_path, f)
+
+                    if is_text_model(f):
+                        monkeypatch_or_replace_lora(
+                            model.text_encoder,
+                            torch.load(lora_file),
+                            target_replace_module=text_encoder_replace_modules,
+                            r=r
+                        )
+                        print("Successfully loaded Text Encoder LoRa.")
+                        continue
+
+                    if is_unet(f):
+                        monkeypatch_or_replace_lora_extended(
+                            model.unet,
+                            torch.load(lora_file),
+                            target_replace_module=unet_replace_modules,
+                            r=r
+                        )
+                        print("Successfully loaded UNET LoRa.")
+                        continue
+
+                    print("Found a .pt file, but doesn't have the correct name format. (unet.pt, text_encoder.pt)")
+
+        except Exception as e:
+            print(e)
+            print("Couldn't inject LoRA's due to an error.")
+
 def extract_lora_ups_down(model, target_replace_module=DEFAULT_TARGET_REPLACE):
 
     loras = []