app.py

import os
import subprocess
import sys

# Allocator config for pixel-space video SR memory spikes — set before torch import.
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
# Disable torch.compile / dynamo before any torch import (unsupported on ZeroGPU).
os.environ["TORCH_COMPILE_DISABLE"] = "1"
os.environ["TORCHDYNAMO_DISABLE"] = "1"

# ── Install xformers (memory-efficient attention) and video I/O deps ──────────
subprocess.run(
    [sys.executable, "-m", "pip", "install", "xformers==0.0.32.post2", "--no-build-isolation"],
    check=False,
)
subprocess.run(
    [sys.executable, "-m", "pip", "install", "imageio[ffmpeg]", "av", "num2words"],
    check=False,
)

# Reinstall torchaudio to match the pre-installed CUDA torch (deps can pull a CPU-only build).
_tv = subprocess.run(
    [sys.executable, "-c", "import torch; print(torch.__version__)"],
    capture_output=True, text=True,
)
if _tv.returncode == 0:
    _full_ver = _tv.stdout.strip()
    _cuda_suffix = _full_ver.split("+")[-1] if "+" in _full_ver else "cu124"
    _base_ver = _full_ver.split("+")[0]
    print(f"Detected torch {_full_ver}, reinstalling matching torchaudio...")
    subprocess.run(
        [sys.executable, "-m", "pip", "install", "--force-reinstall", "--no-deps",
         f"torchaudio=={_base_ver}",
         "--index-url", f"https://download.pytorch.org/whl/{_cuda_suffix}"],
        check=False,
    )

# ── Clone LTX-2 repo at a pinned commit and install ltx-core / ltx-pipelines ──
LTX_REPO_URL = "https://github.com/Lightricks/LTX-2.git"
LTX_REPO_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "LTX-2")
LTX_COMMIT = "ae855f8538843825f9015a419cf4ba5edaf5eec2"

if os.path.exists(LTX_REPO_DIR):
    subprocess.run(["rm", "-rf", LTX_REPO_DIR], check=True)

print(f"Cloning {LTX_REPO_URL}...")
subprocess.run(["git", "clone", LTX_REPO_URL, LTX_REPO_DIR], check=True)
print(f"Checking out commit {LTX_COMMIT}...")
subprocess.run(["git", "-C", LTX_REPO_DIR, "checkout", LTX_COMMIT], check=True)

print("Installing ltx-core and ltx-pipelines from pinned repo commit...")
subprocess.run(
    [sys.executable, "-m", "pip", "install", "--force-reinstall", "--no-deps",
     "-e", os.path.join(LTX_REPO_DIR, "packages", "ltx-core"),
     "-e", os.path.join(LTX_REPO_DIR, "packages", "ltx-pipelines")],
    check=True,
)
sys.path.insert(0, os.path.join(LTX_REPO_DIR, "packages", "ltx-pipelines", "src"))
sys.path.insert(0, os.path.join(LTX_REPO_DIR, "packages", "ltx-core", "src"))

import logging
import random
import tempfile

import torch
torch._dynamo.config.suppress_errors = True
torch._dynamo.config.disable = True

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

from ltx_core.loader import LoraPathStrengthAndSDOps, LTXV_LORA_COMFY_RENAMING_MAP
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
from ltx_core.quantization import QuantizationPolicy
from ltx_pipelines.ic_lora import ICLoraPipeline
from ltx_pipelines.utils.media_io import encode_video

# ── Disable xformers FA3 dispatch (Blackwell / RTX PRO 6000 is sm_120; FA3 is Hopper-only) ──
try:
    from xformers.ops.fmha import _set_use_fa3
    _set_use_fa3(False)
    print("[ATTN] xformers FA3 dispatch disabled (Blackwell-incompatible)")
except Exception as e:
    print(f"[ATTN] FA3 disable skipped: {type(e).__name__}: {e}")

logging.getLogger().setLevel(logging.INFO)

MAX_SEED = np.iinfo(np.int32).max
DEFAULT_FRAME_RATE = 24.0

