handler/handler.py

"""RunPod LTX-2.5 handler pinned to upstream v1.4.2.

Modes: i2v (distilled), premium (guided HQ), keyframe (guided interpolation),
retake (video region), dfr (distilled + spatial detailing IC-LoRA), refine
(tiled native Refine-Details video-to-video), selftest (no render).
All pipelines return video, audio, num_frames, resolved tiling configuration.
Multi-photo conditioning accepts images[{url, frame_idx, strength}].
DFR spatial_upscalings=2 requires dimensions divisible by 128; other two-stage
modes require 64. Frame counts are 8k+1. Temporal upsampling adjusts output fps.
On RunPod, weights resolve only from its provisioned HF cache. Pin
LTX_MODEL_REVISION to the same revision as the endpoint's cached Model field.
DFR/refine remain experimental until measured on the deployed GPU.
"""
import io
import json
import os
import uuid

_RUNPOD_HF_CACHE = "/runpod-volume/huggingface-cache/hub"
if os.path.isdir(_RUNPOD_HF_CACHE):
    os.environ.setdefault("HF_HUB_CACHE", _RUNPOD_HF_CACHE)

# Same fragmentation-mitigation reasoning as ../Wan2.2-runpod/handler.py --
# set before torch is imported, the caching allocator reads this once at
# CUDA init. Both names on purpose: torch 2.9 (this repo's real resolved
# pin) renamed the variable to PYTORCH_ALLOC_CONF and logs a deprecation
# warning for the old one (seen in the real 2026-07-31 worker log) but still
# honors it; older torch only reads the old name. Same value either way.
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
os.environ.setdefault("PYTORCH_ALLOC_CONF", "expandable_segments:True")

import logging

import runpod
import requests
import torch
from PIL import Image

# Upstream's own entry points (distilled.py's main()) do exactly this, and our
# 2.3 handler did NOT -- which is why `ltx_pipelines` ran silently. Everything
# it logs sits at INFO: which model is being built for each stage, and
# critically "Enhancing with separate Gemma ..." when the prompt enhancer runs.
# Without this line the enhancer question had to be answered by reading source
# instead of by reading a worker log.
logging.basicConfig(level=logging.INFO)

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# --- Model locations ---------------------------------------------------
# One repo holds every component. Defaults point at Lightricks' own repo so
# this handler runs standalone with no configuration; the endpoint sets
# LTX_MODEL_REPO to the serving bundle, whose paths are identical.
# Defaults to the serving bundle, NOT upstream. The reverse (upstream default,
# bundle via env) was the original design so the handler would run standalone --
# its only real-world effect was that a missing LTX_MODEL_REPO on the endpoint
# silently downloaded 42GB from upstream and died on container disk. A default
# that fails safe beats a default that is theoretically more portable.
MODEL_REPO = os.environ.get("LTX_MODEL_REPO", "tugrulv89/LTX_2.5-handler")
TRANSFORMER_FILE = os.environ.get(
    "LTX_TRANSFORMER_FILE",
    "diffusion_models/ltx-2.5-22b-distilled-transformer-bf16.safetensors",
)
TEXT_ENCODER_FILE = os.environ.get(
    "LTX_TEXT_ENCODER_FILE",
    "text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors",
)
VIDEO_VAE_FILE = os.environ.get("LTX_VIDEO_VAE_FILE", "vae/ltx-2.5-video-vae-bf16.safetensors")
AUDIO_VAE_FILE = os.environ.get("LTX_AUDIO_VAE_FILE", "vae/ltx-2.5-audio-vae-bf16.safetensors")
# The prompt enhancer is a SECOND, separate Gemma -- and it is not optional
# plumbing, it is a hard requirement of 2.5. LTX-2.3's text encoder was
# `google/gemma-3-12b-it`, a generative instruct model, so `PromptEncoder` could
# reuse it for enhancement. LTX-2.5's encoder is `gemma4_unified`: encoder-only,
# it cannot generate text. utils/blocks.py's call path enforces this:
#     elif enhance_first_prompt and self._encode_model_type != "gemma3":
#         raise ValueError("... requires --prompt-enhancer-gemma-root pointing at
#                           a generative instruct checkpoint ...")
# Proven by a real job, not inferred: enhance_prompt=true returned ok:false with
# exactly that ValueError. The bundle therefore carries `gemma/` -- the same
# gemma-3-12b-it snapshot the 2.3 deployment already used -- copied across, not
# re-downloaded. Set to "" to disable enhancement entirely.
PROMPT_ENHANCER_SUBFOLDER = os.environ.get("LTX_PROMPT_ENHANCER_SUBFOLDER", "gemma")

UPSCALER_FILE = os.environ.get(
    "LTX_UPSCALER_FILE",
    "latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors",
)

# --- DFR: distilled base plus the required spatial detailing IC-LoRA. ---
# Both pipelines diffuse at full resolution in stage 2; DFR additionally uses
# a reference-conditioned detailing pass and optional tiled spatial epilogue.
TEMPORAL_UPSCALER_FILE = os.environ.get(
    "LTX_TEMPORAL_UPSCALER_FILE",
    "latent_upscale_models/ltx-2.5-latent-temporal-upscaler-x2-bf16-1.0.safetensors",
)
DETAILING_LORA_FILE = os.environ.get(
    "LTX_DETAILING_LORA_FILE", "loras/ltx-2.5-22b-ic-lora-pixel-spatial-upscaler-x2-1.0.safetensors")
DETAILING_LORA_STRENGTH = float(os.environ.get("LTX_DETAILING_LORA_STRENGTH", 0.5))
# DFR's own quantization, defaulting to whatever every other mode runs so that
# turning this mode on changes nothing else. "" runs it at full bf16 precision
# without touching the modes whose prompts and prices were measured on fp8.
LTX_DFR_QUANTIZATION = os.environ.get(
    "LTX_DFR_QUANTIZATION", os.environ.get("LTX_QUANTIZATION", "fp8-cast"))

# --- Phase 2 weights: already in the bundle, downloaded only on first use ----
# Both files are listed by the HF API on tugrulv89/LTX_2.5-handler (checked
# 2026-08-28): the dev transformer at 42.02 GB and the distilled LoRA at 8.90 GB.
# They are what every non-distilled mode needs -- `premium` and `keyframe` build
# the FULL transformer and apply the distilled LoRA on top, which is exactly
# what upstream's ti2vid_two_stages_hq.py / keyframe_interpolation.py main() do.
# Nothing else in the pack changes: text encoder, both VAEs and the spatial
# upsampler are shared by all four modes.
DEV_TRANSFORMER_FILE = os.environ.get(
    "LTX_DEV_TRANSFORMER_FILE",
    "diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors",
)
DISTILLED_LORA_FILE = os.environ.get(
    "LTX_DISTILLED_LORA_FILE",
    "loras/ltx-2.5-22b-distilled-lora-450-bf16.safetensors",
)

# HQ's two distilled-LoRA strengths. Upstream's own CLI defaults
# (hq_2_stage_arg_parser: 0.25 for stage 1, 0.5 for stage 2) -- the LoRA is
# applied at a different strength per stage, which is why the pipeline takes
# them separately rather than reading the strength off the LoRA tuple.
HQ_LORA_STRENGTH_STAGE_1 = float(os.environ.get("LTX_HQ_LORA_STRENGTH_STAGE_1", 0.25))
HQ_LORA_STRENGTH_STAGE_2 = float(os.environ.get("LTX_HQ_LORA_STRENGTH_STAGE_2", 0.5))

# Upper bound on a caller-supplied step count. Steps are the price: at ~2 s of
# GPU per step per stage, a payload asking for 200 would run for the better part
# of an hour and hit the endpoint's executionTimeout with nothing to show. The
# real defaults come from upstream's own params objects (see _steps_for).
MAX_INFERENCE_STEPS = int(os.environ.get("LTX_MAX_INFERENCE_STEPS", 60))

# Guidance passes batched into one transformer forward. Upstream's CLI default is
# 1; a guided denoiser can batch up to 4. Worth having on a 96 GB card, worth
# NOT defaulting to on: it trades VRAM for speed and nothing here has measured
# the trade. Per-job overridable.
DEFAULT_MAX_BATCH_SIZE = int(os.environ.get("LTX_MAX_BATCH_SIZE", 1))

# libx264 settings for the delivered mp4. Upstream's encode_video defaults are
# crf 19 / veryfast; kept as the defaults here so nothing changes silently, and
# both are per-job overridable. Lower crf = better and bigger; a slower preset
# buys compression at the same crf, and the encode is CPU work measured in
# seconds against minutes of GPU, so it is nearly free to slow it down.
DEFAULT_OUTPUT_CRF = int(os.environ.get("LTX_OUTPUT_CRF", 19))
DEFAULT_OUTPUT_PRESET = os.environ.get("LTX_OUTPUT_PRESET", "veryfast")

