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

# Tải model từ Hugging Face
pipe = StableVideoDiffusionPipeline.from_pretrained(
    "stabilityai/stable-video-diffusion-img2vid-xt",
    torch_dtype=torch.float16
)

# Đảm bảo chạy trên GPU nếu có
device = "cuda" if torch.cuda.is_available() else "cpu"
pipe.to(device)

# Hàm xử lý đầu vào và tạo video
def generate_video(input_image):
    if input_image.mode != "RGB":
        input_image = input_image.convert("RGB")
    result = pipe(image=input_image)
    video_frames = result.frames  # hoặc result.videos tùy pipeline version
    return video_frames

# Tạo giao diện Gradio
interface = gr.Interface(
    fn=generate_video,
    inputs=gr.Image(type="pil", label="Upload an image"),
    outputs=gr.Video(label="Generated Video"),
    title="Stable Video Diffusion (Img2Vid XT)",
    description="Upload an image to generate a video using stable-video-diffusion-img2vid-xt"
)

interface.launch()