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
import gc

# --- INITIAL SETUP ---
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")

model_id = "Wan-AI/Wan2.1-T2V-14B-Diffusers"

# --- PRE-FUSION STRATEGY ---
# 1. Load everything on the CPU first. Do NOT use device_map or torch_dtype here yet.
#    This prevents ZeroGPU/accelerate from taking over too early.
print("Loading all components on CPU for pre-fusion...")
vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae")
pipe = WanPipeline.from_pretrained(model_id, vae=vae)

# 2. Load and FUSE the LoRA while the entire pipeline is still on the CPU.
print("Loading and fusing base LoRA on CPU...")
CAUSVID_LORA_REPO = "Kijai/WanVideo_comfy"
CAUSVID_LORA_FILENAME = "Wan21_CausVid_14B_T2V_lora_rank32.safetensors"
try:
    causvid_path = hf_hub_download(repo_id=CAUSVID_LORA_REPO, filename=CAUSVID_LORA_FILENAME)
    pipe.load_lora_weights(causvid_path)
    print("✅ Base LoRA loaded on CPU.")
    
    pipe.fuse_lora()
    print("✅ Base LoRA fused on CPU.")
    
    # After fusing, we must unload the original LoRA weights to save memory
    pipe.unload_lora_weights()
    print("✅ Original LoRA weights unloaded.")
    
except Exception as e:
    print(f"⚠️ Could not pre-fuse the LoRA: {e}. The model will run without it.")

# 3. NOW, move the final, fused model to the target device and set the dtype.
#    ZeroGPU will now take over a single, coherent model.
print("Moving the final, fused pipeline to device and setting dtype...")
pipe.to(device=device, dtype=torch.bfloat16)

# 4. Configure the scheduler
flow_shift = 1.0
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=flow_shift)

# Clean up memory before starting the server
del vae
gc.collect()
torch.cuda.empty_cache()

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

@spaces.GPU()
def generate(prompt, negative_prompt, width=1024, height=1024, num_inference_steps=30, progress=gr.Progress(track_tqdm=True)):
    # The LoRA is permanently fused. No dynamic loading is possible or needed.
    
    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:
        # No cleanup needed as the state is static
        pass

# Interface is simplified as the custom LoRA option is removed for stability.
iface = gr.Interface(
    fn=generate,
    inputs=[
        gr.Textbox(label="Input prompt"),
        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),
    ],
    outputs=gr.Image(label="output"),
)

iface.launch()