app.py

import gradio as gr
import torch
import random
from diffusers import LTXPipeline
from diffusers.utils import export_to_video
import spaces

# ============================================
# DARK CINEMATIC THEME CSS
# ============================================
CUSTOM_CSS = """
/* Global dark theme */
.gradio-container {
    background: linear-gradient(135deg, #0a0a0f 0%, #1a1a2e 50%, #16213e 100%) !important;
    min-height: 100vh;
}

/* Header styling */
.header-container {
    text-align: center;
    padding: 2rem 1rem;
    margin-bottom: 1rem;
}

.header-title {
    font-size: 3rem;
    font-weight: 800;
    background: linear-gradient(135deg, #8b5cf6 0%, #3b82f6 50%, #06b6d4 100%);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
    margin-bottom: 0.5rem;
    text-shadow: 0 0 40px rgba(139, 92, 246, 0.3);
}

.header-subtitle {
    color: #94a3b8;
    font-size: 1.1rem;
    font-weight: 400;
}

/* Video output container */
.video-container {
    background: linear-gradient(145deg, #1e1e2e 0%, #2a2a3e 100%);
    border: 1px solid rgba(139, 92, 246, 0.3);
    border-radius: 16px;
    padding: 1rem;
    box-shadow: 0 8px 32px rgba(0, 0, 0, 0.4), 0 0 60px rgba(139, 92, 246, 0.1);
}

/* Prompt input styling */
.prompt-input textarea {
    background: linear-gradient(145deg, #1a1a2e 0%, #252540 100%) !important;
    border: 1px solid rgba(139, 92, 246, 0.4) !important;
    border-radius: 12px !important;
    color: #e2e8f0 !important;
    font-size: 1rem !important;
    padding: 1rem !important;
    transition: all 0.3s ease !important;
}

.prompt-input textarea:focus {
    border-color: #8b5cf6 !important;
    box-shadow: 0 0 20px rgba(139, 92, 246, 0.3) !important;
}

.prompt-input textarea::placeholder {
    color: #64748b !important;
}

/* Generate button */
.generate-btn {
    background: linear-gradient(135deg, #8b5cf6 0%, #6366f1 50%, #3b82f6 100%) !important;
    border: none !important;
    border-radius: 12px !important;
    color: white !important;
    font-size: 1.1rem !important;
    font-weight: 600 !important;
    padding: 0.875rem 2rem !important;
    cursor: pointer !important;
    transition: all 0.3s ease !important;
    box-shadow: 0 4px 20px rgba(139, 92, 246, 0.4) !important;
}

.generate-btn:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 6px 30px rgba(139, 92, 246, 0.6) !important;
}

/* Settings accordion */
.settings-accordion {
    background: rgba(30, 30, 46, 0.6) !important;
    border: 1px solid rgba(139, 92, 246, 0.2) !important;
    border-radius: 12px !important;
    margin-top: 1rem !important;
}

/* Sliders */
input[type="range"] {
    accent-color: #8b5cf6 !important;
}

/* Example prompts */
.example-prompt {
    background: linear-gradient(145deg, #1e1e2e 0%, #2a2a3e 100%);
    border: 1px solid rgba(139, 92, 246, 0.2);
    border-radius: 10px;
    padding: 0.75rem 1rem;
    color: #cbd5e1;
    cursor: pointer;
    transition: all 0.3s ease;
    font-size: 0.9rem;
}

.example-prompt:hover {
    border-color: #8b5cf6;
    background: linear-gradient(145deg, #252540 0%, #303050 100%);
    transform: translateY(-1px);
    box-shadow: 0 4px 15px rgba(139, 92, 246, 0.2);
}

/* Labels */
.gradio-container label {
    color: #94a3b8 !important;
    font-weight: 500 !important;
}

/* Accordion labels */
.gradio-accordion > button {
    background: transparent !important;
    color: #e2e8f0 !important;
}

/* Status text */
.status-text {
    color: #8b5cf6;
    font-size: 0.9rem;
    text-align: center;
    padding: 0.5rem;
}

/* Section dividers */
.section-title {
    color: #8b5cf6;
    font-size: 0.875rem;
    font-weight: 600;
    text-transform: uppercase;
    letter-spacing: 0.1em;
    margin-bottom: 0.75rem;
    padding-bottom: 0.5rem;
    border-bottom: 1px solid rgba(139, 92, 246, 0.2);
}
"""

