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
import config as cfg

# Load the pre-trained model
pipe = MochiPipeline.from_pretrained(cfg.MODEL_PRE_TRAINED_ID, variant="bf16", torch_dtype=torch.bfloat16)

# Enable memory-saving optimizations
pipe.enable_model_cpu_offload()
pipe.enable_vae_tiling()

@spaces.GPU(duration=600)
def generate_video(prompt, num_frames=84, fps=30):
    # Generate video frames
    print("Generating video frames...")
    frames = pipe(prompt, num_frames=num_frames).frames[0]

    # Export frames as video
    video_path = "mochi.mp4"
    export_to_video(frames, video_path, fps=fps)
    return video_path

# Create the Gradio interface
interface = gr.Interface(
    fn=generate_video,
    inputs=[
        gr.Textbox(lines=2, placeholder="Enter your text prompt here... 💡"),
        gr.Slider(minimum=1, maximum=240, value=84, label="Number of frames 🎞️"),
        gr.Slider(minimum=1, maximum=60, value=30, label="FPS (Frames per second) ⏱️")
    ],
    outputs=gr.outputs.Video(),
    title=cfg.TITLE,
    description=cfg.DESCRIPTION,
    examples=cfg.EXAMPLES,
    article=cfg.BUY_ME_A_COFFEE

)

# Launch the application
if __name__ == "__main__":
    interface.launch()