import gradio as gr
import torch
from diffusers import StableVideoDiffusionPipeline
from PIL import Image

pipeline = StableVideoDiffusionPipeline.from_pretrained(
    "stabilityai/stable-video-diffusion-img2vid-xt",
    torch_dtype=torch.float16,
    variant="fp16"
)
pipeline.to("cuda")
pipeline.unet = torch.compile(pipeline.unet, mode="reduce-overhead", fullgraph=True)
pipeline.enable_model_cpu_offload()

def generate(image, seed, motion_bucket_id, fps):
    if image is None:
        return None
    image = image.convert("RGB")
    image = image.resize((1024, 576))
    generator = torch.manual_seed(int(seed))
    frames = pipeline(
        image,
        motion_bucket_id=int(motion_bucket_id),
        fps=int(fps),
        generator=generator
    ).frames[0]
    return frames

with gr.Blocks(title="Stable Video Diffusion") as demo:
    gr.Markdown("# Stable Video Diffusion\nUpload an image to generate a short video (3-4 seconds).")
    with gr.Row():
        with gr.Column():
            input_img = gr.Image(type="pil", label="Input Image")
            seed = gr.Number(value=42, label="Seed")
            motion_bucket = gr.Slider(1, 255, value=127, step=1, label="Motion Bucket ID")
            fps = gr.Slider(5, 30, value=6, step=1, label="FPS")
            btn = gr.Button("Generate Video")
        with gr.Column():
            output_vid = gr.Video(label="Generated Video")
    btn.click(fn=generate, inputs=[input_img, seed, motion_bucket, fps], outputs=output_vid)

demo.launch()