app.py

"""
LTX-2.3 Turbo — ZeroGPU Edition
Generates synchronized audio-video content using Lightricks/LTX-2.3 on
free ZeroGPU hardware via Hugging Face Spaces.

Architecture (following alexnasa/ltx-2-TURBO's proven ZeroGPU pattern):
  1. Vendored ltx-core and ltx-pipelines added to sys.path before any imports.
  2. Model files downloaded at module startup (CPU, no GPU lease).
  3. ModelLedger constructed at module level (CPU-only dataclass, no CUDA init).
  4. Text encoder loaded at module level (kept in memory for reuse).
  5. DistilledPipeline constructed with gemma_root=None (no text encoder in pipeline).
  6. Video encoder and transformer pre-loaded at module level via pipeline cache.
  7. @spaces.GPU() on encode_prompt — encodes text, returns .detach().cpu() tensors.
  8. @spaces.GPU(duration=callable) on generate_video — runs pipeline with pre-encoded
     contexts passed as video_context/audio_context kwargs.
  9. FP8 quantization fits the 22B transformer on ZeroGPU's A10G (40GB VRAM).

Based on the official LTX-2 codebase: https://github.com/Lightricks/LTX-2
Architecture inspired by alexnasa/ltx-2-TURBO.
"""

# ───────────────────────────────────────────────────────────────────────────
# 0) Add vendored packages to sys.path BEFORE any ltx imports
# ───────────────────────────────────────────────────────────────────────────
import sys
from pathlib import Path

_here = Path(__file__).parent
sys.path.insert(0, str(_here / "packages" / "ltx-pipelines" / "src"))
sys.path.insert(0, str(_here / "packages" / "ltx-core" / "src"))

# ───────────────────────────────────────────────────────────────────────────
# Standard library & third-party imports
# ───────────────────────────────────────────────────────────────────────────
import logging
import os
import random
import tempfile
import time
import traceback

import gradio as gr
import numpy as np
import spaces
import torch
from huggingface_hub import hf_hub_download, snapshot_download

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# ───────────────────────────────────────────────────────────────────────────
# LTX imports (from vendored packages)
# ───────────────────────────────────────────────────────────────────────────
from ltx_core.model.video_vae import TilingConfig
from ltx_core.quantization import QuantizationPolicy
from ltx_pipelines.distilled import DistilledPipeline
from ltx_pipelines.utils import ModelLedger
from ltx_pipelines.utils.args import ImageConditioningInput
from ltx_pipelines.utils.helpers import generate_enhanced_prompt

# ───────────────────────────────────────────────────────────────────────────
# Constants
# ───────────────────────────────────────────────────────────────────────────
MAX_SEED = np.iinfo(np.int32).max
LTX_REPO = "Lightricks/LTX-2.3"
GEMMA_REPO = "google/gemma-3-12b-it-qat-q4_0-unquantized"
CKPT_DISTILLED = "ltx-2.3-22b-distilled.safetensors"
CKPT_UPSCALER = "ltx-2.3-spatial-upscaler-x2-1.0.safetensors"

RESOLUTION_PRESETS = {
    "16:9 (768x512)": (768, 512),
    "1:1 (512x512)": (512, 512),
    "9:16 (512x768)": (512, 768),
}

# ───────────────────────────────────────────────────────────────────────────
# 1) Download model files at module startup (CPU, no GPU lease)
# ───────────────────────────────────────────────────────────────────────────
logger.info("Downloading LTX model files...")
checkpoint_path = hf_hub_download(repo_id=LTX_REPO, filename=CKPT_DISTILLED)
logger.info(f"  Distilled checkpoint: {checkpoint_path}")

spatial_upsampler_path = hf_hub_download(repo_id=LTX_REPO, filename=CKPT_UPSCALER)
logger.info(f"  Upscaler: {spatial_upsampler_path}")

logger.info("Downloading Gemma text encoder...")
HF_TOKEN = os.environ.get("HF_TOKEN")
gemma_root = snapshot_download(repo_id=GEMMA_REPO, token=HF_TOKEN)
logger.info(f"  Gemma root: {gemma_root}")
logger.info("All model files ready on disk.")