# ─────────────────────────────────────────────────────────────────────────────
# Model / checkpoint locations
# ─────────────────────────────────────────────────────────────────────────────
LTX_MODEL_REPO = "Lightricks/LTX-2.3"
DISTILLED_FILENAME = "ltx-2.3-22b-distilled-1.1.safetensors"
SPATIAL_UPSAMPLER_FILENAME = "ltx-2.3-spatial-upscaler-x2-1.1.safetensors"
GEMMA_REPO = "google/gemma-3-12b-it-qat-q4_0-unquantized"

UPSCALER_REPO = "Lightricks/LTX-2.3-22b-IC-LoRA-Pixel-Spatial-Upscaler"
UPSCALER_X4_FILENAME = "ltx-2.3-22b-ic-lora-pixel-spatial-upscaler-x4-0.9.safetensors"

print("=" * 80)
print("Downloading LTX-2.3 distilled base + spatial upsampler + Gemma + upscaler IC-LoRA...")
print("=" * 80)

checkpoint_path = hf_hub_download(repo_id=LTX_MODEL_REPO, filename=DISTILLED_FILENAME)
spatial_upsampler_path = hf_hub_download(repo_id=LTX_MODEL_REPO, filename=SPATIAL_UPSAMPLER_FILENAME)
gemma_root = snapshot_download(repo_id=GEMMA_REPO)
upscaler_lora_path = hf_hub_download(repo_id=UPSCALER_REPO, filename=UPSCALER_X4_FILENAME)

print(f"Checkpoint:        {checkpoint_path}")
print(f"Spatial upsampler: {spatial_upsampler_path}")
print(f"Gemma root:        {gemma_root}")
print(f"Upscaler IC-LoRA:  {upscaler_lora_path}")

# ─────────────────────────────────────────────────────────────────────────────
# IC-LoRA pipeline (4x spatial upscaler) — built lazily inside the GPU worker.
#
# The LTX loader opens safetensors with `safe_open(device="cuda")` and streams
# tensors straight onto the GPU, which bypasses the ZeroGPU `.to("cuda")` hijack
# and fails at module scope ("No CUDA GPUs are available"). So we construct the
# pipeline on the first `@spaces.GPU` call, where a real GPU is attached, and
# cache it in a module global for subsequent calls on the same warm worker.
# ─────────────────────────────────────────────────────────────────────────────
ic_loras = [
    LoraPathStrengthAndSDOps(upscaler_lora_path, 1.0, LTXV_LORA_COMFY_RENAMING_MAP)
]

_PIPELINE = None
REFERENCE_DOWNSCALE_FACTOR = 4  # x4 upscaler LoRA (from its safetensors metadata)


def get_pipeline():
    """Build + cache the ICLoraPipeline. Must be called inside a GPU worker."""
    global _PIPELINE
    if _PIPELINE is None:
        print("[Pipeline] Building ICLoraPipeline (first GPU call)...")
        _PIPELINE = ICLoraPipeline(
            distilled_checkpoint_path=checkpoint_path,
            spatial_upsampler_path=spatial_upsampler_path,
            gemma_root=gemma_root,
            loras=ic_loras,
            quantization=QuantizationPolicy.fp8_cast(),
        )
        print(f"[Pipeline] reference_downscale_factor = {_PIPELINE.reference_downscale_factor}")
        print("[Pipeline] Ready.")
    return _PIPELINE


# ─────────────────────────────────────────────────────────────────────────────
# Resolution helpers
# ─────────────────────────────────────────────────────────────────────────────
def _round_to(x: int, base: int) -> int:
    return max(base, int(round(x / base)) * base)


def _probe_video(video_path: str):
    """Return (width, height, fps, duration_seconds) of a video."""
    import av
    with av.open(str(video_path)) as container:
        s = container.streams.video[0]
        w = s.codec_context.width
        h = s.codec_context.height
        fps = float(s.average_rate) if s.average_rate else DEFAULT_FRAME_RATE
        if s.duration is not None and s.time_base is not None:
            dur = float(s.duration * s.time_base)
        elif container.duration is not None:
            dur = float(container.duration) / 1_000_000.0
        else:
            dur = 3.0
    return w, h, fps, dur


# VAE spatial compression, read directly from the checkpoint's embedded VAE
# config (encoder_blocks): patch_size(4) * compress_space_res(2) *
# compress_all_res(2) * compress_all_res(2) = 32. NOT taken from
# ltx_core.types.SpatioTemporalScaleFactors.default() -- that "default" is a
# hardcoded literal on VideoDecoder unrelated to the actual decoder_blocks it's
# built with, so it can't be trusted to reflect any given checkpoint.
VAE_SPATIAL_COMPRESSION = 32

