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 ---
# 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"
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" 
# 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}")


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,
    )
    
    try:
        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)
    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()


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
    ],
    outputs=gr.Image(label="output"),
    title="Wan 2.1 Generator with Two Hard-coded LoRAs",
)

iface.launch()