app.py

import os
import subprocess
import sys

# ZeroGPU: torch.compile / dynamo unsupported — disable before any torch import.
os.environ["TORCH_COMPILE_DISABLE"] = "1"
os.environ["TORCHDYNAMO_DISABLE"] = "1"

# memory-efficient attention
subprocess.run([sys.executable, "-m", "pip", "install", "xformers==0.0.32.post2", "--no-build-isolation"], check=False)

# --- clone + install the NATIVE LTX-2 codebase at the pinned commit the working ZeroGPU spaces use ---
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 not os.path.exists(LTX_REPO_DIR):
    subprocess.run(["git", "clone", LTX_REPO_URL, LTX_REPO_DIR], check=True)
    subprocess.run(["git", "-C", LTX_REPO_DIR, "checkout", LTX_COMMIT], check=True)
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 numpy as np
import imageio.v3 as iio
from PIL import Image, ImageOps

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

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


# Import LTX modules in the proven order — importing ltx_core.quantization/loader FIRST hits a
# circular import (fp8_cast <-> loader.fuse_loras). Importing the model modules first forces the
# correct init order (mirrors the working reference Space).
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number, decode_video as _vae_decode_video  # noqa: F401
from ltx_core.model.upsampler import upsample_video as _upsample_video  # noqa: F401
from ltx_core.model.audio_vae import encode_audio as _vae_encode_audio  # noqa: F401
from ltx_core.quantization import QuantizationPolicy
from ltx_core.loader import LoraPathStrengthAndSDOps, LTXV_LORA_COMFY_RENAMING_MAP
from ltx_pipelines.ic_lora import ICLoraPipeline
from ltx_pipelines.utils.media_io import encode_video

# --- ZeroGPU loader patch -------------------------------------------------------------
# The native loader opens safetensors directly on the CUDA device:
#     safetensors.safe_open(path, framework="pt", device="cuda")
# which performs the host->device copy inside safetensors' own C++ (cudaMemcpy), BYPASSING
# torch.Tensor.to — the exact call ZeroGPU patches to virtualise + pack weights at module
# scope. The result: "No CUDA GPUs are available" at startup and nothing gets packed.
# Patch the loader to open on CPU and move via torch.Tensor.to (which ZeroGPU virtualises),
# so the module-scope preload packs correctly — matching diffusers / the reference Space.
import safetensors as _safetensors
import ltx_core.loader.sft_loader as _sft
from ltx_core.loader.primitives import StateDict as _StateDict

def _zerogpu_safe_load(self, path, sd_ops, device=None):
    device = device or torch.device("cpu")
    sd, size, dtype = {}, 0, set()
    model_paths = path if isinstance(path, list) else [path]
    for shard_path in model_paths:
        with _safetensors.safe_open(shard_path, framework="pt", device="cpu") as f:
            for name in f.keys():
                expected = name if sd_ops is None else sd_ops.apply_to_key(name)
                if expected is None:
                    continue
                value = f.get_tensor(name).to(device=device)  # torch path → ZeroGPU-virtualised
                kvs = ((expected, value),)
                if sd_ops is not None:
                    kvs = sd_ops.apply_to_key_value(expected, value)
                for k, v in kvs:
                    size += v.nbytes
                    dtype.add(v.dtype)
                    sd[k] = v
    return _StateDict(sd=sd, device=device, size=size, dtype=dtype)

_sft.SafetensorsStateDictLoader.load = _zerogpu_safe_load
print("[PATCH] safetensors loader → CPU-open + torch.to (ZeroGPU-virtualisable)")
# --------------------------------------------------------------------------------------

# --- attention backend patch (FA3 crashes on Blackwell ZeroGPU; use xformers/SDPA) ---
import torch.nn.functional as F
from ltx_core.model.transformer import attention as _attn_mod

def _sdpa_as_mea(query, key, value, attn_bias=None, scale=None, **kwargs):
    q, k, v = query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2)
    return F.scaled_dot_product_attention(q, k, v, scale=scale).transpose(1, 2)

