Update app.py

app.py

@@ -7,6 +7,7 @@ from PIL import Image
 import numpy as np
 import gradio as gr
 import spaces
+import gc
 
 # --- INITIAL SETUP ---
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -28,7 +29,6 @@ print("Moving pipeline to device (ZeroGPU will handle offloading)...")
 pipe.to(device)
 
 # --- LORA SETUP ---
-# We will NOT fuse anything. Everything will be handled dynamically.
 CAUSVID_LORA_REPO = "Kijai/WanVideo_comfy"
 CAUSVID_LORA_FILENAME = "Wan21_CausVid_14B_T2V_lora_rank32.safetensors"
 BASE_LORA_NAME = "causvid_lora"
@@ -44,66 +44,48 @@ except Exception as e:
 
 print("Initialization complete. Gradio is starting...")
 
-def move_adapter_to_device(pipe, adapter_name, device):
-    """
-    Surgically moves only the parameters of a specific LoRA adapter to the target device.
-    This avoids touching the base model's meta tensors.
-    """
-    print(f"Moving adapter '{adapter_name}' to {device}...")
-    for param in pipe.transformer.parameters():
-        if hasattr(param, "adapter_name") and param.adapter_name == adapter_name:
-            param.data = param.data.to(device, non_blocking=True)
-            if param.grad is not None:
-                param.grad.data = param.grad.data.to(device, non_blocking=True)
-    print(f"✅ Adapter '{adapter_name}' moved.")
 
 @spaces.GPU()
 def generate(prompt, negative_prompt, width=1024, height=1024, num_inference_steps=30, lora_id=None, progress=gr.Progress(track_tqdm=True)):
     
     # --- DYNAMIC LORA MANAGEMENT FOR EACH RUN ---
-    # Start with a clean slate by disabling any active adapters from previous runs
-    pipe.disable_lora()
-    
     active_adapters = []
     adapter_weights = []
 
-    # 1. Load the Base LoRA
+    # 1. Load the Base LoRA directly onto the correct device
     if causvid_path:
         try:
-            # We load it for every run to ensure a clean state
             print(f"Loading base LoRA '{BASE_LORA_NAME}'...")
-            pipe.load_lora_weights(causvid_path, adapter_name=BASE_LORA_NAME)
-            
-            # THE CRITICAL FIX: Move only this adapter's weights to the GPU
-            move_adapter_to_device(pipe, BASE_LORA_NAME, device)
-
+            # THE CORRECT FIX: Use device_map to load the LoRA directly to the GPU.
+            pipe.load_lora_weights(causvid_path, adapter_name=BASE_LORA_NAME, device_map={"":device})
             active_adapters.append(BASE_LORA_NAME)
             adapter_weights.append(1.0)
+            print("✅ Base LoRA loaded to device.")
         except Exception as e:
             print(f"⚠️ Failed to load base LoRA: {e}")
 
-    # 2. Load the Custom LoRA if provided
+    # 2. Load the Custom LoRA if provided, also directly to the device
     clean_lora_id = lora_id.strip() if lora_id else ""
     if clean_lora_id:
         try:
             print(f"Loading custom LoRA '{CUSTOM_LORA_NAME}' from '{clean_lora_id}'...")
-            pipe.load_lora_weights(clean_lora_id, adapter_name=CUSTOM_LORA_NAME)
-            
-            # THE CRITICAL FIX: Move only this adapter's weights to the GPU
-            move_adapter_to_device(pipe, CUSTOM_LORA_NAME, device)
-
+            # THE CORRECT FIX: Also use device_map for the custom LoRA.
+            pipe.load_lora_weights(clean_lora_id, adapter_name=CUSTOM_LORA_NAME, device_map={"":device})
             active_adapters.append(CUSTOM_LORA_NAME)
             adapter_weights.append(1.0)
+            print("✅ Custom LoRA loaded to device.")
         except Exception as e:
             print(f"⚠️ Failed to load custom LoRA '{clean_lora_id}': {e}")
-            # If it fails, delete the adapter config to prevent issues
-            if CUSTOM_LORA_NAME in pipe.transformer.peft_config:
+            if CUSTOM_LORA_NAME in getattr(pipe.transformer, 'peft_config', {}):
                 del pipe.transformer.peft_config[CUSTOM_LORA_NAME]
 
     # 3. Activate the successfully loaded adapters
     if active_adapters:
         print(f"Activating adapters: {active_adapters} with weights: {adapter_weights}")
         pipe.set_adapters(active_adapters, adapter_weights)
+    else:
+        # Ensure LoRA is disabled if no adapters were loaded
+        pipe.disable_lora()
     
     apply_cache_on_pipe(pipe)
     
@@ -122,16 +104,13 @@ def generate(prompt, negative_prompt, width=1024, height=1024, num_inference_ste
         return Image.fromarray(image)
     finally:
         # --- PROPER CLEANUP ---
-        print("Cleaning up LoRAs for this run...")
-        # Disable adapters to stop them from being used
-        pipe.disable_lora()
-        
-        # Delete the LoRA configs from the model to truly unload them
-        if BASE_LORA_NAME in pipe.transformer.peft_config:
-            del pipe.transformer.peft_config[BASE_LORA_NAME]
-        if CUSTOM_LORA_NAME in pipe.transformer.peft_config:
-            del pipe.transformer.peft_config[CUSTOM_LORA_NAME]
-        print("✅ LoRAs cleaned up.")
+        # The most reliable way to clean up in this complex environment is to unload ALL LoRAs.
+        # This avoids leaving dangling configs.
+        print("Unloading all LoRAs to ensure a clean state...")
+        pipe.unload_lora_weights()
+        gc.collect() # Force garbage collection
+        torch.cuda.empty_cache() # Clear CUDA cache
+        print("✅ LoRAs unloaded and memory cleaned.")
 
 
 iface = gr.Interface(