import torch
import gradio as gr
from diffusers import DiffusionPipeline
from diffusers.utils import export_to_video

# Charger le modèle de diffusion sans utiliser CUDA
pipe = DiffusionPipeline.from_pretrained("damo-vilab/text-to-video-ms-1.7b", torch_dtype=torch.float16, variant="fp16")
pipe = pipe.to("cpu")  # Utiliser le CPU au lieu du GPU

# Fonction pour générer la vidéo
def generate_video(prompt):
    video_frames = pipe(prompt).frames[0]
    video_path = export_to_video(video_frames)
    return video_path

# Interface Gradio
iface = gr.Interface(
    fn=generate_video, 
    inputs=gr.Textbox(label="Enter Prompt", placeholder="e.g., Spiderman is surfing"), 
    outputs=gr.Video(label="Generated Video"),
    title="Text-to-Video Generation",
    description="Generate a video based on a textual prompt."
)

# Lancer l'interface
iface.launch()