import gradio as gr
import torch
from diffusers import DiffusionPipeline

# Load the model and move to GPU if available
model_name = "ByteDance/AnimateDiff-Lightning"
device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = DiffusionPipeline.from_pretrained(model_name)
pipe = pipe.to(device)

def generate_video(prompt):
    # Generate video (replace with actual method for your model)
    result = pipe(prompt)
    # Assuming the model returns a list of videos or a video object
    # Adjust this based on the actual output
    video = result.videos[0]  # or result['videos'][0]
    # Save video to file
    video_path = "generated_video.mp4"
    video.save(video_path)
    return video_path

with gr.Blocks() as demo:
    gr.Markdown("# Text to Video Generation")
    prompt_input = gr.Textbox(label="Enter your prompt", lines=2, placeholder="A spaceship in space, neon colors")
    generate_button = gr.Button("Generate Video")
    video_output = gr.Video(label="Generated Video")
    
    generate_button.click(
        fn=generate_video,
        inputs=prompt_input,
        outputs=video_output
    )

demo.launch()