# Install required packages
#!pip install accelerate torch gradio transformers git+https://github.com/huggingface/diffusers sentencepiece opencv-python

import os

# Define a fallback for environments without GPU
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 DiffusionPipeline
from diffusers.utils import export_to_video
import gradio as gr

# Application configuration
TITLE = "AI Video Generator 🌟"
DESCRIPTION = """\
🌈 Transform your imagination into stunning videos using advanced AI technology with Mochi-1-preview.\
Experience the magic of generative art! 🎥
"""
BUY_ME_A_COFFEE = """
<a href="https://buymeacoffee.com/ruslanmv" target="_blank">
    <button style="background-color: #FFDD00; border: none; color: black; 
                    padding: 10px 20px; text-align: center; 
                    text-decoration: none; display: inline-block; 
                    font-size: 16px; margin: 4px 2px; cursor: pointer; 
                    border-radius: 10px;">\
        ☕ Buy Me a Coffee
    </button>
</a>
"""
MODEL_PRE_TRAINED_ID = "genmo/mochi-1-preview"
EXAMPLES = [
    [
        "Depict a futuristic metropolis at night. Include glowing skyscrapers with holographic advertisements, flying cars zipping through the air, and bustling streets filled with humanoid robots and diverse pedestrians. Add subtle background music and neon lighting effects",
        30,
        30,
    ],
    [
        "Follow a spacecraft traveling through the galaxy. Start with a launch from Earth, then transition to the ship gliding past planets, asteroid belts, and a colorful nebula. End with the ship docking at a massive space station orbiting a distant planet",
        30,
        24,
    ],
    [
        "Present an enchanted forest at dawn. Include glowing mushrooms, magical creatures like fairies and unicorns, and trees with twinkling lights. Show a stream flowing with sparkling water and a serene, mystical atmosphere.",
        30,
        25,
    ],
    [
        "Display a bustling city from sunrise to nighttime. Include scenes of traffic, pedestrians, and iconic landmarks transitioning through different times of the day. Finish with the city skyline glowing under the night sky.",
        30,
        30,
    ],
    [
        "A vast alien desert with shimmering sands of gold and silver, \"\
        punctuated by colossal crystal spires. Twin suns set in the distance, \"\
        casting long, surreal shadows across the dunes.",
        30,
        28,
    ],
]

# Load the pre-trained model
pipe = DiffusionPipeline.from_pretrained(
    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=60 * 3)
def generate_video(prompt, num_frames=84, fps=30, high_quality=False):
    """Generate a video based on the input prompt."""
    if high_quality:
        print("High quality option selected. Requires 42GB VRAM.")
        if os.environ.get("SPACES_ZERO_GPU") is not None:
            raise RuntimeError("High quality option may fail on ZeroGPU environments.")
        with torch.autocast("cuda", torch.bfloat16, cache_enabled=False):
            frames = pipe(prompt, num_frames=num_frames).frames[0]
    else:
        print("Standard quality option selected.")
        frames = pipe(prompt, num_frames=num_frames).frames[0]

    video_path = "generated_video.mp4"
    export_to_video(frames, video_path, fps=fps)
    return video_path

# Define the Gradio interface
interface = gr.Interface(
    fn=generate_video,
    inputs=[
        gr.Textbox(lines=2, placeholder="Enter a vivid text prompt... 🔍"),
        gr.Slider(minimum=1, maximum=240, value=84, label="Frames 🎥"),
        gr.Slider(minimum=1, maximum=60, value=30, label="FPS (Frames Per Second) ⏱"),
        gr.Checkbox(label="High Quality (Requires 42GB VRAM) 🛠"),
    ],
    outputs=gr.Video(label="Generated Video"),
    title=TITLE,
    description=DESCRIPTION,
    examples=EXAMPLES,
    article=BUY_ME_A_COFFEE,
)

# Apply custom CSS for better alignment
interface.css = """
.interface-title {
    text-align: center;
    font-size: 2em;
    color: #4A90E2;
    font-family: 'Arial', sans-serif;
}
.interface-description {
    text-align: center;
    font-size: 1.2em;
    color: #333333;
    margin-bottom: 20px;
}
"""

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