app.py

# IMPORTANT: spaces must be imported first to avoid CUDA initialization issues
import spaces

import os
import numpy as np
from PIL import Image
import gradio as gr

import torch
from diffusers import WanPipeline, AutoencoderKLWan
from diffusers.utils import export_to_video

# ────────────────────────────────────────────────
#  Model + LoRA configuration
# ────────────────────────────────────────────────

MODEL_ID = "Wan-AI/Wan2.2-TI2V-5B-Diffusers"
dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"

AVAILABLE_LORAS = [
    {
        "name": "Lightning (Fast 4-step)",
        "repo_id": "lightx2v/Wan2.2-Distill-Loras",
        "filename": "wan2.2_i2v_A14b_high_noise_lora_rank64_lightx2v_4step_1022.safetensors",
        "default_strength": 1.0,
    },
    {
        "name": "General NSFW",
        "repo_id": "lopi999/Wan2.2-I2V_General-NSFW-LoRA",
        "filename": "pytorch_lora_weights.safetensors",
        "default_strength": 0.8,
    },
    # Add more LoRAs here — they will be pre-loaded automatically
]

# Global pipeline + pre-loaded adapter info
pipe = None
lora_adapters = {}          # name → {"adapter_name": str, "strength": float}

def initialize_pipeline():
    global pipe, lora_adapters

    if pipe is not None:
        return pipe

    print("Loading Wan2.2-TI2V-5B base model...")
    vae = AutoencoderKLWan.from_pretrained(
        MODEL_ID,
        subfolder="vae",
        torch_dtype=torch.float32
    )
    pipe = WanPipeline.from_pretrained(
        MODEL_ID,
        vae=vae,
        torch_dtype=dtype
    )
    pipe.to(device)
    print("Base model loaded.")

    # ── Pre-load ALL available LoRAs once ───────────────
    print("Pre-loading LoRAs...")
    for lora in AVAILABLE_LORAS:
        name = lora["name"]
        try:
            print(f"  → {name}")
            pipe.load_lora_weights(
                lora["repo_id"],
                weight_name=lora["filename"],
                adapter_name=name,           # unique identifier
            )
            # Store for later hot-swapping
            lora_adapters[name] = {
                "adapter_name": name,
                "strength": lora["default_strength"]
            }
        except Exception as e:
            print(f"    Failed to load {name}: {e}")

    # Fuse once → best inference performance
    if lora_adapters:
        pipe.fuse_lora()
        print("All LoRAs fused.")

    print("Pipeline fully initialized.")
    return pipe

@spaces.GPU(duration=180)
def generate_video(
    prompt: str,
    image: Image.Image = None,
    width: int = 1280,
    height: int = 704,
    num_frames: int = 73,
    num_inference_steps: int = 35,
    guidance_scale: float = 5.0,
    seed: int = -1,
    enabled_loras: list = None,
    lora_strength_multiplier: float = 1.0,
    progress=gr.Progress()
):
    try:

        pipeline = initialize_pipeline()

        # ── Hot-swap / enable only selected LoRAs ───────
        active_adapters = []
        active_strengths = []

        enabled = enabled_loras or []

        for lora_name in enabled:
            if lora_name in lora_adapters:
                strength = lora_adapters[lora_name]["strength"] * lora_strength_multiplier
                # Prevent values below 1.0 - many diffusers versions reject them
                strength = max(1.0, strength)
                active_adapters.append(lora_name)
                active_strengths.append(strength)

        # ── Only call set_adapters when we actually want to use LoRAs ────────────────
        if active_adapters:
            print(f"Activating LoRAs: {active_adapters} with strengths {active_strengths}")
            pipeline.set_adapters(active_adapters, adapter_strengths=active_strengths)
        else:
            print("No LoRAs enabled → disabling LoRA")
            # Safely disable LoRA - some pipelines throw on empty set_adapters([])
            try:
                pipeline.disable_lora()
            except Exception:
                # If disable_lora doesn't exist or fails, just continue
                # (some custom pipelines don't have it)
                pass
        
        # Lightning auto-optimization
        if "Lightning (Fast 4-step)" in enabled and num_inference_steps > 8:
            num_inference_steps = 4
            print("Lightning LoRA → reduced to 4 steps")

        # Seed
        if seed == -1:
            seed = torch.randint(0, 2**32 - 1, (1,)).item()
        generator = torch.Generator(device=device).manual_seed(seed)

        gen_params = {
            "prompt": prompt,
            "height": height,
            "width": width,
            "num_frames": num_frames,
            "guidance_scale": guidance_scale,
            "num_inference_steps": num_inference_steps,
            "generator": generator,
        }

        if image is not None:
            gen_params["image"] = image

        print(f"Generating: {width}x{height}, {num_frames} frames, steps={num_inference_steps}")
        #output = pipeline(**gen_params).frames[0]

        progress(0, desc="Starting generation...")

        def step_callback(step, timestep, latents):
            progress(step / num_inference_steps, desc=f"Step {step+1}/{num_inference_steps}")
        
        output = pipeline(
            **gen_params,
            callback=step_callback,
            callback_steps=1,
        ).frames[0] 

        output_path = "output.mp4"
        export_to_video(output, output_path, fps=24)

        status = f"Done! Seed: {seed}"
        if active_adapters:
            status += f"\nLoRAs: {', '.join(active_adapters)} @ {lora_strength_multiplier:.2f}x"

        return output_path, status

    except Exception as e:
        msg = f"Error: {str(e)}"
        print(msg)
        return None, msg