# ───────────────────────────────────────────────────────────────────────────
# 2) Construct ModelLedger (CPU — no model weights loaded to GPU)
#    We construct a separate ledger WITH gemma_root for text encoding.
# ───────────────────────────────────────────────────────────────────────────
logger.info("Constructing ModelLedger (with Gemma for text encoding)...")

fp8_quantization = QuantizationPolicy.fp8_cast()

model_ledger = ModelLedger(
    dtype=torch.bfloat16,
    device="cuda",
    checkpoint_path=checkpoint_path,
    gemma_root_path=gemma_root,
    spatial_upsampler_path=spatial_upsampler_path,
    loras=(),
    quantization=fp8_quantization,
)

logger.info("ModelLedger constructed.")

# ───────────────────────────────────────────────────────────────────────────
# 3) Load text encoder at module level (kept in memory for reuse)
# ───────────────────────────────────────────────────────────────────────────
logger.info("Loading Gemma text encoder...")
text_encoder = model_ledger.text_encoder()
logger.info("Text encoder loaded and ready!")

# ───────────────────────────────────────────────────────────────────────────
# 4) Construct DistilledPipeline WITHOUT text encoder (gemma_root=None)
#    Text encoding is handled externally via encode_prompt().
# ───────────────────────────────────────────────────────────────────────────
logger.info("Constructing DistilledPipeline (gemma_root=None)...")

pipeline = DistilledPipeline(
    device=torch.device("cuda"),
    checkpoint_path=checkpoint_path,
    spatial_upsampler_path=spatial_upsampler_path,
    gemma_root=None,
    loras=[],
    quantization=fp8_quantization,
)

# ───────────────────────────────────────────────────────────────────────────
# 5) Pre-load video encoder and transformer at module level
# ───────────────────────────────────────────────────────────────────────────
logger.info("Pre-loading video encoder and transformer...")
pipeline._video_encoder = pipeline.model_ledger.video_encoder()
pipeline._transformer = pipeline.model_ledger.transformer()

logger.info("=" * 60)
logger.info("Pipeline fully loaded and ready!")
logger.info("=" * 60)


# ───────────────────────────────────────────────────────────────────────────
# Helpers
# ───────────────────────────────────────────────────────────────────────────
def calc_frames(duration: float, fps: float) -> int:
    """Compute num_frames = 8k + 1, frames >= 9."""
    raw = int(duration * fps) + 1
    raw = max(raw, 9)
    k = (raw - 1 + 7) // 8
    return k * 8 + 1


def get_duration(
    first_frame,
    prompt,
    duration,
    enhance_prompt,
    seed,
    randomize_seed,
    resolution,
    *args,
    **kwargs,
):
    """Estimate GPU lease duration for @spaces.GPU(duration=...)."""
    dur = float(duration)
    if dur <= 2:
        return 120
    elif dur <= 4:
        return 180
    else:
        return 240


# ───────────────────────────────────────────────────────────────────────────
# Phase 1: Text Encoding (separate GPU lease)
# ───────────────────────────────────────────────────────────────────────────
@spaces.GPU()
def encode_prompt(
    prompt: str,
    enhance_prompt: bool = True,
    input_image=None,
    seed: int = 42,
):
    """
    Encode prompt using the module-level text_encoder + embeddings_processor.
    Returns a dict with video_context and audio_context tensors on CPU.
    """
    logger.info(f"[encode_prompt] prompt='{prompt[:80]}...', enhance={enhance_prompt}")

    final_prompt = prompt
    if enhance_prompt:
        final_prompt = generate_enhanced_prompt(
            text_encoder=text_encoder,
            prompt=prompt,
            image_path=input_image if input_image is not None else None,
            seed=seed,
        )
        logger.info(f"[encode_prompt] Enhanced prompt: '{final_prompt[:120]}...'")

    with torch.inference_mode():
        # Step 1: Get raw hidden states from text encoder
        hidden_states, attention_mask = text_encoder.encode(final_prompt)

        # Step 2: Process through embeddings processor to get video/audio contexts
        embeddings_processor = model_ledger.gemma_embeddings_processor()
        result = embeddings_processor.process_hidden_states(
            hidden_states, attention_mask
        )
        del embeddings_processor

        video_context = result.video_encoding
        audio_context = result.audio_encoding

    embedding_data = {
        "video_context": video_context.detach().cpu(),
        "audio_context": audio_context.detach().cpu(),
        "prompt": final_prompt,
    }

    logger.info("[encode_prompt] Done.")
    return embedding_data, final_prompt