# ============================================
# VIDEO GENERATION FUNCTION
# ============================================
@spaces.GPU(duration=180)
def generate_video(
    prompt,
    negative_prompt,
    num_frames,
    guidance_scale,
    num_inference_steps,
    seed,
    height,
    width,
    fps,
    progress=gr.Progress(track_tqdm=True)
):
    import gc

    if not prompt or prompt.strip() == "":
        gr.Warning("Please enter a prompt to generate a video.")
        return None

    progress(0, desc="Loading model...")

    # Clear memory before loading
    gc.collect()
    torch.cuda.empty_cache() if torch.cuda.is_available() else None

    # Load model with memory optimizations
    pipe = LTXPipeline.from_pretrained(
        "Lightricks/LTX-Video",
        torch_dtype=torch.float16,
        low_cpu_mem_usage=True,
        use_safetensors=True,
    )

    # Use CPU offload for memory efficiency (don't call .to("cuda") with this)
    pipe.enable_sequential_cpu_offload()
    pipe.vae.enable_slicing()
    pipe.vae.enable_tiling()

    # Handle seed
    if seed == -1:
        seed = random.randint(0, 2**32 - 1)

    generator = torch.Generator(device="cpu").manual_seed(int(seed))

    progress(0.1, desc="Generating video frames...")

    # Generate video with memory-efficient settings
    with torch.inference_mode():
        output = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt if negative_prompt else None,
            num_frames=int(num_frames),
            guidance_scale=float(guidance_scale),
            num_inference_steps=int(num_inference_steps),
            generator=generator,
            height=int(height),
            width=int(width),
        )

    progress(0.9, desc="Exporting video...")

    # Export to video file
    video_path = export_to_video(output.frames[0], fps=int(fps))

    # Cleanup
    del pipe
    gc.collect()
    torch.cuda.empty_cache() if torch.cuda.is_available() else None

    progress(1.0, desc="Complete!")

    return video_path

# ============================================
# EXAMPLE PROMPTS
# ============================================
EXAMPLE_PROMPTS = [
    "A cinematic shot of a golden retriever running through autumn leaves, golden hour lighting, slow motion, 4K quality",
    "Timelapse of a flower blooming in soft morning light, macro photography, dreamy atmosphere",
    "Ocean waves crashing on rocks at sunset, aerial drone shot, cinematic colors, peaceful mood",
    "A futuristic city with flying cars and neon lights, rain falling, cyberpunk aesthetic, night scene",
    "Hot air balloons floating over mountains at sunrise, misty atmosphere, vibrant colors, peaceful",
    "A cozy cabin in the woods during snowfall, warm light from windows, winter wonderland, magical",
]

