import os
if os.environ.get("SPACES_ZERO_GPU") is not None:
    import spaces
else:
    class spaces:
        @staticmethod
        def GPU(func):
            def wrapper(*args, **kwargs):
                return func(*args, **kwargs)
            return wrapper

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

# Caricare il modello pre-addestrato
pipe = MochiPipeline.from_pretrained("genmo/mochi-1-preview", variant="bf16", torch_dtype=torch.bfloat16)

# Abilitare le ottimizzazioni per il risparmio di memoria
pipe.enable_model_cpu_offload()
pipe.enable_vae_tiling()


@spaces.GPU(duration=80)
def generate_video(prompt):
    # Generare i frame del video
    frames = pipe(prompt, num_frames=84).frames[0]

    # Esportare i frame come video
    video_path = "mochi.mp4"
    export_to_video(frames, video_path, fps=30)

    return video_path


# Creare l'interfaccia Gradio
interface = gr.Interface(
    fn=generate_video,
    inputs="text",
    outputs="video",
    title="Mochi Video Generator",
    description="Genera un video basato su un prompt di testo utilizzando MochiPipeline."
)

# Avviare l'applicazione
if __name__ == "__main__":
    interface.launch()