# IMPORTANT (ZeroGPU): do NOT query CUDA at module scope. torch.cuda.get_device_capability()
# forces torch._C._cuda_init() in the GPU-less main process, which poisons ZeroGPU's CUDA
# virtualization — the module-scope model preload then fails with "No CUDA GPUs are available"
# and ZeroGPU can't pack the weights. SDPA works on every GPU (incl. Blackwell ZeroGPU, where
# FA3 crashes), so patch it unconditionally without ever touching torch.cuda here.
_attn_mod.memory_efficient_attention = _sdpa_as_mea
print("[ATTN] SDPA (patched at module scope, no CUDA query)")

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

# =========================== PER-LORA CONFIG (colorize) ===========================
TITLE = "LTX-2.3 Colorize (native LTX-2)"
LORA_REPO = "Lightricks/LTX-2.3-22b-IC-LoRA-Colorization"
LORA_FILE = "ltx-2.3-22b-ic-lora-colorization-0.9.safetensors"
LORA_SCALE = 1.0
SKIP_STAGE_2 = True          # restoration LoRA: stage-1-only native hi-res (per card)
GRAYSCALE_REF = True         # colorize conditions on a B&W reference
RES_PRESETS = {"960×544 (recommended)": (960, 544), "768×448 (fast)": (768, 448)}
DEFAULT_PRESET = "960×544 (recommended)"
FRAME_CHOICES = [49, 73, 97, 121]
DEFAULT_FRAMES = 121

def build_prompt(p):
    return (
        "Reference shows the same scene in high-contrast monochrome with soft natural daylight. "
        "Edited shows the same scene with natural colors restored. "
        f"COLORIZE {p.strip()}. "
        "Subject identity, framing, and background geometry are identical to the reference; "
        "only color information differs between reference and edited."
    )

EXAMPLES = [
    ["examples/rabbit_rocks_gray.mp4",
     "a young brown cottontail rabbit with warm tan and grey-brown fur, a pale cream underside and soft pink inner ears, perched on weathered grey granite boulders flecked with green and ochre lichen in warm late-afternoon sun; gentle wind, distant birdsong and soft rustling grass",
     "960×544 (recommended)", 121, 42, False],
    ["examples/surfing_gray.mp4",
     "a surfer in a black wetsuit riding a curling turquoise ocean wave, bright white foam spraying off the crest, deep blue sky and sunlit teal water; powerful ocean waves crashing, rushing water, wind and distant seagulls",
     "960×544 (recommended)", 121, 42, False],
]
# =================================================================================

FPS = 24.0
MAX_SEED = np.iinfo(np.int32).max
HF_TOKEN = os.environ.get("HF_TOKEN")
LTX_MODEL_REPO = "Lightricks/LTX-2.3"
GEMMA_REPO = "google/gemma-3-12b-it-qat-q4_0-unquantized"


def _src_fps(path, default=FPS):
    try:
        return float(iio.immeta(path, plugin="pyav").get("fps", default)) or default
    except Exception:
        return default


def _prep_reference(path, width, height, num_frames):
    """Resample to 24fps, aspect-fit/crop to WxH, NF frames; (optionally grayscale); write temp mp4."""
    vid = iio.imread(path, plugin="pyav")
    src_fps = _src_fps(path)
    n = len(vid)
    out = []
    for i in range(num_frames):
        idx = min(int(round(i / FPS * src_fps)), n - 1)
        im = Image.fromarray(vid[idx]).convert("RGB")
        im = ImageOps.fit(im, (width, height), Image.LANCZOS)
        if GRAYSCALE_REF:
            im = im.convert("L").convert("RGB")
        out.append(np.array(im))
    tmp = tempfile.mktemp(suffix=".mp4")
    iio.imwrite(tmp, np.stack(out), fps=FPS, plugin="pyav", codec="libx264")
    return tmp


def _pick_resolution(path, preset):
    w, h = RES_PRESETS[preset]
    try:
        f0 = iio.imread(path, plugin="pyav", index=0)
        if f0.shape[0] > f0.shape[1]:  # portrait
            w, h = h, w
    except Exception:
        pass
    return w, h