# ============================================
# BUILD GRADIO UI
# ============================================
def create_ui():
    with gr.Blocks(css=CUSTOM_CSS, theme=gr.themes.Base()) as demo:

        # Header
        gr.HTML("""
            <div class="header-container">
                <h1 class="header-title">LTX Video Studio</h1>
                <p class="header-subtitle">Transform your imagination into cinematic videos with AI</p>
            </div>
        """)

        with gr.Row():
            with gr.Column(scale=3):
                # Video Output
                video_output = gr.Video(
                    label="Generated Video",
                    elem_classes=["video-container"],
                    height=400,
                    interactive=False
                )

                # Main Prompt
                prompt = gr.Textbox(
                    label="Describe your video",
                    placeholder="A cinematic shot of a mountain landscape at golden hour, with clouds slowly drifting by, dramatic lighting...",
                    lines=3,
                    elem_classes=["prompt-input"]
                )

                # Negative Prompt (collapsible)
                with gr.Accordion("Negative Prompt (Optional)", open=False, elem_classes=["settings-accordion"]):
                    negative_prompt = gr.Textbox(
                        label="What to avoid",
                        placeholder="blurry, low quality, distorted, ugly, bad anatomy...",
                        lines=2,
                        elem_classes=["prompt-input"]
                    )

                # Generation Settings
                with gr.Accordion("Generation Settings", open=True, elem_classes=["settings-accordion"]):
                    with gr.Row():
                        with gr.Column(scale=1):
                            height = gr.Dropdown(
                                label="Height",
                                choices=[("480p", 480), ("512", 512), ("576", 576), ("720p", 720)],
                                value=480,
                                interactive=True
                            )
                        with gr.Column(scale=1):
                            width = gr.Dropdown(
                                label="Width",
                                choices=[("480", 480), ("512", 512), ("704", 704), ("720", 720), ("848", 848)],
                                value=704,
                                interactive=True
                            )
                        with gr.Column(scale=1):
                            num_frames = gr.Slider(
                                label="Frames",
                                minimum=9,
                                maximum=57,
                                value=25,
                                step=8,
                                info="More frames = longer video"
                            )

                    with gr.Row():
                        with gr.Column(scale=1):
                            guidance_scale = gr.Slider(
                                label="Guidance Scale",
                                minimum=1.0,
                                maximum=20.0,
                                value=7.5,
                                step=0.5,
                                info="Higher = more prompt adherence"
                            )
                        with gr.Column(scale=1):
                            num_inference_steps = gr.Slider(
                                label="Inference Steps",
                                minimum=10,
                                maximum=40,
                                value=20,
                                step=5,
                                info="More steps = higher quality"
                            )
                        with gr.Column(scale=1):
                            fps = gr.Slider(
                                label="FPS",
                                minimum=8,
                                maximum=24,
                                value=16,
                                step=1,
                                info="Frames per second"
                            )

                    with gr.Row():
                        seed = gr.Number(
                            label="Seed (-1 for random)",
                            value=-1,
                            precision=0
                        )

                # Generate Button
                generate_btn = gr.Button(
                    "Generate Video",
                    variant="primary",
                    elem_classes=["generate-btn"],
                    size="lg"
                )

                # Example Prompts
                gr.HTML('<div class="section-title">Example Prompts</div>')

                with gr.Row():
                    for i in range(3):
                        example_btn = gr.Button(
                            EXAMPLE_PROMPTS[i][:60] + "...",
                            elem_classes=["example-prompt"],
                            size="sm"
                        )
                        example_btn.click(
                            fn=lambda x=EXAMPLE_PROMPTS[i]: x,
                            outputs=prompt
                        )

                with gr.Row():
                    for i in range(3, 6):
                        example_btn = gr.Button(
                            EXAMPLE_PROMPTS[i][:60] + "...",
                            elem_classes=["example-prompt"],
                            size="sm"
                        )
                        example_btn.click(
                            fn=lambda x=EXAMPLE_PROMPTS[i]: x,
                            outputs=prompt
                        )

        # Event handlers
        generate_btn.click(
            fn=generate_video,
            inputs=[
                prompt,
                negative_prompt,
                num_frames,
                guidance_scale,
                num_inference_steps,
                seed,
                height,
                width,
                fps
            ],
            outputs=video_output
        )

        # Footer
        gr.HTML("""
            <div style="text-align: center; padding: 2rem; color: #64748b; font-size: 0.875rem;">
                Powered by <a href="https://huggingface.co/Lightricks/LTX-Video" target="_blank" style="color: #8b5cf6;">Lightricks LTX-Video</a>
                &bull; Built with Gradio
            </div>
        """)

    return demo

# ============================================
# MAIN
# ============================================
if __name__ == "__main__":
    demo = create_ui()
    demo.launch()