# ICLoraPipeline runs stage 1 at half the requested output resolution
# (ic_lora.py: `VideoPixelShape(width=width // 2, height=height // 2, ...)`),
# and stage 1's IC-LoRA reference-conditioning video is encoded at that
# half-resolution divided again by reference_downscale_factor. So the video
# actually reaching the VAE encoder is:
#   height_full / STAGE_1_RESOLUTION_DIVISOR / reference_downscale_factor
# and for that to be divisible by VAE_SPATIAL_COMPRESSION, height_full (and
# width_full) must be divisible by their product.
STAGE_1_RESOLUTION_DIVISOR = 2


def _reference_resolution_multiple(scale: int) -> int:
    """Full output resolution must be a multiple of this for the IC-LoRA
    reference-conditioning video (see constants above) to survive VAE encode."""
    return STAGE_1_RESOLUTION_DIVISOR * scale * VAE_SPATIAL_COMPRESSION


def compute_output_resolution(video_path: str, scale: int = REFERENCE_DOWNSCALE_FACTOR):
    """Target output = input resolution * scale, snapped to a multiple of
    _reference_resolution_multiple(scale)."""
    w, h, _, _ = _probe_video(video_path)
    res_multiple = _reference_resolution_multiple(scale)
    out_w = _round_to(w * scale, res_multiple)
    out_h = _round_to(h * scale, res_multiple)
    # Cap the longest side to keep runtime + VRAM bounded.
    max_side = 1536
    if max(out_w, out_h) > max_side:
        ratio = max_side / max(out_w, out_h)
        out_w = _round_to(out_w * ratio, res_multiple)
        out_h = _round_to(out_h * ratio, res_multiple)
    return int(out_w), int(out_h)


def on_video_upload(video_path):
    if video_path is None:
        return gr.update(), gr.update()
    try:
        out_w, out_h = compute_output_resolution(video_path, scale=REFERENCE_DOWNSCALE_FACTOR)
        return gr.update(value=out_w), gr.update(value=out_h)
    except Exception as e:
        print(f"[probe] failed: {e}")
        return gr.update(), gr.update()


