app.py

import torch
import gradio as gr
from diffusers import StableVideoDiffusionPipeline
from diffusers.utils import load_image, export_to_video
import os
import spaces

# Check if GPU is available
device = "cuda" if torch.cuda.is_available() else "cpu"

# Load the pipeline with optimizations
pipeline = StableVideoDiffusionPipeline.from_pretrained(
    "stabilityai/stable-video-diffusion-img2vid-xt", torch_dtype=torch.float16, variant="fp16"
)
pipeline.to(device)

# Enable forward chunking to reduce memory usage
pipeline.unet.enable_forward_chunking(chunk_size=1)

@spaces.GPU
# Define the video generation function
def generate_video(image_path, seed):
    # Load and preprocess the image
    image = load_image(image_path)
    image = image.resize((1024, 576))

    # Set the generator seed
    generator = torch.Generator(device=device).manual_seed(seed)

    # Generate the video frames
    frames = pipeline(image, decode_chunk_size=8, generator=generator).frames[0]

    # Export the frames to a video file
    output_video_path = "generated.mp4"
    export_to_video(frames, output_video_path, fps=25)

    return output_video_path

# Create the Gradio interface
iface = gr.Interface(
    fn=generate_video,
    inputs=[
        gr.Image(type="filepath", label="Upload Image"),
        gr.Number(label="Seed", value=666666)
    ],
    outputs=gr.Video(label="Generated Video"),
    title="Stable Video Diffusion",
    description="Generate a video from an uploaded image using Stable Video Diffusion.",
    examples=[
        ["example.jpeg", 666666]
    ],
    cache_examples=True  # Enable caching of examples
)

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