import gradio as gr
import torch
from diffusers import DiffusionPipeline

# 🟢 Load model WAN 2.2 Video Generation dari Hugging Face Hub
# Pastikan ID model benar sesuai rilis resmi (misal: "BAAI/Wan2.2-videogen")
model_id = "BAAI/Wan2.2-videogen"

pipe = DiffusionPipeline.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    variant="fp16"
).to("cuda")

# Fungsi generate video
def generate_video(prompt, num_frames=24, fps=8, seed=42, progress=gr.Progress(track_tqdm=True)):
    generator = torch.manual_seed(seed)

    progress(0, desc="🚀 Mulai generate video...")

    # Jalankan pipeline untuk generate frame
    output = pipe(
        prompt=prompt,
        num_frames=num_frames,
        generator=generator
    )

    progress(0.7, desc="📸 Menyusun frame menjadi video...")

    # Konversi output frames → video format
    video_frames = output.frames[0]
    out_path = "output.mp4"

    # Simpan video mp4 ke file
    pipe.save_pretrained_video(video_frames, out_path, fps=fps)

    progress(1, desc="✅ Selesai!")

    return out_path, out_path   # 🎯 untuk Video preview + File download

# UI pakai Gradio Blocks
with gr.Blocks() as demo:
    gr.Markdown("## 🎬 WAN 2.2 Video Generator (Hugging Face Space)")

    with gr.Row():
        prompt_inp = gr.Textbox(
            label="Prompt",
            value="Seekor naga robot terbang melintasi kota futuristik",
            placeholder="Deskripsi video..."
        )

    with gr.Row():
        frame_slider = gr.Slider(8, 64, step=8, value=24, label="Jumlah Frame")
        fps_slider = gr.Slider(4, 16, step=1, value=8, label="FPS (kecepatan video)")

    btn = gr.Button("🚀 Generate Video")

    with gr.Row():
        video_out = gr.Video(label="Hasil Video")
        download_link = gr.File(label="Unduh Video", type="file")

    btn.click(
        fn=generate_video,
        inputs=[prompt_inp, frame_slider, fps_slider],
        outputs=[video_out, download_link]
    )

if __name__ == "__main__":
    demo.queue(max_size=5)  # antrian agar multi-user tetap stabil
    demo.launch()