Update models.py

models.py

@@ -1,7 +1,7 @@
 """
 Models.py - Following examplewithface.py EXACTLY
 NO MultiControlNetModel wrapper!
-Using fuse_lora with scale (examplewithface.py line 267)
+Using Kohya-style LoRA from lora.py (examplewithface.py lines 223-235)
 """
 import torch
 import time
@@ -160,8 +160,8 @@ def load_lora(pipe):
 
 def fuse_lora_with_scale(pipe, lora_scale):
     """
-    Following examplewithface.py lines 266-267:
-    Load LoRA weights and FUSE them into the model
+    Following examplewithface.py lines 223-235:
+    Use the Kohya-style LoRA loader from lora.py (NOT diffusers built-in)
     """
     global lora_path_cached
     
@@ -169,25 +169,39 @@ def fuse_lora_with_scale(pipe, lora_scale):
         return False
     
     try:
-        # Unload and unfuse previous LoRA if exists
-        try:
-            pipe.unfuse_lora()
-            pipe.unload_lora_weights()
-        except:
-            pass
+        # Import the local lora module (Kohya-style)
+        import lora
+        
+        print(f"  [LORA] Creating network from weights...")
+        
+        # examplewithface.py lines 223-229
+        # Note: SDXL has two text encoders, pass both as a list
+        text_encoders = [pipe.text_encoder, pipe.text_encoder_2]
+        
+        lora_model, weights_sd = lora.create_network_from_weights(
+            lora_scale,  # multiplier
+            lora_path_cached,  # file path
+            pipe.vae,
+            text_encoders,  # Both SDXL text encoders
+            pipe.unet,
+            for_inference=True,
+        )
+        
+        # examplewithface.py lines 231-233
+        print(f"  [LORA] Merging to model with scale {lora_scale}...")
+        lora_model.merge_to(
+            text_encoders, pipe.unet, weights_sd, torch.float16, "cuda"
+        )
         
-        # Load LoRA weights (examplewithface.py line 266)
-        print(f"  [LORA] Loading weights...")
-        pipe.load_lora_weights(lora_path_cached)
+        # Cleanup
+        del weights_sd
+        del lora_model
         
-        # CRITICAL: Fuse LoRA into model (examplewithface.py line 267)
-        print(f"  [LORA] Fusing with scale {lora_scale}...")
-        pipe.fuse_lora(lora_scale=lora_scale)
-        print(f"  [OK] LoRA fused into model")
+        print(f"  [OK] LoRA merged into model using Kohya loader")
         
         return True
     except Exception as e:
-        print(f"  [ERROR] LoRA fusion failed: {e}")
+        print(f"  [ERROR] LoRA merge failed: {e}")
         import traceback
         traceback.print_exc()
         return False