import gradio as gr
import torch
import subprocess
import os

CKPT = "Wan-AI/Wan2.2-S2V-14B"

def setup_model():
    if not os.path.exists("model"):
        subprocess.run(["huggingface-cli", "download", CKPT, "--local-dir", "model"])
setup_model()

def generate(prompt, image, audio):
    cmd = [
        "python3", "generate.py",
        "--task", "s2v-14B",
        "--size", "1024*704",
        "--ckpt_dir", "model",
        "--offload_model", "True",
        "--prompt", prompt,
        "--audio", audio
    ]

    if image is not None:
        cmd += ["--image", image]

    subprocess.run(cmd)
    return "output.mp4"

with gr.Blocks() as demo:
    gr.Markdown("# 🎬 Wan2.2 S2V-14B — Online Speech-to-Video Generator")

    prompt = gr.Textbox(label="Prompt")
    image = gr.Image(type="filepath", label="Optional Reference Image")
    audio = gr.Audio(type="filepath", label="Audio (speech)")
    output = gr.Video(label="Generated Video")

    btn = gr.Button("Generate")

    btn.click(generate, inputs=[prompt, image, audio], outputs=output)

demo.launch()