import torch
from diffusers import MochiPipeline
import gradio as gr

# Carica il modello MochiPipeline
pipe = MochiPipeline.from_pretrained("genmo/mochi-1-preview")

# Funzione per generare video (usa meno frame per risparmiare memoria)
def generate_video(prompt):
    with torch.autocast("cuda", torch.float16):
        frames = pipe(prompt, num_frames=30).frames  # Riduci i frame a 30
    frames[0].save("mochi.mp4", format="mp4", save_all=True, duration=100)
    return "mochi.mp4"

# Configura l'interfaccia Gradio
interface = gr.Interface(
    fn=generate_video,
    inputs="text",
    outputs="video",
    title="Video Generator",
    description="Genera un video breve basato sul prompt fornito."
)

# Avvia l'interfaccia
interface.launch()