Update app.py

app.py

@@ -65,26 +65,20 @@ def generate(prompt, negative_prompt, width=1024, height=1024, num_inference_ste
     #     pipe.unload_lora_weights()
     #     pipe.load_lora_weights(lora_id.strip())
 
-    clean_lora_id = lora_id.strip() if lora_id else ""
-
-    # try:
-    # --- IF LORA ID IS NONE, USE ONE LORA ---
-    if not clean_lora_id:
-        print("No custom LoRA provided. Using ONLY the base LoRA.")
-        # Activate just the default LoRA
-        pipe.set_adapters([DEFAULT_LORA_NAME], adapter_weights=[1.0])
+    print("Loading base LoRA for this run...")
+    causvid_path = hf_hub_download(repo_id=CAUSVID_LORA_REPO, filename=CAUSVID_LORA_FILENAME)
+    pipe.load_lora_weights(causvid_path, adapter_name=DEFAULT_LORA_NAME)
 
-    # --- OTHERWISE, LOAD AND USE TWO LORAS ---
-    else:
-        print(f"Custom LoRA provided. Loading '{clean_lora_id}' and combining with base LoRA.")
-        # Load the custom LoRA fresh for this run
+    # If a custom LoRA is provided, load it as well.
+    if clean_lora_id:
+        print(f"Loading custom LoRA '{clean_lora_id}' for this run...")
         pipe.load_lora_weights(clean_lora_id, adapter_name=CUSTOM_LORA_NAME)
-
-        # Activate BOTH LoRAs together
-        pipe.set_adapters(
-            [DEFAULT_LORA_NAME, CUSTOM_LORA_NAME],
-            adapter_weights=[1.0, 1.0]  # Strength for base, strength for custom
-        )
+        # If a custom LoRA is present, activate both.
+        pipe.set_adapters([DEFAULT_LORA_NAME, CUSTOM_LORA_NAME], adapter_weights=[1.0, 1.0])
+    else:
+        # If no custom LoRA, just activate the base one.
+        print("Activating base LoRA only.")
+        pipe.set_adapters([DEFAULT_LORA_NAME], adapter_weights=[1.0])
 
 
     pipe.to("cuda")
@@ -110,12 +104,15 @@ def generate(prompt, negative_prompt, width=1024, height=1024, num_inference_ste
         # if lora_id and lora_id.strip() != "":
         #     pass
             # pipe.unload_lora_weights()
-        if clean_lora_id:
-            print(f"Unloading '{CUSTOM_LORA_NAME}' from this run.")
-            pipe.unload_lora_weights(CUSTOM_LORA_NAME)
+        # if clean_lora_id:
+        #     print(f"Unloading '{CUSTOM_LORA_NAME}' from this run.")
+        #     pipe.unload_lora_weights(CUSTOM_LORA_NAME)
+
+        # # Always disable all active LoRAs to reset the state.
+        # pipe.disable_lora()
+        print("Unloading all LoRAs to clean up.")
+        pipe.unload_lora_weights()
 
-        # Always disable all active LoRAs to reset the state.
-        pipe.disable_lora()
 
 iface = gr.Interface(
     fn=generate,