import torch
from PIL import Image
import imageio
from diffusers import StableVideoDiffusionPipeline
from diffusers.utils import load_image, export_to_video
import gradio as gr
import spaces

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

@spaces.GPU(duration=300)
def generate_video(image, seed=42, fps=7, motion_bucket_id=180, noise_aug_strength=0.1):
    # Resize the image
    image = image.resize((1024, 576))

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

    # Generate the frames
    frames = pipe(image, decode_chunk_size=2, generator=generator, num_frames=25, motion_bucket_id=motion_bucket_id, noise_aug_strength=noise_aug_strength).frames[0]

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

    return output_path

# Create the Gradio interface
iface = gr.Interface(
    fn=generate_video,
    inputs=[
        gr.Image(type="pil", label="Upload Image"),
        gr.Number(label="Seed", value=42),
        gr.Number(label="FPS", value=7),
        gr.Number(label="Motion Bucket ID", value=180),
        gr.Number(label="Noise Aug Strength", value=0.1)
    ],
    outputs=gr.Video(label="Generated Video"),
    title="Stable Video Diffusion",
    description="Generate a video from an uploaded image using Stable Video Diffusion."
)

# Launch the interface
iface.launch()