# ─────────────────────────────────────────────────────────────────────────────
# Generation
# ─────────────────────────────────────────────────────────────────────────────
@spaces.GPU(duration=300)
@torch.inference_mode()
def upscale_video(
    input_video,
    prompt: str = "",
    lora_strength: float = 1.0,
    enhance_prompt: bool = True,
    seed: int = 42,
    randomize_seed: bool = True,
    width: int = 1536,
    height: int = 768,
    duration_seconds: float = 3.0,
    progress=gr.Progress(track_tqdm=True),
):
    """Creatively upscale a low-resolution video 4x with the LTX-2.3 IC-LoRA Pixel Spatial Upscaler.

    Args:
        input_video: Path to the low-resolution source video to upscale.
        prompt: Text describing the scene (helps the model synthesize detail).
        lora_strength: IC-LoRA reference strength. Lower stays closer to the source, higher hallucinates more detail.
        enhance_prompt: Auto-enhance the prompt with the built-in text model.
        seed: RNG seed.
        randomize_seed: Pick a fresh random seed each run.
        width: Target output width in pixels (rounded to a multiple of
            _reference_resolution_multiple(scale), 256 for the default 4x LoRA).
        height: Target output height in pixels (same rounding as width).
        duration_seconds: Length of the output clip in seconds.
    """
    if input_video is None:
        raise gr.Error("Please provide a low-resolution input video to upscale.")

    pipeline = get_pipeline()
    current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)

    frame_rate = DEFAULT_FRAME_RATE
    try:
        _, _, src_fps, src_dur = _probe_video(str(input_video))
        if src_fps and 1.0 <= src_fps <= 60.0:
            frame_rate = float(src_fps)
        duration_seconds = min(float(duration_seconds), float(src_dur) + 0.1)
    except Exception as e:
        print(f"[probe] failed, using defaults: {e}")

    num_frames = int(duration_seconds * frame_rate) + 1
    num_frames = ((num_frames - 1 + 7) // 8) * 8 + 1  # 8k + 1
    num_frames = max(9, min(num_frames, 121))

    scale = pipeline.reference_downscale_factor  # 4, but read live from the pipeline
    res_multiple = _reference_resolution_multiple(scale)
    width = _round_to(int(width), res_multiple)
    height = _round_to(int(height), res_multiple)

    print(f"[Upscale x{scale}] -> {width}x{height}, {num_frames} frames @ {frame_rate}fps, "
          f"seed={current_seed}, strength={lora_strength}")

    tiling_config = TilingConfig.default()
    video_chunks_number = get_video_chunks_number(num_frames, tiling_config)

    video, audio = pipeline(
        prompt=prompt or "",
        seed=current_seed,
        height=int(height),
        width=int(width),
        num_frames=num_frames,
        frame_rate=frame_rate,
        images=[],
        video_conditioning=[(str(input_video), float(lora_strength))],
        enhance_prompt=enhance_prompt,
        tiling_config=tiling_config,
        conditioning_attention_strength=1.0,
    )

    output_path = tempfile.mktemp(suffix=".mp4")
    encode_video(
        video=video,
        fps=frame_rate,
        audio=audio,
        output_path=output_path,
        video_chunks_number=video_chunks_number,
    )
    return str(output_path), current_seed


# ─────────────────────────────────────────────────────────────────────────────
# UI
# ─────────────────────────────────────────────────────────────────────────────
CSS = """
#col-container { max-width: 1150px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
video { object-fit: contain !important; }
"""

with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
# LTX-2.3 Pixel Spatial Upscaler (4×) 🎞️→✨
Creatively upscale a **low-resolution video 4×** with the
[LTX-2.3-22B IC-LoRA Pixel Spatial Upscaler](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-Pixel-Spatial-Upscaler)
on top of [LTX-2.3](https://huggingface.co/Lightricks/LTX-2.3).
This is a **generative** upscaler — it *synthesizes* new fine detail rather than interpolating,
so the result is a high-resolution re-render of your clip. Add a prompt describing the scene to guide the detail.
            """
        )
        with gr.Row():
            with gr.Column(scale=1):
                input_video = gr.Video(label="Low-resolution input video")
                prompt = gr.Textbox(
                    label="Prompt (optional but recommended)",
                    placeholder="describe the scene, e.g. 'a jellyfish gliding through deep blue water, bioluminescent detail'",
                    lines=2,
                )
                run_btn = gr.Button("Upscale 4×", variant="primary", size="lg")

                with gr.Accordion("Advanced settings", open=False):
                    lora_strength = gr.Slider(
                        label="Reference strength",
                        info="Lower = closer to source, higher = more synthesized detail",
                        minimum=0.5, maximum=1.0, value=1.0, step=0.05,
                    )
                    enhance_prompt = gr.Checkbox(label="Enhance prompt", value=True)
                    duration_seconds = gr.Slider(
                        label="Output duration (s)", minimum=0.5, maximum=5.0, value=3.0, step=0.5,
                    )
                    seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, value=42, step=1)
                    randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
                    with gr.Row():
                        width = gr.Number(label="Output width", value=1536, precision=0)
                        height = gr.Number(label="Output height", value=768, precision=0)

            with gr.Column(scale=1):
                output_video = gr.Video(label="Upscaled result (4×)", autoplay=True)

        gr.Examples(
            examples=[
                ["jellyfish_glide.mp4", "a jellyfish gliding gracefully through deep blue water, translucent bioluminescent detail"],
                ["waterfall_forest.mp4", "a waterfall cascading through a lush green forest, misty spray and wet mossy rocks"],
                ["squirrel_eating.mp4", "a squirrel eating a nut, detailed fur, soft natural daylight"],
            ],
            inputs=[input_video, prompt],
            outputs=[output_video, seed],
            fn=upscale_video,
            cache_examples=True,
            cache_mode="lazy",
        )

    input_video.change(fn=on_video_upload, inputs=[input_video], outputs=[width, height])
    run_btn.click(
        fn=upscale_video,
        inputs=[input_video, prompt, lora_strength, enhance_prompt, seed,
                randomize_seed, width, height, duration_seconds],
        outputs=[output_video, seed],
        api_name="upscale",
    )


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