# RetakePipeline's own `distilled` switch. True keeps retake on the distilled
# transformer (no LoRA, fixed schedule, no guidance) -- the cheapest phase-2
# mode. False runs it guided, at which point the negative prompt and step count
# this handler already sends start being used.
RETAKE_DISTILLED = os.environ.get("LTX_RETAKE_DISTILLED", "1") == "1"

# DiffVAE video-decoder backend (ltx_core.model.video_vae.transformer.DiffVAEMode):
#   chunked_eager (default) | chunked_compile | combined_compile | blackwell_dsl
#
# blackwell_dsl is the interesting one on this hardware -- upstream's own tiling
# table gives it a stage-5 memory coefficient of 2.5 against chunked_eager's 5,
# i.e. HALF the decode working set, which is exactly what runs out at 481
# frames. It is NOT the default here for one concrete reason: apply_dsl.py
# raises unless `ltx_kernels.vae` is importable, and that package is the
# `--group kernels` build (nvcc + cccl, compiled from source) which this
# Dockerfile deliberately does not install against RunPod's 30-minute build cap.
# Left as an env switch so enabling it later is a rebuild, not a code change.
LTX_DIFFVAE_MODE = os.environ.get("LTX_DIFFVAE_MODE", "chunked_eager")

# The duration head, so `num_frames` MAY be left to the model. We always send it
# explicitly (billing is per requested duration), but resolving the path costs
# 3.8 MB and removes a redeploy if auto-duration is ever wanted -- upstream's
# require_num_frames_source() raises without it.
DURATION_HEAD_FILE = os.environ.get(
    "LTX_DURATION_HEAD_FILE", "model_patches/ltx-2.5-duration-head-bf16.safetensors",
)

# OffloadMode.NONE/CPU/DISK -- confirmed against the real enum in
# ltx_pipelines/utils/types.py. NONE is fastest but needs the most VRAM;
# CPU/DISK trade speed for headroom. Left at NONE by default since this
# endpoint is deployed on the Blackwell/96GB tier.
LTX_OFFLOAD_MODE = os.environ.get("LTX_OFFLOAD_MODE", "none")

# One of QuantizationKind's real values (utils/quantization_factory.py):
# fp8-cast / fp8-scaled-mm / nvfp4-cast / nvfp4-prequant, or "" to disable.
#
# fp8-cast is the default because it is what the 2.3 deployment ran
# successfully and it needs nothing beyond the bf16 checkpoint already being
# downloaded. nvfp4-prequant is the interesting one for COGS -- 18.72GB of
# weights instead of 42.02GB -- but it requires (a) LTX_TRANSFORMER_FILE
# pointed at the *-nvfp4.safetensors file, (b) Blackwell SM>=10, and (c) the
# `ltx-kernels` build group, which the Dockerfile deliberately does not
# install by default. Its output quality versus fp8-cast has never been
# compared on our product photos. Do not flip this without doing that.
LTX_QUANTIZATION = os.environ.get("LTX_QUANTIZATION", "fp8-cast")

# --- Generation defaults -----------------------------------------------
# Only ever used for a hand-made payload that omits a field; shoupapp's
# Django side always sends all of these explicitly.
DEFAULT_HEIGHT = int(os.environ.get("LTX_DEFAULT_HEIGHT", 1536))
DEFAULT_WIDTH = int(os.environ.get("LTX_DEFAULT_WIDTH", 1152))
DEFAULT_NUM_FRAMES = int(os.environ.get("LTX_DEFAULT_NUM_FRAMES", 121))
DEFAULT_FPS = float(os.environ.get("LTX_DEFAULT_FPS", 24.0))
DEFAULT_SEED = int(os.environ.get("LTX_DEFAULT_SEED", 10))

# DistilledPipeline.__call__'s own default is False -- matched here rather
# than the demo UI's default (True), since enabling it means an extra Gemma
# pass to rewrite the prompt before generation starts, and this handler's
# caller (shoupapp) sends hand-crafted, detailed prompts rather than short
# ones meant to be expanded. Overridable per-job.
DEFAULT_ENHANCE_PROMPT = os.environ.get("LTX_DEFAULT_ENHANCE_PROMPT", "0") == "1"

# LTX generates audio synchronized to the video natively -- the pipeline call
# always produces an `audio` return value and there is no flag to skip
# computing it. This only controls whether _run_one() passes that audio
# through to encode_video() or discards it. For a silent product photo there
# is no real sound event to condition on, so the model's audio output is
# closer to ambient noise than something clearly desirable -- off by default,
# overridable per-request since that judgment may not hold for every prompt.
DEFAULT_INCLUDE_AUDIO = os.environ.get("LTX_DEFAULT_INCLUDE_AUDIO", "0") == "1"

# --- MODES ------------------------------------------------------------------
# One endpoint, four pipelines, selected by the payload's `mode`. Absent means
# "i2v", i.e. exactly what this handler did before phase 2, so every existing
# caller keeps working without sending anything new.
#
# WHY ONE ENDPOINT AND NOT FOUR: a pipeline object holds FILE PATHS, not
# weights. Proven by a real 2026-08-25 worker log, where the transformer is
# built twice inside a single job -- once per stage -- because gpu_model() moves
# every parameter to the `meta` device when its context exits. So keeping four
# pipelines resident costs almost no VRAM, and premium/keyframe share the same
# dev transformer and LoRA with each other. What it costs is TIME: 8 fixed
# distilled steps versus 15-30 CFG-guided ones. Price accordingly (the Django
# side multiplies credits per mode), do not budget it as memory.
MODE_I2V = "i2v"
MODE_PREMIUM = "premium"
MODE_KEYFRAME = "keyframe"
MODE_RETAKE = "retake"
# DFRPipeline: full-resolution stage 2 + generated keyframe slots. See the
# TEMPORAL_UPSCALER_FILE block above for why it exists.
MODE_DFR = "dfr"
MODE_REFINE = "refine"
MODE_MOTION = "motion"
MODE_UNION = "union"
CONTROL_LORA_FILES = {
    MODE_MOTION: "loras/ltx-2.3-22b-ic-lora-motion-track-control-ref0.5.safetensors",
    MODE_UNION: "loras/ltx-2.3-22b-ic-lora-union-control-ref0.5.safetensors",
}
REFINE_LORA_FILE = "loras/ltx-2.5-22b-ic-lora-refine-details-1.0.safetensors"
# Not a render: reports what this worker actually is. See _selftest().
MODE_SELFTEST = "selftest"
VALID_MODES = (MODE_I2V, MODE_PREMIUM, MODE_KEYFRAME, MODE_RETAKE, MODE_DFR, MODE_REFINE, MODE_MOTION, MODE_UNION, MODE_SELFTEST)

# mode -> built pipeline. Lazily populated: a worker that only ever serves i2v
# never downloads the 42 GB dev transformer.
_PIPES: dict = {}
_PATHS = None
# Set by _resolve_paths(); read by _run_one() to degrade instead of failing a paid job.
_ENHANCER_ROOT = None


def _hf_token() -> str | None:
    return os.environ.get("HF_TOKEN")