# ───────────────────────────────────────────────────────────────────────────
# Phase 2: Video Generation (separate GPU lease, dynamic duration)
# ───────────────────────────────────────────────────────────────────────────
@spaces.GPU(duration=get_duration)
def generate_video(
    first_frame,
    prompt: str,
    duration: float,
    enhance_prompt: bool,
    seed: int,
    randomize_seed: bool,
    resolution: str,
    progress=gr.Progress(track_tqdm=True),
):
    """
    Full generation: encode prompt (nested GPU call) then run pipeline with
    pre-encoded contexts.
    """
    if not prompt or not prompt.strip():
        raise gr.Error("Please enter a prompt.")

    current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
    width, height = RESOLUTION_PRESETS.get(resolution, (768, 512))
    num_frames = calc_frames(duration, 24.0)
    frame_rate = 24.0

    logger.info(
        f"[generate_video] seed={current_seed}, {width}x{height}, "
        f"frames={num_frames}, duration={duration}s, enhance={enhance_prompt}"
    )

    # --- Handle input image ---
    images = []
    image_path_for_enhance = None
    if first_frame is not None:
        tmp_dir = tempfile.mkdtemp()
        temp_image_path = os.path.join(tmp_dir, f"input_{int(time.time())}.png")
        if hasattr(first_frame, "save"):
            first_frame.save(temp_image_path)
        else:
            from PIL import Image as PILImage

            PILImage.open(first_frame).save(temp_image_path)
        images = [
            ImageConditioningInput(path=temp_image_path, frame_idx=0, strength=1.0)
        ]
        image_path_for_enhance = temp_image_path

    t0 = time.time()
    try:
        # Phase 1: Encode prompt
        embeddings, final_prompt = encode_prompt(
            prompt=prompt,
            enhance_prompt=enhance_prompt,
            input_image=image_path_for_enhance,
            seed=current_seed,
        )

        video_context = embeddings["video_context"].to("cuda", non_blocking=True)
        audio_context = embeddings["audio_context"].to("cuda", non_blocking=True)
        del embeddings
        torch.cuda.empty_cache()

        # Phase 2: Run pipeline with pre-encoded contexts
        with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile:
            output_path = tmpfile.name

        with torch.inference_mode():
            pipeline(
                prompt=prompt,
                output_path=output_path,
                seed=current_seed,
                height=height,
                width=width,
                num_frames=num_frames,
                frame_rate=frame_rate,
                images=images,
                tiling_config=TilingConfig.default(),
                video_context=video_context,
                audio_context=audio_context,
            )

        del video_context, audio_context
        torch.cuda.empty_cache()

        elapsed = time.time() - t0
        logger.info(f"[generate_video] Done in {elapsed:.1f}s")

    except torch.cuda.OutOfMemoryError:
        elapsed = time.time() - t0
        logger.error(f"OOM after {elapsed:.1f}s")
        raise gr.Error("Out of GPU memory. Try a shorter duration or lower resolution.")
    except Exception as e:
        elapsed = time.time() - t0
        tb = traceback.format_exc()
        logger.error(f"Generation failed after {elapsed:.1f}s:\n{tb}")
        raise gr.Error(f"Generation failed: {type(e).__name__}: {e}")

    info_text = (
        f"Seed: {current_seed}\n"
        f"Resolution: {width}x{height} (upscaled from {width // 2}x{height // 2})\n"
        f"Frames: {num_frames} @ {int(frame_rate)} fps\n"
        f"Duration: {duration}s\n"
        f"Pipeline: Distilled 2-stage (8+4 steps, FP8 quantized)\n"
        f"Total time: {elapsed:.1f}s\n"
        f"Hardware: ZeroGPU"
    )

    return output_path, info_text, current_seed


# ───────────────────────────────────────────────────────────────────────────
# UI toggle
# ───────────────────────────────────────────────────────────────────────────
def toggle_image(mode: str):
    return gr.update(visible=(mode == "Image to Video"))


def on_mode_change(mode: str):
    """Update image visibility and first_frame value based on mode."""
    if mode == "Image to Video":
        return gr.update(visible=True)
    return gr.update(visible=False, value=None)