# ────────────────────────────────────────────────
#               Gradio UI
# ────────────────────────────────────────────────

with gr.Blocks(title="Wan2.2 Video + Fast LoRA") as demo:
    gr.Markdown("""
    # Wan2.2-TI2V-5B Video Generation  
    **Optimized LoRA loading** — all LoRAs pre-loaded at startup, then hot-swapped instantly.
    """)

    with gr.Row():
        with gr.Column():
            prompt_input = gr.Textbox(
                label="Prompt", lines=3,
                value="Two anthropomorphic cats in comfy boxing gear fight on stage"
            )
            image_input = gr.Image(label="Input Image (optional)", type="pil", sources=["upload"])

            with gr.Accordion("Advanced Settings", open=False):
                with gr.Row():
                    width_input  = gr.Slider(512, 1920, step=64, value=1280, label="Width")
                    height_input = gr.Slider(512, 1080, step=64, value=704, label="Height")
                num_frames_input = gr.Slider(25, 145, step=24, value=73, label="Frames")
                num_steps_input  = gr.Slider(4, 60, step=1, value=4, label="Inference Steps",
                                             info="Lightning LoRA → try 4–8 steps")
                guidance_scale_input = gr.Slider(1.0, 15.0, 1.0, value=5.0, label="Guidance Scale")
                seed_input = gr.Number(label="Seed (-1 = random)", value=-1, precision=0)

            with gr.Accordion("LoRA Controls", open=True):
                lora_checkbox = gr.CheckboxGroup(
                    choices=[l["name"] for l in AVAILABLE_LORAS],
                    label="Enable LoRAs",
                    value=[]
                )
                lora_strength = gr.Slider(0.1, 1.5, 0.05, value=1.0,
                                          label="Global Strength Multiplier")

            generate_btn = gr.Button("Generate Video", variant="primary", size="lg")

        with gr.Column():
            video_output = gr.Video(label="Generated Video", autoplay=True)
            status_output = gr.Textbox(label="Status", lines=3)

    # Examples with LoRA usage
    gr.Examples(
        examples=[
            ["Two anthropomorphic cats in comfy boxing gear fight on stage", None, 1280, 704, 73, 35, 5.0, 42, [], 1.0],
            ["A serene underwater scene with colorful coral reefs...", None, 1280, 704, 73, 4, 5.0, 123, ["Lightning (Fast 4-step)"], 1.0],
            ["Explicit adult scene, detailed", None, 1280, 704, 73, 30, 6.0, 999, ["General NSFW"], 0.9],
        ],
        inputs=[prompt_input, image_input, width_input, height_input, num_frames_input,
                num_steps_input, guidance_scale_input, seed_input, lora_checkbox, lora_strength],
        outputs=[video_output, status_output],
        fn=generate_video,
        cache_examples=False,
    )

    generate_btn.click(
        generate_video,
        inputs=[prompt_input, image_input, width_input, height_input, num_frames_input,
                num_steps_input, guidance_scale_input, seed_input, lora_checkbox, lora_strength],
        outputs=[video_output, status_output]
    )

    gr.Markdown("""
    ## Performance Notes
    - LoRAs are **pre-loaded once** → first generation may take ~10–30s longer, later ones are fast.
    - Lightning LoRA: use **4–8 steps** → generation can finish in <60s.
    - Add new LoRAs by appending to `AVAILABLE_LORAS` — they auto-load at startup.
    """)

if __name__ == "__main__":
    demo.queue(max_size=20).launch()