import gradio as gr
import torch
from diffusers import AnimateDiffPipeline
from PIL import Image
import tempfile
import os

pipe = AnimateDiffPipeline.from_pretrained(
    "ByteDance/AnimateDiff-Lightning",
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
)
pipe.to("cuda" if torch.cuda.is_available() else "cpu")

def generate_video(prompt):
    frames = pipe(prompt).frames

    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as temp_video:
        pipe.export_to_video(frames, temp_video.name)
        return temp_video.name

demo = gr.Interface(
    fn=generate_video,
    inputs=gr.Textbox(label="Prompt"),
    outputs=gr.Video(label="Generated Animation"),
    title="AnimateDiff Demo",
    description="Generate animations from text prompts using AnimateDiff."
)

demo.launch()