import os
import tempfile
import uuid

import spaces
import gradio as gr
import torch
from diffusers import AutoencoderKLWan, WanPipeline
from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler

# ── Load model on CPU at startup (ZeroGPU moves to GPU automatically) ───
MODEL_ID = "Wan-AI/Wan2.1-T2V-1.3B-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)
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=3.0)

# ── Resolution presets ──────────────────────────────────────────────────
RESOLUTIONS = {
    "9:16 (480p)": (480, 832),
    "16:9 (480p)": (832, 480),
    "1:1 (480p)": (624, 624),
}

# ── Generation function ─────────────────────────────────────────────────
@spaces.GPU(duration=180)
def generate(prompt, negative_prompt, resolution, steps, guidance, seed, progress=gr.Progress(track_tqdm=True)):
    width, height = RESOLUTIONS.get(resolution, (480, 832))

    pipe.to("cuda")

    generator = torch.Generator("cuda").manual_seed(int(seed)) if seed >= 0 else None

    output = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt or None,
        num_inference_steps=int(steps),
        guidance_scale=float(guidance),
        height=height,
        width=width,
        num_frames=81,
        generator=generator,
    )

    # Export frames to video file
    from diffusers.utils import export_to_video
    out_path = os.path.join(tempfile.gettempdir(), f"wan_{uuid.uuid4().hex[:8]}.mp4")
    export_to_video(output.frames[0], out_path, fps=16)
    return out_path


# ── Gradio UI ────────────────────────────────────────────────────────────
with gr.Blocks(title="Wan 2.1 Video Generator") as demo:
    gr.Markdown("# Wan 2.1 T2V (1.3B) — ZeroGPU")

    with gr.Row():
        with gr.Column(scale=3):
            prompt = gr.Textbox(label="Prompt", lines=3, placeholder="A cat walking on a beach at sunset...")
            negative = gr.Textbox(label="Negative prompt", value="blurry, low quality, distorted, watermark, text")
        with gr.Column(scale=1):
            resolution = gr.Dropdown(list(RESOLUTIONS.keys()), value="9:16 (480p)", label="Resolution")
            steps = gr.Slider(10, 50, value=20, step=1, label="Steps")
            guidance = gr.Slider(1.0, 10.0, value=5.0, step=0.5, label="Guidance scale")
            seed = gr.Number(value=-1, label="Seed (-1 = random)", precision=0)
            btn = gr.Button("Generate", variant="primary")

    video = gr.Video(label="Result")
    btn.click(generate, [prompt, negative, resolution, steps, guidance, seed], video, api_name="generate")

demo.launch()