app.py

import torch
from diffusers import UniPCMultistepScheduler
from diffusers import WanPipeline, AutoencoderKLWan
from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe
from huggingface_hub import hf_hub_download
from PIL import Image
import numpy as np
import gradio as gr
import spaces

device = "cuda" if torch.cuda.is_available() else "cpu"

# --- MODEL SETUP ---
model_id = "Wan-AI/Wan2.1-T2V-14B-Diffusers"
vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32)
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 ---
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"
ROSMX2_NAME = "rosmx2_optional"
PERSONVID2_NAME = "personvid2_optional"

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),
    ROSMX2_NAME: ("ovi054/rosmxVid2000", None),
    PERSONVID2_NAME: ("ovi054/p3r5onVid1900", None)
}

# --- THIS ORDERED LIST IS NOW CRITICAL ---
# It defines the consistent order for the weight vector.
ALL_ADAPTER_NAMES = []
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.")
        ALL_ADAPTER_NAMES.append(name)
    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,
    "ovi054/rosmxVid2000": ROSMX2_NAME,
    "ovi054/p3r5onVid1900": PERSONVID2_NAME,
}
# Filter choices to only include LoRAs that actually loaded
OPTIONAL_LORA_CHOICES = {k: v for k, v in OPTIONAL_LORA_MAP.items() if v in ALL_ADAPTER_NAMES}


# --- SET INITIAL STATE AT STARTUP ---
# Set ALL adapters as active, but control them with weights.
if ALL_ADAPTER_NAMES:
    print(f"Setting up all {len(ALL_ADAPTER_NAMES)} loaded adapters in the pipeline.")
    
    # Start with all weights at 0.0
    initial_weights = [0.0] * len(ALL_ADAPTER_NAMES)
    
    # Set the base LoRA's weight to 1.0
    try:
        base_lora_index = ALL_ADAPTER_NAMES.index(CAUSVID_NAME)
        initial_weights[base_lora_index] = 1.0
    except ValueError:
        print(f"Warning: Base LoRA '{CAUSVID_NAME}' not found in the loaded list. All weights start at 0.")

    print(f"Setting initial state: adapters={ALL_ADAPTER_NAMES}, weights={initial_weights}")
    pipe.set_adapters(ALL_ADAPTER_NAMES, adapter_weights=initial_weights)
else:
    print("No LoRAs were loaded.")

print("Initialization complete. Gradio is starting...")

@spaces.GPU()
def generate(prompt, negative_prompt, width, height, num_inference_steps, optional_lora_id, progress=gr.Progress(track_tqdm=True)):
    
    # --- Step 1: ALWAYS build the full weight vector from scratch for THIS run ---
    
    # Start with the default state: base LoRA on, others off.
    adapter_weights = [0.0] * len(ALL_ADAPTER_NAMES)
    try:
        base_lora_index = ALL_ADAPTER_NAMES.index(CAUSVID_NAME)
        adapter_weights[base_lora_index] = 1.0
    except ValueError:
        pass # Base lora was not loaded, so its weight remains 0.

    # If an optional LoRA is selected, turn its weight on.
    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:
            try:
                optional_lora_index = ALL_ADAPTER_NAMES.index(internal_name_to_add)
                adapter_weights[optional_lora_index] = 1.0
            except ValueError:
                print(f"Warning: Could not find index for selected LoRA '{internal_name_to_add}'. It will not be applied.")

    # --- Step 2: Apply the calculated state, OVERWRITING any previous state ---
    # We always pass the FULL list of adapters, just with different weights.
    print(f"Setting weights for this run: {list(zip(ALL_ADAPTER_NAMES, adapter_weights))}")
    pipe.set_adapters(ALL_ADAPTER_NAMES, adapter_weights=adapter_weights)

    apply_cache_on_pipe(pipe)
    
    # --- Step 3: Run inference ---
    output = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        height=height,
        width=width,
        num_frames=1,
        num_inference_steps=num_inference_steps,
        guidance_scale=1.0, 
    )
    image = output.frames[0][0]
    image = (image * 255).astype(np.uint8)
    return Image.fromarray(image)


# --- 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="Optional LoRA",
        )
    ],
    outputs=gr.Image(label="output"),
)

iface.launch(debug=True)