# --- Load native pipeline + IC-LoRA once at module scope (ZeroGPU packs weights here) ---
print("Downloading checkpoints…")
checkpoint_path = hf_hub_download(LTX_MODEL_REPO, "ltx-2.3-22b-distilled-1.1.safetensors", token=HF_TOKEN)
spatial_upsampler_path = hf_hub_download(LTX_MODEL_REPO, "ltx-2.3-spatial-upscaler-x2-1.1.safetensors", token=HF_TOKEN)
gemma_root = snapshot_download(GEMMA_REPO, token=HF_TOKEN)
lora_path = hf_hub_download(LORA_REPO, LORA_FILE, token=HF_TOKEN)

print("Building ICLoraPipeline…")
pipeline = ICLoraPipeline(
    distilled_checkpoint_path=checkpoint_path,
    spatial_upsampler_path=spatial_upsampler_path,
    gemma_root=gemma_root,
    loras=[LoraPathStrengthAndSDOps(lora_path, LORA_SCALE, LTXV_LORA_COMFY_RENAMING_MAP)],
    # bf16 (NOT fp8): the IC-LoRA must be fused into the transformer at MODULE SCOPE (the GPU
    # worker can't re-open the checkpoint file). fp8_cast()'s fusion runs a custom CUDA kernel
    # that can't be ZeroGPU-virtualised ("CUDA error: no CUDA-capable device"), but the bf16
    # fuse rule is pure torch matmul/add → virtualisable + packable. ~53GB pack (fits H200).
    quantization=None,
)


# All components (incl. the bf16-fused transformer) load + pin at MODULE SCOPE so ZeroGPU
# packs them (~53GB) and transfers them into each GPU worker — the worker can't re-open the
# checkpoint file, so nothing may be built there. The CPU-open loader patch above makes the
# host->device moves virtualisable; bf16 keeps the LoRA fusion virtualisable too.
def _preload_pin(ledger, tag):
    if ledger is None:
        return
    for name in ["video_encoder", "video_decoder", "audio_encoder", "audio_decoder",
                 "vocoder", "spatial_upsampler", "text_encoder", "gemma_embeddings_processor",
                 "transformer"]:
        fn = getattr(ledger, name, None)
        if callable(fn):
            try:
                obj = fn()
                setattr(ledger, name, (lambda o=obj: o))
                print(f"[preload {tag}] {name} ✓")
            except Exception as e:
                print(f"[preload {tag}] {name} skipped: {e}")


_preload_pin(getattr(pipeline, "stage_1_model_ledger", None), "stage1")
if not SKIP_STAGE_2:
    _preload_pin(getattr(pipeline, "stage_2_model_ledger", None), "stage2")
print("Pipeline ready (all components preloaded + pinned for ZeroGPU packing).")


def _duration(*args, **kwargs):
    nf = next((a for a in args if isinstance(a, int) and a in FRAME_CHOICES), DEFAULT_FRAMES)
    return int(60 + nf * 1.2)


@spaces.GPU(duration=_duration)
@torch.inference_mode()
def colorize(video, prompt, preset, num_frames, seed, randomize, progress=gr.Progress(track_tqdm=True)):
    if video is None:
        raise gr.Error("Please upload a video.")
    if not prompt.strip():
        raise gr.Error("Describe the result (e.g. 'a brown rabbit on grey rocks, soft birdsong').")
    seed = random.randint(0, MAX_SEED) if randomize else int(seed)
    num_frames = int(num_frames)
    width, height = _pick_resolution(video, preset)
    ref_path = _prep_reference(video, width, height, num_frames)
    tiling = TilingConfig.default()
    # skip_stage_2 outputs at half the passed dims (height//2, width//2) — pass 2× so the
    # final video matches the chosen preset. (Two-stage demos pass the preset directly.)
    gen_w, gen_h = (width * 2, height * 2) if SKIP_STAGE_2 else (width, height)
    video_out, audio_out = pipeline(
        prompt=build_prompt(prompt),
        seed=seed, height=gen_h, width=gen_w,
        num_frames=num_frames, frame_rate=FPS,
        images=[], video_conditioning=[(ref_path, 1.0)],
        skip_stage_2=SKIP_STAGE_2, tiling_config=tiling,
    )
    out_path = tempfile.mktemp(suffix=".mp4")
    encode_video(video=video_out, fps=FPS, audio=audio_out, output_path=out_path,
                 video_chunks_number=get_video_chunks_number(num_frames, tiling))
    return out_path, seed


