Update app.py

app.py

@@ -16,70 +16,93 @@ vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=to
 pipe = WanPipeline.from_pretrained(model_id, vae=vae, torch_dtype=torch.bfloat16)
 flow_shift = 1.0
 pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=flow_shift)
-
 pipe.to(device)
 
 # --- LORA SETUP ---
-# Define unique names for our adapters
-LORA_1_NAME = "causvid_lora"
-LORA_2_NAME = "person_lora"
-
-# 1. Load the first base LoRA ONCE at startup
-print("Loading first LoRA (CausVid)...")
-# LORA_1_REPO = "Kijai/WanVideo_comfy"
-# LORA_1_FILENAME = "Wan21_CausVid_14B_T2V_lora_rank32_v2.safetensors"
-
-# LORA_1_REPO = "vrgamedevgirl84/Wan14BT2VFusioniX"
-# LORA_1_FILENAME = "FusionX_LoRa/Wan2.1_T2V_14B_FusionX_LoRA.safetensors"
-
-LORA_1_REPO = "joerose/Wan21_T2V_14B_lightx2v_cfg_step_distill_lora_rank32"
-LORA_1_FILENAME = "Wan21_T2V_14B_lightx2v_cfg_step_distill_lora_rank32.safetensors"
-
-# LORA_1_REPO = "lightx2v/Wan2.1-T2V-14B-StepDistill-CfgDistill-Lightx2v"
-# LORA_1_FILENAME = "loras/Wan21_T2V_14B_lightx2v_cfg_step_distill_lora_rank64.safetensors"
-
-try:
-    lora_1_path = hf_hub_download(repo_id=LORA_1_REPO, filename=LORA_1_FILENAME)
-    # The `device_map="auto"` can sometimes help in tricky environments
-    pipe.load_lora_weights(lora_1_path, adapter_name=LORA_1_NAME, device_map="auto")
-    print(f"✅ Default LoRA '{LORA_1_NAME}' loaded successfully.")
-except Exception as e:
-    print(f"⚠️ Default LoRA '{LORA_1_NAME}' could not be loaded: {e}")
-
-# 2. Load the second hard-coded LoRA ONCE at startup
-print("Loading second LoRA (Person)...")
-# LORA_2_REPO = "ovi054/p3r5onVid1900" 
-# LORA_2_REPO = "ovi054/rosmxVid1500" 
-LORA_2_REPO = "ovi054/ovimxVid1750" 
-# Assuming the file is named "pytorch_lora_weights.safetensors" which is standard.
-# If it has a different name, you must specify it with the `filename` argument.
-try:
-    # We load the whole repository and diffusers will find the correct file
-    pipe.load_lora_weights(LORA_2_REPO, adapter_name=LORA_2_NAME, device_map="auto")
-    print(f"✅ Second LoRA '{LORA_2_NAME}' loaded successfully.")
-except Exception as e:
-    print(f"⚠️ Second LoRA '{LORA_2_NAME}' could not be loaded: {e}")
-
-pipe.set_adapters([LORA_1_NAME, LORA_2_NAME], adapter_weights=[1.0, 1.0])
-
+CAUSVID_NAME = "causvid_base"
+PERSONVID_NAME = "personvid_optional"
+ROSMX_NAME = "rosmx_optional"
+OVIMX_NAME = "ovimx_optional"
+OVIMX2_NAME = "ovimx2_optional"
+OVIMX3_NAME = "ovimx3_optional"
+PRTHE_NAME = "prthe_optional"
+
+SUCCESSFULLY_LOADED_LORAS = {}
+
+lora_definitions = {
+    CAUSVID_NAME: ("joerose/Wan21_T2V_14B_lightx2v_cfg_step_distill_lora_rank32", "Wan21_T2V_14B_lightx2v_cfg_step_distill_lora_rank32.safetensors"),
+    PERSONVID_NAME: ("ovi054/p3r5onVid1000", None),
+    ROSMX_NAME: ("ovi054/rosmxVid1500", None),
+    OVIMX_NAME: ("ovi054/ovimxVid1750", None),
+    OVIMX2_NAME: ("ovi054/ovimxVid2250", None),
+    OVIMX3_NAME: ("ovi054/ovimxVid2500", None),
+    PRTHE_NAME: ("ovi054/prwthxVid", None)
+}
+
+for name, (repo, filename) in lora_definitions.items():
+    print(f"Attempting to load LoRA '{name}'...")
+    try:
+        if filename:
+            path = hf_hub_download(repo_id=repo, filename=filename)
+            pipe.load_lora_weights(path, adapter_name=name, device_map="auto")
+        else:
+            pipe.load_lora_weights(repo, adapter_name=name, device_map="auto")
+        print(f"✅ LoRA '{name}' loaded successfully.")
+        SUCCESSFULLY_LOADED_LORAS[name] = repo
+    except Exception as e:
+        print(f"⚠️ LoRA '{name}' could not be loaded: {e}")
+
+OPTIONAL_LORA_MAP = {
+    "ovi054/p3r5onVid1000": PERSONVID_NAME,
+    "ovi054/rosmxVid1500": ROSMX_NAME,
+    "ovi054/ovimxVid1750": OVIMX_NAME,
+    "ovi054/ovimxVid2250": OVIMX2_NAME,
+    "ovi054/ovimxVid2500": OVIMX3_NAME,
+    "ovi054/prwthxVid": PRTHE_NAME,
+}
+OPTIONAL_LORA_CHOICES = {k: v for k, v in OPTIONAL_LORA_MAP.items() if v in SUCCESSFULLY_LOADED_LORAS}
+
+# At startup, disable all adapters. They will be selectively enabled during each run.
+print("Disabling all LoRAs at startup. They will be activated on-demand.")
+pipe.disable_lora()
 
 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)):