# ───────────────────────────────────────────────────────────────────────────
# Gradio UI
# ───────────────────────────────────────────────────────────────────────────
CSS = """
.gradio-container { max-width: 1200px !important; }
.header { text-align: center; margin-bottom: 1rem; }
.generate-btn { min-height: 50px; }
"""

with gr.Blocks(css=CSS, title="LTX-2.3 Turbo", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # LTX-2.3 Turbo (ZeroGPU)
        Generate synchronized **video + audio** from text or images using
        [Lightricks/LTX-2.3](https://huggingface.co/Lightricks/LTX-2.3) —
        a 22B parameter DiT-based audio-video foundation model.

        Running on **free ZeroGPU** with FP8 quantization.
        Distilled pipeline (8+4 denoising steps, two-stage with 2x spatial upscaling).
        """,
        elem_classes="header",
    )

    with gr.Row():
        # --- Left: Controls ---
        with gr.Column(scale=1):
            mode = gr.Radio(
                ["Text to Video", "Image to Video"],
                value="Text to Video",
                label="Mode",
            )
            first_frame = gr.Image(
                type="pil",
                label="Input Image (first frame)",
                visible=False,
            )
            prompt = gr.Textbox(
                label="Prompt",
                lines=3,
                placeholder="Describe the scene, motion, and audio...",
                value=(
                    "A golden retriever puppy plays in fresh snow, "
                    "tossing it up with its paws, soft winter sunlight, "
                    "gentle wind sounds and playful barking"
                ),
            )

            with gr.Row():
                resolution = gr.Dropdown(
                    choices=list(RESOLUTION_PRESETS.keys()),
                    value="16:9 (768x512)",
                    label="Resolution",
                )
                duration = gr.Slider(1, 5, value=2, step=0.5, label="Duration (sec)")

            with gr.Row():
                enhance_prompt = gr.Checkbox(value=True, label="Enhance prompt")
                randomize_seed = gr.Checkbox(value=True, label="Random seed")

            seed = gr.Slider(0, MAX_SEED, value=42, step=1, label="Seed")

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

        # --- Right: Output ---
        with gr.Column(scale=1):
            output_video = gr.Video(label="Generated Video", autoplay=True)
            run_info = gr.Textbox(label="Generation Info", lines=7, interactive=False)

    # --- Events ---
    mode.change(fn=on_mode_change, inputs=mode, outputs=[first_frame])

    _inputs = [
        first_frame,
        prompt,
        duration,
        enhance_prompt,
        seed,
        randomize_seed,
        resolution,
    ]
    _outputs = [output_video, run_info, seed]

    generate_btn.click(fn=generate_video, inputs=_inputs, outputs=_outputs)

    # --- Examples ---
    gr.Examples(
        examples=[
            [
                None,
                "Aerial drone shot of a coastal city at sunset, golden light "
                "reflecting off glass buildings, gentle ocean waves, seagulls "
                "calling, cinematic ambient soundtrack",
                3.0,
                True,
                42,
                True,
                "16:9 (768x512)",
            ],
            [
                None,
                "Close-up of a barista pouring latte art in slow motion, "
                "steam rising from the cup, coffee shop ambience with soft jazz",
                2.0,
                True,
                123,
                True,
                "1:1 (512x512)",
            ],
            [
                None,
                "A cat sits on a windowsill watching rain fall outside, "
                "soft indoor lighting, raindrops on glass, gentle rain sounds",
                3.0,
                True,
                7,
                True,
                "16:9 (768x512)",
            ],
        ],
        fn=generate_video,
        inputs=_inputs,
        outputs=_outputs,
        cache_examples=False,
        label="Example Prompts",
    )

    gr.Markdown(
        """
        ---
        **Notes:**
        - ZeroGPU provides limited GPU time per request. Shorter durations are more reliable.
        - Max duration is capped at 5 seconds to stay within GPU time limits.
        - FP8 quantization reduces VRAM usage by ~50% with minimal quality impact.
        - The 2x spatial upscaler doubles the initial generation resolution.

        Built with [Lightricks/LTX-2.3](https://huggingface.co/Lightricks/LTX-2.3)
        | [GitHub](https://github.com/Lightricks/LTX-2)
        """
    )

if __name__ == "__main__":
    demo.launch(show_error=True)