def _resolve_paths():
    """Locate every shared component once and cache the result.

    Split out of _get_pipe() when the second pipeline arrived: all four modes
    need the same text encoder, VAEs and upsampler, and re-running the resolution
    per mode would re-print the config block and re-do the enhancer lookup for
    no reason. The dev transformer and the distilled LoRA are deliberately NOT
    fetched here -- see _phase2_paths().
    """
    global _PATHS, _ENHANCER_ROOT
    if _PATHS is not None:
        return _PATHS

    from huggingface_hub import snapshot_download

    # NO local_dir= on ANY download here. local_dir= bypasses the standard HF
    # cache layout, which is exactly where RunPod's cached-model feature
    # pre-provisions files -- at /runpod-volume/huggingface-cache/hub, unbilled
    # ("You aren't billed for worker time while your model is being downloaded",
    # docs.runpod.io/serverless/endpoints/model-caching). Default cache
    # resolution (HF_HUB_CACHE, set at this file's top when /runpod-volume
    # exists) finds those pre-provisioned files and downloads nothing. Without
    # the Model field it still works, just downloading into that same cache path
    # on billed time -- which for this pack is ~71GB, so configure the Model
    # field.
    transformer_path = _fetch(TRANSFORMER_FILE)
    text_encoder_path = _fetch(TEXT_ENCODER_FILE)
    video_vae_path = _fetch(VIDEO_VAE_FILE)
    audio_vae_path = _fetch(AUDIO_VAE_FILE)
    upsampler_path = _fetch(UPSCALER_FILE)

    # Optional, and never fatal: a missing duration head only means
    # auto-duration is unavailable, and we always send num_frames.
    duration_head_path = None
    if DURATION_HEAD_FILE:
        try:
            duration_head_path = _fetch(DURATION_HEAD_FILE)
        except Exception as exc:
            print(f"!! duration head unavailable ({type(exc).__name__}: {exc}); "
                  "auto-duration will be refused, explicit num_frames still works", flush=True)

    # allow_patterns is NOT optional here: a bare snapshot_download on the
    # bundle repo would pull all ~165GB. This fetches only the enhancer
    # subfolder (~24GB), and on the endpoint finds it already provisioned by
    # the Model field, downloading nothing.
    if PROMPT_ENHANCER_SUBFOLDER:
        try:
            snap = snapshot_download(
                repo_id=MODEL_REPO, token=_hf_token(),
                revision=os.environ.get("LTX_MODEL_REVISION") or None,
                local_files_only=os.path.isdir(_RUNPOD_HF_CACHE),
                allow_patterns=[f"{PROMPT_ENHANCER_SUBFOLDER}/*"],
            )
            candidate = os.path.join(snap, PROMPT_ENHANCER_SUBFOLDER)
            if os.path.isdir(candidate):
                _ENHANCER_ROOT = candidate
            else:
                print(
                    f"!! prompt enhancer DISABLED: {PROMPT_ENHANCER_SUBFOLDER!r} is not a "
                    f"directory inside {MODEL_REPO} (looked at {candidate}). Jobs asking for "
                    "enhance_prompt will run WITHOUT enhancement rather than fail.",
                    flush=True,
                )
        except Exception as exc:
            print(
                f"!! prompt enhancer DISABLED: could not resolve {PROMPT_ENHANCER_SUBFOLDER!r} "
                f"from {MODEL_REPO} ({type(exc).__name__}: {exc}). Jobs asking for "
                "enhance_prompt will run WITHOUT enhancement rather than fail.",
                flush=True,
            )

    # One block, printed once per worker, naming every resolved input. Both
    # real failures of this migration -- a missing LTX_MODEL_REPO silently
    # pointing at upstream, and a missing prompt enhancer -- would have been
    # visible in the first log line instead of costing a job each.
    print(
        "[LTX config] repo=%s\n"
        "  transformer   = %s\n  text_encoder  = %s\n  video_vae     = %s\n"
        "  audio_vae     = %s\n  upsampler     = %s\n  enhancer_root = %s\n"
        "  quantization  = %s | offload = %s"
        % (MODEL_REPO, transformer_path, text_encoder_path, video_vae_path,
           audio_vae_path, upsampler_path, _ENHANCER_ROOT or "DISABLED",
           LTX_QUANTIZATION or "off", LTX_OFFLOAD_MODE),
        flush=True,
    )

    _PATHS = {
        "transformer": transformer_path,
        "text_encoder": text_encoder_path,
        "video_vae": video_vae_path,
        "audio_vae": audio_vae_path,
        "upsampler": upsampler_path,
        "duration_head": duration_head_path,
    }
    return _PATHS


def _fetch(filename: str) -> str:
    from huggingface_hub import hf_hub_download
    print(f"Downloading/locating {filename}", flush=True)
    return hf_hub_download(
        repo_id=MODEL_REPO, filename=filename, token=_hf_token(),
        revision=os.environ.get("LTX_MODEL_REVISION") or None,
        local_files_only=os.path.isdir(_RUNPOD_HF_CACHE),
    )


def _phase2_paths():
    """The dev transformer + distilled LoRA, fetched on FIRST USE only.

    Kept out of _resolve_paths() on purpose: they are 51 GB together, and a
    worker that only ever serves i2v must not pay for locating them. On the
    endpoint the Model field has already provisioned the whole bundle, so this
    is a cache lookup, not a download -- but that is a property of the
    deployment, not something this code should assume.
    """
    paths = _resolve_paths()
    if "dev_transformer" not in paths:
        paths["dev_transformer"] = _fetch(DEV_TRANSFORMER_FILE)
        paths["distilled_lora"] = _fetch(DISTILLED_LORA_FILE)
        print(
            "[LTX config] phase-2 weights resolved\n"
            "  dev_transformer = %s\n  distilled_lora  = %s"
            % (paths["dev_transformer"], paths["distilled_lora"]),
            flush=True,
        )
    return paths


def _model_paths(transformer_path):
    """ModelPaths for one transformer + the shared components.

    duration_head_path IS filled in (changed 2026-08-28). We still send
    num_frames on every job -- credits are billed per requested duration, so the
    frame count must come from the caller, not from the prompt -- but having the
    path resolved is what makes `auto_duration` a payload field instead of a
    redeploy. Without it `require_num_frames_source()` refuses the request.
    """
    from ltx_pipelines.utils.model_paths import ModelPaths
    paths = _resolve_paths()
    return ModelPaths.from_split(
        transformer_path=transformer_path,
        text_encoder_path=paths["text_encoder"],
        video_vae_path=paths["video_vae"],
        audio_vae_path=paths["audio_vae"],
        duration_head_path=paths.get("duration_head"),
    )


def _quantization(transformer_path, kind=None):
    from ltx_pipelines.utils.quantization_factory import QuantizationKind
    kind = LTX_QUANTIZATION if kind is None else kind
    if not kind:
        return None
    return QuantizationKind(kind).to_policy(transformer_path)


def _distilled_lora(paths):
    """The LoRA tuple every non-distilled pipeline takes.

    LTXV_LORA_COMFY_RENAMING_MAP is not decoration: it is the sd_ops upstream's
    own LoraAction attaches to every --lora / --distilled-lora it parses, and it
    is what maps the published (Comfy-named) LoRA tensors onto this codebase's
    parameter names. Building the tuple without it loads a LoRA that silently
    matches nothing.

    The strength here is ignored by TI2VidTwoStagesHQPipeline, which rebuilds the
    tuple twice at its own per-stage strengths; KeyframeInterpolationPipeline
    uses it as given, and DEFAULT_LORA_STRENGTH (1.0) is upstream's default.
    """
    from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps
    from ltx_pipelines.utils.constants import DEFAULT_LORA_STRENGTH
    return [LoraPathStrengthAndSDOps(paths["distilled_lora"], DEFAULT_LORA_STRENGTH,
                                     LTXV_LORA_COMFY_RENAMING_MAP)]


