import tempfile, spaces, torch, gradio as gr
from diffusers import WanImageToVideoPipeline, AutoencoderKLWan
from diffusers.utils import export_to_video, load_image
from PIL import Image

MODEL = "Wan-AI/Wan2.2-TI2V-5B-Diffusers"
# Load at module startup (CPU) so the 10GB download/load happens at Space boot,
# NOT inside the GPU-time budget. ZeroGPU only gives GPU inside @spaces.GPU.
vae = AutoencoderKLWan.from_pretrained(MODEL, subfolder="vae", torch_dtype=torch.float32)
PIPE = WanImageToVideoPipeline.from_pretrained(MODEL, vae=vae, torch_dtype=torch.bfloat16)

def _r32(x): return max(32, (int(x)//32)*32)

@spaces.GPU(duration=200)
def animate(image, prompt, num_frames=49, steps=25, max_side=704):
    PIPE.to("cuda")
    img = (load_image(image) if isinstance(image, str) else image).convert("RGB")
    w, h = img.size; s = float(max_side)/max(w, h)
    W, H = _r32(w*s), _r32(h*s)
    img = img.resize((W, H), Image.LANCZOS)
    nf = int(num_frames); nf = nf - ((nf-1) % 4)
    frames = PIPE(image=img, prompt=prompt or "natural realistic motion, cinematic",
                  negative_prompt="static, still, frozen, deformed, distorted, jittery, flickering, low quality",
                  height=H, width=W, num_frames=nf, guidance_scale=5.0,
                  num_inference_steps=int(steps)).frames[0]
    path = tempfile.mktemp(suffix=".mp4")
    export_to_video(frames, path, fps=16)
    return path

gr.Interface(
    fn=animate,
    inputs=[gr.Image(type="filepath", label="Photo"), gr.Text(label="Motion prompt"),
            gr.Slider(17, 81, value=49, step=4, label="frames"),
            gr.Slider(15, 40, value=25, step=1, label="steps"),
            gr.Slider(448, 832, value=704, step=32, label="max_side")],
    outputs=gr.Video(label="Video"), title="Nova i2v — Wan 2.2 full-body image-to-video",
    api_name="animate",
).queue(max_size=8).launch()