import gradio as gr
import torch
from diffusers.utils import export_to_video
from diffusers import AutoencoderKLWan, WanPipeline
from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
import os
from uuid import uuid4

# Load model on startup
model_id = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32)
scheduler = UniPCMultistepScheduler(
    prediction_type='flow_prediction',
    use_flow_sigmas=True,
    num_train_timesteps=1000,
    flow_shift=5.0
)
pipe = WanPipeline.from_pretrained(model_id, vae=vae, torch_dtype=torch.bfloat16)
pipe.scheduler = scheduler
pipe.to("cuda")

# Define the generation function
def generate_video(prompt, negative_prompt="", height=720, width=1280, num_frames=81, guidance_scale=5.0):
    output = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        height=height,
        width=width,
        num_frames=num_frames,
        guidance_scale=guidance_scale,
    ).frames[0]

    output_filename = f"{uuid4()}.mp4"
    output_path = os.path.join("outputs", output_filename)
    os.makedirs("outputs", exist_ok=True)
    export_to_video(output, output_path, fps=16)

    return output_path  # Gradio returns this as downloadable file/video

# Gradio Interface with API support
iface = gr.Interface(
    fn=generate_video,
    inputs=[
        gr.Textbox(label="Prompt"),
        gr.Textbox(label="Negative Prompt", value=""),
        gr.Number(label="Height", value=720),
        gr.Number(label="Width", value=1280),
        gr.Number(label="Number of Frames", value=81),
        gr.Number(label="Guidance Scale", value=5.0)
    ],
    outputs=gr.File(label="Generated Video"),
    title="Wan2.1 Video Generator",
    description="Generate realistic videos from text prompts using the Wan2.1 T2V model.",
    api=True  # This enables the API
)

iface.launch(share=True)  # `share=True` will allow others to access your app via a public link