def _build_pipe(mode):
    """Construct one pipeline. Argument lists mirror the matching upstream
    main() at the pinned commit, component for component."""
    from ltx_core.model.video_vae.transformer import DiffVAEMode
    from ltx_pipelines.utils.types import OffloadMode

    offload = OffloadMode(LTX_OFFLOAD_MODE)
    # Constructed here rather than passed as a string: an unknown value must
    # fail at pipeline build with the enum's own error, not silently fall back.
    diffvae = DiffVAEMode(LTX_DIFFVAE_MODE)

    if mode == MODE_I2V:
        from ltx_pipelines.distilled import DistilledPipeline
        paths = _resolve_paths()
        return DistilledPipeline(
            model_paths=_model_paths(paths["transformer"]),
            spatial_upsampler_path=paths["upsampler"],
            loras=(),
            quantization=_quantization(paths["transformer"]),
            offload_mode=offload,
            prompt_enhancer_gemma_root=_ENHANCER_ROOT,
            diffvae_optimization=diffvae,
        )

    if mode == MODE_PREMIUM:
        from ltx_pipelines.ti2vid_two_stages_hq import TI2VidTwoStagesHQPipeline
        paths = _phase2_paths()
        return TI2VidTwoStagesHQPipeline(
            model_paths=_model_paths(paths["dev_transformer"]),
            distilled_lora=_distilled_lora(paths),
            distilled_lora_strength_stage_1=HQ_LORA_STRENGTH_STAGE_1,
            distilled_lora_strength_stage_2=HQ_LORA_STRENGTH_STAGE_2,
            spatial_upsampler_path=paths["upsampler"],
            loras=(),
            quantization=_quantization(paths["dev_transformer"]),
            offload_mode=offload,
            prompt_enhancer_gemma_root=_ENHANCER_ROOT,
            diffvae_optimization=diffvae,
        )

    if mode == MODE_KEYFRAME:
        from ltx_pipelines.keyframe_interpolation import KeyframeInterpolationPipeline
        paths = _phase2_paths()
        return KeyframeInterpolationPipeline(
            model_paths=_model_paths(paths["dev_transformer"]),
            distilled_lora=_distilled_lora(paths),
            spatial_upsampler_path=paths["upsampler"],
            loras=[],
            quantization=_quantization(paths["dev_transformer"]),
            offload_mode=offload,
            prompt_enhancer_gemma_root=_ENHANCER_ROOT,
            diffvae_optimization=diffvae,
        )

    if mode == MODE_DFR:
        from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps
        from ltx_pipelines.dfr_pipeline import DFRPipeline
        paths = _resolve_paths()
        detailing = None
        if DETAILING_LORA_FILE:
            detailing = [LoraPathStrengthAndSDOps(
                _fetch(DETAILING_LORA_FILE), DETAILING_LORA_STRENGTH, LTXV_LORA_COMFY_RENAMING_MAP)]
        return DFRPipeline(
            model_paths=_model_paths(paths["transformer"]),
            spatial_upsampler_path=paths["upsampler"],
            loras=[],
            detailing_lora=detailing,
            # Located on first DFR use, like the dev transformer. Without it a
            # job asking for temporal_upsample_rounds fails with upstream's own
            # error instead of rendering at the wrong frame rate.
            temporal_upsampler_path=_fetch(TEMPORAL_UPSCALER_FILE) if TEMPORAL_UPSCALER_FILE else None,
            quantization=_quantization(paths["transformer"], LTX_DFR_QUANTIZATION),
            offload_mode=offload,
            prompt_enhancer_gemma_root=_ENHANCER_ROOT,
            diffvae_optimization=diffvae,
        )

    if mode in (MODE_REFINE, MODE_MOTION, MODE_UNION):
        from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps
        from ltx_pipelines.ic_lora import ICLoraPipeline, ICLoraStageConfig
        from ltx_pipelines.utils.constants import DISTILLED_SIGMA_VALUES
        from ltx_pipelines.utils.tiled_diffusion import FixedSizeSpatialTiling
        paths = _resolve_paths()
        # shortcut: native trapezoidal tile blending differs from ComfyUI's
        # Gaussian TiledFusion; compare GPU outputs before making this a default.
        stages = [ICLoraStageConfig(
            resolution_divisor=1, apply_ic_lora=True, sigmas=tuple(DISTILLED_SIGMA_VALUES),
            tiling=FixedSizeSpatialTiling(height=576, width=1024,
                                          height_overlap=288, width_overlap=512)
                   if mode == MODE_REFINE else None)]
        lora_file = REFINE_LORA_FILE if mode == MODE_REFINE else CONTROL_LORA_FILES[mode]
        return ICLoraPipeline(
            model_paths=_model_paths(paths["transformer"]),
            spatial_upsampler_path=paths["upsampler"],
            loras=[LoraPathStrengthAndSDOps(_fetch(lora_file), 1.0,
                                           LTXV_LORA_COMFY_RENAMING_MAP)],
            stages=stages, quantization=_quantization(paths["transformer"]),
            offload_mode=offload, diffvae_optimization=diffvae,
            prompt_enhancer_gemma_root=_ENHANCER_ROOT,
        )

    if mode == MODE_RETAKE:
        # distilled=True, so this runs on the DISTILLED transformer with no LoRA
        # at all -- the cheapest phase-2 mode, and the reason retake is worth
        # having: regenerating three seconds of an existing clip costs a
        # fraction of re-rendering it.
        from ltx_pipelines.retake import RetakePipeline
        paths = _resolve_paths()
        # distilled=True (the default) runs the DISTILLED transformer with a
        # fixed schedule and no guidance -- the cheapest mode here. Setting
        # LTX_RETAKE_DISTILLED=0 switches it to the guided path, which is when
        # negative_prompt / num_inference_steps / the guider params start
        # mattering; they are sent either way so that switch stays an env
        # change rather than a payload change.
        return RetakePipeline(
            model_paths=_model_paths(paths["transformer"]),
            loras=(),
            quantization=_quantization(paths["transformer"]),
            distilled=RETAKE_DISTILLED,
            offload_mode=offload,
            prompt_enhancer_gemma_root=_ENHANCER_ROOT,
            diffvae_optimization=diffvae,
        )

    raise ValueError(f"unknown mode {mode!r}; expected one of {VALID_MODES}")


def _get_pipe(mode=MODE_I2V):
    if mode not in VALID_MODES:
        raise ValueError(f"unknown mode {mode!r}; expected one of {VALID_MODES}")
    if mode not in _PIPES:
        print(f"[LTX] building the {mode!r} pipeline", flush=True)
        _PIPES[mode] = _build_pipe(mode)
    return _PIPES[mode]



def _num(job, key, default, cast):
    """Read a numeric job field, treating an explicit JSON null exactly like
    an absent key. See the call site in _run_one() for why this exists."""
    value = job.get(key)
    return cast(default if value is None else value)


# EVERY key this handler understands, per mode. An unknown key is an ERROR, not
# a shrug -- see _validate_job_keys. The whole reason this exists is that the
# previous handler ignored anything it did not read, which is how a payload can
# ask for premium quality, get a standard render, and still be billed as
# premium: silence looks identical to success.
_COMMON_KEYS = {
    # identity / routing -- video_job_id and segment_index are sent by Django and
    # deliberately unread here (they travel in the webhook URL), but they are
    # legal, so they belong in this set rather than tripping the guard.
    "job_id", "video_job_id", "segment_index", "mode",
    # read by every mode
    "prompt", "seed", "fps", "include_audio", "enhance_prompt",
    "enhance_static_cache", "color_space",
    # delivery (encode_video)
    "output_crf", "output_preset",
    # num_frames is read before the mode branch: the generative modes pass it to
    # the pipeline, retake uses it only for the caller's own timeout sizing.
    "num_frames",
}

# Per mode, and derived from each pipeline's REAL __call__ signature -- not from
# what looks harmless. `video_guider` on an i2v job is the case that matters:
# DistilledPipeline has no guidance at all, so accepting it would mean a caller
# paying for guidance that was never applied. That is the same silent-success
# failure the whole guard exists to prevent, so an unusable field is an error in
# the mode that cannot use it.
_CONDITIONING_KEYS = {"image_url", "images"}
_SIZE_KEYS = {"width", "height"}
_GUIDED_KEYS = {"negative_prompt", "num_inference_steps", "video_guider",
                "audio_guider", "max_batch_size"}
_TWO_STAGE_SIGMA_KEYS = {"stage_1_sigmas", "stage_2_sigmas"}

_MODE_KEYS = {
    # DistilledPipeline: two distilled schedules, generated keyframe slots, and
    # auto-duration via the duration head. No guidance, no step count.
    MODE_I2V: _CONDITIONING_KEYS | _SIZE_KEYS | _TWO_STAGE_SIGMA_KEYS
              | {"generated_keyframes", "decode_with_keyframes", "auto_duration"},
    # TI2VidTwoStagesHQPipeline: everything distilled has, plus guidance.
    MODE_PREMIUM: _CONDITIONING_KEYS | _SIZE_KEYS | _TWO_STAGE_SIGMA_KEYS | _GUIDED_KEYS
                  | {"generated_keyframes", "decode_with_keyframes", "auto_duration"},
    # KeyframeInterpolationPipeline: guided, but takes a plain int frame count
    # (no duration head) and has NO generated_keyframes parameter.
    MODE_KEYFRAME: _CONDITIONING_KEYS | _SIZE_KEYS | _TWO_STAGE_SIGMA_KEYS | _GUIDED_KEYS,
    # RetakePipeline: a source clip and a window. Its schedule argument is
    # `sigmas` (one stage), not the two-stage pair.
    MODE_RETAKE: _GUIDED_KEYS | {"video_url", "start_time", "end_time",
                                 "regenerate_video", "sigmas"},
    # DFRPipeline at the pinned commit: no guidance, no negative prompt, no
    # generated_keyframes field (it derives its own slots) -- and one field of
    # its own, the number of frame-rate doublings.
    MODE_DFR: _CONDITIONING_KEYS | _SIZE_KEYS | _TWO_STAGE_SIGMA_KEYS
              | {"temporal_upsample_rounds", "spatial_upscalings", "auto_duration"},
    MODE_REFINE: _SIZE_KEYS | {"video_url"},
    # Reference video must already contain sparse tracks or canny/depth/pose.
    MODE_MOTION: _CONDITIONING_KEYS | _SIZE_KEYS | {"video_url", "conditioning_strength"},
    MODE_UNION: _CONDITIONING_KEYS | _SIZE_KEYS | {"video_url", "conditioning_strength"},
    MODE_SELFTEST: {"resolve_models", "check_gcs"},
}