-    
-    # --- Activate both hard-coded LoRAs for this run ---
-    # We set the adapters at the start of every generation to ensure the state is correct.
-    print("Activating both LoRAs for inference...")
-    # You can adjust the weights here to change the intensity of each LoRA.
-    # For example, [1.0, 0.8] would make the second LoRA less strong.
-    # pipe.set_adapters([LORA_1_NAME, LORA_2_NAME], adapter_weights=[1.0, 1.0])
-
-    apply_cache_on_pipe(
-        pipe,
-    )
+def generate(prompt, negative_prompt, width, height, num_inference_steps, optional_lora_id, progress=gr.Progress(track_tqdm=True)):
     
+    # Using a try...finally block is robust for state management in apps.
     try:
+        # --- Step 1: Build the list of ACTIVE adapters and their weights for THIS run ---
+        
+        active_adapters = []
+        adapter_weights = []
+        
+        # Always include the base LoRA if it was loaded successfully
+        if CAUSVID_NAME in SUCCESSFULLY_LOADED_LORAS:
+            active_adapters.append(CAUSVID_NAME)
+            adapter_weights.append(1.0)
+            
+        # If an optional LoRA is selected, add it to the list
+        if optional_lora_id and optional_lora_id != "None":
+            internal_name_to_add = OPTIONAL_LORA_CHOICES.get(optional_lora_id)
+            if internal_name_to_add:
+                active_adapters.append(internal_name_to_add)
+                adapter_weights.append(1.0)
+
+        # --- Step 2: Apply the adapters and weights for this run using the correct function ---
+        if active_adapters:
+            print(f"Activating adapters: {active_adapters} with weights: {adapter_weights}")
+            # This is the correct, modern way to set adapters and their weights.
+            pipe.set_adapters(active_adapters, adapter_weights=adapter_weights)
+        else:
+            print("No LoRAs are active for this run.")
+            # ensure all are disabled if for some reason none were selected
+            pipe.disable_lora()
+
+        apply_cache_on_pipe(pipe)
+        
+        # --- Step 3: Run inference ---
         output = pipe(
             prompt=prompt,
             negative_prompt=negative_prompt,
@@ -92,25 +115,23 @@ def generate(prompt, negative_prompt, width=1024, height=1024, num_inference_ste
         image = output.frames[0][0]
         image = (image * 255).astype(np.uint8)
         return Image.fromarray(image)
-    finally:
-        # It's good practice to disable the adapters after the run,
-        # although set_adapters() at the start also handles this.
-        # print("Disabling LoRAs after run.")
-        # pipe.disable_lora()
-        pass
 
+    finally:
+        print("Disabling LoRAs after run to reset state.")
+        pipe.disable_lora()
 
+# --- Gradio Interface ---
 iface = gr.Interface(
     fn=generate,
     inputs=[
         gr.Textbox(label="Input prompt"),
-    ],
-    additional_inputs = [
         gr.Textbox(label="Negative prompt", value = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"),
         gr.Slider(label="Width", minimum=480, maximum=1280, step=16, value=1024),
         gr.Slider(label="Height", minimum=480, maximum=1280, step=16, value=1024),
         gr.Slider(minimum=1, maximum=80, step=1, label="Inference Steps", value=10),
-        gr.Textbox(label="LoRA ID", visible=False), # Hiding the dynamic LoRA input for now
+        gr.Textbox(
+            label="Optional LoRA",
+        )
     ],
     outputs=gr.Image(label="output"),
 )