import gradio as gr
import torch
from diffusers import DiffusionPipeline
import imageio
import tempfile

# Load the lighter 5B model
pipe = DiffusionPipeline.from_pretrained(
    "Wan-AI/Wan2.2-TI2V-5B-Diffusers",
    torch_dtype=torch.float16
).to("cuda")

def generate_video(prompt, steps=25):
    video = pipe(prompt, num_inference_steps=steps).videos[0]
    tmpfile = tempfile.NamedTemporaryFile(suffix=".gif", delete=False)
    imageio.mimsave(tmpfile.name, video, fps=8)
    return tmpfile.name

demo = gr.Interface(
    fn=generate_video,
    inputs=[gr.Textbox(label="Prompt"), gr.Slider(1, 50, value=25, label="Steps")],
    outputs=gr.Video(label="Generated Video")
)

if __name__ == "__main__":
    demo.launch()