def _validate_job_keys(job, mode):
    """Refuse a payload carrying anything this handler does not act on.

    Deliberately strict, and deliberately per mode: `negative_prompt` on an i2v
    job is not a harmless extra, it is a caller believing it bought guidance
    from a pipeline that has none. Better a failed job than a quiet one.
    """
    allowed = _COMMON_KEYS | _MODE_KEYS.get(mode, set())
    unknown = sorted(set(job) - allowed)
    if unknown:
        raise ValueError(
            f"unknown payload field(s) for mode {mode!r}: {unknown}. "
            f"Accepted: {sorted(allowed)}"
        )



def _control_strength(job):
    if not isinstance(job.get("video_url"), str) or not job["video_url"].strip():
        raise ValueError("control needs a prepared guide video_url")
    value = float(job.get("conditioning_strength", 1.0))
    if not 0.0 <= value <= 1.0:
        raise ValueError("conditioning_strength must be finite and in [0, 1]")
    return value


def _image_specs(job):
    """Normalise a job's image conditioning into [{url, frame_idx, strength}].

    Two accepted shapes, and the single-image one is not deprecated: it is what
    /image-to-video/, the photo editor and the video editor all send, and
    "animate this photo" genuinely has one image at frame 0.

        {"image_url": "..."}                       -> [frame 0, strength 1.0]
        {"images": [{"url": ..., "frame_idx": 30, "strength": 1.0}, ...]}

    Defaults per entry: frame_idx 0, strength 1.0 -- the same pinning the single
    image gets, so an entry that omits them behaves like the old payload rather
    than like a soft hint. Nothing upstream validates keyframe count, spacing or
    frame_idx bounds (checked in ImageAction, VideoConditionByKeyframeIndex and
    combined_image_conditionings), so the caller owns that mapping; this only
    fills in shape.
    """
    entries = job.get("images")
    if entries:
        specs = []
        for entry in entries:
            if isinstance(entry, str):
                entry = {"url": entry}
            url = entry.get("url") or entry.get("image_url")
            if not url:
                raise ValueError("every images[] entry needs a url")
            specs.append({
                "url": url,
                "frame_idx": entry.get("frame_idx") or 0,
                "strength": 1.0 if entry.get("strength") is None else entry["strength"],
            })
        return specs
    if job.get("image_url"):
        return [{"url": job["image_url"], "frame_idx": 0, "strength": 1.0}]
    return []


def _steps_for(job, default):
    """Denoising steps for a guided mode, clamped to something payable.

    `default` is upstream's own number for that pipeline -- 15 for HQ
    (LTX_2_3_HQ_PARAMS, a plain constant because the HQ preset overrides every
    generation-dependent knob) and whatever detect_params() reads off the
    checkpoint for keyframe interpolation (30 on a 2.5 checkpoint). Not
    hardcoded here, so a checkpoint change moves it.
    """
    value = job.get("num_inference_steps")
    if value is None:
        return int(default)
    return max(1, min(MAX_INFERENCE_STEPS, int(value)))


def _guider_params(job, key, base):
    """One MultiModalGuiderParams, starting from the pipeline's own defaults and
    applying only the fields the payload actually names.

    The base comes from upstream (LTX_2_3_HQ_PARAMS for premium,
    detect_params(checkpoint) for keyframe/retake) -- never invented here. The
    override exists so guidance can be tuned from Django without a redeploy;
    a payload that names nothing gets exactly what upstream's own main() uses.

    Field names are the dataclass's real ones: cfg_scale, stg_scale, stg_blocks,
    rescale_scale, modality_scale, skip_step.
    """
    from dataclasses import replace
    overrides = job.get(key) or {}
    if not isinstance(overrides, dict) or not overrides:
        return base
    allowed = {"cfg_scale", "stg_scale", "stg_blocks", "rescale_scale",
               "modality_scale", "skip_step"}
    unknown = set(overrides) - allowed
    if unknown:
        raise ValueError(f"{key}: unknown field(s) {sorted(unknown)}; expected {sorted(allowed)}")
    clean = {}
    for field, value in overrides.items():
        if value is None:
            continue
        if field == "stg_blocks":
            clean[field] = [int(b) for b in value]
        elif field == "skip_step":
            clean[field] = int(value)
        else:
            clean[field] = float(value)
    return replace(base, **clean) if clean else base


def _sigmas(job, key):
    """An explicit sigma schedule as a list of floats, or None for the
    pipeline's own. Exposed because a schedule is the one quality knob that is
    neither a step count nor a guidance scale, and hardcoding it here would make
    every experiment a redeploy. Validated as descending, in [0, 1]."""
    values = job.get(key)
    if not values:
        return None
    floats = [float(v) for v in values]
    if len(floats) < 2:
        raise ValueError(f"{key}: needs at least two sigmas, got {floats}")
    if any(not (0.0 <= v <= 1.0) for v in floats):
        raise ValueError(f"{key}: sigmas must be within [0, 1], got {floats}")
    if any(a < b for a, b in zip(floats, floats[1:])):
        raise ValueError(f"{key}: sigmas must be non-increasing, got {floats}")
    return torch.tensor(floats)


def _color_space(job):
    """HDR colour space by name, or None (SDR).

    Only meaningful with scene-linear EXR conditioning input; every image this
    app sends is an SDR JPEG/PNG, so this is None in practice. Present because
    it is a real upstream parameter and adding it later would be a redeploy.
    """
    name = (job.get("color_space") or "").strip()
    if not name:
        return None
    from ltx_pipelines.utils.media_io import EXRColorSpace
    try:
        return EXRColorSpace(name.lower())
    except ValueError:
        valid = [c.value for c in EXRColorSpace]
        raise ValueError(f"color_space {name!r} is not one of {valid}")


def _generated_keyframes(job):
    """`generated_keyframes`: an int (how many slots) or an explicit sequence of
    frame indices. Upstream gates it on the transformer config
    (assert_generated_keyframes_supported), so an unsupported checkpoint fails
    loudly rather than ignoring it."""
    value = job.get("generated_keyframes")
    if value is None:
        return 0
    if isinstance(value, (list, tuple)):
        return [int(v) for v in value]
    return int(value)


def _num_frames_arg(job, explicit):
    """Either the caller's frame count, or an AutoDuration window when the
    payload asks for one.

    Auto-duration is OFF by default and should stay that way for billed work:
    the customer is charged per requested duration, so the length cannot be the
    model's choice. It exists because the duration head is in the bundle and
    refusing it in code would mean a redeploy to change our minds.
    """
    auto = job.get("auto_duration")
    if not auto:
        return explicit
    from ltx_pipelines.utils.types import AutoDuration
    if isinstance(auto, dict):
        return AutoDuration(
            min_seconds=float(auto.get("min_seconds", 1.0)),
            max_seconds=float(auto.get("max_seconds", 20.0)),
        )
    return AutoDuration()


def _params_for(checkpoint_path):
    """PipelineParams detected from the checkpoint -- steps and both guider
    parameter sets. detect_params() is upstream's own resolution and is why
    `crf` came out right on 2.5 (it reads model_version from the safetensors
    metadata); guider values must come from here rather than be hardcoded."""
    from ltx_pipelines.utils.constants import detect_params
    return detect_params(checkpoint_path)


class _LazyHQParams:
    """LTX_2_3_HQ_PARAMS, imported on first use.

    A module-level `from ltx_pipelines...import LTX_2_3_HQ_PARAMS` would drag
    torch and the whole pipelines package into import time, which this handler
    deliberately avoids -- every heavy import in this file sits inside the
    function that needs it.
    """

    def __getattr__(self, item):
        from ltx_pipelines.utils.constants import LTX_2_3_HQ_PARAMS
        return getattr(LTX_2_3_HQ_PARAMS, item)


HQ_PARAMS = _LazyHQParams()


def _download_image(url: str) -> Image.Image:
    resp = requests.get(url, timeout=30)
    resp.raise_for_status()
    return Image.open(io.BytesIO(resp.content)).convert("RGB")


def _download_file(url: str, dest_path: str) -> None:
    """Stream a (potentially large) file to disk -- the retake mode's source
    clip, which is a video and must not be held in memory the way
    _download_image holds a still."""
    resp = requests.get(url, stream=True, timeout=120)
    resp.raise_for_status()
    with open(dest_path, "wb") as f:
        for chunk in resp.iter_content(chunk_size=1 << 20):
            f.write(chunk)


_gcs_bucket = None


