import os
import torch
import gradio as gr
from diffusers import AutoencoderKLWan, WanPipeline
from diffusers.utils import export_to_video

# ---- MODEL SETUP ----

# Model ID on Hugging Face
MODEL_ID = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"

# Load the pipeline and VAE
def load_model():
    print("Loading model... (this may take a while)")
    vae = AutoencoderKLWan.from_pretrained(MODEL_ID, subfolder="vae", torch_dtype=torch.float32)
    pipe = WanPipeline.from_pretrained(MODEL_ID, vae=vae, torch_dtype=torch.float16)
    pipe = pipe.to("cuda")  # Ensure using GPU
    return pipe

video_pipe = load_model()

# ---- GENERATION FUNCTION ----

def generate_video(prompt):
    try:
        # Text to video generation
        result = video_pipe(
            prompt=prompt,
            num_frames=24,
            guidance_scale=7.5
        )

        # Convert list of PIL frames to MP4
        frames = result.frames
        output_path = "generated_video.mp4"
        export_to_video(frames, output_path, fps=8)

        return output_path

    except Exception as e:
        return f"Error: {str(e)}"

# ---- GRADIO UI ----

with gr.Blocks() as app:
    gr.Markdown("# 🧠 Text‑to‑Video with Wan2.1‑T2V")
    prompt = gr.Textbox(label="Enter your video prompt", placeholder="e.g. A dragon flying over mountains")
    btn = gr.Button("Generate Video")
    video_output = gr.Video()

    btn.click(fn=generate_video, inputs=prompt, outputs=video_output)

app.launch()