import gradio as gr
from transformers import pipeline
from spaces import GPU  # Import the GPU decorator for ZeroGPU

# Use @GPU to allocate a GPU when running the video generation function
@GPU
def generate_video(prompt, image):
    # Load the Kandinsky video model
    video_model = pipeline("video-generation", model="ai-forever/KandinskyVideo_1_1", device=0)
    
    # Generate video based on the prompt and image
    video = video_model(prompt=prompt, init_image=image)
    return video  # Ensure this returns the video file or frames

# Create the Gradio interface with text and image inputs
interface = gr.Interface(
    fn=generate_video,
    inputs=[
        gr.Textbox(label="Enter your prompt here"),
        gr.Image(label="Upload an initial image")  # Image upload input
    ],
    outputs="video",
    title="AI Video Generator",
    description="This app generates a short video based on your input prompt and initial image using Kandinsky 1.1.",
    theme="dark"
)

if __name__ == "__main__":
    interface.launch()