Update models.py

models.py

@@ -177,9 +177,11 @@ def load_lora(pipe):
 
 def fuse_lora_with_scale(pipe, lora_scale):
     """
-    Following examplewithface.py lines 223-235:
-    Use the Kohya-style LoRA loader from lora.py
-    CRITICAL: Pass pipe.text_encoder (not a list) for the example to work
+    Following examplewithface.py lines 266-267 EXACTLY:
+    pipe.load_lora_weights(loaded_state_dict)
+    pipe.fuse_lora(lora_scale)
+    
+    Uses DIFFUSERS built-in LoRA (NOT Kohya lora.py!)
     """
     global lora_path_cached
     
@@ -187,38 +189,39 @@ def fuse_lora_with_scale(pipe, lora_scale):
         return False
     
     try:
-        # Import the local lora module (Kohya-style)
-        import lora
-        
-        print(f"  [LORA] Creating network from weights...")
+        # Unfuse previous LoRA (example line 259)
+        try:
+            pipe.unfuse_lora()
+        except:
+            pass
         
-        # examplewithface.py lines 223-229
-        # IMPORTANT: Example passes pipe.text_encoder (singular), not a list!
-        # The LoRA network will handle SDXL's dual encoders internally
-        lora_model, weights_sd = lora.create_network_from_weights(
-            lora_scale,  # multiplier
-            lora_path_cached,  # file path  
-            pipe.vae,
-            pipe.text_encoder,  # Single text encoder (example line 227)
-            pipe.unet,
-            for_inference=True,
-        )
+        # Unload previous LoRA (example line 260)
+        try:
+            pipe.unload_lora_weights()
+        except:
+            pass
         
-        # examplewithface.py lines 231-233
-        print(f"  [LORA] Merging to model with scale {lora_scale}...")
-        lora_model.merge_to(
-            pipe.text_encoder, pipe.unet, weights_sd, torch.float16, "cuda"
-        )
+        print(f"  [LORA] Loading state dict from file...")
+        # Load state dict like example (lines 75-78)
+        if lora_path_cached.endswith('.safetensors'):
+            from safetensors.torch import load_file
+            state_dict = load_file(lora_path_cached)
+        else:
+            state_dict = torch.load(lora_path_cached, map_location="cpu")
         
-        # Cleanup
-        del weights_sd
-        del lora_model
+        print(f"  [LORA] Loading weights into pipeline...")
+        # examplewithface.py line 266
+        pipe.load_lora_weights(state_dict)
         
-        print(f"  [OK] LoRA merged into model using Kohya loader")
+        # examplewithface.py line 267
+        print(f"  [LORA] Fusing with scale {lora_scale}...")
+        pipe.fuse_lora(lora_scale)
         
+        print(f"  [OK] LoRA fused into model (diffusers method)")
         return True
+        
     except Exception as e:
-        print(f"  [ERROR] LoRA merge failed: {e}")
+        print(f"  [ERROR] LoRA fusion failed: {e}")
         import traceback
         traceback.print_exc()
         return False