# --- UI config (match the public Space exactly) ---
RES_PRESETS = {"960×544 (recommended)": (960, 544), "768×448 (fast)": (768, 448)}
FRAME_CHOICES = [49, 73, 97, 121]


with gr.Blocks(title="LTX-2.3 Colorize") as demo:
    gr.Markdown(
        "# 🎨 LTX-2.3 Video Colorization\n"
        "Restore natural color to black-and-white or desaturated footage while keeping subject, framing and "
        "motion identity. Using [LTX 2.3 Distilled](https://huggingface.co/Lightricks/LTX-2.3) "
        "with the [Colorization IC-LoRA](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-Colorization)."
    )
    with gr.Row():
        with gr.Column():
            video_in = gr.Video(label="Input video (any clip — recolored as B&W)")
            prompt = gr.Textbox(
                label="Prompt — describe the colorized scene, plus any sounds", lines=3,
                placeholder="a young brown rabbit with cream underside and pink inner ears on grey granite rocks flecked with green lichen, warm afternoon light, soft golden dry grass; gentle wind and distant birdsong",
            )
            with gr.Accordion("Settings", open=False):
                preset = gr.Dropdown(list(RES_PRESETS), value="960×544 (recommended)", label="Resolution")
                num_frames = gr.Dropdown(FRAME_CHOICES, value=121, label="Frames (24fps)")
                randomize = gr.Checkbox(True, label="Randomize seed")
                seed = gr.Slider(0, MAX_SEED, value=42, step=1, label="Seed")
            run = gr.Button("Colorize", variant="primary")
        with gr.Column():
            video_out = gr.Video(label="Colorized result")

    run.click(colorize, inputs=[video_in, prompt, preset, num_frames, seed, randomize],
              outputs=[video_out, seed])

    gr.Examples(
        examples=[
            ["examples/rabbit_rocks_gray.mp4",
             "a young brown cottontail rabbit with warm tan and grey-brown fur, a pale cream underside and soft pink inner ears, perched on weathered grey granite boulders flecked with green and ochre lichen, with bleached driftwood and clumps of golden dry grass in warm late-afternoon sun; gentle wind, distant birdsong and soft rustling grass",
             "960×544 (recommended)", 121, 42, False],
            ["examples/slicing_veggie_gray.mp4",
             "close-up of hands slicing a fresh green zucchini into thin rounds on a light wooden cutting board, the bright green skin and pale interior of the courgette, a stainless-steel knife catching warm kitchen light; crisp rhythmic chopping on the board and a faint kitchen ambience",
             "960×544 (recommended)", 121, 42, False],
            ["examples/surfing_gray.mp4",
             "a surfer in a black wetsuit riding a curling turquoise ocean wave, bright white foam spraying off the crest, deep blue sky and sunlit teal water; powerful ocean waves crashing, rushing water, wind and distant seagulls",
             "960×544 (recommended)", 121, 42, False],
            ["examples/cat_on_a_tree_gray.mp4",
             "a tabby cat with warm brown and grey striped fur and bright green eyes, clinging to rough red-brown tree bark surrounded by lush sunlit green leaves; soft leaves rustling in the breeze, faint birdsong and a quiet meow",
             "960×544 (recommended)", 121, 42, False],
        ],
        inputs=[video_in, prompt, preset, num_frames, seed, randomize],
        outputs=[video_out, seed], fn=colorize, cache_examples=True, cache_mode="lazy",
    )

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