def _get_gcs_bucket():
    """Same JSON_PATH_KEY / GCS_BUCKET_NAME convention as
    ../Wan2.2-runpod/handler.py -- including the file-path-OR-raw-JSON
    detection that repo's 2026-07-30 real deploy failure required (RunPod
    Secrets inject the key's JSON content directly into the env var, not a
    path to a file, and this endpoint has no Network Volume either -- same
    reasoning applies here, copied rather than re-derived)."""
    global _gcs_bucket
    if _gcs_bucket is None:
        import json
        from google.cloud import storage
        from google.oauth2 import service_account
        raw = os.environ["JSON_PATH_KEY"]
        bucket_name = os.environ.get("GCS_BUCKET_NAME", "shoupapp")
        try:
            key_info = json.loads(raw)
        except json.JSONDecodeError:
            credentials = service_account.Credentials.from_service_account_file(raw)
        else:
            credentials = service_account.Credentials.from_service_account_info(key_info)
        client = storage.Client(credentials=credentials)
        _gcs_bucket = client.get_bucket(bucket_name, timeout=60)
    return _gcs_bucket


def _upload_video(local_path: str) -> str:
    blob_name = f"videos/runpod-ltx/{uuid.uuid4().hex}.mp4"
    blob = _get_gcs_bucket().blob(blob_name)
    blob.upload_from_filename(local_path, content_type="video/mp4")
    return blob.public_url


# @torch.inference_mode() is REQUIRED, not an optimization -- its absence
# caused the real 2026-07-31 OOM (94.04 GiB allocated on a 94.97 GiB GPU,
# died inside denoising step 1/8): without it, autograd retains every
# intermediate activation of the 22B transformer's forward pass for a
# backward that never comes. Every upstream entry point wraps inference this
# way -- distilled.py's own main() carries the same decorator at the pinned
# commit, and multigpu/runner.py states outright that pipelines do not wrap
# themselves, the caller must. Decorating _run_one (not just the pipe() call)
# is deliberate: pipe() returns a LAZY video iterator whose actual VAE decode
# runs inside encode_video() below, so the decode must be covered too
# (community confirmation: huggingface.co/Lightricks/LTX-2/discussions/47 --
# undecorated tiled VAE decode builds autograd graphs across all tiles).
def _selftest(job):
    """Report the worker instead of rendering on it. Costs seconds, no GPU work.

    WHY IT EXISTS: everything this handler decides -- offload mode, quantization,
    whether the Blackwell DSL decoder is even importable, whether natten is
    installed, how much VRAM a 481-frame decode may use -- depends on facts about
    the box that are otherwise only visible in the RunPod console, and this
    deployment's API keys are endpoint-scoped and cannot read that console. One
    `{"mode": "selftest"}` job answers all of them from inside the container, so
    a wrong assumption costs a 5-second job instead of a second deploy.

    Reports the PRESENCE of secrets, never their values.
    """
    import importlib.util
    import platform
    import hashlib

    def have(module):
        try:
            return importlib.util.find_spec(module) is not None
        except Exception:
            return False

    info = {
        "job_id": job.get("job_id"),
        "ok": True,
        "mode": MODE_SELFTEST,
        "handler_sha256": hashlib.sha256(open(__file__, "rb").read()).hexdigest(),
        "python": platform.python_version(),
        "torch": torch.__version__,
        "cuda_available": torch.cuda.is_available(),
        "packages": {
            # natten is the DiffVAE decoder's attention backend -- without it the
            # chain falls back to Triton, which throws CUDA: invalid argument at
            # 481 frames. This is the single most important line here.
            "natten": have("natten"),
            # ltx_kernels gates DiffVAEMode.BLACKWELL_DSL (half the decode
            # working set) and nvfp4-prequant. Built by `uv sync --group kernels`.
            "ltx_kernels": have("ltx_kernels"),
            "flash_attn": have("flash_attn"),
        },
        "config": {
            "model_repo": MODEL_REPO,
            "upstream_commit": os.environ.get("LTX_UPSTREAM_COMMIT"),
            "model_revision": os.environ.get("LTX_MODEL_REVISION") or "main",
            "cache_only": os.path.isdir(_RUNPOD_HF_CACHE),
            "supported_modes": list(_MODE_KEYS),
            "offload_mode": LTX_OFFLOAD_MODE,
            "quantization": LTX_QUANTIZATION or "off",
            "diffvae_mode": LTX_DIFFVAE_MODE,
            "transformer_file": TRANSFORMER_FILE,
            "dev_transformer_file": DEV_TRANSFORMER_FILE,
            "distilled_lora_file": DISTILLED_LORA_FILE,
            "temporal_upscaler_file": TEMPORAL_UPSCALER_FILE,
            "detailing_lora_file": DETAILING_LORA_FILE,
            "dfr_quantization": LTX_DFR_QUANTIZATION,
            "output_crf": DEFAULT_OUTPUT_CRF,
            "max_batch_size": DEFAULT_MAX_BATCH_SIZE,
            "default_enhance_prompt": DEFAULT_ENHANCE_PROMPT,
            "skip_upload": os.environ.get("LTX_SKIP_UPLOAD") == "1",
        },
        # Presence only. A missing one here explains a failure that would
        # otherwise only show up as a broken result URL.
        "secrets_present": {
            "HF_TOKEN": bool(os.environ.get("HF_TOKEN")),
            "JSON_PATH_KEY": bool(os.environ.get("JSON_PATH_KEY")),
            "GCS_BUCKET_NAME": bool(os.environ.get("GCS_BUCKET_NAME")),
        },
        "stale_2_3_env_present": sorted(
            k for k in ("LTX_CHECKPOINT_REPO", "LTX_CHECKPOINT_FILE",
                        "LTX_GEMMA_REPO", "LTX_GEMMA_SUBFOLDER")
            if os.environ.get(k)
        ),
    }

    if torch.cuda.is_available():
        props = torch.cuda.get_device_properties(0)
        info["gpu"] = {
            "name": props.name,
            "total_gib": round(props.total_memory / 1024 ** 3, 2),
            "capability": f"{props.major}.{props.minor}",
            # Datacenter Blackwell is sm_100+; BLACKWELL_DSL needs it AND
            # ltx_kernels. Reported so "can we turn DSL on" stops being a guess.
            "blackwell_class": props.major >= 10,
            "multi_processor_count": props.multi_processor_count,
        }

    # Resolving paths also proves the model cache is populated and warms it --
    # a cold worker doing this in a selftest is a download NOT paid for during a
    # customer's render.
    if job.get("resolve_models", True):
        try:
            paths = _resolve_paths()
            info["resolved"] = {k: bool(v) for k, v in paths.items()}
            info["enhancer_root"] = bool(_ENHANCER_ROOT)
            info["control_adapters_cached"] = {mode: bool(_fetch(path))
                                                for mode, path in CONTROL_LORA_FILES.items()}
        except Exception as exc:
            info["ok"] = False
            info["error"] = f"model resolution failed: {type(exc).__name__}: {exc}"

    # GCS is where every result lands; a credential problem here is otherwise
    # only discovered after a full render.
    if job.get("check_gcs", True):
        try:
            bucket = _get_gcs_bucket()
            info["gcs_bucket_reachable"] = bool(bucket.name)
        except Exception as exc:
            info["ok"] = False
            info["gcs_error"] = f"{type(exc).__name__}: {exc}"

    print("[LTX selftest] " + json.dumps(info), flush=True)
    return info


