import gradio as gr
import torch
from diffusers import AutoencoderKLWan, WanPipeline
from diffusers.utils import export_to_video
import spaces  # ZeroGPU integration

@spaces.GPU  # This decorator requests a GPU when the function is called (ZeroGPU config)
def generate_video(prompt, negative_prompt=""):
    model_id = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
    
    # Load the VAE and pipeline with proper data types per model requirements
    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.to("cuda")
    
    # Generate video frames at 480p resolution (480x832) with the desired settings
    output = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        height=480,        # 480p height
        width=832,         # width set to match 480p (adjust as needed)
        num_frames=81,     # Number of frames (adjust for desired video length)
        guidance_scale=5.0 # Recommended for the 1.3B model
    ).frames[0]
    
    # Save the generated frames as a video file
    video_path = "output.mp4"
    export_to_video(output, video_path, fps=15)
    return video_path

# Create a Gradio interface for the video generation
iface = gr.Interface(
    fn=generate_video,
    inputs=[
        gr.Textbox(label="Prompt", placeholder="Enter your video prompt here"),
        gr.Textbox(label="Negative Prompt", placeholder="Optional negative prompt", value="")
    ],
    outputs=gr.Video(label="Generated Video"),
    title="Wan2.1-T2V-1.3B Video Generator",
    description="Generate 480p videos using the Wan2.1-T2V-1.3B diffusers pipeline with ZeroGPU support."
)

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