import torch
import gradio as gr
from diffusers import DiffusionPipeline

# Load the pipeline from Hugging Face's model hub
pipe = DiffusionPipeline.from_pretrained("THUDM/CogVideoX-2b")

# Optional: Move the model to GPU for faster processing if available
if torch.cuda.is_available():
    pipe.to("cuda")

# Define the function that generates the video from the user prompt
def generate_video(prompt):
    # Generate the video
    video = pipe(prompt).videos[0]
    # Save the video to a file (Gradio needs a file path to display video)
    video_path = "generated_video.mp4"
    video.save(video_path)
    return video_path

# Create a Gradio interface with text input for the prompt and video output
interface = gr.Interface(
    fn=generate_video,           # The function to generate the video
    inputs="text",               # Text input for the prompt
    outputs="video",             # Video output
    title="Video Generator",     # Title of the Gradio app
    description="Enter a prompt to generate a video using the diffusion model"
)

# Launch the Gradio app
interface.launch()