@torch.inference_mode()
def _run_one(job):
    """Render one job, whichever mode it asks for.

    Takes the JOB, not a pipe: the mode is a per-job field, so the pipeline has
    to be selected in here rather than built once by handler(). _get_pipe()
    caches, so a batch of same-mode jobs still builds nothing twice.
    """
    job_id = job.get("job_id")
    mode = (job.get("mode") or MODE_I2V).strip().lower()
    local_path = f"/tmp/{uuid.uuid4().hex}.mp4"
    scratch = []
    try:
        from ltx_core.model.video_vae import AUTO_TILING, get_video_chunks_number
        from ltx_pipelines.utils.args import ImageConditioningInput
        from ltx_pipelines.utils.constants import (
            DEFAULT_NEGATIVE_PROMPT, DISTILLED_SIGMAS, STAGE_2_DISTILLED_SIGMAS,
        )
        from ltx_pipelines.utils.media_io import encode_video

        if mode not in VALID_MODES:
            raise ValueError(f"unknown mode {mode!r}; expected one of {VALID_MODES}")
        _validate_job_keys(job, mode)
        if mode == MODE_SELFTEST:
            return _selftest(job)
        control_strength = _control_strength(job) if mode in CONTROL_LORA_FILES else 1.0
        pipe = _get_pipe(mode)

        # _num() (not a bare int(job.get(k, default))) because an EXPLICIT
        # JSON null must fall back to the default too, not crash. dict.get's
        # default only fires when the key is ABSENT -- a payload carrying
        # {"seed": null} returns None and int(None) raises TypeError. That is
        # not hypothetical: shoupapp's own form leaves `seed` optional and
        # serializes an unset one as null, so every seedless job would have
        # failed here.
        height = _num(job, "height", DEFAULT_HEIGHT, int)
        width = _num(job, "width", DEFAULT_WIDTH, int)
        num_frames = _num(job, "num_frames", DEFAULT_NUM_FRAMES, int)
        frame_rate = _num(job, "fps", DEFAULT_FPS, float)
        seed = _num(job, "seed", DEFAULT_SEED, int)
        # `is not None`, not truthiness -- so an explicit False in a hand-made
        # payload isn't confused with "unset". (Same null-safety concern the
        # _num() helper above solves for the numeric fields.)
        include_audio = bool(job["include_audio"]) if job.get("include_audio") is not None else DEFAULT_INCLUDE_AUDIO
        enhance_prompt = bool(job["enhance_prompt"]) if job.get("enhance_prompt") is not None else DEFAULT_ENHANCE_PROMPT
        # Degrade, never fail. Without an enhancer root, upstream raises
        # ValueError from inside pipe() and the whole render is lost -- on a job
        # the customer has already been charged for. A missing optional feature
        # must not cost a paid render, so the video is produced from the
        # caller's own prompt and the downgrade is logged loudly.
        if enhance_prompt and _ENHANCER_ROOT is None:
            print(
                f"[LTX] job {job_id}: enhance_prompt requested but no enhancer root is "
                "configured -- rendering with the original prompt instead.",
                flush=True,
            )
            enhance_prompt = False

        prompt = job.get("prompt", "")
        negative_prompt = job.get("negative_prompt") or DEFAULT_NEGATIVE_PROMPT
        # Everything below is optional and defaults to what upstream's own
        # main() would use. They exist as payload fields so tuning any of them
        # is a Django change, not a rebuild of this image.
        enhance_static_cache = bool(job.get("enhance_static_cache") or False)
        max_batch_size = int(job.get("max_batch_size") or DEFAULT_MAX_BATCH_SIZE)
        color_space = _color_space(job)
        # Derived, exactly as every upstream main() derives it -- float32 when a
        # colour space is set, bf16 otherwise (vae_dtype_for_hdr). Not a payload
        # field: it is a consequence of color_space, and letting a caller set
        # the two independently is how you get an HDR render decoded at bf16.
        from ltx_pipelines.utils.media_io import vae_dtype_for_hdr
        vae_dtype = vae_dtype_for_hdr(color_space, torch.bfloat16)
        generated_keyframes = _generated_keyframes(job)
        frames_arg = _num_frames_arg(job, num_frames)

        if DEVICE == "cuda":
            torch.cuda.reset_peak_memory_stats()

        # ---------------- retake: a time region of an existing clip ----------
        if mode == MODE_RETAKE:
            from ltx_pipelines.utils.media_io import get_videostream_metadata

            source_video_path = f"/tmp/{uuid.uuid4().hex}-source.mp4"
            scratch.append(source_video_path)
            _download_file(job["video_url"], source_video_path)

            start_time = _num(job, "start_time", 0.0, float)
            end_time = _num(job, "end_time", 0.0, float)
            if start_time >= end_time:
                raise ValueError(f"start_time ({start_time}) must be less than end_time ({end_time})")

            # Upstream's CLI validates both of these before touching the GPU and
            # so does this: a video whose frame count is not 8k+1, or whose
            # dimensions are not multiples of 32, fails deep inside the
            # patchifier with a shape error that says nothing about the input.
            # Our own renders satisfy both by construction; a concatenated one
            # does not, which is one more reason chaining is gone.
            src = get_videostream_metadata(source_video_path)
            if (src.frames - 1) % 8 != 0:
                snapped = ((src.frames - 1) // 8) * 8 + 1
                raise ValueError(
                    f"retake needs a video whose frame count is 8k+1; this one has {src.frames} "
                    f"(nearest valid: {snapped})"
                )
            if src.width % 32 != 0 or src.height % 32 != 0:
                raise ValueError(
                    f"retake needs width and height divisible by 32; this video is {src.width}x{src.height}"
                )

            # Guider params and step count are IGNORED in distilled mode (the
            # pipeline's own docstring says so) and are passed anyway, exactly
            # as upstream's retake main() does -- so flipping LTX to the
            # non-distilled retake later is an env change, not a payload change.
            params = _params_for(_resolve_paths()["transformer"])
            video, audio, out_frames, tiling_config = pipe(
                video_path=source_video_path,
                prompt=prompt,
                start_time=start_time,
                end_time=end_time,
                seed=seed,
                negative_prompt=negative_prompt,
                num_inference_steps=_steps_for(job, params.num_inference_steps),
                video_guider_params=_guider_params(job, "video_guider", params.video_guider_params),
                audio_guider_params=_guider_params(job, "audio_guider", params.audio_guider_params),
                # Two independent switches: replace the picture in the window,
                # the sound in the window, or both. `regenerate_video: false`
                # with audio on is a re-scored clip, which is a real product.
                regenerate_video=bool(job.get("regenerate_video", True)),
                regenerate_audio=include_audio,
                enhance_prompt=enhance_prompt,
                enhance_static_cache=enhance_static_cache,
                tiling_config=AUTO_TILING,
                vae_dtype=vae_dtype,
                max_batch_size=max_batch_size,
                sigmas=_sigmas(job, "sigmas"),
                color_space=color_space,
            )
            frame_rate = float(src.fps)
        elif mode in (MODE_REFINE, MODE_MOTION, MODE_UNION):
            from ltx_pipelines.utils.media_io import get_videostream_metadata
            source_video_path = f"/tmp/{uuid.uuid4().hex}-source.mp4"
            scratch.append(source_video_path)
            _download_file(job["video_url"], source_video_path)
            src = get_videostream_metadata(source_video_path)
            if src.frames != num_frames or abs(float(src.fps) - frame_rate) > 0.01:
                raise ValueError(f"{mode} needs the source video's exact frame count and frame rate")
            if abs(src.width / src.height - width / height) > 0.02:
                raise ValueError(f"{mode} must preserve the source video's aspect ratio")
            images = []
            if mode in CONTROL_LORA_FILES:
                for spec in _image_specs(job):
                    path = f"/tmp/{uuid.uuid4().hex}.png"
                    scratch.append(path)
                    _download_image(spec["url"]).save(path, format="PNG")
                    images.append(ImageConditioningInput(
                        path=path, frame_idx=int(spec["frame_idx"]), strength=float(spec["strength"])))
            video, audio, out_frames, tiling_config = pipe(
                prompt=prompt, seed=seed, height=height, width=width,
                num_frames=num_frames, frame_rate=frame_rate, images=images,
                video_conditioning=[(source_video_path, 1.0)],
                conditioning_attention_strength=control_strength,
                tiling_config=AUTO_TILING, vae_dtype=vae_dtype,
                enhance_prompt=enhance_prompt, enhance_static_cache=enhance_static_cache,
                color_space=color_space,
            )
        else:
            # ---------------- image conditioning ----------------------------
            # `images` (a list) is the 360-degree / keyframe shape; `image_url`
            # is the single-image shape every existing caller sends. Both end up
            # as ImageConditioningInput, which is what all three generative
            # pipelines take.
            #
            # frame_idx is a TARGET FRAME INDEX in pixel frames (upstream's own
            # --image PATH FRAME_IDX STRENGTH help text). combined_image_conditionings()
            # turns frame_idx 0 into the opening-frame latent replacement and
            # every other index into a keyframe conditioning, with no cap on how
            # many -- which is what makes a turntable possible on the pipeline
            # that is already in production.
            images = []
            for spec in _image_specs(job):
                path = f"/tmp/{uuid.uuid4().hex}.png"
                scratch.append(path)
                _download_image(spec["url"]).save(path, format="PNG")
                # crf deliberately NOT passed -- see point 4 of the module
                # docstring: ImageConditioner.resolve_crf reads the right value
                # (18 for 2.5) off the checkpoint.
                images.append(ImageConditioningInput(
                    path=path, frame_idx=int(spec["frame_idx"]), strength=float(spec["strength"]),
                ))
            if not images:
                raise ValueError("no conditioning image: send image_url, or images[]")

            if mode == MODE_I2V:
                # Parameter set and ordering mirror distilled.py's own main() at
                # the pinned commit. num_frames is passed explicitly rather than
                # left to DEFAULT_AUTO_DURATION: the caller bills per requested
                # duration.
                video, audio, out_frames, tiling_config = pipe(
                    prompt=prompt,
                    seed=seed,
                    height=height,
                    width=width,
                    num_frames=frames_arg,
                    frame_rate=frame_rate,
                    images=images,
                    tiling_config=AUTO_TILING,
                    vae_dtype=vae_dtype,
                    enhance_prompt=enhance_prompt,
                    enhance_static_cache=enhance_static_cache,
                    # No negative_prompt, no step count, no guidance: this
                    # pipeline HAS none of them. Sending them would be a payload
                    # that lies. What it does have is the distilled schedules and
                    # generated keyframe slots.
                    stage_1_sigmas=_sigmas(job, "stage_1_sigmas") if job.get("stage_1_sigmas") else DISTILLED_SIGMAS,
                    stage_2_sigmas=_sigmas(job, "stage_2_sigmas") if job.get("stage_2_sigmas") else STAGE_2_DISTILLED_SIGMAS,
                    generated_keyframes=generated_keyframes,
                    decode_with_keyframes=bool(job.get("decode_with_keyframes", False)),
                    color_space=color_space,
                )
            elif mode == MODE_PREMIUM:
                steps = _steps_for(job, HQ_PARAMS.num_inference_steps)
                video, audio, out_frames, tiling_config = pipe(
                    prompt=prompt,
                    negative_prompt=negative_prompt,
                    seed=seed,
                    height=height,
                    width=width,
                    num_frames=frames_arg,
                    frame_rate=frame_rate,
                    num_inference_steps=steps,
                    video_guider_params=_guider_params(job, "video_guider", HQ_PARAMS.video_guider_params),
                    audio_guider_params=_guider_params(job, "audio_guider", HQ_PARAMS.audio_guider_params),
                    images=images,
                    tiling_config=AUTO_TILING,
                    vae_dtype=vae_dtype,
                    enhance_prompt=enhance_prompt,
                    enhance_static_cache=enhance_static_cache,
                    max_batch_size=max_batch_size,
                    # None means "build the schedule from num_inference_steps
                    # with LTX2Scheduler", which is what upstream does.
                    stage_1_sigmas=_sigmas(job, "stage_1_sigmas"),
                    stage_2_sigmas=_sigmas(job, "stage_2_sigmas") or STAGE_2_DISTILLED_SIGMAS,
                    generated_keyframes=generated_keyframes,
                    decode_with_keyframes=bool(job.get("decode_with_keyframes", False)),
                    color_space=color_space,
                )
            elif mode == MODE_DFR:
                # DFR's __call__ at the pinned commit has no vae_dtype and no
                # color_space, so an HDR request cannot be honoured here -- and
                # must not be silently rendered as SDR.
                if color_space is not None:
                    raise ValueError("color_space is not supported in dfr mode")
                rounds = int(job.get("temporal_upsample_rounds") or 0)
                video, audio, out_frames, tiling_config = pipe(
                    prompt=prompt,
                    seed=seed,
                    height=height,
                    width=width,
                    frame_rate=frame_rate,
                    images=images,
                    num_frames=frames_arg,
                    temporal_upscalings=rounds,
                    spatial_upscalings=int(job.get("spatial_upscalings") or 1),
                    tiling_config=AUTO_TILING,
                    enhance_prompt=enhance_prompt,
                    enhance_static_cache=enhance_static_cache,
                    stage_1_sigmas=_sigmas(job, "stage_1_sigmas") if job.get("stage_1_sigmas") else DISTILLED_SIGMAS,
                    stage_2_sigmas=_sigmas(job, "stage_2_sigmas") if job.get("stage_2_sigmas") else STAGE_2_DISTILLED_SIGMAS,
                )
                # Each round doubles the PLAYBACK rate, not the duration: the
                # pipeline returns that many more frames for the same seconds,
                # and encoding them at the request's fps would play the clip in
                # slow motion.
                frame_rate = frame_rate * 2 ** rounds
            else:  # MODE_KEYFRAME
                params = _params_for(_phase2_paths()["dev_transformer"])
                steps = _steps_for(job, params.num_inference_steps)
                video, audio, out_frames, tiling_config = pipe(
                    prompt=prompt,
                    negative_prompt=negative_prompt,
                    seed=seed,
                    height=height,
                    width=width,
                    # int only: this pipeline has no duration head path and
                    # takes a plain frame count, so auto_duration does not apply.
                    num_frames=num_frames,
                    frame_rate=frame_rate,
                    num_inference_steps=steps,
                    video_guider_params=_guider_params(job, "video_guider", params.video_guider_params),
                    audio_guider_params=_guider_params(job, "audio_guider", params.audio_guider_params),
                    images=images,
                    tiling_config=AUTO_TILING,
                    vae_dtype=vae_dtype,
                    enhance_prompt=enhance_prompt,
                    enhance_static_cache=enhance_static_cache,
                    max_batch_size=max_batch_size,
                    stage_1_sigmas=_sigmas(job, "stage_1_sigmas"),
                    stage_2_sigmas=_sigmas(job, "stage_2_sigmas") or STAGE_2_DISTILLED_SIGMAS,
                    color_space=color_space,
                )

        # out_frames/tiling_config are the pipeline's RESOLVED values, not the
        # requested ones -- encoding with anything else risks a chunk count
        # that disagrees with the frames actually decoded.
        # OUTPUT quality, which is a different crf from the conditioning image's.
        # ImageConditioningInput.crf (resolved to 18 from the 2.5 checkpoint) is
        # how the SOURCE PHOTO is re-compressed before it conditions frame 0;
        # encode_video's crf is the h264 quantizer of the mp4 the customer
        # downloads, and upstream defaults it to 19. Nothing here ever set it,
        # so every clip this endpoint has produced is a crf-19 encode. Exposed
        # rather than pinned: it is a real quality/size lever on a product video
        # and changing it should not need a rebuild.
        encode_video(
            video=video,
            fps=frame_rate,
            audio=audio if include_audio else None,
            output_path=local_path,
            video_chunks_number=get_video_chunks_number(out_frames, tiling_config),
            crf=int(job.get("output_crf") or DEFAULT_OUTPUT_CRF),
            preset=job.get("output_preset") or DEFAULT_OUTPUT_PRESET,
            # HDR writes an EXR sequence plus an HLG master instead of H.264, so
            # it has to reach the encoder too -- passing it only to the pipeline
            # would decode HDR and then encode it as SDR.
            color_space=color_space,
        )

        # Pod-testing convenience: with LTX_SKIP_UPLOAD=1 the generated file is
        # kept locally and returned as a path -- no GCS credentials needed on
        # the test pod, and a several-minute generation isn't thrown away by an
        # upload-side failure. Never set on the real endpoint.
        if os.environ.get("LTX_SKIP_UPLOAD") == "1":
            kept_path = f"/tmp/ltx-out-{uuid.uuid4().hex}.mp4"
            os.replace(local_path, kept_path)
            return {"job_id": job_id, "ok": True, "mode": mode, "output_path": kept_path}

        output_url = _upload_video(local_path)
        return {"job_id": job_id, "ok": True, "mode": mode, "output_url": output_url}
    except Exception as e:
        # Free whatever the failed run left cached so a later job in the same
        # batch/worker starts from a clean allocator instead of inheriting a
        # nearly-full one (the exact state the 2026-07-31 OOM died in).
        if DEVICE == "cuda":
            torch.cuda.empty_cache()
        return {"job_id": job_id, "ok": False, "mode": mode, "error": f"{type(e).__name__}: {e}"}
    finally:
        # In finally, not after pipe(): the video iterator's VAE decode only
        # runs inside encode_video(), so printing earlier would understate the
        # real peak -- and printing here still reports it when a job fails
        # (the OOM case is exactly the one where this number matters most).
        if DEVICE == "cuda":
            print(f"peak CUDA memory: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GiB", flush=True)
        for path in [local_path, *scratch]:
            if os.path.exists(path):
                os.remove(path)



def handler(event):
    """
    Input (event["input"]): {"jobs": [{...}]}, one dict per render.

      every mode: job_id, prompt, seed, width, height, num_frames, fps,
                  include_audio, enhance_prompt, and `mode` (default "i2v")
      i2v / premium / keyframe: image_url, OR images=[{url, frame_idx, strength}]
                  for a keyframed (e.g. 360-degree turntable) render
      premium / keyframe: negative_prompt, num_inference_steps (both optional;
                  defaults come from upstream's own params for that pipeline)
      retake:     video_url, start_time, end_time

    Output: {"results": [{"job_id":..., "ok": true, "mode": ..., "output_url": "..."}]}

    The i2v shape is byte-for-byte what the 2.3 handler took, so every existing
    shoupapp caller keeps working without sending a `mode` at all.
    """
    job_input = event["input"]
    jobs = job_input.get("jobs") or []

    # No _get_pipe() here any more: the pipeline depends on each job's `mode`,
    # and _run_one() resolves (and caches) its own. A batch of same-mode jobs
    # still builds exactly one.
    results = [_run_one(job) for job in jobs]

    if DEVICE == "cuda":
        torch.cuda.empty_cache()
    return {"results": results}


runpod.serverless.start({"handler": handler})