Update app.py

app.py

@@ -36,23 +36,67 @@ pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow
 #     flow_shift=flow_shift           # Retain flow_shift for WanPipeline compatibility
 # )
 
+# --- LoRA State Management ---
+# Define unique names for our adapters
+DEFAULT_LORA_NAME = "causvid_lora"
+CUSTOM_LORA_NAME = "custom_lora"
+# Track which custom LoRA is currently loaded to avoid reloading
+CURRENTLY_LOADED_CUSTOM_LORA = None
+
+# Load the default base LoRA ONCE at startup
+print("Loading base LoRA...")
 CAUSVID_LORA_REPO = "WanVideo_comfy"
 CAUSVID_LORA_FILENAME = "Wan21_CausVid_bidirect2_T2V_1_3B_lora_rank32.safetensors"
-
 try:
     causvid_path = hf_hub_download(repo_id=CAUSVID_LORA_REPO, filename=CAUSVID_LORA_FILENAME)
-    pipe.load_lora_weights(causvid_path, adapter_name="causvid_lora")
-    print("✅ CausVid LoRA loaded (strength: 1.0)")
+    pipe.load_lora_weights(causvid_path, adapter_name=DEFAULT_LORA_NAME)
+    print(f"✅ Default LoRA '{DEFAULT_LORA_NAME}' loaded successfully.")
 except Exception as e:
-    print(f"⚠️ CausVid LoRA not loaded: {e}")
-    causvid_path = None
+    print(f"⚠️ Default LoRA could not be loaded: {e}")
+    DEFAULT_LORA_NAME = None
+
+print("Initialization complete. Gradio is starting...")
+
 
 
 @spaces.GPU()
 def generate(prompt, negative_prompt, width=1024, height=1024, num_inference_steps=30, lora_id=None, progress=gr.Progress(track_tqdm=True)):
-    if lora_id and lora_id.strip() != "":
-        # pipe.unload_lora_weights()
-        pipe.load_lora_weights(lora_id.strip())
+    # if lora_id and lora_id.strip() != "":
+    #     pipe.unload_lora_weights()
+    #     pipe.load_lora_weights(lora_id.strip())
+
+    global CURRENTLY_LOADED_CUSTOM_LORA
+
+    active_adapters = []
+    adapter_weights = []
+
+    # Always add the default LoRA if it was loaded successfully
+    if DEFAULT_LORA_NAME:
+        active_adapters.append(DEFAULT_LORA_NAME)
+        adapter_weights.append(1.0) # Strength for base LoRA
+
+    # Handle the user-provided custom LoRA
+    clean_lora_id = lora_id.strip() if lora_id else ""
+    if clean_lora_id:
+        try:
+            # If the requested LoRA is different from the one in memory, swap it
+            if clean_lora_id != CURRENTLY_LOADED_CUSTOM_LORA:
+                print(f"Switching custom LoRA to: {clean_lora_id}")
+                # Unload the old custom LoRA to save memory
+                if CURRENTLY_LOADED_CUSTOM_LORA is not None:
+                    pipe.unload_lora_weights(CUSTOM_LORA_NAME)
+                
+                # Load the new one with its unique name
+                pipe.load_lora_weights(clean_lora_id, adapter_name=CUSTOM_LORA_NAME)
+                CURRENTLY_LOADED_CUSTOM_LORA = clean_lora_id
+            
+            # Add the custom LoRA to the active list for this generation
+            active_adapters.append(CUSTOM_LORA_NAME)
+            adapter_weights.append(1.0) # Strength for custom LoRA
+
+        except Exception as e:
+            print(f"⚠️ Failed to load custom LoRA '{clean_lora_id}'. Error: {e}")
+
     pipe.to("cuda")
     # apply_first_block_cache(pipe.transformer, FirstBlockCacheConfig(threshold=0.2))
     apply_cache_on_pipe(