import gradio as gr
from diffusers import StableVideoDiffusionPipeline, EulerDiscreteScheduler
import torch

# Load the Stable Video Diffusion model
model_id = "stabilityai/stable-video-diffusion-img2vid-xt"
pipe = StableVideoDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16, revision="main")
pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config)
pipe.to("cuda")

def generate_video(image, prompt, num_frames=25, resolution=(576, 1024)):
    # Generate the video
    video_frames = pipe(prompt, image=image, num_frames=num_frames, height=resolution[0], width=resolution[1]).frames
    return video_frames

# Create the Gradio interface
iface = gr.Interface(
    fn=generate_video,
    inputs=[
        gr.inputs.Image(type="pil", label="Upload Image"),
        gr.inputs.Textbox(lines=2, placeholder="Enter prompt...", label="Prompt"),
        gr.inputs.Slider(1, 50, step=1, default=25, label="Number of Frames"),
        gr.inputs.Number(label="Resolution Height", default=576),
        gr.inputs.Number(label="Resolution Width", default=1024)
    ],
    outputs=gr.outputs.Video(label="Generated Video"),
    title="Image to Video with Stable Diffusion XT",
    description="Upload an image and enter a prompt to generate a video."
)

# Launch the interface
if __name__ == "__main__":
    iface.launch()