import gradio as gr
from diffusers import CogVideoXPipeline
import torch

model_name = "THUDM/CogVideoX-5b"
device = "cuda" if torch.cuda.is_available() else "cpu"

if device == "cuda":
    pipe = CogVideoXPipeline.from_pretrained(model_name,torch_dtype=torch.float16).to(device)
else:
    pipe = CogVideoXPipeline.from_pretrained(model_name).to(device)

def generate_video(prompt):
    video = pipe(prompt, num_inference_steps=50, guidance_scale=9.5).videos

    video_path = "generated_video.mp4"

    video[0].save(video_path)

    return video_path

iface = gr.Interface(
    fn=generate_video,
    inputs="text",
    outputs="video",
    title="BL's T2V Generator"
)

iface.launch()