import os
os.system('pip install --upgrade --pre --extra-index-url https://download.pytorch.org/whl/nightly/cu126 "torch<2.9"')
import spaces
import os
import gradio as gr
import torch
from diffusers import WanPipeline, AutoencoderKLWan
from diffusers.utils import export_to_video

# Model setup
dtype = torch.bfloat16
device = "cuda"
model_id = "FastDM/Wan2.2-T2V-A14B-Merge-Lightning-V1.0-Diffusers"

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=dtype
).to(device)

# Default values
DEFAULT_PROMPT = "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage."
DEFAULT_NEGATIVE = "bad quality, blurry, distorted, extra limbs, watermark, text"

@spaces.GPU(duration=60)
def generate_video(prompt, negative_prompt, height, width, num_frames, steps, guidance):
    video = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        height=height,
        width=width,
        num_frames=num_frames,
        guidance_scale=guidance,
        num_inference_steps=steps,
    ).frames[0]

    output_path = "t2v_out.mp4"
    export_to_video(video, output_path, fps=16)
    return output_path


with gr.Blocks() as demo:
    gr.Markdown("# 🎬 Wan2.2 Text-to-Video Demo")
    gr.Markdown("Generate short AI videos from text prompts.")

    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Prompt", value=DEFAULT_PROMPT, lines=3)
            negative_prompt = gr.Textbox(label="Negative Prompt", value=DEFAULT_NEGATIVE, lines=2)
            height = gr.Slider(256, 1280, value=720, step=64, label="Height")
            width = gr.Slider(256, 1280, value=1280, step=64, label="Width")
            num_frames = gr.Slider(16, 128, value=81, step=1, label="Number of Frames")
            steps = gr.Slider(1, 20, value=4, step=1, label="Inference Steps")
            guidance = gr.Slider(0.1, 10.0, value=1.0, step=0.1, label="Guidance Scale")
            generate_btn = gr.Button("🚀 Generate Video")

        with gr.Column():
            video_output = gr.Video(label="Generated Video")

    generate_btn.click(
        fn=generate_video,
        inputs=[prompt, negative_prompt, height, width, num_frames, steps, guidance],
        outputs=[video_output],
    )

demo.launch()