Automated PR - 2026-06-17

Public sync - 2026-05-28

.claude/skills/train-model/SKILL.md

@@ -0,0 +1,277 @@
+---
+name: train-model
+description: End-to-end agent for training LTX-2 models. Probes filesystem and GPU, picks the right conditioning mode from the user's intent, prepares the dataset (scenes, captions, references), preprocesses, autotunes, launches, and monitors training. Use when the user wants to train, fine-tune, LoRA, or otherwise produce a custom LTX-2 model.
+argument-hint: [optional source path or run name]
+user-invocable: true
+allowed-tools: Bash, Read, Grep, Glob, Edit, Write, Agent, AskUserQuestion, TodoWrite
+---
+
+# Train Model — Orchestrator
+
+Take the user from "I want to train something" to a running, monitored training job — automating the mechanical glue (dataset layout, captioning, preprocessing, config patching, launch, monitoring) without making silent decisions on their behalf.
+
+> **Source of truth for the trainer:** [`packages/ltx-trainer/docs/`](../../../packages/ltx-trainer/docs/). This skill orchestrates those scripts; it does not duplicate their reference content.
+
+## Hard Invariants
+
+These are non-negotiable. Re-read them before every action that touches the filesystem or starts a process.
+
+1. **No file mutation outside the run workspace without explicit user approval.** The workspace is `./projects/<run-name>/`. Never overwrite, move, delete, or modify any user file or directory outside it without surfacing an explicit ask. Hours of dataset work must never be silently destroyed.
+2. **No heavy work before plan approval.** Captioning, preprocessing, training, and autotune do **not** start until the user approves `plan.md` (Phase 4). Probing the filesystem and running `nvidia-smi` is fine; encoding videos or downloading models is not.
+3. **No silent assumptions.** Every non-trivial default appears under "Assumptions" in `plan.md`.
+4. **No code changes to the trainer package without explicit consent.** If the user's intent doesn't map to a supported configuration, follow the **Escape Hatch** section below — do not unilaterally edit `packages/ltx-trainer/`.
+5. **No fabricated claims about training outcomes or data sufficiency.** Do not assert how well something *will* train, whether a dataset is "too small," how many samples/seconds of audio are "enough," which modality will "learn better," expected quality, or any similar prediction — you have no grounded basis for these, they're frequently wrong, and they mislead users. Stick to facts you can substantiate: what the trainer/docs actually say, observed numbers (loss, step time, VRAM), counts, and the user's own stated goals. If the user asks for a recommendation that depends on such judgment, you may share it **only** as an explicitly-flagged uncertainty ("I'm not sure — you'd have to try it"), never as authoritative fact. When in doubt, say less.
+
+## Keep the User Informed
+
+Most users don't know how this skill works under the hood — they don't know what "preprocessing," "a one-sample sanity check," or "autotune" mean or why they're happening. Narrate the run in plain language so it never feels like a black box:
+
+- **Entering a phase:** one or two sentences on *what you're about to do and why* — in user terms, not jargon.
+- **Leaving a phase:** one line on *what came out of it* (e.g. "captioned 9 clips," "found the fastest stable config: batch 1, ~3s/step").
+- **Explain the non-obvious phases explicitly** — these are the ones that confuse people:
+  - *Sanity check (Phase 6):* "Before the full run, I do a quick dry run on a single clip at your target resolution. It catches out-of-memory or config problems in a couple of minutes instead of failing hours into training."
+  - *Autotune (Phase 6):* "Then I try a few configuration variants on that one clip to pick the fastest one that still fits your GPU — so the full run is as fast as it can be."
+  - *Preprocess (Phase 7):* "I'm encoding your videos into the compressed latents the trainer reads. One-time step; the trained model never sees the raw videos directly."
+- Keep it concise — a sentence or two per transition, not walls of text or raw logs. This is running commentary, not a replacement for the upfront plan (Phase 4) or the status reports (Phase 8).
+- Long-running steps (preprocess, training): say roughly how long it'll take and that you'll report back, so silence doesn't read as "stuck."
+- **Describe what you're doing — don't editorialize about how it'll turn out.** Narration covers *what's happening*; it must not drift into unfounded predictions about training quality or data sufficiency (e.g. "26s of audio is too little," "voice won't learn well"). Those are fabricated claims — see Hard Invariant #5. State facts and the user's choices; leave the "will it be good?" judgment to the user watching the results.
+
+## Phase 0 — Set Up
+
+Create the workspace and todos.
+
+1. Pick a workspace root in this order (use first writable):
+   - `$LTX_TRAININGS_DIR`
+   - `/data/ltx-trainings/`
+   - `/workspace/ltx-trainings/`
+   - `./projects/` (repo-relative — preferred default in this repo)
+2. Derive a tentative `<run-name>` from the user's words; finalise after Phase 1 once the mode is known. Format: `<mode>-<dataset-name>-<YYYYMMDD-HHMM>`.
+3. Create `<workspace>/<run-name>/` and seed empty subdirs: `dataset/`, `outputs/`, `overfit/`.
+4. Create a todo list covering Phases 1–9 so the user can see progress.
+
+## Phase 1 — Intent
+
+Ask one question, framed in user terms (not jargon):
+
+> What do you want the model to learn? Examples: "generate videos from text," "make a LoRA of a specific style," "extend a video forward in time," "add sound effects to a silent video," "fill in masked regions of a video."
+
+Map the answer to one or more conditioning modes via `references/mode-selector.md`. If the requested capability has no mapping, go to **Escape Hatch**.
+
+### Plain concept/style LoRA → ask how it'll be used, default to I2V
+
+A "train a LoRA on X" request (a character/style/concept LoRA, no specific conditioning task) maps to either T2V or I2V. These aren't locked to inference: LoRA weights are pipeline-agnostic (the same checkpoint loads in both T2V and I2V inference), and the `i2v_lora` config applies first-frame conditioning with **`probability: 0.5`** — so it learns **both** conditioned (I2V) and unconditioned (T2V) generation in one run, and the first frame is taken automatically from each training clip (no extra data prep). I2V is therefore a versatile superset.
+
+Ask how they intend to use the result:
+
+> Will you generate videos from **text alone** (T2V), from a **starting image** (I2V), or **both / not sure**?
+
+- **Both / not sure (default):** use `i2v_lora` (probabilistic first-frame) — works for both at inference.
+- **I2V:** `i2v_lora`.
+- **Text only:** `t2v_lora`.
+
+(This only applies to plain concept/style LoRAs. A specific task — extension, inpainting, foley, IC-LoRA, etc. — maps directly to its mode via `mode-selector.md`; no usage question needed.)
+
+### Confirm the mode before proceeding
+
+Once the mode is determined, **state it plainly and confirm it** before doing any probing or work — a wrong inference is cheap to fix here and expensive later:
+
+> "Got it — I'll train an **I2V LoRA** (usable for both image-to-video and text-to-video at inference). Sound right?"
+
+The mode also appears in the plan (Phase 4), but confirm it here so the rest of the flow isn't built on a wrong guess.
+
+## Phase 2 — Probe
+
+No questions in this phase. Inspect what's already there. Use `references/onboarding.md` as the source of truth for the prerequisite checklist and what to do when something is missing.
+
+### Filesystem probe
+- If the user pointed at a path, classify: directory of raw videos, single long video, directory with a metadata file (CSV/JSON/JSONL), existing `.precomputed/`, partial outputs from a prior run.
+- For metadata files, identify columns: `video`/`media_path`, `caption`, `audio`, `reference_video`, `video_mask`, etc. (see `packages/ltx-trainer/docs/dataset-preparation.md`).
+- **Clip lengths (small datasets only):** for datasets up to a few hundred clips, `ffprobe` each clip's frame count and note the **minimum**. Clips shorter than the target frame bucket are silently skipped by `process_dataset.py`, so the shortest clip caps the achievable frame count — feed this into the Phase 3 resolution/frame-count choice (pick a bucket the clips support, or plan multi-bucket). Skip this per-clip probe for large datasets (too slow); rely instead on the post-preprocess reconciliation in Phase 7, which flags any dropped clips regardless of dataset size.
+- Check for an existing `<workspace>/<run-name>/` and whether `outputs/checkpoints/` contains a prior checkpoint (`lora_weights_step_*.safetensors` or `model_weights_step_*.safetensors`, plus a matching `training_state_step_*.pt` when resume state is enabled). This is a **resume candidate** — but note the trainer does *not* auto-resume from the output dir; resuming requires explicitly setting `model.load_checkpoint` in `config.yaml` to that checkpoint path. See `phases/launch-and-monitor.md` for the resume flow.
+- Check disk space at the workspace root. Preprocessed latents, checkpoints, and validation samples add up across a run; surface the available space alongside a rough sense of what one run consumes (one preprocessed bucket scales with sample count and resolution; each checkpoint is several GB), and warn the user if free space looks tight given their dataset size.
+
+### Hardware probe
+- `nvidia-smi --query-gpu=name,memory.total,driver_version --format=csv,noheader` → GPU model, count, VRAM. Stop the run if no CUDA GPU.
+- W&B login state — use wandb's **own** credential resolution (source-agnostic: covers env var, netrc, and the wandb settings file), not a hand-rolled netrc grep and **not** `wandb status` (which misleadingly reports `api_key: null` even when logged in):
+
+  ```bash
+  uv run python -c "import wandb; print(bool(wandb.Api().api_key))"   # True => logged in
+  ```
+
+  `True` → W&B is available, enable it. `False` → genuinely not logged in. If the check errors or is ambiguous, **ask the user** rather than silently disabling — a wrong "W&B off" assumption can incorrectly disable expected tracking.
+- Apply `references/hardware-profiles.md` to derive defaults (32GB / 40–60GB / 80GB+ VRAM tier).
+
+### Prerequisite probe (first-run sanity)
+- `uv` installed (`command -v uv`).
+- Workspace synced (lockfile present + `ltx-trainer` import works).
+- LTX-2 `.safetensors` and Gemma text encoder dir present in `/models/`, `~/models/`, or `$LTX_MODELS_DIR`.
+- Captioner availability: Gemini auth (`GEMINI_API_KEY`/`GOOGLE_API_KEY` or gcloud/Vertex) OR a ≥40 GiB GPU to host the Qwen3-Omni-30B vLLM server (FP8; bf16 needs ≥66 GiB). On typical consumer GPUs (24/32 GB), Gemini is effectively the only local-free option — see `references/onboarding.md`.
+
+For any missing prerequisite, **do not silently fail in a later phase**. Present the finding in chat with the specific next step from `references/onboarding.md`. The skill may offer to auto-install / auto-download missing pieces — but only ever after asking the user explicitly, one item at a time (model downloads are tens of GB each). Never auto-modify shell rc files or system config without consent.
+
+When everything (or what the user agreed to set up) is in place, fold the resolved paths into the plan's Assumptions section.
+
+### Pre-existing artifacts in the run dir
+
+If the run dir (or a user-supplied path) already contains artifacts from a prior session, classify each into one of three buckets — only the third prompts the user:
+
+1. **Deterministically verifiable → verify, then reuse silently (or stop).** `.precomputed/` latents: load a sample, check tensor shapes + modality coverage against the target (Phase 7). Match → reuse, no question. Mismatch/incomplete → stop and ask (reuse-at-old-spec / re-preprocess to a new dir / abort).
+2. **Cheap, fully-derived intermediates → regenerate silently.** `overfit/`, eval renders, sanity/temp configs, one-sample metadata. Delete and redo; don't ask.
+3. **Expensive AND not deterministically verifiable → ask.** Captions (`dataset.json`) and trained checkpoints/outputs. We can't programmatically decide whether existing captions or a half-finished run are what the user wants now, so surface what was found (counts, and how/when produced if knowable) and ask: reuse vs regenerate (for checkpoints: resume vs fresh).
+
+Principle: only ask when reuse-vs-regenerate is a genuine judgment call with cost either way. Never silently delete user-supplied data (hard invariant #1).
+
+## Phase 3 — Ask (minimum viable set)
+
+Ask only what cannot be inferred. Use `AskUserQuestion` with multiple-choice when possible. Typical questions:
+
+- **Target resolution / frame count** — propose a default per mode (e.g., `768x512x49` for T2V LoRA on consumer GPUs); offer overrides.
+- **Training steps** — if dataset size doesn't pin it, propose a default (e.g., 2000 for small LoRA datasets).
+- **LoRA trigger word / concept name** — only for style/concept LoRAs. Ask **only** for the word itself (or whether they want one). **Never** ask or mention *how* it's injected — it's always the `--lora-trigger` flag (passed to `process_dataset.py`, which forwards it to `process_captions.py` where the prepend happens); this is a fixed implementation detail. Presenting caption-injection as an option creates unnecessary confusion.
+- **Captioner backend** — only if more than one path is viable (e.g. a ≥40 GiB GPU can host the Qwen3-Omni-30B server *and* Gemini auth is available). On typical consumer GPUs, default to `gemini_flash` and surface that Gemini auth is required rather than asking.
+- **Model paths** — only if not found in the probe.
+
+Never ask anything answerable by `ls`, `nvidia-smi`, or the W&B credential check above.
+
+## Phase 4 — Plan
+
+Write the plan to `<workspace>/<run-name>/plan.md` using `references/plan-template.md`. Present it to the user in chat (don't just dump the file path). Wait for explicit approval before proceeding.
+
+If the user requests changes, edit the plan and re-present. Do not start Phase 5 until approval.
+
+## Phase 5 — Prepare Dataset
+
+If captioned metadata already exists with all required columns for the chosen mode, skip this phase. Otherwise follow `phases/prepare-dataset.md` — re-read it before acting.
+
+**Captioning gate:** caption a 3-sample spot-check first, show the captions in full, and **STOP for explicit user approval** before captioning the full set. The user must approve or give tuning instructions — never auto-proceed to the full pass. (Details in `phases/prepare-dataset.md`.)
+
+**Conditioning-inputs gate:** modes that need a reference (V2V/A2A/AV2AV IC-LoRA) or a mask (video/audio inpainting) require a per-sample input that encodes the user's specific idea. **Ask the user to provide it** — never invent the method (no defaulting to Canny/depth/generic masks). Only help generate it if the user explicitly asks, following *their* approach. Don't enter preprocessing for these modes without the input present. (Details in `phases/prepare-dataset.md` Step 4.)
+
+## Phase 6 — Sanity Check + Autotune (always run)
+
+**Tell the user what this phase is before starting it** — it's the most opaque to someone who doesn't know the design (see "Keep the User Informed"). In plain terms: a quick single-clip dry run at the target resolution to catch OOM/config errors cheaply, followed by trying a few config variants to pick the fastest stable one.
+
+Run **at the full target resolution** on **one sample** before the full preprocess. Purpose:
+1. Catch OOM / config errors before paying the full preprocessing cost.
+2. Empirically pick the fastest stable config via a small sweep.
+
+Steps:
+1. Pick one sample from the dataset metadata. Preprocess just that sample to `<workspace>/<run-name>/overfit/.precomputed/` (see `phases/preprocess-dataset.md` — use it in "one-sample" mode).
+2. Generate a temp config matching the planned full-run config but with `data.preprocessed_data_root: overfit/.precomputed`, `optimization.steps: 50`, `validation.interval: 50`, `checkpoints.interval: null`.
+3. Run the **baseline trial**: the conservative config from the matched VRAM tier (32GB tier = `t2v_lora_low_vram.yaml` defaults; 80GB+ tier = `t2v_lora.yaml` defaults — see `references/hardware-profiles.md`).
+4. **Success criteria** (all required):
+   - No OOM, no NaN loss, no crash.
+   - All 50 training steps complete.
+   - Validation sample at step 50 generates successfully (validation pass is a real OOM risk — do not skip).
+   - **For audio runs:** the one-sample `audio_latents/` is non-empty (the trainer log should report audio enabled). A joint/audio run that silently produced no audio latents is a failure even if steps complete — see the audio gate in `phases/preprocess-dataset.md`.
+   - **Loss is NOT a success criterion.** Loss can be non-monotonic even when training is healthy.
+5. **Autotune sweep** — incremental, capped at 5 trials total. Each trial = current best + one change. Stop on OOM (revert), no step-time improvement, or 5 trials:
+   - Trial 2: `quantization: null` (disable transformer quantization) if VRAM headroom allows.
+   - Trial 3: `optimizer_type: adamw` (disable 8-bit optimizer) if headroom allows.
+   - Trial 4: `batch_size` up (1 → 2 → 4). Adjust `gradient_accumulation_steps` proportionally to keep effective batch constant. **Note:** batch size can't be meaningfully tested on the one-sample set — defer it (test on the full set later, or just keep `batch_size: 1`, which is preferable for small concept-LoRA datasets anyway).
+   - **Do not sweep:** resolution (user decision), `acceleration.load_text_encoder_in_8bit` (one-time, no step-time impact), `enable_gradient_checkpointing` (the trainer's example configs ship with it on; on the 80GB+ tier you *may* try it off, but for the 22B model it usually OOMs even with tens of GB of apparent headroom — don't expect a win; never turn it off on the 32GB tier).
+6. Collect per trial: step time and peak VRAM. **Prefer the trainer's own end-of-run stats** (it prints total time / step time and peak GPU memory) — no external timing tool is needed (`/usr/bin/time` is often not installed). Append results to `<workspace>/<run-name>/autotune.log`.
+7. Winning trial's deltas are patched into the main `config.yaml`. Summarise the sweep to the user (one line per trial + winner).
+
+If the baseline trial fails, consult `references/troubleshooting.md`, propose a fix, re-run. Never push forward to the full preprocess after a failed sanity check.
+
+## Phase 7 — Full Preprocess
+
+Follow `phases/preprocess-dataset.md`. Re-read it before acting.
+
+**If `.precomputed/` already exists at the target path** (user-supplied or prior run), the phase verifies shapes and modality coverage before reuse. On mismatch it stops and asks — never silently overwrites.
+
+## Phase 8 — Launch & Monitor
+
+Follow `phases/launch-and-monitor.md`. Re-read it before acting. Surface the W&B URL (if enabled) and produce periodic status reports. At completion, the phase writes `<workspace>/<run-name>/outputs/run-summary.md`.
+
+## Phase 9 — Post-Train Validate
+
+After training finishes, follow `phases/post-train-validate.md`. Re-read it before acting. The phase renders the final LoRA against in-distribution, out-of-distribution, and held-out prompts; surfaces the MP4 paths to the user and exits.
+
+**Important constraint:** the post-train validate phase does not solicit a verdict and does not coach on causes for "soft" failures. Soft training quality has no reliable if/then rule book — the user inspects the renders and decides for themselves whether to ship, iterate, or change course. The orchestrator's job ends after Phase 9; iteration is a new invocation with a new run-name.
+
+## Monitor-Only Entry
+
+If invoked against an existing `<workspace>/<run-name>/` that already has a training process running or completed, **skip to Phase 8 in monitor-only mode** instead of restarting anything: report current step, recent loss, ETA, checkpoint list, W&B URL, and the resume command. Distinguish "live process" vs "stopped run with checkpoints" and offer the appropriate next action.
+
+## Escape Hatch: Unsupported Modes
+
+If Phase 1 intent doesn't map to any combination of supported `flexible`-strategy conditions:
+
+1. Stop. Do not edit `packages/ltx-trainer/` on your own.
+2. Explain what's missing in concrete terms: "you want X. The trainer supports A, B, C via conditions D, E, F. X requires a new condition / strategy."
+3. Identify the code change needed (typically a new `Condition` subclass in `ltx_trainer/training_strategies/flexible.py` plus schema wiring in `config.py`).
+4. Ask the user explicitly: "Proceed with the code change, do it yourself, or abort?"
+5. Only on explicit consent, drop out of this skill's orchestrator and edit code as a normal agent task. After the change lands and is tested, return here.
+
+## Ask-vs-Assume Cheat Sheet
+
+| Decision | How |
+|----------|-----|
+| Precision, quantization, optimizer, grad checkpointing | Assume from matched VRAM tier. List in plan's "Assumptions". |
+| Checkpoint/validation interval, seed, W&B project name, output dir | Assume sensible defaults. List in "Assumptions". |
+| Target resolution / frame count | Ask if not given; propose mode-appropriate default. |
+| Step count | Ask if dataset size doesn't pin it. |
+| LoRA trigger word / concept name | Ask for the word only (style/concept LoRAs). Never ask *how* it's injected — always `--lora-trigger`. |
+| Captioner backend | Ask only if multiple backends are viable. |
+| Anything answerable by `ls`, `nvidia-smi`, or the W&B credential check | Never ask. Probe. |
+
+## The Two `load_text_encoder_in_8bit` Flags
+
+Same name, different layers — do not conflate:
+
+| Flag | Layer | Effect |
+|------|-------|--------|
+| `process_dataset.py --load-text-encoder-in-8bit` | Preprocessing CLI | Memory during caption-embedding precompute (Phase 7). One-time per dataset. |
+| `acceleration.load_text_encoder_in_8bit` (trainer YAML) | Trainer config | Memory during validation-sample prompt-embedding caching at training start (Phase 8). One-time per run. |
+
+Both are one-time costs and neither affects per-step training speed. Default per the trainer's shipped configs: **ON** on the 32GB tier (matches `t2v_lora_low_vram.yaml`), **OFF** on the 80GB+ tier (matches `t2v_lora.yaml`). No measured guidance for the 40–60GB tier beyond starting from the 32GB tier and autotuning.
+
+## Workspace Layout
+
+```
+<workspace>/<run-name>/
+  plan.md                  # the approved plan
+  config.yaml              # generated training config (NOT in packages/ltx-trainer/configs/)
+  autotune.log             # per-trial sweep results
+  dataset/
+    dataset.json           # captions + media paths (training split)
+    holdout.jsonl          # held-out split (if reserved)
+    videos/                # source media copies (NO derived files here)
+    .precomputed/          # latents/ audio_latents/ conditions/ (+ references/masks per mode)
+  outputs/
+    checkpoints/           # training checkpoints + states
+    samples/               # in-training validation samples (step_*)
+    eval/                  # Phase 9: in-distribution/ out-of-distribution/ held-out/ + prompts.json
+    run-summary.md         # written at completion
+  logs/                    # all run logs
+  overfit/                 # Phase 6 scratch (one-sample preprocess + sanity/autotune runs)
+```
+
+`<run-name>` default: `<mode>-<dataset-name>-<YYYYMMDD-HHMM>`. Surface in the plan; user may rename.
+
+### Workspace hygiene (keep it clean)
+
+- **Don't create undocumented directories** (e.g. an ad-hoc `scratch/`). Intermediates belong under `overfit/` (sanity/autotune scratch) or a `/tmp` tempdir — not loose in the run dir or the repo root.
+- **Never write derived files into `dataset/videos/`** (the source media dir). Latents go under `.precomputed/`; one-sample/eval metadata goes under `overfit/`, not `dataset/`.
+- **One canonical manifest per artifact** — don't leave duplicate `*_prompts.json` / metadata copies.
+- **Clean up phase byproducts:** the Phase 9 eval must delete its validate-only trainer cruft (see `phases/post-train-validate.md`); `overfit/` is scratch and may be removed after a successful run. The final tree should look like the layout above — no stray `.pt`/`.wav` files, no `eval/checkpoints/`, no duplicate manifests.
+
+## References
+
+- `references/mode-selector.md` — user intent → conditioning mode mapping, LoRA rank guidance (read in Phase 1).
+- `references/onboarding.md` — first-run prerequisite checklist, model download paths, captioner graceful degradation (read in Phase 2).
+- `references/hardware-profiles.md` — GPU VRAM → tier + config defaults (read in Phase 2).
+- `references/config-patching.md` — safe YAML edits + schema constraints (read whenever editing `config.yaml`).
+- `references/troubleshooting.md` — OOM, NaN, validation failures, resume (read on any failure).
+- `references/plan-template.md` — exact plan.md format (read in Phase 4).
+
+## Phase Procedures
+
+These are procedure documents the orchestrator reads when entering each phase. They are not standalone skills — they're never invoked by Claude's skill-discovery system. The orchestrator opens them via the `Read` tool and follows the instructions inline.
+
+- `phases/prepare-dataset.md` — Phase 5: scenes, captioner iteration, IC-LoRA references, metadata, holdout split.
+- `phases/preprocess-dataset.md` — Phases 6 (one-sample) & 7 (full): `process_dataset.py` orchestration, existing-data verification.
+- `phases/launch-and-monitor.md` — Phase 8: launch command, accelerate, W&B, status reports, run-summary writing.
+- `phases/post-train-validate.md` — Phase 9: render the final LoRA against three prompt categories; surface paths only.

.claude/skills/train-model/phases/launch-and-monitor.md

@@ -0,0 +1,235 @@
+# Phase 8 — Launch & Monitor
+
+Procedure document for the `train-model` orchestrator (Phase 8 + monitor-only re-entry). Read this file in full before acting on launch/monitor.
+
+Goal: start the training job, surface the W&B URL, produce periodic status reports, write `run-summary.md` at completion.
+
+The orchestrator's hard invariants apply (see `../SKILL.md`).
+
+## Launch — Single GPU
+
+```bash
+cd packages/ltx-trainer
+uv run python scripts/train.py "<workspace>/<run-name>/config.yaml"
+```
+
+## Launch — Multi-GPU
+
+Use Accelerate. For LoRA, DDP (default) is fine. For full FT, use FSDP.
+
+```bash
+cd packages/ltx-trainer
+
+# DDP (LoRA, multi-GPU)
+uv run accelerate launch scripts/train.py "<workspace>/<run-name>/config.yaml"
+
+# FSDP (full FT, multi-GPU)
+uv run accelerate launch \
+  --config_file configs/accelerate/fsdp.yaml \
+  scripts/train.py "<workspace>/<run-name>/config.yaml"
+```
+
+**Pass `--disable-progress-bars` to `train.py` whenever stdout is redirected to a log file** (every background run) or running multi-GPU. The Rich progress bar rewrites a single line with carriage returns and does **not** flush parseable newlines to a redirected log, so without this flag the log shows no step/loss lines and you're forced to poll `nvidia-smi`. With it, step/loss lines are written normally and the log is greppable.
+
+## Pre-Launch Checks (every launch)
+
+1. Run the self-check from `references/config-patching.md` (paths exist, frame/resolution constraints, generated modalities have matching latents dirs, references/masks dirs present for conditional modes).
+2. Confirm `nvidia-smi` shows expected GPUs available (not occupied by another process).
+3. If `wandb.enabled: true`, confirm credentials still resolve: `uv run python -c "import wandb; print(bool(wandb.Api().api_key))"` → `True`. (Don't use `wandb status`.) If `False`, surface to the user before launching — they may want to `wandb login` or run without tracking.
+
+## Run In Background, Monitor Foreground
+
+Long training runs should not block the agent's response loop, and they must survive past the launching turn.
+
+**Prefer the agent's managed/native background-shell mechanism** (the harness facility for long-running background commands — output streaming + PID/exit tracking, survives across turns). It's the reliable way to launch training: it stays alive, streams to a log the agent can poll, and reports completion. Launch the training command through that mechanism, writing to `<workspace>/<run-name>/logs/train.log` with `--disable-progress-bars`, using **absolute paths** for the config and log (`uv run --directory packages/ltx-trainer` changes the cwd, so a relative config path won't resolve).
+
+**`nohup ... &` is a last-resort fallback only.** Detached jobs can be harder to track and may not survive environment/session cleanup, so only use it if no managed background mechanism is available, and verify the PID is still alive afterward:
+
+```bash
+# Fallback ONLY — prefer the managed background shell above.
+mkdir -p "<workspace>/<run-name>/logs"
+nohup uv run --directory packages/ltx-trainer python scripts/train.py \
+  "<ABSOLUTE path>/<run-name>/config.yaml" --disable-progress-bars \
+  > "<ABSOLUTE path>/<run-name>/logs/train.log" 2>&1 &
+echo $! > "<workspace>/<run-name>/logs/train.pid"
+```
+
+## Status Report
+
+Produce on user request (or every <interval> automatically). Pull from:
+
+- **W&B run URL:** First lines of `train.log` after init, or `wandb.run.url` from a `wandb` Python snippet. Surface as a clickable URL.
+- **Latest step:** `tail -n 200 "<workspace>/<run-name>/logs/train.log" | grep -oE "step [0-9]+" | tail -1`.
+- **Recent loss:** `tail -n 200 "<workspace>/<run-name>/logs/train.log" | grep -oE "loss[: ]+[0-9.]+" | tail -5`.
+- **Checkpoints saved:** `ls -1t "<workspace>/<run-name>/outputs/checkpoints/" 2>/dev/null`.
+- **Validation samples:** `ls -1t "<workspace>/<run-name>/outputs/samples/" 2>/dev/null`.
+- **GPU utilization snapshot:** `nvidia-smi --query-gpu=name,utilization.gpu,memory.used,memory.total --format=csv,noheader`.
+- **ETA:** only report one **grounded in real numbers the trainer has actually emitted** — do not invent or "educated-guess" an ETA before the trainer has produced per-step timings. Compute as `(total_steps - current_step) * recent_avg_step_time`, where `recent_avg_step_time` is measured from **steady-state training steps** (the trainer's reported per-step time or log timestamps), **excluding** one-time setup that doesn't repeat per step: model loading, the step-0 validation pass, and periodic validation passes. The trainer's own early ETA projection is skewed high by the slow step-0 validation and settles after a few steps — wait for it to settle rather than quoting the inflated early figure. Until real step timings exist, say "measuring step time…" rather than guessing a duration.
+
+Format the report tightly — one block, no fluff:
+
+```
+Step 1240 / 2000 (62%) — loss ~0.072 (last 5: 0.071, 0.073, 0.069, 0.075, 0.072)
+ETA: ~1h 24m  |  GPU: 91% util, 39.2 / 48.0 GB
+Latest checkpoint: lora_weights_step_01000.safetensors
+Latest validation: samples/step_01200_*.mp4
+W&B: https://wandb.ai/<entity>/<project>/runs/<id>
+```
+
+## Monitor-Only Mode
+
+Invoked when the orchestrator detects an existing `<workspace>/<run-name>/` with checkpoints or a running process.
+
+1. Check if the training process is live: `[ -f .../logs/train.pid ] && kill -0 $(cat .../logs/train.pid) 2>/dev/null && echo LIVE || echo STOPPED`.
+2. Produce the same status report as above.
+3. If STOPPED:
+   - Compute step from last checkpoint.
+   - Find the latest checkpoint pair (`lora_weights_step_*.safetensors` or `model_weights_step_*.safetensors`, plus a matching `training_state_step_*.pt` when resume state is enabled) under `<workspace>/<run-name>/outputs/checkpoints/`.
+   - Patch `model.load_checkpoint` in `config.yaml` to point at that checkpoint file (this is the only way the trainer knows to resume — there's no auto-detection from `output_dir`).
+   - Surface the resume command: `uv run python scripts/train.py "<workspace>/<run-name>/config.yaml"`.
+   - Ask the user to confirm both the patch and the launch before applying.
+
+## Resume
+
+The trainer **does not auto-resume from `output_dir`**. To resume an interrupted run:
+
+1. Set `model.load_checkpoint` in `config.yaml` to the latest checkpoint file (e.g. `<workspace>/<run-name>/outputs/checkpoints/lora_weights_step_02000.safetensors`).
+2. Launch normally. The trainer loads those weights, then looks for a matching `training_state_step_*.pt` **next to the loaded checkpoint** and restores optimizer/scheduler/step state from it. If the state file is missing, weights load but training starts from step 0.
+3. To load the weights but skip the state restore (e.g. for branching off into a new run from a known-good checkpoint), set `checkpoints.no_resume: true`.
+
+Always ask the user before patching `model.load_checkpoint` or setting `no_resume` — checkpoints are precious.
+
+## Failure During Training
+
+If the training process exits non-zero:
+
+1. Tail the log and identify the error type.
+2. Cross-reference `references/troubleshooting.md`.
+3. Propose a config fix (with the exact diff to `config.yaml`).
+4. Ask the user before applying. Then resume with the patched config.
+
+Never silently restart a failed training run without acknowledging the failure to the user.
+
+## After Training Completes
+
+1. Write a **run summary** to `<workspace>/<run-name>/outputs/run-summary.md` (see below) so the user can find their bearings months later without rereading the trainer docs.
+2. Show final checkpoint path and step count.
+3. Show W&B URL if enabled.
+4. Return control to the orchestrator (Phase 9 — post-train validate runs next).
+
+### Writing `run-summary.md`
+
+The summary is the **landing page** for this run. Anyone (including the user months from now) should be able to read it and understand: what was trained, on what data, with what config, where everything lives, how to use the result, and how to continue. The trainer doesn't produce this — the skill does.
+
+Template (fill from `plan.md`, `config.yaml`, `autotune.log`, dataset metadata, training log):
+
+```markdown
+# <run-name>
+
+**Trained:** <YYYY-MM-DD HH:MM> on <GPU(s)>
+**Final checkpoint:** `outputs/checkpoints/<filename>.safetensors` (step <N>)
+
+## What this LoRA does
+
+<One paragraph from the plan's Goal section — restating the user's intent.>
+
+## Trigger word
+
+`<trigger>` — include in prompts at inference time. (Omit this section if no trigger word.)
+
+## Mode
+
+<Mode name> (<lora|full>). Conditioning: <list of conditions, or "none">.
+
+## Dataset
+
+- Source: `<absolute path>`
+- Captioning: <`qwen_omni` (Qwen3-Omni-30B via vLLM) | `gemini_flash` | "user-supplied">
+  - Captioner instruction used: <verbatim string, or "default">
+- Samples: <N training> + <K held-out> at <W>x<H>x<F>
+- Preprocessed to: `dataset/.precomputed/`
+
+## Training config
+
+Final values after autotune (deltas from baseline noted in `autotune.log`):
+
+| Field | Value |
+|-------|-------|
+| Optimizer | <...> |
+| Mixed precision | <bf16/fp16> |
+| Quantization | <...> |
+| Gradient checkpointing | <on/off> |
+| Batch size × grad accum | <B> × <A> (effective <BxA>) |
+| LoRA rank / alpha | <R> / <A> (or "full FT") |
+| LoRA target modules | <list> (or "n/a") |
+| Steps | <N> |
+| Learning rate | <value> |
+| Step time (final) | ~<T>s |
+| Peak VRAM | ~<V> GB |
+
+Full config: `<workspace>/<run-name>/config.yaml`
+
+## Outputs
+
+- Checkpoints: `outputs/checkpoints/`
+- Validation samples (during training): `outputs/samples/`
+- Post-train eval renders (if Phase 9 ran): `outputs/eval/`
+- W&B run: <url, or "(W&B not enabled)">
+
+## How to use this checkpoint
+
+For inference, point `packages/ltx-pipelines/` at the final checkpoint. Example invocation:
+
+\`\`\`bash
+# (Minimal sketch — adapt to the pipeline you're using.)
+# load base LTX-2 model + apply this LoRA from outputs/checkpoints/<filename>.safetensors
+\`\`\`
+
+## How to continue training
+
+To resume from the final checkpoint (e.g. more steps, different LR), edit `config.yaml`:
+
+\`\`\`yaml
+model:
+  load_checkpoint: "<absolute path to outputs/checkpoints/<filename>.safetensors>"
+optimization:
+  steps: <new total>          # trainer resumes optimizer/scheduler/step from the training_state_step_*.pt sitting next to the checkpoint above
+\`\`\`
+
+Then re-launch with the same command in the "Launched with" section below.
+
+## How this was launched
+
+\`\`\`
+<exact command used, with the workspace's absolute config path>
+\`\`\`
+
+## Reproducibility
+
+- Seed: <value>
+- Workspace: `<absolute path>`
+- Repo commit at launch: `<git rev-parse HEAD output>`
+- LTX-2 model: `<model.model_path from config>`
+- Text encoder: `<model.text_encoder_path from config>`
+```
+
+Write the file using the `Write` tool. Don't embed it in a heredoc — the markdown nested in this skill is illustrative; fill the template with real values from the run's artifacts.
+
+### Next steps (surface to user)
+
+After writing the summary, point the user at:
+- The summary file path.
+- The final checkpoint path.
+- The W&B URL if enabled.
+- The upcoming Phase 9 (post-train validate) — the orchestrator handles the transition.
+
+Suggest, but don't run:
+- Test inference with `packages/ltx-pipelines/`.
+- Push to HF Hub via the trainer's `hub.push_to_hub` config (a separate, lightweight re-launch).
+- Continue training from the final checkpoint (see summary's "How to continue training" section).
+
+## Do Not
+
+- Do not modify `output_dir` contents after a run completes — checkpoints belong to the user now.
+- Do not start a second training run into the same `output_dir` without explicit user approval. Resume requires patching `model.load_checkpoint` (the trainer does not auto-detect prior checkpoints); a true fresh-start from the same dir additionally needs `checkpoints.no_resume: true`.
+- Do not auto-restart a failed run without diagnosis and user approval.

.claude/skills/train-model/phases/post-train-validate.md

@@ -0,0 +1,171 @@
+# Phase 9 — Post-Train Validate
+
+Procedure document for the `train-model` orchestrator (Phase 9). Read this file in full before acting on post-train validation.
+
+Goal: render the final checkpoint against three prompt categories so the user can inspect the result and form their own judgement. Save outputs in an organized layout. **Do not** prompt for pass/fail verdicts, do not infer causes for failures, do not suggest fixes — soft training failures don't have a clean if/then rule book, and pretending otherwise wastes the user's time.
+
+The orchestrator's hard invariants apply (see `../SKILL.md`).
+
+## What this phase does
+
+1. Collect prompts for three categories.
+2. Render the final LoRA against all collected prompts.
+3. Save outputs under `<workspace>/<run-name>/outputs/eval/<category>/`.
+4. Print the paths and exit.
+
+## Categories
+
+### 1 — In-distribution
+
+A few captions from the training set itself. Tests whether the model learned what it was shown.
+
+- Default: 3 random captions from the dataset metadata (seed 42 for reproducibility).
+- Source: `<workspace>/<run-name>/dataset/dataset.json` (the captions used for training, after the held-out split).
+
+### 2 — Out-of-distribution
+
+Prompts the model has never seen, but in the same domain. Tests whether the model generalizes the concept beyond memorized phrasings.
+
+Ask the user once:
+
+> "For out-of-distribution validation, paste 2–3 prompts you'd realistically want to generate at inference time. (If you don't have any specific ones in mind, reply 'default' — I'll use a few generic prompts that include the trigger word.)"
+
+- On `default`: synthesize 3 short prompts using the LoRA's trigger word and a generic scene context (e.g., "<trigger> walking in a forest at dawn"). Note these are generic — they're better than nothing, but real user-style prompts make a stronger test.
+- Otherwise: use the user's prompts verbatim.
+- **For a run with a generated audio modality, the synthesized prompts must describe the audio** — matching how the training captions describe it (inspect a few from `dataset.json` first). A prompt with no audio direction leaves the audio branch unguided and the generated audio comes out poor. E.g. for the talking-head case include spoken-voice/room-tone direction; for music/ambience/foley describe the sound character. (Categories 1 and 3 reuse the real captions verbatim, so they already carry audio description — this only applies to the synthesized Category 2 prompts.) If the user pasted their own prompts and it's an audio run, and they omitted audio direction, note that the audio may be weak without it.
+
+### 3 — Held-out
+
+Captions from samples that were never seen during training. Tests true generalization, not memorization.
+
+- Source: `<workspace>/<run-name>/dataset/holdout.jsonl` (written by `prepare-dataset` Step 5).
+- Use all entries if there are ≤5; otherwise sample 5 with seed 42.
+- **If holdout doesn't exist** (dataset was too small, or user-skipped during prepare): print a clear note in the output summary — *"Held-out evaluation skipped: no holdout set was reserved for this run. The post-train eval only covers Categories 1 and 2."* Don't synthesize substitutes.
+
+## Rendering mechanism
+
+Use the trainer's existing validation infrastructure rather than wiring up `ltx-pipelines` from scratch. Create a temporary "validate-only" config and run the trainer with it.
+
+### Step 1 — Build eval config
+
+Copy `<workspace>/<run-name>/config.yaml` to `<workspace>/<run-name>/eval-config.yaml`. Patch:
+
+```yaml
+model:
+  load_checkpoint: "<absolute path to outputs/checkpoints/<final-lora>.safetensors>"
+
+optimization:
+  steps: 1                     # we don't want to train; we want validation to fire
+  # Keep batch_size/grad_accum at the run's autotuned values to match its VRAM footprint.
+
+validation:
+  skip_initial_validation: false
+  interval: 1                  # run validation at step 0 (and at the only training step)
+  samples:
+    # Inject all collected prompts here, tagged by category in the prompt itself
+    # so the output filenames make the category obvious.
+    - prompt: "[CAT1-IND] <caption from dataset.json>"
+    # ... repeat for each prompt in all three categories
+  # Keep video_dims, frame_rate, guidance/STG settings as the trained config.
+  # Keep generate_audio consistent with the trained modality config.
+
+checkpoints:
+  interval: null               # do not save more checkpoints
+  no_resume: true              # load the LoRA's weights but do not restore optimizer/scheduler/step state
+
+output_dir: "<workspace>/<run-name>/outputs/eval"
+```
+
+The `[CAT1-IND]`, `[CAT2-OOD]`, `[CAT3-HELDOUT]` tags in the prompt strings make the output MP4 filenames self-describing in the trainer's validation sample directory.
+
+**Attach the mode's conditions to each sample.** A bare `prompt` only validates a pure text-to-X mode (T2V, T2A). For any conditioned mode, the trained model expects the same conditioning at validation time — a prompt with no conditions tests a different task than what was trained, and conditioned modes may fail outright. Add the `conditions` list that matches the run's mode (the trained `config.yaml` `training_strategy` and the example config for the mode are the reference):
+
+| Mode | Add to each sample |
+|------|--------------------|
+| I2V | `conditions: [{type: first_frame, image_or_video: <frame/clip path>}]` |
+| Video extension / suffix | `conditions: [{type: prefix|suffix, ...}]` |
+| V2V / AV2AV IC-LoRA | `conditions: [{type: reference, ...}]` (point at a held-out reference) |
+| V2A (foley) | `conditions: [{type: video_to_audio, ...}]` |
+| A2V | `conditions: [{type: audio_to_video, ...}]` |
+| Inpainting (video/audio) | `conditions: [{type: mask, ...}]` |
+| Outpainting | `conditions: [{type: spatial_crop, ...}]` |
+| A2A IC-LoRA | `conditions: [{type: reference, ...}]` (held-out reference audio) |
+| T2V, T2A | none — a bare `prompt` is correct |
+
+Mirror the condition shapes used in the mode's example config under `packages/ltx-trainer/configs/`. For held-out (Category 3) and OOD (Category 2) samples on conditioned modes, draw the conditioning media from the held-out set so the eval stays out-of-distribution.
+
+### Step 2 — Run the trainer in validate-only mode
+
+```bash
+cd packages/ltx-trainer
+uv run python scripts/train.py "<workspace>/<run-name>/eval-config.yaml"
+```
+
+The trainer will load the LoRA, run initial validation against all the prompts, do one trivial training step (which we discard), and exit. Validation samples land in `<workspace>/<run-name>/outputs/eval/samples/`.
+
+### Step 3 — Organize outputs and clean up trainer cruft
+
+The validate-only run is a trainer run, so it inevitably writes throwaway artifacts: an indexed `samples/` dir, a forced final checkpoint (the trainer **always** saves one at the end, regardless of `checkpoints.interval`), and a `training_config.yaml`. Don't leave these around or duplicate the renders.
+
+1. **Move** (don't copy) each generated MP4 from the trainer's indexed `samples/` dir into the category layout, naming by category + a slug of the prompt. Use the index→category mapping you built when constructing `validation.samples`.
+2. **Delete the trainer cruft** from the eval dir once the renders are moved: the indexed `samples/` dir, the forced `checkpoints/` dir, and `training_config.yaml`. (These are byproducts of the validate-only hack — there's no config flag to suppress the final-checkpoint save, so clean it up here.)
+3. Write **one** manifest, `outputs/eval/prompts.json` (filename → full prompt + category). Don't leave a second copy elsewhere.
+4. Remove the temporary `eval-config.yaml` (or keep it under the run's scratch, not in `outputs/`).
+
+Final eval layout — exactly this, nothing else:
+
+```
+<workspace>/<run-name>/outputs/eval/
+  in-distribution/    <NN>_<prompt-slug>.mp4 ...
+  out-of-distribution/ <NN>_<prompt-slug>.mp4 ...
+  held-out/           <NN>_<prompt-slug>.mp4 ...   # only if a holdout set existed
+  prompts.json        # filename -> full prompt + category (single manifest)
+```
+
+No `eval/samples/`, no `eval/checkpoints/`, no `eval/training_config.yaml`, no duplicate manifest.
+
+### Step 4 — Surface paths
+
+Print a tight block, no judgement, no follow-up question:
+
+```
+Post-train evaluation complete.
+
+In-distribution renders (<K> samples):
+  <workspace>/<run-name>/outputs/eval/in-distribution/
+
+Out-of-distribution renders (<M> samples):
+  <workspace>/<run-name>/outputs/eval/out-of-distribution/
+
+Held-out renders (<N> samples):
+  <workspace>/<run-name>/outputs/eval/held-out/    # or: "(skipped — no holdout set)"
+
+Open the MP4s and decide for yourself whether the model is good. Soft
+training quality is judged by watching the videos, not by a checklist —
+there's no substitute for your own eyes here.
+```
+
+Return control to the orchestrator. The orchestrator's run is now complete.
+
+## What this phase does NOT do
+
+- Does not ask "is this good?" / "pass / partial / fail?".
+- Does not infer failure causes.
+- Does not suggest fixes, follow-up runs, hyperparameter changes, dataset changes.
+- Does not write any verdict to `run-summary.md` or elsewhere.
+- Does not delete or modify training checkpoints.
+- Does not push to any remote / cloud / registry.
+
+The user looks at the videos and makes their own call. If they want to iterate, they re-invoke the orchestrator with a new run-name.
+
+## Failure modes
+
+- **Final checkpoint missing.** Surface and stop. Don't render against an intermediate checkpoint without explicit user consent.
+- **`load_checkpoint` OOM at inference time.** Lower `validation.video_dims` in the eval config (smaller renders are still useful for a sanity look). Retry once. If still OOM, surface the failure and let the user run inference manually via `packages/ltx-pipelines/`.
+- **All renders look broken/black.** May be an inference-pipeline-side issue rather than a training failure. Mention in the output block: *"If renders look broken across all categories, try `packages/ltx-pipelines/` directly to rule out a pipeline issue."* Then exit. Do not investigate further.
+
+## Do not
+
+- Do not skip Category 1 or 2. They're cheap and informative.
+- Do not invent a held-out set if `holdout.jsonl` is missing — the prepare-dataset step decides that.
+- Do not coach the user on what "good" means for their use case.

.claude/skills/train-model/phases/prepare-dataset.md

@@ -0,0 +1,245 @@
+# Phase 5 — Prepare Dataset
+
+Procedure document for the `train-model` orchestrator (Phase 5). Read this file in full before acting on the prepare-dataset phase.
+
+Goal: produce a captioned, complete dataset metadata file at `<workspace>/<run-name>/dataset/dataset.json` consumable by `process_dataset.py`. Idempotent — re-runs skip work already done.
+
+The orchestrator's hard invariants apply (see `../SKILL.md`), especially: **no file mutation outside the workspace without explicit user approval.**
+
+## Inputs
+
+The orchestrator passes:
+- Source path (directory of videos / single video / pre-existing metadata file).
+- Target mode (T2V, I2V, V2V IC-LoRA, V2A, etc.) — determines which columns are required.
+- Captioner backend choice (Qwen3-Omni local vLLM server / Gemini Flash cloud / skip).
+- Workspace path `<workspace>/<run-name>/`.
+
+## Required Columns by Mode
+
+`process_dataset.py` detects columns by convention and resolves each to a role. The media column may be `video` **or** `audio`. When a dataset has a `video` column with an audio track and no separate `audio` column, audio is **auto-extracted** from the video (unless `--skip-audio`), so an explicit `audio` column is only needed when the audio lives in separate files.
+
+**Video-generating modes** (need a `video` column):
+
+| Mode | Required | Optional |
+|------|----------|----------|
+| T2V, I2V, video extension/suffix | `video`, `caption` | `audio` (else auto-extracted) |
+| Video outpainting | `video`, `caption` | |
+| Video inpainting | `video`, `caption`, `video_mask` | |
+| V2A (foley) | `video`, `caption` | `audio` (target; else auto-extracted) |
+| A2V | `video`, `caption` | `audio` (else auto-extracted from video) |
+| V2V IC-LoRA | `video`, `caption`, `reference_video` | |
+| AV2AV IC-LoRA | `video`, `caption`, `reference_video`, `reference_audio` | `audio` (else auto-extracted) |
+
+**Audio-only modes** (no `video` column — the media column is `audio`):
+
+| Mode | Required | Optional |
+|------|----------|----------|
+| T2A | `audio`, `caption` | |
+| Audio extension/suffix | `audio`, `caption` | |
+| Audio inpainting | `audio`, `caption`, `audio_mask` | |
+| A2A IC-LoRA | `audio`, `caption`, `reference_audio` | |
+
+Aliases: `media_path` for `video`, `ref_media_path` for `reference_video`.
+
+## Workflow
+
+### Step 1 — Classify source
+
+```bash
+# If source is a file, identify type:
+file "<source>"
+# If source is a directory, count media:
+find "<source>" -maxdepth 1 -type f \( -name "*.mp4" -o -name "*.mov" -o -name "*.webm" \) | wc -l
+```
+
+Cases:
+- **Pre-existing metadata file** (CSV/JSON/JSONL) → copy to `<workspace>/<run-name>/dataset/dataset.json`, audit columns. Skip to Step 4.
+- **Directory of short scenes** → skip Step 2, go to Step 3.
+- **Directory containing long videos** → run Step 2.
+- **Single long video** → run Step 2.
+
+**Stage the media under `dataset/` before captioning — don't discover this by failing.** Both `caption_videos.py` and `process_dataset.py` reference media by paths **relative to the metadata file's own directory**, so the media must live under `<workspace>/<run-name>/dataset/`. Stage it up front into `dataset/videos/` (and write metadata paths relative to `dataset/`, e.g. `videos/1.mp4`):
+- Prefer **symlinks** (instant, no disk cost): `ln -s <abs-source>/<clip> <workspace>/<run-name>/dataset/videos/<clip>`. Symlinks pointing at the original source location work correctly.
+- Use a **copy** instead if the workspace and source are on different filesystems or the source may move/change.
+- Never caption or preprocess directly against an out-of-tree source path (e.g. `/path/to/source-videos`) — it will fail the relative-path resolution. The original source is left untouched either way.
+(Scene-split output in Step 2 already lands under `dataset/scenes/`, which satisfies this.)
+
+### Step 2 — Scene splitting (only for long videos)
+
+`split_scenes.py` takes **one video file** at a time (`video_path` and `output_dir` are both positional arguments). When the source is a directory of long videos, iterate over each file. To drop scenes shorter than 2 seconds, use `--filter-shorter-than 2s` (the `--min-scene-length` option is an integer **frame** count, not seconds — don't pass a float).
+
+```bash
+cd packages/ltx-trainer
+# Single file:
+uv run python scripts/split_scenes.py \
+  "<video-file>" \
+  "<workspace>/<run-name>/dataset/scenes" \
+  --filter-shorter-than 2s
+
+# Directory of long videos — iterate:
+for f in "<source>"/*.mp4 "<source>"/*.mov "<source>"/*.webm; do
+  [ -e "$f" ] || continue
+  uv run python scripts/split_scenes.py "$f" \
+    "<workspace>/<run-name>/dataset/scenes" \
+    --filter-shorter-than 2s
+done
+```
+
+Result: scenes saved to `<workspace>/<run-name>/dataset/scenes/`. Pass that directory to Step 3.
+
+### Step 3 — Captioning
+
+Skip entirely if a metadata file with all required `caption` entries already exists.
+
+**Use the captioner's default instruction.** `caption_videos.py` ships a well-tuned default caption prompt — use it as-is (do **not** pass `--instruction`). Captioning runs in two phases: a small **spot-check pass** so the user can confirm the captions look sane, then a **full pass** on the rest.
+
+A custom `--instruction` is the exception, not the norm. Only use one when:
+- the **nature of the dataset genuinely demands it** (e.g. a narrow domain the default prompt won't describe well), or
+- the **user, after seeing the spot-check captions, explicitly asks** for a change (e.g. "too much background detail").
+
+Do not invent a custom instruction pre-emptively, and in particular **do not bake a subject name / trigger word into the captions via `--instruction`** — the trigger word is handled separately at preprocessing (see "Trigger word" below).
+
+#### Choosing a backend
+
+Two backends, with very different hardware needs:
+
+- **`qwen_omni` (local, default):** Qwen3-Omni-30B-A3B-Thinking served by a local vLLM HTTP server (`serve_captioner.py`). ~65 GiB model download. Default **FP8** quantization uses ~31 GiB of weights and **fits on a 40 GiB GPU** (plus KV cache); **bf16** uses ~60 GiB and needs **≥66 GiB free VRAM**.
+- **`gemini_flash` (cloud):** Google `gemini-3.5-flash`. No local model, runs anywhere, parallelisable with `--num-workers`.
+
+**Steer modest hardware to Gemini.** If the GPU is below ~40 GiB (i.e. typical consumer cards — 24 GB / 32 GB), it can't host even the FP8 server, so the local captioner isn't an option — recommend `gemini_flash` and tell the user they'll need Gemini auth: either a `GEMINI_API_KEY`/`GOOGLE_API_KEY` (get one at <https://aistudio.google.com/apikey>) or working gcloud/Vertex AI credentials. If they can't or won't set that up and the hardware can't run Qwen3, the only remaining path is bringing their own captions in the dataset metadata (skip captioning entirely). On a 40 GiB+ GPU the local server is viable (FP8); bf16 needs an 80GB-class card.
+
+#### Qwen server prerequisite (qwen_omni only)
+
+The local backend talks to a vLLM server that must already be running. Launch it once in a **separate terminal** (it stays loaded across captioning runs):
+
+```bash
+cd packages/ltx-trainer
+uv run python scripts/serve_captioner.py        # FP8 by default, serves on http://127.0.0.1:8001/v1
+# bf16 (needs >= 66 GiB free VRAM):  --quantization bf16
+# different port/interface:           --port 9000 --host 0.0.0.0
+```
+
+First launch downloads the model (~65 GiB). `caption_videos.py` reaches it via `--vllm-url` (default `http://127.0.0.1:8001/v1`). Skip this entirely when using `gemini_flash`.
+
+#### 3a — Spot-check pass (3 samples)
+
+Caption 3 samples with the **default prompt** (no `--instruction`) to confirm the captioner is producing sane output before committing to the whole set.
+
+```bash
+cd packages/ltx-trainer
+# qwen_omni (server from the previous step must be running):
+uv run python scripts/caption_videos.py \
+  "<workspace>/<run-name>/dataset/videos/<one-staged-clip>" \
+  --output "<workspace>/<run-name>/dataset/preview-captions.json" \
+  --captioner-type qwen_omni
+  # Point at 3 staged clips under dataset/videos/ (a small subdir or 3 explicit files) — not the out-of-tree source.
+  # Optional: --vllm-url http://127.0.0.1:9000/v1   (if the server uses a non-default port)
+```
+
+Print the 3 captions **in full** to the user, then **STOP and wait** for their explicit verdict:
+
+> "Here are sample captions from the default prompt. Please review them — reply 'good' to caption the rest, or tell me what to change."
+
+**This is a hard gate. Do NOT caption the full set until the user explicitly approves the samples.** Do not auto-proceed, do not assume "looks fine," do not batch this with other questions. The user must either approve or give tuning instructions first — the whole point of the spot-check is to let them judge caption quality and content before paying for the full pass.
+
+If the user requests changes, introduce a custom `--instruction` (or switch backend), re-run the spot-check on the same 3 samples, show the new captions, and **stop for approval again**. Loop until the user approves. If a custom instruction still isn't converging after a few rounds, switch captioner backend or have the user supply a few manual captions as examples — but still don't proceed to the full set without their OK.
+
+#### 3b — Full pass
+
+Run on the **staged media dir** (`dataset/videos/` from Step 1, or `dataset/scenes/` from Step 2) with the **default prompt** (or the same `--instruction` only if one was explicitly agreed in 3a). Caption the staged in-tree media — not the original out-of-tree source path.
+
+**Qwen3-Omni (local — server must be running):**
+
+```bash
+cd packages/ltx-trainer
+uv run python scripts/caption_videos.py \
+  "<workspace>/<run-name>/dataset/videos" \
+  --output "<workspace>/<run-name>/dataset/dataset.json" \
+  --captioner-type qwen_omni
+```
+
+**Gemini Flash (cloud — runs anywhere, parallelisable):**
+
+```bash
+# Auth: GEMINI_API_KEY / GOOGLE_API_KEY env var, or gcloud / Vertex AI credentials.
+cd packages/ltx-trainer
+uv run python scripts/caption_videos.py \
+  "<workspace>/<run-name>/dataset/videos" \
+  --output "<workspace>/<run-name>/dataset/dataset.json" \
+  --captioner-type gemini_flash \
+  --num-workers 5
+```
+
+Output: JSON list of `{caption, media_path}` with paths **relative to the output file location**. The 3 spot-check captions can be merged in to avoid re-captioning them.
+
+#### Trigger word (handled at preprocessing, not in captions)
+
+For style/concept LoRAs the trigger word is **not** written into the captions here. It is prepended to every caption at preprocessing: pass `--lora-trigger "<word>"` to `process_dataset.py` in Phase 7, which forwards it to the caption-processing step (`process_captions.py`, where the prepend actually happens). That is the canonical mechanism — keep the captions describing what's actually on screen (via the default prompt), and let the trigger flag bind the concept to the token. Record the chosen trigger word in the plan so Phase 7 passes it through. Do not also bake the word into captions (it would double up).
+
+**The injection mechanism is a fixed implementation detail — never make it a user-facing question.** The *only* trigger-word thing to ask the user is the **word itself** (or whether they want a trigger word at all). Do **not** ask, mention, or present as an option *how* it gets injected (e.g. "inject into the caption vs via `process_dataset`") — it is always `--lora-trigger`, full stop. Surfacing the method as a choice creates unnecessary confusion.
+
+### Step 4 — Conditioning inputs (modes that need references or masks)
+
+Some modes need a per-sample conditioning input beyond the video/audio and caption:
+
+| Mode | Required extra input | Column |
+|------|----------------------|--------|
+| V2V IC-LoRA | reference video | `reference_video` (alias `ref_media_path`) |
+| AV2AV IC-LoRA | reference video + reference audio | `reference_video`, `reference_audio` |
+| A2A IC-LoRA | reference audio | `reference_audio` |
+| Video inpainting | per-frame video mask | `video_mask` |
+| Audio inpainting | audio mask | `audio_mask` |
+
+These inputs encode **the user's specific idea** for the LoRA (what the reference represents, which regions the mask covers). There is no universal recipe, so **do not invent or default to a particular method** (e.g. don't assume Canny edges, depth, pose, or some generic box/border mask). The agent must not pick the conditioning semantics for the user.
+
+Workflow when the chosen mode needs one of these and the dataset doesn't already provide it:
+
+1. **Check first** — if the user already supplied the column (and the files exist), use it as-is and move on.
+2. **Otherwise, ask the user to provide it**, explaining concretely what's needed: the column name, that it's one file per sample aligned to each clip, and that the *content/semantics are their call* (what the reference should depict, what the mask should cover). Make clear this reflects their specific use-case — you won't guess it.
+3. **Help generate only if the user asks.** If they say "can you generate the references/masks by doing X" (X = their described method), then help: write or run a small script for *their* approach, or use a repo tool if it fits. One such tool exists — `scripts/compute_reference.py` generates **Canny edge** reference videos — but only mention/use it if the user specifically wants Canny; never offer it as the default.
+4. **Hard gate:** do not proceed to preprocessing for a conditioning mode until the required column is present with real files. Surface clearly if it's missing.
+
+**Column-naming note (if references are generated):** `compute_reference.py` writes a `reference_video` field, which
+`process_dataset.py` detects automatically. Legacy datasets using `ref_media_path` also work.
+
+### Step 5 — Holdout split
+
+Reserve a subset of samples as a **held-out set** never seen during training. This is what Phase 9 (post-train validate) renders against to test true generalization rather than memorization.
+
+Decision tree:
+
+1. **User already supplied a held-out set** (separate file or directory they explicitly nominated): do not split. Copy/reference their file to `<workspace>/<run-name>/dataset/holdout.jsonl` and leave `dataset.json` as-is.
+2. **Small dataset** (judge qualitatively; tens of samples or fewer): holding samples out meaningfully reduces training capacity. Ask the user:
+   > "Dataset has <N> samples. Reserving a holdout meaningfully reduces what's available for training. Options: (a) reserve 1–2 for holdout, (b) skip holdout — post-train eval will only test in-distribution. Your call."
+3. **Otherwise:** auto-split. Reserve a small fraction (this skill's default: roughly 10% of samples, bounded so the holdout doesn't grow huge — a handful of held-out samples is usually enough). Use seed 42 for the split so it's reproducible. Surface the count and the picked IDs in the plan.
+
+After splitting, write `<workspace>/<run-name>/dataset/holdout.jsonl` (one JSON object per line with the same columns as `dataset.json`). **Remove the held-out entries from `dataset.json`** so they don't enter preprocessing or training.
+
+Always print:
+> "Held out <K> of <N> samples for post-train evaluation. Held-out IDs: <list>."
+
+If holdout is skipped, surface in the plan: *"Skipping holdout — dataset is too small. Post-train eval will only render in-distribution prompts; true generalization isn't testable for this run."*
+
+### Step 6 — Audit
+
+Before returning to the orchestrator, verify the metadata file has all required columns for the chosen mode. Print a one-line summary:
+
+> "Prepared <N> training samples (+ <K> held out) for <mode>. Columns: <list>. Saved to `<workspace>/<run-name>/dataset/dataset.json` (+ `holdout.jsonl`)."
+
+## Idempotency
+
+- If `dataset.json` already exists and all required columns are present: skip captioning. Confirm reuse with the user only if the file was supplied by them outside the workspace (per the orchestrator's file-safety invariant).
+- If captioning was partial (some entries missing `caption`), re-run captioning only on the missing entries by filtering the metadata file before passing to `caption_videos.py`.
+
+## Failure Modes
+
+- Qwen server won't start / OOMs on launch → use the default `--quantization fp8` (not `bf16`), lower `--gpu-memory-utilization`, or reduce `--max-model-len` on `serve_captioner.py`. If the GPU simply can't host a 30B model, switch to `gemini_flash`.
+- `caption_videos.py` can't connect (qwen_omni) → the vLLM server isn't running or `--vllm-url` is wrong. Start `serve_captioner.py` first and confirm the URL/port match.
+- Gemini rate-limit → reduce `--num-workers`, retry.
+- Scene splitter produces 0 scenes → the detector found no cuts, or `--filter-shorter-than` removed everything. Lower/remove `--filter-shorter-than`, or adjust the detector threshold. (`--min-scene-length` is an integer minimum-frames-per-scene for the detector, not a short-scene filter.)
+- IC-LoRA reference compute fails on some frames → script logs the failures; report counts to the user and ask whether to proceed with the remaining samples or stop.
+
+## Do Not
+
+- Do not move or rename the user's source files. The skill reads them in place; the workspace contains only **derived** artifacts.
+- Do not delete `scenes/` or partial captioning outputs without approval — they may be expensive to regenerate.

.claude/skills/train-model/phases/preprocess-dataset.md

@@ -0,0 +1,130 @@
+# Phases 6 (one-sample) & 7 (full) — Preprocess Dataset
+
+Procedure document for the `train-model` orchestrator. Read this file in full before acting on the preprocess phase.
+
+Goal: run `process_dataset.py` to produce VAE latents, audio latents, and text embeddings. Two modes:
+1. **One-sample** (Phase 6 sanity check) — preprocess a single sample to `<workspace>/<run-name>/overfit/.precomputed/`.
+2. **Full** (Phase 7) — preprocess the whole dataset to `<workspace>/<run-name>/dataset/.precomputed/`.
+
+The orchestrator's hard invariants apply (see `../SKILL.md`), especially: **never silently overwrite existing user data.**
+
+## Required Subdirectories by Mode
+
+Under `.precomputed/`:
+
+| Subdir | Required for |
+|--------|--------------|
+| `latents/` | Video-bearing modes only (T2V, I2V, video extend/inpaint/outpaint, V2V/AV2AV IC-LoRA, A2V, V2A). **Not** produced for audio-only modes. |
+| `conditions/` | Always (text embeddings) |
+| `audio_latents/` | Any mode with audio: video modes carrying audio, plus all audio-only modes (T2A, audio extend/suffix/inpaint, A2A IC-LoRA) |
+| `reference_latents/` | V2V IC-LoRA, AV2AV IC-LoRA |
+| `reference_audio_latents/` | A2A IC-LoRA, AV2AV IC-LoRA |
+| `video_masks/` | Video inpainting |
+| `audio_masks/` | Audio inpainting |
+
+**Audio-only modes** (T2A, audio extend/suffix, audio inpainting, A2A IC-LoRA) produce `audio_latents/` + `conditions/` (plus `audio_masks/` or `reference_audio_latents/` as applicable) and **no `latents/`**. Do not flag a missing `latents/` as incomplete for these modes.
+
+## Workflow — Full Preprocess
+
+### Step 1 — Verify existing `.precomputed/` (if present)
+
+If `<workspace>/<run-name>/dataset/.precomputed/` already exists:
+
+1. List subdirectories present. Confirm all required for the chosen mode are present.
+2. Load one sample per modality and check tensor shapes:
+   ```bash
+   uv run python -c "import torch; t = torch.load('<path>'); print(t.shape if hasattr(t, 'shape') else {k: v.shape for k, v in t.items()})"
+   ```
+3. Compare shapes against the target resolution from the plan.
+
+**On any mismatch or missing subdirectory: STOP. Do not run `process_dataset.py`.** Ask the user via `AskUserQuestion`:
+- Reuse the existing data at its current resolution (update plan + config accordingly).
+- Re-preprocess to a new directory (`<workspace>/<run-name>/dataset/.precomputed-v2/` etc.) — preserves the existing data.
+- Abort.
+
+**Never pass `--overwrite` without explicit user approval** for this exact action.
+
+### Step 2 — Invoke `process_dataset.py`
+
+```bash
+cd packages/ltx-trainer
+uv run python scripts/process_dataset.py \
+  "<workspace>/<run-name>/dataset/dataset.json" \
+  --resolution-buckets "<W>x<H>x<F>" \
+  --model-path "<absolute-model-path>" \
+  --text-encoder-path "<absolute-gemma-path>" \
+  --output-dir "<workspace>/<run-name>/dataset/.precomputed" \
+  --load-text-encoder-in-8bit         # on 32GB tier (low-VRAM config), per t2v_lora_low_vram.yaml
+```
+
+Add as needed:
+- `--skip-audio` — if mode doesn't use audio (T2V video-only variants).
+- `--audio-durations "<list>"` — for T2A from a captions-only file.
+- `--lora-trigger "<trigger>"` — for style/concept LoRAs.
+- `--reference-downscale-factor <N>` — for IC-LoRA modes if downscaled references are desired.
+- `--video-column`, `--caption-column` — only if the metadata file uses non-standard column names.
+
+**Do not pass `--overwrite`** unless re-preprocessing was explicitly approved in Step 1.
+
+### Step 3 — Audit output
+
+After completion, verify:
+
+```bash
+ls "<workspace>/<run-name>/dataset/.precomputed/"
+# Expected subdirs per the mode table above.
+
+# Count files in each:
+for d in latents conditions audio_latents reference_latents video_masks audio_masks; do
+  if [ -d "<workspace>/<run-name>/dataset/.precomputed/$d" ]; then
+    echo "$d: $(ls "<workspace>/<run-name>/dataset/.precomputed/$d" | wc -l)"
+  fi
+done
+```
+
+Counts in each required subdir should equal the dataset sample count.
+
+**Reconcile counts — do this for every run, any dataset size.** Compare the `latents/` (and `audio_latents/`) count against the caption/sample count. If fewer latents were produced, `process_dataset.py` **silently skipped** clips — most commonly because they were **shorter than the target frame bucket** (it logs each skip). When counts don't match:
+
+1. Identify which clips were dropped (grep the preprocess log for skip/"fewer frames" lines, or diff the produced `.pt` stems against the metadata).
+2. **Surface it to the user** with the count and the specific clips — never silently proceed on a shrunk dataset.
+3. Offer options: re-preprocess at a **smaller frame bucket** the clips support, add a **second (shorter) bucket** to keep the short clips (multi-bucket requires `batch_size: 1`), or accept the loss. Let the user decide.
+
+**Audio gate (hard stop for audio runs).** For any run with an audio modality (joint audio+video, A2V, V2A, T2A, audio-only modes), verify `audio_latents/` is **present and non-empty** with one `.pt` per sample. `process_dataset.py` **swallows audio-decode errors and continues** — it logs "0 videos with audio" and produces empty `audio_latents/` rather than failing. If an audio run produced no audio latents, **stop** — do not proceed to training (it would silently train audio-free). The usual cause is a broken audio decode path (e.g. missing/incompatible `torchcodec`); confirm `uv run python -c "import torchaudio; torchaudio.load('<a clip>')"` works (see `references/troubleshooting.md`), fix it, then re-preprocess with `--overwrite`.
+
+## Workflow — One-Sample (Phase 6)
+
+Same as full preprocess, but operate on a single-sample metadata file:
+
+1. Pick the first sample from `<workspace>/<run-name>/dataset/dataset.json` and write a one-sample metadata file **inside the dataset dir** — e.g. `<workspace>/<run-name>/dataset/_one_sample.json` — copying the entry **verbatim, keeping its relative `media_path`**. `process_dataset.py` resolves media paths relative to the metadata file's own directory, so the one-sample file must sit beside the real media (i.e. in `dataset/`, the same dir as `dataset.json`). **Do not** place it in `overfit/` and **do not** rewrite the path to an absolute one — an absolute path produces mirrored nested output dirs (`.precomputed/latents/absolute/path/.../x.pt`) instead of a clean `latents/x.pt`.
+2. Run `process_dataset.py` on that file with `--output-dir "<workspace>/<run-name>/overfit/.precomputed"` (output still goes to `overfit/`, only the metadata lives in `dataset/`).
+3. Use the **same `--resolution-buckets`** as the planned full run. Critical: a small-shape sanity check is misleading because resolution is the dominant memory factor.
+4. Clean up the temporary `dataset/_one_sample.json` afterward (it's scratch; don't leave it in the dataset dir).
+
+## Decode-and-Verify (optional debug aid)
+
+If the user reports validation samples look wrong or training diverges, decode one preprocessed sample back to media:
+
+`decode_latents.py` takes the **latents directory** and an **output directory** as positional arguments (it decodes the whole directory, not a single `.pt` file). Add `--with-audio` and `--audio-latents-dir` if the dataset has audio.
+
+```bash
+cd packages/ltx-trainer
+uv run python scripts/decode_latents.py \
+  "<workspace>/<run-name>/dataset/.precomputed/latents" \
+  "<workspace>/<run-name>/dataset/.precomputed/decoded_check" \
+  --model-path "<absolute-model-path>"
+```
+
+If decoded output is garbled, preprocessing itself is suspect (wrong model, wrong VAE).
+
+## Failure Modes
+
+- **"shape mismatch" on resume:** Step 1's verification check. Ask user before any mutation.
+- **`frames % 8 != 1`** error from process_dataset.py: the requested frame count is invalid; correct in the plan and re-launch.
+- **VRAM OOM during preprocessing:** add `--load-text-encoder-in-8bit`. If still OOM, reduce `--batch-size`.
+- **Disk full:** preprocessed latents can be large (especially audio). Surface to user with a `du -sh` summary of `.precomputed/`.
+
+## Do Not
+
+- Do not delete or overwrite existing `.precomputed/` data without explicit user approval for that exact action.
+- Do not preprocess at a smaller resolution to "save time" — the sanity check exists specifically to validate the planned resolution.

.claude/skills/train-model/references/config-patching.md

@@ -0,0 +1,73 @@
+# Config Patching
+
+How to safely produce `<workspace>/<run-name>/config.yaml` from an example in `packages/ltx-trainer/configs/`. The trainer's config schema is Pydantic with `extra="forbid"` — unknown fields are rejected. Full field reference: [`packages/ltx-trainer/docs/configuration-reference.md`](../../../../packages/ltx-trainer/docs/configuration-reference.md).
+
+## Workflow
+
+1. Copy the example config matching the selected mode (see `mode-selector.md`) to `<workspace>/<run-name>/config.yaml`.
+2. Patch fields as described below. Preserve YAML comments where possible — they help the user audit the run later.
+3. **Never** edit the example config in `packages/ltx-trainer/configs/`. That's the user's reference library.
+
+## Required Patches (every run)
+
+| Field | Value |
+|-------|-------|
+| `model.model_path` | Absolute path to local `.safetensors` (from probe or user). |
+| `model.text_encoder_path` | Absolute path to local Gemma directory (from probe or user). |
+| `data.preprocessed_data_root` | `<workspace>/<run-name>/dataset/.precomputed` (absolute). |
+| `output_dir` | `<workspace>/<run-name>/outputs` (absolute). |
+
+## Hardware-Driven Patches
+
+Apply per the matched VRAM tier in `references/hardware-profiles.md`. After autotune (Phase 6), patch the winning trial's deltas in.
+
+## Schema Constraints (validate before launch)
+
+These will cause Pydantic errors or runtime failures; check before invoking the trainer.
+
+- **Frame count:** `validation.video_dims[2]` must satisfy `frames % 8 == 1` (1, 9, 17, 25, 33, 41, 49, 57, 65, 73, 81, 89, ...).
+- **Resolution:** `validation.video_dims[0]` and `[1]` must be divisible by 32.
+- **Multi-bucket training:** if dataset uses multiple resolution buckets, set `optimization.batch_size: 1`.
+- **At least one generated modality:** `training_strategy` must have at least one of `video.is_generated` or `audio.is_generated` set to `true`.
+- **Audio condition restrictions:** the audio modality cannot use `first_frame` or `spatial_crop` conditions.
+- **Strategy name:** prefer `training_strategy.name: "flexible"`. `text_to_video` and `video_to_video` still work but emit deprecation warnings.
+
+## LoRA Patches
+
+For style/concept LoRAs:
+
+- `lora.rank` and `lora.alpha`: set from the matched VRAM tier and use case. **32GB tier** pins rank 16 per `t2v_lora_low_vram.yaml`; **80GB+ tier** uses rank 32 per `t2v_lora.yaml`. Keep `alpha == rank`. See `mode-selector.md` for use-case-driven rank guidance.
+- `lora.target_modules`: short patterns like `"to_k"`, `"to_q"`, `"to_v"`, `"to_out.0"` match all attention modules (video + audio + cross-modal). Add `"ff.net.0.proj"`, `"ff.net.2"` only if user explicitly wants higher capacity.
+- **Audio-only LoRA targets** (T2A, audio inpainting): use `"audio_attn1.to_*"`, `"audio_attn2.to_*"` patterns to avoid touching video weights. See `configs/t2a_lora.yaml` for the exact list.
+
+## Validation Sample Prompts
+
+The example configs ship with placeholder validation prompts. Validation condition fields are documented in
+[`configuration-reference.md#validation-condition-types`](../../../../packages/ltx-trainer/docs/configuration-reference.md#validation-condition-types).
+For style/concept LoRAs:
+- Replace at least one `validation.samples[].prompt` with a prompt that uses the user's trigger word or describes the target concept. Tells the user something useful at the first validation interval.
+- Keep `validation.video_dims` consistent with the training resolution to make samples comparable.
+- **Describe the audio, for any run with a generated audio modality** (joint audio+video, T2A, V2A, etc.). The validation prompts must describe the audio the **same way the training captions do** — if the training captions transcribe speech or characterise sound (e.g. *"he says: ‘…’"*, *"calm spoken voice, quiet room tone"*, *"upbeat acoustic guitar"*), the validation prompts must include comparable audio direction. A prompt with no audio description gives the model no guidance for the audio branch and the generated audio comes out poor. This is not speech-specific — any audio (music, ambience, foley) needs describing. Mirror the structure/level of audio detail found in the dataset captions (inspect a few before writing the prompts).
+
+## W&B Patches
+
+- If the W&B credential check passes (`uv run python -c "import wandb; print(bool(wandb.Api().api_key))"` → `True`): `wandb.enabled: true`, `wandb.project` = `ltx2-<mode>`, `wandb.tags` includes the mode. (Do not use `wandb status` — it falsely reports `api_key: null` when logged in via netrc.)
+- If `False`: `wandb.enabled: false`. Surface in plan: *"Not logged in to W&B — run `wandb login` before training to enable tracking."* If the check errored/was ambiguous, ask the user rather than assuming off.
+
+## Output Dir Behaviour
+
+The trainer resumes optimizer/scheduler/step state **only when `model.load_checkpoint` is set** to a checkpoint file; it then looks for a matching `training_state_step_*.pt` next to that file. It does **not** auto-detect prior checkpoints in `output_dir/checkpoints/`. To resume, patch `model.load_checkpoint` to the latest checkpoint. To load weights but skip state restore: `checkpoints.no_resume: true`.
+
+For the orchestrator's resume flow: when resuming an interrupted run, patch `model.load_checkpoint` to the latest checkpoint under `output_dir/checkpoints/`. Leaving it unset starts a fresh run from step 0 even if checkpoints exist on disk.
+
+## Self-Check Before Launch
+
+Before any `python scripts/train.py` invocation:
+
+1. All `model_path`, `text_encoder_path`, `preprocessed_data_root` exist on disk.
+2. Frame and resolution constraints satisfied (see above).
+3. Generated modalities have matching latents directories under `.precomputed/`.
+4. For modes with `reference` condition: `reference_latents/` (and/or `reference_audio_latents/`) exists.
+5. For modes with `mask` condition: `video_masks/` (and/or `audio_masks/`) exists.
+
+A failed check at this point is much cheaper than a failed training start.

.claude/skills/train-model/references/hardware-profiles.md

@@ -0,0 +1,98 @@
+# VRAM Tiers
+
+Map probed GPU(s) to a starting training config. The autotune sweep in Phase 6 then empirically improves on this baseline. Use these **tier names** in `plan.md` and user-facing chat — not letter codes.
+
+Source of truth: the two configs shipped in the trainer repo —
+`packages/ltx-trainer/configs/t2v_lora.yaml` (standard) and
+`packages/ltx-trainer/configs/t2v_lora_low_vram.yaml` (low VRAM).
+Per `packages/ltx-trainer/docs/quick-start.md`, the trainer documents
+**80GB recommended** and **32GB minimum**. Anything below 32GB is
+unsupported by the project.
+
+## Probe
+
+```bash
+nvidia-smi --query-gpu=name,memory.total --format=csv,noheader
+```
+
+Pick the **smallest** VRAM tier across the visible GPUs. Multi-GPU only adds throughput at the same per-GPU memory budget — it doesn't relax per-GPU limits.
+
+## Minimum Gate
+
+If per-GPU VRAM is **< 32 GB**, stop the run. Surface to the user:
+
+> "This GPU has <N>GB VRAM. The LTX-2 trainer requires a minimum of 32GB (see `packages/ltx-trainer/docs/quick-start.md`). Training is unlikely to fit even with maximum memory savings, and we don't ship a tested config below 32GB. Options: (a) abort, (b) try anyway with the low-VRAM config and accept it may OOM — purely at your own risk."
+
+Do not invent a sub-32GB tier. The trainer team doesn't ship one.
+
+## 32GB tier — low-VRAM config
+
+**VRAM range:** 32 GB per GPU (trainer minimum).
+
+**Typical GPUs:** RTX 5090, V100 32GB.
+
+Start from `packages/ltx-trainer/configs/t2v_lora_low_vram.yaml` verbatim. Key choices already in that file (do not re-specify in `<workspace>/<run-name>/config.yaml` — copy the file and patch only the paths from `references/config-patching.md`):
+
+- `optimizer_type: "adamw8bit"`
+- `enable_gradient_checkpointing: true`
+- `batch_size: 1`, `gradient_accumulation_steps: 1`
+- `quantization: "int8-quanto"`
+- `load_text_encoder_in_8bit: true`
+- `offload_optimizer_during_validation: true`
+- `lora.rank: 16`, `lora.alpha: 16`
+
+Autotune (Phase 6) will sweep `quantization` off, `optimizer_type` → adamw, and `batch_size` up — but at 32GB the sweep often hits OOM on trial 2 or 3. That's fine; the conservative baseline still works.
+
+## 80GB+ tier — standard config
+
+**VRAM range:** 80 GB per GPU and above (trainer recommended).
+
+**Typical GPUs:** A100 80GB, H100 80GB, H200, B200.
+
+Start from `packages/ltx-trainer/configs/t2v_lora.yaml` verbatim. Key choices already in that file:
+
+- `optimizer_type: "adamw"`
+- `enable_gradient_checkpointing: true` (autotune may turn it off if headroom allows)
+- `batch_size: 1`, `gradient_accumulation_steps: 1`
+- `quantization: null`
+- `load_text_encoder_in_8bit: false`
+- `lora.rank: 32`, `lora.alpha: 32`
+
+On the 80GB+ tier, the autotune baseline already equals this config (adamw, no quantization), so the quantization/optimizer trials are no-ops; the only real lever is gradient checkpointing off — but for the 22B model that **usually OOMs even with tens of GB of apparent headroom**, so treat a win there as unlikely. The trainer reports its own step-time and peak-VRAM at the end of each run — use those rather than an external timer.
+
+For ≥140GB GPUs (H200, B200), the same 80GB+ tier baseline applies. FA3/FA4 attention backends are viable on Hopper/Blackwell and can speed up training, but they're optional — the trainer's defaults work on PyTorch SDPA without extra setup.
+
+## 40–60GB tier — mid-range (autotune from low-VRAM)
+
+**VRAM range:** 40–60 GB per GPU. The trainer doesn't ship a tested config for this range.
+
+**Typical GPUs:** A40, A6000 48GB, L40, RTX 6000 Ada.
+
+Start from the **32GB tier** (low-VRAM config) and let autotune relax `quantization`, `optimizer_type`, and `batch_size` based on actual headroom. Don't pre-bake intermediate YAML values that haven't been measured. Surface this as **40–60GB tier** in the plan.
+
+## Multi-GPU
+
+If `nvidia-smi` reports N ≥ 2 GPUs of the same model:
+
+- Launch with `uv run accelerate launch scripts/train.py <config>`.
+- Use `packages/ltx-trainer/configs/accelerate/fsdp.yaml` for full fine-tune.
+- DDP (default `accelerate launch` without a config file) is fine for LoRA.
+- Effective batch = `batch_size * gradient_accumulation_steps * num_gpus`. Reduce `gradient_accumulation_steps` proportionally to keep the effective batch consistent with the plan.
+
+## Full Fine-Tune
+
+If the user chose full fine-tune (`model.training_mode: "full"`):
+
+- Require multi-GPU + FSDP on 80GB+ tier GPUs. Otherwise warn in the plan that single-GPU full FT is unlikely to fit and propose LoRA instead.
+- Set `acceleration.offload_optimizer_during_validation: true` always (optimizer state is huge under full FT).
+
+## Model Path Constraints
+
+- `model.model_path`: local `.safetensors` only. No URLs.
+- `model.text_encoder_path`: local Gemma model directory. No URLs.
+
+If probe didn't find these in conventional locations (`/models/`, `~/models/`, `$LTX_MODELS_DIR`), ask the user in Phase 3 (or offer to download per `references/onboarding.md`).
+
+## Notes on Loss-of-Generality
+
+The two anchor tiers (32GB and 80GB+) correspond directly to the two configs the trainer ships. The autotune sweep is the empirical layer — if a particular GPU consistently lands on a different stable config, **update the relevant trainer config first**, not this file. This skill follows the trainer's choices, not the other way around.

.claude/skills/train-model/references/mode-selector.md

@@ -0,0 +1,74 @@
+# Mode Selector
+
+Map the user's stated intent to a `flexible`-strategy configuration. All modes are supported via a single strategy (`training_strategy.name: "flexible"`); the difference is which modality is generated and which `conditions` are attached.
+
+> Full reference: [`packages/ltx-trainer/docs/training-modes.md`](../../../../packages/ltx-trainer/docs/training-modes.md). This file is the **lookup table** for translating user intent.
+
+## Decision Table
+
+| User says (roughly)... | Mode | Example config | Modalities | Conditions |
+|------------------------|------|----------------|------------|------------|
+| "generate videos from text", "T2V LoRA" | T2V | `configs/t2v_lora.yaml` | video gen, audio gen | none |
+| "generate videos from a starting image", "I2V" | I2V | `configs/i2v_lora.yaml` | video gen, audio gen | `first_frame` (video) |
+| **plain concept/style LoRA** ("train a LoRA on X", no specific task) | **I2V by default** (see note) | `configs/i2v_lora.yaml` | video gen, audio gen | `first_frame` (video), `probability: 0.5` |
+| "extend a video forward in time" | Video extension (prefix) | `configs/video_extend_lora.yaml` | video gen, audio gen | `prefix` (video) |
+| "extend a video backward in time" | Video extension (suffix) | `configs/video_suffix_lora.yaml` | video gen, audio gen | `suffix` (video) |
+| "fill in masked regions of a video" | Video inpainting | `configs/video_inpainting_lora.yaml` | video gen | `mask` (video) |
+| "expand a video beyond its borders" | Video outpainting | `configs/video_outpainting_lora.yaml` | video gen | `spatial_crop` (video) |
+| "style transfer from reference video", "IC-LoRA", "depth/pose/canny control" | V2V IC-LoRA | `configs/v2v_ic_lora.yaml` | video gen | `reference` (video) |
+| "generate video to match an audio track" | A2V | `configs/a2v_lora.yaml` | video gen, audio frozen | none (audio `is_generated: false`) |
+| "add sound effects to silent video", "foley", "V2A" | V2A | `configs/v2a_lora.yaml` | video frozen, audio gen | none (video `is_generated: false`) |
+| "generate audio from text", "T2A" | T2A | `configs/t2a_lora.yaml` | audio gen | none |
+| "extend audio forward / backward" | Audio extension | `configs/audio_extend_lora.yaml`, `configs/audio_suffix_lora.yaml` | audio gen | `prefix` / `suffix` (audio) |
+| "fill in masked regions of audio" | Audio inpainting | `configs/audio_inpainting_lora.yaml` | audio gen | `mask` (audio) |
+| "audio style transfer from reference", "A2A IC-LoRA" | A2A IC-LoRA | `configs/a2a_ic_lora.yaml` | audio gen | `reference` (audio) |
+| "joint video+audio reference control" | AV2AV IC-LoRA | `configs/av2av_ic_lora.yaml` | video gen, audio gen | `reference` (both) |
+| Any of the above with full fine-tune | Full FT variant | as above, set `model.training_mode: "full"` | (mode-specific) | (mode-specific) |
+
+### Why I2V is the default for a plain concept/style LoRA
+
+A "train a LoRA on X" request isn't tied to one inference mode: **LoRA weights are pipeline-agnostic** — the same checkpoint loads in both T2V and I2V inference (both use `TI2VidOneStagePipeline`/`TwoStages`). The `i2v_lora` config trains `first_frame` with **`probability: 0.5`**, so the model learns both first-frame-conditioned (I2V) and unconditioned (T2V) generation in one run, and the first frame comes from each training clip automatically (no extra data prep). That makes I2V a versatile **superset** — usable for both at inference at no extra cost — which is why it's the default for a plain LoRA. Ask the user how they'll use it (text-only / from an image / both) and only drop to `t2v_lora` if they're sure it's text-only. (See the orchestrator `SKILL.md` Phase 1.)
+
+## Disambiguation Questions
+
+When the user's first answer is ambiguous, ask **one** follow-up:
+
+- "extend a video" → forward or backward in time?
+- "control with a reference" → video reference (depth/pose/canny/etc.) or audio reference?
+- "fill in regions" → video regions (masked frames) or audio regions (masked time)?
+- "T2V" → joint video+audio (default) or video-only?
+- LoRA or full fine-tune? Default to LoRA unless the user has multi-GPU + clear reason for full.
+
+## Combining Modes
+
+The flexible strategy allows stacking conditions. Common combinations:
+
+- **I2V + V2A** (start frame + generate audio for the resulting video) — not directly expressible; would need two passes.
+- **Video extension + audio extension** — both modalities generate, both have `prefix` condition. Express in one config.
+- **IC-LoRA + I2V** — `first_frame` + `reference` conditions on the video modality.
+
+If the user asks for a combination not listed, check `packages/ltx-trainer/src/ltx_trainer/training_strategies/flexible.py` for which conditions can co-exist on a modality. Audio modality cannot use `first_frame` or `spatial_crop`.
+
+## When the Intent Doesn't Map
+
+If after disambiguation there is no entry in the table and no combination of `flexible` conditions covers the user's request, go to the **Escape Hatch** section in the orchestrator `SKILL.md`. Do not silently pick the closest mode.
+
+## LoRA Rank by Use Case
+
+Once the mode is picked, choose `lora.rank` (and matching `lora.alpha`) based on what the LoRA is supposed to capture. These are starting points; autotune doesn't sweep rank because it's a quality knob, not a step-time one.
+
+| Use case | Suggested rank | Notes |
+|----------|----------------|-------|
+| Single character, single object, single style | 32–64 | Default for most concept LoRAs. Start at 32; bump to 64 if validation samples underfit. |
+| Multi-character world, dense series, complex multi-concept | 96–128 | More capacity for distinguishing several concepts inside one LoRA. |
+| Camera move, motion, transition (i.e. behavioural, not visual) | 8–16 | Motion is a thin signal — high ranks just memorise frame content. |
+| IC-LoRA control (V2V depth/pose/Canny/etc., A2A audio reference) | 16–32 | Start at 16 for structural control (depth, pose, edges); 24–32 if the reference carries richer style/texture. Video IC-LoRA often lands lower than concept LoRAs. |
+| LTX-2 trainer's default if unsure | 32 | Safe baseline. |
+
+On the **32GB tier** (`hardware-profiles.md`), the low-VRAM config already pins `lora.rank: 16`. If the user picks a higher-rank use case on 32GB, surface the trade-off in the plan but let them decide — the autotune sweep doesn't touch rank, so a too-high rank will simply OOM at training time.
+
+Keep `alpha == rank` unless the user has a specific reason otherwise (effective scaling = `alpha / rank`).
+
+## Starting Config
+
+Copy the matching example config — the exact filename from the **Example config** column of the decision table above (e.g. T2V → `configs/t2v_lora.yaml`, I2V → `configs/i2v_lora.yaml`, A2A IC-LoRA → `configs/a2a_ic_lora.yaml`) — into `<workspace>/<run-name>/config.yaml` as the starting point. Then patch per `references/config-patching.md`.

.claude/skills/train-model/references/onboarding.md

@@ -0,0 +1,115 @@
+# First-Run Onboarding
+
+What the orchestrator's Phase 2 probe checks for, what to do when something is missing, and what the skill is allowed to set up automatically (with explicit user approval).
+
+## Prerequisites Checked in Phase 2
+
+| Prerequisite | How to detect | If missing |
+|--------------|---------------|------------|
+| CUDA GPU visible | `nvidia-smi` returns ≥1 GPU | Stop. Training requires CUDA — point the user at non-LTX-2 docs. |
+| Linux | `uname -s` returns `Linux` | Stop. Trainer uses Triton (Linux-only). |
+| `uv` installed | `command -v uv` | Offer to install (see "Auto-setup" below). |
+| Workspace synced | `[ -f uv.lock ] && uv pip list \| grep -q ltx-trainer` | Offer to run `uv sync` from repo root. |
+| LTX-2 model weights | Search `/models/`, `~/models/`, `$LTX_MODELS_DIR` for a `.safetensors` matching `*ltx*2*` | Offer to download (see "Model downloads"). |
+| Gemma text encoder dir | Search same locations for a directory containing Gemma config | Offer to download. |
+| Captioner backend | Gemini auth (`GEMINI_API_KEY`/`GOOGLE_API_KEY` or gcloud/Vertex), OR a ≥40 GiB GPU to host the Qwen3-Omni-30B vLLM server (FP8), OR captions already in the dataset | See "Captioner graceful degradation" below. **Check the HF cache for an already-downloaded Qwen model before assuming a download is needed** (see note below the table). |
+| W&B login (optional) | `uv run python -c "import wandb; print(bool(wandb.Api().api_key))"` → `True` means logged in. Uses wandb's own credential resolution (env/netrc/settings). **Don't** use `wandb status` (reports `api_key: null` even when logged in). | Not a blocker. If `False`: disabled in config + flagged in plan. If the check errors/ambiguous: ask the user, don't assume off. |
+| Disk space | `df -h $WORKSPACE` | Surface available space alongside what a run consumes (preprocessed latents, several-GB checkpoints, validation samples). Flag concerns to the user; don't enforce a hard threshold. |
+
+Surface findings as a compact table in chat. For each missing item, present the user with a concrete next step (download command, install command, or "skip this — here's the consequence").
+
+**Check the HF cache before declaring the local captioner unavailable.** The Qwen3-Omni model is served from the HuggingFace cache, not `~/models/`. Before concluding it must be downloaded (or ruling it out on free-disk grounds), check whether it's already cached:
+
+```bash
+ls -d "${HF_HOME:-$HOME/.cache/huggingface}"/hub/models--Qwen--Qwen3-Omni* 2>/dev/null \
+  && du -sh "${HF_HOME:-$HOME/.cache/huggingface}"/hub/models--Qwen--Qwen3-Omni* 2>/dev/null
+```
+
+If it's cached (~60 GiB, all shards present), no download is needed — don't rule out the local captioner because of low free disk on the *home* partition; the weights already exist. Only the GPU-VRAM constraint (≥40 GiB for FP8) then applies.
+
+## Auto-Setup (with explicit user approval)
+
+The skill may, only after the user explicitly says yes, do these setup actions. Each action is a single discrete question.
+
+### Run `uv sync`
+
+```bash
+cd <repo-root>
+uv sync
+```
+
+Ask: *"Repo not synced. Run `uv sync` now? It will download the project's Python dependencies."*
+
+### Install `uv`
+
+Ask: *"`uv` not installed. Install via the official one-liner now?"*
+
+```bash
+curl -LsSf https://astral.sh/uv/install.sh | sh
+```
+
+After install, ask the user to restart their shell or `source ~/.bashrc` before continuing.
+
+### Model downloads
+
+Use `huggingface-cli` (comes with `huggingface-hub`, transitively pulled by `uv sync`). Default destination: `$LTX_MODELS_DIR` if set, else `~/models/` (create if missing). Surface destination in the prompt — never download into the repo or into the workspace.
+
+**LTX-2 base model:**
+
+```bash
+huggingface-cli download Lightricks/LTX-2.3 \
+  ltx-2.3-22b-dev.safetensors \
+  --local-dir ~/models/ltx-2.3
+```
+
+Public reference: <https://huggingface.co/Lightricks/LTX-2.3>.
+
+**Gemma text encoder:**
+
+```bash
+huggingface-cli download google/gemma-3-12b-it-qat-q4_0-unquantized \
+  --local-dir ~/models/gemma-3-12b
+```
+
+**Qwen3-Omni captioner (only if the GPU can host it):**
+
+The local captioner is Qwen3-Omni-30B-A3B-Thinking served by a vLLM server (`scripts/serve_captioner.py`), which downloads the model (~65 GiB) on first launch via `uvx vllm` — there's no separate `huggingface-cli` step. Default **FP8** (~31 GiB weights) fits on a **40 GiB** GPU; **bf16** (~60 GiB) needs **≥66 GiB free VRAM**. On a GPU below ~40 GiB (typical consumer 24/32 GB cards), don't use it — use Gemini instead (see "Captioner graceful degradation").
+
+The base model + text encoder are large (multi-GB) downloads. Ask the user before each one — `huggingface-cli` reports the actual size at the start of the transfer. Do **not** batch them into a single "yes/no"; the user may want only what's missing.
+
+### Hugging Face login (if any downloads fail with 401)
+
+Ask: *"Hugging Face download requires login (some Lightricks models are gated). Run `huggingface-cli login` now? You'll need a token from <https://huggingface.co/settings/tokens>."*
+
+## Captioner Graceful Degradation
+
+The captioner is the trickiest prerequisite. The local backend (`qwen_omni`) is now a **30B model served by a vLLM server** — ~65 GiB download; FP8 fits on a 40 GiB GPU, bf16 needs ≥66 GiB. Typical consumer cards (24/32 GB) can't host it, so for most users **prefer Gemini**.
+
+Decision tree:
+
+1. **User already has captions in their dataset metadata** → skip the captioner entirely (Step 3 of `prepare-dataset` skips when captions are present).
+2. **GPU ≥40 GiB (FP8) / ≥66 GiB (bf16)** → `qwen_omni` is viable: launch `scripts/serve_captioner.py` first, then caption. Gemini is still fine here too.
+3. **GPU below ~40 GiB (the common consumer case)** → `qwen_omni` is not an option. Recommend **`gemini_flash`** and tell the user they need Gemini auth: a `GEMINI_API_KEY`/`GOOGLE_API_KEY` (get one at <https://aistudio.google.com/apikey>) **or** working gcloud/Vertex AI credentials.
+4. **No Gemini auth and can't run Qwen3** → the only remaining path is bringing their own captions: add a `caption` column to the dataset metadata, then re-invoke the skill.
+
+Wait for the user's choice. Don't pick one automatically — but make the hardware reality explicit so they don't try to run the server on a card that can't host it.
+
+## What Auto-Setup Does NOT Touch
+
+- Does not modify the user's shell rc files except by explicit instruction (e.g., "add `export LTX_MODELS_DIR=...` to your `.bashrc`?" with the user agreeing).
+- Does not modify `~/.gitconfig`, `~/.ssh/`, or any auth-related files.
+- Does not install GPU drivers, CUDA, or system packages.
+- Does not delete or move existing model files. If a model is found but at an unexpected path, surface and let the user decide.
+- Does not download into the repo or workspace. Models live at `~/models/` (or `$LTX_MODELS_DIR`).
+
+## Configuration Inheritance
+
+After downloads, the skill records the resolved paths into the plan's Assumptions section and into `config.yaml`:
+
+```yaml
+model:
+  model_path: "/home/<user>/models/ltx-2.3/ltx-2.3-22b-dev.safetensors"
+  text_encoder_path: "/home/<user>/models/gemma-3-12b"
+```
+
+Suggest (don't enforce) setting `LTX_MODELS_DIR=~/models` in their shell rc for future runs.

.claude/skills/train-model/references/plan-template.md

@@ -0,0 +1,107 @@
+# Plan Template
+
+Write `<workspace>/<run-name>/plan.md` following this template. The plan is the user's contract with the agent — it's the gate before any heavy work runs.
+
+Scale each section to its relevance. Skip subsections that don't apply (e.g., no captioning if data is already captioned), but never collapse "Assumptions" or "Cost/time estimate".
+
+```markdown
+# Training Plan — <run-name>
+
+## Goal
+
+<One paragraph restating the user's intent in their own terms.>
+
+## Mode
+
+**<Mode name>** — <one-line rationale linking user intent to mode>.
+
+- Config base: the concrete example config selected by the mode (e.g. `packages/ltx-trainer/configs/t2v_lora.yaml`) — use the actual filename, not a placeholder
+- Conditions: <list, or "none">
+- Training mode: <`lora` | `full`>
+
+## Dataset
+
+- Source: `<absolute path>` (<N samples>)
+- Captions: <`already present` | `will be generated with <backend>`>
+- Audio: <`present` | `absent — using --skip-audio` | `to be paired with --audio-durations`>
+- IC-LoRA references: <`present` | `to be generated via compute_reference.py` | `n/a`>
+
+## Preprocessing
+
+- Target resolution buckets: `<W>x<H>x<F>` (frames satisfy `frames % 8 == 1`; W,H divisible by 32)
+- Estimated time: ~<duration> on detected hardware
+- Output: `<workspace>/<run-name>/dataset/.precomputed/`
+
+## Training Config
+
+| Field | Value |
+|-------|-------|
+| Optimizer | <adamw / adamw8bit> |
+| Mixed precision | <bf16 / fp16> |
+| Quantization | <null / int8-quanto / ...> |
+| Gradient checkpointing | <on / off> |
+| Batch size | <N> |
+| Gradient accumulation | <N> (effective batch = <N>) |
+| Steps | <N> |
+| Learning rate | <value> |
+| LoRA rank / alpha | <N / N> (or "full FT") |
+| LoRA target modules | <list> (or "n/a") |
+| LoRA trigger word | <word> (or "n/a") |
+| Validation interval | every <N> steps |
+| Checkpoint interval | every <N> steps |
+
+## Hardware
+
+- GPU(s): <name> x <count>, <VRAM>GB per GPU
+- Launch: <`python scripts/train.py` (single) | `accelerate launch` (multi)>
+- VRAM tier: <32GB tier | 40–60GB tier | 80GB+ tier> — <low-VRAM config (`t2v_lora_low_vram.yaml`) | standard config (`t2v_lora.yaml`) | mid-range, autotuned from low-VRAM>
+
+## Sanity Check + Autotune
+
+*In plain terms: before committing to the full run, I do a quick dry run on a single clip at your target resolution to catch out-of-memory or config errors in ~2 minutes (rather than failing hours in), then try a few config variants to pick the fastest one that fits your GPU.*
+
+Mechanics:
+- 1 sample, full target resolution, 50 steps + 1 validation pass.
+- Autotune sweep: up to 5 trials varying quantization / optimizer / batch size.
+- Stops at first OOM or no-improvement.
+
+## Monitoring
+
+<One of:>
+- W&B: enabled, project `<name>`, entity `<entity-or-default>`. URL will be surfaced once training starts.
+- W&B: **not logged in** — run `wandb login` before training to enable tracking. Otherwise training proceeds without remote logging.
+
+## Outputs
+
+- Training config: `<workspace>/<run-name>/config.yaml`
+- Checkpoints: `<workspace>/<run-name>/outputs/checkpoints/`
+- Validation samples: `<workspace>/<run-name>/outputs/samples/`
+- Autotune log: `<workspace>/<run-name>/autotune.log`
+
+## Assumptions
+
+Defaults the agent chose silently. Override any by replying with the new value.
+
+- <list every non-trivial assumed value: precision, scheduler type, seed, validation prompts, checkpoint retention, etc.>
+
+## Cost / Time Estimate
+
+Give only estimates you can ground; label anything not yet measured as rough. Do **not** state a confident training duration before the sanity check has measured a real step time — say "training duration TBD until the sanity check measures step time" and fill it in afterward (per Hard Invariant #5: no fabricated predictions).
+
+- Captioning: ~<duration> (rough)
+- Preprocessing: ~<duration> (rough)
+- Sanity check + autotune: ~<duration> (rough)
+- Full training: **measured after sanity check** — then `<measured step-time> × <steps>`
+- **Total wall-clock estimate:** rough until step time is measured; refine after the sanity check.
+
+## Approve to Proceed
+
+Reply "approve" (or with edits) to start. No captioning, preprocessing, autotune, or training will run before approval.
+```
+
+## Notes on Writing the Plan
+
+- Show numbers, not adjectives. "~3 hours" beats "fairly long."
+- Surface every assumption that, if wrong, would cost the user time. Better to over-list than under-list — the user can skim.
+- If a section reveals you need to ask another question, **stop and ask** before finalizing the plan. The plan is the last gate, not the first.
+- If the user's hardware can't reasonably support the requested mode (e.g., single-GPU full FT on a 32GB consumer card), say so plainly in the plan and propose the alternative (LoRA, multi-GPU, etc.), rather than silently downgrading.

.claude/skills/train-model/references/troubleshooting.md

@@ -0,0 +1,84 @@
+# Troubleshooting
+
+Quick lookup for failures during sanity check, preprocessing, or training. For deeper coverage see `packages/ltx-trainer/docs/troubleshooting.md`.
+
+## OOM During Training Step
+
+Order of operations (cheapest first):
+
+1. `optimization.enable_gradient_checkpointing: true` (if not already on).
+2. `optimization.batch_size: 1` and increase `gradient_accumulation_steps` to preserve effective batch.
+3. `optimization.optimizer_type: "adamw8bit"`.
+4. `acceleration.quantization: "int8-quanto"`.
+5. Reduce `lora.rank` (32 → 16 → 8). Alpha follows rank.
+6. Reduce target resolution (`validation.video_dims` and re-preprocess the dataset at the new resolution).
+
+The last option is expensive — flag it clearly to the user before re-preprocessing.
+
+## OOM During Validation Sample Generation
+
+The validation pass loads decoders + runs CFG/STG inference; it can OOM even when the training step fits.
+
+1. `acceleration.load_text_encoder_in_8bit: true` (trainer config, not the dataset script).
+2. `acceleration.offload_optimizer_during_validation: true` (especially for full FT or high-rank LoRA).
+3. Reduce `validation.video_dims` (smaller validation than training is fine — it's only for visual feedback).
+4. Reduce `validation.inference_steps` (e.g. 30 → 20).
+5. Increase `validation.interval` to validate less often.
+
+## NaN Loss
+
+1. Check `acceleration.mixed_precision_mode`: prefer `"bf16"`. If `"fp16"`, switch.
+2. Verify dataset latents are well-formed: `uv run python scripts/decode_latents.py <latents-dir> <output-dir> --model-path <model>` (it decodes a whole latents directory, not a single `.pt`) should reconstruct sensibly.
+3. Lower `optimization.learning_rate` by 5x.
+4. Add `optimization.max_grad_norm: 1.0` (default; verify it's set).
+5. If using `quantization`, try `null` — INT8/INT4 quantization can interact badly with poorly-conditioned LoRA inits at high LR.
+
+## Validation Samples Look Wrong but Loss Is Fine
+
+Often not a bug — validation uses simplified inference. For real quality assessment, run a checkpoint through `packages/ltx-pipelines/` after training.
+
+## Trainer Won't Start: Config Validation Error
+
+Pydantic `extra="forbid"` means typos in field names fail loudly. Read the error carefully — it names the offending field and path. Fix and re-launch.
+
+Common offenders:
+- `latents_dir` typo or wrong relative path.
+- A required field genuinely missing after copying an example (most fields have defaults; check the error message for the exact field path).
+- `target_modules` listed at wrong nesting level (must be under `lora:`).
+
+## Trainer Won't Start: Missing Files
+
+- `model_path` not found → re-probe `/models/`, `~/models/`, `$LTX_MODELS_DIR`, or ask the user.
+- `text_encoder_path` directory missing the Gemma config → ensure the path is to the Gemma model dir, not its parent.
+- `preprocessed_data_root` doesn't contain expected subdirs → re-verify Phase 7 ran for the chosen mode (see `phases/preprocess-dataset.md`).
+
+## Resume Stops Working
+
+The trainer **does not auto-resume from `output_dir`**. Resume happens only when `model.load_checkpoint` is explicitly set to a checkpoint file; the trainer then loads those weights and looks for a `training_state_step_*.pt` next to that file to restore optimizer/scheduler/step. Common pitfalls:
+- `model.load_checkpoint` not set or set to the wrong path → fresh run from step 0 even when `outputs/checkpoints/` is full of artifacts. Patch `model.load_checkpoint` to the latest checkpoint.
+- `checkpoints.no_resume: true` is set → weights load but state is discarded. Remove the flag if you want a proper resume.
+- `training_state_step_*.pt` missing from next to the loaded checkpoint → weights load but step counter resets. Make sure the state file accompanies the checkpoint.
+- `training_state_step_*.pt` corrupted (size 0, fails `torch.load`) → trainer falls back to step 0 with a warning.
+
+## Autotune Trial Failed Mid-Sweep
+
+- If trial 2 (quantization off) OOMs: revert to trial 1 and stop the sweep. The 32GB tier baseline is correctly aggressive.
+- If trial 3 (adamw) OOMs: revert to adamw8bit. Continue with trial 4 if VRAM headroom allows.
+- If trial 4 (batch_size up) OOMs: revert and stop. We've found the ceiling.
+
+Never carry over a failing trial's deltas. Always revert to the last-known-good before the next change.
+
+## Captioning Is Slow
+
+- Local `qwen_omni` is a 30B model served by `serve_captioner.py` (vLLM). If the server won't start or OOMs on launch: keep the default `--quantization fp8` (don't use `bf16` unless ≥66 GiB free VRAM), lower `--gpu-memory-utilization`, or reduce `--max-model-len`. If the GPU can't host a 30B model at all, switch to `gemini_flash`.
+- `caption_videos.py --captioner-type qwen_omni` errors connecting → the vLLM server isn't running or `--vllm-url` doesn't match. Start `serve_captioner.py` first and confirm the port.
+- For most hardware and for larger datasets, prefer `--captioner-type gemini_flash --num-workers <N>` (needs Gemini auth: `GEMINI_API_KEY`/`GOOGLE_API_KEY` or gcloud/Vertex) — runs anywhere and parallelises; local Qwen needs a heavy GPU and a running server.
+
+## Process_Dataset Errors
+
+- "frames divisible by..." → the video doesn't have enough frames at the requested temporal resolution. Either shorten the requested frame count or use `split_scenes.py` to break long videos.
+- "shape mismatch" on existing `.precomputed/` → user requested a different resolution than the existing data. Per the invariants, **stop and ask** — do not overwrite. Offer: reuse at old resolution / re-preprocess to a new dir / abort.
+
+## When To Give Up and Ask The User
+
+If a fix isn't obvious from this file or `packages/ltx-trainer/docs/troubleshooting.md` within two attempts, stop and surface the full error + the steps already tried to the user. Don't loop indefinitely on autonomous fixes — the user has context the agent doesn't (which checkpoints are precious, what they care about preserving, etc.).

.gitattributes

@@ -2,8 +2,11 @@
 *.safetensors filter=lfs diff=lfs merge=lfs -text
 *.sft filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
+*.mp3 filter=lfs diff=lfs merge=lfs -text
 *.mp4 filter=lfs diff=lfs merge=lfs -text
+packages/ltx-pipelines/tests/assets/*.wav filter=lfs diff=lfs merge=lfs -text
 *.png filter=lfs diff=lfs merge=lfs -text
 *.jpeg filter=lfs diff=lfs merge=lfs -text
 *.jpg filter=lfs diff=lfs merge=lfs -text
 *.webp filter=lfs diff=lfs merge=lfs -text
+*.exr filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -17,8 +17,10 @@ checkpoints/
 
 # Other files
 .DS_Store
-tmp
 .wandb
+projects/
+tmp
+wandb/
 
 # Model checkpoints
 *.ckpt
@@ -27,6 +29,7 @@ tmp
 *.sft
 
 # Media files
+*.exr
 *.gif
 *.heic
 *.heif
@@ -35,10 +38,32 @@ tmp
 *.json
 *.m4a
 *.mov
+*.mp3
 *.mp4
 *.png
 *.wav
 *.webp
 
+# HDR IC-LoRA e2e test baseline (checked in via Git LFS)
+!packages/ltx-pipelines/tests/assets/expected_hdr_ic_lora_exr/frame_*.exr
+!packages/ltx-pipelines/tests/assets/hdr_ic_lora_test_input.mp4
+
+# Text-to-audio (T2A) e2e test baseline (checked in via Git LFS)
+!packages/ltx-pipelines/tests/assets/expected_t2a_one_stage_ltx2_3.wav
+
+# ltx-bench Grafana dashboards (source of truth in the repo)
+!packages/ltx-bench/grafana/*.json
+
+# ltx-trainer E2E test dataset (committed via Git LFS).
+# The target_audio/ and reference_audio/ subdirectories are populated lazily at test runtime
+# (high-bitrate MP3 extracted from target_videos, plus a low-bitrate MP3 reference); their
+# contents stay gitignored via the root *.mp3 rule above.
+!packages/ltx-trainer/tests/assets/test_dataset/dataset.json
+!packages/ltx-trainer/tests/assets/test_dataset/target_videos/*.mp4
+!packages/ltx-trainer/tests/assets/test_dataset/reference_videos/*.mp4
+!packages/ltx-trainer/tests/assets/test_dataset/conditioning/first_frame.jpg
+!packages/ltx-trainer/tests/assets/test_dataset/conditioning/video_mask.png
+!packages/ltx-trainer/tests/assets/test_dataset/conditioning/audio_mask.pt
+
 # Binary files
 *.so

README.md

@@ -39,7 +39,7 @@ Download the following models from the [LTX-2.3 HuggingFace repository](https://
 **Temporal Upscaler** - Supported by the model and will be required for future pipeline implementations
   * [`ltx-2.3-temporal-upscaler-x2-1.0.safetensors`](https://huggingface.co/Lightricks/LTX-2.3/blob/main/ltx-2.3-temporal-upscaler-x2-1.0.safetensors) - [Download](https://huggingface.co/Lightricks/LTX-2.3/resolve/main/ltx-2.3-temporal-upscaler-x2-1.0.safetensors)
 
-**Distilled LoRA** - Required for current two-stage pipeline implementations in this repository (except DistilledPipeline and ICLoraPipeline)
+**Distilled LoRA** - Required for current two-stage pipeline implementations in this repository (except DistilledPipeline, ICLoraPipeline, and LipDubPipeline)
   * [`ltx-2.3-22b-distilled-lora-384-1.1.safetensors`](https://huggingface.co/Lightricks/LTX-2.3/blob/main/ltx-2.3-22b-distilled-lora-384-1.1.safetensors) - [Download](https://huggingface.co/Lightricks/LTX-2.3/resolve/main/ltx-2.3-22b-distilled-lora-384-1.1.safetensors)
 
 **Gemma Text Encoder** (download all assets from the repository)
@@ -57,6 +57,8 @@ Download the following models from the [LTX-2.3 HuggingFace repository](https://
   * [`LTX-2-19b-LoRA-Camera-Control-Jib-Down`](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Jib-Down) - [Download](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Jib-Down/resolve/main/ltx-2-19b-lora-camera-control-jib-down.safetensors)
   * [`LTX-2-19b-LoRA-Camera-Control-Jib-Up`](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Jib-Up) - [Download](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Jib-Up/resolve/main/ltx-2-19b-lora-camera-control-jib-up.safetensors)
   * [`LTX-2-19b-LoRA-Camera-Control-Static`](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Static) - [Download](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Static/resolve/main/ltx-2-19b-lora-camera-control-static.safetensors)
+  * [`LTX-2.3-22b-IC-LoRA-HDR`](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-HDR) - HDR IC-LoRA and pre-computed text embeddings for `HDRICLoraPipeline`
+  * [`LTX-2.3-22b-IC-LoRA-LipDub`](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-LipDub) - [Download](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-LipDub/resolve/main/ltx-2.3-22b-ic-lora-lipdub-0.9.safetensors)
 
 ### Available Pipelines
 
@@ -68,12 +70,14 @@ Download the following models from the [LTX-2.3 HuggingFace repository](https://
 * **[KeyframeInterpolationPipeline](packages/ltx-pipelines/src/ltx_pipelines/keyframe_interpolation.py)** - Interpolate between keyframe images
 * **[A2VidPipelineTwoStage](packages/ltx-pipelines/src/ltx_pipelines/a2vid_two_stage.py)** - Audio-to-video generation conditioned on an input audio file
 * **[RetakePipeline](packages/ltx-pipelines/src/ltx_pipelines/retake.py)** - Regenerate a specific time region of an existing video
+* **[HDRICLoraPipeline](packages/ltx-pipelines/src/ltx_pipelines/hdr_ic_lora.py)** - Video-to-video with HDR output (linear float frames via LogC3 inverse decode, suitable for EXR export and tonemapping)
+* **[LipDubPipeline](packages/ltx-pipelines/src/ltx_pipelines/lipdub.py)** - Lip dubbing, rephrasing, matching speaker identity (distilled model, single IC-LoRA, Two stages).
 
 ### ⚡ Optimization Tips
 
 * **Use DistilledPipeline** - Fastest inference with only 8 predefined sigmas (8 steps stage 1, 4 steps stage 2)
 * **Enable FP8 quantization** - Enables lower memory footprint: `--quantization fp8-cast` (CLI) or `quantization=QuantizationPolicy.fp8_cast()` (Python). Fp8-cast should be used with bf16 checkpoints, it shall downcast them on the fly. For Hopper GPUs with TensorRT-LLM, use `--quantization fp8-scaled-mm` for FP8 scaled matrix multiplication. Fp8-scaled-mm should be used with fp8 checkpoints.
-* **Install attention optimizations** - Use xFormers (`uv sync --extra xformers`) or [Flash Attention 3](https://github.com/Dao-AILab/flash-attention) for Hopper GPUs
+* **Install attention optimizations** - On datacenter Blackwell GPUs (B200), install FlashAttention 4 manually: `uv pip install 'flash-attn-4==4.0.0b9'` (this specific revision is the one we have verified against torch 2.9.1+cu128; newer betas have known issues on consumer Blackwell). On other CUDA GPUs (including Hopper), use xFormers (`uv sync --extra xformers`).
 * **Use gradient estimation** - Reduce inference steps from 40 to 20-30 while maintaining quality (see [pipeline documentation](packages/ltx-pipelines/README.md#denoising-loop-optimization))
 * **Skip memory cleanup** - If you have sufficient VRAM, disable automatic memory cleanup between stages for faster processing
 * **Choose single-stage pipeline** - Use `TI2VidOneStagePipeline` for faster generation when high resolution isn't required

packages/ltx-core/README.md

@@ -8,6 +8,7 @@ The foundational library for the LTX-2 Audio-Video generation model. This packag
 - **`conditioning/`**: Tools for preparing latent states and applying conditioning (image, video, keyframes)
 - **`guidance/`**: Perturbation system for fine-grained control over attention mechanisms
 - **`loader/`**: Utilities for loading weights from `.safetensors`, fusing LoRAs, and managing memory
+- **`block_streaming/`**: Memory-efficient inference that streams transformer blocks through the GPU one at a time (from pinned CPU buffers or directly from disk)
 - **`model/`**: PyTorch implementations of the LTX-2 Transformer, Video VAE, Audio VAE, Vocoder and Upscaler
 - **`text_encoders/gemma`**: Gemma text encoder implementation with tokenizers, feature extractors, and separate encoders for audio-video and video-only generation
 - **`quantization/`**: FP8 quantization backends (FP8-TensorRT-LLM scaled MM, FP8 cast) for reduced memory footprint.
@@ -55,6 +56,7 @@ pip install -e packages/ltx-core
 
 - **Loader** ([`loader/`](src/ltx_core/loader/)): Model loading from `.safetensors`, LoRA fusion, weight remapping, and memory management
 - **Quantization** ([`quantization/`](src/ltx_core/quantization/)): FP8 quantization backends for reduced memory footprint and faster inference
+- **Block Streaming** ([`block_streaming/`](src/ltx_core/block_streaming/)): Streams transformer blocks through the GPU one block at a time, so the full model runs on machines without enough memory to hold all its weights at once
 
 ### Loader
 
@@ -77,13 +79,17 @@ model = builder.build(device=torch.device("cuda"))
 Use the `.lora()` method to attach one or more LoRA adapters before calling `.build()`:
 
 ```python
+from ltx_core.loader import SDOps
+
+lora_sd_ops = SDOps(name="identity").with_matching()  # or a model-specific key-renaming SDOps
+
 builder = (
     SingleGPUModelBuilder(
         model_class_configurator=MyModelConfigurator,
         model_path="/path/to/model.safetensors",
     )
-    .lora("/path/to/lora_a.safetensors", strength=0.8)
-    .lora("/path/to/lora_b.safetensors", strength=0.5)
+    .lora("/path/to/lora_a.safetensors", 0.8, lora_sd_ops)
+    .lora("/path/to/lora_b.safetensors", 0.5, lora_sd_ops)
 )
 model = builder.build(device=torch.device("cuda"))
 ```
@@ -103,7 +109,7 @@ builder = SingleGPUModelBuilder(
     model_class_configurator=MyModelConfigurator,
     model_path="/path/to/model.safetensors",
     lora_load_device=torch.device("cuda"),
-).lora("/path/to/lora.safetensors", strength=1.0)
+).lora("/path/to/lora.safetensors", 1.0, lora_sd_ops)
 
 model = builder.build(device=torch.device("cuda"))
 ```
@@ -121,39 +127,26 @@ Uses NVIDIA TensorRT-LLM's `cublas_scaled_mm` for efficient FP8 matrix multiplic
 **Usage with QuantizationPolicy:**
 
 ```python
-from ltx_core.quantization import QuantizationPolicy
+from ltx_core.quantization.fp8_scaled_mm import build_policy as build_fp8_scaled_mm_policy
 
-# Dynamic input quantization (no calibration needed)
-policy = QuantizationPolicy.fp8_scaled_mm()
-
-# Static input quantization with calibration file
-policy = QuantizationPolicy.fp8_scaled_mm(calibration_amax_path="/path/to/amax.json")
+# Discovers the layer set from the checkpoint's .weight_scale tensors
+policy = build_fp8_scaled_mm_policy("/path/to/checkpoint.safetensors")
 ```
 
-The policy provides `sd_ops` and `module_ops` that can be passed to the model builder:
+The policy carries `sd_ops`, `module_ops`, and `fuse_rule` that are passed to the model builder:
 
 ```python
+import torch
 from ltx_core.loader import SingleGPUModelBuilder
 
 builder = SingleGPUModelBuilder(
-    model=model,
-    device=device,
-    sd_ops=policy.sd_ops,
+    model_class_configurator=MyModelConfigurator,
+    model_path="/path/to/checkpoint.safetensors",
+    model_sd_ops=policy.sd_ops,
     module_ops=policy.module_ops,
+    fuse_rule=policy.fuse_rule,
 )
-builder.load(checkpoint_path)
-```
-
-**Calibration File Format** (for static input quantization):
-
-```json
-{
-  "amax_values": {
-    "transformer_blocks.0.attn.to_q.input_quantizer": 12.5,
-    "transformer_blocks.0.attn.to_k.input_quantizer": 8.3,
-    ...
-  }
-}
+model = builder.build(device=torch.device("cuda"))
 ```
 
 #### FP8 Cast
@@ -161,7 +154,42 @@ builder.load(checkpoint_path)
 A simpler approach that casts weights to FP8 for storage and upcasts during inference:
 
 ```python
-policy = QuantizationPolicy.fp8_cast()
+from ltx_core.quantization.fp8_cast import build_policy as build_fp8_cast_policy
+
+policy = build_fp8_cast_policy("/path/to/checkpoint.safetensors")
+```
+
+### Block Streaming
+
+The `block_streaming/` module ([`src/ltx_core/block_streaming/`](src/ltx_core/block_streaming/)) lets the full model run on machines that lack the memory to hold all of its weights at once. It streams the transformer's blocks through a small rolling set of GPU buffers, loading each block's weights just before it runs and recycling them afterwards, so only a few blocks are resident on the GPU at any moment. Construct it with `StreamingModelBuilder`, which returns a `BlockStreamingWrapper` -- an `nn.Module` drop-in for the wrapped model.
+
+#### Strategies
+
+The strategy is chosen automatically from `cpu_slots_count` relative to the number of blocks:
+
+- **RAM streaming** (default, `cpu_slots_count` omitted or `>= num_blocks`): all blocks are pre-loaded into pinned CPU buffers (with LoRA fusion) at build time, then copied to the GPU on demand. Fast; higher CPU memory.
+- **Disk streaming** (`cpu_slots_count < num_blocks`): blocks are read from the `.safetensors` file on demand on a background worker thread. Slower; lowest CPU memory.
+
+#### Basic usage
+
+```python
+import torch
+from ltx_core.block_streaming import StreamingModelBuilder
+
+builder = StreamingModelBuilder(
+    model_class_configurator=MyModelConfigurator,
+    model_path="/path/to/model.safetensors",
+    blocks_attr="transformer_blocks",      # dotted path to the nn.ModuleList
+    blocks_prefix="transformer_blocks",    # state-dict key prefix for block weights
+)
+
+# Omit cpu_slots_count for RAM streaming; pass a value < num_blocks for disk streaming.
+model = builder.build(
+    device=torch.device("cuda"),
+    dtype=torch.bfloat16,
+    cpu_slots_count=4,
+    gpu_slots_count=2,
+)
 ```
 
 For complete, production-ready pipeline implementations that combine these building blocks, see the [`ltx-pipelines`](../ltx-pipelines/) package.

packages/ltx-core/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ltx-core"
-version = "1.1.1"
+version = "1.1.6"
 description = "Core implementation of Lightricks' LTX-2 model"
 readme = "README.md"
 requires-python = ">=3.10"

packages/ltx-core/src/ltx_core/batch_split.py

@@ -1,5 +1,5 @@
 """Batch-splitting adapter for the transformer.
-Wraps an ``X0Model`` (or ``LayerStreamingWrapper``) and splits batched inputs
+Wraps an ``X0Model`` (or ``BlockStreamingWrapper``) and splits batched inputs
 into smaller chunks before forwarding, then concatenates the results. This
 controls peak activation memory at the cost of more forward passes.
 The adapter is transparent — it has the same ``forward`` signature as
@@ -42,7 +42,7 @@ class BatchSplitAdapter(nn.Module):
     Has the same ``forward`` signature as ``X0Model``:
     ``(video, audio, perturbations) -> (denoised_video, denoised_audio)``.
     Args:
-        model: The model to wrap (``X0Model``, ``LayerStreamingWrapper``, etc.).
+        model: The model to wrap (``X0Model``, ``BlockStreamingWrapper``, etc.).
         max_batch_size: Maximum batch size per forward pass. Input batches
             larger than this are split into sequential chunks.
     """

packages/ltx-core/src/ltx_core/block_streaming/__init__.py

@@ -0,0 +1,20 @@
+"""Block streaming: memory-efficient sequential-block inference.
+Streams transformer blocks from safetensors to GPU one at a time.
+Block weights are provided by a :class:`WeightsProvider` which handles
+CPU-to-GPU copies, caching, and stream synchronization.  Two weight
+source strategies are available:
+- **RAM streaming** (default): all blocks pre-loaded into pinned CPU
+  buffers with LoRA fusion at build time.  Fast, higher CPU memory.
+- **Disk streaming** (``cpu_slots < blocks_number``): blocks are read from
+  disk on demand by a :class:`DiskWeightSource`, on a background worker
+  thread.  Slower, lower CPU memory.
+"""
+
+from ltx_core.block_streaming.builder import DISK_CPU_SLOTS, StreamingModelBuilder
+from ltx_core.block_streaming.wrapper import BlockStreamingWrapper
+
+__all__ = [
+    "DISK_CPU_SLOTS",
+    "BlockStreamingWrapper",
+    "StreamingModelBuilder",
+]

packages/ltx-core/src/ltx_core/block_streaming/block_fetcher.py

@@ -0,0 +1,83 @@
+"""BlockFetcher: async disk reads on a worker thread."""
+
+from __future__ import annotations
+
+import logging
+import queue
+import threading
+from dataclasses import dataclass
+
+import torch
+
+from ltx_core.block_streaming.disk import DiskBlockReader
+
+logger = logging.getLogger(__name__)
+
+Buffer = dict[str, torch.Tensor]
+
+
+@dataclass(slots=True)
+class FetchHandle:
+    """Caller-facing handle for an outstanding read returned by :meth:`BlockFetcher.submit`.
+    Carries only what the caller needs: a completion event the worker sets and the
+    read's error. The worker updates these once the read finishes.
+    """
+
+    _done: threading.Event
+    _error: BaseException | None = None
+
+    def wait(self) -> BaseException | None:
+        """Block until the read finishes; return its error, or ``None`` on success."""
+        self._done.wait()
+        return self._error
+
+
+@dataclass(slots=True)
+class _ReadRequest:
+    """One outstanding read, internal to :class:`BlockFetcher`.
+    The fetcher's worker reads block ``idx`` into the caller-carved ``buffer`` and
+    updates ``handle`` (its error, then its event) once the read has finished.
+    """
+
+    idx: int
+    buffer: Buffer
+    handle: FetchHandle
+
+
+class BlockFetcher:
+    """Fills caller-supplied buffers on a worker thread."""
+
+    def __init__(self, reader: DiskBlockReader) -> None:
+        self._reader = reader
+        self._request_queue: queue.SimpleQueue[_ReadRequest | None] = queue.SimpleQueue()
+        self._worker = threading.Thread(target=self._run, name="BlockFetcher-IO", daemon=True)
+        self._worker.start()
+
+    def submit(self, idx: int, buffer: Buffer) -> FetchHandle:
+        """Enqueue a read of block *idx* into the caller-carved *buffer*, return its handle."""
+        handle = FetchHandle(_done=threading.Event())
+        request = _ReadRequest(idx=idx, buffer=buffer, handle=handle)
+        self._request_queue.put(request)
+        return handle
+
+    def cleanup(self) -> None:
+        """Drain pending reads, join the worker, close the reader."""
+        self._request_queue.put(None)
+        self._worker.join()
+        self._reader.cleanup()
+
+    def _run(self) -> None:
+        # Pinned buffers are allocated under the caller's inference_mode, so
+        # in-place copy_ from this thread requires inference_mode here too.
+        with torch.inference_mode():
+            while True:
+                request = self._request_queue.get()
+                if request is None:
+                    return
+
+                try:
+                    self._reader.read_into(request.buffer, request.idx)
+                except Exception as exc:
+                    logger.exception("BlockFetcher: fetch failed for item %d", request.idx)
+                    request.handle._error = exc
+                request.handle._done.set()

packages/ltx-core/src/ltx_core/block_streaming/builder.py

@@ -0,0 +1,476 @@
+"""Builder that constructs a BlockStreamingWrapper from safetensors checkpoints."""
+
+from __future__ import annotations
+
+import copy
+import logging
+from dataclasses import replace
+from typing import TYPE_CHECKING, Final, Generic
+
+import safetensors
+import torch
+from torch import nn
+
+from ltx_core.block_streaming import utils as bs_utils
+from ltx_core.block_streaming.block_fetcher import BlockFetcher
+from ltx_core.block_streaming.disk import DiskBlockReader, DiskTensorReader, LoraSource
+from ltx_core.block_streaming.pool import BufferPool
+from ltx_core.block_streaming.provider import WeightsProvider
+from ltx_core.block_streaming.source import DiskWeightSource, PinnedBlock, PinnedWeightSource, WeightSource
+from ltx_core.block_streaming.utils import (
+    carve_buffer,
+    derive_layout,
+    layout_nbytes,
+    make_block_key,
+    resolve_attr,
+)
+from ltx_core.block_streaming.wrapper import BlockStreamingWrapper
+from ltx_core.loader.fuse_loras import FuseRule, bf16_fuse_rule, fuse_lora_weights
+from ltx_core.loader.helpers import create_meta_model, load_state_dict, read_model_config
+from ltx_core.loader.module_ops import ModuleOps
+from ltx_core.loader.primitives import (
+    LoraPathStrengthAndSDOps,
+    LoraStateDictWithStrength,
+    ModelBuilderProtocol,
+    StateDict,
+    StateDictLoader,
+    TensorLayout,
+)
+from ltx_core.loader.registry import DummyRegistry, Registry
+from ltx_core.loader.sd_ops import SDOps
+from ltx_core.loader.sft_loader import SafetensorsModelStateDictLoader
+from ltx_core.model.model_protocol import ModelConfigurator, ModelType
+
+if TYPE_CHECKING:
+    from typing_extensions import Self
+
+logger = logging.getLogger(__name__)
+
+DISK_CPU_SLOTS = 2
+_DEFAULT_GPU_SLOTS = 2
+_PREFETCH_DEPTH = 2
+
+
+class StreamingModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType]):
+    """Immutable builder for :class:`BlockStreamingWrapper`.
+    Reads block weights from safetensors on demand.  ``cpu_slots`` and
+    ``gpu_slots`` control the memory/speed trade-off (see :meth:`build`).
+    The builder is immutable (``with_*`` return modified copies) and exposes
+    its state via read-only properties backed by private attributes.
+    Args:
+        model_class_configurator: Creates the model from a config dict.
+        model_path: One or more ``.safetensors`` checkpoint paths.
+        model_sd_ops: Key remapping applied to safetensors keys.
+        module_ops: Module-level mutations for the meta model.
+        loras: LoRA adapters fused into weights at load time.
+        model_loader: Strategy for reading checkpoint metadata.
+        registry: Shared cache for loaded state dicts.
+        fuse_rule: Per-policy LoRA merge rule. Defaults to ``bf16_fuse_rule``;
+            use ``fp8_cast_fuse_rule`` for fp8_cast streaming so the pinned
+            buffers receive correctly-quantized weights.
+        blocks_attr: Dotted path to the ``nn.ModuleList`` (e.g.
+            ``"transformer_blocks"``).
+        blocks_prefix: State-dict key prefix for block weights
+            (e.g. ``"transformer_blocks"``).
+    """
+
+    def __init__(
+        self,
+        model_class_configurator: type[ModelConfigurator[ModelType]],
+        model_path: str | tuple[str, ...],
+        model_sd_ops: SDOps | None = None,
+        module_ops: tuple[ModuleOps, ...] = (),
+        loras: tuple[LoraPathStrengthAndSDOps, ...] = (),
+        model_loader: StateDictLoader | None = None,
+        registry: Registry | None = None,
+        fuse_rule: FuseRule = bf16_fuse_rule,
+        blocks_attr: str = "",
+        blocks_prefix: str = "",
+    ) -> None:
+        # Read-only: typed with the covariant ModelType, so it must not be a mutable attribute.
+        self._model_class_configurator: Final = model_class_configurator
+        self._model_path = model_path
+        self._model_sd_ops = model_sd_ops
+        self._module_ops = module_ops
+        self._loras = loras
+        self._model_loader = model_loader if model_loader is not None else SafetensorsModelStateDictLoader()
+        self._registry = registry if registry is not None else DummyRegistry()
+        self._fuse_rule = fuse_rule
+        self._blocks_attr = blocks_attr
+        self._blocks_prefix = blocks_prefix
+
+    @property
+    def model_class_configurator(self) -> type[ModelConfigurator[ModelType]]:
+        return self._model_class_configurator
+
+    @property
+    def model_path(self) -> str | tuple[str, ...]:
+        return self._model_path
+
+    @property
+    def model_sd_ops(self) -> SDOps | None:
+        return self._model_sd_ops
+
+    @property
+    def module_ops(self) -> tuple[ModuleOps, ...]:
+        return self._module_ops
+
+    @property
+    def loras(self) -> tuple[LoraPathStrengthAndSDOps, ...]:
+        return self._loras
+
+    @property
+    def model_loader(self) -> StateDictLoader:
+        return self._model_loader
+
+    @property
+    def registry(self) -> Registry:
+        return self._registry
+
+    @property
+    def fuse_rule(self) -> FuseRule:
+        return self._fuse_rule
+
+    @property
+    def blocks_attr(self) -> str:
+        return self._blocks_attr
+
+    @property
+    def blocks_prefix(self) -> str:
+        return self._blocks_prefix
+
+    def with_sd_ops(self, sd_ops: SDOps | None) -> Self:
+        clone = copy.copy(self)
+        clone._model_sd_ops = sd_ops
+        return clone
+
+    def with_module_ops(self, module_ops: tuple[ModuleOps, ...]) -> Self:
+        clone = copy.copy(self)
+        clone._module_ops = module_ops
+        return clone
+
+    def with_loras(self, loras: tuple[LoraPathStrengthAndSDOps, ...]) -> Self:
+        clone = copy.copy(self)
+        clone._loras = loras
+        return clone
+
+    def with_registry(self, registry: Registry) -> Self:
+        clone = copy.copy(self)
+        clone._registry = registry
+        return clone
+
+    def with_lora_load_device(self, device: torch.device) -> Self:
+        # Streaming fuses LoRAs into pinned CPU buffers; no other staging device is meaningful.
+        raise NotImplementedError("StreamingModelBuilder loads LoRA weights on CPU only.")
+
+    def with_fuse_rule(self, fuse_rule: FuseRule) -> Self:
+        clone = copy.copy(self)
+        clone._fuse_rule = fuse_rule
+        return clone
+
+    def model_config(self) -> dict:
+        """Read model configuration from the checkpoint metadata."""
+        return read_model_config(self.model_path, self.model_loader)
+
+    def meta_model(self, config: dict, module_ops: tuple[ModuleOps, ...]) -> ModelType:
+        """Create a model on the meta device and apply module operations."""
+        return create_meta_model(self.model_class_configurator, config, module_ops)
+
+    def build(
+        self,
+        device: torch.device | None = None,
+        dtype: torch.dtype | None = None,
+        cpu_slots_count: int | None = None,
+        gpu_slots_count: int | None = None,
+        **_kwargs: object,
+    ) -> BlockStreamingWrapper:
+        """Build and return a ready-to-use :class:`BlockStreamingWrapper`.
+        Args:
+            device: GPU device for compute. ``None`` defaults to ``cuda``.
+            dtype: Weight dtype (e.g. ``torch.bfloat16``). Required.
+            cpu_slots_count: Number of pinned CPU buffer slots.
+                ``None`` = RAM streaming (all blocks pre-loaded with LoRA fusion).
+            gpu_slots_count: Number of GPU buffer slots.
+                ``None`` = ``_DEFAULT_GPU_SLOTS`` (2).
+        """
+        if not self.blocks_prefix:
+            raise ValueError("blocks_prefix must be non-empty for streaming")
+        if dtype is None:
+            raise ValueError("StreamingModelBuilder.build requires an explicit dtype")
+        device = device if device is not None else torch.device("cuda")
+
+        config = read_model_config(self.model_path, self.model_loader)
+        meta_model: nn.Module = create_meta_model(self.model_class_configurator, config, self.module_ops)
+        meta_model.eval()
+
+        blocks = resolve_attr(meta_model, self.blocks_attr)
+
+        checkpoint_paths = list(self.model_path) if isinstance(self.model_path, tuple) else [self.model_path]
+        block_key_map, non_block_keys = _scan_checkpoint_keys(checkpoint_paths, self.model_sd_ops, self.blocks_prefix)
+        expected_indices = set(range(len(blocks)))
+        if set(block_key_map) != expected_indices:
+            missing = sorted(expected_indices - set(block_key_map))
+            extra = sorted(set(block_key_map) - expected_indices)
+            raise ValueError(
+                f"Block weights under prefix '{self.blocks_prefix}.' do not match the {len(blocks)} model blocks: "
+                f"missing indices {missing}, unexpected indices {extra}"
+            )
+
+        cpu_slots_count = cpu_slots_count if cpu_slots_count is not None else len(blocks)
+        gpu_slots_count = gpu_slots_count if gpu_slots_count is not None else _DEFAULT_GPU_SLOTS
+
+        if cpu_slots_count >= len(blocks):
+            lora_sd_and_strengths = self._load_lora_sds()
+            source, lora_sources = self._build_pinned_source(
+                blocks, dtype, cpu_slots_count, block_key_map, lora_sd_and_strengths
+            )
+            non_block_loras = lora_sd_and_strengths
+        else:
+            reader = DiskTensorReader(checkpoint_paths)
+            source, lora_sources = self._build_disk_source(
+                blocks, dtype, cpu_slots_count, reader, block_key_map, prefetch_depth=_PREFETCH_DEPTH
+            )
+            non_block_loras = [src.as_state_dict_with_strength() for src in lora_sources]
+
+        self._load_non_block_weights(meta_model, non_block_keys, device, dtype, non_block_loras)
+
+        copy_stream = torch.cuda.Stream(device=device)
+        gpu_pool = BufferPool(
+            source.slot_nbytes,
+            gpu_slots_count,
+            device,
+            reuse_barrier=lambda event: copy_stream.wait_event(event),
+        )
+        provider = WeightsProvider(
+            gpu_pool,
+            copy_stream,
+            device,
+            source,
+            lora_sources,
+            self.blocks_prefix,
+            fuse_rule=self.fuse_rule,
+        )
+        return BlockStreamingWrapper(
+            model=meta_model,
+            blocks=blocks,
+            provider=provider,
+            target_device=device,
+        )
+
+    def _load_lora_sds(self) -> list[LoraStateDictWithStrength]:
+        """Load each configured LoRA into a state dict for fusion (pinned path)."""
+        return [
+            LoraStateDictWithStrength(
+                load_state_dict([lora.path], self.model_loader, self.registry, torch.device("cpu"), lora.sd_ops),
+                lora.strength,
+            )
+            for lora in self.loras
+        ]
+
+    def _filtered_sd_ops(self, name_suffix: str, allowed_model_keys: frozenset[str]) -> SDOps:
+        """``model_sd_ops`` restricted to *allowed_model_keys* (post-rename keys).
+        The loader skips keys filtered to None before reading them, so a restricted
+        load never materializes the excluded partition. The distinct ``name`` avoids
+        a registry cache-id collision with the other partition.
+        """
+        base = self.model_sd_ops if self.model_sd_ops is not None else SDOps("streaming").with_matching()
+        allowed = allowed_model_keys if base.allowed_keys is None else (allowed_model_keys & base.allowed_keys)
+        return replace(base, name=f"{base.name}__{name_suffix}", allowed_keys=allowed)
+
+    def _build_pinned_source(
+        self,
+        blocks: nn.ModuleList,
+        dtype: torch.dtype,
+        cpu_slots_count: int,
+        block_key_map: dict[int, list[tuple[str, str]]],
+        lora_sd_and_strengths: list[LoraStateDictWithStrength],
+    ) -> tuple[WeightSource, list[LoraSource]]:
+        """Pre-load each block into its own contiguous pinned CPU buffer with LoRA fusion."""
+        for block_idx in block_key_map:
+            if block_idx >= cpu_slots_count:
+                raise ValueError(
+                    f"Pinned source requires one CPU slot per block; "
+                    f"got block index {block_idx} with only {cpu_slots_count} slots."
+                )
+
+        # One contiguous pinned buffer per block, carved into per-param views. The
+        # views (flattened by full key) are filled in place; the source then keeps
+        # only the contiguous buffer and the layout to re-carve it on read.
+        pinned_buffers: dict[int, torch.Tensor] = {}
+        block_layouts: dict[int, TensorLayout] = {}
+        fill_views: dict[str, torch.Tensor] = {}
+        for block_idx, entries in block_key_map.items():
+            block_state = _block_state(blocks[block_idx])
+            layout = derive_layout({param_name: block_state[param_name] for _sft_key, param_name in entries}, dtype)
+            buffer = bs_utils.alloc_buffer(layout_nbytes(layout), torch.device("cpu"), pin_memory=True)
+            views = carve_buffer(buffer, layout)
+            pinned_buffers[block_idx] = buffer
+            block_layouts[block_idx] = layout
+            for param_name, view in views.items():
+                fill_views[make_block_key(self.blocks_prefix, block_idx, param_name)] = view
+
+        block_sd = load_state_dict(
+            self.model_path,
+            self.model_loader,
+            self.registry,
+            torch.device("cpu"),
+            self._filtered_sd_ops("blocks", frozenset(fill_views)),
+        )
+
+        should_sync = False
+        for key, fused in fuse_lora_weights(
+            block_sd, lora_sd_and_strengths, fuse_rule=self.fuse_rule, preserve_input_device=False
+        ):
+            if key not in fill_views:
+                raise ValueError(f"Block-restricted load produced {key!r}, which is not a pinned block weight")
+            fill_views[key].copy_(fused, non_blocking=True)
+            block_sd.sd[key] = None
+            should_sync = True
+        if should_sync:
+            torch.cuda.synchronize()
+
+        # Fill remaining pinned keys from the source state dict.
+        for key, view in fill_views.items():
+            if block_sd.sd[key] is None:
+                continue
+            view.copy_(block_sd.sd[key])
+            block_sd.sd[key] = None
+
+        pinned = {idx: PinnedBlock(pinned_buffers[idx], block_layouts[idx]) for idx in pinned_buffers}
+        return PinnedWeightSource(pinned), []
+
+    def _build_disk_source(
+        self,
+        blocks: nn.ModuleList,
+        dtype: torch.dtype,
+        cpu_slots_count: int,
+        reader: DiskTensorReader,
+        block_key_map: dict[int, list[tuple[str, str]]],
+        prefetch_depth: int,
+    ) -> tuple[WeightSource, list[LoraSource]]:
+        """Create a DiskWeightSource backed by a DiskBlockReader.
+        Pool slots are sized to the largest block and carved per block on read, so
+        heterogeneous blocks (e.g. layers with differing attention layouts) share
+        one pool. Pool capacity is ``cpu_slots_count + prefetch_depth`` so the
+        lookahead loop in ``DiskWeightSource.get`` never evicts its own target.
+        Layouts come from the meta model; assumes module_ops keep the meta param
+        dtype in sync with the post-sd_ops checkpoint dtype.
+        """
+        block_layouts = _block_layouts(blocks, block_key_map, dtype)
+        slot_nbytes = max(layout_nbytes(layout) for layout in block_layouts.values())
+
+        cpu_pool = BufferPool(
+            slot_nbytes,
+            cpu_slots_count + prefetch_depth,
+            torch.device("cpu"),
+            reuse_barrier=lambda event: event.synchronize(),
+            pin_memory=True,
+        )
+        block_reader = DiskBlockReader(
+            reader=reader,
+            block_key_map=block_key_map,
+            sd_ops=self.model_sd_ops,
+            blocks_prefix=self.blocks_prefix,
+        )
+        fetcher = BlockFetcher(block_reader)
+        source = DiskWeightSource(
+            cpu_pool,
+            fetcher,
+            block_layouts,
+            blocks_number=len(blocks),
+            prefetch_depth=prefetch_depth,
+        )
+        lora_sources = [LoraSource(lora.path, lora.sd_ops, lora.strength) for lora in self.loras]
+
+        return source, lora_sources
+
+    @torch.inference_mode()
+    def _load_non_block_weights(
+        self,
+        model: nn.Module,
+        non_block_keys: list[tuple[str, str]],
+        device: torch.device,
+        dtype: torch.dtype,
+        lora_sd_and_strengths: list[LoraStateDictWithStrength],
+    ) -> None:
+        """Load the non-block weights onto *device* and fuse LoRAs -- both paths.
+        Reads through the loader with ``model_sd_ops`` restricted to the
+        non-block keys, so ``sd_ops`` (incl. kv-ops such as Gemma's ``lm_head``
+        duplication) is applied exactly once and block tensors are never read.
+        """
+        non_block_sd_ops = self._filtered_sd_ops("non_block", frozenset(mk for _sft_key, mk in non_block_keys))
+        loaded = load_state_dict(self.model_path, self.model_loader, self.registry, device, non_block_sd_ops)
+        non_block_sd = {key: tensor.to(dtype=dtype) for key, tensor in loaded.sd.items()}
+
+        if lora_sd_and_strengths:
+            non_block_state = StateDict(sd=non_block_sd, device=device, size=0, dtype={dtype})
+            for key, fused in fuse_lora_weights(
+                non_block_state,
+                lora_sd_and_strengths,
+                fuse_rule=self.fuse_rule,
+                preserve_input_device=True,
+            ):
+                non_block_sd[key] = fused
+
+        model.load_state_dict(non_block_sd, strict=False, assign=True)
+
+
+def _block_state(block: nn.Module) -> dict[str, torch.Tensor]:
+    """Streamed-eligible tensors of a block: parameters then buffers.
+    Block streaming swaps both params and checkpoint-backed buffers (e.g. Gemma4's
+    per-layer ``layer_scalar``), so the layout, pinned packing, and meta-ordering
+    all consult parameters and buffers together. Non-checkpoint (computed) buffers
+    are harmless here -- only keys present in ``block_key_map`` are ever streamed.
+    """
+    return {**dict(block.named_parameters()), **dict(block.named_buffers())}
+
+
+def _block_layouts(
+    blocks: nn.ModuleList,
+    block_key_map: dict[int, list[tuple[str, str]]],
+    dtype: torch.dtype,
+) -> dict[int, TensorLayout]:
+    """Per-block layout of the streamed tensors, taken from the meta model.
+    Blocks may differ in shape and even in which tensors they have (e.g. Gemma4's
+    full-attention layers drop ``v_proj``), so each block gets its own layout in
+    ``block_key_map`` order. The pinned packing, the disk reader, and the GPU carve
+    all key off this same per-block layout, so the provider's contiguous H2D copy
+    is valid for any entry order (no cross-block ordering required).
+    """
+    layouts: dict[int, TensorLayout] = {}
+    for idx, entries in block_key_map.items():
+        state = _block_state(blocks[idx])
+        layouts[idx] = derive_layout({param_name: state[param_name] for _sft_key, param_name in entries}, dtype)
+    return layouts
+
+
+def _scan_checkpoint_keys(
+    checkpoint_paths: list[str],
+    sd_ops: SDOps | None,
+    blocks_prefix: str,
+) -> tuple[dict[int, list[tuple[str, str]]], list[tuple[str, str]]]:
+    """Partition checkpoint keys into per-block and non-block lists.
+    Opens the safetensors files for header-only key enumeration; no tensor data
+    is read.
+    """
+    block_key_map: dict[int, list[tuple[str, str]]] = {}
+    non_block_keys: list[tuple[str, str]] = []
+    prefix_dot = blocks_prefix + "."
+    for path in checkpoint_paths:
+        with safetensors.safe_open(path, framework="pt", device="cpu") as handle:
+            for sft_key in handle.keys():  # noqa: SIM118
+                model_key = sd_ops.apply_to_key(sft_key) if sd_ops else sft_key
+                if model_key is None:
+                    continue
+                if model_key.startswith(prefix_dot):
+                    rest = model_key[len(prefix_dot) :]
+                    idx_str, _, param_name = rest.partition(".")
+                    try:
+                        block_idx = int(idx_str)
+                    except ValueError:
+                        non_block_keys.append((sft_key, model_key))
+                        continue
+                    block_key_map.setdefault(block_idx, []).append((sft_key, param_name))
+                else:
+                    non_block_keys.append((sft_key, model_key))
+    return block_key_map, non_block_keys

packages/ltx-core/src/ltx_core/block_streaming/disk.py

@@ -0,0 +1,165 @@
+"""Safetensors I/O and LoRA fusion for block streaming."""
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+
+import safetensors
+import torch
+
+from ltx_core.block_streaming.utils import allocate_layout_views, make_block_key
+from ltx_core.loader.fuse_loras import LoraProduct
+from ltx_core.loader.primitives import LoraStateDictWithStrength, StateDict
+from ltx_core.loader.sd_ops import SDOps
+
+_SAFETENSORS_DTYPE_TO_TORCH: dict[str, torch.dtype] = {
+    "F64": torch.float64,
+    "F32": torch.float32,
+    "F16": torch.float16,
+    "BF16": torch.bfloat16,
+}
+
+
+class DiskTensorReader:
+    """Key-based tensor accessor over one or more safetensors files."""
+
+    def __init__(self, paths: list[str]) -> None:
+        self._handles: list[safetensors.safe_open] = []
+        self._key_to_handle_idx: dict[str, int] = {}
+        for path in paths:
+            handle = safetensors.safe_open(path, framework="pt", device="cpu")
+            handle_idx = len(self._handles)
+            self._handles.append(handle)
+            for sft_key in handle.keys():  # noqa: SIM118
+                self._key_to_handle_idx[sft_key] = handle_idx
+
+    def get_tensor(self, key: str) -> torch.Tensor:
+        return self._handles[self._key_to_handle_idx[key]].get_tensor(key)
+
+    def close(self) -> None:
+        self._handles.clear()
+        self._key_to_handle_idx.clear()
+
+    def __contains__(self, key: str) -> bool:
+        return key in self._key_to_handle_idx
+
+    def __iter__(self) -> Iterator[str]:
+        return iter(self._key_to_handle_idx)
+
+
+class DiskBlockReader:
+    """Reads one block at a time from safetensors into provided buffers."""
+
+    def __init__(
+        self,
+        reader: DiskTensorReader,
+        block_key_map: dict[int, list[tuple[str, str]]],
+        sd_ops: SDOps | None = None,
+        blocks_prefix: str = "",
+    ) -> None:
+        self._reader = reader
+        self._block_key_map = block_key_map
+        self._sd_ops = sd_ops
+        self._blocks_prefix = blocks_prefix
+
+    def read_into(self, target: dict[str, torch.Tensor], block_idx: int) -> None:
+        block_prefix = make_block_key(self._blocks_prefix, block_idx, "")
+        for sft_key, param_name in self._block_key_map[block_idx]:
+            tensor = self._reader.get_tensor(sft_key)
+            if self._sd_ops is None:
+                target[param_name].copy_(tensor)
+                continue
+            full_key = make_block_key(self._blocks_prefix, block_idx, param_name)
+            for result in self._sd_ops.apply_to_key_value(full_key, tensor):
+                if not result.new_key.startswith(block_prefix):
+                    raise ValueError(
+                        f"SDOps output key '{result.new_key}' is outside block {block_idx} "
+                        f"(expected prefix '{block_prefix}'); cannot route to a per-block buffer."
+                    )
+                target[result.new_key[len(block_prefix) :]].copy_(result.new_value)
+
+    def cleanup(self) -> None:
+        self._reader.close()
+
+
+class LoraSource:
+    """Pinned-memory cache of matched LoRA A/B factors backed by a single buffer."""
+
+    def __init__(self, path: str, sd_ops: SDOps | None, strength: float) -> None:
+        self.strength = strength
+        self._pinned_ab: dict[str, tuple[torch.Tensor, torch.Tensor]] = {}
+
+        a_keys: dict[str, str] = {}
+        b_keys: dict[str, str] = {}
+        with safetensors.safe_open(path, framework="pt", device="cpu") as handle:
+            for sft_key in handle.keys():  # noqa: SIM118
+                model_key = sd_ops.apply_to_key(sft_key) if sd_ops is not None else sft_key
+                if model_key is None:
+                    continue
+                if model_key.endswith(".lora_A.weight"):
+                    a_keys[model_key[: -len(".lora_A.weight")]] = sft_key
+                elif model_key.endswith(".lora_B.weight"):
+                    b_keys[model_key[: -len(".lora_B.weight")]] = sft_key
+
+            matched_prefixes = list(a_keys.keys() & b_keys.keys())
+
+            # Build the layout from safetensors header metadata only — no tensor data is read.
+            layout: dict[str, tuple[torch.Size, torch.dtype]] = {}
+            for prefix in matched_prefixes:
+                a_slice_view = handle.get_slice(a_keys[prefix])
+                b_slice_view = handle.get_slice(b_keys[prefix])
+                layout[f"{prefix}.A"] = (
+                    torch.Size(a_slice_view.get_shape()),
+                    _SAFETENSORS_DTYPE_TO_TORCH[a_slice_view.get_dtype()],
+                )
+                layout[f"{prefix}.B"] = (
+                    torch.Size(b_slice_view.get_shape()),
+                    _SAFETENSORS_DTYPE_TO_TORCH[b_slice_view.get_dtype()],
+                )
+
+            all_views = allocate_layout_views(layout, pin_memory=True)
+
+            for prefix in matched_prefixes:
+                a_view = all_views[f"{prefix}.A"]
+                b_view = all_views[f"{prefix}.B"]
+                a_view.copy_(handle.get_tensor(a_keys[prefix]))
+                b_view.copy_(handle.get_tensor(b_keys[prefix]))
+                self._pinned_ab[prefix] = (a_view, b_view)
+
+    def as_state_dict_with_strength(self) -> LoraStateDictWithStrength:
+        """Return a :class:`LoraStateDictWithStrength` view of the pinned A/B factors.
+        Lets non-block fusion consume the already-loaded disk-streaming LoRA
+        without re-reading the safetensors file or re-applying ``sd_ops``.
+        """
+        sd: dict[str, torch.Tensor] = {}
+        for prefix, (a, b) in self._pinned_ab.items():
+            sd[f"{prefix}.lora_A.weight"] = a
+            sd[f"{prefix}.lora_B.weight"] = b
+        size = sum(t.numel() * t.element_size() for t in sd.values())
+        dtypes = {t.dtype for t in sd.values()}
+        return LoraStateDictWithStrength(
+            StateDict(sd=sd, device=torch.device("cpu"), size=size, dtype=dtypes),
+            self.strength,
+        )
+
+    def get_ab(
+        self,
+        param_prefix: str,
+        device: torch.device | None = None,
+        dtype: torch.dtype | None = None,
+    ) -> LoraProduct | None:
+        """Return the :class:`LoraProduct` for *param_prefix*, or ``None``."""
+        pair = self._pinned_ab.get(param_prefix)
+        if pair is None:
+            return None
+        a, b = pair
+        if device is not None and device.type == "cuda":
+            a = a.to(device=device, non_blocking=True)
+            b = b.to(device=device, non_blocking=True)
+        if dtype is not None:
+            a = a.to(dtype=dtype)
+            b = b.to(dtype=dtype)
+        return LoraProduct(a, b, self.strength)
+
+    def cleanup(self) -> None:
+        self._pinned_ab.clear()

packages/ltx-core/src/ltx_core/block_streaming/pool.py

@@ -0,0 +1,70 @@
+"""Raw buffer pool for block streaming."""
+
+from __future__ import annotations
+
+from collections import deque
+from typing import Callable
+
+import torch
+
+from ltx_core.block_streaming import utils
+
+
+class BufferPool:
+    """Fixed pool of pre-allocated raw buffer slots with event-based reuse.
+    Slots are carved from a single contiguous ``uint8`` buffer; each is
+    ``slot_nbytes`` long and handed out as a raw 1-D ``uint8`` tensor.
+    Args:
+        slot_nbytes: Byte size of each slot.
+        capacity: Number of slots to pre-allocate.
+        device: Device for allocation.
+        reuse_barrier: Called with the pending event before a slot is reused.
+        pin_memory: Pin buffers (for async H2D copies from CPU).
+    """
+
+    def __init__(
+        self,
+        slot_nbytes: int,
+        capacity: int,
+        device: torch.device,
+        reuse_barrier: Callable[[torch.cuda.Event], None],
+        pin_memory: bool = False,
+    ) -> None:
+        self._slot_nbytes = slot_nbytes
+        self._capacity = capacity
+        self._free: deque[torch.Tensor] = deque()
+        self._events: dict[int, torch.cuda.Event] = {}
+        self._reuse_barrier = reuse_barrier
+        buffer = utils.alloc_buffer(max(slot_nbytes * capacity, 1), device, pin_memory)
+        for slot in range(capacity):
+            self._free.append(buffer[slot * slot_nbytes : (slot + 1) * slot_nbytes])
+
+    @property
+    def capacity(self) -> int:
+        return self._capacity
+
+    @property
+    def slot_nbytes(self) -> int:
+        return self._slot_nbytes
+
+    def acquire(self) -> torch.Tensor:
+        """Take a free raw slot, waiting any pending event before returning.
+        Raises :class:`RuntimeError` if every slot is currently in use.
+        """
+        if not self._free:
+            raise RuntimeError(f"BufferPool exhausted: all {self._capacity} buffers are in use")
+        buffer = self._free.popleft()
+        event = self._events.pop(id(buffer), None)
+        if event is not None:
+            self._reuse_barrier(event)
+        return buffer
+
+    def release(self, buffer: torch.Tensor, event: torch.cuda.Event | None = None) -> None:
+        """Return a raw slot to the free list.
+        The *buffer* must be the exact tensor object returned by :meth:`acquire`
+        (reuse is keyed on its identity). If *event* is given it is waited on the
+        next :meth:`acquire` of this slot, ensuring the prior operation finished.
+        """
+        if event is not None:
+            self._events[id(buffer)] = event
+        self._free.append(buffer)

packages/ltx-core/src/ltx_core/block_streaming/provider.py

@@ -0,0 +1,146 @@
+"""GPU weights provider for block streaming."""
+
+from __future__ import annotations
+
+from collections import OrderedDict
+from typing import NamedTuple
+
+import torch
+
+from ltx_core.block_streaming.disk import LoraSource
+from ltx_core.block_streaming.pool import BufferPool
+from ltx_core.block_streaming.source import WeightSource
+from ltx_core.block_streaming.utils import carve_buffer, layout_nbytes
+from ltx_core.loader.fuse_loras import FuseRule, aggregate_lora_products, bf16_fuse_rule
+from ltx_core.loader.primitives import StateDict
+
+_EMPTY_STATE_DICT = StateDict(sd={}, device=torch.device("cpu"), size=0, dtype=set())
+
+
+class CachedBlock(NamedTuple):
+    """A cached GPU block: the raw pool slot plus the carved per-key views.
+    The raw slot is what is returned to the pool on eviction; the views are
+    what callers consume.
+    """
+
+    raw: torch.Tensor
+    views: dict[str, torch.Tensor]
+
+
+class WeightsProvider:
+    """Provides GPU-ready block weights via H2D copy from a pinned CPU weight source.
+    Args:
+        pool: Pre-allocated GPU weight buffer pool.
+        copy_stream: Dedicated CUDA stream for async H2D copies.
+        target_device: GPU device for compute.
+        source: Pinned CPU weight source.
+        lora_sources: LoRA adapters fused on H2D copy.
+        blocks_prefix: State-dict prefix for LoRA key matching.
+        fuse_rule: Per-policy LoRA merge rule (must be streaming-compatible:
+            no companion-key emission). Defaults to ``bf16_fuse_rule``.
+    """
+
+    def __init__(
+        self,
+        pool: BufferPool,
+        copy_stream: torch.cuda.Stream,
+        target_device: torch.device,
+        source: WeightSource,
+        lora_sources: list[LoraSource] | None = None,
+        blocks_prefix: str = "",
+        fuse_rule: FuseRule = bf16_fuse_rule,
+    ) -> None:
+        self._copy_stream = copy_stream
+        self._pool = pool
+        self._cache: OrderedDict[int, CachedBlock] = OrderedDict()
+        self._events: dict[int, torch.cuda.Event] = {}
+        self._target_device = target_device
+        self._source = source
+        self._lora_sources = lora_sources or []
+        self._blocks_prefix = blocks_prefix
+        self._fuse_rule = fuse_rule
+
+    def get(self, idx: int) -> dict[str, torch.Tensor]:
+        """Return GPU weights for block *idx*. Does H2D copy on miss."""
+        if idx in self._cache:
+            return self._cache[idx].views
+
+        # Evict oldest GPU buffer if at capacity.
+        if len(self._cache) >= self._pool.capacity:
+            evicted_idx, evicted = self._cache.popitem(last=False)
+            self._pool.release(evicted.raw, event=self._events.pop(evicted_idx, None))
+
+        layout = self._source.block_layout(idx)
+        raw = self._pool.acquire()
+        gpu_weights = carve_buffer(raw, layout)
+        cpu_buffer = self._source.get(idx)
+
+        h2d_event = self._copy_to_gpu(idx, raw, gpu_weights, cpu_buffer, layout_nbytes(layout))
+        self._source.release(idx, event=h2d_event)
+
+        self._cache[idx] = CachedBlock(raw, gpu_weights)
+        return gpu_weights
+
+    def _copy_to_gpu(
+        self,
+        idx: int,
+        raw: torch.Tensor,
+        gpu_weights: dict[str, torch.Tensor],
+        cpu_buffer: torch.Tensor,
+        nbytes: int,
+    ) -> torch.cuda.Event:
+        """Enqueue H2D copy + LoRA fusion on the copy stream and wait on compute.
+        *cpu_buffer* is one contiguous source buffer carved by the same layout as
+        *raw*, so a single byte copy of its leading *nbytes* reproduces every view
+        in *gpu_weights*. The wait is intentionally inside this method so callers --
+        and instrumentation regions wrapping it -- observe the full transfer time.
+        """
+        if not cpu_buffer.is_contiguous() or cpu_buffer.dtype != torch.uint8 or cpu_buffer.numel() < nbytes:
+            raise ValueError(
+                f"source buffer for block {idx} must be a contiguous uint8 buffer of >= {nbytes} bytes, "
+                f"got {cpu_buffer.dim()}-D {cpu_buffer.dtype} with {cpu_buffer.numel()} elements"
+            )
+        with torch.cuda.stream(self._copy_stream):
+            raw[:nbytes].copy_(cpu_buffer[:nbytes], non_blocking=True)
+            if self._lora_sources:
+                self._fuse_block_loras(idx, gpu_weights)
+            h2d_event = torch.cuda.Event()
+            h2d_event.record(self._copy_stream)
+
+        torch.cuda.current_stream(self._target_device).wait_event(h2d_event)
+        return h2d_event
+
+    def release(self, idx: int, event: torch.cuda.Event) -> None:
+        """Attach a compute-done event -- waited before this buffer is recycled."""
+        self._events[idx] = event
+
+    def cleanup(self) -> None:
+        """Synchronize streams and release all resources."""
+        self._copy_stream.synchronize()
+        torch.cuda.current_stream(self._target_device).synchronize()
+        self._cache.clear()
+        self._events.clear()
+        self._source.cleanup()
+        for lora in self._lora_sources:
+            lora.cleanup()
+
+    def __len__(self) -> int:
+        return len(self._cache)
+
+    def _fuse_block_loras(self, idx: int, weights: dict[str, torch.Tensor]) -> None:
+        """Fuse LoRA deltas directly into GPU block weights via ``fuse_rule``."""
+        agg_dtype = self._fuse_rule.aggregation_dtype
+        for name, tensor in weights.items():
+            if not name.endswith(".weight"):
+                continue
+            prefix = f"{self._blocks_prefix}.{idx}.{name}".removesuffix(".weight")
+            products = (
+                ab
+                for ab in (s.get_ab(prefix, device=self._target_device, dtype=agg_dtype) for s in self._lora_sources)
+                if ab is not None
+            )
+            deltas = aggregate_lora_products(products, agg_dtype)
+            if deltas is None:
+                continue
+            fused = self._fuse_rule(name, tensor, deltas, _EMPTY_STATE_DICT)
+            tensor.copy_(fused[name])

packages/ltx-core/src/ltx_core/block_streaming/source.py

@@ -0,0 +1,174 @@
+"""Weight sources for block streaming: protocol and implementations."""
+
+from __future__ import annotations
+
+from typing import NamedTuple, Protocol
+
+import torch
+
+from ltx_core.block_streaming.block_fetcher import BlockFetcher, FetchHandle
+from ltx_core.block_streaming.pool import BufferPool
+from ltx_core.block_streaming.utils import carve_buffer, layout_nbytes
+from ltx_core.loader.primitives import TensorLayout
+
+
+class WeightSource(Protocol):
+    """Provides pinned CPU weights for a given block index.
+    Blocks may be heterogeneous: each has its own layout, so the source exposes a
+    per-block layout and the byte size of the largest block (which sizes the
+    pool slots -- a smaller block is carved into the front of a max-sized slot).
+    The source is the single source of truth for each block's layout.
+    """
+
+    def block_layout(self, idx: int) -> TensorLayout:
+        """Per-block buffer layout (shape + dtype for each param)."""
+        ...
+
+    @property
+    def slot_nbytes(self) -> int:
+        """Byte size of the largest block; sizes a pool slot (16-byte aligned)."""
+        ...
+
+    def get(self, idx: int) -> torch.Tensor:
+        """Return one contiguous CPU buffer for block *idx*."""
+        ...
+
+    def release(self, idx: int, event: torch.cuda.Event | None) -> None:
+        """Signal that an async operation using these weights is guarded by *event*."""
+        ...
+
+    def cleanup(self) -> None:
+        """Release all resources (buffers, readers, events)."""
+        ...
+
+
+class _Scheduled(NamedTuple):
+    """A scheduled (possibly in-flight) read: the raw pool slot and its fetch handle."""
+
+    raw: torch.Tensor
+    status: FetchHandle
+
+
+class DiskWeightSource(WeightSource):
+    """WeightSource that streams blocks from disk via a :class:`BlockFetcher`.
+    ``get(idx)`` must be paired with a ``release(idx)`` before *idx* is fetched
+    again; getting a block that is already in flight raises. Each read acquires a
+    raw pool slot, carves it to block *idx*'s layout, and hands the carved views to
+    the fetcher to fill, so one max-sized slot can serve blocks of differing shapes.
+    ``get`` returns that same contiguous slot.
+    """
+
+    def __init__(
+        self,
+        pool: BufferPool,
+        fetcher: BlockFetcher,
+        block_layouts: dict[int, TensorLayout],
+        blocks_number: int,
+        prefetch_depth: int = 0,
+    ) -> None:
+        if blocks_number <= 0:
+            raise ValueError(f"blocks_number must be > 0, got {blocks_number}")
+        if prefetch_depth < 0:
+            raise ValueError(f"prefetch_depth must be >= 0, got {prefetch_depth}")
+        max_layout_nbytes = max((layout_nbytes(layout) for layout in block_layouts.values()), default=0)
+        if pool.slot_nbytes < max_layout_nbytes:
+            raise ValueError(
+                f"pool slot is too small for the largest block: slot {pool.slot_nbytes} bytes < {max_layout_nbytes}"
+            )
+
+        self._pool = pool
+        self._blocks_number = blocks_number
+        self._prefetch_depth = prefetch_depth
+        self._fetcher = fetcher
+        self._block_layouts = block_layouts
+        self._scheduled: dict[int, _Scheduled] = {}
+        self._in_flight: dict[int, torch.Tensor] = {}
+
+    def block_layout(self, idx: int) -> TensorLayout:
+        return self._block_layouts[idx]
+
+    @property
+    def slot_nbytes(self) -> int:
+        return self._pool.slot_nbytes
+
+    def get(self, idx: int) -> torch.Tensor:
+        if idx in self._in_flight:
+            raise RuntimeError(f"Block {idx} is already in flight; release it before getting it again")
+
+        scheduled = self._scheduled.pop(idx, None)
+        if scheduled is None:
+            scheduled = self._schedule(idx)
+        error = scheduled.status.wait()
+        if error is not None:
+            self._pool.release(scheduled.raw)
+            raise error
+
+        self._in_flight[idx] = scheduled.raw
+        for k in range(1, self._prefetch_depth + 1):
+            self._ensure_scheduled((idx + k) % self._blocks_number)
+        return scheduled.raw
+
+    def release(self, idx: int, event: torch.cuda.Event | None) -> None:
+        raw_buffer = self._in_flight.pop(idx)
+        self._pool.release(raw_buffer, event=event)
+
+    def cleanup(self) -> None:
+        self._fetcher.cleanup()
+        while self._in_flight:
+            _, raw_buffer = self._in_flight.popitem()
+            self._pool.release(raw_buffer)
+        while self._scheduled:
+            _, scheduled = self._scheduled.popitem()
+            scheduled.status.wait()
+            self._pool.release(scheduled.raw)
+
+    def _ensure_scheduled(self, idx: int) -> None:
+        """Schedule a read for *idx* if one is not already pending."""
+        if idx not in self._scheduled:
+            self._scheduled[idx] = self._schedule(idx)
+
+    def _schedule(self, idx: int) -> _Scheduled:
+        """Acquire a raw slot, carve it to block *idx*, enqueue a read, return the handle.
+        The raw slot and its fetch status are returned together so the caller can
+        track both as one unit; the fetcher only receives the carved views to fill.
+        """
+        raw_buffer = self._pool.acquire()
+        carved = carve_buffer(raw_buffer, self._block_layouts[idx])
+        status = self._fetcher.submit(idx, carved)
+        return _Scheduled(raw_buffer, status)
+
+
+class PinnedBlock(NamedTuple):
+    """A pre-loaded pinned block: its single contiguous buffer and the layout to carve it with."""
+
+    buffer: torch.Tensor
+    layout: TensorLayout
+
+
+class PinnedWeightSource(WeightSource):
+    """Pre-loaded pinned CPU weights, one contiguous (possibly heterogeneous) buffer per block."""
+
+    def __init__(self, blocks: dict[int, PinnedBlock]) -> None:
+        if not blocks:
+            raise ValueError("PinnedWeightSource requires at least one block")
+        self._blocks = blocks
+        self._slot_nbytes = max(layout_nbytes(block.layout) for block in blocks.values())
+
+    def block_layout(self, idx: int) -> TensorLayout:
+        return self._blocks[idx].layout
+
+    @property
+    def slot_nbytes(self) -> int:
+        return self._slot_nbytes
+
+    def get(self, idx: int) -> torch.Tensor:
+        return self._blocks[idx].buffer
+
+    def release(self, idx: int, event: torch.cuda.Event | None) -> None:
+        pass
+
+    def cleanup(self) -> None:
+        self._blocks.clear()
+
+    def __len__(self) -> int:
+        return len(self._blocks)

packages/ltx-core/src/ltx_core/block_streaming/utils.py

@@ -0,0 +1,168 @@
+"""Shared utilities for the block_streaming package."""
+
+from __future__ import annotations
+
+import math
+import weakref
+from dataclasses import dataclass
+from typing import Any, NamedTuple
+
+import torch
+from torch import nn
+
+from ltx_core.loader.primitives import TensorLayout
+
+FP8_DTYPES = frozenset({torch.float8_e4m3fn, torch.float8_e5m2})
+
+_BUFFER_ALIGN = 16
+
+
+def make_block_key(blocks_prefix: str, block_idx: int, param_name: str) -> str:
+    """Return the state-dict key for *param_name* under block *block_idx*."""
+    return f"{blocks_prefix}.{block_idx}.{param_name}"
+
+
+def resolve_attr(module: nn.Module, dotted_path: str) -> nn.ModuleList:
+    """Resolve a dotted attribute path like ``'model.language_model.layers'``."""
+    obj: Any = module
+    for part in dotted_path.split("."):
+        obj = getattr(obj, part)
+    if not isinstance(obj, nn.ModuleList):
+        raise TypeError(f"Expected nn.ModuleList at '{dotted_path}', got {type(obj).__name__}")
+    return obj
+
+
+def assign_tensor_to_module(root: nn.Module, dotted_name: str, tensor: torch.Tensor) -> None:
+    """Assign *tensor* to the parameter/buffer at *dotted_name* inside *root*.
+    Unlike ``param.data = tensor``, this works even when the existing parameter
+    lives on the ``meta`` device (which has an incompatible storage type).
+    """
+    parts = dotted_name.split(".")
+    parent = root
+    for part in parts[:-1]:
+        parent = getattr(parent, part)
+    leaf = parts[-1]
+    if leaf in parent._parameters:
+        parent._parameters[leaf] = nn.Parameter(tensor, requires_grad=False)
+    elif leaf in parent._buffers:
+        parent._buffers[leaf] = tensor
+    else:
+        raise AttributeError(f"{leaf} is not a parameter or buffer of {type(parent).__name__}")
+
+
+def derive_layout(tensors: dict[str, torch.Tensor], dtype: torch.dtype | None = None) -> TensorLayout:
+    """Derive a layout from a ``{name: tensor}`` dict.
+    If ``dtype`` is given, non-FP8 dtypes are coerced to it (FP8 preserved). If
+    ``None``, the source dtype is preserved as-is.
+    """
+    return {
+        name: (t.shape, t.dtype if dtype is None or t.dtype in FP8_DTYPES else dtype) for name, t in tensors.items()
+    }
+
+
+def _align_up(offset: int, alignment: int) -> int:
+    return (offset + alignment - 1) & ~(alignment - 1)
+
+
+def _alloc_pinned_exact(nbytes: int) -> torch.Tensor | None:
+    """Allocate exactly ``nbytes`` of pinned host memory via ``cudaHostRegister``.
+    Bypasses PyTorch's ``CachingHostAllocator``, which rounds every
+    ``pin_memory=True`` request up to ``PowerOf2Ceil(N)`` (see
+    ``aten/src/ATen/core/CachingHostAllocator.h``). Returns ``None`` if
+    registration fails. The unregister hook is bound to the storage (not the
+    tensor) so views of the buffer keep the registration alive until the
+    memory is actually freed. Caller is responsible for ensuring CUDA is
+    available.
+    """
+    cudart = torch.cuda.cudart()
+    buf = torch.empty(nbytes, dtype=torch.uint8)
+    ptr = buf.data_ptr()
+    err = int(cudart.cudaHostRegister(ptr, nbytes, 0))
+    if err != 0:
+        return None
+    weakref.finalize(buf.untyped_storage(), lambda p=ptr: cudart.cudaHostUnregister(p))
+    return buf
+
+
+def alloc_buffer(nbytes: int, device: torch.device | None, pin_memory: bool) -> torch.Tensor:
+    """Allocate one ``uint8`` buffer for :func:`allocate_layout_views`.
+    For pinned host buffers, prefer ``cudaHostRegister`` to dodge the caching
+    allocator's power-of-2 rounding. Falls back to the caching allocator if
+    registration fails. Raises if pinning is requested without a CUDA runtime,
+    since pinning is fundamentally a CUDA driver operation.
+    """
+    if pin_memory and (device is None or torch.device(device).type == "cpu"):
+        if not torch.cuda.is_available():
+            raise RuntimeError("pin_memory=True requires CUDA, which is not available")
+        buf = _alloc_pinned_exact(nbytes)
+        if buf is not None:
+            return buf
+    return torch.empty(nbytes, dtype=torch.uint8, device=device, pin_memory=pin_memory)
+
+
+@dataclass(frozen=True)
+class _TensorSlice:
+    """Location of a single tensor view within the buffer."""
+
+    offset: int
+    shape: torch.Size
+    dtype: torch.dtype
+
+    def size(self) -> int:
+        return math.prod(self.shape) * self.dtype.itemsize
+
+
+class LayoutSlices(NamedTuple):
+    """Per-key tensor slices of a layout plus the total aligned buffer size."""
+
+    slices: dict[str, _TensorSlice]
+    nbytes: int
+
+
+def _layout_slices(layout: TensorLayout) -> LayoutSlices:
+    """Compute the byte offset of each key in *layout* and the total aligned size.
+    The size is at least one byte so empty layouts still produce a valid buffer.
+    """
+    slices: dict[str, _TensorSlice] = {}
+    cursor = 0
+    for key, (shape, dtype) in layout.items():
+        cursor = _align_up(cursor, _BUFFER_ALIGN)
+        slices[key] = _TensorSlice(offset=cursor, shape=shape, dtype=dtype)
+        cursor += slices[key].size()
+    return LayoutSlices(slices, max(_align_up(cursor, _BUFFER_ALIGN), 1))
+
+
+def layout_nbytes(layout: TensorLayout) -> int:
+    """Byte size of one contiguous, 16-byte-aligned buffer holding *layout* (>= 1)."""
+    return _layout_slices(layout).nbytes
+
+
+def carve_buffer(buffer: torch.Tensor, layout: TensorLayout) -> dict[str, torch.Tensor]:
+    """Carve per-key tensor views for *layout* into the front of *buffer*.
+    *buffer* is a 1-D ``uint8`` tensor at least :func:`layout_nbytes` long. Each
+    returned tensor is a non-overlapping slice of its leading bytes reinterpreted
+    at the requested shape and dtype; any trailing bytes are left unused. That
+    slack is what lets one max-sized pool slot hold a smaller (heterogeneous)
+    block. The views keep *buffer*'s storage alive via PyTorch refcounting.
+    """
+    if buffer.dtype != torch.uint8 or buffer.dim() != 1:
+        raise ValueError(f"carve_buffer expects a 1-D uint8 buffer, got {buffer.dim()}-D {buffer.dtype}")
+    slices, nbytes = _layout_slices(layout)
+    if buffer.numel() < nbytes:
+        raise ValueError(f"buffer too small to carve layout: need {nbytes} bytes, got {buffer.numel()}")
+    return {key: buffer[s.offset : s.offset + s.size()].view(s.dtype).view(s.shape) for key, s in slices.items()}
+
+
+def allocate_layout_views(
+    layout: TensorLayout,
+    device: torch.device | None = None,
+    pin_memory: bool = False,
+) -> dict[str, torch.Tensor]:
+    """Allocate a single ``uint8`` buffer and return per-key tensor views into it.
+    All keys in *layout* live in one contiguous allocation; each returned
+    tensor is a non-overlapping slice of that buffer reinterpreted at the
+    requested shape and dtype. The views keep the underlying storage alive
+    via PyTorch refcounting — drop them all to release the memory.
+    """
+    buffer = alloc_buffer(layout_nbytes(layout), device, pin_memory)
+    return carve_buffer(buffer, layout)

packages/ltx-core/src/ltx_core/block_streaming/wrapper.py

@@ -0,0 +1,96 @@
+"""Block streaming wrapper: streams transformer blocks through a WeightsProvider."""
+
+from __future__ import annotations
+
+import itertools
+from typing import Any
+
+import torch
+from torch import nn
+
+from ltx_core.block_streaming.provider import WeightsProvider
+from ltx_core.block_streaming.utils import assign_tensor_to_module
+
+
+class BlockStreamingWrapper(nn.Module):
+    """Streams sequential model blocks through GPU buffer caches.
+    The wrapper delegates all weight management to a :class:`WeightsProvider`
+    which handles CPU-to-GPU copies, caching, LoRA fusion, and stream
+    synchronization internally.
+    Use :class:`StreamingModelBuilder` to construct this wrapper -- it
+    handles checkpoint parsing, source selection, and provider creation.
+    Args:
+        model: The wrapped model (non-block params already on GPU).
+        blocks: Sequential blocks to stream (``nn.ModuleList``).
+        provider: Provides GPU-ready weights on demand.
+        target_device: GPU device for compute.
+    """
+
+    def __init__(
+        self,
+        model: nn.Module,
+        blocks: nn.ModuleList,
+        provider: WeightsProvider,
+        target_device: torch.device,
+    ) -> None:
+        super().__init__()
+        self._model = model
+        self._blocks = blocks
+        self._target_device = target_device
+        self._provider = provider
+
+        self._hooks: list[torch.utils.hooks.RemovableHandle] = []
+        self._register_hooks()
+
+    # ------------------------------------------------------------------
+    # Hook registration
+    # ------------------------------------------------------------------
+
+    def _pre_hook(self, block_idx: int) -> None:
+        """Load GPU weights for a block and inject them into its parameters."""
+        gpu_weights = self._provider.get(block_idx)
+
+        block = self._blocks[block_idx]
+        for name, _param in itertools.chain(block.named_parameters(), block.named_buffers()):
+            assign_tensor_to_module(block, name, gpu_weights[name])
+
+    def _post_hook(self, block_idx: int) -> None:
+        """Record a compute-done event and release the block weights."""
+        compute_done = torch.cuda.Event()
+        compute_done.record(torch.cuda.current_stream(self._target_device))
+        self._provider.release(block_idx, event=compute_done)
+
+    def _register_hooks(self) -> None:
+        for idx, block in enumerate(self._blocks):
+            pre = block.register_forward_pre_hook(
+                lambda _mod, _args, *, idx=idx: self._pre_hook(idx),
+            )
+            post = block.register_forward_hook(
+                lambda _mod, _args, _out, *, idx=idx: self._post_hook(idx),
+            )
+            self._hooks.extend([pre, post])
+
+    # ------------------------------------------------------------------
+    # Teardown
+    # ------------------------------------------------------------------
+
+    def teardown(self) -> None:
+        """Remove hooks and release all resources."""
+        for h in self._hooks:
+            h.remove()
+        self._hooks.clear()
+        self._provider.cleanup()
+
+    # ------------------------------------------------------------------
+    # Forward and attribute delegation
+    # ------------------------------------------------------------------
+
+    def forward(self, *args: Any, **kwargs: Any) -> Any:  # noqa: ANN401
+        return self._model(*args, **kwargs)
+
+    def __getattr__(self, name: str) -> Any:  # noqa: ANN401
+        """Proxy attribute access to the wrapped model."""
+        try:
+            return super().__getattr__(name)
+        except AttributeError:
+            return getattr(self._model, name)

packages/ltx-core/src/ltx_core/components/diffusion_steps.py

@@ -4,6 +4,24 @@ from ltx_core.components.protocols import DiffusionStepProtocol
 from ltx_core.utils import to_velocity
 
 
+def _get_ancestral_step(
+    sigma_from: torch.Tensor,
+    sigma_to: torch.Tensor,
+    eta: float = 1.0,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Compute ``(sigma_down, sigma_up)`` for one DDIM ancestral sampling step.
+    Both inputs are in the rescaled parameterization ``sigma / alpha``.
+    Returns ``sigma_down`` (deterministic component) and ``sigma_up``
+    (stochastic component) in the same rescaled space.
+    """
+    if not eta:
+        return sigma_to, torch.zeros_like(sigma_to)
+    variance = sigma_to**2 * (sigma_from**2 - sigma_to**2).clamp(min=0) / sigma_from**2
+    sigma_up = (eta * variance**0.5).clamp(max=sigma_to)
+    sigma_down = (sigma_to**2 - sigma_up**2).clamp(min=0) ** 0.5
+    return sigma_down, sigma_up
+
+
 class EulerDiffusionStep(DiffusionStepProtocol):
     """
     First-order Euler method for diffusion sampling.
@@ -104,3 +122,65 @@ class Res2sDiffusionStep(DiffusionStepProtocol):
         # Mix deterministic and stochastic components
         x_noised = alpha_ratio * (denoised_next + sigma_down * eps_next) + sigma_up * noise
         return x_noised.to(output_dtype)
+
+
+class EulerCfgPpDiffusionStep(DiffusionStepProtocol):
+    """Euler step using the CFG++ correction for the ODE derivative.
+    Instead of the standard velocity formula, the ODE derivative is computed
+    from the unconditioned prediction, keeping the conditioned prediction as
+    the target denoised state.  Ancestral (DDIM) noise injection is applied
+    in the rescaled sigma parameterization (sigma / alpha).
+    All diffusion quantities (alpha, ODE derivative, ancestral coefficients)
+    are computed internally from ``sigmas`` and ``uncond_denoised``.
+    Reference: CFG++ (https://arxiv.org/abs/2406.08070).
+    """
+
+    def __init__(self, eta: float = 1.0, s_noise: float = 1.0) -> None:
+        self.eta = eta
+        self.s_noise = s_noise
+
+    def step(
+        self,
+        sample: torch.Tensor,
+        denoised_sample: torch.Tensor,
+        sigmas: torch.Tensor,
+        step_index: int,
+        uncond_denoised: torch.Tensor,
+        noise: torch.Tensor | None = None,
+        **_kwargs,
+    ) -> torch.Tensor:
+        """Advance one CFG++ Euler step.
+        Args:
+            sample: Current noisy latent x_t.
+            denoised_sample: Conditioned denoised prediction x_0^cond.
+            sigmas: Full sigma schedule tensor.
+            step_index: Current step index.
+            uncond_denoised: Unconditioned denoised prediction x_0^uncond,
+                used to compute the ODE derivative direction.
+            noise: Noise tensor for stochastic injection; ignored when
+                ``eta=0`` or ``s_noise=0``.
+        Returns:
+            Updated latent x_{t-1}.
+        """
+        sigma_s = sigmas[step_index].to(torch.float32)
+        sigma_t = sigmas[step_index + 1].to(torch.float32)
+        _eps = torch.finfo(torch.float32).eps
+        # Clamp to avoid division by zero when sigma == 1.0 exactly.
+        alpha_s = (1.0 - sigma_s).clamp(min=_eps)
+        alpha_t = (1.0 - sigma_t).clamp(min=_eps)
+
+        x = sample.to(torch.float32)
+        denoised = denoised_sample.to(torch.float32)
+        uncond = uncond_denoised.to(torch.float32)
+
+        # ODE derivative: direction toward noise using uncond prediction (CFG++ correction)
+        d = (x - alpha_s * uncond) / sigma_s
+
+        # Ancestral step in rescaled sigma space (sigma / alpha)
+        sigma_down, sigma_up = _get_ancestral_step(sigma_s / alpha_s, sigma_t / alpha_t, eta=self.eta)
+        sigma_down = alpha_t * sigma_down
+
+        x_next = alpha_t * denoised + sigma_down * d
+        if noise is not None and self.eta > 0 and self.s_noise > 0:
+            x_next = x_next + alpha_t * noise.to(torch.float32) * self.s_noise * sigma_up
+        return x_next.to(sample.dtype)

packages/ltx-core/src/ltx_core/components/guiders.py

@@ -253,6 +253,11 @@ class MultiModalGuider:
         and as scale * (cond - uncond) for stg, steering the denoising process away from the unconditioned
         prediction.
         """
+        dtype = cond.dtype
+        cond = cond.float()
+        uncond_text = uncond_text.float() if isinstance(uncond_text, torch.Tensor) else uncond_text
+        uncond_perturbed = uncond_perturbed.float() if isinstance(uncond_perturbed, torch.Tensor) else uncond_perturbed
+        uncond_modality = uncond_modality.float() if isinstance(uncond_modality, torch.Tensor) else uncond_modality
         pred = (
             cond
             + (self.params.cfg_scale - 1) * (cond - uncond_text)
@@ -265,7 +270,7 @@ class MultiModalGuider:
             factor = self.params.rescale_scale * factor + (1 - self.params.rescale_scale)
             pred = pred * factor
 
-        return pred
+        return pred.to(dtype)
 
     def do_unconditional_generation(self) -> bool:
         """Returns True if the guider is doing unconditional generation."""

packages/ltx-core/src/ltx_core/components/noisers.py

@@ -17,18 +17,20 @@ class GaussianNoiser(Noiser):
 
     def __init__(self, generator: torch.Generator):
         super().__init__()
-
         self.generator = generator
 
-    def __call__(self, latent_state: LatentState, noise_scale: float = 1.0) -> LatentState:
-        noise = torch.randn(
+    def _sample_noise(self, latent_state: LatentState) -> torch.Tensor:
+        return torch.randn(
             *latent_state.latent.shape,
             device=latent_state.latent.device,
             dtype=latent_state.latent.dtype,
             generator=self.generator,
         )
-        scaled_mask = latent_state.denoise_mask * noise_scale
-        latent = noise * scaled_mask + latent_state.latent * (1 - scaled_mask)
+
+    def __call__(self, latent_state: LatentState, noise_scale: float = 1.0) -> LatentState:
+        noise = self._sample_noise(latent_state)
+        latent = torch.lerp(latent_state.latent.float(), noise.float(), noise_scale)
+        latent = torch.lerp(latent_state.clean_latent.float(), latent, latent_state.denoise_mask)
         return replace(
             latent_state,
             latent=latent.to(latent_state.latent.dtype),

packages/ltx-core/src/ltx_core/components/patchifiers.py

@@ -151,17 +151,22 @@ def get_pixel_coords(
             that treat frame zero differently still yield non-negative timestamps.
     """
     # Broadcast the VAE scale factors so they align with the `(batch, axis, patch, bound)` layout.
+    # Axis 1 of `latent_coords` is ordered (frame/time, height, width) — match that explicitly by
+    # pulling fields from the NamedTuple rather than relying on tuple iteration order.
     broadcast_shape = [1] * latent_coords.ndim
     broadcast_shape[1] = -1  # axis dimension corresponds to (frame/time, height, width)
-    scale_tensor = torch.tensor(scale_factors, device=latent_coords.device).view(*broadcast_shape)
+    scale_tensor = torch.tensor(
+        [scale_factors.time, scale_factors.height, scale_factors.width],
+        device=latent_coords.device,
+    ).view(*broadcast_shape)
 
     # Apply per-axis scaling to convert latent bounds into pixel-space coordinates.
     pixel_coords = latent_coords * scale_tensor
 
     if causal_fix:
-        # VAE temporal stride for the very first frame is 1 instead of `scale_factors[0]`.
+        # VAE temporal stride for the very first frame is 1 instead of `scale_factors.time`.
         # Shift and clamp to keep the first-frame timestamps causal and non-negative.
-        pixel_coords[:, 0, ...] = (pixel_coords[:, 0, ...] + 1 - scale_factors[0]).clamp(min=0)
+        pixel_coords[:, 0, ...] = (pixel_coords[:, 0, ...] + 1 - scale_factors.time).clamp(min=0)
 
     return pixel_coords
 

packages/ltx-core/src/ltx_core/conditioning/__init__.py

@@ -3,17 +3,21 @@
 from ltx_core.conditioning.exceptions import ConditioningError
 from ltx_core.conditioning.item import ConditioningItem
 from ltx_core.conditioning.types import (
+    AudioConditionByReferenceLatent,
     ConditioningItemAttentionStrengthWrapper,
     VideoConditionByKeyframeIndex,
     VideoConditionByLatentIndex,
+    VideoConditionByMask,
     VideoConditionByReferenceLatent,
 )
 
 __all__ = [
+    "AudioConditionByReferenceLatent",
     "ConditioningError",
     "ConditioningItem",
     "ConditioningItemAttentionStrengthWrapper",
     "VideoConditionByKeyframeIndex",
     "VideoConditionByLatentIndex",
+    "VideoConditionByMask",
     "VideoConditionByReferenceLatent",
 ]

packages/ltx-core/src/ltx_core/conditioning/types/__init__.py

@@ -3,11 +3,15 @@
 from ltx_core.conditioning.types.attention_strength_wrapper import ConditioningItemAttentionStrengthWrapper
 from ltx_core.conditioning.types.keyframe_cond import VideoConditionByKeyframeIndex
 from ltx_core.conditioning.types.latent_cond import VideoConditionByLatentIndex
+from ltx_core.conditioning.types.mask_cond import VideoConditionByMask
+from ltx_core.conditioning.types.reference_audio_cond import AudioConditionByReferenceLatent
 from ltx_core.conditioning.types.reference_video_cond import VideoConditionByReferenceLatent
 
 __all__ = [
+    "AudioConditionByReferenceLatent",
     "ConditioningItemAttentionStrengthWrapper",
     "VideoConditionByKeyframeIndex",
     "VideoConditionByLatentIndex",
+    "VideoConditionByMask",
     "VideoConditionByReferenceLatent",
 ]

packages/ltx-core/src/ltx_core/conditioning/types/keyframe_cond.py

@@ -10,19 +10,28 @@ from ltx_core.types import LatentState, VideoLatentShape
 class VideoConditionByKeyframeIndex(ConditioningItem):
     """
     Conditions video generation on keyframe latents at a specific frame index.
-    Appends keyframe tokens to the latent state with positions offset by frame_idx,
-    and sets denoise strength according to the strength parameter.
+    Appends keyframe tokens to the sequence with positions offset by frame_idx: the keyframe
+    latents become clean-latent tokens (placeholder zeros in the noisy latent) and the denoise
+    mask is set from the strength parameter.
     To add attention masking, wrap with :class:`ConditioningItemAttentionStrengthWrapper`.
     Args:
         keyframes: Keyframe latents [B, C, F, H, W].
         frame_idx: Frame index offset for positional encoding.
         strength: Conditioning strength (1.0 = clean, 0.0 = fully denoised).
+        num_pixel_frames: Number of pixel frames the keyframe latent originally encodes.
     """
 
-    def __init__(self, keyframes: torch.Tensor, frame_idx: int, strength: float):
+    def __init__(
+        self,
+        keyframes: torch.Tensor,
+        frame_idx: int,
+        strength: float,
+        num_pixel_frames: int = 1,
+    ):
         self.keyframes = keyframes
         self.frame_idx = frame_idx
         self.strength = strength
+        self.num_pixel_frames = num_pixel_frames
 
     def apply_to(
         self,
@@ -41,6 +50,11 @@ class VideoConditionByKeyframeIndex(ConditioningItem):
         )
 
         positions[:, 0, ...] += self.frame_idx
+        # If the keyframe latent encodes a single pixel frame,
+        # narrow the temporal end to [start, start + 1) instead of the
+        # VAE-scaled range.
+        if self.num_pixel_frames == 1:
+            positions[:, 0, ..., 1:] = positions[:, 0, ..., :1] + 1
         positions = positions.to(dtype=torch.float32)
         positions[:, 0, ...] /= latent_tools.fps
 
@@ -62,7 +76,7 @@ class VideoConditionByKeyframeIndex(ConditioningItem):
         )
 
         return LatentState(
-            latent=torch.cat([latent_state.latent, tokens], dim=1),
+            latent=torch.cat([latent_state.latent, torch.zeros_like(tokens)], dim=1),
             denoise_mask=torch.cat([latent_state.denoise_mask, denoise_mask], dim=1),
             positions=torch.cat([latent_state.positions, positions], dim=2),
             clean_latent=torch.cat([latent_state.clean_latent, tokens], dim=1),

packages/ltx-core/src/ltx_core/conditioning/types/latent_cond.py

@@ -9,8 +9,8 @@ from ltx_core.types import LatentState
 class VideoConditionByLatentIndex(ConditioningItem):
     """
     Conditions video generation by injecting latents at a specific latent frame index.
-    Replaces tokens in the latent state at positions corresponding to latent_idx,
-    and sets denoise strength according to the strength parameter.
+    Sets the clean latents at positions corresponding to latent_idx to the injected latents,
+    sets denoise strength according to the strength parameter.
     """
 
     def __init__(self, latent: torch.Tensor, strength: float, latent_idx: int):
@@ -37,7 +37,6 @@ class VideoConditionByLatentIndex(ConditioningItem):
 
         latent_state = latent_state.clone()
 
-        latent_state.latent[:, start_token:stop_token] = tokens
         latent_state.clean_latent[:, start_token:stop_token] = tokens
         latent_state.denoise_mask[:, start_token:stop_token] = 1.0 - self.strength
 

packages/ltx-core/src/ltx_core/conditioning/types/mask_cond.py

@@ -0,0 +1,49 @@
+"""Mask-based conditioning for inpainting and spatial conditioning."""
+
+from dataclasses import replace
+
+import torch
+
+from ltx_core.conditioning.item import ConditioningItem
+from ltx_core.tools import LatentTools
+from ltx_core.types import LatentState
+
+
+class VideoConditionByMask(ConditioningItem):
+    """Condition video generation using a binary mask over latent frames.
+    Masked positions (mask=1) receive the provided clean latent values and are
+    excluded from denoising (denoise_mask set to ``1 - strength``). Unmasked
+    positions (mask=0) are left unchanged and denoised normally.
+    The mask operates in **unpatchified latent** space — it should have shape
+    ``[B, F, H, W]`` matching the latent dimensions (after VAE encoding,
+    before patchification). This is consistent with the latent input format
+    used by all other conditioning items.
+    Args:
+        latent: Clean conditioning latents in unpatchified format [B, C, F, H, W].
+            Must match the target shape of the latent tools.
+        mask: Binary mask [B, F, H, W] in unpatchified latent space.
+            1 = conditioning position (clean, excluded from denoising),
+            0 = generated position (noised, denoised normally).
+        strength: Conditioning strength for masked positions. 1.0 = fully clean
+            (no denoising), 0.0 = no conditioning effect. Default 1.0.
+    """
+
+    def __init__(self, latent: torch.Tensor, mask: torch.Tensor, strength: float = 1.0):
+        self.latent = latent
+        self.mask = mask
+        self.strength = strength
+
+    def apply_to(self, latent_state: LatentState, latent_tools: LatentTools) -> LatentState:
+        """Apply mask-based conditioning to the latent state."""
+        tokens = latent_tools.patchifier.patchify(self.latent)
+
+        mask = latent_tools.patchifier.patchify(self.mask.unsqueeze(1))
+
+        m = mask.to(dtype=latent_state.latent.dtype)
+        inv = 1 - m
+
+        return replace(
+            latent_state,
+            clean_latent=latent_state.clean_latent * inv + tokens * m,
+            denoise_mask=latent_state.denoise_mask * inv + (1.0 - self.strength) * m,
+        )

packages/ltx-core/src/ltx_core/conditioning/types/reference_audio_cond.py

@@ -0,0 +1,59 @@
+"""Audio reference conditioning items."""
+
+from __future__ import annotations
+
+import torch
+
+from ltx_core.conditioning.mask_utils import update_attention_mask
+from ltx_core.tools import LatentTools
+from ltx_core.types import LatentState
+
+
+class AudioConditionByReferenceLatent:
+    """Append patchified reference audio tokens after the target audio sequence.
+    Mirrors :class:`ltx_core.conditioning.types.reference_video_cond.VideoConditionByReferenceLatent`
+    but for audio. The reference tokens are appended so the target audio tokens stay
+    in the first ``num_noisy_tokens`` positions and can be kept by
+    :meth:`ltx_core.tools.LatentTools.clear_conditioning`.
+    Args:
+        patchified: Patchified reference latent ``[B, T_ref, C]``.
+        positions: RoPE positions for reference tokens, ``[B, 1, T_ref, 2]``.
+        strength: 1.0 keeps reference clean; 0.0 would fully denoise it.
+    """
+
+    def __init__(
+        self,
+        patchified: torch.Tensor,
+        positions: torch.Tensor,
+        strength: float = 1.0,
+    ) -> None:
+        self.patchified = patchified
+        self.positions = positions.to(dtype=torch.float32)
+        self.strength = strength
+
+    def apply_to(self, latent_state: LatentState, latent_tools: LatentTools) -> LatentState:
+        tokens = self.patchified
+        denoise_mask = torch.full(
+            size=(*tokens.shape[:2], 1),
+            fill_value=1.0 - self.strength,
+            device=tokens.device,
+            dtype=tokens.dtype,
+        )
+
+        new_attention_mask = update_attention_mask(
+            latent_state=latent_state,
+            attention_mask=None,
+            num_noisy_tokens=latent_tools.patchifier.get_token_count(latent_tools.target_shape),
+            num_new_tokens=tokens.shape[1],
+            batch_size=tokens.shape[0],
+            device=tokens.device,
+            dtype=tokens.dtype,
+        )
+
+        return LatentState(
+            latent=torch.cat([latent_state.latent, torch.zeros_like(tokens)], dim=1),
+            denoise_mask=torch.cat([latent_state.denoise_mask, denoise_mask], dim=1),
+            positions=torch.cat([latent_state.positions, self.positions], dim=2),
+            clean_latent=torch.cat([latent_state.clean_latent, tokens], dim=1),
+            attention_mask=new_attention_mask,
+        )

packages/ltx-core/src/ltx_core/conditioning/types/reference_video_cond.py

@@ -14,7 +14,8 @@ class VideoConditionByReferenceLatent(ConditioningItem):
     Conditions video generation on a reference video latent for IC-LoRA inference.
     IC-LoRAs are trained by concatenating reference (control signal) and target tokens,
     learning to attend across both. This class replicates that setup at inference by
-    appending reference tokens to the latent sequence.
+    appending the reference tokens to the sequence as clean latents (with placeholder zeros
+    in the noisy latent).
     IC-LoRAs can be trained with lower-resolution references than the target (e.g., 384px
     reference for 768px output) for efficiency and better generalization. The
     `downscale_factor` scales reference positions to match target coordinates, preserving
@@ -22,9 +23,9 @@ class VideoConditionByReferenceLatent(ConditioningItem):
     (stored in LoRA metadata).
     To add attention masking, wrap with :class:`ConditioningItemAttentionStrengthWrapper`.
     Args:
-        latent: Reference video latents [B, C, F, H, W]
-        downscale_factor: Target/reference resolution ratio (e.g., 2 = half-resolution
-            reference). Spatial positions are scaled by this factor.
+        latent: Reference video latents [B, C, F, H, W].
+        downscale_factor: Target/reference spatial ratio (e.g. 2 = half-res ref).
+        temporal_scale_factor: Target/reference temporal ratio S (e.g. 4 = ref at 1/4 fps).
         strength: Conditioning strength. 1.0 = full (reference kept clean),
             0.0 = none (reference denoised). Default 1.0.
     """
@@ -33,10 +34,12 @@ class VideoConditionByReferenceLatent(ConditioningItem):
         self,
         latent: torch.Tensor,
         downscale_factor: int = 1,
+        temporal_scale_factor: int = 1,
         strength: float = 1.0,
     ):
         self.latent = latent
         self.downscale_factor = downscale_factor
+        self.temporal_scale_factor = temporal_scale_factor
         self.strength = strength
 
     def apply_to(
@@ -44,10 +47,9 @@ class VideoConditionByReferenceLatent(ConditioningItem):
         latent_state: LatentState,
         latent_tools: VideoLatentTools,
     ) -> LatentState:
-        """Append reference video tokens with scaled positions."""
+        """Append reference video tokens with positions translated into the target frame."""
         tokens = latent_tools.patchifier.patchify(self.latent)
 
-        # Compute positions for the reference video's actual dimensions
         latent_coords = latent_tools.patchifier.get_patch_grid_bounds(
             output_shape=VideoLatentShape.from_torch_shape(self.latent.shape),
             device=self.latent.device,
@@ -58,9 +60,18 @@ class VideoConditionByReferenceLatent(ConditioningItem):
             causal_fix=latent_tools.causal_fix,
         )
         positions = positions.to(dtype=torch.float32)
-        positions[:, 0, ...] /= latent_tools.fps
 
-        # Scale spatial positions to match target coordinate space
+        # Place ref tokens on their own time spacing (= target_fps / S).
+        positions[:, 0, ...] /= latent_tools.fps / self.temporal_scale_factor
+
+        # Translate into the target's frame so ref's last patch ends with target's last
+        # patch; clamp the causal patch's negative start back to [0, 1/target_fps).
+        if self.temporal_scale_factor != 1:
+            t_target = latent_state.positions[:, 0, 0:1, 1:2].to(dtype=torch.float32)  # = 1/target_fps
+            positions[:, 0, ...] = torch.clamp(
+                positions[:, 0, ...] - (self.temporal_scale_factor - 1) * t_target,
+                min=0,
+            )
         if self.downscale_factor != 1:
             positions[:, 1, ...] *= self.downscale_factor  # height axis
             positions[:, 2, ...] *= self.downscale_factor  # width axis
@@ -83,7 +94,7 @@ class VideoConditionByReferenceLatent(ConditioningItem):
         )
 
         return LatentState(
-            latent=torch.cat([latent_state.latent, tokens], dim=1),
+            latent=torch.cat([latent_state.latent, torch.zeros_like(tokens)], dim=1),
             denoise_mask=torch.cat([latent_state.denoise_mask, denoise_mask], dim=1),
             positions=torch.cat([latent_state.positions, positions], dim=2),
             clean_latent=torch.cat([latent_state.clean_latent, tokens], dim=1),

packages/ltx-core/src/ltx_core/hdr.py

@@ -0,0 +1,71 @@
+"""HDR utilities: LogC3 compression for HDR IC-LoRA training and inference.
+Provides compress/decompress and postprocess helpers for HDR video generation.
+Used by ltx-pipelines for HDR IC-LoRA and by ltx-trainer for HDR validation.
+"""
+
+from __future__ import annotations
+
+from typing import Literal
+
+import torch
+from torch import Tensor
+
+
+class LogC3:
+    """ARRI LogC3 (EI 800) HDR compression.
+    Maps linear [0, ∞) <-> LogC3 [0, 1] via the camera log curve. The log
+    curve allocates more precision to shadows/midtones and compresses
+    highlights smoothly. Callers are responsible for mapping the [0, 1]
+    output to the VAE's [-1, 1] input range.
+    """
+
+    name = "LogC3"
+    A = 5.555556
+    B = 0.052272
+    C = 0.247190
+    D = 0.385537
+    E = 5.367655
+    F = 0.092809
+    CUT = 0.010591
+
+    def compress(self, hdr: Tensor) -> Tensor:
+        """Compress linear HDR [0, ∞) → LogC3 [0, 1]."""
+        x = torch.clamp(hdr, min=0.0)
+        log_part = self.C * torch.log10(self.A * x + self.B) + self.D
+        lin_part = self.E * x + self.F
+        logc = torch.where(x >= self.CUT, log_part, lin_part)
+        return torch.clamp(logc, 0.0, 1.0)
+
+    def compress_ldr(self, ldr: Tensor) -> Tensor:
+        """Compress LDR [0, 1] → [0, 1] (no log curve, just clamp)."""
+        return torch.clamp(ldr, 0.0, 1.0)
+
+    def decompress(self, logc: Tensor) -> Tensor:
+        """Decompress LogC3 [0, 1] → linear HDR [0, ∞)."""
+        logc = torch.clamp(logc, 0.0, 1.0)
+        cut_log = self.E * self.CUT + self.F
+        lin_from_log = (torch.pow(10.0, (logc - self.D) / self.C) - self.B) / self.A
+        lin_from_lin = (logc - self.F) / self.E
+        return torch.where(logc >= cut_log, lin_from_log, lin_from_lin)
+
+    def decompress_ldr(self, logc: Tensor) -> Tensor:
+        """Decompress [0, 1] → LDR [0, 1] (identity clamp)."""
+        return torch.clamp(logc, 0.0, 1.0)
+
+
+def apply_hdr_decode_postprocess(
+    decoded_video: Tensor,
+    transform: Literal["logc3"] = "logc3",
+) -> Tensor:
+    """Apply HDR decompress to VAE decode output for HDR recovery.
+    Args:
+        decoded_video: Tensor from VAE decode in [0, 1], shape [B, C, F, H, W].
+            Must be float32 for sufficient color resolution.
+        transform: "logc3".
+    Returns:
+        HDR video tensor float32.
+    """
+    decoded_video = decoded_video.float()
+    if transform == "logc3":
+        return LogC3().decompress(decoded_video)
+    raise ValueError(f"Unsupported HDR transform: {transform}")

packages/ltx-core/src/ltx_core/layer_streaming.py

@@ -1,306 +0,0 @@
-"""Layer streaming wrapper for memory-efficient inference.
-Keeps most transformer/decoder layers on CPU pinned memory and streams them
-to GPU on demand, using a secondary CUDA stream to prefetch upcoming layers
-so that data transfer overlaps with compute.
-General-purpose: works with any ``nn.Module`` whose forward iterates over a
-``nn.ModuleList`` attribute (e.g. ``transformer_blocks``, ``layers``).
-Each layer is evicted back to CPU immediately after its forward completes,
-and prefetch uses modular indexing so the last layer's prefetch wraps around
-to prepare early layers for the next forward pass.
-Example
--------
->>> model = build_my_model(device=torch.device("cpu"))
->>> model = LayerStreamingWrapper(
-...     model,
-...     layers_attr="transformer_blocks",
-...     target_device=torch.device("cuda:0"),
-...     prefetch_count=2,
-... )
->>> out = model(inputs)            # hooks handle layer streaming
->>> model.teardown()               # move everything back to CPU
-"""
-
-from __future__ import annotations
-
-import functools
-import itertools
-import logging
-from typing import Any
-
-import torch
-from torch import nn
-
-logger = logging.getLogger(__name__)
-
-
-def _resolve_attr(module: nn.Module, dotted_path: str) -> nn.ModuleList:
-    """Resolve a dotted attribute path like ``'model.language_model.layers'``."""
-    obj: Any = module
-    for part in dotted_path.split("."):
-        obj = getattr(obj, part)
-    if not isinstance(obj, nn.ModuleList):
-        raise TypeError(f"Expected nn.ModuleList at '{dotted_path}', got {type(obj).__name__}")
-    return obj
-
-
-class _LayerStore:
-    """Manages CPU-pinned copies of layer parameters/buffers.
-    Tracks which layers currently reside on GPU so the prefetcher and evictor
-    can make correct decisions.
-    """
-
-    def __init__(self, layers: nn.ModuleList, target_device: torch.device) -> None:
-        self.target_device = target_device
-        self.num_layers = len(layers)
-
-        # CPU-pinned copies keyed by (layer_idx, param_name)
-        self._pinned: list[dict[str, torch.Tensor]] = []
-        self._on_gpu: set[int] = set()
-
-        for layer in layers:
-            pinned: dict[str, torch.Tensor] = {}
-            for name, tensor in itertools.chain(layer.named_parameters(), layer.named_buffers()):
-                pinned_tensor = tensor.data.pin_memory()
-                tensor.data = pinned_tensor
-                pinned[name] = pinned_tensor
-            self._pinned.append(pinned)
-
-    def _check_idx(self, idx: int) -> None:
-        if idx < 0 or idx >= self.num_layers:
-            raise IndexError(f"Layer index {idx} out of range [0, {self.num_layers})")
-
-    def is_on_gpu(self, idx: int) -> bool:
-        return idx in self._on_gpu
-
-    def move_to_gpu(self, idx: int, layer: nn.Module, *, non_blocking: bool = False) -> None:
-        """Move layer *idx* parameters from pinned CPU to *target_device*."""
-        self._check_idx(idx)
-        if idx in self._on_gpu:
-            return
-        pinned = self._pinned[idx]
-        for name, param in itertools.chain(layer.named_parameters(), layer.named_buffers()):
-            param.data = pinned[name].to(self.target_device, non_blocking=non_blocking)
-        self._on_gpu.add(idx)
-
-    def evict_to_cpu(self, idx: int, layer: nn.Module) -> None:
-        """Swap layer *idx* parameters back to their pinned CPU copies."""
-        self._check_idx(idx)
-        if idx not in self._on_gpu:
-            return
-        pinned = self._pinned[idx]
-        for name, param in itertools.chain(layer.named_parameters(), layer.named_buffers()):
-            param.data = pinned[name]
-        self._on_gpu.discard(idx)
-
-    def cleanup(self) -> None:
-        """Release all pinned memory references.
-        After this call, the pinned tensors can be garbage-collected once
-        the layer parameters (which still reference them via ``.data``) are
-        also released (e.g. via ``.to("meta")``).
-        """
-        for pinned_dict in self._pinned:
-            pinned_dict.clear()
-        self._pinned.clear()
-
-
-class _AsyncPrefetcher:
-    """Issues H2D transfers on a dedicated CUDA stream.
-    Uses per-layer CUDA events so that the compute stream only waits for the
-    specific layer it needs, not all pending transfers.
-    """
-
-    def __init__(self, store: _LayerStore, layers: nn.ModuleList) -> None:
-        self._store = store
-        self._layers = layers
-        self._stream = torch.cuda.Stream(device=store.target_device)
-        self._events: dict[int, torch.cuda.Event] = {}
-
-    def prefetch(self, idx: int) -> None:
-        """Begin async transfer of layer *idx* to GPU (no-op if already there)."""
-        if self._store.is_on_gpu(idx) or idx in self._events:
-            return
-        with torch.cuda.stream(self._stream):
-            self._store.move_to_gpu(idx, self._layers[idx], non_blocking=True)
-            event = torch.cuda.Event()
-            event.record(self._stream)
-            self._events[idx] = event
-
-    def wait(self, idx: int) -> None:
-        """Block the compute stream until layer *idx* transfer is complete."""
-        event = self._events.pop(idx, None)
-        if event is not None:
-            torch.cuda.current_stream(self._store.target_device).wait_event(event)
-
-    def cleanup(self) -> None:
-        """Drain pending work and release CUDA stream/event resources."""
-        self._events.clear()
-        self._stream = None
-        self._layers = None
-        self._store = None
-
-
-class LayerStreamingWrapper(nn.Module):
-    """Wraps a model to stream its sequential layers between CPU and GPU.
-    Each layer is evicted immediately after its forward completes, and
-    prefetch wraps around using modular indexing so the end of one forward
-    pass prepares early layers for the next.
-    Parameters
-    ----------
-    model:
-        The model to wrap, with all parameters on **CPU**.
-    layers_attr:
-        Dotted attribute path to the ``nn.ModuleList`` of sequential layers
-        (e.g. ``"transformer_blocks"`` or ``"model.language_model.layers"``).
-    target_device:
-        The GPU device to use for compute.
-    prefetch_count:
-        How many layers ahead to prefetch.  The maximum number of layers on
-        GPU at once is ``1 + prefetch_count``.  Must be >= 1.
-    """
-
-    def __init__(
-        self,
-        model: nn.Module,
-        layers_attr: str,
-        target_device: torch.device,
-        prefetch_count: int = 2,
-    ) -> None:
-        if prefetch_count < 1:
-            raise ValueError("prefetch_count must be >= 1")
-        super().__init__()
-        # Store the wrapped model as a submodule so parameters are discoverable.
-        self._model = model
-        self._layers = _resolve_attr(model, layers_attr)
-        self._target_device = target_device
-        # Clamp: no point prefetching more than num_layers - 1 (the rest are evicted).
-        self._prefetch_count = min(prefetch_count, len(self._layers) - 1)
-        self._hooks: list[torch.utils.hooks.RemovableHandle] = []
-
-        self._setup()
-
-    # ------------------------------------------------------------------
-    # Setup / teardown
-    # ------------------------------------------------------------------
-
-    def _setup(self) -> None:
-        # 1. Build the pinned CPU store (copies all layer tensors to pinned memory).
-        self._store = _LayerStore(self._layers, self._target_device)
-
-        # 2. Move all NON-layer params/buffers to GPU.
-        layer_tensor_ids: set[int] = set()
-        for layer in self._layers:
-            for t in itertools.chain(layer.parameters(), layer.buffers()):
-                layer_tensor_ids.add(id(t))
-
-        for p in self._model.parameters():
-            if id(p) not in layer_tensor_ids:
-                p.data = p.data.to(self._target_device)
-        for b in self._model.buffers():
-            if id(b) not in layer_tensor_ids:
-                b.data = b.data.to(self._target_device)
-
-        # 3. Pre-load the first (1 + prefetch_count) layers synchronously.
-        for idx in range(min(self._prefetch_count + 1, len(self._layers))):
-            self._store.move_to_gpu(idx, self._layers[idx])
-
-        # 4. Create the async prefetcher and register hooks.
-        self._prefetcher = _AsyncPrefetcher(self._store, self._layers)
-        self._register_hooks()
-
-    def _register_hooks(self) -> None:
-        idx_map: dict[int, int] = {id(layer): idx for idx, layer in enumerate(self._layers)}
-        num_layers = len(self._layers)
-
-        def _pre_hook(
-            module: nn.Module,
-            _args: Any,  # noqa: ANN401
-            *,
-            idx: int,
-        ) -> None:
-            # Wait only for THIS layer's H2D transfer (not all pending ones).
-            self._prefetcher.wait(idx)
-            if not self._store.is_on_gpu(idx):
-                self._store.move_to_gpu(idx, module)
-
-            # Record that the compute stream will read these weight tensors.
-            # They were allocated on the prefetch stream, so without this the
-            # caching allocator would allow the prefetch stream to reuse their
-            # memory immediately after eviction — even if the compute kernel
-            # that reads them hasn't finished yet.
-            compute_stream = torch.cuda.current_stream(self._target_device)
-            for param in itertools.chain(module.parameters(), module.buffers()):
-                param.data.record_stream(compute_stream)
-
-            # Kick off prefetch for upcoming layers (wraps around for next pass).
-            for offset in range(1, self._prefetch_count + 1):
-                self._prefetcher.prefetch((idx + offset) % num_layers)
-
-        def _post_hook(
-            module: nn.Module,
-            _args: Any,  # noqa: ANN401
-            _output: Any,  # noqa: ANN401
-            *,
-            idx: int,
-        ) -> None:
-            # Evict this layer immediately — its computation is done.
-            self._store.evict_to_cpu(idx, module)
-
-        for layer in self._layers:
-            idx = idx_map[id(layer)]
-            h1 = layer.register_forward_pre_hook(functools.partial(_pre_hook, idx=idx))
-            h2 = layer.register_forward_hook(functools.partial(_post_hook, idx=idx))
-            self._hooks.extend([h1, h2])
-
-    def teardown(self) -> None:
-        """Remove hooks, release pinned memory, and move parameters back to CPU.
-        After this call the wrapper is inert: hooks are removed, the prefetch
-        stream is drained and destroyed, all parameters reside on regular
-        (non-pinned) CPU memory, and the ``_LayerStore`` pinned-tensor cache is
-        cleared.  Callers should still follow up with ``.to("meta")`` to release
-        the CPU copies if the model is no longer needed.
-        """
-        for h in self._hooks:
-            h.remove()
-        self._hooks.clear()
-
-        # Drain all in-flight async H2D copies, then release stream resources.
-        # Without the synchronize, clearing the stream/events can trigger
-        # use-after-free at the CUDA driver level.
-        torch.cuda.synchronize(device=self._target_device)
-        if self._prefetcher is not None:
-            self._prefetcher.cleanup()
-            self._prefetcher = None
-
-        # Move everything to CPU.
-        for idx, layer in enumerate(self._layers):
-            self._store.evict_to_cpu(idx, layer)
-
-        for p in self._model.parameters():
-            p.data = p.data.to("cpu")
-        for b in self._model.buffers():
-            b.data = b.data.to("cpu")
-
-        # Release pinned memory.  After evict_to_cpu() the layer parameters
-        # still reference the pinned tensors (since .to("cpu") on a pinned
-        # tensor is a no-op).  The caller is expected to follow up with
-        # .to("meta") to drop the param refs; cleanup() drops the store's refs.
-        self._store.cleanup()
-
-    # ------------------------------------------------------------------
-    # Forward and attribute delegation
-    # ------------------------------------------------------------------
-
-    def forward(self, *args: Any, **kwargs: Any) -> Any:  # noqa: ANN401
-        return self._model(*args, **kwargs)
-
-    def __getattr__(self, name: str) -> Any:  # noqa: ANN401
-        """Proxy attribute access to the wrapped model.
-        This allows calling methods like ``encode()`` on a wrapped
-        GemmaTextEncoder without the caller needing to know about the wrapper.
-        ``nn.Module.__getattr__`` is only called when normal attribute lookup
-        fails, so ``_model``, ``_store``, etc. are found first via ``__dict__``.
-        """
-        try:
-            return super().__getattr__(name)
-        except AttributeError:
-            return getattr(self._model, name)

packages/ltx-core/src/ltx_core/loader/__init__.py

@@ -1,6 +1,11 @@
 """Loader utilities for model weights, LoRAs, and safetensor operations."""
 
 from ltx_core.loader.fuse_loras import apply_loras
+from ltx_core.loader.helpers import (
+    create_meta_model,
+    load_state_dict,
+    read_model_config,
+)
 from ltx_core.loader.module_ops import ModuleOps
 from ltx_core.loader.primitives import (
     LoRAAdaptableProtocol,
@@ -45,4 +50,7 @@ __all__ = [
     "StateDictLoader",
     "StateDictRegistry",
     "apply_loras",
+    "create_meta_model",
+    "load_state_dict",
+    "read_model_config",
 ]

packages/ltx-core/src/ltx_core/loader/attention_ops.py

@@ -0,0 +1,42 @@
+"""Builder ops for swapping attention backends on a meta model before load."""
+
+import torch
+
+from ltx_core.loader.module_ops import ModuleOps
+from ltx_core.model.transformer.attention import (
+    Attention,
+    AttentionCallable,
+    AttentionFunction,
+    MaskedAttentionCallable,
+    MaskedAttentionFunction,
+)
+
+
+def set_attention_module_op(
+    attention: AttentionFunction | AttentionCallable | None = None,
+    masked_attention: MaskedAttentionFunction | MaskedAttentionCallable | None = None,
+) -> ModuleOps:
+    """Build a ``ModuleOps`` that overrides the attention callables on every
+    ``Attention`` submodule of a model. Applied via ``create_meta_model`` so
+    the meta model is mutated before weight loading. Matcher returns False
+    for models with no ``Attention`` submodules, so the op is a no-op there.
+    Either or both slots may be supplied; *None* leaves that slot untouched.
+    """
+    fn = attention.to_callable() if isinstance(attention, AttentionFunction) else attention
+    masked_fn = (
+        masked_attention.to_callable() if isinstance(masked_attention, MaskedAttentionFunction) else masked_attention
+    )
+
+    def matcher(model: torch.nn.Module) -> bool:
+        return any(isinstance(m, Attention) for m in model.modules())
+
+    def mutator(model: torch.nn.Module) -> torch.nn.Module:
+        for module in model.modules():
+            if isinstance(module, Attention):
+                if fn is not None:
+                    module.attention_function = fn
+                if masked_fn is not None:
+                    module.masked_attention_function = masked_fn
+        return model
+
+    return ModuleOps(name="set_attention_backend", matcher=matcher, mutator=mutator)

packages/ltx-core/src/ltx_core/loader/fuse_loras.py

@@ -1,10 +1,74 @@
-from collections.abc import Iterator
+from collections.abc import Callable, Iterable, Iterator
+from dataclasses import dataclass
+from typing import NamedTuple
 
 import torch
 
 from ltx_core.loader.primitives import LoraStateDictWithStrength, StateDict
-from ltx_core.quantization.fp8_cast import _fused_add_round_launch
-from ltx_core.quantization.fp8_scaled_mm import quantize_weight_to_fp8_per_tensor
+
+
+class LoraProduct(NamedTuple):
+    """A LoRA's ``A``, ``B`` factors and its strength scalar."""
+
+    a: torch.Tensor
+    b: torch.Tensor
+    strength: float
+
+
+#: Signature for a fuse callable used by :class:`FuseRule`.
+#:
+#: Args:
+#:     key: The state-dict key being fused (e.g. ``"...layers.0.attn.q.weight"``).
+#:     weight: The current value at ``key`` from ``model_sd``, on the fusion device.
+#:     deltas: The pre-aggregated LoRA delta for ``key``, in ``aggregation_dtype``.
+#:     model_sd: The full state dict, for rules that need companion keys
+#:         (e.g. an existing ``.weight_scale``).
+#:
+#: Returns a dict of state-dict keys to overwrite -- at minimum ``{key: new_weight}``,
+#: plus any companion keys (e.g. an updated ``.weight_scale``) the policy needs to
+#: keep in sync.
+FuseFn = Callable[[str, torch.Tensor, torch.Tensor, StateDict], dict[str, torch.Tensor]]
+
+
+@dataclass(frozen=True)
+class FuseRule:
+    """Fuse an aggregated LoRA delta into one weight key.
+    Each policy supplies its own rule (see ``QuantizationPolicy.fuse_rule``);
+    ``fuse_lora_weights`` is policy-agnostic boilerplate around it.
+    Attributes:
+        aggregation_dtype: Dtype callers must pre-aggregate LoRA deltas in
+            before invoking the rule.
+        fuse_fn: Callable that applies the pre-aggregated deltas to the weight
+            (and any companion keys) and returns a dict of keys to overwrite —
+            at minimum ``{key: new_weight}``, plus any companion keys (e.g. an
+            updated ``.weight_scale`` for scaled-FP8 layouts) the policy needs
+            to keep in sync.
+    """
+
+    aggregation_dtype: torch.dtype
+    fuse_fn: FuseFn
+
+    def __call__(
+        self,
+        key: str,
+        weight: torch.Tensor,
+        deltas: torch.Tensor,
+        model_sd: StateDict,
+    ) -> dict[str, torch.Tensor]:
+        return self.fuse_fn(key, weight, deltas, model_sd)
+
+
+def _bf16_fuse(
+    key: str,
+    weight: torch.Tensor,
+    deltas: torch.Tensor,
+    model_sd: StateDict,  # noqa: ARG001
+) -> dict[str, torch.Tensor]:
+    deltas.add_(weight)
+    return {key: deltas.to(dtype=weight.dtype)}
+
+
+bf16_fuse_rule = FuseRule(aggregation_dtype=torch.bfloat16, fuse_fn=_bf16_fuse)
 
 
 def _get_device() -> torch.device:
@@ -13,121 +77,102 @@ def _get_device() -> torch.device:
     return torch.device("cpu")
 
 
+def aggregate_lora_products(
+    products: Iterable[LoraProduct],
+    dtype: torch.dtype,
+) -> torch.Tensor | None:
+    """Accumulate ``sum((B * strength) @ A)`` across :class:`LoraProduct` items.
+    The first product materializes a freshly-allocated aggregator via
+    ``torch.matmul(B * strength, A).to(dtype)`` -- preserving the
+    ``(B * strength) @ A`` rounding pattern. Subsequent products use
+    ``addmm_`` to avoid allocating the full intermediate delta.
+    Returns the aggregator, or ``None`` if ``products`` was empty.
+    """
+    aggregated: torch.Tensor | None = None
+    for product in products:
+        if aggregated is None:
+            aggregated = torch.matmul(product.b * product.strength, product.a).to(dtype=dtype)
+        else:
+            aggregated.addmm_(product.b, product.a, alpha=product.strength)
+    return aggregated
+
+
 def fuse_lora_weights(
     model_sd: StateDict,
     lora_sd_and_strengths: list[LoraStateDictWithStrength],
-    dtype: torch.dtype | None = None,
+    fuse_rule: FuseRule = bf16_fuse_rule,
+    preserve_input_device: bool = True,
 ) -> Iterator[tuple[str, torch.Tensor]]:
     """Yield ``(key, fused_tensor)`` for each weight modified by at least one LoRA.
-    For scaled-FP8 weights, this includes both the updated ``.weight`` tensor
-    and its corresponding ``.weight_scale`` tensor.
+    The fusion math is delegated to ``fuse_rule``.
+    Output dtypes are the rule's responsibility.
+    When ``preserve_input_device`` is False, fused tensors are yielded on the device
+    used for fusion; caller is responsible for moving them to their final
+    destination.
     """
-    for key, original_weight in model_sd.sd.items():
-        if original_weight is None or key.endswith(".weight_scale"):
+    fusion_device = _get_device()
+    for key in _affected_weight_keys(lora_sd_and_strengths):
+        original_weight = model_sd.sd.get(key)
+        if original_weight is None:
             continue
-        original_device = original_weight.device
-        weight = original_weight.to(device=_get_device())
-        target_dtype = dtype if dtype is not None else weight.dtype
-        deltas_dtype = target_dtype if target_dtype not in [torch.float8_e4m3fn, torch.float8_e5m2] else torch.bfloat16
 
-        deltas = _prepare_deltas(lora_sd_and_strengths, key, deltas_dtype, weight.device)
+        products = _products_for_sd_key(lora_sd_and_strengths, key, fuse_rule.aggregation_dtype, fusion_device)
+        deltas = aggregate_lora_products(products, fuse_rule.aggregation_dtype)
         if deltas is None:
             continue
 
-        scale_key = key.replace(".weight", ".weight_scale") if key.endswith(".weight") else None
-        is_scaled_fp8 = scale_key is not None and scale_key in model_sd.sd
+        original_device = original_weight.device
+        weight = original_weight.to(device=fusion_device)
 
-        if weight.dtype == torch.float8_e4m3fn:
-            if is_scaled_fp8:
-                fused = _fuse_delta_with_scaled_fp8(deltas, weight, key, scale_key, model_sd)
-            else:
-                fused = _fuse_delta_with_cast_fp8(deltas, weight, key, target_dtype)
-        elif weight.dtype == torch.bfloat16:
-            fused = _fuse_delta_with_bfloat16(deltas, weight, key, target_dtype)
-        else:
-            raise ValueError(f"Unsupported dtype: {weight.dtype}")
+        fused = fuse_rule(key, weight, deltas, model_sd)
 
         for k, v in fused.items():
-            yield k, v.to(device=original_device)
+            yield k, v.to(device=original_device) if preserve_input_device else v
 
 
 def apply_loras(
     model_sd: StateDict,
     lora_sd_and_strengths: list[LoraStateDictWithStrength],
-    dtype: torch.dtype | None = None,
+    fuse_rule: FuseRule = bf16_fuse_rule,
     destination_sd: StateDict | None = None,
 ) -> StateDict:
+    """Fuse LoRAs into ``model_sd`` and place the results in ``destination_sd``.
+    When ``destination_sd`` is provided, the fused tensors are placed directly into it.
+    """
+    fused_iter = fuse_lora_weights(
+        model_sd,
+        lora_sd_and_strengths,
+        fuse_rule=fuse_rule,
+    )
     if destination_sd is not None:
-        sd = destination_sd.sd
-        for key, tensor in fuse_lora_weights(model_sd, lora_sd_and_strengths, dtype):
-            sd[key] = tensor
+        for key, fused in fused_iter:
+            destination_sd.sd[key] = fused
         return destination_sd
 
-    fused = dict(fuse_lora_weights(model_sd, lora_sd_and_strengths, dtype))
+    fused = dict(fused_iter)
     sd = {k: (fused[k] if k in fused else v.clone()) for k, v in model_sd.sd.items()}
     return StateDict(sd, model_sd.device, model_sd.size, model_sd.dtype)
 
 
-def _prepare_deltas(
-    lora_sd_and_strengths: list[LoraStateDictWithStrength], key: str, dtype: torch.dtype, device: torch.device
-) -> torch.Tensor | None:
-    deltas = []
+def _affected_weight_keys(lora_sd_and_strengths: list[LoraStateDictWithStrength]) -> set[str]:
+    """Return the set of ``.weight`` keys touched by at least one LoRA in the list."""
+    suffix = ".lora_A.weight"
+    return {k[: -len(suffix)] + ".weight" for lsd, _ in lora_sd_and_strengths for k in lsd.sd if k.endswith(suffix)}
+
+
+def _products_for_sd_key(
+    lora_sd_and_strengths: list[LoraStateDictWithStrength],
+    key: str,
+    dtype: torch.dtype,
+    device: torch.device,
+) -> Iterator[LoraProduct]:
+    """Yield :class:`LoraProduct` items matching *key* across state-dict-backed LoRAs."""
     prefix = key[: -len(".weight")]
     key_a = f"{prefix}.lora_A.weight"
     key_b = f"{prefix}.lora_B.weight"
     for lsd, coef in lora_sd_and_strengths:
         if key_a not in lsd.sd or key_b not in lsd.sd:
             continue
-        a = lsd.sd[key_a].to(device=device)
-        b = lsd.sd[key_b].to(device=device)
-        product = torch.matmul(b * coef, a)
-        del a, b
-        deltas.append(product.to(dtype=dtype))
-    if len(deltas) == 0:
-        return None
-    elif len(deltas) == 1:
-        return deltas[0]
-    return torch.sum(torch.stack(deltas, dim=0), dim=0)
-
-
-def _fuse_delta_with_scaled_fp8(
-    deltas: torch.Tensor,
-    weight: torch.Tensor,
-    key: str,
-    scale_key: str,
-    model_sd: StateDict,
-) -> dict[str, torch.Tensor]:
-    """Dequantize scaled FP8 weight, add LoRA delta, and re-quantize."""
-    weight_scale = model_sd.sd[scale_key]
-
-    original_weight = weight.t().to(torch.float32) * weight_scale
-
-    new_weight = original_weight + deltas.to(torch.float32)
-
-    new_fp8_weight, new_weight_scale = quantize_weight_to_fp8_per_tensor(new_weight)
-    return {key: new_fp8_weight, scale_key: new_weight_scale}
-
-
-def _fuse_delta_with_cast_fp8(
-    deltas: torch.Tensor,
-    weight: torch.Tensor,
-    key: str,
-    target_dtype: torch.dtype,
-) -> dict[str, torch.Tensor]:
-    """Fuse LoRA delta with cast-only FP8 weight (no scale factor)."""
-    if str(weight.device).startswith("cuda"):
-        _fused_add_round_launch(deltas, weight, seed=0)
-    else:
-        deltas.add_(weight.to(dtype=deltas.dtype))
-    return {key: deltas.to(dtype=target_dtype)}
-
-
-def _fuse_delta_with_bfloat16(
-    deltas: torch.Tensor,
-    weight: torch.Tensor,
-    key: str,
-    target_dtype: torch.dtype,
-) -> dict[str, torch.Tensor]:
-    """Fuse LoRA delta with bfloat16 weight."""
-    deltas.add_(weight)
-    return {key: deltas.to(dtype=target_dtype)}
+        a = lsd.sd[key_a].to(device=device, dtype=dtype, non_blocking=True)
+        b = lsd.sd[key_b].to(device=device, dtype=dtype, non_blocking=True)
+        yield LoraProduct(a, b, coef)

packages/ltx-core/src/ltx_core/loader/helpers.py

@@ -0,0 +1,61 @@
+"""Shared model-construction helpers used by both SingleGPUModelBuilder and StreamingModelBuilder."""
+
+from __future__ import annotations
+
+from typing import TypeVar
+
+import torch
+from torch import nn
+
+from ltx_core.loader.module_ops import ModuleOps
+from ltx_core.loader.primitives import StateDict, StateDictLoader
+from ltx_core.loader.registry import Registry
+from ltx_core.loader.sd_ops import SDOps
+from ltx_core.model.model_protocol import ModelConfigurator
+
+_M = TypeVar("_M", bound=nn.Module)
+
+
+def load_state_dict(
+    paths: str | tuple[str, ...] | list[str],
+    loader: StateDictLoader,
+    registry: Registry,
+    device: torch.device | None,
+    sd_ops: SDOps | None = None,
+) -> StateDict:
+    """Load a state dict from disk, using registry caching."""
+    if isinstance(paths, str):
+        path_list = [paths]
+    elif isinstance(paths, tuple):
+        path_list = list(paths)
+    else:
+        path_list = paths
+    cached = registry.get(path_list, sd_ops)
+    if cached is not None:
+        return cached
+    result = loader.load(path_list, sd_ops=sd_ops, device=device)
+    registry.add(path_list, sd_ops=sd_ops, state_dict=result)
+    return result
+
+
+def read_model_config(
+    model_path: str | tuple[str, ...],
+    loader: StateDictLoader,
+) -> dict:
+    """Read metadata from the first shard of a checkpoint."""
+    first = model_path[0] if isinstance(model_path, tuple) else model_path
+    return loader.metadata(first)
+
+
+def create_meta_model(
+    configurator: type[ModelConfigurator[_M]],
+    config: dict,
+    module_ops: tuple[ModuleOps, ...] = (),
+) -> _M:
+    """Create a model on the meta device and apply module operations."""
+    with torch.device("meta"):
+        model = configurator.from_config(config)
+    for op in module_ops:
+        if op.matcher(model):
+            model = op.mutator(model)
+    return model

packages/ltx-core/src/ltx_core/loader/kernels.py

@@ -1,72 +1,79 @@
 # ruff: noqa: ANN001, ANN201, ERA001, N803, N806
-import triton
-import triton.language as tl
+try:
+    import triton
+    import triton.language as tl
+
+    TRITON_AVAILABLE = True
+except (ImportError, OSError):
+    TRITON_AVAILABLE = False
 
 
-@triton.jit
-def fused_add_round_kernel(
-    x_ptr,
-    output_ptr,  # contents will be added to the output
-    seed,
-    n_elements,
-    EXPONENT_BIAS,
-    MANTISSA_BITS,
-    BLOCK_SIZE: tl.constexpr,
-):
-    """
-    A kernel to upcast 8bit quantized weights to bfloat16 with stochastic rounding
-    and add them to bfloat16 output weights. Might be used to upcast original model weights
-    and to further add them to precalculated deltas coming from LoRAs.
-    """
-    # Get program ID and compute offsets
-    pid = tl.program_id(axis=0)
-    block_start = pid * BLOCK_SIZE
-    offsets = block_start + tl.arange(0, BLOCK_SIZE)
-    mask = offsets < n_elements
+if TRITON_AVAILABLE:
 
-    # Load data
-    x = tl.load(x_ptr + offsets, mask=mask)
-    rand_vals = tl.rand(seed, offsets) - 0.5
+    @triton.jit
+    def fused_add_round_kernel(
+        x_ptr,
+        output_ptr,  # contents will be added to the output
+        seed,
+        n_elements,
+        EXPONENT_BIAS,
+        MANTISSA_BITS,
+        BLOCK_SIZE: tl.constexpr,
+    ):
+        """
+        A kernel to upcast 8bit quantized weights to bfloat16 with stochastic rounding
+        and add them to bfloat16 output weights. Might be used to upcast original model weights
+        and to further add them to precalculated deltas coming from LoRAs.
+        """
+        # Get program ID and compute offsets
+        pid = tl.program_id(axis=0)
+        block_start = pid * BLOCK_SIZE
+        offsets = block_start + tl.arange(0, BLOCK_SIZE)
+        mask = offsets < n_elements
 
-    x = tl.cast(x, tl.float16)
-    delta = tl.load(output_ptr + offsets, mask=mask)
-    delta = tl.cast(delta, tl.float16)
-    x = x + delta
+        # Load data
+        x = tl.load(x_ptr + offsets, mask=mask)
+        rand_vals = tl.rand(seed, offsets) - 0.5
 
-    x_bits = tl.cast(x, tl.int16, bitcast=True)
+        x = tl.cast(x, tl.float16)
+        delta = tl.load(output_ptr + offsets, mask=mask)
+        delta = tl.cast(delta, tl.float16)
+        x = x + delta
 
-    # Calculate the exponent. Unbiased fp16 exponent is ((x_bits & 0x7C00) >> 10) - 15 for
-    # normal numbers and -14 for subnormals.
-    fp16_exponent_bits = (x_bits & 0x7C00) >> 10
-    fp16_normals = fp16_exponent_bits > 0
-    fp16_exponent = tl.where(fp16_normals, fp16_exponent_bits - 15, -14)
+        x_bits = tl.cast(x, tl.int16, bitcast=True)
 
-    # Add the target dtype's exponent bias and clamp to the target dtype's exponent range.
-    exponent = fp16_exponent + EXPONENT_BIAS
-    MAX_EXPONENT = 2 * EXPONENT_BIAS + 1
-    exponent = tl.where(exponent > MAX_EXPONENT, MAX_EXPONENT, exponent)
-    exponent = tl.where(exponent < 0, 0, exponent)
+        # Calculate the exponent. Unbiased fp16 exponent is ((x_bits & 0x7C00) >> 10) - 15 for
+        # normal numbers and -14 for subnormals.
+        fp16_exponent_bits = (x_bits & 0x7C00) >> 10
+        fp16_normals = fp16_exponent_bits > 0
+        fp16_exponent = tl.where(fp16_normals, fp16_exponent_bits - 15, -14)
 
-    # Normal ULP exponent, expressed as an fp16 exponent field:
-    # (exponent - EXPONENT_BIAS - MANTISSA_BITS) + 15
-    # Simplifies to: fp16_exponent - MANTISSA_BITS + 15
-    # See https://en.wikipedia.org/wiki/Unit_in_the_last_place
-    eps_exp = tl.maximum(0, tl.minimum(31, exponent - EXPONENT_BIAS - MANTISSA_BITS + 15))
+        # Add the target dtype's exponent bias and clamp to the target dtype's exponent range.
+        exponent = fp16_exponent + EXPONENT_BIAS
+        MAX_EXPONENT = 2 * EXPONENT_BIAS + 1
+        exponent = tl.where(exponent > MAX_EXPONENT, MAX_EXPONENT, exponent)
+        exponent = tl.where(exponent < 0, 0, exponent)
 
-    # Calculate epsilon in the target dtype
-    eps_normal = tl.cast(tl.cast(eps_exp << 10, tl.int16), tl.float16, bitcast=True)
+        # Normal ULP exponent, expressed as an fp16 exponent field:
+        # (exponent - EXPONENT_BIAS - MANTISSA_BITS) + 15
+        # Simplifies to: fp16_exponent - MANTISSA_BITS + 15
+        # See https://en.wikipedia.org/wiki/Unit_in_the_last_place
+        eps_exp = tl.maximum(0, tl.minimum(31, exponent - EXPONENT_BIAS - MANTISSA_BITS + 15))
 
-    # Subnormal ULP: 2^(1 - EXPONENT_BIAS - MANTISSA_BITS) ->
-    # fp16 exponent bits: (1 - EXPONENT_BIAS - MANTISSA_BITS) + 15 =
-    # 16 - EXPONENT_BIAS - MANTISSA_BITS
-    eps_subnormal = tl.cast(tl.cast((16 - EXPONENT_BIAS - MANTISSA_BITS) << 10, tl.int16), tl.float16, bitcast=True)
-    eps = tl.where(exponent > 0, eps_normal, eps_subnormal)
+        # Calculate epsilon in the target dtype
+        eps_normal = tl.cast(tl.cast(eps_exp << 10, tl.int16), tl.float16, bitcast=True)
 
-    # Apply zero mask to epsilon
-    eps = tl.where(x == 0, 0.0, eps)
+        # Subnormal ULP: 2^(1 - EXPONENT_BIAS - MANTISSA_BITS) ->
+        # fp16 exponent bits: (1 - EXPONENT_BIAS - MANTISSA_BITS) + 15 =
+        # 16 - EXPONENT_BIAS - MANTISSA_BITS
+        eps_subnormal = tl.cast(tl.cast((16 - EXPONENT_BIAS - MANTISSA_BITS) << 10, tl.int16), tl.float16, bitcast=True)
+        eps = tl.where(exponent > 0, eps_normal, eps_subnormal)
 
-    # Apply stochastic rounding
-    output = tl.cast(x + rand_vals * eps, tl.bfloat16)
+        # Apply zero mask to epsilon
+        eps = tl.where(x == 0, 0.0, eps)
 
-    # Store the result
-    tl.store(output_ptr + offsets, output, mask=mask)
+        # Apply stochastic rounding
+        output = tl.cast(x + rand_vals * eps, tl.bfloat16)
+
+        # Store the result
+        tl.store(output_ptr + offsets, output, mask=mask)

packages/ltx-core/src/ltx_core/loader/primitives.py

@@ -1,17 +1,25 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, NamedTuple, Protocol
+from typing import TYPE_CHECKING, Any, NamedTuple, Protocol, TypeVar
 
 import torch
 
 from ltx_core.loader.module_ops import ModuleOps
 from ltx_core.loader.sd_ops import SDOps
-from ltx_core.model.model_protocol import ModelType
 
 if TYPE_CHECKING:
+    from typing_extensions import Self
+
+    from ltx_core.loader.fuse_loras import FuseRule
     from ltx_core.loader.registry import Registry
 
+BuiltType = TypeVar("BuiltType", covariant=True)  # noqa: PLC0105
+
+
+# Per-key shape and dtype description for a flat collection of tensors.
+TensorLayout = dict[str, tuple[torch.Size, torch.dtype]]
+
 
 @dataclass(frozen=True)
 class StateDict:
@@ -45,62 +53,75 @@ class StateDictLoader(Protocol):
         """
         Load metadata from path
         """
+        ...
 
     def load(self, path: str | list[str], sd_ops: SDOps | None = None, device: torch.device | None = None) -> StateDict:
         """
         Load state dict from path or paths (for sharded model storage) and apply sd_ops
         """
-
-
-class ModelBuilderProtocol(Protocol[ModelType]):
-    """
-    Protocol for building PyTorch models from configuration dictionaries.
-    Implementations must provide:
-    - meta_model: Create a model from configuration dictionary and apply module operations
-    - build: Create and initialize a model from state dictionary and apply dtype transformations
-    """
-
-    model_sd_ops: SDOps | None
-    module_ops: tuple[ModuleOps, ...]
-    loras: tuple["LoraPathStrengthAndSDOps", ...]
-    registry: "Registry"
-
-    def meta_model(self, config: dict, module_ops: list[ModuleOps] | None = None) -> ModelType:
-        """
-        Create a model on the meta device from a configuration dictionary.
-        This decouples model creation from weight loading, allowing the model
-        architecture to be instantiated without allocating memory for parameters.
-        Args:
-            config: Model configuration dictionary.
-            module_ops: Optional list of module operations to apply (e.g., quantization).
-        Returns:
-            Model instance on meta device (no actual memory allocated for parameters).
-        """
         ...
 
-    def with_sd_ops(self, sd_ops: SDOps | None) -> "ModelBuilderProtocol[ModelType]":
-        """Return a copy of this builder with the given state-dict key remapping ops."""
-        ...
 
-    def with_module_ops(self, module_ops: tuple[ModuleOps, ...]) -> "ModelBuilderProtocol[ModelType]":
-        """Return a copy of this builder with the given module operations (e.g. quantization)."""
-        ...
+class BuilderProtocol(Protocol[BuiltType]):
+    """Protocol for model builders that produce a model via ``build()``."""
 
-    def with_loras(self, loras: tuple["LoraPathStrengthAndSDOps", ...]) -> "ModelBuilderProtocol[ModelType]":
-        """Return a copy of this builder with the given LoRAs to fuse at build time."""
-        ...
+    def build(
+        self,
+        device: torch.device | None = None,
+        dtype: torch.dtype | None = None,
+        **kwargs: Any,  # noqa: ANN401
+    ) -> BuiltType: ...
 
-    def with_registry(self, registry: "Registry") -> "ModelBuilderProtocol[ModelType]":
+    @property
+    def registry(self) -> "Registry": ...
+
+    def with_registry(self, registry: "Registry") -> "Self":
         """Return a copy of this builder using the given weight registry for allocation."""
         ...
 
-    def with_lora_load_device(self, device: torch.device) -> "ModelBuilderProtocol[ModelType]":
+
+class ModelBuilderProtocol(BuilderProtocol[BuiltType], Protocol[BuiltType]):
+    """
+    Protocol for building PyTorch models from configuration dictionaries.
+    Implementations must provide:
+    - build: Create and initialize a model from state dictionary and apply dtype transformations
+    """
+
+    @property
+    def model_sd_ops(self) -> SDOps | None: ...
+
+    @property
+    def module_ops(self) -> tuple[ModuleOps, ...]: ...
+
+    @property
+    def loras(self) -> tuple["LoraPathStrengthAndSDOps", ...]: ...
+
+    def with_sd_ops(self, sd_ops: SDOps | None) -> "Self":
+        """Return a copy of this builder with the given state-dict key remapping ops."""
+        ...
+
+    def with_module_ops(self, module_ops: tuple[ModuleOps, ...]) -> "Self":
+        """Return a copy of this builder with the given module operations (e.g. quantization)."""
+        ...
+
+    def with_loras(self, loras: tuple["LoraPathStrengthAndSDOps", ...]) -> "Self":
+        """Return a copy of this builder with the given LoRAs to fuse at build time."""
+        ...
+
+    def with_lora_load_device(self, device: torch.device) -> "Self":
         """Return a copy of this builder that loads LoRA weights onto the given device."""
         ...
 
+    def with_fuse_rule(self, fuse_rule: "FuseRule") -> "Self":
+        """Return a copy of this builder with the given LoRA fuse rule (e.g. from a quantization policy)."""
+        ...
+
     def build(
-        self, device: torch.device | None = None, dtype: torch.dtype | None = None, **kwargs: object
-    ) -> ModelType:
+        self,
+        device: torch.device | None = None,
+        dtype: torch.dtype | None = None,
+        **kwargs: Any,  # noqa: ANN401
+    ) -> BuiltType:
         """
         Build the model
         Args:
@@ -123,8 +144,7 @@ class LoRAAdaptableProtocol(Protocol):
     - lora: Add a LoRA to the model
     """
 
-    def lora(self, lora_path: str, strength: float) -> "LoRAAdaptableProtocol":
-        pass
+    def lora(self, lora_path: str, strength: float, sd_ops: SDOps) -> "LoRAAdaptableProtocol": ...
 
 
 class LoraPathStrengthAndSDOps(NamedTuple):

packages/ltx-core/src/ltx_core/loader/sd_ops.py

@@ -24,6 +24,7 @@ class ContentMatching:
 
     prefix: str = ""
     suffix: str = ""
+    contains: str = ""
 
 
 class KeyValueOperationResult(NamedTuple):
@@ -72,10 +73,10 @@ class SDOps:
         new_mapping = (*self.mapping, ContentReplacement(content, replacement))
         return replace(self, mapping=new_mapping)
 
-    def with_matching(self, prefix: str = "", suffix: str = "") -> "SDOps":
-        """Create a new SDOps instance with the specified prefix and suffix matching added to the mapping."""
+    def with_matching(self, prefix: str = "", suffix: str = "", contains: str = "") -> "SDOps":
+        """Create a new SDOps instance with the specified prefix, suffix and contains matching added to the mapping."""
 
-        new_mapping = (*self.mapping, ContentMatching(prefix, suffix))
+        new_mapping = (*self.mapping, ContentMatching(prefix, suffix, contains))
         return replace(self, mapping=new_mapping)
 
     def with_additional_allowed_keys(self, keys: frozenset[str]) -> "SDOps":
@@ -100,7 +101,12 @@ class SDOps:
     def apply_to_key(self, key: str) -> str | None:
         """Apply the mapping to the given name."""
         matchers = [content for content in self.mapping if isinstance(content, ContentMatching)]
-        valid = any(key.startswith(f.prefix) and key.endswith(f.suffix) for f in matchers)
+        valid = any(
+            key.startswith(matcher.prefix)
+            and key.endswith(matcher.suffix)
+            and (not matcher.contains or matcher.contains in key)
+            for matcher in matchers
+        )
         if not valid:
             return None
 

packages/ltx-core/src/ltx_core/loader/single_gpu_model_builder.py

@@ -1,10 +1,14 @@
+from __future__ import annotations
+
+import copy
 import logging
-from dataclasses import dataclass, field, replace
-from typing import Generic
+from typing import TYPE_CHECKING, Final, Generic
 
 import torch
+from torch import nn
 
-from ltx_core.loader.fuse_loras import apply_loras
+from ltx_core.loader.fuse_loras import FuseRule, apply_loras, bf16_fuse_rule
+from ltx_core.loader.helpers import create_meta_model, load_state_dict, read_model_config
 from ltx_core.loader.module_ops import ModuleOps
 from ltx_core.loader.primitives import (
     LoRAAdaptableProtocol,
@@ -19,13 +23,72 @@ from ltx_core.loader.sd_ops import SDOps
 from ltx_core.loader.sft_loader import SafetensorsModelStateDictLoader
 from ltx_core.model.model_protocol import ModelConfigurator, ModelType
 
+if TYPE_CHECKING:
+    from typing_extensions import Self
+
 logger: logging.Logger = logging.getLogger(__name__)
 
 
-@dataclass(frozen=True)
+def _check_uninitialized(model: nn.Module) -> list[str]:
+    """Return names of any parameters/buffers still on meta device."""
+    names = []
+    for name, param in model.named_parameters():
+        if str(param.device) == "meta":
+            names.append(name)
+    for name, buf in model.named_buffers():
+        if str(buf.device) == "meta":
+            names.append(name)
+    return names
+
+
+def _load_model_weights(
+    meta_model: nn.Module,
+    model_path: str | tuple[str, ...],
+    loras: tuple[LoraPathStrengthAndSDOps, ...],
+    loader: StateDictLoader,
+    registry: Registry,
+    device: torch.device,
+    dtype: torch.dtype | None,
+    model_sd_ops: SDOps | None = None,
+    lora_load_device: torch.device | None = None,
+    fuse_rule: FuseRule = bf16_fuse_rule,
+) -> None:
+    """Load base weights and fuse LoRAs into *meta_model* in-place."""
+    if lora_load_device is None:
+        lora_load_device = device
+
+    model_sd = load_state_dict(model_path, loader, registry, device, model_sd_ops)
+
+    lora_strengths = [lora.strength for lora in loras]
+    if not lora_strengths or (min(lora_strengths) == 0 and max(lora_strengths) == 0):
+        sd = model_sd.sd
+        if dtype is not None:
+            sd = {key: value.to(dtype=dtype) for key, value in sd.items()}
+        meta_model.load_state_dict(sd, strict=False, assign=True)
+        return
+
+    lora_state_dicts = [load_state_dict([lora.path], loader, registry, lora_load_device, lora.sd_ops) for lora in loras]
+    lora_sd_and_strengths = [
+        LoraStateDictWithStrength(sd, strength) for sd, strength in zip(lora_state_dicts, lora_strengths, strict=True)
+    ]
+    final_sd = apply_loras(
+        model_sd=model_sd,
+        lora_sd_and_strengths=lora_sd_and_strengths,
+        fuse_rule=fuse_rule,
+        destination_sd=model_sd if isinstance(registry, DummyRegistry) else None,
+    )
+    fused_sd = final_sd.sd
+    if dtype is not None:
+        fused_sd = {key: value.to(dtype=dtype) for key, value in fused_sd.items()}
+    meta_model.load_state_dict(fused_sd, strict=False, assign=True)
+
+
 class SingleGPUModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType], LoRAAdaptableProtocol):
     """
     Builder for PyTorch models residing on a single GPU.
+    The builder is immutable: ``with_*``/``lora`` return modified copies. The
+    ``ModelBuilderProtocol`` surface is exposed via read-only properties backed
+    by private attributes.
     Attributes:
         model_class_configurator: Class responsible for constructing the model from a config dict.
         model_path: Path (or tuple of shard paths) to the model's `.safetensors` checkpoint(s).
@@ -39,64 +102,109 @@ class SingleGPUModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType],
             ``torch.device("cpu")``, which keeps LoRA weights in CPU memory and transfers them to
             the target GPU sequentially during fusion, reducing peak GPU memory usage compared to
             loading all LoRA weights directly onto the GPU at once.
+        fuse_rule: Per-policy LoRA merge rule. Defaults to ``bf16_fuse_rule``;
     """
 
-    model_class_configurator: type[ModelConfigurator[ModelType]]
-    model_path: str | tuple[str, ...]
-    model_sd_ops: SDOps | None = None
-    module_ops: tuple[ModuleOps, ...] = field(default_factory=tuple)
-    loras: tuple[LoraPathStrengthAndSDOps, ...] = field(default_factory=tuple)
-    model_loader: StateDictLoader = field(default_factory=SafetensorsModelStateDictLoader)
-    registry: Registry = field(default_factory=DummyRegistry)
-    lora_load_device: torch.device = field(default_factory=lambda: torch.device("cpu"))
+    def __init__(
+        self,
+        model_class_configurator: type[ModelConfigurator[ModelType]],
+        model_path: str | tuple[str, ...],
+        model_sd_ops: SDOps | None = None,
+        module_ops: tuple[ModuleOps, ...] = (),
+        loras: tuple[LoraPathStrengthAndSDOps, ...] = (),
+        model_loader: StateDictLoader | None = None,
+        registry: Registry | None = None,
+        lora_load_device: torch.device | None = None,
+        fuse_rule: FuseRule = bf16_fuse_rule,
+    ) -> None:
+        # Read-only: typed with the covariant ModelType, so it must not be a mutable attribute.
+        self._model_class_configurator: Final = model_class_configurator
+        self._model_path = model_path
+        self._model_sd_ops = model_sd_ops
+        self._module_ops = module_ops
+        self._loras = loras
+        self._model_loader = model_loader if model_loader is not None else SafetensorsModelStateDictLoader()
+        self._registry = registry if registry is not None else DummyRegistry()
+        self._lora_load_device = lora_load_device if lora_load_device is not None else torch.device("cpu")
+        self._fuse_rule = fuse_rule
 
-    def lora(self, lora_path: str, strength: float = 1.0, sd_ops: SDOps | None = None) -> "SingleGPUModelBuilder":
-        return replace(self, loras=(*self.loras, LoraPathStrengthAndSDOps(lora_path, strength, sd_ops)))
+    @property
+    def model_sd_ops(self) -> SDOps | None:
+        return self._model_sd_ops
 
-    def with_sd_ops(self, sd_ops: SDOps | None) -> "SingleGPUModelBuilder":
-        return replace(self, model_sd_ops=sd_ops)
+    @property
+    def module_ops(self) -> tuple[ModuleOps, ...]:
+        return self._module_ops
 
-    def with_module_ops(self, module_ops: tuple[ModuleOps, ...]) -> "SingleGPUModelBuilder":
-        return replace(self, module_ops=module_ops)
+    @property
+    def loras(self) -> tuple[LoraPathStrengthAndSDOps, ...]:
+        return self._loras
 
-    def with_loras(self, loras: tuple[LoraPathStrengthAndSDOps, ...]) -> "SingleGPUModelBuilder":
-        return replace(self, loras=loras)
+    @property
+    def registry(self) -> Registry:
+        return self._registry
 
-    def with_registry(self, registry: Registry) -> "SingleGPUModelBuilder":
-        return replace(self, registry=registry)
+    @property
+    def model_path(self) -> str | tuple[str, ...]:
+        return self._model_path
 
-    def with_lora_load_device(self, device: torch.device) -> "SingleGPUModelBuilder":
-        return replace(self, lora_load_device=device)
+    @property
+    def model_loader(self) -> StateDictLoader:
+        return self._model_loader
+
+    @property
+    def lora_load_device(self) -> torch.device:
+        return self._lora_load_device
+
+    @property
+    def fuse_rule(self) -> FuseRule:
+        return self._fuse_rule
+
+    def lora(self, lora_path: str, strength: float, sd_ops: SDOps) -> Self:
+        clone = copy.copy(self)
+        clone._loras = (*self._loras, LoraPathStrengthAndSDOps(lora_path, strength, sd_ops))
+        return clone
+
+    def with_sd_ops(self, sd_ops: SDOps | None) -> Self:
+        clone = copy.copy(self)
+        clone._model_sd_ops = sd_ops
+        return clone
+
+    def with_module_ops(self, module_ops: tuple[ModuleOps, ...]) -> Self:
+        clone = copy.copy(self)
+        clone._module_ops = module_ops
+        return clone
+
+    def with_loras(self, loras: tuple[LoraPathStrengthAndSDOps, ...]) -> Self:
+        clone = copy.copy(self)
+        clone._loras = loras
+        return clone
+
+    def with_registry(self, registry: Registry) -> Self:
+        clone = copy.copy(self)
+        clone._registry = registry
+        return clone
+
+    def with_lora_load_device(self, device: torch.device) -> Self:
+        clone = copy.copy(self)
+        clone._lora_load_device = device
+        return clone
+
+    def with_fuse_rule(self, fuse_rule: FuseRule) -> Self:
+        clone = copy.copy(self)
+        clone._fuse_rule = fuse_rule
+        return clone
 
     def model_config(self) -> dict:
-        first_shard_path = self.model_path[0] if isinstance(self.model_path, tuple) else self.model_path
-        return self.model_loader.metadata(first_shard_path)
+        return read_model_config(self._model_path, self._model_loader)
 
     def meta_model(self, config: dict, module_ops: tuple[ModuleOps, ...]) -> ModelType:
-        with torch.device("meta"):
-            model = self.model_class_configurator.from_config(config)
-        for module_op in module_ops:
-            if module_op.matcher(model):
-                model = module_op.mutator(model)
-        return model
+        return create_meta_model(self._model_class_configurator, config, module_ops)
 
     def load_sd(
         self, paths: list[str], registry: Registry, device: torch.device | None, sd_ops: SDOps | None = None
     ) -> StateDict:
-        state_dict = registry.get(paths, sd_ops)
-        if state_dict is None:
-            state_dict = self.model_loader.load(paths, sd_ops=sd_ops, device=device)
-            registry.add(paths, sd_ops=sd_ops, state_dict=state_dict)
-        return state_dict
-
-    def _return_model(self, meta_model: ModelType, device: torch.device) -> ModelType:
-        uninitialized_params = [name for name, param in meta_model.named_parameters() if str(param.device) == "meta"]
-        uninitialized_buffers = [name for name, buffer in meta_model.named_buffers() if str(buffer.device) == "meta"]
-        if uninitialized_params or uninitialized_buffers:
-            logger.warning(f"Uninitialized parameters or buffers: {uninitialized_params + uninitialized_buffers}")
-            return meta_model
-        retval = meta_model.to(device)
-        return retval
+        return load_state_dict(paths, self._model_loader, registry, device, sd_ops)
 
     def build(
         self,
@@ -106,31 +214,23 @@ class SingleGPUModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType],
     ) -> ModelType:
         device = torch.device("cuda") if device is None else device
         config = self.model_config()
-        meta_model = self.meta_model(config, self.module_ops)
-        model_paths = list(self.model_path) if isinstance(self.model_path, tuple) else [self.model_path]
-        model_state_dict = self.load_sd(model_paths, sd_ops=self.model_sd_ops, registry=self.registry, device=device)
+        meta_model = self.meta_model(config, self._module_ops)
 
-        lora_strengths = [lora.strength for lora in self.loras]
-        if not lora_strengths or (min(lora_strengths) == 0 and max(lora_strengths) == 0):
-            sd = model_state_dict.sd
-            if dtype is not None:
-                sd = {key: value.to(dtype=dtype) for key, value in model_state_dict.sd.items()}
-            meta_model.load_state_dict(sd, strict=False, assign=True)
-            return self._return_model(meta_model, device)
-
-        lora_state_dicts = [
-            self.load_sd([lora.path], sd_ops=lora.sd_ops, registry=self.registry, device=self.lora_load_device)
-            for lora in self.loras
-        ]
-        lora_sd_and_strengths = [
-            LoraStateDictWithStrength(sd, strength)
-            for sd, strength in zip(lora_state_dicts, lora_strengths, strict=True)
-        ]
-        final_sd = apply_loras(
-            model_sd=model_state_dict,
-            lora_sd_and_strengths=lora_sd_and_strengths,
+        _load_model_weights(
+            meta_model=meta_model,
+            model_path=self._model_path,
+            loras=self._loras,
+            loader=self._model_loader,
+            registry=self._registry,
+            device=device,
             dtype=dtype,
-            destination_sd=model_state_dict if isinstance(self.registry, DummyRegistry) else None,
+            model_sd_ops=self._model_sd_ops,
+            lora_load_device=self._lora_load_device,
+            fuse_rule=self._fuse_rule,
         )
-        meta_model.load_state_dict(final_sd.sd, strict=False, assign=True)
-        return self._return_model(meta_model, device)
+
+        uninitialized = _check_uninitialized(meta_model)
+        if uninitialized:
+            logger.warning(f"Uninitialized parameters or buffers: {uninitialized}")
+            return meta_model
+        return meta_model.to(device)

packages/ltx-core/src/ltx_core/modality_tiling.py

@@ -21,11 +21,14 @@ from ltx_core.types import VideoLatentShape
 @dataclass(frozen=True)
 class TilingContext:
     """Opaque context produced by :meth:`VideoModalityTilingHelper.tile_modality`.
-    Carries the token-level keep mask and per-conditioning-token blend
+    Carries the token-level keep indices and per-conditioning-token blend
     weights needed by :meth:`~VideoModalityTilingHelper.blend`.
     """
 
-    keep_mask: torch.Tensor
+    keep_indices: torch.Tensor
+    """``(num_kept,)`` int64 — sorted indices of tokens the tile processes."""
+    num_total_tokens: int
+    """Total number of tokens in the full (untiled) sequence."""
     cond_blend_weights: torch.Tensor | None
     """``(num_kept_cond,)`` — weight for each kept conditioning token,
     equal to ``1 / num_tiles_that_keep_this_token``.  ``None`` when
@@ -81,14 +84,32 @@ class VideoModalityTilingHelper:
             A ``(tiled_modality, context)`` tuple.  Pass *context* to
             :meth:`blend` together with the model output.
         """
-        keep_mask = self._keep_mask(modality, tile)
+        device = modality.positions.device
+        gen_indices = self._generated_token_indices(tile, device=device)
+        num_total = modality.latent.shape[1]
+
+        cond_blend_weights: torch.Tensor | None = None
+        if num_total > self._num_generated_tokens:
+            keep_per_tile_cond = self._all_tiles_cond_keep(modality)  # (num_tiles, num_cond) bool
+            tile_idx = next((i for i, t in enumerate(self._tiles) if t.in_coords == tile.in_coords), None)
+            if tile_idx is None:
+                raise ValueError(
+                    f"Tile with in_coords={tile.in_coords} is not in this helper's tile set; "
+                    f"pass a tile obtained from `helper.tiles`."
+                )
+            my_cond_keep = keep_per_tile_cond[tile_idx]
+            cond_indices = self._num_generated_tokens + my_cond_keep.nonzero(as_tuple=False).squeeze(1)
+            keep_indices = torch.cat([gen_indices, cond_indices])
+            total_keepers = keep_per_tile_cond.sum(dim=0).float()  # (num_cond,)
+            cond_blend_weights = 1.0 / total_keepers[my_cond_keep]
+        else:
+            keep_indices = gen_indices
 
         tile_attention_mask = None
         if modality.attention_mask is not None:
-            keep_indices = keep_mask.nonzero(as_tuple=False).squeeze(1)
             tile_attention_mask = modality.attention_mask[:, keep_indices, :][:, :, keep_indices]
 
-        positions = modality.positions[:, :, keep_mask, :]
+        positions = modality.positions[:, :, keep_indices, :]
         if normalize_positions:
             num_tile_gen = self._tile_generated_token_count(tile)
             gen_pos = positions[:, :, :num_tile_gen, :]  # (B, 3, num_tile_gen, 2)
@@ -97,26 +118,15 @@ class VideoModalityTilingHelper:
 
         tiled = replace(
             modality,
-            latent=modality.latent[:, keep_mask, :],
-            timesteps=modality.timesteps[:, keep_mask],
+            latent=modality.latent[:, keep_indices, :],
+            timesteps=modality.timesteps[:, keep_indices],
             positions=positions,
             attention_mask=tile_attention_mask,
         )
 
-        cond_blend_weights = None
-        num_total = modality.latent.shape[1]
-        if num_total > self._num_generated_tokens:
-            cond_keep = keep_mask[self._num_generated_tokens :]
-            # Count how many tiles keep each conditioning token.
-            cond_counts = torch.zeros(cond_keep.sum(), dtype=torch.float32)
-            for t in self._tiles:
-                other_mask = self._keep_mask(modality, t)
-                other_cond = other_mask[self._num_generated_tokens :]
-                # Map other tile's kept cond tokens into this tile's kept subset.
-                cond_counts += other_cond[cond_keep].float()
-            cond_blend_weights = 1.0 / cond_counts
-
-        return tiled, TilingContext(keep_mask=keep_mask, cond_blend_weights=cond_blend_weights)
+        return tiled, TilingContext(
+            keep_indices=keep_indices, num_total_tokens=num_total, cond_blend_weights=cond_blend_weights
+        )
 
     # -- blend -------------------------------------------------------------
 
@@ -147,9 +157,9 @@ class VideoModalityTilingHelper:
         """
         batch, _, dim = tile_to_blend.shape
         num_tile_gen = self._tile_generated_token_count(tile)
-        gen_indices = self._generated_token_indices(tile)
+        gen_indices = self._generated_token_indices(tile, device=tile_to_blend.device)
 
-        num_total_tokens = context.keep_mask.shape[0]
+        num_total_tokens = context.num_total_tokens
         expected_shape = (batch, num_total_tokens, dim)
 
         if output is not None:
@@ -168,8 +178,7 @@ class VideoModalityTilingHelper:
         # Scatter kept conditioning tokens, weighted by 1/N where N is
         # the number of tiles that keep each token (so they sum to 1).
         if num_total_tokens > self._num_generated_tokens and context.cond_blend_weights is not None:
-            cond_keep = context.keep_mask[self._num_generated_tokens :]
-            cond_indices = self._num_generated_tokens + cond_keep.nonzero(as_tuple=False).squeeze(1)
+            cond_indices = context.keep_indices[context.keep_indices >= self._num_generated_tokens]
             weights = context.cond_blend_weights.to(device=tile_to_blend.device, dtype=tile_to_blend.dtype)
             result[:, cond_indices, :] += tile_to_blend[:, num_tile_gen:, :] * weights[None, :, None]
 
@@ -189,46 +198,42 @@ class VideoModalityTilingHelper:
         )
         return self._patchifier.get_token_count(tile_shape)
 
-    def _generated_token_indices(self, tile: Tile) -> torch.Tensor:
+    def _generated_token_indices(self, tile: Tile, device: torch.device | None = None) -> torch.Tensor:
         """Flat token indices of *tile*'s generated tokens in the full sequence."""
         frame_slice, height_slice, width_slice = tile.in_coords
-        f = torch.arange(frame_slice.start, frame_slice.stop)
-        h = torch.arange(height_slice.start, height_slice.stop)
-        w = torch.arange(width_slice.start, width_slice.stop)
+        f = torch.arange(frame_slice.start, frame_slice.stop, device=device)
+        h = torch.arange(height_slice.start, height_slice.stop, device=device)
+        w = torch.arange(width_slice.start, width_slice.stop, device=device)
         return (
             f[:, None, None] * self._latent_shape.height * self._latent_shape.width
             + h[None, :, None] * self._latent_shape.width
             + w[None, None, :]
         ).reshape(-1)
 
-    def _keep_mask(self, modality: Modality, tile: Tile) -> torch.Tensor:
-        """Boolean mask ``(num_total_tokens,)`` — True for tokens the tile processes.
-        Generated tokens are selected by grid position.  Conditioning
-        tokens are kept when their ``[start, end)`` intervals overlap
-        the tile in all three dimensions, or when they have a negative
-        time coordinate (reference tokens).
+    def _all_tiles_cond_keep(self, modality: Modality) -> torch.Tensor:
+        """Vectorized (num_tiles, num_cond) bool: which tiles keep each conditioning token.
+        A conditioning token is kept by a tile when its ``[start, end)`` interval
+        overlaps the tile in all three dimensions, or when it has a negative time
+        coordinate (reference token).
         """
-        num_total = modality.latent.shape[1]
-        mask = torch.zeros(num_total, dtype=torch.bool)
+        cond_positions = modality.positions[:, :, self._num_generated_tokens :, :]  # (B, 3, num_cond, 2)
+        device = cond_positions.device
 
-        gen_indices = self._generated_token_indices(tile)
-        mask[gen_indices] = True
+        # Per-tile (start, end) bounds along each axis; small Python loop (num_tiles <= ~16).
+        starts_list: list[torch.Tensor] = []
+        ends_list: list[torch.Tensor] = []
+        for t in self._tiles:
+            gen_idx = self._generated_token_indices(t, device=device)
+            gen_positions = modality.positions[:, :, gen_idx, :]  # (B, 3, num_tile_gen, 2)
+            starts_list.append(gen_positions[..., 0].amin(dim=2))  # (B, 3)
+            ends_list.append(gen_positions[..., 1].amax(dim=2))  # (B, 3)
+        tile_starts = torch.stack(starts_list, dim=0)  # (num_tiles, B, 3)
+        tile_ends = torch.stack(ends_list, dim=0)  # (num_tiles, B, 3)
 
-        if num_total > self._num_generated_tokens:
-            gen_positions = modality.positions[:, :, gen_indices, :]  # (B, 3, num_tile_gen, 2)
-            tile_start = gen_positions[..., 0].amin(dim=2)  # (B, 3)
-            tile_end = gen_positions[..., 1].amax(dim=2)  # (B, 3)
-
-            cond_positions = modality.positions[:, :, self._num_generated_tokens :, :]  # (B, 3, num_cond, 2)
-
-            overlaps = (cond_positions[..., 0] < tile_end.unsqueeze(2)) & (
-                cond_positions[..., 1] > tile_start.unsqueeze(2)
-            )  # (B, 3, num_cond)
-            overlaps_all_dims = overlaps.all(dim=1)  # (B, num_cond)
-
-            has_negative_time = cond_positions[:, 0, :, 0] < 0  # (B, num_cond)
-
-            keep_cond = (overlaps_all_dims | has_negative_time).any(dim=0)  # (num_cond,)
-            mask[self._num_generated_tokens :] = keep_cond
-
-        return mask
+        cond_starts = cond_positions[..., 0]  # (B, 3, num_cond)
+        cond_ends = cond_positions[..., 1]  # (B, 3, num_cond)
+        # Broadcast: (1, B, 3, num_cond) vs (num_tiles, B, 3, 1) -> (num_tiles, B, 3, num_cond).
+        overlaps = (cond_starts[None] < tile_ends[..., None]) & (cond_ends[None] > tile_starts[..., None])
+        overlaps_all_dims = overlaps.all(dim=2)  # (num_tiles, B, num_cond)
+        has_negative_time = (cond_positions[:, 0, :, 0] < 0)[None]  # (1, B, num_cond)
+        return (overlaps_all_dims | has_negative_time).any(dim=1)  # (num_tiles, num_cond)

packages/ltx-core/src/ltx_core/model/model_protocol.py

@@ -1,6 +1,14 @@
-from typing import Protocol, TypeVar
+from __future__ import annotations
 
-ModelType = TypeVar("ModelType")
+from typing import TYPE_CHECKING, Protocol, TypeVar
+
+import torch
+
+if TYPE_CHECKING:
+    from ltx_core.guidance.perturbations import BatchedPerturbationConfig
+    from ltx_core.model.transformer.modality import Modality
+
+ModelType = TypeVar("ModelType", covariant=True, bound=torch.nn.Module)  # noqa: PLC0105
 
 
 class ModelConfigurator(Protocol[ModelType]):
@@ -8,3 +16,24 @@ class ModelConfigurator(Protocol[ModelType]):
 
     @classmethod
     def from_config(cls, config: dict) -> ModelType: ...
+
+
+class LTXModelProtocol(Protocol):
+    """Velocity-model forward interface shared by ``LTXModel`` and its multi-GPU wrappers.
+    ``forward`` pins the real signature (enforced structurally); ``__call__`` mirrors it
+    so protocol-typed values stay callable via ``model(...)``.
+    """
+
+    def forward(
+        self,
+        video: Modality | None,
+        audio: Modality | None,
+        perturbations: BatchedPerturbationConfig,
+    ) -> tuple[torch.Tensor | None, torch.Tensor | None]: ...
+
+    def __call__(
+        self,
+        video: Modality | None,
+        audio: Modality | None,
+        perturbations: BatchedPerturbationConfig,
+    ) -> tuple[torch.Tensor | None, torch.Tensor | None]: ...

packages/ltx-core/src/ltx_core/model/transformer/__init__.py

@@ -3,13 +3,17 @@
 from ltx_core.model.transformer.modality import Modality
 from ltx_core.model.transformer.model import LTXModel, X0Model
 from ltx_core.model.transformer.model_configurator import (
+    LTXV_AUDIO_ONLY_MODEL_COMFY_RENAMING_MAP,
     LTXV_MODEL_COMFY_RENAMING_MAP,
+    LTXAudioOnlyModelConfigurator,
     LTXModelConfigurator,
     LTXVideoOnlyModelConfigurator,
 )
 
 __all__ = [
+    "LTXV_AUDIO_ONLY_MODEL_COMFY_RENAMING_MAP",
     "LTXV_MODEL_COMFY_RENAMING_MAP",
+    "LTXAudioOnlyModelConfigurator",
     "LTXModel",
     "LTXModelConfigurator",
     "LTXVideoOnlyModelConfigurator",

packages/ltx-core/src/ltx_core/model/transformer/attention.py

@@ -1,12 +1,35 @@
+import functools
+from dataclasses import dataclass, field
 from enum import Enum
 from typing import Protocol
 
 import torch
+from torch.nn.attention import SDPBackend, sdpa_kernel
+
+from ltx_core.model.transformer.ops import (
+    GatedAttentionCallable,
+    PreAttentionCallable,
+    PytorchGatedAttention,
+    PytorchPreAttention,
+)
+from ltx_core.model.transformer.rope import LTXRopeType
+
+
+def _torch_default_sdpa_priority() -> list[SDPBackend]:
+    """Fetch torch's current default SDPA priority order at runtime.
+    Used as the default for ``PytorchAttention`` so the wrapper-always
+    code path matches torch's native dispatch order without hard-coding it
+    (which would drift if torch updates the default).
+    ``torch._C._get_sdp_priority_order`` is a private API; we accept that
+    risk because the project pins ``torch`` in the lockfile, so any
+    rename/removal surfaces on a controlled torch bump rather than silently.
+    """
+    return [SDPBackend(p) for p in torch._C._get_sdp_priority_order()]
 
-from ltx_core.model.transformer.rope import LTXRopeType, apply_rotary_emb
 
 memory_efficient_attention = None
 flash_attn_interface = None
+flash_attn_4_func = None
 try:
     from xformers.ops import memory_efficient_attention
 except ImportError:
@@ -17,15 +40,44 @@ try:
         import flash_attn_interface
 except ImportError:
     flash_attn_interface = None
+try:
+    from flash_attn.cute import flash_attn_func as flash_attn_4_func
+except ImportError:
+    flash_attn_4_func = None
 
 
 class AttentionCallable(Protocol):
+    """Unmasked attention. Backends without a mask kernel (FA3/FA4) implement only
+    this protocol; backends that support masks too (Pytorch/SDPA, xFormers) are
+    structurally usable here and as :class:`MaskedAttentionCallable`."""
+
+    def __call__(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, heads: int) -> torch.Tensor: ...
+
+
+class MaskedAttentionCallable(Protocol):
+    """Masked attention. Mask is required (not optional) -- the caller has already
+    decided this is the masked path and chosen a backend that can serve it. Used
+    by :class:`Attention` when its forward receives a non-None ``mask``."""
+
     def __call__(
-        self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, heads: int, mask: torch.Tensor | None = None
+        self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, heads: int, mask: torch.Tensor
     ) -> torch.Tensor: ...
 
 
 class PytorchAttention(AttentionCallable):
+    def __init__(self, priority: list[SDPBackend] | None = None) -> None:
+        # priority=None -> snapshot torch's default SDPA priority at construction.
+        # Always passed through ``sdpa_kernel(..., set_priority=True)`` so the
+        # call site is uniform regardless of how the priority was chosen.
+        self._priority = priority if priority is not None else _torch_default_sdpa_priority()
+
+    @property
+    def label(self) -> str:
+        """Human-readable identifier for this backend. Encodes the SDPA priority
+        list so a single-backend pin reads differently from the full-priority
+        dispatcher walk."""
+        return f"SDPA[{'>'.join(b.name for b in self._priority)}]"
+
     def __call__(
         self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, heads: int, mask: torch.Tensor | None = None
     ) -> torch.Tensor:
@@ -41,12 +93,17 @@ class PytorchAttention(AttentionCallable):
             if mask.ndim == 3:
                 mask = mask.unsqueeze(1)
 
-        out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False)
+        with sdpa_kernel(self._priority, set_priority=True):
+            out = torch.nn.functional.scaled_dot_product_attention(
+                q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False
+            )
         out = out.transpose(1, 2).reshape(b, -1, heads * dim_head)
         return out
 
 
 class XFormersAttention(AttentionCallable):
+    label = "xFormers"
+
     def __call__(
         self,
         q: torch.Tensor,
@@ -92,13 +149,14 @@ class XFormersAttention(AttentionCallable):
 
 
 class FlashAttention3(AttentionCallable):
+    label = "FlashAttention3"
+
     def __call__(
         self,
         q: torch.Tensor,
         k: torch.Tensor,
         v: torch.Tensor,
         heads: int,
-        mask: torch.Tensor | None = None,
     ) -> torch.Tensor:
         if flash_attn_interface is None:
             raise RuntimeError("FlashAttention3 was selected but `FlashAttention3` is not installed.")
@@ -108,32 +166,276 @@ class FlashAttention3(AttentionCallable):
 
         q, k, v = (t.view(b, -1, heads, dim_head) for t in (q, k, v))
 
-        if mask is not None:
-            raise NotImplementedError("Mask is not supported for FlashAttention3")
-
         out = flash_attn_interface.flash_attn_func(q.to(v.dtype), k.to(v.dtype), v)
         out = out.reshape(b, -1, heads * dim_head)
         return out
 
 
+class FlashAttention4(AttentionCallable):
+    label = "FlashAttention4"
+
+    def __call__(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        heads: int,
+    ) -> torch.Tensor:
+        if flash_attn_4_func is None:
+            raise RuntimeError("FlashAttention4 was selected but `flash-attn-4` is not installed.")
+
+        b, _, dim_head = q.shape
+        dim_head //= heads
+
+        q, k, v = (t.view(b, -1, heads, dim_head) for t in (q, k, v))
+
+        out, _ = flash_attn_4_func(q.to(v.dtype), k.to(v.dtype), v)
+        out = out.reshape(b, -1, heads * dim_head)
+        return out
+
+
+# --- Automatic selection -----------------------------------------------------
+# AUTOMATIC inspects installed extras and the GPU arch and returns the fastest
+# usable callable for each path. The selection runs once per process (cached).
+# The unmasked and masked picks are independent: each calls its own helper and
+# may end up on different backends (e.g. FA3 unmasked + xFormers masked on H100).
+
+
+def _sdpa_can_use(backend: SDPBackend, *, with_mask: bool) -> bool:
+    """Ask torch whether *backend* can run with the given mask shape.
+    ``MATH`` is the universal SDPA fallback (pure PyTorch ops, no kernel
+    requirements) so it returns True everywhere, CPU included. The other
+    backends use ``torch.backends.cuda.can_use_*`` capability checks (no GPU
+    compute, no synchronization) and are False without CUDA. The probe shapes
+    are small but realistic enough to surface constraints (head dim, dtype)
+    that the per-backend rules care about.
+    """
+    if backend is SDPBackend.MATH:
+        return True
+    if not torch.cuda.is_available():
+        return False
+    q = torch.empty(1, 4, 128, 64, device="cuda", dtype=torch.bfloat16)
+    k = torch.empty(1, 4, 128, 64, device="cuda", dtype=torch.bfloat16)
+    v = torch.empty(1, 4, 128, 64, device="cuda", dtype=torch.bfloat16)
+    mask = torch.zeros(1, 4, 128, 128, device="cuda", dtype=torch.bfloat16) if with_mask else None
+    params = torch.backends.cuda.SDPAParams(q, k, v, mask, 0.0, False, False)
+    if backend is SDPBackend.CUDNN_ATTENTION:
+        return torch.backends.cuda.can_use_cudnn_attention(params, debug=False)
+    if backend is SDPBackend.FLASH_ATTENTION:
+        return torch.backends.cuda.can_use_flash_attention(params, debug=False)
+    if backend is SDPBackend.EFFICIENT_ATTENTION:
+        return torch.backends.cuda.can_use_efficient_attention(params, debug=False)
+    return False
+
+
+_SDPA_FULL_PRIORITY: tuple[SDPBackend, ...] = (
+    SDPBackend.CUDNN_ATTENTION,
+    SDPBackend.FLASH_ATTENTION,
+    SDPBackend.EFFICIENT_ATTENTION,
+    SDPBackend.MATH,
+)
+
+
+def _sdpa_full_priority() -> PytorchAttention:
+    """Hand SDPA the full backend priority order; let torch's dispatcher pick at call time.
+    ``sdpa_kernel(_SDPA_FULL_PRIORITY, set_priority=True)`` enables all four
+    backends and orders them; torch then walks the order at call time and picks
+    the first backend whose ``can_use_*`` check passes for the actual
+    shapes/dtype/mask. FLASH is rejected automatically when a mask is present;
+    CUDNN may be rejected under deterministic mode; MATH is the universal
+    fallback. Probing per-backend usability up front from generic probe shapes
+    cannot anticipate the variety of real call sites (e.g. broadcast key-only
+    masks, large head dim), so we defer the choice to the dispatcher.
+    """
+    return PytorchAttention(priority=list(_SDPA_FULL_PRIORITY))
+
+
+def _select_primary_attention() -> AttentionCallable:
+    """Pick the fastest unmasked attention based on installed extras and GPU arch.
+    Priority by arch:
+    - Hopper (sm_90, H100): FA3 / xFormers (mutually exclusive at import) > FA4 > SDPA.
+    - Datacenter Blackwell (sm_100, B200): FA4 > SDPA. FA4 is intentionally *not*
+      picked on consumer Blackwell (sm_120) -- known regressions in newer
+      FA4 betas; users who want it on sm_120 must opt in explicitly.
+    - Everywhere else (Ada, Ampere, CPU): SDPA with the full backend priority
+      list -- torch's runtime dispatcher picks the best fit at call time.
+    """
+    if torch.cuda.is_available():
+        major, _ = torch.cuda.get_device_capability(0)
+        if major == 9:
+            if flash_attn_interface is not None:
+                return FlashAttention3()
+            if memory_efficient_attention is not None:
+                return XFormersAttention()
+            if flash_attn_4_func is not None:
+                return FlashAttention4()
+        if major == 10 and flash_attn_4_func is not None:
+            return FlashAttention4()
+    return _sdpa_full_priority()
+
+
+def _select_masked_attention() -> MaskedAttentionCallable:
+    """Pick a mask-aware attention. Prefers xFormers when installed; else SDPA with
+    the full priority list (the dispatcher rejects FLASH automatically when a
+    mask is present and walks past it)."""
+    if memory_efficient_attention is not None:
+        return XFormersAttention()
+    return _sdpa_full_priority()
+
+
+@functools.cache
+def automatic_attention() -> AttentionCallable:
+    """Cached AUTOMATIC pick for the unmasked path.
+    Cached so every ``AttentionOps`` in the process shares one instance."""
+    return _select_primary_attention()
+
+
+@functools.cache
+def automatic_masked_attention() -> MaskedAttentionCallable:
+    """Cached AUTOMATIC pick for the masked path. See :func:`automatic_attention`."""
+    return _select_masked_attention()
+
+
+def attention_label(fn: AttentionCallable | MaskedAttentionCallable) -> str:
+    """Best-effort human-readable backend name.
+    Built-in callables expose ``.label`` (encoding the SDPA priority list for the
+    Pytorch backends); fall back to the class name for custom or wrapped callables
+    (e.g. the multi-GPU All2All wrappers) that don't define one."""
+    return getattr(fn, "label", type(fn).__name__)
+
+
+def _resolve_sdpa_variant(backend: SDPBackend, name: str, *, with_mask: bool) -> PytorchAttention:
+    """Build a single-backend ``PytorchAttention`` pin, raising if the backend
+    can't actually serve the call on this machine. Used by both
+    :meth:`AttentionFunction.to_callable` and :meth:`MaskedAttentionFunction.to_callable`;
+    ``with_mask`` differs between the two so the capability check considers
+    the protocol the caller intends to use. Not used for ``MATH`` -- MATH is
+    the universal fallback and would falsely fail the CUDA-only probe on CPU.
+    """
+    if not _sdpa_can_use(backend, with_mask=with_mask):
+        raise RuntimeError(
+            f"{name} selected but the SDPA {backend.name} backend is not usable on this machine "
+            "(either no CUDA, the backend rejected the probe shapes, or "
+            "torch.use_deterministic_algorithms(True) excluded it)."
+        )
+    return PytorchAttention(priority=[backend])
+
+
 class AttentionFunction(Enum):
     PYTORCH = "pytorch"
     XFORMERS = "xformers"
     FLASH_ATTENTION_3 = "flash_attention_3"
-    DEFAULT = "default"
+    FLASH_ATTENTION_4 = "flash_attention_4"
+    SDPA_CUDNN = "sdpa_cudnn"
+    SDPA_FLASH = "sdpa_flash"
+    SDPA_EFFICIENT = "sdpa_efficient"
+    SDPA_MATH = "sdpa_math"
+    # Pick the fastest unmasked backend for the current GPU/extras combo; see
+    # :func:`automatic_attention`. Default for :class:`AttentionOps`.
+    AUTOMATIC = "automatic"
 
-    def to_callable(self) -> AttentionCallable:
+    def to_callable(self) -> AttentionCallable:  # noqa: PLR0911
         """Resolve to a concrete callable. Use this at module init time so that
-        torch.compile can trace through the attention call without graph breaks."""
-        if self is AttentionFunction.PYTORCH:
-            return PytorchAttention()
-        elif self is AttentionFunction.XFORMERS:
-            return XFormersAttention()
-        elif self is AttentionFunction.FLASH_ATTENTION_3:
-            return FlashAttention3()
-        else:
-            # Default behavior: XFormers if installed else - PyTorch
-            return XFormersAttention() if memory_efficient_attention is not None else PytorchAttention()
+        torch.compile can trace through the attention call without graph breaks.
+        Every non-AUTOMATIC variant raises :class:`RuntimeError` when the backend
+        isn't usable on this machine -- missing package or SDPA backend rejected
+        on this hardware (e.g. cuDNN under ``torch.use_deterministic_algorithms``).
+        Opting in means "this kernel or fail loudly". ``AUTOMATIC`` returns the
+        cached :func:`automatic_attention` instance so every build shares one callable.
+        """
+        match self:
+            case AttentionFunction.AUTOMATIC:
+                return automatic_attention()
+            case AttentionFunction.PYTORCH:
+                return PytorchAttention()
+            case AttentionFunction.XFORMERS:
+                if memory_efficient_attention is None:
+                    raise RuntimeError("AttentionFunction.XFORMERS selected but `xformers` is not installed.")
+                return XFormersAttention()
+            case AttentionFunction.FLASH_ATTENTION_3:
+                if flash_attn_interface is None:
+                    raise RuntimeError(
+                        "AttentionFunction.FLASH_ATTENTION_3 selected but `flash-attn-3` is not installed."
+                    )
+                return FlashAttention3()
+            case AttentionFunction.FLASH_ATTENTION_4:
+                if flash_attn_4_func is None:
+                    raise RuntimeError(
+                        "AttentionFunction.FLASH_ATTENTION_4 selected but `flash-attn-4` is not installed."
+                    )
+                return FlashAttention4()
+            case AttentionFunction.SDPA_MATH:
+                return PytorchAttention(priority=[SDPBackend.MATH])
+            case AttentionFunction.SDPA_CUDNN:
+                return _resolve_sdpa_variant(
+                    SDPBackend.CUDNN_ATTENTION, "AttentionFunction.SDPA_CUDNN", with_mask=False
+                )
+            case AttentionFunction.SDPA_FLASH:
+                return _resolve_sdpa_variant(
+                    SDPBackend.FLASH_ATTENTION, "AttentionFunction.SDPA_FLASH", with_mask=False
+                )
+            case AttentionFunction.SDPA_EFFICIENT:
+                return _resolve_sdpa_variant(
+                    SDPBackend.EFFICIENT_ATTENTION, "AttentionFunction.SDPA_EFFICIENT", with_mask=False
+                )
+
+
+class MaskedAttentionFunction(Enum):
+    """Backends usable on the masked path. Mirrors :class:`AttentionFunction` minus
+    the variants the torch SDPA dispatcher (or the wrapped kernel) rejects with a
+    mask: ``SDPA_FLASH`` -- FLASH kernel cannot serve an additive ``attn_mask``;
+    ``FLASH_ATTENTION_3``/``FLASH_ATTENTION_4`` -- neither has a mask kernel at all.
+    Keeping them out makes "this backend cannot mask" a type error, not a runtime one."""
+
+    PYTORCH = "pytorch"
+    XFORMERS = "xformers"
+    SDPA_CUDNN = "sdpa_cudnn"
+    SDPA_EFFICIENT = "sdpa_efficient"
+    SDPA_MATH = "sdpa_math"
+    # Pick the fastest mask-capable backend for the current extras combo; see
+    # :func:`automatic_masked_attention`. Default for the masked slot of
+    # :class:`AttentionOps`.
+    AUTOMATIC = "automatic"
+
+    def to_callable(self) -> MaskedAttentionCallable:
+        """Resolve to a concrete masked callable. Same backend classes as
+        :meth:`AttentionFunction.to_callable`; the protocol returned just exposes
+        the masked call signature.
+        Non-AUTOMATIC variants raise :class:`RuntimeError` when the backend isn't
+        usable for the masked path on this machine. SDPA probes run with
+        ``with_mask=True`` so the capability check considers the protocol the
+        caller will actually use."""
+        match self:
+            case MaskedAttentionFunction.AUTOMATIC:
+                return automatic_masked_attention()
+            case MaskedAttentionFunction.PYTORCH:
+                return PytorchAttention()
+            case MaskedAttentionFunction.XFORMERS:
+                if memory_efficient_attention is None:
+                    raise RuntimeError("MaskedAttentionFunction.XFORMERS selected but `xformers` is not installed.")
+                return XFormersAttention()
+            case MaskedAttentionFunction.SDPA_MATH:
+                return PytorchAttention(priority=[SDPBackend.MATH])
+            case MaskedAttentionFunction.SDPA_CUDNN:
+                return _resolve_sdpa_variant(
+                    SDPBackend.CUDNN_ATTENTION, "MaskedAttentionFunction.SDPA_CUDNN", with_mask=True
+                )
+            case MaskedAttentionFunction.SDPA_EFFICIENT:
+                return _resolve_sdpa_variant(
+                    SDPBackend.EFFICIENT_ATTENTION, "MaskedAttentionFunction.SDPA_EFFICIENT", with_mask=True
+                )
+
+
+@dataclass(frozen=True)
+class AttentionOps:
+    """Pluggable callables consumed by :class:`Attention`."""
+
+    attention_function: AttentionCallable = field(default_factory=lambda: AttentionFunction.AUTOMATIC.to_callable())
+    masked_attention_function: MaskedAttentionCallable = field(
+        default_factory=lambda: MaskedAttentionFunction.AUTOMATIC.to_callable()
+    )
+    preattention_function: PreAttentionCallable = field(default_factory=PytorchPreAttention)
+    gated_attention_function: GatedAttentionCallable = field(default_factory=PytorchGatedAttention)
 
 
 class Attention(torch.nn.Module):
@@ -144,17 +446,18 @@ class Attention(torch.nn.Module):
         heads: int = 8,
         dim_head: int = 64,
         norm_eps: float = 1e-6,
-        rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
-        attention_function: AttentionCallable | AttentionFunction = AttentionFunction.DEFAULT,
+        rope_type: LTXRopeType = LTXRopeType.SPLIT,
+        ops: AttentionOps | None = None,
         apply_gated_attention: bool = False,
     ) -> None:
         super().__init__()
+        if ops is None:
+            ops = AttentionOps()
         self.rope_type = rope_type
-        self.attention_function = (
-            attention_function.to_callable()
-            if isinstance(attention_function, AttentionFunction)
-            else attention_function
-        )
+        self.attention_function = ops.attention_function
+        self.masked_attention_function = ops.masked_attention_function
+        self.preattention_function = ops.preattention_function
+        self.gated_attention_function = ops.gated_attention_function
 
         inner_dim = dim_head * heads
         context_dim = query_dim if context_dim is None else context_dim
@@ -196,7 +499,9 @@ class Attention(torch.nn.Module):
             context: Key/value context tensor of shape ``(B, S, context_dim)``.
                 Falls back to ``x`` (self-attention) when *None*.
             mask: Optional attention mask. Interpretation depends on the attention
-                backend (additive bias for xformers/PyTorch SDPA).
+                backend (additive bias for xformers/PyTorch SDPA). A non-None
+                ``mask`` routes to ``masked_attention_function``; ``None`` keeps
+                the unmasked path.
             pe: Rotary positional embeddings applied to both ``q`` and ``k``.
             k_pe: Separate rotary positional embeddings for ``k`` only. When
                 *None*, ``pe`` is reused for keys.
@@ -221,29 +526,17 @@ class Attention(torch.nn.Module):
         else:
             q = self.to_q(x)
             k = self.to_k(context)
-
-            q = self.q_norm(q)
-            k = self.k_norm(k)
-
-            if pe is not None:
-                q = apply_rotary_emb(q, pe, self.rope_type)
-                k = apply_rotary_emb(k, pe if k_pe is None else k_pe, self.rope_type)
-
-            out = self.attention_function(q, k, v, self.heads, mask)  # (B, T, H*D)
+            q, k = self.preattention_function(q, k, self, mask, pe, k_pe)
+            if mask is None:
+                out = self.attention_function(q, k, v, self.heads)  # (B, T, H*D)
+            else:
+                out = self.masked_attention_function(q, k, v, self.heads, mask)
 
             if perturbation_mask is not None:
                 out = out * perturbation_mask + v * (1 - perturbation_mask)
 
         # Apply per-head gating if enabled
         if self.to_gate_logits is not None:
-            gate_logits = self.to_gate_logits(x)  # (B, T, H)
-            b, t, _ = out.shape
-            # Reshape to (B, T, H, D) for per-head gating
-            out = out.view(b, t, self.heads, self.dim_head)
-            # Apply gating: 2 * sigmoid(x) so that zero-init gives identity (2 * 0.5 = 1.0)
-            gates = 2.0 * torch.sigmoid(gate_logits)  # (B, T, H)
-            out = out * gates.unsqueeze(-1)  # (B, T, H, D) * (B, T, H, 1)
-            # Reshape back to (B, T, H*D)
-            out = out.view(b, t, self.heads * self.dim_head)
+            out = self.gated_attention_function(x, out, self)
 
         return self.to_out(out)

packages/ltx-core/src/ltx_core/model/transformer/compiling.py

@@ -1,19 +1,112 @@
+from dataclasses import dataclass, field
+from typing import Any
+
 import torch
 
+from ltx_core.guidance.perturbations import BatchedPerturbationConfig, PerturbationType
 from ltx_core.loader.module_ops import ModuleOps
 from ltx_core.loader.sd_ops import SDOps
 from ltx_core.model.transformer.model import LTXModel
+from ltx_core.model.transformer.transformer_args import BlockPerturbationsProcessor, TransformerArgs
+
+# Defaults applied inside the patched forward. Overriding via CompilationConfig
+# replaces these wholesale; it does not merge.
+_DEFAULT_INDUCTOR_CONFIG: dict[str, Any] = {}
+_DEFAULT_DYNAMO_CONFIG: dict[str, Any] = {"inline_inbuilt_nn_modules": True, "cache_size_limit": 256}
 
 
-def compile_transformer(model: LTXModel) -> LTXModel:
-    model.transformer_blocks = torch.nn.ModuleList(torch.compile(m) for m in model.transformer_blocks)
+@dataclass(frozen=True)
+class CompilationConfig:
+    """``torch.compile`` configuration for transformer blocks. ``None`` keeps eager."""
+
+    mode: str | None = None
+    backend: str = "inductor"
+    fullgraph: bool = False
+    dynamic: bool | None = None
+    inductor_config: dict[str, Any] = field(default_factory=lambda: dict(_DEFAULT_INDUCTOR_CONFIG))
+    dynamo_config: dict[str, Any] = field(default_factory=lambda: dict(_DEFAULT_DYNAMO_CONFIG))
+
+
+class _SeqDynamicMarkingProcessor:
+    """Marks the per-block seq dim dynamic, then delegates to an inner processor.
+    Installed by ``compile_transformer`` so the per-block compile artifact stays
+    shape-polymorphic. Wraps whatever ``block_input_processor`` was already on
+    the model -- callers that customised the processor keep their customisation;
+    only the seq-dim marking is layered on top. Lives outside the compiled
+    region, so ``mark_dynamic`` runs in eager mode on the tensors that are
+    about to cross into the trace.
+    """
+
+    def __init__(self, inner: BlockPerturbationsProcessor) -> None:
+        self.inner = inner
+
+    def __call__(
+        self,
+        args: TransformerArgs,
+        perturbations: BatchedPerturbationConfig,
+        block_idx: int,
+        self_attn_type: PerturbationType,
+        cross_attn_type: PerturbationType,
+    ) -> TransformerArgs:
+        # Positional embeddings are second-from-last regardless of rope type:
+        # split rope is (B, H, T, D//2) -- dim -2 == 2; interleaved rope is (B, T, D)
+        # -- dim -2 == 1. Both work via the negative index.
+        torch._dynamo.mark_dynamic(args.x, 1)
+        cos, sin = args.positional_embeddings
+        torch._dynamo.mark_dynamic(cos, cos.ndim - 2)
+        torch._dynamo.mark_dynamic(sin, sin.ndim - 2)
+        if args.cross_positional_embeddings is not None:
+            cross_cos, cross_sin = args.cross_positional_embeddings
+            torch._dynamo.mark_dynamic(cross_cos, cross_cos.ndim - 2)
+            torch._dynamo.mark_dynamic(cross_sin, cross_sin.ndim - 2)
+        if args.self_attention_mask is not None:
+            # Dense form is (B, 1, T, T); key-padding form (from the SP wrapper)
+            # is (B, 1, 1, T) -- leave the size-1 query dim static so Dynamo
+            # keeps the broadcast.
+            if args.self_attention_mask.shape[2] > 1:
+                torch._dynamo.mark_dynamic(args.self_attention_mask, 2)
+            torch._dynamo.mark_dynamic(args.self_attention_mask, 3)
+        if args.context_mask is not None:
+            torch._dynamo.mark_dynamic(args.context_mask, 2)
+        # `timesteps` / `embedded_timestep` are per-token when conditioning sets a
+        # per-position denoise mask, in which case their dim 1 equals the seq length
+        # and must vary with it. When they're a single timestep broadcast across the
+        # sequence (dim 1 == 1), leaving them static lets Dynamo keep the size-1
+        # broadcast.
+        if args.timesteps.shape[1] > 1:
+            torch._dynamo.mark_dynamic(args.timesteps, 1)
+        if args.embedded_timestep.shape[1] > 1:
+            torch._dynamo.mark_dynamic(args.embedded_timestep, 1)
+        # `cross_scale_shift_timestep` is the cross-attn AdaLN scale/shift input
+        # derived from the own-modality per-token timesteps (denoise_mask * sigma),
+        # so its dim 1 equals the seq length when conditioning is per-token.
+        # `cross_gate_timestep` is the cross-modality sigma scalar -- dim 1 is 1
+        # and broadcasts, leave it static. Same guard pattern as `timesteps`.
+        if args.cross_scale_shift_timestep is not None and args.cross_scale_shift_timestep.shape[1] > 1:
+            torch._dynamo.mark_dynamic(args.cross_scale_shift_timestep, 1)
+        return self.inner(args, perturbations, block_idx, self_attn_type, cross_attn_type)
+
+
+def compile_transformer(model: LTXModel, config: CompilationConfig) -> LTXModel:
+    """Compile each transformer block via ``torch.compile`` with the given settings.
+    The patched forward emits ``torch.compiler.cudagraph_mark_step_begin()`` once
+    per step. Under CUDA-graph-enabling modes (``"reduce-overhead"`` /
+    ``"max-autotune"``) this overrides Dynamo's per-invocation auto-mark
+    heuristic, which would otherwise fire once per compiled block call (48 per
+    forward) and treat each block call as a fresh iteration. Under other modes
+    the mark is a no-op (decrements an unread counter).
+    """
+    model.transformer_blocks = torch.nn.ModuleList(
+        torch.compile(m, mode=config.mode, backend=config.backend, fullgraph=config.fullgraph, dynamic=config.dynamic)
+        for m in model.transformer_blocks
+    )
+    model.block_input_processor = _SeqDynamicMarkingProcessor(inner=model.block_input_processor)
 
     def patched_dynamo_forward(*args, **kwargs) -> tuple[torch.Tensor, torch.Tensor]:
+        torch.compiler.cudagraph_mark_step_begin()
         with (
-            torch._inductor.config.patch(unsafe_skip_cache_dynamic_shape_guards=True),
-            torch._dynamo.config.patch(  # type: ignore[attr-defined]
-                inline_inbuilt_nn_modules=True, cache_size_limit=256, allow_unspec_int_on_nn_module=True
-            ),
+            torch._inductor.config.patch(**config.inductor_config),
+            torch._dynamo.config.patch(**config.dynamo_config),  # type: ignore[attr-defined]
         ):
             return model.forward_without_compilation(*args, **kwargs)
 
@@ -22,11 +115,13 @@ def compile_transformer(model: LTXModel) -> LTXModel:
     return model
 
 
-COMPILE_TRANSFORMER = ModuleOps(
-    name="compile_transformer",
-    matcher=lambda model: isinstance(model, LTXModel),
-    mutator=lambda model: compile_transformer(model),
-)
+def build_compile_transformer_op(config: CompilationConfig) -> ModuleOps:
+    """Build a ``ModuleOps`` that compiles transformer blocks with the given settings."""
+    return ModuleOps(
+        name="compile_transformer",
+        matcher=lambda model: isinstance(model, LTXModel),
+        mutator=lambda model: compile_transformer(model, config),
+    )
 
 
 def modify_sd_ops_for_compilation(original_sd_ops: SDOps, number_of_blocks: int = 48) -> SDOps:

packages/ltx-core/src/ltx_core/model/transformer/modality.py

@@ -17,8 +17,19 @@ class Modality:
             the batch size, *T* is the total number of tokens (noisy +
             conditioning), and *D* is the input dimension.
         timesteps: Per-token timestep embeddings, shape ``(B, T)``.
-        positions: Positional coordinates, shape ``(B, 3, T)`` for video
-            (time, height, width) or ``(B, 1, T)`` for audio.
+        positions: Per-token patch coordinates used to build the RoPE
+            frequencies. With the default ``use_middle_indices_grid=True``,
+            shape is ``(B, n_pos_dims, T, 2)`` where ``n_pos_dims=3`` for
+            video (time, height, width) and ``n_pos_dims=1`` for audio
+            (time); the last dim of size 2 holds the ``[start, end)``
+            index bounds of each patch, and RoPE is evaluated at the
+            *middle* of that range -- hence the flag name. Taking the
+            patch midpoint produces a smoother and more accurate
+            positional signal than indexing by the patch's start when
+            patches span more than one spatial / temporal unit.
+            When ``use_middle_indices_grid=False``, the legacy 3-D form
+            ``(B, n_pos_dims, T)`` of integer positional indices is
+            accepted instead and used as-is (no midpoint derivation).
         context: Text conditioning embeddings from the prompt encoder.
         enabled: Whether this modality is active in the current forward pass.
         context_mask: Optional mask for the text context tokens.
@@ -34,9 +45,10 @@ class Modality:
     )  # Shape: (B, T, D) where B is the batch size, T is the number of tokens, and D is input dimension
     sigma: torch.Tensor  # Shape: (B,). Current sigma value, used for cross-attention timestep calculation.
     timesteps: torch.Tensor  # Shape: (B, T) where T is the number of timesteps
-    positions: (
-        torch.Tensor
-    )  # Shape: (B, 3, T) for video, where 3 is the number of dimensions and T is the number of tokens
+    # Shape: (B, n_pos_dims, T, 2) by default (use_middle_indices_grid=True);
+    # n_pos_dims=3 for video, 1 for audio; last dim holds [start, end) patch bounds.
+    # Legacy form (B, n_pos_dims, T) when use_middle_indices_grid=False.
+    positions: torch.Tensor
     context: torch.Tensor
     enabled: bool = True
     context_mask: torch.Tensor | None = None

packages/ltx-core/src/ltx_core/model/transformer/model.py

@@ -1,20 +1,30 @@
+import logging
 from enum import Enum
 
 import torch
 
-from ltx_core.guidance.perturbations import BatchedPerturbationConfig
+from ltx_core.guidance.perturbations import BatchedPerturbationConfig, PerturbationType
+from ltx_core.model.model_protocol import LTXModelProtocol
 from ltx_core.model.transformer.adaln import AdaLayerNormSingle, adaln_embedding_coefficient
-from ltx_core.model.transformer.attention import AttentionCallable, AttentionFunction
+from ltx_core.model.transformer.attention import attention_label
 from ltx_core.model.transformer.modality import Modality
 from ltx_core.model.transformer.rope import LTXRopeType
-from ltx_core.model.transformer.transformer import BasicAVTransformerBlock, TransformerConfig
+from ltx_core.model.transformer.transformer import (
+    DEFAULT_TRANSFORMER_OPS,
+    BasicAVTransformerBlock,
+    TransformerConfig,
+    TransformerOpsConfig,
+)
 from ltx_core.model.transformer.transformer_args import (
+    BlockPerturbationsProcessor,
     MultiModalTransformerArgsPreprocessor,
     TransformerArgs,
     TransformerArgsPreprocessor,
 )
 from ltx_core.utils import to_denoised
 
+logger = logging.getLogger(__name__)
+
 
 class LTXModelType(Enum):
     AudioVideo = "ltx av model"
@@ -45,7 +55,7 @@ class LTXModel(torch.nn.Module):
         num_layers: int = 48,
         cross_attention_dim: int = 4096,
         norm_eps: float = 1e-06,
-        attention_type: AttentionFunction | AttentionCallable = AttentionFunction.DEFAULT,
+        ops: TransformerOpsConfig = DEFAULT_TRANSFORMER_OPS,
         positional_embedding_theta: float = 10000.0,
         positional_embedding_max_pos: list[int] | None = None,
         timestep_scale_multiplier: int = 1000,
@@ -57,7 +67,7 @@ class LTXModel(torch.nn.Module):
         audio_cross_attention_dim: int = 2048,
         audio_positional_embedding_max_pos: list[int] | None = None,
         av_ca_timestep_scale_multiplier: int = 1,
-        rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
+        rope_type: LTXRopeType = LTXRopeType.SPLIT,
         double_precision_rope: bool = False,
         apply_gated_attention: bool = False,
         caption_projection: torch.nn.Module | None = None,
@@ -65,6 +75,15 @@ class LTXModel(torch.nn.Module):
         cross_attention_adaln: bool = False,
     ):
         super().__init__()
+        # Log the attention backends this transformer is built with. Reading the resolved
+        # ``label`` off the ops reports whatever was selected -- AUTOMATIC, an explicit pin
+        # (PYTORCH/XFORMERS/FA3/FA4/SDPA_*), or a directly supplied callable -- so this is the
+        # single source of truth for which kernel a build uses. Fires once per build.
+        logger.info(
+            "Building transformer with attention backends -- self: %s, masked: %s",
+            attention_label(ops.attention_ops.attention_function),
+            attention_label(ops.attention_ops.masked_attention_function),
+        )
         self._enable_gradient_checkpointing = False
         self.cross_attention_adaln = cross_attention_adaln
         self.use_middle_indices_grid = use_middle_indices_grid
@@ -115,9 +134,13 @@ class LTXModel(torch.nn.Module):
             audio_attention_head_dim=audio_attention_head_dim if model_type.is_audio_enabled() else 0,
             audio_cross_attention_dim=audio_cross_attention_dim,
             norm_eps=norm_eps,
-            attention_type=attention_type,
+            ops=ops,
             apply_gated_attention=apply_gated_attention,
         )
+        # Hook for per-block input prep. Compile transforms in `compiling.py`
+        # wrap (not replace) this with a processor that also marks the seq dim
+        # dynamic, so any caller customisation here is preserved as the inner.
+        self.block_input_processor = BlockPerturbationsProcessor()
 
     @property
     def _adaln_embedding_coefficient(self) -> int:
@@ -284,7 +307,7 @@ class LTXModel(torch.nn.Module):
         audio_attention_head_dim: int,
         audio_cross_attention_dim: int,
         norm_eps: float,
-        attention_type: AttentionFunction | AttentionCallable,
+        ops: TransformerOpsConfig,
         apply_gated_attention: bool,
     ) -> None:
         """Initialize transformer blocks for LTX."""
@@ -315,14 +338,13 @@ class LTXModel(torch.nn.Module):
         self.transformer_blocks = torch.nn.ModuleList(
             [
                 BasicAVTransformerBlock(
-                    idx=idx,
                     video=video_config,
                     audio=audio_config,
                     rope_type=self.rope_type,
                     norm_eps=norm_eps,
-                    attention_function=attention_type,
+                    ops=ops,
                 )
-                for idx in range(num_layers)
+                for _ in range(num_layers)
             ]
         )
 
@@ -340,29 +362,44 @@ class LTXModel(torch.nn.Module):
         self,
         video: TransformerArgs | None,
         audio: TransformerArgs | None,
-        perturbations: BatchedPerturbationConfig,
-    ) -> tuple[TransformerArgs, TransformerArgs]:
-        """Process transformer blocks for LTXAV."""
+        perturbations: BatchedPerturbationConfig | None,
+    ) -> tuple[TransformerArgs | None, TransformerArgs | None]:
+        """Process transformer blocks for LTXAV.
+        Per-block perturbation masks are precomputed here and attached to each
+        modality's ``TransformerArgs`` so the block forward has no per-block
+        identity to specialise on — all blocks share a single Dynamo cache slot.
+        """
+        if perturbations is None:
+            batch_size = (video or audio).x.shape[0]
+            perturbations = BatchedPerturbationConfig.empty(batch_size)
+
+        for block_idx, block in enumerate(self.transformer_blocks):
+            if video is not None:
+                video = self.block_input_processor(
+                    video,
+                    perturbations,
+                    block_idx,
+                    self_attn_type=PerturbationType.SKIP_VIDEO_SELF_ATTN,
+                    cross_attn_type=PerturbationType.SKIP_A2V_CROSS_ATTN,
+                )
+            if audio is not None:
+                audio = self.block_input_processor(
+                    audio,
+                    perturbations,
+                    block_idx,
+                    self_attn_type=PerturbationType.SKIP_AUDIO_SELF_ATTN,
+                    cross_attn_type=PerturbationType.SKIP_V2A_CROSS_ATTN,
+                )
 
-        # Process transformer blocks
-        for block in self.transformer_blocks:
             if self._enable_gradient_checkpointing and self.training:
-                # Use gradient checkpointing to save memory during training.
-                # With use_reentrant=False, we can pass dataclasses directly -
-                # PyTorch will track all tensor leaves in the computation graph.
                 video, audio = torch.utils.checkpoint.checkpoint(
                     block,
                     video,
                     audio,
-                    perturbations,
                     use_reentrant=False,
                 )
             else:
-                video, audio = block(
-                    video=video,
-                    audio=audio,
-                    perturbations=perturbations,
-                )
+                video, audio = block(video=video, audio=audio)
 
         return video, audio
 
@@ -388,7 +425,7 @@ class LTXModel(torch.nn.Module):
 
     def forward(
         self, video: Modality | None, audio: Modality | None, perturbations: BatchedPerturbationConfig
-    ) -> tuple[torch.Tensor, torch.Tensor]:
+    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
         """
         Forward pass for LTX models.
         Returns:
@@ -436,7 +473,7 @@ class LegacyX0Model(torch.nn.Module):
     Returns fully denoised output based on the velocities produced by the base model.
     """
 
-    def __init__(self, velocity_model: LTXModel):
+    def __init__(self, velocity_model: LTXModelProtocol):
         super().__init__()
         self.velocity_model = velocity_model
 
@@ -465,7 +502,7 @@ class X0Model(torch.nn.Module):
     Applies scaled denoising to the video and audio according to the timesteps = sigma * denoising_mask.
     """
 
-    def __init__(self, velocity_model: LTXModel):
+    def __init__(self, velocity_model: LTXModelProtocol):
         super().__init__()
         self.velocity_model = velocity_model
 

packages/ltx-core/src/ltx_core/model/transformer/model_configurator.py

@@ -2,10 +2,10 @@ import torch
 
 from ltx_core.loader.sd_ops import SDOps
 from ltx_core.model.model_protocol import ModelConfigurator
-from ltx_core.model.transformer.attention import AttentionFunction
 from ltx_core.model.transformer.model import LTXModel, LTXModelType
 from ltx_core.model.transformer.rope import LTXRopeType
 from ltx_core.model.transformer.text_projection import create_caption_projection
+from ltx_core.model.transformer.transformer import DEFAULT_TRANSFORMER_OPS, TransformerOpsConfig
 from ltx_core.utils import check_config_value
 
 
@@ -16,7 +16,7 @@ class LTXModelConfigurator(ModelConfigurator[LTXModel]):
     """
 
     @classmethod
-    def from_config(cls: type[LTXModel], config: dict) -> LTXModel:
+    def from_config(cls, config: dict, ops: TransformerOpsConfig = DEFAULT_TRANSFORMER_OPS) -> LTXModel:
         # Build caption projections for 19B models (projection handled in transformer).
         caption_projection, audio_caption_projection = _build_caption_projections(config, is_av=True)
 
@@ -40,6 +40,7 @@ class LTXModelConfigurator(ModelConfigurator[LTXModel]):
         check_config_value(config, "share_ff", False)
         check_config_value(config, "av_cross_ada_norm", True)
         check_config_value(config, "use_middle_indices_grid", True)
+        check_config_value(config, "num_attention_heads", config.get("audio_num_attention_heads", float("nan")))
 
         return LTXModel(
             model_type=LTXModelType.AudioVideo,
@@ -50,7 +51,7 @@ class LTXModelConfigurator(ModelConfigurator[LTXModel]):
             num_layers=config.get("num_layers", 48),
             cross_attention_dim=config.get("cross_attention_dim", 4096),
             norm_eps=config.get("norm_eps", 1e-06),
-            attention_type=AttentionFunction(config.get("attention_type", "default")),
+            ops=ops,
             positional_embedding_theta=config.get("positional_embedding_theta", 10000.0),
             positional_embedding_max_pos=config.get("positional_embedding_max_pos", [20, 2048, 2048]),
             timestep_scale_multiplier=config.get("timestep_scale_multiplier", 1000),
@@ -62,7 +63,7 @@ class LTXModelConfigurator(ModelConfigurator[LTXModel]):
             audio_cross_attention_dim=config.get("audio_cross_attention_dim", 2048),
             audio_positional_embedding_max_pos=config.get("audio_positional_embedding_max_pos", [20]),
             av_ca_timestep_scale_multiplier=config.get("av_ca_timestep_scale_multiplier", 1),
-            rope_type=LTXRopeType(config.get("rope_type", "interleaved")),
+            rope_type=LTXRopeType(config.get("rope_type", "split")),
             double_precision_rope=config.get("frequencies_precision", False) == "float64",
             apply_gated_attention=config.get("apply_gated_attention", False),
             caption_projection=caption_projection,
@@ -78,7 +79,7 @@ class LTXVideoOnlyModelConfigurator(ModelConfigurator[LTXModel]):
     """
 
     @classmethod
-    def from_config(cls: type[LTXModel], config: dict) -> LTXModel:
+    def from_config(cls, config: dict, ops: TransformerOpsConfig = DEFAULT_TRANSFORMER_OPS) -> LTXModel:
         # Build caption projection for 19B model (projection handled in transformer).
         caption_projection, _ = _build_caption_projections(config, is_av=False)
 
@@ -109,12 +110,12 @@ class LTXVideoOnlyModelConfigurator(ModelConfigurator[LTXModel]):
             num_layers=config.get("num_layers", 48),
             cross_attention_dim=config.get("cross_attention_dim", 4096),
             norm_eps=config.get("norm_eps", 1e-06),
-            attention_type=AttentionFunction(config.get("attention_type", "default")),
+            ops=ops,
             positional_embedding_theta=config.get("positional_embedding_theta", 10000.0),
             positional_embedding_max_pos=config.get("positional_embedding_max_pos", [20, 2048, 2048]),
             timestep_scale_multiplier=config.get("timestep_scale_multiplier", 1000),
             use_middle_indices_grid=config.get("use_middle_indices_grid", True),
-            rope_type=LTXRopeType(config.get("rope_type", "interleaved")),
+            rope_type=LTXRopeType(config.get("rope_type", "split")),
             double_precision_rope=config.get("frequencies_precision", False) == "float64",
             apply_gated_attention=config.get("apply_gated_attention", False),
             caption_projection=caption_projection,
@@ -122,6 +123,58 @@ class LTXVideoOnlyModelConfigurator(ModelConfigurator[LTXModel]):
         )
 
 
+class LTXAudioOnlyModelConfigurator(ModelConfigurator[LTXModel]):
+    """
+    Configurator for LTX audio only model.
+    Builds an audio-only LTX model (``model_type=AudioOnly``) so the video
+    transformer weights are never instantiated or loaded. Useful for
+    text-to-audio inference where the video branch is unused.
+    """
+
+    @classmethod
+    def from_config(cls, config: dict, ops: TransformerOpsConfig = DEFAULT_TRANSFORMER_OPS) -> LTXModel:
+        # Build audio caption projection for 19B models (projection handled in transformer).
+        _, audio_caption_projection = _build_caption_projections(config, is_av=True)
+
+        config = config.get("transformer", {})
+
+        check_config_value(config, "dropout", 0.0)
+        check_config_value(config, "attention_bias", True)
+        check_config_value(config, "num_vector_embeds", None)
+        check_config_value(config, "activation_fn", "gelu-approximate")
+        check_config_value(config, "num_embeds_ada_norm", 1000)
+        check_config_value(config, "use_linear_projection", False)
+        check_config_value(config, "only_cross_attention", False)
+        check_config_value(config, "cross_attention_norm", True)
+        check_config_value(config, "double_self_attention", False)
+        check_config_value(config, "upcast_attention", False)
+        check_config_value(config, "standardization_norm", "rms_norm")
+        check_config_value(config, "norm_elementwise_affine", False)
+        check_config_value(config, "qk_norm", "rms_norm")
+        check_config_value(config, "positional_embedding_type", "rope")
+        check_config_value(config, "use_middle_indices_grid", True)
+
+        return LTXModel(
+            model_type=LTXModelType.AudioOnly,
+            num_layers=config.get("num_layers", 48),
+            norm_eps=config.get("norm_eps", 1e-06),
+            ops=ops,
+            timestep_scale_multiplier=config.get("timestep_scale_multiplier", 1000),
+            use_middle_indices_grid=config.get("use_middle_indices_grid", True),
+            audio_num_attention_heads=config.get("audio_num_attention_heads", 32),
+            audio_attention_head_dim=config.get("audio_attention_head_dim", 64),
+            audio_in_channels=config.get("audio_in_channels", 128),
+            audio_out_channels=config.get("audio_out_channels", 128),
+            audio_cross_attention_dim=config.get("audio_cross_attention_dim", 2048),
+            audio_positional_embedding_max_pos=config.get("audio_positional_embedding_max_pos", [20]),
+            rope_type=LTXRopeType(config.get("rope_type", "split")),
+            double_precision_rope=config.get("frequencies_precision", False) == "float64",
+            apply_gated_attention=config.get("apply_gated_attention", False),
+            audio_caption_projection=audio_caption_projection,
+            cross_attention_adaln=config.get("cross_attention_adaln", False),
+        )
+
+
 def _build_caption_projections(
     config: dict,
     is_av: bool,
@@ -150,3 +203,16 @@ LTXV_MODEL_COMFY_RENAMING_MAP = (
     .with_matching(prefix="model.diffusion_model.")
     .with_replacement("model.diffusion_model.", "")
 )
+
+LTXV_AUDIO_ONLY_MODEL_COMFY_RENAMING_MAP = (
+    SDOps("LTXV_AUDIO_ONLY_MODEL_COMFY_MAP")
+    .with_matching(prefix="model.diffusion_model.", contains="audio_attn1")
+    .with_matching(prefix="model.diffusion_model.", contains="audio_attn2")
+    .with_matching(prefix="model.diffusion_model.", contains="audio_ff")
+    .with_matching(prefix="model.diffusion_model.", contains="audio_patchify")
+    .with_matching(prefix="model.diffusion_model.", contains="audio_proj_out")
+    .with_matching(prefix="model.diffusion_model.", contains="audio_adaln_single")
+    .with_matching(prefix="model.diffusion_model.", contains="audio_prompt")
+    .with_matching(prefix="model.diffusion_model.", contains="audio_scale_shift_table")
+    .with_replacement("model.diffusion_model.", "")
+)

packages/ltx-core/src/ltx_core/model/transformer/ops.py

@@ -0,0 +1,106 @@
+from typing import List, Protocol
+
+import torch
+from torch import nn
+
+from ltx_core.model.transformer.rope import apply_rotary_emb
+from ltx_core.utils import rms_norm
+
+
+class PreAttentionCallable(Protocol):
+    def __call__(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        attn_module: nn.Module,
+        mask: torch.Tensor | None,
+        pe: torch.Tensor | None,
+        k_pe: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor]: ...
+
+
+class PytorchPreAttention(PreAttentionCallable):
+    def __call__(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        attn_module: nn.Module,
+        mask: torch.Tensor | None,  # noqa: ARG002
+        pe: torch.Tensor | None,
+        k_pe: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        q = attn_module.q_norm(q)
+        k = attn_module.k_norm(k)
+        if pe is not None:
+            q = apply_rotary_emb(q, pe, attn_module.rope_type)
+            k = apply_rotary_emb(k, pe if k_pe is None else k_pe, attn_module.rope_type)
+        return q, k
+
+
+class AdaZeroCallable(Protocol):
+    def __call__(
+        self,
+        x: torch.Tensor,
+        eps: float,
+        scale: torch.Tensor,
+        shift: torch.Tensor,
+    ) -> torch.Tensor: ...
+
+
+class PytorchAdaZeroFunction(AdaZeroCallable):
+    def __call__(
+        self,
+        x: torch.Tensor,
+        eps: float,
+        scale: torch.Tensor,
+        shift: torch.Tensor,
+    ) -> torch.Tensor:
+        return rms_norm(x, eps=eps) * (1 + scale) + shift
+
+
+class PostSACallable(Protocol):
+    def __call__(
+        self,
+        x: torch.Tensor,
+        y: torch.Tensor,
+        norm_weights: torch.Tensor | None,
+        eps: float,
+        gate: torch.Tensor,
+    ) -> List[torch.Tensor]: ...
+
+
+class PytorchPostSAFunction(PostSACallable):
+    def __call__(
+        self,
+        x: torch.Tensor,
+        y: torch.Tensor,
+        norm_weights: torch.Tensor | None,
+        eps: float,
+        gate: torch.Tensor,
+    ) -> List[torch.Tensor]:
+        x_fma = x + y * gate
+        return x_fma, rms_norm(x_fma, norm_weights, eps=eps)
+
+
+class GatedAttentionCallable(Protocol):
+    def __call__(
+        self,
+        x: torch.Tensor,
+        attn_out: torch.Tensor,
+        attn_module: nn.Module,
+    ) -> torch.Tensor: ...
+
+
+class PytorchGatedAttention(GatedAttentionCallable):
+    def __call__(
+        self,
+        x: torch.Tensor,
+        attn_out: torch.Tensor,
+        attn_module: nn.Module,
+    ) -> torch.Tensor:
+        gate_logits = attn_module.to_gate_logits(x)  # (B, T, H)
+        b, t, _ = attn_out.shape
+        out = attn_out.view(b, t, attn_module.heads, attn_module.dim_head)
+        gates = 2.0 * torch.sigmoid(gate_logits)  # (B, T, H)
+        out = out * gates.unsqueeze(-1)  # (B, T, H, D) * (B, T, H, 1)
+        return out.view(b, t, attn_module.heads * attn_module.dim_head)

packages/ltx-core/src/ltx_core/model/transformer/rope.py

@@ -16,9 +16,10 @@ class LTXRopeType(Enum):
 def apply_rotary_emb(
     input_tensor: torch.Tensor,
     freqs_cis: Tuple[torch.Tensor, torch.Tensor],
-    rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
+    rope_type: LTXRopeType = LTXRopeType.SPLIT,
 ) -> torch.Tensor:
     if rope_type == LTXRopeType.INTERLEAVED:
+        # Note: INTERLEAVED rope is a legacy mode. Prefer SPLIT instead.
         return apply_interleaved_rotary_emb(input_tensor, *freqs_cis)
     elif rope_type == LTXRopeType.SPLIT:
         return apply_split_rotary_emb(input_tensor, *freqs_cis)
@@ -42,11 +43,26 @@ def apply_interleaved_rotary_emb(
 def apply_split_rotary_emb(
     input_tensor: torch.Tensor, cos_freqs: torch.Tensor, sin_freqs: torch.Tensor
 ) -> torch.Tensor:
-    needs_reshape = False
-    if input_tensor.ndim != 4 and cos_freqs.ndim == 4:
-        b, h, t, _ = cos_freqs.shape
-        input_tensor = input_tensor.reshape(b, t, h, -1).swapaxes(1, 2)
-        needs_reshape = True
+    if sin_freqs.shape != cos_freqs.shape:
+        raise ValueError(
+            f"apply_split_rotary_emb: sin_freqs.shape {tuple(sin_freqs.shape)} must equal "
+            f"cos_freqs.shape {tuple(cos_freqs.shape)}."
+        )
+    needs_reshape = input_tensor.ndim != 4 and cos_freqs.ndim == 4
+    if needs_reshape:
+        b_freq = cos_freqs.shape[0]
+        h = cos_freqs.shape[1]
+        b_in = input_tensor.shape[0]
+        if b_freq not in (1, b_in):
+            raise ValueError(
+                f"apply_split_rotary_emb: cos_freqs batch ({b_freq}) must be 1 "
+                f"(broadcast) or equal input_tensor batch ({b_in})."
+            )
+        # `unflatten` only touches the last dim, keeping the batch and seq dims as
+        # the input tensor's own symbolic ints under torch.compile. `reshape(b_in,
+        # t, h, -1)` would have forced Dynamo to specialise those dims because it
+        # cannot prove `b_in == cos_freqs.shape[0]` and `seq == t` across tensors.
+        input_tensor = input_tensor.unflatten(-1, (h, -1)).transpose(1, 2)
 
     split_input = rearrange(input_tensor, "... (d r) -> ... d r", d=2)
     first_half_input = split_input[..., :1, :]
@@ -61,7 +77,9 @@ def apply_split_rotary_emb(
 
     output = rearrange(output, "... d r -> ... (d r)")
     if needs_reshape:
-        output = output.swapaxes(1, 2).reshape(b, t, -1)
+        # `transpose(1, 2).flatten(-2)` keeps the batch and seq dims symbolic; using
+        # `reshape(b_in, t, -1)` would force Dynamo to specialise both axes.
+        output = output.transpose(1, 2).flatten(-2)
 
     return output
 
@@ -183,7 +201,7 @@ def precompute_freqs_cis(
     max_pos: list[int] | None = None,
     use_middle_indices_grid: bool = False,
     num_attention_heads: int = 32,
-    rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
+    rope_type: LTXRopeType = LTXRopeType.SPLIT,
     freq_grid_generator: Callable[[float, int, int, torch.device], torch.Tensor] = generate_freq_grid_pytorch,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     if max_pos is None:

packages/ltx-core/src/ltx_core/model/transformer/transformer.py

@@ -1,14 +1,29 @@
-from dataclasses import dataclass, replace
+from dataclasses import dataclass, field, replace
 
 import torch
 
-from ltx_core.guidance.perturbations import BatchedPerturbationConfig, PerturbationType
 from ltx_core.model.transformer.adaln import adaln_embedding_coefficient
-from ltx_core.model.transformer.attention import Attention, AttentionCallable, AttentionFunction
+from ltx_core.model.transformer.attention import (
+    Attention,
+    AttentionCallable,
+    AttentionFunction,
+    AttentionOps,
+    MaskedAttentionCallable,
+    MaskedAttentionFunction,
+)
 from ltx_core.model.transformer.feed_forward import FeedForward
+from ltx_core.model.transformer.ops import (
+    AdaZeroCallable,
+    GatedAttentionCallable,
+    PostSACallable,
+    PreAttentionCallable,
+    PytorchAdaZeroFunction,
+    PytorchGatedAttention,
+    PytorchPostSAFunction,
+    PytorchPreAttention,
+)
 from ltx_core.model.transformer.rope import LTXRopeType
 from ltx_core.model.transformer.transformer_args import TransformerArgs
-from ltx_core.utils import rms_norm
 
 
 @dataclass
@@ -21,19 +36,68 @@ class TransformerConfig:
     cross_attention_adaln: bool = False
 
 
+@dataclass(frozen=True)
+class TransformerOpsConfig:
+    """Pluggable ops for :class:`BasicAVTransformerBlock`.
+    Use :meth:`from_functions` to construct from enum values or partial overrides
+    without spelling out a full :class:`AttentionOps`.
+    """
+
+    attention_ops: AttentionOps = field(default_factory=AttentionOps)
+    ada_zero_function: AdaZeroCallable = field(default_factory=PytorchAdaZeroFunction)
+    post_sa_function: PostSACallable = field(default_factory=PytorchPostSAFunction)
+
+    @classmethod
+    def from_functions(
+        cls,
+        attention: AttentionFunction | AttentionCallable = AttentionFunction.AUTOMATIC,
+        masked_attention: MaskedAttentionFunction | MaskedAttentionCallable = MaskedAttentionFunction.AUTOMATIC,
+        preattention: PreAttentionCallable | None = None,
+        gated_attention: GatedAttentionCallable | None = None,
+        ada_zero: AdaZeroCallable | None = None,
+        post_sa: PostSACallable | None = None,
+    ) -> "TransformerOpsConfig":
+        """Build a config from individual functions or enums. Each *None* slot
+        falls back to the standard PyTorch implementation."""
+        attention_callable = attention.to_callable() if isinstance(attention, AttentionFunction) else attention
+        masked_callable = (
+            masked_attention.to_callable()
+            if isinstance(masked_attention, MaskedAttentionFunction)
+            else masked_attention
+        )
+        attention_ops = AttentionOps(
+            attention_function=attention_callable,
+            masked_attention_function=masked_callable,
+            preattention_function=preattention if preattention is not None else PytorchPreAttention(),
+            gated_attention_function=(gated_attention if gated_attention is not None else PytorchGatedAttention()),
+        )
+        return cls(
+            attention_ops=attention_ops,
+            ada_zero_function=ada_zero if ada_zero is not None else PytorchAdaZeroFunction(),
+            post_sa_function=post_sa if post_sa is not None else PytorchPostSAFunction(),
+        )
+
+
+# Frozen, so safe to share as a default argument across callers that want the
+# stock PyTorch ops without explicit construction.
+DEFAULT_TRANSFORMER_OPS = TransformerOpsConfig()
+
+
 class BasicAVTransformerBlock(torch.nn.Module):
     def __init__(
         self,
-        idx: int,
         video: TransformerConfig | None = None,
         audio: TransformerConfig | None = None,
-        rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
+        rope_type: LTXRopeType = LTXRopeType.SPLIT,
         norm_eps: float = 1e-6,
-        attention_function: AttentionFunction | AttentionCallable = AttentionFunction.DEFAULT,
+        ops: TransformerOpsConfig | None = None,
     ):
         super().__init__()
 
-        self.idx = idx
+        if ops is None:
+            ops = TransformerOpsConfig()
+        self.ada_zero_function = ops.ada_zero_function
+        self.post_sa_function = ops.post_sa_function
         if video is not None:
             self.attn1 = Attention(
                 query_dim=video.dim,
@@ -42,7 +106,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 context_dim=None,
                 rope_type=rope_type,
                 norm_eps=norm_eps,
-                attention_function=attention_function,
+                ops=ops.attention_ops,
                 apply_gated_attention=video.apply_gated_attention,
             )
             self.attn2 = Attention(
@@ -52,7 +116,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 dim_head=video.d_head,
                 rope_type=rope_type,
                 norm_eps=norm_eps,
-                attention_function=attention_function,
+                ops=ops.attention_ops,
                 apply_gated_attention=video.apply_gated_attention,
             )
             self.ff = FeedForward(video.dim, dim_out=video.dim)
@@ -67,7 +131,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 context_dim=None,
                 rope_type=rope_type,
                 norm_eps=norm_eps,
-                attention_function=attention_function,
+                ops=ops.attention_ops,
                 apply_gated_attention=audio.apply_gated_attention,
             )
             self.audio_attn2 = Attention(
@@ -77,7 +141,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 dim_head=audio.d_head,
                 rope_type=rope_type,
                 norm_eps=norm_eps,
-                attention_function=attention_function,
+                ops=ops.attention_ops,
                 apply_gated_attention=audio.apply_gated_attention,
             )
             self.audio_ff = FeedForward(audio.dim, dim_out=audio.dim)
@@ -93,7 +157,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 dim_head=audio.d_head,
                 rope_type=rope_type,
                 norm_eps=norm_eps,
-                attention_function=attention_function,
+                ops=ops.attention_ops,
                 apply_gated_attention=video.apply_gated_attention,
             )
 
@@ -105,7 +169,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 dim_head=audio.d_head,
                 rope_type=rope_type,
                 norm_eps=norm_eps,
-                attention_function=attention_function,
+                ops=ops.attention_ops,
                 apply_gated_attention=audio.apply_gated_attention,
             )
 
@@ -157,7 +221,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
 
     def _apply_text_cross_attention(
         self,
-        x: torch.Tensor,
+        x_normed: torch.Tensor,
         context: torch.Tensor,
         attn: AttentionCallable,
         scale_shift_table: torch.Tensor,
@@ -167,11 +231,14 @@ class BasicAVTransformerBlock(torch.nn.Module):
         context_mask: torch.Tensor | None,
         cross_attention_adaln: bool = False,
     ) -> torch.Tensor:
-        """Apply text cross-attention, with optional AdaLN modulation."""
+        """Apply text cross-attention, with optional AdaLN modulation.
+        ``x_normed`` is the RMS-normalized self-attention output produced by
+        ``post_sa_function`` -- this method does not normalize again.
+        """
         if cross_attention_adaln:
-            shift_q, scale_q, gate = self.get_ada_values(scale_shift_table, x.shape[0], timestep, slice(6, 9))
+            shift_q, scale_q, gate = self.get_ada_values(scale_shift_table, x_normed.shape[0], timestep, slice(6, 9))
             return apply_cross_attention_adaln(
-                x,
+                x_normed,
                 context,
                 attn,
                 shift_q,
@@ -180,24 +247,17 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 prompt_scale_shift_table,
                 prompt_timestep,
                 context_mask,
-                self.norm_eps,
             )
-        return attn(rms_norm(x, eps=self.norm_eps), context=context, mask=context_mask)
+        return attn(x_normed, context=context, mask=context_mask)
 
     def forward(  # noqa: PLR0915
         self,
         video: TransformerArgs | None,
         audio: TransformerArgs | None,
-        perturbations: BatchedPerturbationConfig | None = None,
     ) -> tuple[TransformerArgs | None, TransformerArgs | None]:
         if video is None and audio is None:
             raise ValueError("At least one of video or audio must be provided")
 
-        batch_size = (video or audio).x.shape[0]
-
-        if perturbations is None:
-            perturbations = BatchedPerturbationConfig.empty(batch_size)
-
         vx = video.x if video is not None else None
         ax = audio.x if audio is not None else None
 
@@ -211,30 +271,20 @@ class BasicAVTransformerBlock(torch.nn.Module):
             vshift_msa, vscale_msa, vgate_msa = self.get_ada_values(
                 self.scale_shift_table, vx.shape[0], video.timesteps, slice(0, 3)
             )
-            norm_vx = rms_norm(vx, eps=self.norm_eps) * (1 + vscale_msa) + vshift_msa
+            norm_vx = self.ada_zero_function(vx, self.norm_eps, vscale_msa, vshift_msa)
             del vshift_msa, vscale_msa
 
-            all_perturbed = perturbations.all_in_batch(PerturbationType.SKIP_VIDEO_SELF_ATTN, self.idx)
-            none_perturbed = not perturbations.any_in_batch(PerturbationType.SKIP_VIDEO_SELF_ATTN, self.idx)
-            v_mask = (
-                perturbations.mask_like(PerturbationType.SKIP_VIDEO_SELF_ATTN, self.idx, vx)
-                if not all_perturbed and not none_perturbed
-                else None
+            vx_msa_out = self.attn1(
+                norm_vx,
+                pe=video.positional_embeddings,
+                mask=video.self_attention_mask,
+                perturbation_mask=video.self_attn_perturbation_mask,
+                all_perturbed=video.self_attn_all_perturbed,
             )
-            vx = (
-                vx
-                + self.attn1(
-                    norm_vx,
-                    pe=video.positional_embeddings,
-                    mask=video.self_attention_mask,
-                    perturbation_mask=v_mask,
-                    all_perturbed=all_perturbed,
-                )
-                * vgate_msa
-            )
-            del vgate_msa, norm_vx, v_mask
+            vx, vx_normed = self.post_sa_function(vx, vx_msa_out, None, self.norm_eps, vgate_msa)
+            del vgate_msa, norm_vx, vx_msa_out
             vx = vx + self._apply_text_cross_attention(
-                vx,
+                vx_normed,
                 video.context,
                 self.attn2,
                 self.scale_shift_table,
@@ -244,35 +294,26 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 video.context_mask,
                 cross_attention_adaln=self.cross_attention_adaln,
             )
+            del vx_normed
 
         if run_ax:
             ashift_msa, ascale_msa, agate_msa = self.get_ada_values(
                 self.audio_scale_shift_table, ax.shape[0], audio.timesteps, slice(0, 3)
             )
 
-            norm_ax = rms_norm(ax, eps=self.norm_eps) * (1 + ascale_msa) + ashift_msa
+            norm_ax = self.ada_zero_function(ax, self.norm_eps, ascale_msa, ashift_msa)
             del ashift_msa, ascale_msa
-            all_perturbed = perturbations.all_in_batch(PerturbationType.SKIP_AUDIO_SELF_ATTN, self.idx)
-            none_perturbed = not perturbations.any_in_batch(PerturbationType.SKIP_AUDIO_SELF_ATTN, self.idx)
-            a_mask = (
-                perturbations.mask_like(PerturbationType.SKIP_AUDIO_SELF_ATTN, self.idx, ax)
-                if not all_perturbed and not none_perturbed
-                else None
+            ax_msa_out = self.audio_attn1(
+                norm_ax,
+                pe=audio.positional_embeddings,
+                mask=audio.self_attention_mask,
+                perturbation_mask=audio.self_attn_perturbation_mask,
+                all_perturbed=audio.self_attn_all_perturbed,
             )
-            ax = (
-                ax
-                + self.audio_attn1(
-                    norm_ax,
-                    pe=audio.positional_embeddings,
-                    mask=audio.self_attention_mask,
-                    perturbation_mask=a_mask,
-                    all_perturbed=all_perturbed,
-                )
-                * agate_msa
-            )
-            del agate_msa, norm_ax, a_mask
+            ax, ax_normed = self.post_sa_function(ax, ax_msa_out, None, self.norm_eps, agate_msa)
+            del agate_msa, norm_ax, ax_msa_out
             ax = ax + self._apply_text_cross_attention(
-                ax,
+                ax_normed,
                 audio.context,
                 self.audio_attn2,
                 self.audio_scale_shift_table,
@@ -282,13 +323,15 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 audio.context_mask,
                 cross_attention_adaln=self.cross_attention_adaln,
             )
+            del ax_normed
 
         # Audio - Video cross attention.
         if run_a2v or run_v2a:
-            vx_norm3 = rms_norm(vx, eps=self.norm_eps)
-            ax_norm3 = rms_norm(ax, eps=self.norm_eps)
-
-            if run_a2v and not perturbations.all_in_batch(PerturbationType.SKIP_A2V_CROSS_ATTN, self.idx):
+            # Snapshot vx/ax before A2V mutates vx; V2A's video keys/values must
+            # use the pre-A2V state so direction order doesn't bias the result.
+            vx_pre_av = vx
+            ax_pre_av = ax
+            if run_a2v and not video.cross_attn_skip_all:
                 scale_ca_video_a2v, shift_ca_video_a2v, gate_out_a2v = self.get_av_ca_ada_values(
                     self.scale_shift_table_a2v_ca_video,
                     vx.shape[0],
@@ -296,7 +339,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
                     video.cross_gate_timestep,
                     slice(0, 2),
                 )
-                vx_scaled = vx_norm3 * (1 + scale_ca_video_a2v) + shift_ca_video_a2v
+                a2v_vx_scaled = self.ada_zero_function(vx_pre_av, self.norm_eps, scale_ca_video_a2v, shift_ca_video_a2v)
                 del scale_ca_video_a2v, shift_ca_video_a2v
 
                 scale_ca_audio_a2v, shift_ca_audio_a2v, _ = self.get_av_ca_ada_values(
@@ -306,22 +349,21 @@ class BasicAVTransformerBlock(torch.nn.Module):
                     audio.cross_gate_timestep,
                     slice(0, 2),
                 )
-                ax_scaled = ax_norm3 * (1 + scale_ca_audio_a2v) + shift_ca_audio_a2v
+                a2v_ax_scaled = self.ada_zero_function(ax_pre_av, self.norm_eps, scale_ca_audio_a2v, shift_ca_audio_a2v)
                 del scale_ca_audio_a2v, shift_ca_audio_a2v
-                a2v_mask = perturbations.mask_like(PerturbationType.SKIP_A2V_CROSS_ATTN, self.idx, vx)
                 vx = vx + (
                     self.audio_to_video_attn(
-                        vx_scaled,
-                        context=ax_scaled,
+                        a2v_vx_scaled,
+                        context=a2v_ax_scaled,
                         pe=video.cross_positional_embeddings,
                         k_pe=audio.cross_positional_embeddings,
                     )
                     * gate_out_a2v
-                    * a2v_mask
+                    * video.cross_attn_perturbation_mask
                 )
-                del gate_out_a2v, a2v_mask, vx_scaled, ax_scaled
+                del gate_out_a2v, a2v_vx_scaled, a2v_ax_scaled
 
-            if run_v2a and not perturbations.all_in_batch(PerturbationType.SKIP_V2A_CROSS_ATTN, self.idx):
+            if run_v2a and not audio.cross_attn_skip_all:
                 scale_ca_audio_v2a, shift_ca_audio_v2a, gate_out_v2a = self.get_av_ca_ada_values(
                     self.scale_shift_table_a2v_ca_audio,
                     ax.shape[0],
@@ -329,7 +371,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
                     audio.cross_gate_timestep,
                     slice(2, 4),
                 )
-                ax_scaled = ax_norm3 * (1 + scale_ca_audio_v2a) + shift_ca_audio_v2a
+                v2a_ax_scaled = self.ada_zero_function(ax_pre_av, self.norm_eps, scale_ca_audio_v2a, shift_ca_audio_v2a)
                 del scale_ca_audio_v2a, shift_ca_audio_v2a
                 scale_ca_video_v2a, shift_ca_video_v2a, _ = self.get_av_ca_ada_values(
                     self.scale_shift_table_a2v_ca_video,
@@ -338,28 +380,26 @@ class BasicAVTransformerBlock(torch.nn.Module):
                     video.cross_gate_timestep,
                     slice(2, 4),
                 )
-                vx_scaled = vx_norm3 * (1 + scale_ca_video_v2a) + shift_ca_video_v2a
+                v2a_vx_scaled = self.ada_zero_function(vx_pre_av, self.norm_eps, scale_ca_video_v2a, shift_ca_video_v2a)
                 del scale_ca_video_v2a, shift_ca_video_v2a
-                v2a_mask = perturbations.mask_like(PerturbationType.SKIP_V2A_CROSS_ATTN, self.idx, ax)
                 ax = ax + (
                     self.video_to_audio_attn(
-                        ax_scaled,
-                        context=vx_scaled,
+                        v2a_ax_scaled,
+                        context=v2a_vx_scaled,
                         pe=audio.cross_positional_embeddings,
                         k_pe=video.cross_positional_embeddings,
                     )
                     * gate_out_v2a
-                    * v2a_mask
+                    * audio.cross_attn_perturbation_mask
                 )
-                del gate_out_v2a, v2a_mask, ax_scaled, vx_scaled
-
-            del vx_norm3, ax_norm3
+                del gate_out_v2a, v2a_vx_scaled, v2a_ax_scaled
+            del vx_pre_av, ax_pre_av
 
         if run_vx:
             vshift_mlp, vscale_mlp, vgate_mlp = self.get_ada_values(
                 self.scale_shift_table, vx.shape[0], video.timesteps, slice(3, 6)
             )
-            vx_scaled = rms_norm(vx, eps=self.norm_eps) * (1 + vscale_mlp) + vshift_mlp
+            vx_scaled = self.ada_zero_function(vx, self.norm_eps, vscale_mlp, vshift_mlp)
             vx = vx + self.ff(vx_scaled) * vgate_mlp
 
             del vshift_mlp, vscale_mlp, vgate_mlp, vx_scaled
@@ -368,7 +408,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
             ashift_mlp, ascale_mlp, agate_mlp = self.get_ada_values(
                 self.audio_scale_shift_table, ax.shape[0], audio.timesteps, slice(3, 6)
             )
-            ax_scaled = rms_norm(ax, eps=self.norm_eps) * (1 + ascale_mlp) + ashift_mlp
+            ax_scaled = self.ada_zero_function(ax, self.norm_eps, ascale_mlp, ashift_mlp)
             ax = ax + self.audio_ff(ax_scaled) * agate_mlp
 
             del ashift_mlp, ascale_mlp, agate_mlp, ax_scaled
@@ -377,7 +417,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
 
 
 def apply_cross_attention_adaln(
-    x: torch.Tensor,
+    x_normed: torch.Tensor,
     context: torch.Tensor,
     attn: AttentionCallable,
     q_shift: torch.Tensor,
@@ -386,13 +426,17 @@ def apply_cross_attention_adaln(
     prompt_scale_shift_table: torch.Tensor,
     prompt_timestep: torch.Tensor,
     context_mask: torch.Tensor | None = None,
-    norm_eps: float = 1e-6,
 ) -> torch.Tensor:
-    batch_size = x.shape[0]
+    """Apply query/key AdaLN modulation then cross-attention.
+    ``x_normed`` is already RMS-normalized by ``post_sa_function``; this only
+    applies the affine (scale/shift) modulation, so the normalization is not
+    repeated here.
+    """
+    batch_size = x_normed.shape[0]
     shift_kv, scale_kv = (
-        prompt_scale_shift_table[None, None].to(device=x.device, dtype=x.dtype)
+        prompt_scale_shift_table[None, None].to(device=x_normed.device, dtype=x_normed.dtype)
         + prompt_timestep.reshape(batch_size, prompt_timestep.shape[1], 2, -1)
     ).unbind(dim=2)
-    attn_input = rms_norm(x, eps=norm_eps) * (1 + q_scale) + q_shift
+    attn_input = x_normed * (1 + q_scale) + q_shift
     encoder_hidden_states = context * (1 + scale_kv) + shift_kv
     return attn(attn_input, context=encoder_hidden_states, mask=context_mask) * q_gate

packages/ltx-core/src/ltx_core/model/transformer/transformer_args.py

@@ -2,6 +2,7 @@ from dataclasses import dataclass, replace
 
 import torch
 
+from ltx_core.guidance.perturbations import BatchedPerturbationConfig, PerturbationType
 from ltx_core.model.transformer.adaln import AdaLayerNormSingle
 from ltx_core.model.transformer.modality import Modality
 from ltx_core.model.transformer.rope import (
@@ -19,8 +20,8 @@ class TransformerArgs:
     context_mask: torch.Tensor
     timesteps: torch.Tensor
     embedded_timestep: torch.Tensor
-    positional_embeddings: torch.Tensor
-    cross_positional_embeddings: torch.Tensor | None
+    positional_embeddings: tuple[torch.Tensor, torch.Tensor]
+    cross_positional_embeddings: tuple[torch.Tensor, torch.Tensor] | None
     cross_scale_shift_timestep: torch.Tensor | None
     cross_gate_timestep: torch.Tensor | None
     enabled: bool
@@ -28,6 +29,59 @@ class TransformerArgs:
     self_attention_mask: torch.Tensor | None = (
         None  # Additive log-space self-attention bias (B, 1, T, T), None = full attention
     )
+    # Per-block perturbation state, precomputed by `LTXModel._process_transformer_blocks`
+    # so the block forward needs no per-block identity. The bool shortcuts
+    # (`*_all_perturbed`, `cross_attn_skip_all`) are Python bools that Dynamo specialises
+    # on — fine because they're stable across denoising steps for a fixed perturbation
+    # config.
+    self_attn_perturbation_mask: torch.Tensor | None = None
+    self_attn_all_perturbed: bool = False
+    cross_attn_perturbation_mask: torch.Tensor | None = None
+    cross_attn_skip_all: bool = False
+
+
+class BlockPerturbationsProcessor:
+    """Per-block preparation of ``TransformerArgs``.
+    The base implementation returns a copy of ``args`` with this block's
+    precomputed perturbation flags and masks attached. Subclasses can layer in
+    operations that must run on each block's inputs but stay outside the
+    compile boundary -- e.g. ``torch._dynamo.mark_dynamic`` for
+    shape-polymorphic block compilation (see ``compiling.py``). Swapping the
+    processor on an ``LTXModel`` instance is how compile transforms opt in to
+    such behaviour without baking it into the model's forward.
+    ``self_attn_perturbation_mask`` is None when all or none of the batch is
+    perturbed (the attention call can take the shortcut path). ``cross_attn_*``
+    is None when every sample skips the cross-attention entirely.
+    """
+
+    def __call__(
+        self,
+        args: "TransformerArgs",
+        perturbations: BatchedPerturbationConfig,
+        block_idx: int,
+        self_attn_type: PerturbationType,
+        cross_attn_type: PerturbationType,
+    ) -> "TransformerArgs":
+        device, dtype = args.x.device, args.x.dtype
+
+        all_self = perturbations.all_in_batch(self_attn_type, block_idx)
+        any_self = perturbations.any_in_batch(self_attn_type, block_idx)
+        self_mask: torch.Tensor | None = None
+        if any_self and not all_self:
+            self_mask = perturbations.mask(self_attn_type, block_idx, device, dtype).view(-1, 1, 1)
+
+        all_cross = perturbations.all_in_batch(cross_attn_type, block_idx)
+        cross_mask: torch.Tensor | None = None
+        if not all_cross:
+            cross_mask = perturbations.mask(cross_attn_type, block_idx, device, dtype).view(-1, 1, 1)
+
+        return replace(
+            args,
+            self_attn_perturbation_mask=self_mask,
+            self_attn_all_perturbed=all_self,
+            cross_attn_perturbation_mask=cross_mask,
+            cross_attn_skip_all=all_cross,
+        )
 
 
 class TransformerArgsPreprocessor:
@@ -98,9 +152,12 @@ class TransformerArgsPreprocessor:
         self, attention_mask: torch.Tensor | None, x_dtype: torch.dtype
     ) -> torch.Tensor | None:
         """Prepare self-attention mask by converting [0,1] values to additive log-space bias.
-        Input shape: (B, T, T) with values in [0, 1].
-        Output shape: (B, 1, T, T) with 0.0 for full attention and a large negative value
-        for masked positions.
+        Input shape: 3D ``(B, T_q, T_k)`` with values in [0, 1]. The dense form
+        is ``(B, T, T)``; broadcastable forms like ``(1, 1, T)`` (key-only
+        padding) or ``(B, 1, T)`` are also valid and yield a correspondingly
+        broadcastable output.
+        Output shape: ``(B, 1, T_q, T_k)`` (heads dim inserted) with 0.0 for
+        full attention and a large negative value for masked positions.
         Positions with attention_mask <= 0 are fully masked (mapped to the dtype's minimum
         representable value). Strictly positive entries are converted via log-space for
         smooth attenuation, with small values clamped for numerical stability.
@@ -120,7 +177,7 @@ class TransformerArgsPreprocessor:
         if positive.any():
             bias[positive] = torch.log(attention_mask[positive].clamp(min=eps)).to(x_dtype)
 
-        return bias.unsqueeze(1)  # (B, 1, T, T) for head broadcast
+        return bias.unsqueeze(1)  # (B, 1, T_q, T_k) for head broadcast
 
     def _prepare_positional_embeddings(
         self,
@@ -244,10 +301,6 @@ class MultiModalTransformerArgsPreprocessor:
             if cross_modality.sigma.ndim != 1:
                 raise ValueError("Cross modality sigma must be a 1D tensor")
 
-        cross_timestep = cross_modality.sigma.view(
-            modality.timesteps.shape[0], 1, *[1] * len(modality.timesteps.shape[2:])
-        )
-
         cross_pe = self.simple_preprocessor._prepare_positional_embeddings(
             positions=modality.positions[:, 0:1, :],
             inner_dim=self.audio_cross_attention_dim,
@@ -258,7 +311,8 @@ class MultiModalTransformerArgsPreprocessor:
         )
 
         cross_scale_shift_timestep, cross_gate_timestep = self._prepare_cross_attention_timestep(
-            timestep=cross_timestep,
+            modality_timesteps=modality.timesteps,
+            cross_modality_sigma=cross_modality.sigma,
             timestep_scale_multiplier=self.simple_preprocessor.timestep_scale_multiplier,
             batch_size=transformer_args.x.shape[0],
             hidden_dtype=modality.latent.dtype,
@@ -273,23 +327,23 @@ class MultiModalTransformerArgsPreprocessor:
 
     def _prepare_cross_attention_timestep(
         self,
-        timestep: torch.Tensor | None,
+        modality_timesteps: torch.Tensor,
+        cross_modality_sigma: torch.Tensor,
         timestep_scale_multiplier: int,
         batch_size: int,
         hidden_dtype: torch.dtype,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        """Prepare cross attention timestep embeddings."""
-        timestep = timestep * timestep_scale_multiplier
-
+        """Prepare A-V cross-attention AdaLN inputs."""
         av_ca_factor = self.av_ca_timestep_scale_multiplier / timestep_scale_multiplier
 
         scale_shift_timestep, _ = self.cross_scale_shift_adaln(
-            timestep.flatten(),
+            (modality_timesteps * timestep_scale_multiplier).flatten(),
             hidden_dtype=hidden_dtype,
         )
         scale_shift_timestep = scale_shift_timestep.view(batch_size, -1, scale_shift_timestep.shape[-1])
+
         gate_noise_timestep, _ = self.cross_gate_adaln(
-            timestep.flatten() * av_ca_factor,
+            (cross_modality_sigma * timestep_scale_multiplier * av_ca_factor).flatten(),
             hidden_dtype=hidden_dtype,
         )
         gate_noise_timestep = gate_noise_timestep.view(batch_size, -1, gate_noise_timestep.shape[-1])

packages/ltx-core/src/ltx_core/model/video_vae/__init__.py

@@ -1,5 +1,6 @@
 """Video VAE package."""
 
+from ltx_core.model.video_vae.memory_efficient_decode import MEMORY_EFFICIENT_DECODE
 from ltx_core.model.video_vae.model_configurator import (
     VAE_DECODER_COMFY_KEYS_FILTER,
     VAE_ENCODER_COMFY_KEYS_FILTER,
@@ -10,6 +11,7 @@ from ltx_core.model.video_vae.tiling import SpatialTilingConfig, TemporalTilingC
 from ltx_core.model.video_vae.video_vae import VideoDecoder, VideoEncoder, get_video_chunks_number
 
 __all__ = [
+    "MEMORY_EFFICIENT_DECODE",
     "VAE_DECODER_COMFY_KEYS_FILTER",
     "VAE_ENCODER_COMFY_KEYS_FILTER",
     "SpatialTilingConfig",

packages/ltx-core/src/ltx_core/model/video_vae/memory_efficient_decode.py

@@ -0,0 +1,662 @@
+"""Memory-efficient VAE decoder operations.
+Reduces peak VRAM usage during video decoding through in-place operations
+and workspace buffer reuse.  The main optimizations are:
+1. **Workspace buffers** -- Pre-allocated tensors with temporal padding replace
+   dynamic padding (``F.pad`` / ``concatenate``) in ``CausalConv3d``.  A
+   workspace of shape ``[B, C, T+2, H, W]`` holds the data in positions
+   ``[1:-1]`` with replicate padding at ``[0]`` and ``[-1]``.
+2. **In-place temporal-chunked Conv3d** *(non-causal only)* -- The convolution
+   output is written back into the workspace buffer, avoiding a separate
+   output allocation.  Temporal chunking with boundary save/restore ensures
+   correct reads despite in-place writes.
+3. **In-place normalization and affine transforms** -- PixelNorm, scale/shift,
+   and SiLU are applied in-place on workspace views.
+4. **Free-before-conv** -- For ``DepthToSpaceUpsample`` blocks the input
+   tensor is freed before the convolution runs so that peak VRAM never holds
+   input *and* output simultaneously.
+Both causal and non-causal modes are supported.  Non-causal mode benefits
+from all four optimizations.  Causal mode benefits from optimizations 1, 3,
+and 4; in-place conv (2) is skipped because the asymmetric causal padding
+layout prevents clean in-place overwrites.
+Usage via the ``ModuleOps`` pattern (preferred)::
+    from ltx_core.model.video_vae import MEMORY_EFFICIENT_DECODE
+    builder = decoder_builder.with_module_ops(
+        (*decoder_builder.module_ops, MEMORY_EFFICIENT_DECODE)
+    )
+Or applied directly to an existing decoder::
+    from ltx_core.model.video_vae.memory_efficient_decode import (
+        enable_memory_efficient_decode,
+    )
+    enable_memory_efficient_decode(decoder)
+"""
+
+from __future__ import annotations
+
+import math
+from typing import TYPE_CHECKING
+
+import torch
+from einops import rearrange
+from torch import nn
+from torch.nn import functional as F
+
+from ltx_core.loader.module_ops import ModuleOps
+from ltx_core.model.common.normalization import PixelNorm
+from ltx_core.model.video_vae.convolution import CausalConv3d
+from ltx_core.model.video_vae.ops import unpatchify
+from ltx_core.model.video_vae.resnet import ResnetBlock3D, UNetMidBlock3D
+from ltx_core.model.video_vae.sampling import DepthToSpaceUpsample
+
+if TYPE_CHECKING:
+    from ltx_core.model.video_vae.video_vae import VideoDecoder
+
+
+# ---------------------------------------------------------------------------
+# Low-level helpers
+# ---------------------------------------------------------------------------
+
+
+def _memory_format_of(t: torch.Tensor, prefer_channels_last_3d: bool = False) -> torch.memory_format:
+    """Pick the memory format for a workspace allocation.
+    When ``prefer_channels_last_3d`` is True and ``t`` is 5D, return
+    ``channels_last_3d`` regardless of ``t``'s current strides -- the
+    workspace's ``.copy_(t)`` will transcribe the data into the new layout.
+    This is needed because intermediate tensors inside the decoder (after
+    ``rearrange`` + slice + residual add in ``_upsample_forward_efficient``)
+    are not NHWC-contiguous, so an auto-detect helper would silently fall
+    back to NCHW for every workspace after the first upsample.
+    Otherwise fall back to inspecting ``t``: ``channels_last_3d`` if ``t``
+    already uses it, else contiguous.
+    """
+    if prefer_channels_last_3d and t.dim() == 5:
+        return torch.channels_last_3d
+    if t.dim() == 5 and t.is_contiguous(memory_format=torch.channels_last_3d):
+        return torch.channels_last_3d
+    return torch.contiguous_format
+
+
+def _find_temporal_split_size(num_frames: int) -> int:
+    """Find chunk size for in-place temporal convolution.
+    The chunk size ensures the last chunk has at least 3 frames
+    (the temporal kernel size), avoiding degenerate chunks.
+    """
+    for s in range(16, 2, -1):
+        remainder = num_frames % s
+        if remainder == 0 or remainder >= 3:
+            return s
+
+    raise ValueError(
+        f"Unable to find a valid temporal split size for num_frames={num_frames}. "
+        "Expected a split size between 3 and 16 such that the final chunk is "
+        "either exact or has at least 3 frames."
+    )
+
+
+def _pad_workspace_temporal(workspace: torch.Tensor) -> None:
+    """Apply non-causal replicate padding to temporal boundaries.
+    Sets ``workspace[:, :, 0]`` to a copy of ``workspace[:, :, 1]`` and
+    ``workspace[:, :, -1]`` to a copy of ``workspace[:, :, -2]``.
+    """
+    workspace[:, :, 0, :, :].copy_(workspace[:, :, 1, :, :])
+    workspace[:, :, -1, :, :].copy_(workspace[:, :, -2, :, :])
+
+
+# ---------------------------------------------------------------------------
+# In-place Conv3d (non-causal only)
+# ---------------------------------------------------------------------------
+
+
+def inplace_conv3d_temporal_chunked(workspace: torch.Tensor, conv: nn.Conv3d) -> None:
+    """Run a 3x3x3 Conv3d in-place on a temporally-padded workspace.
+    The workspace has shape ``[B, C, T+2, H, W]`` where positions ``[1:-1]``
+    hold the real data and positions ``[0]`` and ``[-1]`` are padding slots.
+    The convolution must have ``kernel_size=(3,3,3)``, ``stride=(1,1,1)``,
+    ``padding=(0,1,1)`` -- no temporal padding, symmetric spatial padding.
+    The output (T frames) overwrites positions ``[1:-1]``.  Temporal chunking
+    with boundary save/restore ensures each chunk reads unmodified input even
+    though earlier chunks already wrote to the same buffer.
+    Only valid for **non-causal** mode (symmetric replicate padding).
+    Args:
+        workspace: Tensor ``[B, max(C_in, C_out), T+2, H, W]``.
+            Modified in-place; after the call ``workspace[:, :C_out, 1:-1]``
+            holds the convolution result.
+        conv: ``nn.Conv3d`` with the constraints above.
+    """
+    if conv.kernel_size != (3, 3, 3):
+        raise ValueError(f"Expected kernel_size=(3,3,3), got {conv.kernel_size}")
+    if conv.stride != (1, 1, 1):
+        raise ValueError(f"Expected stride=(1,1,1), got {conv.stride}")
+    if conv.padding != (0, 1, 1):
+        raise ValueError(f"Expected padding=(0,1,1), got {conv.padding}")
+
+    _pad_workspace_temporal(workspace)
+
+    total_frames = workspace.shape[2]
+    out_channels = conv.out_channels
+    in_channels = conv.in_channels
+
+    if total_frames > 16:
+        split_size = _find_temporal_split_size(total_frames)
+        num_splits = (total_frames + split_size - 1) // split_size
+    else:
+        split_size = total_frames - 1
+        num_splits = 1
+
+    # 1-frame buffers for saving / restoring boundary frames across chunks.
+    x_buf = torch.empty(
+        workspace.shape[0],
+        workspace.shape[1],
+        1,
+        workspace.shape[3],
+        workspace.shape[4],
+        device=workspace.device,
+        dtype=workspace.dtype,
+        memory_format=_memory_format_of(workspace),
+    )
+    o_buf = torch.empty_like(x_buf)
+
+    # Helper: extract a chunk and make it contiguous.  Workspace views can
+    # inherit strides > 2^31 from the full buffer, which makes Conv3d's
+    # reflect-padding path (F.pad) crash with "input tensor must fit into
+    # 32-bit index math".  A small .clone() per chunk avoids this.
+    needs_clone = workspace.untyped_storage().nbytes() > (2**31 - 1) * workspace.element_size()
+
+    def _chunk(t_start: int, t_end: int) -> torch.Tensor:
+        s = workspace[:, :in_channels, t_start:t_end]
+        return s.clone() if needs_clone else s
+
+    # --- First chunk ---
+    if num_splits > 1:
+        # Save the boundary now so the loop below can restore it.  Skipped
+        # when there is only one chunk: the loop never runs, and the save
+        # would be a wasted full HW slice copy.
+        x_buf[:, :, 0] = workspace[:, :, split_size - 1].clone()
+    workspace[:, :out_channels, 1:split_size] = conv(_chunk(0, split_size + 1))
+
+    # --- Remaining chunks ---
+    for i in range(1, num_splits):
+        start = i * split_size
+        end = min((i + 1) * split_size, total_frames - 1)
+
+        # Save the value at start-1 (now holds previous chunk's output).
+        o_buf[:, :, 0] = workspace[:, :, start - 1].clone()
+        # Restore the original input value needed by this chunk's conv.
+        workspace[:, :, start - 1] = x_buf[:, :, 0]
+        # Save the boundary for the *next* chunk before we overwrite it.
+        x_buf[:, :, 0] = workspace[:, :, end - 1].clone()
+
+        workspace[:, :out_channels, start:end] = conv(_chunk(start - 1, end + 1))
+
+        # Put back the previous chunk's output at the boundary.
+        workspace[:, :, start - 1] = o_buf[:, :, 0]
+
+
+# ---------------------------------------------------------------------------
+# Causal conv helper (free-before-conv)
+# ---------------------------------------------------------------------------
+
+
+def _causal_pad(x: torch.Tensor, pad_size: int) -> torch.Tensor:
+    """Build a causal-padded buffer of shape ``[B, C, T+pad_size, H, W]``.
+    Copies ``x`` into ``padded[:, :, pad_size:]`` and replicates the first
+    real frame into the leading ``pad_size`` slots.  The caller still owns
+    ``x`` after this returns.
+    """
+    padded = torch.empty(
+        x.shape[0],
+        x.shape[1],
+        x.shape[2] + pad_size,
+        x.shape[3],
+        x.shape[4],
+        device=x.device,
+        dtype=x.dtype,
+        memory_format=_memory_format_of(x),
+    )
+    padded[:, :, pad_size:].copy_(x)
+    for i in range(pad_size):
+        padded[:, :, i] = padded[:, :, pad_size]
+    return padded
+
+
+def _causal_pad_free_and_conv(x: torch.Tensor, causal_conv: CausalConv3d) -> torch.Tensor:
+    """Causal-pad *x*, free it, then run the raw ``nn.Conv3d``.
+    This avoids the peak where both the original and padded tensors are
+    live simultaneously (as happens inside ``CausalConv3d.forward``).
+    Args:
+        x: Input ``[B, C_in, T, H, W]``.  **Deleted** inside this function;
+            the caller must not use it afterwards.
+    Returns:
+        Convolution output ``[B, C_out, T, H, W]``.
+    """
+    padded = _causal_pad(x, causal_conv.time_kernel_size - 1)
+    del x
+    result = causal_conv.conv(padded)
+    del padded
+    return result
+
+
+# ---------------------------------------------------------------------------
+# In-place normalization
+# ---------------------------------------------------------------------------
+
+
+def _pixel_norm_inplace(x: torch.Tensor, eps: float = 1e-8) -> None:
+    """In-place RMS (pixel) normalization along the channel dimension."""
+    rms = torch.sqrt(torch.mean(x**2, dim=1, keepdim=True) + eps)
+    x.div_(rms)
+
+
+def _norm_inplace(norm: nn.Module, x: torch.Tensor) -> None:
+    """Apply *norm* in-place, using an optimised path for ``PixelNorm``."""
+    if isinstance(norm, PixelNorm):
+        _pixel_norm_inplace(x, eps=norm.eps)
+    else:
+        # GroupNorm or other -- fall back to allocating a temporary.
+        result = norm(x)
+        x.copy_(result)
+        del result
+
+
+# ---------------------------------------------------------------------------
+# Per-block efficient forwards
+# ---------------------------------------------------------------------------
+
+
+def _resnet_block_forward_inplace(
+    resnet: ResnetBlock3D,
+    workspace: torch.Tensor,
+    causal: bool,
+    timestep: torch.Tensor | None,
+    generator: torch.Generator | None,
+) -> None:
+    """Run a ``ResnetBlock3D`` in-place on a workspace buffer.
+    The workspace has shape ``[B, C, T+2, H, W]`` with real data in
+    ``[1:-1]``.  After this call ``workspace[:, :, 1:-1]`` holds the
+    residual-branch output ``F(x)`` (without the skip connection --
+    the caller adds it back to the hidden state).
+    Only valid when ``in_channels == out_channels`` (true for all
+    ``ResnetBlock3D`` instances inside a ``UNetMidBlock3D``).
+    """
+    if resnet.in_channels != resnet.out_channels:
+        raise ValueError(
+            "In-place resnet forward requires in_channels == out_channels, "
+            f"got {resnet.in_channels} != {resnet.out_channels}"
+        )
+
+    interior = workspace[:, :, 1:-1]
+
+    # --- norm1 + [ada scaling] + SiLU + conv1 ---
+    _norm_inplace(resnet.norm1, interior)
+
+    if resnet.timestep_conditioning and timestep is not None:
+        ada = resnet.scale_shift_table[None, ..., None, None, None].to(
+            device=interior.device, dtype=interior.dtype
+        ) + timestep.reshape(
+            interior.shape[0],
+            4,
+            -1,
+            timestep.shape[-3],
+            timestep.shape[-2],
+            timestep.shape[-1],
+        )
+        shift1, scale1, shift2, scale2 = ada.unbind(dim=1)
+        interior.mul_(1 + scale1).add_(shift1)
+
+    F.silu(interior, inplace=True)
+
+    if causal:
+        result = resnet.conv1(interior, causal=True)
+        interior.copy_(result)
+        del result
+    else:
+        inplace_conv3d_temporal_chunked(workspace, resnet.conv1.conv)
+
+    if resnet.inject_noise:
+        spatial_shape = interior.shape[-2:]
+        scale = resnet.per_channel_scale1.to(device=interior.device, dtype=interior.dtype)
+        noise = torch.randn(spatial_shape, device=interior.device, dtype=interior.dtype, generator=generator)
+        interior.add_((noise * scale)[None, :, None, ...])
+
+    # --- norm2 + [ada scaling] + SiLU + conv2 ---
+    _norm_inplace(resnet.norm2, interior)
+
+    if resnet.timestep_conditioning and timestep is not None:
+        interior.mul_(1 + scale2).add_(shift2)  # type: ignore[possibly-undefined]
+
+    F.silu(interior, inplace=True)
+    # dropout is always 0.0 during inference -- skip.
+
+    if causal:
+        result = resnet.conv2(interior, causal=True)
+        interior.copy_(result)
+        del result
+    else:
+        inplace_conv3d_temporal_chunked(workspace, resnet.conv2.conv)
+
+    if resnet.inject_noise:
+        spatial_shape = interior.shape[-2:]
+        scale = resnet.per_channel_scale2.to(device=interior.device, dtype=interior.dtype)
+        noise = torch.randn(spatial_shape, device=interior.device, dtype=interior.dtype, generator=generator)
+        interior.add_((noise * scale)[None, :, None, ...])
+
+
+def _midblock_forward_efficient(
+    block: UNetMidBlock3D,
+    hidden_states: torch.Tensor,
+    causal: bool,
+    timestep: torch.Tensor | None,
+    generator: torch.Generator | None,
+    prefer_channels_last_3d: bool = False,
+) -> torch.Tensor:
+    """Memory-efficient ``UNetMidBlock3D`` forward.
+    Allocates a single workspace buffer that is reused across all
+    ``ResnetBlock3D`` iterations.  For each block the workspace is
+    populated with the current hidden state, processed in-place, and
+    the result is added back (residual connection).
+    """
+    timestep_embed = None
+    if block.timestep_conditioning:
+        if timestep is None:
+            raise ValueError("'timestep' required when timestep_conditioning=True")
+        batch_size = hidden_states.shape[0]
+        timestep_embed = block.time_embedder(
+            timestep=timestep.flatten(),
+            hidden_dtype=hidden_states.dtype,
+        )
+        timestep_embed = timestep_embed.view(batch_size, timestep_embed.shape[-1], 1, 1, 1)
+
+    workspace = torch.empty(
+        hidden_states.shape[0],
+        hidden_states.shape[1],
+        hidden_states.shape[2] + 2,
+        hidden_states.shape[3],
+        hidden_states.shape[4],
+        device=hidden_states.device,
+        dtype=hidden_states.dtype,
+        memory_format=_memory_format_of(hidden_states, prefer_channels_last_3d),
+    )
+
+    for resnet in block.res_blocks:
+        workspace[:, :, 1:-1].copy_(hidden_states)
+        _resnet_block_forward_inplace(resnet, workspace, causal, timestep_embed, generator)
+        hidden_states.add_(workspace[:, :, 1:-1])
+
+    del workspace
+    return hidden_states
+
+
+def _upsample_forward_efficient(
+    block: DepthToSpaceUpsample,
+    x: torch.Tensor,
+    causal: bool,
+    prefer_channels_last_3d: bool = False,
+) -> torch.Tensor:
+    """Memory-efficient ``DepthToSpaceUpsample`` forward.
+    For non-causal mode the input is copied into a workspace and the
+    convolution runs in-place.  For causal mode the input is manually
+    padded and freed before the convolution runs.  Both paths avoid
+    the peak where input *and* output coexist.
+    """
+    if block.residual:
+        x_in = rearrange(
+            x,
+            "b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)",
+            p1=block.stride[0],
+            p2=block.stride[1],
+            p3=block.stride[2],
+        )
+        num_repeat = math.prod(block.stride) // block.out_channels_reduction_factor
+        x_in = x_in.repeat(1, num_repeat, 1, 1, 1)
+        if block.stride[0] == 2:
+            x_in = x_in[:, :, 1:, :, :]
+
+    conv = block.conv.conv  # underlying nn.Conv3d inside CausalConv3d
+    in_channels = x.shape[1]
+    out_channels = conv.out_channels
+
+    if causal:
+        x = _causal_pad_free_and_conv(x, block.conv)
+    else:
+        mem_fmt = _memory_format_of(x, prefer_channels_last_3d)
+        workspace = torch.empty(
+            x.shape[0],
+            max(in_channels, out_channels),
+            x.shape[2] + 2,
+            x.shape[3],
+            x.shape[4],
+            device=x.device,
+            dtype=x.dtype,
+            memory_format=mem_fmt,
+        )
+        workspace[:, :in_channels, 1:-1].copy_(x)
+        del x
+        inplace_conv3d_temporal_chunked(workspace, conv)
+        x = workspace[:, :out_channels, 1:-1].contiguous(memory_format=mem_fmt)
+        del workspace
+
+    x = rearrange(
+        x,
+        "b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)",
+        p1=block.stride[0],
+        p2=block.stride[1],
+        p3=block.stride[2],
+    )
+    if block.stride[0] == 2:
+        x = x[:, :, 1:, :, :]
+    if block.residual:
+        x.add_(x_in)
+        del x_in
+    return x
+
+
+# ---------------------------------------------------------------------------
+# Final norm + conv_out
+# ---------------------------------------------------------------------------
+
+
+def _final_norm_and_conv_out(
+    decoder: VideoDecoder,
+    sample: torch.Tensor,
+    causal: bool,
+    scaled_timestep: torch.Tensor | None,
+    batch_size: int,
+    prefer_channels_last_3d: bool = False,
+) -> torch.Tensor:
+    """Workspace-based final norm + [ada] + SiLU + conv_out + unpatchify."""
+    conv_out_mod: CausalConv3d = decoder.conv_out  # type: ignore[assignment]
+    conv_out = conv_out_mod.conv
+    feature_channels = sample.shape[1]
+
+    mem_fmt = _memory_format_of(sample, prefer_channels_last_3d)
+    workspace = torch.empty(
+        sample.shape[0],
+        max(feature_channels, conv_out.out_channels),
+        sample.shape[2] + 2,
+        sample.shape[3],
+        sample.shape[4],
+        device=sample.device,
+        dtype=sample.dtype,
+        memory_format=mem_fmt,
+    )
+    workspace[:, :feature_channels, 1:-1].copy_(sample)
+    del sample
+
+    interior = workspace[:, :feature_channels, 1:-1]
+    _norm_inplace(decoder.conv_norm_out, interior)
+
+    if decoder.timestep_conditioning:
+        embedded_timestep = decoder.last_time_embedder(
+            timestep=scaled_timestep.flatten(),
+            hidden_dtype=interior.dtype,
+        )
+        embedded_timestep = embedded_timestep.view(batch_size, embedded_timestep.shape[-1], 1, 1, 1)
+        ada_values = decoder.last_scale_shift_table[None, ..., None, None, None].to(
+            device=interior.device, dtype=interior.dtype
+        ) + embedded_timestep.reshape(
+            batch_size,
+            2,
+            -1,
+            embedded_timestep.shape[-3],
+            embedded_timestep.shape[-2],
+            embedded_timestep.shape[-1],
+        )
+        shift, scale = ada_values.unbind(dim=1)
+        interior.mul_(1 + scale).add_(shift)
+
+    F.silu(interior, inplace=True)
+
+    if causal:
+        # Causal: build padded tensor directly from the interior view,
+        # then free the workspace before running the conv.
+        padded = _causal_pad(interior, conv_out_mod.time_kernel_size - 1)
+        del workspace, interior
+        result = conv_out(padded)
+        del padded
+    else:
+        inplace_conv3d_temporal_chunked(workspace, conv_out)
+        result = workspace[:, : conv_out.out_channels, 1:-1].contiguous(memory_format=mem_fmt)
+        del workspace, interior
+
+    return unpatchify(result, patch_size_hw=decoder.patch_size, patch_size_t=1)
+
+
+# ---------------------------------------------------------------------------
+# Top-level efficient decoder forward
+# ---------------------------------------------------------------------------
+
+
+def _memory_efficient_forward(
+    decoder: VideoDecoder,
+    sample: torch.Tensor,
+    timestep: torch.Tensor | None = None,
+    generator: torch.Generator | None = None,
+) -> torch.Tensor:
+    """Full memory-efficient ``VideoDecoder.forward`` replacement.
+    Orchestrates the entire decode through workspace-based operations:
+    ``UNetMidBlock3D`` and ``DepthToSpaceUpsample`` blocks use efficient
+    paths; standalone ``ResnetBlock3D`` blocks fall back to the standard
+    forward.  The final norm + ada + SiLU + conv_out is also workspace-based.
+    All workspaces are allocated ``channels_last_3d`` so cuDNN's NHWC 3D
+    conv kernels run end-to-end. The caller (:func:`enable_memory_efficient_decode`)
+    is responsible for converting the input sample and decoder weights to NHWC.
+    """
+    causal = decoder.causal
+    batch_size = sample.shape[0]
+    sample = sample.to(next(decoder.parameters()).dtype)
+
+    # --- Noise injection and de-normalisation (identical to standard path) ---
+    if decoder.timestep_conditioning:
+        noise = (
+            torch.randn(sample.size(), generator=generator, dtype=sample.dtype, device=sample.device)
+            * decoder.decode_noise_scale
+        )
+        sample = noise + (1.0 - decoder.decode_noise_scale) * sample
+
+    sample = decoder.per_channel_statistics.un_normalize(sample)
+
+    if timestep is None and decoder.timestep_conditioning:
+        timestep = torch.full((batch_size,), decoder.decode_timestep, device=sample.device, dtype=sample.dtype)
+
+    # --- conv_in (latent tensor is small -- standard path is fine) ---
+    sample = decoder.conv_in(sample, causal=causal)
+
+    upscale_dtype = next(iter(decoder.up_blocks.parameters())).dtype
+    sample = sample.to(upscale_dtype)
+
+    scaled_timestep = None
+    if decoder.timestep_conditioning:
+        if timestep is None:
+            raise ValueError("'timestep' required when timestep_conditioning=True")
+        scaled_timestep = timestep * decoder.timestep_scale_multiplier.to(sample)
+
+    # Workspaces are unconditionally NHWC: rearrange + slice + residual-add
+    # inside _upsample_forward_efficient produces NCHW-default output, so
+    # per-tensor inspection would silently fall back to NCHW for every
+    # workspace after the first upsample.
+
+    # --- Up blocks (dispatch to efficient path per block type) ---
+    for up_block in decoder.up_blocks:
+        if isinstance(up_block, UNetMidBlock3D):
+            sample = _midblock_forward_efficient(
+                up_block,
+                sample,
+                causal=causal,
+                timestep=scaled_timestep if decoder.timestep_conditioning else None,
+                generator=generator,
+                prefer_channels_last_3d=True,
+            )
+        elif isinstance(up_block, DepthToSpaceUpsample):
+            sample = _upsample_forward_efficient(up_block, sample, causal=causal, prefer_channels_last_3d=True)
+        elif isinstance(up_block, ResnetBlock3D):
+            sample = up_block(sample, causal=causal, generator=generator)
+        else:
+            sample = up_block(sample, causal=causal)
+
+    return _final_norm_and_conv_out(decoder, sample, causal, scaled_timestep, batch_size, prefer_channels_last_3d=True)
+
+
+# ---------------------------------------------------------------------------
+# Public API
+# ---------------------------------------------------------------------------
+
+
+def enable_memory_efficient_decode(decoder: nn.Module) -> nn.Module:
+    """Patch a ``VideoDecoder`` to use the memory-efficient forward path.
+    The mem-efficient path runs the decoder in ``channels_last_3d`` memory
+    format: weights and inputs are converted on first call so cuDNN's NHWC
+    3D conv kernels are used (~2x faster, avoids the large vol2col scratch
+    buffer of the NCHW path).
+    The original ``forward`` is saved as ``decoder._original_forward`` so
+    that it can be restored later with :func:`disable_memory_efficient_decode`.
+    """
+    # Import here to avoid circular dependency at module level.
+    from ltx_core.model.video_vae.video_vae import VideoDecoder  # noqa: PLC0415
+
+    if not isinstance(decoder, VideoDecoder):
+        raise TypeError(f"Expected VideoDecoder, got {type(decoder).__name__}")
+
+    if hasattr(decoder, "_original_forward"):
+        return decoder
+
+    original_forward = decoder.forward
+    weights_converted = False
+
+    def efficient_forward(
+        sample: torch.Tensor,
+        timestep: torch.Tensor | None = None,
+        generator: torch.Generator | None = None,
+    ) -> torch.Tensor:
+        nonlocal weights_converted
+        if sample.dim() == 5:
+            if not weights_converted:
+                # Lazy: weights are real by first-call time (meta -> loader -> here).
+                decoder.to(memory_format=torch.channels_last_3d)
+                weights_converted = True
+            sample = sample.to(memory_format=torch.channels_last_3d)
+        return _memory_efficient_forward(decoder, sample, timestep, generator)
+
+    decoder._original_forward = original_forward  # type: ignore[attr-defined]
+    decoder.forward = efficient_forward  # type: ignore[assignment]
+    return decoder
+
+
+def disable_memory_efficient_decode(decoder: nn.Module) -> nn.Module:
+    """Restore the original ``forward`` method on a patched ``VideoDecoder``."""
+    if hasattr(decoder, "_original_forward"):
+        decoder.forward = decoder._original_forward  # type: ignore[attr-defined]
+        del decoder._original_forward  # type: ignore[attr-defined]
+    return decoder
+
+
+def _is_video_decoder(model: nn.Module) -> bool:
+    """Matcher for the ``MEMORY_EFFICIENT_DECODE`` module op."""
+    from ltx_core.model.video_vae.video_vae import VideoDecoder  # noqa: PLC0415
+
+    return isinstance(model, VideoDecoder)
+
+
+MEMORY_EFFICIENT_DECODE = ModuleOps(
+    name="memory_efficient_vae_decode",
+    matcher=_is_video_decoder,
+    mutator=enable_memory_efficient_decode,
+)

packages/ltx-core/src/ltx_core/model/video_vae/tiling.py

@@ -64,6 +64,6 @@ class TilingConfig:
     @classmethod
     def default(cls) -> "TilingConfig":
         return cls(
-            spatial_config=SpatialTilingConfig(tile_size_in_pixels=512, tile_overlap_in_pixels=64),
-            temporal_config=TemporalTilingConfig(tile_size_in_frames=64, tile_overlap_in_frames=24),
+            spatial_config=SpatialTilingConfig(tile_size_in_pixels=768, tile_overlap_in_pixels=64),
+            temporal_config=TemporalTilingConfig(tile_size_in_frames=80, tile_overlap_in_frames=24),
         )

packages/ltx-core/src/ltx_core/model/video_vae/video_vae.py

@@ -259,6 +259,7 @@ class VideoEncoder(nn.Module):
         Args:
             sample: Input video (B, C, F, H, W). F should be 1 + 8*k (e.g., 1, 9, 17, 25, 33...).
                 If not, the encoder crops the last frames to the nearest valid length.
+                Should be normalized to [-1, 1] range before encoding.
         Returns:
             Normalized latent means (B, 128, F', H', W') where F' = 1+(F-1)/8, H' = H/32, W' = W/32.
             Example: (B, 3, 33, 512, 512) -> (B, 128, 5, 16, 16).
@@ -605,8 +606,8 @@ class VideoDecoder(nn.Module):
         # many video frames and pixels correspond to a single latent cell.
         self.video_downscale_factors = SpatioTemporalScaleFactors(
             time=8,
-            width=32,
             height=32,
+            width=32,
         )
 
         self.patch_size = patch_size
@@ -698,6 +699,9 @@ class VideoDecoder(nn.Module):
             When causal=False, allows future frame dependencies in convolutions but maintains same output shape.
         """
         batch_size = sample.shape[0]
+        output_dtype = sample.dtype
+        weights_dtype = next(self.parameters()).dtype
+        sample = sample.to(weights_dtype)
 
         # Add noise if timestep conditioning is enabled
         if self.timestep_conditioning:
@@ -770,7 +774,7 @@ class VideoDecoder(nn.Module):
         # Example: (B, 48, F, 128, 128) -> (B, 3, F, 512, 512) with patch_size=4
         sample = unpatchify(sample, patch_size_hw=self.patch_size, patch_size_t=1)
 
-        return sample
+        return sample.to(output_dtype)
 
     def _prepare_tiles(
         self,
@@ -903,22 +907,21 @@ class VideoDecoder(nn.Module):
         tiling_config: TilingConfig | None = None,
         generator: torch.Generator | None = None,
     ) -> Iterator[torch.Tensor]:
-        """Decode a video latent tensor, yielding uint8 chunks ``[f, h, w, c]``.
+        """Decode a video latent tensor, yielding float chunks ``[f, h, w, c]`` in ``[0, 1]``.
         Subclasses (e.g. ``DistributedVideoDecoder``) may override this to
         control eagerness or distribution across ranks.
         """
 
-        def convert_to_uint8(frames: torch.Tensor) -> torch.Tensor:
-            frames = (((frames + 1.0) / 2.0).clamp(0.0, 1.0) * 255.0).to(torch.uint8)
-            frames = rearrange(frames[0], "c f h w -> f h w c")
-            return frames
+        def to_rgb(frames: torch.Tensor) -> torch.Tensor:
+            video = rearrange(frames[0], "c f h w -> f h w c")
+            return video.add_(1.0).mul_(0.5).clamp_(0.0, 1.0)
 
         if tiling_config is not None:
             for frames in self.tiled_decode(latent, tiling_config, generator=generator):
-                yield convert_to_uint8(frames)
+                yield to_rgb(frames)
         else:
             decoded = self(latent, generator=generator)
-            yield convert_to_uint8(decoded)
+            yield to_rgb(decoded)
 
     def _group_tiles_by_temporal_slice(self, tiles: List[Tile]) -> List[List[Tile]]:
         """Group tiles by their temporal output slice."""

packages/ltx-core/src/ltx_core/quantization/__init__.py

@@ -2,15 +2,16 @@ from ltx_core.quantization.fp8_cast import (
     TRANSFORMER_LINEAR_DOWNCAST_MAP,
     UPCAST_DURING_INFERENCE,
     UpcastWithStochasticRounding,
+    fp8_cast_fuse_rule,
 )
-from ltx_core.quantization.fp8_scaled_mm import FP8_PREPARE_MODULE_OPS, FP8_TRANSPOSE_SD_OPS
+from ltx_core.quantization.fp8_scaled_mm import fp8_scaled_mm_fuse_rule
 from ltx_core.quantization.policy import QuantizationPolicy
 
 __all__ = [
-    "FP8_PREPARE_MODULE_OPS",
-    "FP8_TRANSPOSE_SD_OPS",
     "TRANSFORMER_LINEAR_DOWNCAST_MAP",
     "UPCAST_DURING_INFERENCE",
     "QuantizationPolicy",
     "UpcastWithStochasticRounding",
+    "fp8_cast_fuse_rule",
+    "fp8_scaled_mm_fuse_rule",
 ]

packages/ltx-core/src/ltx_core/quantization/fp8_cast.py

@@ -1,14 +1,26 @@
+from pathlib import Path
+
+import safetensors
 import torch
 
+from ltx_core.loader.fuse_loras import FuseRule, bf16_fuse_rule
+from ltx_core.loader.kernels import TRITON_AVAILABLE
 from ltx_core.loader.module_ops import ModuleOps
+from ltx_core.loader.primitives import StateDict
 from ltx_core.loader.sd_ops import KeyValueOperationResult, SDOps
 from ltx_core.model.transformer.model import LTXModel
+from ltx_core.quantization.policy import QuantizationPolicy
 
 BLOCK_SIZE = 1024
 
 
-def _fused_add_round_launch(target_weight: torch.Tensor, original_weight: torch.Tensor, seed: int) -> torch.Tensor:
-    # Lazy import triton - only available on CUDA platforms
+def fused_add_round_launch(target_weight: torch.Tensor, original_weight: torch.Tensor, seed: int) -> torch.Tensor:
+    if not TRITON_AVAILABLE:
+        raise RuntimeError(
+            "fused_add_round_launch requires Triton, which is not available on this platform. "
+            "Callers should gate on ltx_core.loader.kernels.TRITON_AVAILABLE and use a "
+            "deterministic-rounding fallback instead."
+        )
     import triton  # noqa: PLC0415
 
     from ltx_core.loader.kernels import fused_add_round_kernel  # noqa: PLC0415
@@ -53,10 +65,13 @@ def _upcast_and_round(
     """
     Upcast the weight to the given dtype and optionally apply stochastic rounding.
     Input weight needs to have float8_e4m3fn or float8_e5m2 dtype.
+    Stochastic rounding is implemented via a Triton kernel. When Triton is not
+    available (e.g., on Windows), this falls back to deterministic (nearest)
+    rounding via ``weight.to(dtype)``.
     """
-    if not with_stochastic_rounding:
+    if not with_stochastic_rounding or not TRITON_AVAILABLE or weight.device.type != "cuda":
         return weight.to(dtype)
-    return _fused_add_round_launch(torch.zeros_like(weight, dtype=dtype), weight, seed)
+    return fused_add_round_launch(torch.zeros_like(weight, dtype=dtype), weight, seed)
 
 
 class Fp8CastLinear(torch.nn.Linear):
@@ -82,68 +97,87 @@ class Fp8CastLinear(torch.nn.Linear):
 def _replace_fwd_with_upcast(layer: torch.nn.Linear, with_stochastic_rounding: bool = False, seed: int = 0) -> None:
     """
     Intended to be applied via __class__ reassignment to existing nn.Linear
-    instances so that their parameter and buffer tensors are preserved in-place,
-    avoiding re-instantiation. Forward remains defined at the class level, which
-    is required for torch.compile compatibility — instance-level closure
-    monkey-patches cause graph breaks.
+    instances. Forward remains defined at the class level, which is required for
+    torch.compile compatibility — instance-level closure monkey-patches cause
+    graph breaks.
+    Also retypes ``weight`` and ``bias`` to fp8 so the meta param dtype matches
+    the post-load tensor dtype (sd_ops downcasts checkpoint bf16 -> fp8 at load).
+    Block streaming relies on this to derive pool buffer layout from the meta
+    model without an eager checkpoint read.
     """
     layer.__class__ = Fp8CastLinear
     layer._with_stochastic_rounding = with_stochastic_rounding
     layer._seed = seed
+    layer.weight = torch.nn.Parameter(
+        torch.empty(layer.weight.shape, dtype=torch.float8_e4m3fn, device=layer.weight.device),
+        requires_grad=layer.weight.requires_grad,
+    )
+    if layer.bias is not None:
+        layer.bias = torch.nn.Parameter(
+            torch.empty(layer.bias.shape, dtype=torch.float8_e4m3fn, device=layer.bias.device),
+            requires_grad=layer.bias.requires_grad,
+        )
+
+
+# Module-name suffixes for the Linears that participate in fp8 cast. Used by
+# both the upcast matcher and the sd_ops downcast map so the two cannot drift.
+# - ``.to_q`` / ``.to_k`` / ``.to_v`` / ``.to_out.0`` have a leading dot so they
+#   only match the attention Linears at ``...attnN.to_q`` etc.
+# - ``ff.net.0.proj`` / ``ff.net.2`` are intentionally **dotless** so they match
+#   both video FF (``...ff.net.0.proj``) and audio FF (``...audio_ff.net.0.proj``).
+_FP8_CAST_KEY_PREFIX = "transformer_blocks."
+_FP8_CAST_LINEAR_SUFFIXES: tuple[str, ...] = (
+    ".to_q",
+    ".to_k",
+    ".to_v",
+    ".to_out.0",
+    "ff.net.0.proj",
+    "ff.net.2",
+)
+
+
+def _is_fp8_cast_linear(module_name: str) -> bool:
+    """Return True if *module_name* names a Linear that should be fp8-cast."""
+    if _FP8_CAST_KEY_PREFIX not in module_name:
+        return False
+    return any(module_name.endswith(suffix) for suffix in _FP8_CAST_LINEAR_SUFFIXES)
 
 
 def _amend_forward_with_upcast(
     model: torch.nn.Module, with_stochastic_rounding: bool = False, seed: int = 0
 ) -> torch.nn.Module:
     """
-    Replace the forward method of the model's Linear layers to forward
-    with upcast and optional stochastic rounding.
+    Replace the forward method of the fp8-cast Linear layers (per
+    :data:`_FP8_CAST_LINEAR_SUFFIXES`) to forward with upcast and optional
+    stochastic rounding.
+    Only the Linears whose weights are downcast by :data:`TRANSFORMER_LINEAR_DOWNCAST_MAP`
+    are retyped. Linears outside that subset (e.g. ``to_gate_logits``) are left as
+    plain ``nn.Linear`` so the meta-model param dtype matches the loaded checkpoint
+    dtype.
     """
-    for m in model.modules():
-        if isinstance(m, (torch.nn.Linear)):
+    for name, m in model.named_modules():
+        if isinstance(m, torch.nn.Linear) and _is_fp8_cast_linear(name):
             _replace_fwd_with_upcast(m, with_stochastic_rounding, seed)
     return model
 
 
-TRANSFORMER_LINEAR_DOWNCAST_MAP = (
-    SDOps("TRANSFORMER_LINEAR_DOWNCAST_MAP")
-    .with_kv_operation(
-        key_prefix="transformer_blocks.", key_suffix=".to_q.weight", operation=_naive_weight_or_bias_downcast
-    )
-    .with_kv_operation(
-        key_prefix="transformer_blocks.", key_suffix=".to_q.bias", operation=_naive_weight_or_bias_downcast
-    )
-    .with_kv_operation(
-        key_prefix="transformer_blocks.", key_suffix=".to_k.weight", operation=_naive_weight_or_bias_downcast
-    )
-    .with_kv_operation(
-        key_prefix="transformer_blocks.", key_suffix=".to_k.bias", operation=_naive_weight_or_bias_downcast
-    )
-    .with_kv_operation(
-        key_prefix="transformer_blocks.", key_suffix=".to_v.weight", operation=_naive_weight_or_bias_downcast
-    )
-    .with_kv_operation(
-        key_prefix="transformer_blocks.", key_suffix=".to_v.bias", operation=_naive_weight_or_bias_downcast
-    )
-    .with_kv_operation(
-        key_prefix="transformer_blocks.", key_suffix=".to_out.0.weight", operation=_naive_weight_or_bias_downcast
-    )
-    .with_kv_operation(
-        key_prefix="transformer_blocks.", key_suffix=".to_out.0.bias", operation=_naive_weight_or_bias_downcast
-    )
-    .with_kv_operation(
-        key_prefix="transformer_blocks.", key_suffix="ff.net.0.proj.weight", operation=_naive_weight_or_bias_downcast
-    )
-    .with_kv_operation(
-        key_prefix="transformer_blocks.", key_suffix="ff.net.0.proj.bias", operation=_naive_weight_or_bias_downcast
-    )
-    .with_kv_operation(
-        key_prefix="transformer_blocks.", key_suffix="ff.net.2.weight", operation=_naive_weight_or_bias_downcast
-    )
-    .with_kv_operation(
-        key_prefix="transformer_blocks.", key_suffix="ff.net.2.bias", operation=_naive_weight_or_bias_downcast
-    )
-)
+def _build_transformer_linear_downcast_map() -> SDOps:
+    """Build the sd_ops downcast map from the same suffix registry as the matcher."""
+    ops = SDOps("TRANSFORMER_LINEAR_DOWNCAST_MAP")
+    for suffix in _FP8_CAST_LINEAR_SUFFIXES:
+        ops = ops.with_kv_operation(
+            key_prefix=_FP8_CAST_KEY_PREFIX,
+            key_suffix=suffix + ".weight",
+            operation=_naive_weight_or_bias_downcast,
+        ).with_kv_operation(
+            key_prefix=_FP8_CAST_KEY_PREFIX,
+            key_suffix=suffix + ".bias",
+            operation=_naive_weight_or_bias_downcast,
+        )
+    return ops
+
+
+TRANSFORMER_LINEAR_DOWNCAST_MAP = _build_transformer_linear_downcast_map()
 
 UPCAST_DURING_INFERENCE = ModuleOps(
     name="upcast_fp8_during_linear_forward",
@@ -165,3 +199,139 @@ class UpcastWithStochasticRounding(ModuleOps):
             matcher=lambda model: isinstance(model, LTXModel),
             mutator=lambda model: _amend_forward_with_upcast(model, True, seed),
         )
+
+
+def fuse_cast_fp8_weight(
+    delta_bf16: torch.Tensor,
+    weight_fp8: torch.Tensor,
+) -> torch.Tensor:
+    """Return ``(delta_bf16 + dequantize(weight_fp8)).to(weight_fp8.dtype)``.
+    CUDA with Triton uses stochastic rounding via the fused kernel; otherwise
+    falls back to a deterministic bf16 add. ``delta_bf16`` is the bf16
+    accumulator and is mutated in place.
+    """
+    if delta_bf16.dtype != torch.bfloat16:
+        raise ValueError(f"delta_bf16 must be bfloat16, got {delta_bf16.dtype}")
+    if str(weight_fp8.device).startswith("cuda") and TRITON_AVAILABLE:
+        fused_add_round_launch(delta_bf16, weight_fp8, seed=0)
+    else:
+        delta_bf16.add_(weight_fp8.to(dtype=torch.bfloat16))
+    return delta_bf16.to(dtype=weight_fp8.dtype)
+
+
+def _fp8_cast_fuse(
+    key: str,
+    weight: torch.Tensor,
+    deltas: torch.Tensor,
+    model_sd: StateDict,
+) -> dict[str, torch.Tensor]:
+    """Cast the dequantized FP8 weight + BF16 deltas back to ``weight.dtype``
+    (FP8) via the fused-add-round kernel on CUDA.
+    Only a subset of linears are FP8-downcast (see ``TRANSFORMER_LINEAR_DOWNCAST_MAP``);
+    LoRAs may also target layers left in BF16 (e.g. audio ``add_q/k/v_proj``, cross-modal
+    projections). For those, fall back to a plain BF16 fuse.
+    """
+    if weight.dtype not in (torch.float8_e4m3fn, torch.float8_e5m2):
+        return bf16_fuse_rule(key, weight, deltas, model_sd)
+    return {key: fuse_cast_fp8_weight(deltas, weight)}
+
+
+fp8_cast_fuse_rule = FuseRule(aggregation_dtype=torch.bfloat16, fuse_fn=_fp8_cast_fuse)
+
+
+# Raw safetensors storage prefix shared by every diffusion-transformer
+# parameter (and every prequant `*_scale` sibling). Verified against
+# ltx-2.3-22b-{dev,distilled}-fp8.safetensors: 2924/2924 and 2992/2992 of
+# the scale keys start with this exact prefix.
+_RAW_DIFFUSION_MODEL_PREFIX = "model.diffusion_model."
+
+
+def _read_scales(checkpoint_path: str | Path) -> dict[str, torch.Tensor]:
+    """Return ``{post_rename_param_key: scale_tensor}`` for every prequant
+    ``*_scale`` sibling in *checkpoint_path*.
+    Keys are returned in the post-rename form the loader will pass to the
+    sd-op (e.g. ``transformer_blocks.0.attn1.to_q.weight``) -- the raw
+    ``model.diffusion_model.`` prefix and the ``_scale`` suffix are both
+    stripped. Catches both ``.weight_scale`` and ``.bias_scale``; the
+    latter is absent in the current LTX-2.3 prequant checkpoints but
+    accepted for forward compatibility.
+    """
+    out: dict[str, torch.Tensor] = {}
+    with safetensors.safe_open(str(checkpoint_path), framework="pt", device="cpu") as h:
+        raw_keys = h.keys()
+        for k in raw_keys:
+            if not k.endswith("_scale"):
+                continue
+            if not k.startswith(_RAW_DIFFUSION_MODEL_PREFIX):
+                raise ValueError(
+                    f"Scale key {k!r} does not start with the expected raw prefix {_RAW_DIFFUSION_MODEL_PREFIX!r}"
+                )
+            param_key = k.removeprefix(_RAW_DIFFUSION_MODEL_PREFIX).removesuffix("_scale")
+            out[param_key] = h.get_tensor(k)
+    return out
+
+
+def _build_prequant_fold_sd_ops(scales: dict[str, torch.Tensor]) -> SDOps:
+    """Build sd-ops that fold prequant ``*_scale`` siblings into their parent
+    tensor at load time.
+    *scales* is keyed by the **post-rename** param key (e.g.
+    ``transformer_blocks.0.attn1.to_q.weight``); see :func:`_read_scales`.
+    Four ``with_kv_operation`` entries (symmetric for ``.weight`` and ``.bias``):
+    * ``.weight`` / ``.bias`` -> if a sibling scale exists in *scales*, fold;
+      then delegate to ``TRANSFORMER_LINEAR_DOWNCAST_MAP`` (downcast covered
+      Linears, pass everything else through). Without a scale, delegate
+      directly.
+    * ``.weight_scale`` / ``.bias_scale`` -> drop (the scale is consumed by
+      the fold). Raises if the scale key doesn't correspond to a known
+      entry in *scales* -- that means the file shipped a scale we didn't
+      pre-register, which would silently desync the fold.
+    """
+
+    def _on_param(param_key: str, value: torch.Tensor) -> list[KeyValueOperationResult]:
+        scale = scales.get(param_key)
+        if scale is None:
+            return TRANSFORMER_LINEAR_DOWNCAST_MAP.apply_to_key_value(param_key, value)
+        scale = scale.to(device=value.device)
+        if scale.ndim != 0:
+            raise ValueError(f"Unsupported scale shape {tuple(scale.shape)} for {param_key}")
+        bf16 = (value.to(torch.float32) * scale).to(torch.bfloat16)
+        # Delegate the final fp8-vs-bf16 decision to the downcast map: Linears
+        # outside the fp8 subset (e.g. to_gate_logits) stay bf16 to match the
+        # plain nn.Linear that the upcast matcher leaves untouched.
+        return TRANSFORMER_LINEAR_DOWNCAST_MAP.apply_to_key_value(param_key, bf16)
+
+    def _drop_scale(scale_key: str, _value: torch.Tensor) -> list[KeyValueOperationResult]:
+        param_key = scale_key.removesuffix("_scale")
+        if param_key not in scales:
+            raise ValueError(
+                f"Scale key {scale_key!r} has no matching entry in the prequant scales dict; "
+                f"_read_scales and the loader's rename map have drifted"
+            )
+        return []
+
+    # Register the drop ops first so the dict-membership sanity check is the
+    # earliest sd-op that can fire on a scale key -- we crash on a stray scale
+    # before any silently mismatched fold has a chance to land in the state
+    # dict. Registration order is irrelevant for correctness (no overlap
+    # between matchers) but communicates intent.
+    return (
+        SDOps("FP8_CAST_PREQUANT_AWARE")
+        .with_kv_operation(key_suffix=".weight_scale", operation=_drop_scale)
+        .with_kv_operation(key_suffix=".bias_scale", operation=_drop_scale)
+        .with_kv_operation(key_suffix=".weight", operation=_on_param)
+        .with_kv_operation(key_suffix=".bias", operation=_on_param)
+    )
+
+
+def build_policy(checkpoint_path: str | Path) -> QuantizationPolicy:
+    """FP8 casting with upcasting during inference.
+    *checkpoint_path* is required (mirroring ``fp8_scaled_mm.build_policy``).
+    For prequantized fp8 checkpoints, sibling ``*_scale`` tensors (weight or
+    bias) are folded into the parent at load time.
+    """
+    scales = _read_scales(checkpoint_path)
+    return QuantizationPolicy(
+        sd_ops=_build_prequant_fold_sd_ops(scales),
+        module_ops=(UPCAST_DURING_INFERENCE,),
+        fuse_rule=fp8_cast_fuse_rule,
+    )

packages/ltx-core/src/ltx_core/quantization/fp8_scaled_mm.py

@@ -1,11 +1,24 @@
+import json
+import struct
 from typing import Callable
 
 import torch
 from torch import nn
 
+from ltx_core.loader.fuse_loras import FuseRule, bf16_fuse_rule
 from ltx_core.loader.module_ops import ModuleOps
-from ltx_core.loader.sd_ops import KeyValueOperationResult, SDOps
+from ltx_core.loader.primitives import StateDict
 from ltx_core.model.transformer import LTXModel
+from ltx_core.quantization.policy import QuantizationPolicy
+from ltx_core.quantization.trtllm_scaled_usable import trtllm_scaled_mm_usable
+
+
+def _read_safetensors_dtypes(path: str) -> dict[str, str]:
+    """Return ``{tensor_name: dtype_string}`` from the safetensors header."""
+    with open(path, "rb") as f:
+        header_size = struct.unpack("<Q", f.read(8))[0]
+        header = json.loads(f.read(header_size).decode("utf-8"))
+    return {k: v["dtype"] for k, v in header.items() if k != "__metadata__"}
 
 
 class FP8Linear(nn.Module):
@@ -25,11 +38,8 @@ class FP8Linear(nn.Module):
         self.in_features = in_features
         self.out_features = out_features
 
-        fp8_shape = (in_features, out_features)
-        self.weight = nn.Parameter(torch.empty(fp8_shape, dtype=torch.float8_e4m3fn, device=device))
-        # Weight scale for FP8 dequantization (shape matches checkpoint format)
+        self.weight = nn.Parameter(torch.empty((out_features, in_features), dtype=torch.float8_e4m3fn, device=device))
         self.weight_scale = nn.Parameter(torch.empty((), dtype=torch.float32, device=device))
-        # Input scale for static quantization (pre-quantized checkpoints)
         self.input_scale = nn.Parameter(torch.empty((), dtype=torch.float32, device=device))
 
         if bias:
@@ -40,31 +50,38 @@ class FP8Linear(nn.Module):
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         origin_shape = x.shape
 
-        # Static quantization: use pre-computed scale
-        qinput, cur_input_scale = torch.ops.tensorrt_llm.static_quantize_e4m3_per_tensor(x, self.input_scale)
+        if trtllm_scaled_mm_usable():
+            qinput, cur_input_scale = torch.ops.tensorrt_llm.static_quantize_e4m3_per_tensor(x, self.input_scale)
+            if qinput.dim() == 3:
+                qinput = qinput.reshape(-1, qinput.shape[-1])
+            output = torch.ops.trtllm.cublas_scaled_mm(
+                qinput,
+                self.weight.t(),
+                scale_a=cur_input_scale,
+                scale_b=self.weight_scale,
+                bias=None,
+                out_dtype=x.dtype,
+            )
+        else:
+            # Clamp before cast: out-of-range values cast to NaN/saturated FP8, which
+            # produces black-screen output on some checkpoints (e.g. ltx-2-19b-dev-fp8).
+            fp8_min = torch.finfo(torch.float8_e4m3fn).min
+            fp8_max = torch.finfo(torch.float8_e4m3fn).max
+            qinput = torch.clamp(x * self.input_scale.reciprocal(), fp8_min, fp8_max).to(torch.float8_e4m3fn)
+            if qinput.dim() == 3:
+                qinput = qinput.reshape(-1, qinput.shape[-1])
+            output = torch._scaled_mm(
+                qinput,
+                self.weight.t(),
+                scale_a=self.input_scale,
+                scale_b=self.weight_scale,
+                out_dtype=x.dtype,
+                use_fast_accum=True,
+            )
 
-        # Flatten to 2D for matmul
-        if qinput.dim() == 3:
-            qinput = qinput.reshape(-1, qinput.shape[-1])
-
-        # FP8 scaled matmul
-        output = torch.ops.trtllm.cublas_scaled_mm(
-            qinput,
-            self.weight,
-            scale_a=cur_input_scale,
-            scale_b=self.weight_scale,
-            bias=None,
-            out_dtype=x.dtype,
-        )
-
-        # Add bias
         if self.bias is not None:
-            bias = self.bias
-            if bias.dtype != output.dtype:
-                bias = bias.to(output.dtype)
-            output = output + bias
+            output = output + self.bias.to(output.dtype)
 
-        # Restore original shape
         if output.dim() != len(origin_shape):
             output_shape = list(origin_shape)
             output_shape[-1] = output.shape[-1]
@@ -74,15 +91,7 @@ class FP8Linear(nn.Module):
 
 
 def quantize_weight_to_fp8_per_tensor(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-    """
-    Quantize a weight tensor to FP8 (float8_e4m3fn) using per-tensor scaling.
-    Args:
-        weight: The weight tensor to quantize (any dtype, will be cast to float32)
-    Returns:
-        Tuple of (quantized_weight, weight_scale):
-        - quantized_weight: FP8 tensor, transposed for cublas_scaled_mm
-        - weight_scale: Per-tensor scale factor (reciprocal of quantization scale)
-    """
+    """Quantize a weight tensor to ``float8_e4m3fn`` with a per-tensor scale."""
     weight_fp32 = weight.to(torch.float32)
 
     fp8_min = torch.finfo(torch.float8_e4m3fn).min
@@ -96,7 +105,6 @@ def quantize_weight_to_fp8_per_tensor(weight: torch.Tensor) -> tuple[torch.Tenso
         weight_fp32: torch.Tensor, scale: torch.Tensor, fp8_min: torch.Tensor, fp8_max: torch.Tensor
     ) -> tuple[torch.Tensor, torch.Tensor]:
         quantized_weight = torch.clamp(weight_fp32 * scale, min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
-        quantized_weight = quantized_weight.t()
         weight_scale = scale.reciprocal()
         return quantized_weight, weight_scale
 
@@ -104,36 +112,8 @@ def quantize_weight_to_fp8_per_tensor(weight: torch.Tensor) -> tuple[torch.Tenso
     return quantized_weight, weight_scale
 
 
-def _should_skip_layer(layer_name: str, excluded_layer_substrings: tuple[str, ...]) -> bool:
-    return any(substring in layer_name for substring in excluded_layer_substrings)
-
-
-EXCLUDED_LAYER_SUBSTRINGS = (
-    "patchify_proj",
-    "adaln_single",
-    "av_ca_video_scale_shift_adaln_single",
-    "av_ca_a2v_gate_adaln_single",
-    "caption_projection",
-    "proj_out",
-    "audio_patchify_proj",
-    "audio_adaln_single",
-    "av_ca_audio_scale_shift_adaln_single",
-    "av_ca_v2a_gate_adaln_single",
-    "audio_caption_projection",
-    "audio_proj_out",
-    "transformer_blocks.0.",
-    *[f"transformer_blocks.{i}." for i in range(43, 48)],
-)
-
-
 def _linear_to_fp8linear(layer: nn.Linear) -> FP8Linear:
-    """
-    Create an FP8Linear layer from an nn.Linear layer.
-    Args:
-        layer: The nn.Linear layer to convert (typically on meta device)
-    Returns:
-        A new FP8Linear with the same configuration
-    """
+    """Create an ``FP8Linear`` matching the shape/bias of *layer*."""
     return FP8Linear(
         in_features=layer.in_features,
         out_features=layer.out_features,
@@ -142,15 +122,14 @@ def _linear_to_fp8linear(layer: nn.Linear) -> FP8Linear:
     )
 
 
-def _apply_fp8_prepare_to_model(model: nn.Module, excluded_layer_substrings: tuple[str, ...]) -> nn.Module:
-    """Replace nn.Linear layers with FP8Linear in the module tree."""
+def _swap_linears_to_fp8(model: nn.Module, should_swap: Callable[[str], bool]) -> nn.Module:
+    """Replace nn.Linear layers with FP8Linear where ``should_swap(name)`` returns True."""
     replacements: list[tuple[nn.Module, str, nn.Linear]] = []
 
     for name, module in model.named_modules():
         if not isinstance(module, nn.Linear) or isinstance(module, FP8Linear):
             continue
-
-        if _should_skip_layer(name, excluded_layer_substrings):
+        if not should_swap(name):
             continue
 
         if "." in name:
@@ -168,40 +147,71 @@ def _apply_fp8_prepare_to_model(model: nn.Module, excluded_layer_substrings: tup
     return model
 
 
-def _create_transpose_kv_operation(
-    excluded_layer_substrings: tuple[str, ...],
-) -> Callable[[str, torch.Tensor], list[KeyValueOperationResult]]:
-    def transpose_if_matches(key: str, value: torch.Tensor) -> list[KeyValueOperationResult]:
-        # Only process .weight keys
-        if not key.endswith(".weight"):
-            return [KeyValueOperationResult(key, value)]
+def get_fp8_swap_module_ops(checkpoint_path: str) -> tuple[ModuleOps, ...]:
+    """Return the FP8 swap ``ModuleOps`` for layers whose ``.weight`` is ``F8_E4M3``
+    and which have a sibling ``.weight_scale`` tensor in the checkpoint.
+    Raises ``ValueError`` if no such layers are found — that combination is ambiguous
+    (a BF16 checkpoint with this policy would load as a no-op).
+    """
+    dtypes = _read_safetensors_dtypes(checkpoint_path)
+    fp8_scale_paths = frozenset(
+        key.removesuffix(".weight_scale")
+        for key in dtypes
+        if key.endswith(".weight_scale") and dtypes.get(key.removesuffix(".weight_scale") + ".weight") == "F8_E4M3"
+    )
+    if not fp8_scale_paths:
+        raise ValueError(
+            f"fp8_scaled_mm requires a pre-quantized checkpoint with F8_E4M3 .weight + .weight_scale "
+            f"tensors, but {checkpoint_path!r} has none. Use QuantizationPolicy.fp8_cast() for BF16 checkpoints."
+        )
 
-        # Only transpose 2D FP8 tensors (Linear weights)
-        if value.dim() != 2 or value.dtype != torch.float8_e4m3fn:
-            return [KeyValueOperationResult(key, value)]
+    def _should_swap(name: str) -> bool:
+        suffix = "." + name
+        return any(p == name or p.endswith(suffix) for p in fp8_scale_paths)
 
-        # Check if the layer is excluded
-        layer_name = key.rsplit(".weight", 1)[0]
-        if _should_skip_layer(layer_name, excluded_layer_substrings):
-            return [KeyValueOperationResult(key, value)]
-
-        # Transpose to cuBLAS layout (in, out)
-        transposed_weight = value.t()
-
-        return [KeyValueOperationResult(key, transposed_weight)]
-
-    return transpose_if_matches
+    return (
+        ModuleOps(
+            name="fp8_swap_linears",
+            matcher=lambda model: isinstance(model, LTXModel),
+            mutator=lambda model: _swap_linears_to_fp8(model, _should_swap),
+        ),
+    )
 
 
-FP8_TRANSPOSE_SD_OPS = SDOps("fp8_transpose_weights").with_kv_operation(
-    _create_transpose_kv_operation(EXCLUDED_LAYER_SUBSTRINGS),
-    key_prefix="transformer_blocks.",
-    key_suffix=".weight",
-)
+def _fp8_scaled_mm_fuse(
+    key: str,
+    weight: torch.Tensor,
+    deltas: torch.Tensor,
+    model_sd: StateDict,
+) -> dict[str, torch.Tensor]:
+    """Dequantize via ``weight.float() * weight_scale``, add the BF16 delta,
+    and re-quantize to FP8 with a fresh per-tensor scale.
+    Layers that were not swapped to scaled FP8 (e.g. small embedder linears
+    excluded from the auto-discovered swap set) stay BF16 and have no
+    ``.weight_scale`` companion -- for those, fall back to a plain bf16 fuse.
+    """
+    scale_key = key.replace(".weight", ".weight_scale")
+    if scale_key not in model_sd.sd:
+        return bf16_fuse_rule(key, weight, deltas, model_sd)
+    weight_scale = model_sd.sd[scale_key]
+    original_weight = weight.to(torch.float32) * weight_scale
+    new_weight = original_weight + deltas.to(torch.float32)
+    new_fp8_weight, new_weight_scale = quantize_weight_to_fp8_per_tensor(new_weight)
+    return {key: new_fp8_weight, scale_key: new_weight_scale}
 
 
-FP8_PREPARE_MODULE_OPS = ModuleOps(
-    name="fp8_prepare_for_loading",
-    matcher=lambda model: isinstance(model, LTXModel),
-    mutator=lambda model: _apply_fp8_prepare_to_model(model, EXCLUDED_LAYER_SUBSTRINGS),
-)
+fp8_scaled_mm_fuse_rule = FuseRule(aggregation_dtype=torch.bfloat16, fuse_fn=_fp8_scaled_mm_fuse)
+
+
+def build_policy(checkpoint_path: str) -> QuantizationPolicy:
+    """FP8 scaled matmul for checkpoints pre-quantized with per-tensor scales.
+    The set of layers to swap to ``FP8Linear`` is discovered from the
+    checkpoint's ``.weight_scale`` tensors via suffix-matching against the
+    model's named modules. Requires a pre-quantized checkpoint; for BF16
+    checkpoints, use :func:`ltx_core.quantization.fp8_cast.build_policy`.
+    """
+    return QuantizationPolicy(
+        sd_ops=None,
+        module_ops=get_fp8_swap_module_ops(checkpoint_path),
+        fuse_rule=fp8_scaled_mm_fuse_rule,
+    )

packages/ltx-core/src/ltx_core/quantization/policy.py

@@ -1,39 +1,24 @@
 from dataclasses import dataclass
 
+from ltx_core.loader.fuse_loras import FuseRule, bf16_fuse_rule
 from ltx_core.loader.module_ops import ModuleOps
 from ltx_core.loader.sd_ops import SDOps
-from ltx_core.quantization.fp8_cast import TRANSFORMER_LINEAR_DOWNCAST_MAP, UPCAST_DURING_INFERENCE
-from ltx_core.quantization.fp8_scaled_mm import FP8_PREPARE_MODULE_OPS, FP8_TRANSPOSE_SD_OPS
+from ltx_core.model.model_protocol import ModelConfigurator
+from ltx_core.model.transformer.model import LTXModel
 
 
 @dataclass(frozen=True)
 class QuantizationPolicy:
     """Configuration for model quantization during loading.
     Attributes:
-        sd_ops: State dict operations for weight transformation.
-        module_ops: Post-load module transformations.
+        sd_ops: State-dict operations applied to each tensor during load.
+        module_ops: Post-load module transformations applied to the meta model.
+        model_configurator: Configurator class to use when constructing the transformer.
+        fuse_rule: How LoRA deltas merge into this policy's weight layout.
+            Default ``bf16_fuse_rule`` is used when no policy is configured.
     """
 
     sd_ops: SDOps | None = None
     module_ops: tuple[ModuleOps, ...] = ()
-
-    @classmethod
-    def fp8_cast(cls) -> "QuantizationPolicy":
-        """Create policy using FP8 casting with upcasting during inference."""
-        return cls(
-            sd_ops=TRANSFORMER_LINEAR_DOWNCAST_MAP,
-            module_ops=(UPCAST_DURING_INFERENCE,),
-        )
-
-    @classmethod
-    def fp8_scaled_mm(cls) -> "QuantizationPolicy":
-        """Create policy using FP8 scaled matrix multiplication."""
-        try:
-            import tensorrt_llm  # noqa: F401, PLC0415
-        except ImportError as e:
-            raise ImportError("tensorrt_llm is not installed, skipping FP8 scaled MM quantization") from e
-
-        return cls(
-            sd_ops=FP8_TRANSPOSE_SD_OPS,
-            module_ops=(FP8_PREPARE_MODULE_OPS,),
-        )
+    model_configurator: type[ModelConfigurator[LTXModel]] | None = None
+    fuse_rule: FuseRule = bf16_fuse_rule

packages/ltx-core/src/ltx_core/quantization/trtllm_scaled_usable.py

@@ -0,0 +1,37 @@
+"""Runtime detection of TensorRT-LLM FP8 scaled-matmul availability.
+When the TRT-LLM ops are usable on the current host (Linux + Hopper-class CUDA
++ tensorrt_llm wheel installed) we use them since they outperform the PyTorch-native
+``torch._scaled_mm`` path. Otherwise we fall back to the native implementation,
+which is portable across platforms (Windows, macOS, AMD GPUs).
+The check runs once and is cached.
+"""
+
+from __future__ import annotations
+
+import platform
+from functools import cache
+
+import torch
+
+
+@cache
+def trtllm_scaled_mm_usable() -> bool:
+    if platform.system() != "Linux":
+        return False
+
+    if not torch.cuda.is_available():
+        return False
+
+    major, minor = torch.cuda.get_device_capability()
+    sm = major * 10 + minor
+
+    if sm < 90 or sm >= 120:
+        return False
+
+    # The import is load-bearing — registers the trtllm torch ops as a side effect.
+    try:
+        import tensorrt_llm  # noqa: F401, PLC0415
+    except Exception:
+        return False
+
+    return True

packages/ltx-core/src/ltx_core/text_encoders/gemma/embeddings_connector.py

@@ -18,7 +18,7 @@ class _BasicTransformerBlock1D(torch.nn.Module):
         dim: int,
         heads: int,
         dim_head: int,
-        rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
+        rope_type: LTXRopeType = LTXRopeType.SPLIT,
         apply_gated_attention: bool = False,
     ):
         super().__init__()
@@ -39,7 +39,7 @@ class _BasicTransformerBlock1D(torch.nn.Module):
     def forward(
         self,
         hidden_states: torch.Tensor,
-        attention_mask: torch.Tensor | None = None,
+        additive_attention_mask: torch.Tensor | None = None,
         pe: torch.Tensor | None = None,
     ) -> torch.Tensor:
         # Notice that normalization is always applied before the real computation in the following blocks.
@@ -49,8 +49,8 @@ class _BasicTransformerBlock1D(torch.nn.Module):
 
         norm_hidden_states = norm_hidden_states.squeeze(1)
 
-        # 2. Self-Attention
-        attn_output = self.attn1(norm_hidden_states, mask=attention_mask, pe=pe)
+        # 2. Self-Attention — `mask` is the kernel-boundary name for the additive mask.
+        attn_output = self.attn1(norm_hidden_states, mask=additive_attention_mask, pe=pe)
 
         hidden_states = attn_output + hidden_states
         if hidden_states.ndim == 4:
@@ -84,7 +84,7 @@ class Embeddings1DConnector(torch.nn.Module):
         causal_temporal_positioning (bool): If True, uses causal attention (default=False).
         num_learnable_registers (int | None): Number of learnable registers to replace padded tokens. If None, disables
             register replacement. (default=128)
-        rope_type (LTXRopeType): The RoPE variant to use (default=DEFAULT_ROPE_TYPE).
+        rope_type (LTXRopeType): The RoPE variant to use.
         double_precision_rope (bool): Use double precision rope calculation (default=False).
     """
 
@@ -99,7 +99,7 @@ class Embeddings1DConnector(torch.nn.Module):
         positional_embedding_max_pos: list[int] | None = None,
         causal_temporal_positioning: bool = False,
         num_learnable_registers: int | None = 128,
-        rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
+        rope_type: LTXRopeType = LTXRopeType.SPLIT,
         double_precision_rope: bool = False,
         apply_gated_attention: bool = False,
     ):
@@ -133,51 +133,40 @@ class Embeddings1DConnector(torch.nn.Module):
             )
 
     def _replace_padded_with_learnable_registers(
-        self, hidden_states: torch.Tensor, attention_mask: torch.Tensor
+        self, hidden_states: torch.Tensor, additive_attention_mask: torch.Tensor
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        assert hidden_states.shape[1] % self.num_learnable_registers == 0, (
-            f"Hidden states sequence length {hidden_states.shape[1]} must be divisible by num_learnable_registers "
-            f"{self.num_learnable_registers}."
-        )
+        batch_size, seq_len, _ = hidden_states.shape
 
-        num_registers_duplications = hidden_states.shape[1] // self.num_learnable_registers
-        learnable_registers = torch.tile(self.learnable_registers, (num_registers_duplications, 1))
-        attention_mask_binary = (attention_mask.squeeze(1).squeeze(1).unsqueeze(-1) >= -9000.0).int()
+        assert seq_len % self.num_learnable_registers == 0
 
-        non_zero_hidden_states = hidden_states[:, attention_mask_binary.squeeze().bool(), :]
-        non_zero_nums = non_zero_hidden_states.shape[1]
-        pad_length = hidden_states.shape[1] - non_zero_nums
-        adjusted_hidden_states = torch.nn.functional.pad(non_zero_hidden_states, pad=(0, 0, 0, pad_length), value=0)
-        flipped_mask = torch.flip(attention_mask_binary, dims=[1])
-        hidden_states = flipped_mask * adjusted_hidden_states + (1 - flipped_mask) * learnable_registers
+        registers = self.learnable_registers.repeat(seq_len // self.num_learnable_registers, 1).to(hidden_states.dtype)
+        registers = registers.unsqueeze(0).expand(batch_size, -1, -1)  # (B, seq_len, hidden_dim)
+        binary_mask = additive_attention_mask[:, 0, 0, :].unsqueeze(-1) >= 0
+        binary_mask = binary_mask.to(hidden_states.dtype)
+        hidden_states = binary_mask * hidden_states + (1 - binary_mask) * registers
 
-        attention_mask = torch.full_like(
-            attention_mask,
-            0.0,
-            dtype=attention_mask.dtype,
-            device=attention_mask.device,
-        )
-
-        return hidden_states, attention_mask
+        return hidden_states, torch.zeros_like(additive_attention_mask)
 
     def forward(
         self,
         hidden_states: torch.Tensor,
-        attention_mask: torch.Tensor | None = None,
+        additive_attention_mask: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        """
-        Forward pass of Embeddings1DConnector.
+        """Forward pass of Embeddings1DConnector.
         Args:
-            hidden_states (torch.Tensor): Input tensor of embeddings (shape [batch, seq_len, feature_dim]).
-            attention_mask (torch.Tensor|None): Optional mask for valid tokens (shape compatible with hidden_states).
+            hidden_states: (B, S, D) input embeddings.
+            additive_attention_mask: optional additive mask of shape (B, 1, 1, S), where
+                valid = 0.0 and padding = -torch.finfo(dtype).max.
         Returns:
-            tuple[torch.Tensor, torch.Tensor]: Processed features and the corresponding (possibly modified) mask.
+            (hidden_states, additive_attention_mask)
         """
         if self.num_learnable_registers:
-            hidden_states, attention_mask = self._replace_padded_with_learnable_registers(hidden_states, attention_mask)
+            hidden_states, additive_attention_mask = self._replace_padded_with_learnable_registers(
+                hidden_states, additive_attention_mask
+            )
 
         indices_grid = torch.arange(hidden_states.shape[1], dtype=torch.float32, device=hidden_states.device)
-        indices_grid = indices_grid[None, None, :]
+        indices_grid = indices_grid[None, None, :].expand(hidden_states.shape[0], -1, -1)
         freq_grid_generator = generate_freq_grid_np if self.double_precision_rope else generate_freq_grid_pytorch
         freqs_cis = precompute_freqs_cis(
             indices_grid=indices_grid,
@@ -191,11 +180,11 @@ class Embeddings1DConnector(torch.nn.Module):
         )
 
         for block in self.transformer_1d_blocks:
-            hidden_states = block(hidden_states, attention_mask=attention_mask, pe=freqs_cis)
+            hidden_states = block(hidden_states, additive_attention_mask=additive_attention_mask, pe=freqs_cis)
 
         hidden_states = rms_norm(hidden_states)
 
-        return hidden_states, attention_mask
+        return hidden_states, additive_attention_mask
 
 
 class Embeddings1DConnectorConfigurator(ModelConfigurator[Embeddings1DConnector]):
@@ -204,7 +193,7 @@ class Embeddings1DConnectorConfigurator(ModelConfigurator[Embeddings1DConnector]
     @classmethod
     def from_config(cls: type[Embeddings1DConnector], config: dict) -> Embeddings1DConnector:
         transformer_config = config.get("transformer", {})
-        rope_type = LTXRopeType(transformer_config.get("rope_type", "interleaved"))
+        rope_type = LTXRopeType(transformer_config.get("rope_type", "split"))
         double_precision_rope = transformer_config.get("frequencies_precision", False) == "float64"
         pe_max_pos = transformer_config.get("connector_positional_embedding_max_pos", [1])
 
@@ -231,7 +220,7 @@ class AudioEmbeddings1DConnectorConfigurator(ModelConfigurator[Embeddings1DConne
     @classmethod
     def from_config(cls: type[Embeddings1DConnector], config: dict) -> Embeddings1DConnector:
         transformer_config = config.get("transformer", {})
-        rope_type = LTXRopeType(transformer_config.get("rope_type", "interleaved"))
+        rope_type = LTXRopeType(transformer_config.get("rope_type", "split"))
         double_precision_rope = transformer_config.get("frequencies_precision", False) == "float64"
         pe_max_pos = transformer_config.get("connector_positional_embedding_max_pos", [1])
 

packages/ltx-core/src/ltx_core/text_encoders/gemma/embeddings_processor.py

@@ -19,12 +19,32 @@ def convert_to_additive_mask(attention_mask: torch.Tensor, dtype: torch.dtype) -
     ) * torch.finfo(dtype).max
 
 
-def _to_binary_mask(encoded: torch.Tensor, encoded_mask: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-    """Convert connector output mask to binary mask and apply to encoded tensor."""
-    binary_mask = (encoded_mask < 0.000001).to(torch.int64)
-    binary_mask = binary_mask.reshape([encoded.shape[0], encoded.shape[1], 1])
-    encoded = encoded * binary_mask
-    return encoded, binary_mask
+def _compute_right_pad_order(additive_mask: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """Compute the index permutation that places valid tokens before pads in each row.
+    Stable sort: valid tokens keep their relative order. Idempotent for inputs already
+    right-padded. The sort and reordered mask depend only on the mask, so they can be
+    computed once and reused across multiple feature tensors that share the mask.
+    Args:
+        additive_mask: (B, 1, 1, S) additive mask, ``0.0`` for valid, ``-finfo.max`` for pad.
+    Returns:
+        ``(sort_idx, reordered_additive_mask)``: ``sort_idx`` is (B, S); the reordered mask
+        has the same shape as the input.
+    """
+    binary = (additive_mask[:, 0, 0, :] >= 0).to(torch.int32)  # (B, S)
+    sort_idx = torch.argsort(binary, dim=-1, descending=True, stable=True)  # (B, S)
+    new_binary = torch.gather(binary, 1, sort_idx)
+    new_additive = (new_binary.to(additive_mask.dtype) - 1) * torch.finfo(additive_mask.dtype).max
+    return sort_idx, new_additive[:, None, None, :]
+
+
+def _apply_right_pad_order(features: torch.Tensor, sort_idx: torch.Tensor) -> torch.Tensor:
+    """Apply a precomputed right-pad permutation (from ``_compute_right_pad_order``) to features."""
+    return torch.gather(features, 1, sort_idx.unsqueeze(-1).expand_as(features))
+
+
+def _to_binary_mask(encoded_mask: torch.Tensor, lead_shape: tuple[int, int]) -> torch.Tensor:
+    """Convert connector output mask to a binary (0/1) mask shaped ``(B, S, 1)`` for broadcasting."""
+    return (encoded_mask < 0.000001).to(torch.int64).reshape([lead_shape[0], lead_shape[1], 1])
 
 
 class EmbeddingsProcessor(nn.Module):
@@ -57,12 +77,19 @@ class EmbeddingsProcessor(nn.Module):
         if self.audio_connector is None and audio_features is not None:
             raise ValueError("Audio features were provided but no audio connector is configured.")
 
-        video_encoded, video_mask = self.video_connector(video_features, additive_attention_mask)
-        video_encoded, binary_mask = _to_binary_mask(video_encoded, video_mask)
+        # Connectors expect right-padded input ([valid, pad]). Normalize layout here so the
+        # upstream tokenizer can keep using either side without coupling to the connector.
+        # The sort index depends only on the mask, so compute it once and reuse for audio.
+        sort_idx, mask_for_connector = _compute_right_pad_order(additive_attention_mask)
+        video_features = _apply_right_pad_order(video_features, sort_idx)
+        video_encoded, video_mask = self.video_connector(video_features, mask_for_connector)
+        binary_mask = _to_binary_mask(video_mask, video_encoded.shape[:2])
+        video_encoded = video_encoded * binary_mask
 
         audio_encoded = None
         if self.audio_connector is not None:
-            audio_encoded, _ = self.audio_connector(audio_features, additive_attention_mask)
+            audio_features = _apply_right_pad_order(audio_features, sort_idx)
+            audio_encoded, _ = self.audio_connector(audio_features, mask_for_connector)
 
         return video_encoded, audio_encoded, binary_mask.squeeze(-1)
 

packages/ltx-core/src/ltx_core/text_encoders/gemma/encoders/base_encoder.py

@@ -30,21 +30,31 @@ class GemmaTextEncoder(torch.nn.Module):
 
     def encode(
         self,
-        text: str,
+        prompts: list[str],
         padding_side: str = "left",  # noqa: ARG002
-    ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]:
-        """Run Gemma LLM and return raw hidden states + attention mask.
-        Calls the inner model (self.model.model) to skip lm_head logits computation (~500 MiB saving).
-        Returns:
-            (hidden_states, attention_mask) where hidden_states is a tuple of per-layer tensors.
+    ) -> list[tuple[tuple[torch.Tensor, ...], torch.Tensor]]:
+        """Run a single fused Gemma forward over a batch of prompts.
+        Calls the inner model (self.model.model) to skip lm_head logits computation
+        (~500 MiB saving). The tokenizer pads every prompt to ``max_length`` (1024),
+        so the inputs stack into a single ``[N, 1024]`` batch with no further padding
+        logic; per-prompt outputs are sliced back to ``[1, 1024, D]`` / ``[1, 1024]``
+        in the original order.
         """
-        token_pairs = self.tokenizer.tokenize_with_weights(text)["gemma"]
-        input_ids = torch.tensor([[t[0] for t in token_pairs]], device=self.model.device)
-        attention_mask = torch.tensor([[w[1] for w in token_pairs]], device=self.model.device)
+        if not prompts:
+            return []
+        tokenized = [self.tokenizer.tokenize_with_weights(t)["gemma"] for t in prompts]
+        input_ids = torch.tensor(
+            [[tok for tok, _ in pairs] for pairs in tokenized],
+            device=self.model.device,
+        )
+        attention_mask = torch.tensor(
+            [[w for _, w in pairs] for pairs in tokenized],
+            device=self.model.device,
+        )
         outputs = self.model.model(input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True)
         hidden_states = outputs.hidden_states
         del outputs
-        return hidden_states, attention_mask
+        return [(tuple(h[i : i + 1] for h in hidden_states), attention_mask[i : i + 1]) for i in range(len(prompts))]
 
     # --- Prompt enhancement methods ---
 
@@ -183,7 +193,7 @@ def module_ops_from_gemma_root(gemma_root: str) -> tuple[ModuleOps, ...]:
         return module
 
     def load_processor(module: GemmaTextEncoder) -> GemmaTextEncoder:
-        image_processor = AutoImageProcessor.from_pretrained(processor_root, local_files_only=True)
+        image_processor = AutoImageProcessor.from_pretrained(processor_root, local_files_only=True, use_fast=False)
         if not module.tokenizer:
             raise ValueError("Tokenizer model operation must be performed before processor model operation")
         module.processor = Gemma3Processor(image_processor=image_processor, tokenizer=module.tokenizer.tokenizer)

packages/ltx-core/src/ltx_core/text_encoders/gemma/encoders/encoder_configurator.py

@@ -1,4 +1,5 @@
 import torch
+import transformers
 from transformers import Gemma3Config
 from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
 from transformers.models.gemma3 import Gemma3ForConditionalGeneration
@@ -19,6 +20,8 @@ from ltx_core.text_encoders.gemma.feature_extractor import (
     FeatureExtractorV2,
 )
 
+_TRANSFORMERS_V5: bool = int(transformers.__version__.split(".", 1)[0]) >= 5
+
 
 class GemmaTextEncoderConfigurator(ModelConfigurator[GemmaTextEncoder]):
     @classmethod
@@ -100,26 +103,45 @@ def _create_feature_extractor(transformer_config: dict) -> torch.nn.Module:
 
 # --- Split SDOps: Gemma LLM keys vs Embeddings Processor keys ---
 
-GEMMA_LLM_KEY_OPS = (
-    SDOps("GEMMA_LLM_KEY_OPS")
-    # 1. Map language model layers (note the double .model prefix)
-    .with_matching(prefix="language_model.model.")
-    .with_replacement("language_model.model.", "model.model.language_model.")
-    # 2. Map the Vision Tower
-    .with_matching(prefix="vision_tower.")
-    .with_replacement("vision_tower.", "model.model.vision_tower.")
-    # 3. Map the Multi-Modal Projector
-    .with_matching(prefix="multi_modal_projector.")
-    .with_replacement("multi_modal_projector.", "model.model.multi_modal_projector.")
-    # 4. Duplicate embed_tokens to lm_head (needed for prompt enhancement via generate())
-    .with_kv_operation(
-        operation=lambda key, value: [
-            KeyValueOperationResult(key, value),
-            KeyValueOperationResult("model.lm_head.weight", value),
-        ],
-        key_prefix="model.model.language_model.embed_tokens.weight",
+
+def _build_gemma_llm_key_ops(*, transformers_v5: bool) -> SDOps:
+    """Build the checkpoint-key remapping for the Gemma multimodal encoder.
+    The vision-tower mapping differs between transformers <5 and >=5 because
+    upstream PR https://github.com/huggingface/transformers/pull/39847 flattened
+    ``Gemma3ForConditionalGeneration.model.vision_tower.vision_model`` into
+    ``model.vision_tower``. Checkpoints continue to ship the legacy
+    ``vision_tower.vision_model.*`` prefix, so we strip the inner ``vision_model.``
+    when targeting v5 and pass it through unchanged for v4.
+    """
+    base = (
+        SDOps("GEMMA_LLM_KEY_OPS")
+        # 1. Map language model layers (note the double .model prefix)
+        .with_matching(prefix="language_model.model.")
+        .with_replacement("language_model.model.", "model.model.language_model.")
+        # 2. Map the Vision Tower (version-dependent — see docstring)
+        .with_matching(prefix="vision_tower.")
     )
-)
+    if transformers_v5:
+        base = base.with_replacement("vision_tower.vision_model.", "model.model.vision_tower.")
+    else:
+        base = base.with_replacement("vision_tower.", "model.model.vision_tower.")
+    return (
+        base
+        # 3. Map the Multi-Modal Projector
+        .with_matching(prefix="multi_modal_projector.")
+        .with_replacement("multi_modal_projector.", "model.model.multi_modal_projector.")
+        # 4. Duplicate embed_tokens to lm_head (needed for prompt enhancement via generate())
+        .with_kv_operation(
+            operation=lambda key, value: [
+                KeyValueOperationResult(key, value),
+                KeyValueOperationResult("model.lm_head.weight", value),
+            ],
+            key_prefix="model.model.language_model.embed_tokens.weight",
+        )
+    )
+
+
+GEMMA_LLM_KEY_OPS = _build_gemma_llm_key_ops(transformers_v5=_TRANSFORMERS_V5)
 
 EMBEDDINGS_PROCESSOR_KEY_OPS = (
     SDOps("EMBEDDINGS_PROCESSOR_KEY_OPS")
@@ -152,24 +174,85 @@ VIDEO_ONLY_EMBEDDINGS_PROCESSOR_KEY_OPS = (
 )
 
 
+def _resolve_local_base_freq(config: object) -> float:
+    rope_parameters = getattr(config, "rope_parameters", None)
+    if isinstance(rope_parameters, dict) and "sliding_attention" in rope_parameters:
+        sliding = rope_parameters["sliding_attention"]
+        if isinstance(sliding, dict) and "rope_theta" in sliding:
+            return float(sliding["rope_theta"])
+    if hasattr(config, "rope_local_base_freq"):
+        return float(config.rope_local_base_freq)
+    raise AttributeError(
+        "Gemma text_config exposes neither rope_local_base_freq nor rope_parameters['sliding_attention']['rope_theta']"
+    )
+
+
+def _resolve_full_rope_type(config: object) -> str:
+    rope_parameters = getattr(config, "rope_parameters", None)
+    if isinstance(rope_parameters, dict) and "full_attention" in rope_parameters:
+        full = rope_parameters["full_attention"]
+        if isinstance(full, dict) and "rope_type" in full:
+            return str(full["rope_type"])
+    rope_scaling = getattr(config, "rope_scaling", None)
+    if rope_scaling is not None:
+        if isinstance(rope_scaling, dict):
+            if "rope_type" in rope_scaling:
+                return str(rope_scaling["rope_type"])
+        elif hasattr(rope_scaling, "rope_type"):
+            return str(rope_scaling.rope_type)
+    raise AttributeError(
+        "Gemma text_config exposes neither rope_scaling.rope_type nor rope_parameters['full_attention']['rope_type']"
+    )
+
+
+def _populate_rotary_v4(l_model: torch.nn.Module, config: object) -> None:
+    """transformers <5 layout: separate ``rotary_emb_local`` + ``rotary_emb``."""
+    dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+    base = _resolve_local_base_freq(config)
+    local_inv = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(dtype=torch.float) / dim))
+    full_inv, _ = ROPE_INIT_FUNCTIONS[_resolve_full_rope_type(config)](config)
+    l_model.rotary_emb_local.register_buffer("inv_freq", local_inv)
+    l_model.rotary_emb.register_buffer("inv_freq", full_inv)
+
+
+def _populate_rotary_v5(l_model: torch.nn.Module, config: object) -> None:
+    """transformers >=5 layout: single ``rotary_emb`` with per-layer-type buffers.
+    Mirrors ``Gemma3PreTrainedModel._init_weights`` for ``Gemma3RotaryEmbedding``
+    so meta-built models reach the same numerical state as a from_pretrained load.
+    """
+    rope_emb = l_model.rotary_emb
+    for layer_type in dict.fromkeys(config.layer_types):
+        rope_params = config.rope_parameters[layer_type]
+        if rope_params is None:
+            continue
+        rope_type = rope_params["rope_type"]
+        if rope_type == "default":
+            inv_freq, attn_scaling = rope_emb.compute_default_rope_parameters(config, layer_type=layer_type)
+        else:
+            inv_freq, attn_scaling = ROPE_INIT_FUNCTIONS[rope_type](config, layer_type=layer_type)
+        rope_emb.register_buffer(f"{layer_type}_inv_freq", inv_freq, persistent=False)
+        rope_emb.register_buffer(f"{layer_type}_original_inv_freq", inv_freq.clone(), persistent=False)
+        setattr(rope_emb, f"{layer_type}_attention_scaling", attn_scaling)
+
+
 def create_and_populate(module: GemmaTextEncoder) -> GemmaTextEncoder:
     model = module.model
-    v_model = model.model.vision_tower.vision_model
+    v_tower = model.model.vision_tower
+    v_model = v_tower.vision_model if hasattr(v_tower, "vision_model") else v_tower
     l_model = model.model.language_model
 
     config = model.config.text_config
-    dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
-    base = config.rope_local_base_freq
-    local_rope_freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(dtype=torch.float) / dim))
-    inv_freqs, _ = ROPE_INIT_FUNCTIONS[config.rope_scaling["rope_type"]](config)
+
+    if hasattr(l_model, "rotary_emb_local"):
+        _populate_rotary_v4(l_model, config)
+    else:
+        _populate_rotary_v5(l_model, config)
 
     positions_length = len(v_model.embeddings.position_ids[0])
     position_ids = torch.arange(positions_length, dtype=torch.long, device="cpu").unsqueeze(0)
     v_model.embeddings.register_buffer("position_ids", position_ids)
-    embed_scale = torch.tensor(model.config.text_config.hidden_size**0.5, device="cpu")
+    embed_scale = torch.tensor(config.hidden_size**0.5, device="cpu")
     l_model.embed_tokens.register_buffer("embed_scale", embed_scale)
-    l_model.rotary_emb_local.register_buffer("inv_freq", local_rope_freqs)
-    l_model.rotary_emb.register_buffer("inv_freq", inv_freqs)
 
     return module
 

packages/ltx-core/src/ltx_core/text_encoders/gemma/feature_extractor.py

@@ -11,37 +11,25 @@ from torch import nn
 
 def _norm_and_concat_padded_batch(
     encoded_text: torch.Tensor,
-    sequence_lengths: torch.Tensor,
-    padding_side: str = "right",
+    attention_mask: torch.Tensor,
 ) -> torch.Tensor:
     """Normalize and flatten multi-layer hidden states, respecting padding.
     Performs per-batch, per-layer normalization using masked mean and range,
-    then concatenates across the layer dimension.
+    then concatenates across the layer dimension. Padding-side agnostic: the
+    binary ``attention_mask`` already encodes which positions are valid.
     Args:
         encoded_text: Hidden states of shape [batch, seq_len, hidden_dim, num_layers].
-        sequence_lengths: Number of valid (non-padded) tokens per batch item.
-        padding_side: Whether padding is on "left" or "right".
+        attention_mask: Binary mask of shape [batch, seq_len], 1 for valid tokens, 0 for padding.
     Returns:
         Normalized tensor of shape [batch, seq_len, hidden_dim * num_layers],
         with padded positions zeroed out.
     """
-    b, t, d, l = encoded_text.shape  # noqa: E741
-    device = encoded_text.device
-
-    token_indices = torch.arange(t, device=device)[None, :]
-
-    if padding_side == "right":
-        mask = token_indices < sequence_lengths[:, None]
-    elif padding_side == "left":
-        start_indices = t - sequence_lengths[:, None]
-        mask = token_indices >= start_indices
-    else:
-        raise ValueError(f"padding_side must be 'left' or 'right', got {padding_side}")
-
-    mask = rearrange(mask, "b t -> b t 1 1")
-
+    b, _, d, l = encoded_text.shape  # noqa: E741
     eps = 1e-6
 
+    sequence_lengths = attention_mask.sum(dim=-1)
+    mask = rearrange(attention_mask.bool(), "b t -> b t 1 1")
+
     masked = encoded_text.masked_fill(~mask, 0.0)
     denom = (sequence_lengths * d).view(b, 1, 1, 1)
     mean = masked.sum(dim=(1, 2), keepdim=True) / (denom + eps)
@@ -51,12 +39,10 @@ def _norm_and_concat_padded_batch(
     range_ = x_max - x_min
 
     normed = 8 * (encoded_text - mean) / (range_ + eps)
-    normed = normed.reshape(b, t, -1)
+    normed = normed.reshape(b, -1, d * l)
 
     mask_flattened = rearrange(mask, "b t 1 1 -> b t 1").expand(-1, -1, d * l)
-    normed = normed.masked_fill(~mask_flattened, 0.0)
-
-    return normed
+    return normed.masked_fill(~mask_flattened, 0.0)
 
 
 def norm_and_concat_per_token_rms(
@@ -97,12 +83,14 @@ class FeatureExtractorV1(nn.Module):
         self.is_av = is_av
 
     def forward(
-        self, hidden_states: torch.Tensor, attention_mask: torch.Tensor, padding_side: str = "left"
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+        padding_side: str = "left",  # noqa: ARG002 — kept for API stability; norm is layout-agnostic
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
         encoded = torch.stack(hidden_states, dim=-1) if isinstance(hidden_states, (list, tuple)) else hidden_states
         dtype = encoded.dtype
-        sequence_lengths = attention_mask.sum(dim=-1)
-        normed = _norm_and_concat_padded_batch(encoded, sequence_lengths, padding_side)
+        normed = _norm_and_concat_padded_batch(encoded, attention_mask)
         features = self.aggregate_embed(normed.to(dtype))
         if self.is_av:
             return features, features

packages/ltx-core/src/ltx_core/text_encoders/gemma/tokenizer.py

@@ -1,6 +1,13 @@
+from enum import Enum
+
 from transformers import AutoTokenizer
 
 
+class PaddingSide(str, Enum):
+    LEFT = "left"
+    RIGHT = "right"
+
+
 class LTXVGemmaTokenizer:
     """
     Tokenizer wrapper for Gemma models compatible with LTXV processes.
@@ -8,18 +15,18 @@ class LTXVGemmaTokenizer:
     ensuring correct settings and output formatting for downstream consumption.
     """
 
-    def __init__(self, tokenizer_path: str, max_length: int = 256):
+    def __init__(self, tokenizer_path: str, max_length: int = 256, padding_side: PaddingSide = PaddingSide.LEFT):
         """
         Initialize the tokenizer.
         Args:
             tokenizer_path (str): Path to the pretrained tokenizer files or model directory.
             max_length (int, optional): Max sequence length for encoding. Defaults to 256.
+            padding_side (PaddingSide, optional): Side to pad on. Defaults to ``PaddingSide.LEFT``.
         """
         self.tokenizer = AutoTokenizer.from_pretrained(
             tokenizer_path, local_files_only=True, model_max_length=max_length
         )
-        # Gemma expects left padding for chat-style prompts; for plain text it doesn't matter much.
-        self.tokenizer.padding_side = "left"
+        self.tokenizer.padding_side = padding_side.value
         if self.tokenizer.pad_token is None:
             self.tokenizer.pad_token = self.tokenizer.eos_token
 

packages/ltx-core/src/ltx_core/tools.py

@@ -138,7 +138,7 @@ class VideoLatentTools(LatentTools):
             LatentState(
                 latent=initial_latent,
                 denoise_mask=denoise_mask,
-                positions=positions.to(dtype),
+                positions=positions,
                 clean_latent=clean_latent,
             )
         )

packages/ltx-core/src/ltx_core/types.py

@@ -20,15 +20,17 @@ class SpatioTemporalScaleFactors(NamedTuple):
     """
     Describes the spatiotemporal downscaling between decoded video space and
     the corresponding VAE latent grid.
+    Field order matches the (frame/time, height, width) axis layout used by
+    latent tensors and meshgrid coordinates elsewhere in the codebase.
     """
 
     time: int
-    width: int
     height: int
+    width: int
 
     @classmethod
     def default(cls) -> "SpatioTemporalScaleFactors":
-        return cls(time=8, width=32, height=32)
+        return cls(time=8, height=32, width=32)
 
 
 VIDEO_SCALE_FACTORS = SpatioTemporalScaleFactors.default()
@@ -74,9 +76,9 @@ class VideoLatentShape(NamedTuple):
         latent_channels: int = 128,
         scale_factors: SpatioTemporalScaleFactors = VIDEO_SCALE_FACTORS,
     ) -> "VideoLatentShape":
-        frames = (shape.frames - 1) // scale_factors[0] + 1
-        height = shape.height // scale_factors[1]
-        width = shape.width // scale_factors[2]
+        frames = (shape.frames - 1) // scale_factors.time + 1
+        height = shape.height // scale_factors.height
+        width = shape.width // scale_factors.width
 
         return VideoLatentShape(
             batch=shape.batch,

packages/ltx-pipelines/CLAUDE.md

@@ -17,12 +17,14 @@ Inference pipelines for LTX-2 audio-video generation. Depends on `ltx-core` for
 | Pipeline | File | Stages | Model | Sampler | Use case |
 |----------|------|--------|-------|---------|----------|
 | `TI2VidOneStagePipeline` | `ti2vid_one_stage.py` | 1 | Full | Euler | Simple text/image-to-video |
+| `T2AOneStagePipeline` | `t2a_one_stage.py` | 1 | Full | Euler | Text-to-audio (audio-only output, no video branch) |
 | `TI2VidTwoStagesPipeline` | `ti2vid_two_stages.py` | 2 | Full + distilled LoRA | Euler | Production quality |
 | `TI2VidTwoStagesHQPipeline` | `ti2vid_two_stages_hq.py` | 2 | Full + distilled LoRA (both stages) | Res2s | Highest quality, fewer steps |
 | `A2VidPipelineTwoStage` | `a2vid_two_stage.py` | 2 | Full + distilled LoRA | Euler | Audio-conditioned video |
 | `KeyframeInterpolationPipeline` | `keyframe_interpolation.py` | 2 | Full + distilled LoRA | Euler | Keyframe interpolation |
 | `DistilledPipeline` | `distilled.py` | 2 | Distilled only | Euler | Fastest inference |
 | `ICLoraPipeline` | `ic_lora.py` | 2 | Distilled only | Euler | Video-to-video with IC-LoRA control |
+| `LipDubPipeline` | `lipdub.py` | 2 | Distilled only | Euler | Lip dubbing with IC-LoRA + audio ref conditioning |
 | `RetakePipeline` | `retake.py` | 1 | Full or distilled | Euler | Video region regeneration |
 
 ## Guidance
@@ -56,7 +58,7 @@ Inference pipelines for LTX-2 audio-video generation. Depends on `ltx-core` for
 ### Memory management
 
 - **Model lifecycle**: All blocks build their model on call and free it on exit. `gpu_model()` moves params to `"meta"` device on exit, immediately releasing storage. No model persists between calls.
-- **Layer streaming**: When `streaming_prefetch_count` is set, `DiffusionStage` wraps the transformer in `LayerStreamingWrapper`. Layers live on pinned CPU memory; only `1 + prefetch_count` layers are on GPU at a time, with async H2D prefetch on a separate CUDA stream.
+- **Block streaming**: When offloading is enabled, `DiffusionStage` wraps the transformer in `BlockStreamingWrapper`. Blocks live on pinned CPU memory; only 2 blocks are buffered on GPU at a time (one for compute, one for async H2D copy on a separate CUDA stream).
 - **Batch splitting**: `BatchSplitAdapter` wraps the transformer and splits inputs exceeding `max_batch_size` into sequential chunks. If guidance needs B=4 but `max_batch_size=1`, it runs 4 sequential B=1 passes. Higher `max_batch_size` reduces layer-streaming PCIe transfers at the cost of peak memory.
 
 ## Denoisers (`utils/denoisers.py`)
@@ -65,6 +67,10 @@ Inference pipelines for LTX-2 audio-video generation. Depends on `ltx-core` for
 - `GuidedDenoiser` -- CFG/STG with static `MultiModalGuider` instances (HQ, A2Vid, Retake non-distilled).
 - `FactoryGuidedDenoiser` -- per-step guider creation via factory (OneStageTI2Vid, TwoStagesTI2Vid, Keyframe).
 
+All denoisers return a `(video_result, audio_result)` tuple of `DenoisedLatentResult` (defined in `utils/types.py`), either element may be `None` for absent modalities. `DenoisedLatentResult.denoised` is the final blended tensor. Guided denoisers additionally populate per-pass fields (`.cond`, `.uncond`, `.ptb`, `.mod`) on each result; `SimpleDenoiser` leaves these `None`.
+
+`GuidedDenoiser` and `FactoryGuidedDenoiser` accept `force_uncond_pass=True` to run the uncond pass even when `cfg_scale=1.0` (required by CFG++ when the guidance scale is 1 but the uncond prediction is still needed for the ODE derivative). Requires `negative_context` to be set on the guider. When enabled, `DenoisedLatentResult.uncond` will be a tensor instead of `None`.
+
 Guided denoisers batch all guidance passes into a **single transformer call**: states are repeated along the batch dimension, contexts concatenated, and a `BatchedPerturbationConfig` controls which attention ops are skipped per sample. Pass count is dynamic: B=2 for CFG-only, up to B=4 with CFG+STG+modality isolation. Results are split back and blended by the guider.
 
 ## Per-pipeline unique features
@@ -72,6 +78,7 @@ Guided denoisers batch all guidance passes into a **single transformer call**: s
 - **HQ**: Res2s second-order sampler for **both** stages, latent-dependent sigma schedule, distilled LoRA on both stages with separate strengths.
 - **A2Vid**: Audio frozen in both stages (`frozen=True, noise_scale=0.0`). Returns original audio (not VAE-decoded); no `AudioDecoder`.
 - **IC-LoRA**: `VideoConditionByReferenceLatent`, `reference_downscale_factor` from LoRA metadata, `skip_stage_2`, attention mask downsampling. Stage 2 is LoRA-free and uses `combined_image_conditionings` (no IC-LoRA conditioning).
+- **LipDub**: Standalone pipeline; IC reference **video** helpers in `iclora_utils.py`, LipDub-only **audio** patchify/negative positions in `lipdub.py`. Appends frozen audio-reference tokens via `AudioConditionByReferenceLatent` (ltx-core), matching video token order (`[target | ref]`) while keeping reference RoPE positions negative (training-compatible). Single IC-LoRA on both stages; full IC-LoRA video conditioning at stage 1 and 2; stage-2 audio is frozen with S1 latent as initial state and uses S1-derived ref. Final audio decoded from stage 1 latent. The LipDub CLI does not expose `--conditioning-attention-mask`; use `ic_lora.py` if you need spatial IC attention masking.
 - **Keyframe**: Uses `image_conditionings_by_adding_guiding_latent` in both stages (all frames as keyframe guidance, no replacement) -- unlike TI2Vid which uses `combined_image_conditionings` (frame_idx=0 replaces, others guide).
 - **Retake**: `TemporalRegionMask` for selective time-window regeneration. `regenerate_video`/`regenerate_audio` flags. Conditional distilled/full behavior.
 - **Distilled**: Single `self.stage` reused for both stages (not `stage_1`/`stage_2`).

packages/ltx-pipelines/README.md

@@ -58,11 +58,14 @@ Available pipeline modules:
 - `ltx_pipelines.ti2vid_two_stages` - Two-stage text/image-to-video (recommended).
 - `ltx_pipelines.ti2vid_two_stages_hq` - Two-stage text/image-to-video (different sampler, better quality).
 - `ltx_pipelines.ti2vid_one_stage` - Single-stage text/image-to-video.
+- `ltx_pipelines.t2a_one_stage` - Single-stage text-to-audio (audio-only output).
 - `ltx_pipelines.distilled` - Fast text/image-to-video pipeline using only the distilled model.
 - `ltx_pipelines.ic_lora` - Video-to-video with IC-LoRA.
 - `ltx_pipelines.keyframe_interpolation` - Keyframe interpolation.
 - `ltx_pipelines.a2vid_two_stage` - Audio-to-video generation conditioned on an input audio.
 - `ltx_pipelines.retake` - Regenerate a time region of an existing video.
+- `ltx_pipelines.hdr_ic_lora` - Video-to-video with HDR output (linear float via LogC3 inverse decode).
+- `ltx_pipelines.lipdub` - Lip dubbing / re-voicing with IC-LoRA and audio reference conditioning.
 
 Use `--help` with any pipeline module to see all available options and parameters.
 
@@ -79,6 +82,9 @@ Do you have an existing video to modify?
 Do you have an audio file to drive generation?
 ├─ YES → Use A2VidPipelineTwoStage (audio-to-video)
 │
+Do you need HDR output (linear float frames for EXR / tonemapping)?
+├─ YES → Use HDRICLoraPipeline (video-to-video with LogC3 inverse decode)
+│
 Do you need to condition on existing images/videos?
 ├─ YES → Do you have reference videos for video-to-video?
 │  ├─ YES → Use ICLoraPipeline
@@ -108,6 +114,8 @@ Do you need to condition on existing images/videos?
 | **KeyframeInterpolationPipeline** | 2 | ✅ | ✅ | Keyframes | Animation, interpolation |
 | **A2VidPipelineTwoStage** | 2 | ✅ | ✅ | Audio + Image | Audio-driven video generation |
 | **RetakePipeline** | 1 | ✅ | ❌ | Source Video | Regenerating a time region of a video |
+| **HDRICLoraPipeline** | 2 | ❌ | ✅ | Video | HDR video-to-video (linear float output for EXR) |
+| **LipDubPipeline** | 2 | ✅ | ✅ | Video + Audio | Lip dubbing with audio ref conditioning |
 
 ---
 
@@ -219,6 +227,48 @@ Single-stage generation that encodes the source video and audio into latents, ap
 
 ---
 
+### 9. HDRICLoraPipeline
+
+**Best for:** Video-to-video generation with HDR output for EXR export and offline tonemapping.
+
+**Source**: [`src/ltx_pipelines/hdr_ic_lora.py`](src/ltx_pipelines/hdr_ic_lora.py)
+
+Two-stage video-to-video on the distilled model with an HDR IC-LoRA. Decoded latents pass through an HDR inverse transform (ARRI LogC3, auto-detected from LoRA metadata) to produce a **linear HDR float** tensor `[f, h, w, c]`. Video-only (audio skipped). Text embeddings are pre-computed externally and loaded from a `.safetensors` file. Tonemapping and EXR saving are the caller's responsibility. LoRA and embeddings: [`Lightricks/LTX-2.3-22b-IC-LoRA-HDR`](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-HDR).
+
+**Extra CLI arguments:** `--input` (mp4 or directory, required), `--output-dir` (required), `--hdr-lora` (required), `--text-embeddings` (pre-computed `.safetensors`, required), `--num-frames`, `--spatial-tile` (tiled VAE decode tile size; reduce on lower-VRAM GPUs), `--skip-mp4` (EXR only, no H.264 preview), `--exr-half` (float16 EXR), `--high-quality` (generates 2x frames internally for smoother output, ~2x slower), `--offload {none,cpu,disk}` (weight offloading; disables FP8 quantization when not `none`).
+
+**Use when:** You need linear HDR float output for EXR export, color grading, or custom tonemapping workflows.
+
+---
+
+### 10. LipDubPipeline
+
+**Best for:** Lip dubbing, rephrasing while keeping the same speaker identity and matching lip movements to new audio.
+
+**Source**: [`src/ltx_pipelines/lipdub.py`](src/ltx_pipelines/lipdub.py)
+
+Uses IC-LoRA on a **distilled** checkpoint with a **single** lip-dub IC-LoRA applied in **both** stages. The reference clip provides video and audio reference tokens whose VAE latents are appended to the target audio sequence as frozen reference tokens. The frame count and frame rate are derived from the reference video (frame count is silently snapped to the nearest `8k+1`), so the CLI does not accept `--num-frames` or `--frame-rate`. Required: `--reference-video`. Optional: `--reference-strength`. LoRA: [`Lightricks/LTX-2.3-22b-IC-LoRA-LipDub`](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-LipDub).
+
+**Note:** Requires a distilled model checkpoint and one lip-dub IC-LoRA (`--lora` exactly once).
+
+**Use when:** Dubbing, rephrasing with matched lips and speaker identity.
+
+---
+
+### 11. T2AOneStagePipeline
+
+**Best for:** Text-to-audio — generating speech/audio only (no video) from a text prompt, e.g. driving an audio-style LoRA such as an accent LoRA.
+
+**Source**: [`src/ltx_pipelines/t2a_one_stage.py`](src/ltx_pipelines/t2a_one_stage.py)
+
+Single-stage, **audio-only** generation: the video branch is absent (`video=None`), so only the audio modality is denoised and decoded through the audio VAE + vocoder, producing a wave file. Audio duration is derived from `--num-frames` / `--frame-rate` (the same `8k+1` frame convention as video). Audio guidance (CFG/STG) is optional — the `--audio-*` flags default to the model's values; the video→audio cross-modal guidance is disabled since there is no video modality.
+
+**Extra CLI arguments (all optional, with sensible defaults):** `--num-frames`, `--frame-rate`, `--negative-prompt`, `--audio-cfg-guidance-scale`, `--audio-stg-guidance-scale`, `--audio-stg-blocks`, `--audio-rescale-scale`, `--audio-skip-step`. No `--height/--width/--image` (audio has no spatial dimensions).
+
+**Use when:** You need speech/audio from text alone, or to evaluate an audio-only LoRA (accent, voice style) without generating video.
+
+---
+
 ## 🎨 Conditioning Types
 
 Pipelines use different conditioning methods from [`ltx-core`](../ltx-core/) for controlling generation. See the [ltx-core conditioning documentation](../ltx-core/README.md#conditioning--control) for details.
@@ -334,7 +384,9 @@ PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True python -m ltx_pipelines.ti2vid_
 When authoring custom scripts, pass a `QuantizationPolicy` to pipeline classes:
 
 ```python
-from ltx_core.quantization import QuantizationPolicy
+from ltx_core.quantization.fp8_cast import build_policy as build_fp8_cast_policy
+# Alternative:
+# from ltx_core.quantization.fp8_scaled_mm import build_policy as build_fp8_scaled_mm_policy
 
 pipeline = TI2VidTwoStagesPipeline(
     checkpoint_path=ltx_model_path,
@@ -342,7 +394,7 @@ pipeline = TI2VidTwoStagesPipeline(
     spatial_upsampler_path=upsampler_path,
     gemma_root=gemma_root_path,
     loras=[],
-    quantization=QuantizationPolicy.fp8_cast(),  # or QuantizationPolicy.fp8_scaled_mm()
+    quantization=build_fp8_cast_policy(ltx_model_path),
 )
 pipeline(...)
 ```
@@ -363,6 +415,69 @@ By default, pipelines clean GPU memory (especially transformer weights) between
 # utils.cleanup_memory()  # Comment out if you have enough VRAM
 ```
 
+### Compilation (`torch.compile`)
+
+Compiling the transformer blocks with `torch.compile` speeds up inference. It is **opt-in and off by default**. The blocks are compiled shape-polymorphically (the sequence dimension is marked dynamic), so one compiled artifact serves any token count without recompiling.
+
+**CLI** — the `--compile` flag maps directly to `CompilationConfig`:
+
+| Form | Result |
+| ---- | ------ |
+| *(flag absent)* | eager, no compilation |
+| `--compile` | compile with defaults |
+| `--compile KEY=VALUE ...` | compile, overriding individual fields |
+
+```bash
+# Defaults
+python -m ltx_pipelines.ti2vid_two_stages --compile --checkpoint-path=...
+
+# reduce-overhead captures CUDA graphs -- the main latency lever for the denoising loop.
+# Off by default because graph capture reserves static memory pools (extra VRAM), so it
+# trades memory for speed; enable it when you have headroom.
+python -m ltx_pipelines.ti2vid_two_stages --compile mode=reduce-overhead --checkpoint-path=...
+
+# Several overrides at once
+python -m ltx_pipelines.ti2vid_two_stages \
+    --compile mode=max-autotune fullgraph=true dynamic=true --checkpoint-path=...
+```
+
+| Field | Values | Default | Notes |
+| ----- | ------ | ------- | ----- |
+| `mode` | `none`, `reduce-overhead`, `max-autotune`, … | `none` | `reduce-overhead`/`max-autotune` enable CUDA graphs |
+| `backend` | `inductor`, `eager`, … | `inductor` | |
+| `fullgraph` | `true`/`false` | `false` | |
+| `dynamic` | `auto`/`true`/`false` | `auto` | the seq dim is marked dynamic regardless |
+| `inductor_config` | JSON object or path to a `.json` | `{}` | `torch._inductor.config` overrides |
+| `dynamo_config` | JSON object or path to a `.json` | `{"inline_inbuilt_nn_modules": true, "cache_size_limit": 256}` | `torch._dynamo.config` overrides |
+
+**Controlling inductor / dynamo configs.** `inductor_config` and `dynamo_config` take either an inline JSON object or a path to a `.json` file, applied via `torch._inductor.config.patch(...)` / `torch._dynamo.config.patch(...)` around the compiled forward. They **replace the defaults wholesale — they do not merge**, so when overriding `dynamo_config` re-include any defaults you want to keep:
+
+```bash
+python -m ltx_pipelines.ti2vid_two_stages \
+    --compile 'inductor_config={"max_autotune": true}' \
+              'dynamo_config={"inline_inbuilt_nn_modules": true, "cache_size_limit": 256, "recompile_limit": 32}' \
+    --checkpoint-path=...
+```
+
+**Programmatically**, pass a `CompilationConfig` to the pipeline:
+
+```python
+from ltx_core.model.transformer.compiling import CompilationConfig
+
+pipeline = TI2VidTwoStagesPipeline(
+    ...,
+    compilation_config=CompilationConfig(mode="reduce-overhead"),
+)
+```
+
+**Faster cache loads: `unsafe_skip_cache_dynamic_shape_guards` (unsafe, opt-in).** Inductor's FX-graph cache re-checks the dynamic-shape guards stored with each entry on every lookup. Setting this flag skips that re-check (every entry is treated as a guard hit), which speeds up warm and cross-process cache loads. It is **not enabled by default** because it is a correctness hazard: a kernel first compiled at a small sequence length keeps int32 address arithmetic, and reusing it at a larger sequence length (roughly **>58k tokens/rank**) overflows int32 and reads out of bounds — surfacing as a CUDA illegal memory access or silently corrupted output. Only enable it when your token counts stay within the range the cached kernels were compiled for:
+
+```bash
+python -m ltx_pipelines.ti2vid_two_stages \
+    --compile 'inductor_config={"unsafe_skip_cache_dynamic_shape_guards": true}' \
+    --checkpoint-path=...
+```
+
 ### Denoising Loop Optimization
 
 **Gradient Estimation Denoising Loop:**
@@ -379,12 +494,13 @@ def denoising_loop(sigmas, video_state, audio_state, stepper):
         video_state=video_state,
         audio_state=audio_state,
         stepper=stepper,
-        denoise_fn=your_denoise_function,
+        transformer=transformer,
+        denoiser=denoiser,
         ge_gamma=2.0,  # Gradient estimation coefficient
     )
 ```
 
-This allows you to use **20-30 steps instead of 40** while maintaining quality. The gradient estimation function is available in [`pipeline_utils.py`](src/ltx_pipelines/utils/helpers.py).
+This allows you to use **20-30 steps instead of 40** while maintaining quality. The gradient estimation function is defined in [`samplers.py`](src/ltx_pipelines/utils/samplers.py).
 
 ---
 
@@ -400,15 +516,18 @@ This allows you to use **20-30 steps instead of 40** while maintaining quality.
 ## 📖 Example: Image-to-Video
 
 ```python
-from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps
-from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline
 from ltx_core.components.guiders import MultiModalGuiderParams
+from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps
+from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
+from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline
+from ltx_pipelines.utils.args import ImageConditioningInput
+from ltx_pipelines.utils.media_io import encode_video
 
 distilled_lora = [
     LoraPathStrengthAndSDOps(
         "/path/to/distilled_lora.safetensors",
         0.6,
-        LTXV_LORA_COMFY_RENAMING_MAP
+        LTXV_LORA_COMFY_RENAMING_MAP,
     ),
 ]
 
@@ -438,19 +557,31 @@ audio_guider_params = MultiModalGuiderParams(
     stg_blocks=[29],
 )
 
-# Generate video from image
-pipeline(
+# Generate video from image. The pipeline returns (video_iterator, audio);
+# the caller is responsible for encoding to file via encode_video().
+num_frames = 121
+frame_rate = 25.0
+tiling_config = TilingConfig.default()
+video, audio = pipeline(
     prompt="A serene landscape with mountains in the background",
-    output_path="output.mp4",
+    negative_prompt="worst quality, low quality, blurry, distorted",
     seed=42,
     height=512,
     width=768,
-    num_frames=121,
-    frame_rate=25.0,
+    num_frames=num_frames,
+    frame_rate=frame_rate,
     num_inference_steps=40,
     video_guider_params=video_guider_params,
     audio_guider_params=audio_guider_params,
-    images=[ImageConditioningInput("input_image.jpg", 0, 1.0, 33)],  # Image at frame 0, strength 1.0, CRF 33
+    images=[ImageConditioningInput("input_image.jpg", 0, 1.0, 33)],  # path, frame_idx=0, strength=1.0, crf=33
+    tiling_config=tiling_config,
+)
+encode_video(
+    video=video,
+    fps=frame_rate,
+    audio=audio,
+    output_path="output.mp4",
+    video_chunks_number=get_video_chunks_number(num_frames, tiling_config),
 )
 ```
 

packages/ltx-pipelines/pyproject.toml

@@ -1,10 +1,10 @@
 [project]
 name = "ltx-pipelines"
-version = "1.1.1"
+version = "1.1.6"
 description = "Pipelines implementation for Lightricks' LTX-2 model"
 readme = "README.md"
 requires-python = ">=3.10"
-dependencies = ["ltx-core", "av", "tqdm", "pillow"]
+dependencies = ["ltx-core", "av", "tqdm", "pillow", "openimageio"]
 
 [build-system]
 requires = ["uv_build>=0.9.8,<0.10.0"]

packages/ltx-pipelines/src/ltx_pipelines/__init__.py

@@ -2,9 +2,11 @@
 LTX-2 Pipelines: High-level video generation pipelines and utilities.
 This package provides ready-to-use pipelines for video generation:
 - TI2VidOneStagePipeline: Text/image-to-video in a single stage
+- T2AOneStagePipeline: Text-to-audio in a single stage (audio-only output)
 - TI2VidTwoStagesPipeline: Two-stage generation with upsampling
 - DistilledPipeline: Fast distilled two-stage generation
 - ICLoraPipeline: Image/video conditioning with distilled LoRA
+- LipDubPipeline: Lip dubbing with IC-LoRA and audio conditioning
 - KeyframeInterpolationPipeline: Keyframe-based video interpolation
 - RetakePipeline: Regenerate a time region (retake) of an existing video
 For more detailed components and utilities, import from specific submodules
@@ -15,7 +17,9 @@ from ltx_pipelines.a2vid_two_stage import A2VidPipelineTwoStage
 from ltx_pipelines.distilled import DistilledPipeline
 from ltx_pipelines.ic_lora import ICLoraPipeline
 from ltx_pipelines.keyframe_interpolation import KeyframeInterpolationPipeline
+from ltx_pipelines.lipdub import LipDubPipeline
 from ltx_pipelines.retake import RetakePipeline
+from ltx_pipelines.t2a_one_stage import T2AOneStagePipeline
 from ltx_pipelines.ti2vid_one_stage import TI2VidOneStagePipeline
 from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline
 
@@ -24,7 +28,9 @@ __all__ = [
     "DistilledPipeline",
     "ICLoraPipeline",
     "KeyframeInterpolationPipeline",
+    "LipDubPipeline",
     "RetakePipeline",
+    "T2AOneStagePipeline",
     "TI2VidOneStagePipeline",
     "TI2VidTwoStagesPipeline",
 ]

packages/ltx-pipelines/src/ltx_pipelines/a2vid_two_stage.py

@@ -9,6 +9,7 @@ from ltx_core.components.schedulers import LTX2Scheduler
 from ltx_core.loader import LoraPathStrengthAndSDOps
 from ltx_core.loader.registry import Registry
 from ltx_core.model.audio_vae import encode_audio as vae_encode_audio
+from ltx_core.model.transformer.compiling import CompilationConfig
 from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
 from ltx_core.quantization import QuantizationPolicy
 from ltx_core.types import Audio, AudioLatentShape, VideoPixelShape
@@ -31,7 +32,7 @@ from ltx_pipelines.utils.helpers import (
     get_device,
 )
 from ltx_pipelines.utils.media_io import decode_audio_from_file, encode_video
-from ltx_pipelines.utils.types import ModalitySpec
+from ltx_pipelines.utils.types import ModalitySpec, OffloadMode
 
 
 class A2VidPipelineTwoStage:
@@ -52,13 +53,16 @@ class A2VidPipelineTwoStage:
         device: torch.device | None = None,
         quantization: QuantizationPolicy | None = None,
         registry: Registry | None = None,
-        torch_compile: bool = False,
+        compilation_config: CompilationConfig | None = None,
+        offload_mode: OffloadMode = OffloadMode.NONE,
     ):
         self.device = device or get_device()
         self.dtype = torch.bfloat16
         self._scheduler = LTX2Scheduler()
 
-        self.prompt_encoder = PromptEncoder(checkpoint_path, gemma_root, self.dtype, self.device, registry=registry)
+        self.prompt_encoder = PromptEncoder(
+            checkpoint_path, gemma_root, self.dtype, self.device, registry=registry, offload_mode=offload_mode
+        )
         self.image_conditioner = ImageConditioner(checkpoint_path, self.dtype, self.device, registry=registry)
         self.audio_conditioner = AudioConditioner(checkpoint_path, self.dtype, self.device, registry=registry)
         self.stage_1 = DiffusionStage(
@@ -68,7 +72,8 @@ class A2VidPipelineTwoStage:
             loras=tuple(loras),
             quantization=quantization,
             registry=registry,
-            torch_compile=torch_compile,
+            compilation_config=compilation_config,
+            offload_mode=offload_mode,
         )
         stage_2_loras = (*tuple(loras), *tuple(distilled_lora))
         self.stage_2 = DiffusionStage(
@@ -78,7 +83,8 @@ class A2VidPipelineTwoStage:
             loras=stage_2_loras,
             quantization=quantization,
             registry=registry,
-            torch_compile=torch_compile,
+            compilation_config=compilation_config,
+            offload_mode=offload_mode,
         )
         self.upsampler = VideoUpsampler(
             checkpoint_path, spatial_upsampler_path, self.dtype, self.device, registry=registry
@@ -102,7 +108,6 @@ class A2VidPipelineTwoStage:
         audio_max_duration: float | None = None,
         tiling_config: TilingConfig | None = None,
         enhance_prompt: bool = False,
-        streaming_prefetch_count: int | None = None,
         max_batch_size: int = 1,
         stage_1_sigmas: torch.Tensor | None = None,
         stage_2_sigmas: torch.Tensor = STAGE_2_DISTILLED_SIGMAS,
@@ -117,7 +122,6 @@ class A2VidPipelineTwoStage:
             [prompt, negative_prompt],
             enhance_first_prompt=enhance_prompt,
             enhance_prompt_image=images[0][0] if len(images) > 0 else None,
-            streaming_prefetch_count=streaming_prefetch_count,
         )
         v_context_p, a_context_p = ctx_p.video_encoding, ctx_p.audio_encoding
         v_context_n, _ = ctx_n.video_encoding, ctx_n.audio_encoding
@@ -183,7 +187,6 @@ class A2VidPipelineTwoStage:
                 noise_scale=0.0,
                 initial_latent=encoded_audio_latent,
             ),
-            streaming_prefetch_count=streaming_prefetch_count,
             max_batch_size=max_batch_size,
         )
 
@@ -223,7 +226,6 @@ class A2VidPipelineTwoStage:
                 noise_scale=0.0,
                 initial_latent=encoded_audio_latent,
             ),
-            streaming_prefetch_count=streaming_prefetch_count,
         )
 
         decoded_video = self.video_decoder(video_state.latent, tiling_config, generator)
@@ -237,7 +239,7 @@ class A2VidPipelineTwoStage:
 
 @torch.inference_mode()
 def main() -> None:
-    logging.getLogger().setLevel(logging.INFO)
+    logging.basicConfig(level=logging.INFO)
     parser = default_2_stage_arg_parser()
     parser.add_argument(
         "--audio-path",
@@ -265,7 +267,8 @@ def main() -> None:
         gemma_root=args.gemma_root,
         loras=tuple(args.lora) if args.lora else (),
         quantization=args.quantization,
-        torch_compile=args.compile,
+        compilation_config=args.compile,
+        offload_mode=args.offload_mode,
     )
     tiling_config = TilingConfig.default()
     video_chunks_number = get_video_chunks_number(args.num_frames, tiling_config)
@@ -294,7 +297,6 @@ def main() -> None:
         audio_max_duration=args.audio_max_duration
         if args.audio_max_duration is not None
         else args.num_frames / args.frame_rate,
-        streaming_prefetch_count=args.streaming_prefetch_count,
         max_batch_size=args.max_batch_size,
     )
 

packages/ltx-pipelines/src/ltx_pipelines/distilled.py

@@ -6,6 +6,7 @@ import torch
 from ltx_core.components.noisers import GaussianNoiser
 from ltx_core.loader import LoraPathStrengthAndSDOps
 from ltx_core.loader.registry import Registry
+from ltx_core.model.transformer.compiling import CompilationConfig
 from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
 from ltx_core.quantization import QuantizationPolicy
 from ltx_core.types import Audio
@@ -34,7 +35,7 @@ from ltx_pipelines.utils.helpers import (
     get_device,
 )
 from ltx_pipelines.utils.media_io import encode_video
-from ltx_pipelines.utils.types import ModalitySpec
+from ltx_pipelines.utils.types import ModalitySpec, OffloadMode
 
 
 class DistilledPipeline:
@@ -53,13 +54,19 @@ class DistilledPipeline:
         device: torch.device | None = None,
         quantization: QuantizationPolicy | None = None,
         registry: Registry | None = None,
-        torch_compile: bool = False,
+        compilation_config: CompilationConfig | None = None,
+        offload_mode: OffloadMode = OffloadMode.NONE,
     ):
         self.device = device or get_device()
         self.dtype = torch.bfloat16
 
         self.prompt_encoder = PromptEncoder(
-            distilled_checkpoint_path, gemma_root, self.dtype, self.device, registry=registry
+            distilled_checkpoint_path,
+            gemma_root,
+            self.dtype,
+            self.device,
+            registry=registry,
+            offload_mode=offload_mode,
         )
         self.image_conditioner = ImageConditioner(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
         self.stage = DiffusionStage(
@@ -69,7 +76,8 @@ class DistilledPipeline:
             loras=tuple(loras),
             quantization=quantization,
             registry=registry,
-            torch_compile=torch_compile,
+            compilation_config=compilation_config,
+            offload_mode=offload_mode,
         )
         self.upsampler = VideoUpsampler(
             distilled_checkpoint_path, spatial_upsampler_path, self.dtype, self.device, registry=registry
@@ -88,7 +96,6 @@ class DistilledPipeline:
         images: list[ImageConditioningInput],
         tiling_config: TilingConfig | None = None,
         enhance_prompt: bool = False,
-        streaming_prefetch_count: int | None = None,
         stage_1_sigmas: torch.Tensor = DISTILLED_SIGMAS,
         stage_2_sigmas: torch.Tensor = STAGE_2_DISTILLED_SIGMAS,
     ) -> tuple[Iterator[torch.Tensor], Audio]:
@@ -102,7 +109,6 @@ class DistilledPipeline:
             [prompt],
             enhance_first_prompt=enhance_prompt,
             enhance_prompt_image=images[0][0] if len(images) > 0 else None,
-            streaming_prefetch_count=streaming_prefetch_count,
         )
         video_context, audio_context = ctx_p.video_encoding, ctx_p.audio_encoding
 
@@ -130,7 +136,6 @@ class DistilledPipeline:
             fps=frame_rate,
             video=ModalitySpec(context=video_context, conditionings=stage_1_conditionings),
             audio=ModalitySpec(context=audio_context),
-            streaming_prefetch_count=streaming_prefetch_count,
         )
 
         # Stage 2: Upsample and refine the video at higher resolution with distilled LORA.
@@ -167,7 +172,6 @@ class DistilledPipeline:
                 noise_scale=stage_2_sigmas[0].item(),
                 initial_latent=audio_state.latent,
             ),
-            streaming_prefetch_count=streaming_prefetch_count,
         )
 
         decoded_video = self.video_decoder(video_state.latent, tiling_config, generator)
@@ -177,7 +181,7 @@ class DistilledPipeline:
 
 @torch.inference_mode()
 def main() -> None:
-    logging.getLogger().setLevel(logging.INFO)
+    logging.basicConfig(level=logging.INFO)
     checkpoint_path = detect_checkpoint_path(distilled=True)
     params = detect_params(checkpoint_path)
     parser = default_2_stage_distilled_arg_parser(params=params)
@@ -188,7 +192,8 @@ def main() -> None:
         gemma_root=args.gemma_root,
         loras=tuple(args.lora) if args.lora else (),
         quantization=args.quantization,
-        torch_compile=args.compile,
+        compilation_config=args.compile,
+        offload_mode=args.offload_mode,
     )
     tiling_config = TilingConfig.default()
     video_chunks_number = get_video_chunks_number(args.num_frames, tiling_config)
@@ -202,7 +207,6 @@ def main() -> None:
         images=args.images,
         tiling_config=tiling_config,
         enhance_prompt=args.enhance_prompt,
-        streaming_prefetch_count=args.streaming_prefetch_count,
     )
 
     encode_video(

packages/ltx-pipelines/src/ltx_pipelines/hdr_ic_lora.py

@@ -0,0 +1,904 @@
+"""HDR IC-LoRA pipeline: two-stage video generation with HDR output.
+Extends the standard IC-LoRA pipeline with HDR decode via LogC3 inverse
+transform.  ``__call__`` returns a **linear HDR float** tensor
+``[f, h, w, c]``; tonemapping and EXR saving are the caller's
+responsibility.
+Text embeddings must be pre-computed externally (e.g. using
+``PromptEncoder`` from ``ltx_pipelines.utils.blocks`` with a Gemma text
+encoder) and saved as a ``.safetensors`` file with ``video_context``
+and ``audio_context`` tensors (via ``safetensors.torch.save_file``).
+The path is passed via ``text_embeddings_path``.
+Run as a script for batch inference::
+    python -m ltx_pipelines.hdr_ic_lora \\
+        --input ./videos/ \\
+        --output-dir ./hdr-output \\
+        --distilled-checkpoint-path /models/ltx-2.3-22b-distilled.safetensors \\
+        --spatial-upsampler-path /models/ltx-2.3-spatial-upscaler-x2-1.0.safetensors \\
+        --hdr-lora /path/to/hdr_lora.safetensors \\
+        --text-embeddings /path/to/hdr_scene_emb.safetensors \\
+        --num-frames 161
+Supports resolutions up to 4K (3840x2160 @ 121 frames on 80 GB,
+49 frames on 48 GB).  The caller is responsible for choosing a resolution
+and frame count that fits in GPU memory.  See ``--help`` for a reference
+table, or use ``ltx_pipelines.utils.vram_budget.max_frames_for_resolution``
+to query your specific configuration.
+"""
+
+import dataclasses
+import logging
+from dataclasses import replace
+from pathlib import Path
+
+import torch
+from einops import rearrange
+from safetensors import safe_open
+
+from ltx_core.components.noisers import GaussianNoiser
+from ltx_core.components.patchifiers import VideoLatentPatchifier
+from ltx_core.conditioning import (
+    ConditioningItem,
+    VideoConditionByReferenceLatent,
+)
+from ltx_core.hdr import apply_hdr_decode_postprocess
+from ltx_core.loader import LoraPathStrengthAndSDOps
+from ltx_core.loader.registry import Registry
+from ltx_core.loader.sd_ops import LTXV_LORA_COMFY_RENAMING_MAP
+from ltx_core.modality_tiling import VideoModalityTilingHelper
+from ltx_core.model.video_vae import TilingConfig, VideoEncoder
+from ltx_core.quantization import QuantizationPolicy
+from ltx_core.tiling import DimensionTilingConfig, TileCountConfig
+from ltx_core.tools import VideoLatentTools
+from ltx_core.types import VideoLatentShape, VideoPixelShape
+from ltx_pipelines.utils.blocks import (
+    DiffusionStage,
+    ImageConditioner,
+    VideoDecoder,
+    VideoUpsampler,
+)
+from ltx_pipelines.utils.constants import DISTILLED_SIGMA_VALUES, STAGE_2_DISTILLED_SIGMA_VALUES
+from ltx_pipelines.utils.denoisers import SimpleDenoiser
+from ltx_pipelines.utils.helpers import get_device, modality_from_latent_state
+from ltx_pipelines.utils.media_io import ResizeMode, align_resolution, load_video_conditioning_hdr
+from ltx_pipelines.utils.quantization_factory import QuantizationKind
+from ltx_pipelines.utils.types import ModalitySpec, OffloadMode
+
+logger = logging.getLogger(__name__)
+
+# ---------------------------------------------------------------------------
+# Constants
+# ---------------------------------------------------------------------------
+
+DEFAULT_NUM_FRAMES = 161
+MIN_RESOLUTION = 64
+ALIGNMENT_DIVISOR = 64
+
+# Conditioning videos whose spatial resolution (H x W) exceeds this value are
+# encoded with the tiled encoder. The default (512 x 768) is suitable for
+# H100-80GB. On lower-VRAM GPUs pass tiled_vae_encode_pixel_threshold=256*256
+# to the pipeline constructor.
+TILED_VAE_ENCODE_PIXEL_THRESHOLD = 512 * 768
+
+_DEFAULT_QUANTIZATION = QuantizationKind.FP8_CAST
+
+# Default stage-2 configuration: one refinement phase with modest 2-way tiling
+# in every dimension and a short 2-step distilled sigma schedule.
+_S2 = STAGE_2_DISTILLED_SIGMA_VALUES
+
+_TILED_2F2H2W_OV8_6 = TileCountConfig(
+    frames=DimensionTilingConfig(2, 8),
+    height=DimensionTilingConfig(2, 6),
+    width=DimensionTilingConfig(2, 6),
+)
+
+STAGE2_TILINGS = [_TILED_2F2H2W_OV8_6]
+STAGE2_SIGMAS = [[_S2[0], _S2[1], 0.0]]
+STAGE2_USE_IC_LORA = [True]
+
+
+def _clamp_dim_tiling(cfg: DimensionTilingConfig, dim_size: int, axis: str) -> DimensionTilingConfig:
+    """Clamp a single dim's tile count and overlap to the latent's extent.
+    ``split_by_count`` requires ``overlap < tile_size``; with
+    ``tile_size = (dim_size + overlap*(n-1)) // n`` this reduces to
+    ``overlap <= dim_size - n``. When the configured overlap exceeds this
+    bound it is clamped; if the latent is too small to hold ``n`` tiles
+    at all, tiling falls back to a single tile on this axis.
+    """
+    n = cfg.num_tiles
+    if n <= 1:
+        return cfg
+    if dim_size < n:
+        logger.warning(
+            "%s tiling: dim_size=%d < num_tiles=%d; falling back to 1 tile on this axis.",
+            axis,
+            dim_size,
+            n,
+        )
+        return DimensionTilingConfig(1, 0)
+    max_overlap = dim_size - n
+    if cfg.overlap <= max_overlap:
+        return cfg
+    logger.warning(
+        "%s tiling: overlap=%d exceeds latent bound (%d); clamping to %d.",
+        axis,
+        cfg.overlap,
+        max_overlap,
+        max_overlap,
+    )
+    return DimensionTilingConfig(n, max_overlap)
+
+
+def _clamp_tile_to_latent(tiling: TileCountConfig, latent_shape: tuple[int, int, int]) -> TileCountConfig:
+    """Clamp frame, height, and width tilings to the latent's extents.
+    ``latent_shape`` is ``(F, H, W)`` in latent units.
+    """
+    f, h, w = latent_shape
+    return replace(
+        tiling,
+        frames=_clamp_dim_tiling(tiling.frames, f, "Frame"),
+        height=_clamp_dim_tiling(tiling.height, h, "Height"),
+        width=_clamp_dim_tiling(tiling.width, w, "Width"),
+    )
+
+
+# Default tiling config (spatial tile 1280 px, overlap 256 px; temporal 32
+# frames, overlap 16). On GPUs with < 80 GB VRAM you may need to shrink
+# the spatial tile size (e.g. 768) to avoid OOM during VAE decode.
+DEFAULT_SPATIAL_TILE = 1280
+DEFAULT_SPATIAL_OVERLAP = 256
+DEFAULT_TEMPORAL_TILE = 32
+DEFAULT_TEMPORAL_OVERLAP = 16
+
+
+# ---------------------------------------------------------------------------
+# HDR LoRA config
+# ---------------------------------------------------------------------------
+
+
+@dataclasses.dataclass(frozen=True)
+class HdrLoraConfig:
+    """Explicit HDR LoRA parameters.
+    Read from LoRA safetensors metadata by :func:`read_hdr_lora_config`, or
+    constructed manually for testing.
+    """
+
+    hdr_transform: str = "logc3"
+    reference_downscale_factor: int = 1
+
+
+def read_hdr_lora_config(lora_path: str) -> HdrLoraConfig | None:
+    """Read HDR config from LoRA safetensors metadata.
+    Returns ``None`` when the LoRA has no HDR metadata.
+    """
+    try:
+        with safe_open(lora_path, framework="pt") as f:
+            metadata = f.metadata() or {}
+    except (OSError, ValueError) as e:
+        logger.warning("Failed to read metadata from LoRA file '%s': %s", lora_path, e)
+        return None
+
+    hdr_transform = metadata.get("hdr_transform", "")
+    has_hdr = bool(hdr_transform or metadata.get("use_hdr_transform"))
+    if not has_hdr:
+        return None
+
+    transform = hdr_transform if hdr_transform and hdr_transform != "true" else "logc3"
+    scale = int(metadata.get("reference_downscale_factor", 1))
+    return HdrLoraConfig(hdr_transform=transform, reference_downscale_factor=scale)
+
+
+# ---------------------------------------------------------------------------
+# Pipeline
+# ---------------------------------------------------------------------------
+
+
+class HDRICLoraPipeline:
+    """Two-stage IC-LoRA pipeline with HDR support.
+    Same two-stage architecture as ICLoraPipeline (half-res generation + 2x
+    upscale refinement), with HDR decode via LogC3 inverse.
+    ``__call__`` returns a **linear HDR float** tensor ``[f, h, w, c]``.
+    Tonemapping and EXR saving are the caller's responsibility.
+    """
+
+    def __init__(
+        self,
+        distilled_checkpoint_path: str,
+        spatial_upsampler_path: str,
+        hdr_lora: str | Path,
+        text_embeddings_path: str | Path,
+        device: torch.device | None = None,
+        quantization: QuantizationPolicy | QuantizationKind | None = _DEFAULT_QUANTIZATION,
+        registry: Registry | None = None,
+        hdr_lora_config: HdrLoraConfig | None = None,
+        tiled_vae_encode_pixel_threshold: int = TILED_VAE_ENCODE_PIXEL_THRESHOLD,
+        offload_mode: OffloadMode = OffloadMode.NONE,
+    ):
+        """
+        Args:
+            distilled_checkpoint_path: Path to the distilled model checkpoint.
+            spatial_upsampler_path: Path to the spatial upsampler checkpoint.
+            hdr_lora: Path to the HDR IC-LoRA ``.safetensors`` file.
+            text_embeddings_path: Path to pre-computed text embeddings
+                (``.safetensors`` file with ``video_context`` and
+                ``audio_context`` tensors).
+            device: Target device. Auto-detected when ``None``.
+            quantization: Quantization policy. Defaults to ``fp8_cast``.
+            registry: Optional model registry for caching loaded components.
+            hdr_lora_config: Explicit HDR LoRA config override. When ``None``,
+                auto-detected from LoRA safetensors metadata.
+            tiled_vae_encode_pixel_threshold: Conditioning videos whose spatial
+                area (H x W) exceeds this value are encoded with the tiled
+                encoder. Default ``512 * 768`` is suitable for 80 GB GPUs.
+                Use ``256 * 256`` on GPUs with less VRAM.
+            offload_mode: Weight offloading strategy for diffusion stages.
+        """
+        self.device = device or get_device()
+        self._tiled_vae_encode_threshold = tiled_vae_encode_pixel_threshold
+        if isinstance(quantization, QuantizationKind):
+            quantization = quantization.to_policy(checkpoint_path=distilled_checkpoint_path)
+        if offload_mode != OffloadMode.NONE and quantization is not None:
+            logger.info("Offload mode enabled — disabling quantization (not supported with layer streaming).")
+            quantization = None
+        self.dtype = torch.bfloat16
+
+        lora_path = str(Path(hdr_lora).resolve())
+        loras = (LoraPathStrengthAndSDOps(lora_path, 1.0, LTXV_LORA_COMFY_RENAMING_MAP),)
+
+        # Load pre-computed text embeddings from safetensors.
+        emb_path = Path(text_embeddings_path)
+        logger.info("Loading text embeddings from %s", emb_path)
+        with safe_open(emb_path, framework="pt", device=str(self.device)) as f:
+            self.text_embeddings: tuple[torch.Tensor, torch.Tensor] = (
+                f.get_tensor("video_context"),
+                f.get_tensor("audio_context"),
+            )
+
+        self.image_conditioner = ImageConditioner(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
+        self.stage_1 = DiffusionStage(
+            distilled_checkpoint_path,
+            self.dtype,
+            self.device,
+            loras=loras,
+            quantization=quantization,
+            registry=registry,
+            offload_mode=offload_mode,
+        )
+        self.stage_2 = DiffusionStage(
+            distilled_checkpoint_path,
+            self.dtype,
+            self.device,
+            loras=loras,
+            quantization=quantization,
+            registry=registry,
+            offload_mode=offload_mode,
+        )
+        self.upsampler = VideoUpsampler(
+            distilled_checkpoint_path, spatial_upsampler_path, self.dtype, self.device, registry=registry
+        )
+        self.video_decoder = VideoDecoder(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
+
+        # HDR config: explicit override, or auto-detect from LoRA metadata.
+        if hdr_lora_config is not None:
+            self._hdr_config: HdrLoraConfig | None = hdr_lora_config
+        else:
+            self._hdr_config = read_hdr_lora_config(lora_path)
+
+        if self._hdr_config is not None:
+            logger.info("[HDR IC-LoRA] HDR mode enabled (%s decode)", self._hdr_config.hdr_transform)
+
+    @property
+    def hdr_transform(self) -> str:
+        """Active HDR transform name (defaults to 'logc3')."""
+        return self._hdr_config.hdr_transform if self._hdr_config is not None else "logc3"
+
+    @property
+    def reference_downscale_factor(self) -> int:
+        """Reference video downscale factor from HDR LoRA config."""
+        return self._hdr_config.reference_downscale_factor if self._hdr_config is not None else 1
+
+    def __call__(  # noqa: PLR0913
+        self,
+        seed: int,
+        height: int,
+        width: int,
+        num_frames: int,
+        frame_rate: float,
+        video_conditioning: list[tuple[str, float]],
+        tiling_config: TilingConfig | None = None,
+        high_quality_hdr: bool = False,
+        stage2_tilings: list[TileCountConfig] | None = None,
+        stage2_sigmas: list[list[float]] | None = None,
+        stage2_use_ic_lora: list[bool] | None = None,
+    ) -> torch.Tensor:
+        """Generate video with IC-LoRA conditioning and HDR output.
+        Returns a linear HDR float tensor ``[f, h, w, c]``.
+        Args:
+            seed: Random seed for reproducibility.
+            height: Desired output video height in pixels. Aligned internally
+                to the nearest multiple of 64 (rounded up). Decoded output is
+                cropped back to this size.
+            width: Desired output video width in pixels. Same alignment rules
+                as *height*.
+            num_frames: Number of frames to generate.
+            frame_rate: Output video frame rate.
+            video_conditioning: List of (path, strength) tuples for IC-LoRA video conditioning.
+            high_quality_hdr: High-quality HDR mode. Duplicates each conditioning
+                frame and generates at 2x frame count, then keeps every other
+                output frame. Reduces temporal artifacts at the cost of ~2x
+                generation time.
+        Returns:
+            Linear HDR float tensor ``[f, h, w, c]``.
+        """
+        # In high-quality HDR mode, generate 2*N - 1 frames internally
+        # (satisfies (n-1)%8==0 when N itself does), then keep every other frame.
+        if high_quality_hdr:
+            gen_num_frames = 2 * num_frames - 1
+            logger.info("[HDR IC-LoRA] High-quality HDR: %d -> %d internal frames", num_frames, gen_num_frames)
+        else:
+            gen_num_frames = num_frames
+        gen_w, gen_h, crop_w, crop_h = align_resolution(
+            width, height, ResizeMode.REFLECT_PAD, divisor=ALIGNMENT_DIVISOR
+        )
+        if gen_h < MIN_RESOLUTION or gen_w < MIN_RESOLUTION:
+            raise ValueError(
+                f"Resolution ({width}x{height}) is too small after alignment "
+                f"(got {gen_w}x{gen_h}, need at least {MIN_RESOLUTION}x{MIN_RESOLUTION})."
+            )
+        needs_crop = crop_w != gen_w or crop_h != gen_h
+        if needs_crop:
+            logger.info(
+                "[HDR IC-LoRA] Aligned %dx%d -> %dx%d, will crop to %dx%d",
+                width,
+                height,
+                gen_w,
+                gen_h,
+                crop_w,
+                crop_h,
+            )
+
+        generator = torch.Generator(device=self.device).manual_seed(seed)
+        noiser = GaussianNoiser(generator=generator)
+
+        video_context, _ = self.text_embeddings
+
+        # Stage 1: Initial low resolution video generation.
+        s1_w, s1_h = gen_w // 2, gen_h // 2
+
+        stage_1_conditionings = self.image_conditioner(
+            lambda enc: self._create_conditionings(
+                video_conditioning=video_conditioning,
+                height=s1_h,
+                width=s1_w,
+                video_encoder=enc,
+                num_frames=gen_num_frames,
+                tiling_config=tiling_config,
+                high_quality_hdr=high_quality_hdr,
+            )
+        )
+
+        stage_1_sigmas = torch.Tensor(DISTILLED_SIGMA_VALUES).to(self.device)
+
+        # HDR is video-only: skip the audio stream to avoid denoising 5B audio params.
+        video_state, _ = self.stage_1(
+            denoiser=SimpleDenoiser(video_context, None),
+            sigmas=stage_1_sigmas,
+            noiser=noiser,
+            width=s1_w,
+            height=s1_h,
+            frames=gen_num_frames,
+            fps=frame_rate,
+            video=ModalitySpec(
+                context=video_context,
+                conditionings=stage_1_conditionings,
+            ),
+        )
+
+        if stage2_tilings is None:
+            stage2_tilings = list(STAGE2_TILINGS)
+        if stage2_sigmas is None:
+            stage2_sigmas = [list(s) for s in STAGE2_SIGMAS]
+        if stage2_use_ic_lora is None:
+            stage2_use_ic_lora = list(STAGE2_USE_IC_LORA)
+        if not (len(stage2_tilings) == len(stage2_sigmas) == len(stage2_use_ic_lora)):
+            raise ValueError("stage2_tilings, stage2_sigmas, and stage2_use_ic_lora must have equal length")
+
+        # Stage 2: Upsample and refine at full resolution.
+        upscaled_video_latent = self.upsampler(video_state.latent[:1])
+
+        stage_2_conditionings = self.image_conditioner(
+            lambda enc: self._create_conditionings(
+                video_conditioning=video_conditioning,
+                height=gen_h,
+                width=gen_w,
+                video_encoder=enc,
+                num_frames=gen_num_frames,
+                tiling_config=tiling_config,
+                high_quality_hdr=high_quality_hdr,
+            )
+        )
+        # video_tools is required by TiledDataParallelBuilder when stage_2 is
+        # wrapped for multi-GPU
+        stage2_video_tools = VideoLatentTools(
+            VideoLatentPatchifier(patch_size=1),
+            VideoLatentShape.from_pixel_shape(
+                VideoPixelShape(
+                    batch=1,
+                    frames=gen_num_frames,
+                    height=gen_h,
+                    width=gen_w,
+                    fps=frame_rate,
+                )
+            ),
+            frame_rate,
+        )
+        with self.stage_2.model_context(video_tools=stage2_video_tools) as transformer:
+            phase_latent = upscaled_video_latent
+            for phase_idx, (tiling, sigmas_list, use_ic) in enumerate(
+                zip(stage2_tilings, stage2_sigmas, stage2_use_ic_lora, strict=True)
+            ):
+                diffusion_tiling = _clamp_tile_to_latent(tiling, tuple(phase_latent.shape[2:5]))
+                conditionings = stage_2_conditionings if use_ic else []
+                sigma_t = torch.tensor(sigmas_list, dtype=torch.float32, device=self.device)
+                logger.info(
+                    "[Stage 2 / phase %d] sigmas=%s ic_lora=%s tiling_h=%s tiling_w=%s",
+                    phase_idx,
+                    sigmas_list,
+                    use_ic,
+                    diffusion_tiling.height,
+                    diffusion_tiling.width,
+                )
+                phase_latent = self._run_stage2_phase(
+                    transformer=transformer,
+                    latent=phase_latent,
+                    conditionings=conditionings,
+                    tiling=diffusion_tiling,
+                    sigmas=sigma_t,
+                    v_ctx=video_context,
+                    frame_rate=frame_rate,
+                    seed=seed,
+                )
+
+        final_video_latent = phase_latent
+
+        crop_size = (crop_w, crop_h) if needs_crop else None
+        return self._decode_video(
+            final_video_latent,
+            tiling_config,
+            generator,
+            crop_size,
+            high_quality_hdr=high_quality_hdr,
+        )
+
+    def _run_stage2_phase(
+        self,
+        transformer: object,
+        latent: torch.Tensor,
+        conditionings: list[ConditioningItem],
+        tiling: TileCountConfig,
+        sigmas: torch.Tensor,
+        v_ctx: torch.Tensor,
+        frame_rate: float,
+        seed: int,
+    ) -> torch.Tensor:
+        """Run one stage-2 denoising phase with optional IC-LoRA conditioning.
+        Each tile calls ``stage_2.run()`` with a tile-sized ``ModalitySpec`` for
+        video only (audio is omitted entirely for HDR). IC-LoRA conditionings
+        are sliced spatially to match each tile's extent.
+        """
+        batch, n_channels, n_frames, n_height, n_width = latent.shape
+        full_shape = VideoLatentShape(batch=batch, channels=n_channels, frames=n_frames, height=n_height, width=n_width)
+        full_tools = VideoLatentTools(VideoLatentPatchifier(patch_size=1), full_shape, frame_rate)
+        helper = VideoModalityTilingHelper(tiling, full_tools)
+
+        ref_initial = full_tools.create_initial_state(device=self.device, dtype=self.dtype)
+        ref_modality = modality_from_latent_state(ref_initial, v_ctx, sigmas[0])
+        n_gen = full_tools.target_shape.token_count()
+        blend_output = torch.zeros(batch, n_gen, n_channels, device=self.device, dtype=self.dtype)
+        patchifier = VideoLatentPatchifier(patch_size=1)
+        df = self.reference_downscale_factor
+
+        for tile_idx, tile in enumerate(helper.tiles):
+            _, ctx = helper.tile_modality(ref_modality, tile, normalize_positions=True)
+            frame_s, height_s, width_s = tile.in_coords
+            tile_h = height_s.stop - height_s.start
+            tile_w = width_s.stop - width_s.start
+            tile_f = frame_s.stop - frame_s.start
+
+            tile_conditionings = [
+                VideoConditionByReferenceLatent(
+                    latent=cond.latent[
+                        :,
+                        :,
+                        frame_s,
+                        slice(height_s.start // df, height_s.stop // df),
+                        slice(width_s.start // df, width_s.stop // df),
+                    ].to(device=self.device, dtype=self.dtype),
+                    downscale_factor=cond.downscale_factor,
+                    strength=cond.strength,
+                )
+                for cond in conditionings
+            ]
+
+            tile_video_state, _ = self.stage_2.run(
+                transformer=transformer,
+                denoiser=SimpleDenoiser(v_ctx, None),
+                sigmas=sigmas,
+                noiser=GaussianNoiser(generator=torch.Generator(device=self.device).manual_seed(seed + tile_idx)),
+                width=tile_w * 32,
+                height=tile_h * 32,
+                frames=(tile_f - 1) * 8 + 1,
+                fps=frame_rate,
+                video=ModalitySpec(
+                    context=v_ctx,
+                    conditionings=tile_conditionings,
+                    noise_scale=sigmas[0].item(),
+                    initial_latent=latent[:, :, frame_s, height_s, width_s].to(device=self.device, dtype=self.dtype),
+                ),
+            )
+
+            tile_tokens = patchifier.patchify(tile_video_state.latent)
+            blend_output = helper.blend(tile_tokens, tile, ctx, blend_output)
+
+        return full_tools.unpatchify(replace(ref_initial, latent=blend_output)).latent
+
+    def _decode_video(
+        self,
+        latent: torch.Tensor,
+        tiling_config: TilingConfig | None,
+        generator: torch.Generator,
+        crop_size: tuple[int, int] | None = None,
+        *,
+        high_quality_hdr: bool = False,
+    ) -> torch.Tensor:
+        """Decode latent to HDR video, optionally cropping to target size.
+        Args:
+            crop_size: ``(width, height)`` to crop decoded frames to, or
+                ``None`` to skip cropping.
+            high_quality_hdr: When True, keep only every other frame (undoes the
+                2x generation applied during high-quality HDR mode).
+        Returns:
+            Linear HDR float tensor ``[f, h, w, c]``.
+        """
+        # Cast to float32 so tiled-decode accumulation buffers and blending
+        # masks run in full precision, avoiding bfloat16 seam artifacts.
+        # apply_hdr_decode_postprocess expects float32 [0, 1].
+        latent = latent.float()
+        decoded = torch.cat(
+            [chunk.float() for chunk in self.video_decoder(latent, tiling_config, generator)],
+            dim=0,
+        )
+        decoded = rearrange(decoded, "f h w c -> 1 c f h w")
+        hdr = apply_hdr_decode_postprocess(decoded, transform=self.hdr_transform)
+        del decoded
+        out = rearrange(hdr[0], "c f h w -> f h w c")
+        if crop_size is not None:
+            out = out[:, : crop_size[1], : crop_size[0], :]
+        if high_quality_hdr:
+            out = out[::2]
+        return out
+
+    def _create_conditionings(
+        self,
+        video_conditioning: list[tuple[str, float]],
+        height: int,
+        width: int,
+        num_frames: int,
+        video_encoder: VideoEncoder,
+        tiling_config: TilingConfig | None = None,
+        high_quality_hdr: bool = False,
+    ) -> list[ConditioningItem]:
+        """Create conditioning items for video generation."""
+        conditionings: list[ConditioningItem] = []
+
+        scale = self.reference_downscale_factor
+        if scale != 1 and (height % scale != 0 or width % scale != 0):
+            raise ValueError(
+                f"Output dimensions ({height}x{width}) must be divisible by reference_downscale_factor ({scale})"
+            )
+        ref_height = height // scale
+        ref_width = width // scale
+
+        # In high-quality HDR mode, load half the frames then duplicate each one.
+        load_frame_cap = (num_frames + 1) // 2 if high_quality_hdr else num_frames
+
+        for video_path, strength in video_conditioning:
+            video = torch.cat(
+                list(
+                    load_video_conditioning_hdr(
+                        video_path=video_path,
+                        height=ref_height,
+                        width=ref_width,
+                        frame_cap=load_frame_cap,
+                        dtype=self.dtype,
+                        device=self.device,
+                        hdr_transform=self.hdr_transform,
+                        resize_mode=ResizeMode.REFLECT_PAD,
+                    )
+                ),
+                dim=2,
+            )
+            if high_quality_hdr:
+                video = video.repeat_interleave(2, dim=2)[:, :, :num_frames, :, :]
+            if tiling_config is not None and ref_height * ref_width > self._tiled_vae_encode_threshold:
+                encoded_video = video_encoder.tiled_encode(video, tiling_config)
+            else:
+                encoded_video = video_encoder(video)
+
+            cond = VideoConditionByReferenceLatent(
+                latent=encoded_video,
+                downscale_factor=scale,
+                strength=strength,
+            )
+            conditionings.append(cond)
+
+        if video_conditioning:
+            logger.info("[HDR IC-LoRA] Added %d video conditioning(s)", len(video_conditioning))
+
+        return conditionings
+
+
+# ---------------------------------------------------------------------------
+# CLI helpers
+# ---------------------------------------------------------------------------
+
+
+def _make_tiling_config(
+    spatial_tile: int = DEFAULT_SPATIAL_TILE,
+    spatial_overlap: int = DEFAULT_SPATIAL_OVERLAP,
+    temporal_tile: int = DEFAULT_TEMPORAL_TILE,
+    temporal_overlap: int = DEFAULT_TEMPORAL_OVERLAP,
+) -> TilingConfig:
+    """Build a TilingConfig from explicit sizes.
+    The defaults (1280 px spatial tile, 256 px overlap; 32 temporal frames,
+    16 overlap) are suitable for H100-80 GB.  On GPUs with less VRAM,
+    reduce the spatial tile size (e.g. ``spatial_tile=768``).
+    """
+    from ltx_core.model.video_vae.tiling import SpatialTilingConfig, TemporalTilingConfig  # noqa: PLC0415
+
+    return TilingConfig(
+        spatial_config=SpatialTilingConfig(tile_size_in_pixels=spatial_tile, tile_overlap_in_pixels=spatial_overlap),
+        temporal_config=TemporalTilingConfig(
+            tile_size_in_frames=temporal_tile,
+            tile_overlap_in_frames=temporal_overlap,
+        ),
+    )
+
+
+_VIDEO_SUFFIXES = {".mp4", ".mov"}
+
+
+def _collect_videos(input_path: Path) -> list[Path]:
+    """Return a list of .mp4/.mov files from *input_path* (file or directory)."""
+    if input_path.is_file():
+        return [input_path]
+    if input_path.is_dir():
+        return sorted(p for p in input_path.iterdir() if p.is_file() and p.suffix.lower() in _VIDEO_SUFFIXES)
+    logger.error("Input %s is not a file or directory", input_path)
+    return []
+
+
+def _process_single_video(  # noqa: PLR0913
+    pipeline: HDRICLoraPipeline,
+    video_path: Path,
+    vid_w: int,
+    vid_h: int,
+    num_frames: int,
+    frame_rate: float,
+    output_dir: Path,
+    tiling_config: TilingConfig,
+    seed: int,
+    skip_mp4: bool,
+    exr_half: bool,
+    exr_executor: "ThreadPoolExecutor",  # noqa: F821
+    exr_futures: list,
+    high_quality_hdr: bool = False,
+) -> None:
+    """Run inference on a single video: generate EXR frames + optional H.264 .mp4 preview."""
+    import gc  # noqa: PLC0415
+    import time  # noqa: PLC0415
+
+    from ltx_pipelines.utils.media_io import encode_exr_sequence_to_mp4, save_exr_tensor  # noqa: PLC0415
+
+    output_mp4 = output_dir / f"{video_path.stem}.mp4"
+    exr_dir = output_dir / f"{video_path.stem}_exr"
+
+    t0 = time.time()
+    hdr_video = pipeline(
+        seed=seed,
+        height=vid_h,
+        width=vid_w,
+        num_frames=num_frames,
+        frame_rate=frame_rate,
+        video_conditioning=[(str(video_path), 1.0)],
+        tiling_config=tiling_config,
+        high_quality_hdr=high_quality_hdr,
+    )
+
+    exr_dir.mkdir(parents=True, exist_ok=True)
+    for j in range(hdr_video.shape[0]):
+        frame_cpu = hdr_video[j].cpu().clone()
+        path = exr_dir / f"frame_{j:05d}.exr"
+        exr_futures.append(exr_executor.submit(save_exr_tensor, frame_cpu, str(path), exr_half))
+
+    del hdr_video
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    if not skip_mp4:
+        # Wait for EXR saves to finish before encoding.
+        for fut in exr_futures:
+            fut.result()
+        logger.info("Encoding H.264 sRGB preview: %s", video_path.name)
+        encode_exr_sequence_to_mp4(exr_dir, output_mp4, frame_rate)
+
+    elapsed = time.time() - t0
+    logger.info("Decode + encode: %.1fs | %s", elapsed, output_mp4)
+
+
+# ---------------------------------------------------------------------------
+# CLI entry point
+# ---------------------------------------------------------------------------
+
+
+def _build_arg_parser() -> "argparse.ArgumentParser":  # noqa: F821
+    """Build the argument parser for HDR IC-LoRA batch inference."""
+    import argparse  # noqa: PLC0415
+
+    parser = argparse.ArgumentParser(
+        description="HDR IC-LoRA inference: EXR frames + tonemapped ProRes .mov.",
+        epilog="""\
+Resolution & frame constraints
+------------------------------
+  * Width and height must each be divisible by 32.
+  * Frame count must satisfy (frames - 1) %% 8 == 0.
+    Valid counts: 1, 9, 17, 25, ..., 121, 129, 137, 145, 153, 161.
+
+Max frames by resolution (fp8_cast, bfloat16 VAE, tiled decode)
+---------------------------------------------------------------
+  Resolution       80 GB (H100)    48 GB (A6000)
+  ------------------------------------------------
+   720p 1280x720    161+ frames     161+ frames
+  1080p 1920x1080   161+ frames     161+ frames
+  2K    2048x1080   161+ frames     161+ frames
+  1440p 2560x1440   161+ frames     137 frames
+  4K    3840x2160   121 frames       49 frames
+  4K    4096x2160   105 frames       49 frames
+
+  Estimates from ltx_pipelines.utils.vram_budget. Run
+    python -c "from ltx_pipelines.utils.vram_budget import \\
+      max_frames_for_resolution as mf; print(mf(W, H, vram_gb=GB))"
+  to check your specific resolution and GPU.
+
+  * The tiled-encode threshold (%(tiled_threshold)s px) and the default
+    tiling config (%(stile)s px spatial tile) are tuned for 80 GB.
+    On lower-VRAM GPUs pass --spatial-tile 768 (or smaller).
+"""
+        % {"tiled_threshold": TILED_VAE_ENCODE_PIXEL_THRESHOLD, "stile": DEFAULT_SPATIAL_TILE},
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    parser.add_argument("--input", required=True, help="Single .mp4 or directory of .mp4 videos.")
+    parser.add_argument("--output-dir", required=True, help="Directory for .mov and EXR folders.")
+    parser.add_argument("--hdr-lora", required=True, help="HDR IC-LoRA .safetensors file.")
+    parser.add_argument("--text-embeddings", required=True, help="Pre-computed text embeddings (.safetensors file).")
+    parser.add_argument("--distilled-checkpoint-path", required=True, help="Distilled model checkpoint (.safetensors).")
+    parser.add_argument("--spatial-upsampler-path", required=True, help="Spatial upsampler (.safetensors).")
+    parser.add_argument(
+        "--num-frames",
+        type=int,
+        default=DEFAULT_NUM_FRAMES,
+        help=f"Number of output frames.  Must satisfy (n-1) %% 8 == 0 (default: {DEFAULT_NUM_FRAMES}).",
+    )
+    parser.add_argument(
+        "--spatial-tile",
+        type=int,
+        default=DEFAULT_SPATIAL_TILE,
+        help=f"Spatial tile size in pixels for tiled VAE decode (default: {DEFAULT_SPATIAL_TILE}). "
+        "Reduce on lower-VRAM GPUs (e.g. 768 for 48 GB).",
+    )
+    parser.add_argument("--skip-mp4", action="store_true", help="Skip H.264 MP4 encoding, only produce EXR.")
+    parser.add_argument("--exr-half", action="store_true", help="Save EXR as float16.")
+    parser.add_argument("--seed", type=int, default=10, help="Random seed (default: 10).")
+    parser.add_argument(
+        "--offload",
+        dest="offload_mode",
+        type=OffloadMode,
+        default=OffloadMode.NONE,
+        choices=list(OffloadMode),
+        help=(
+            "Weight offloading strategy. "
+            "'none' keeps all weights on GPU (default). "
+            "'cpu' pins weights in CPU RAM, streams to GPU per layer. "
+            "'disk' reads weights from disk on demand (lowest memory). "
+            "Example: --offload cpu"
+        ),
+    )
+    parser.add_argument(
+        "--high-quality",
+        action="store_true",
+        help="High-quality HDR mode. Generates at 2x frame count internally "
+        "and keeps every other frame for smoother output. ~2x slower.",
+    )
+    return parser
+@torch.inference_mode()
+def main() -> None:
+    """Batch HDR IC-LoRA inference: per-frame EXR + tonemapped ProRes .mov."""
+    import time  # noqa: PLC0415
+    from concurrent.futures import ThreadPoolExecutor  # noqa: PLC0415
+
+    from ltx_pipelines.utils.media_io import get_videostream_metadata  # noqa: PLC0415
+
+    logging.basicConfig(level=logging.INFO)
+
+    args = _build_arg_parser().parse_args()
+    high_quality = args.high_quality
+    num_frames = args.num_frames
+
+    tiling_config = _make_tiling_config(spatial_tile=args.spatial_tile)
+
+    input_path = Path(args.input)
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    videos = _collect_videos(input_path)
+    if not videos:
+        logger.error("No valid videos to process.")
+        return
+    logger.info("Found %d video(s), generating %d frames each", len(videos), num_frames)
+
+    logger.info("Loading pipeline...")
+    pipeline = HDRICLoraPipeline(
+        distilled_checkpoint_path=args.distilled_checkpoint_path,
+        spatial_upsampler_path=args.spatial_upsampler_path,
+        hdr_lora=args.hdr_lora,
+        text_embeddings_path=args.text_embeddings,
+        offload_mode=args.offload_mode,
+    )
+    logger.info("Pipeline loaded.")
+
+    exr_executor = ThreadPoolExecutor(max_workers=4)
+    exr_futures: list = []
+
+    total_t0 = time.time()
+    successes = 0
+
+    for i, video_path in enumerate(videos, 1):
+        meta = get_videostream_metadata(str(video_path))
+        vid_w, vid_h = meta.width, meta.height
+        logger.info("%s", "=" * 60)
+        logger.info("[%d/%d] %s  (%dx%d, %df)", i, len(videos), video_path.name, vid_w, vid_h, num_frames)
+
+        _process_single_video(
+            pipeline=pipeline,
+            video_path=video_path,
+            vid_w=vid_w,
+            vid_h=vid_h,
+            num_frames=num_frames,
+            frame_rate=meta.fps,
+            output_dir=output_dir,
+            tiling_config=tiling_config,
+            seed=args.seed,
+            skip_mp4=args.skip_mp4,
+            exr_half=args.exr_half,
+            exr_executor=exr_executor,
+            exr_futures=exr_futures,
+            high_quality_hdr=high_quality,
+        )
+        successes += 1
+
+    infer_elapsed = time.time() - total_t0
+    logger.info("%s", "=" * 60)
+    logger.info("All inference done in %.0fs  (%d/%d OK)", infer_elapsed, successes, len(videos))
+
+    if exr_futures:
+        t0 = time.time()
+        logger.info("Waiting for %d EXR saves...", len(exr_futures))
+        for fut in exr_futures:
+            fut.result()
+        exr_wait = time.time() - t0
+        if exr_wait > 0.1:
+            logger.info("EXR save wait: %.1fs", exr_wait)
+
+    logger.info("Total wall time: %.0fs", time.time() - total_t0)
+
+
+if __name__ == "__main__":
+    main()

packages/ltx-pipelines/src/ltx_pipelines/ic_lora.py

@@ -2,20 +2,20 @@ import logging
 from collections.abc import Iterator
 
 import torch
-from einops import rearrange
-from safetensors import safe_open
 
 from ltx_core.components.noisers import GaussianNoiser
-from ltx_core.conditioning import (
-    ConditioningItem,
-    ConditioningItemAttentionStrengthWrapper,
-    VideoConditionByReferenceLatent,
-)
+from ltx_core.conditioning import ConditioningItem
 from ltx_core.loader import LoraPathStrengthAndSDOps
 from ltx_core.loader.registry import Registry
+from ltx_core.model.transformer.compiling import CompilationConfig
 from ltx_core.model.video_vae import TilingConfig, VideoEncoder, get_video_chunks_number
 from ltx_core.quantization import QuantizationPolicy
-from ltx_core.types import Audio, VideoLatentShape, VideoPixelShape
+from ltx_core.types import Audio, VideoPixelShape
+from ltx_pipelines.iclora_utils import (
+    append_ic_lora_reference_video_conditionings,
+    read_lora_reference_downscale_factor,
+    read_lora_reference_temporal_scale_factor,
+)
 from ltx_pipelines.utils.args import (
     ImageConditioningInput,
     VideoConditioningAction,
@@ -39,7 +39,7 @@ from ltx_pipelines.utils.constants import (
 from ltx_pipelines.utils.denoisers import SimpleDenoiser
 from ltx_pipelines.utils.helpers import assert_resolution, combined_image_conditionings, get_device
 from ltx_pipelines.utils.media_io import decode_video_by_frame, encode_video, video_preprocess
-from ltx_pipelines.utils.types import ModalitySpec
+from ltx_pipelines.utils.types import ModalitySpec, OffloadMode
 
 
 class ICLoraPipeline:
@@ -62,13 +62,19 @@ class ICLoraPipeline:
         device: torch.device | None = None,
         quantization: QuantizationPolicy | None = None,
         registry: Registry | None = None,
-        torch_compile: bool = False,
+        compilation_config: CompilationConfig | None = None,
+        offload_mode: OffloadMode = OffloadMode.NONE,
     ):
         self.device = device or get_device()
         self.dtype = torch.bfloat16
 
         self.prompt_encoder = PromptEncoder(
-            distilled_checkpoint_path, gemma_root, self.dtype, self.device, registry=registry
+            distilled_checkpoint_path,
+            gemma_root,
+            self.dtype,
+            self.device,
+            registry=registry,
+            offload_mode=offload_mode,
         )
         self.image_conditioner = ImageConditioner(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
         self.stage_1 = DiffusionStage(
@@ -78,7 +84,8 @@ class ICLoraPipeline:
             loras=tuple(loras),
             quantization=quantization,
             registry=registry,
-            torch_compile=torch_compile,
+            compilation_config=compilation_config,
+            offload_mode=offload_mode,
         )
         self.stage_2 = DiffusionStage(
             distilled_checkpoint_path,
@@ -87,7 +94,8 @@ class ICLoraPipeline:
             loras=(),
             quantization=quantization,
             registry=registry,
-            torch_compile=torch_compile,
+            compilation_config=compilation_config,
+            offload_mode=offload_mode,
         )
         self.upsampler = VideoUpsampler(
             distilled_checkpoint_path, spatial_upsampler_path, self.dtype, self.device, registry=registry
@@ -95,12 +103,13 @@ class ICLoraPipeline:
         self.video_decoder = VideoDecoder(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
         self.audio_decoder = AudioDecoder(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
 
-        # Read reference downscale factor from LoRA metadata.
-        # IC-LoRAs trained with low-resolution reference videos store this factor
-        # so inference can resize reference videos to match training conditions.
+        # Read reference scale factors from LoRA metadata.
+        # IC-LoRAs trained with scaled reference videos store these factors
+        # so inference can resize/subsample reference videos to match training conditions.
         self.reference_downscale_factor = 1
+        self.reference_temporal_scale_factor = 1
         for lora in loras:
-            scale = _read_lora_reference_downscale_factor(lora.path)
+            scale = read_lora_reference_downscale_factor(lora.path)
             if scale != 1:
                 if self.reference_downscale_factor not in (1, scale):
                     raise ValueError(
@@ -109,6 +118,15 @@ class ICLoraPipeline:
                         f"specifies {scale}. Cannot combine LoRAs with different reference scales."
                     )
                 self.reference_downscale_factor = scale
+            temporal = read_lora_reference_temporal_scale_factor(lora.path)
+            if temporal != 1:
+                if self.reference_temporal_scale_factor not in (1, temporal):
+                    raise ValueError(
+                        f"Conflicting reference_temporal_scale_factor values in LoRAs: "
+                        f"already have {self.reference_temporal_scale_factor}, but {lora.path} "
+                        f"specifies {temporal}. Cannot combine LoRAs with different temporal scales."
+                    )
+                self.reference_temporal_scale_factor = temporal
 
     def __call__(  # noqa: PLR0913
         self,
@@ -125,7 +143,6 @@ class ICLoraPipeline:
         conditioning_attention_strength: float = 1.0,
         skip_stage_2: bool = False,
         conditioning_attention_mask: torch.Tensor | None = None,
-        streaming_prefetch_count: int | None = None,
         stage_1_sigmas: torch.Tensor = DISTILLED_SIGMAS,
         stage_2_sigmas: torch.Tensor = STAGE_2_DISTILLED_SIGMAS,
     ) -> tuple[Iterator[torch.Tensor], Audio]:
@@ -175,7 +192,6 @@ class ICLoraPipeline:
             enhance_first_prompt=enhance_prompt,
             enhance_prompt_image=images[0][0] if len(images) > 0 else None,
             enhance_prompt_seed=seed,
-            streaming_prefetch_count=streaming_prefetch_count,
         )
         video_context, audio_context = ctx_p.video_encoding, ctx_p.audio_encoding
 
@@ -219,7 +235,6 @@ class ICLoraPipeline:
             audio=ModalitySpec(
                 context=audio_context,
             ),
-            streaming_prefetch_count=streaming_prefetch_count,
         )
 
         if skip_stage_2:
@@ -264,7 +279,6 @@ class ICLoraPipeline:
                 noise_scale=stage_2_sigmas[0].item(),
                 initial_latent=audio_state.latent,
             ),
-            streaming_prefetch_count=streaming_prefetch_count,
         )
 
         decoded_video = self.video_decoder(video_state.latent, tiling_config, generator)
@@ -305,108 +319,31 @@ class ICLoraPipeline:
             device=self.device,
         )
 
-        # Calculate scaled dimensions for reference video conditioning.
-        # IC-LoRAs trained with downscaled reference videos expect the same ratio at inference.
-        scale = self.reference_downscale_factor
-        if scale != 1 and (height % scale != 0 or width % scale != 0):
-            raise ValueError(
-                f"Output dimensions ({height}x{width}) must be divisible by reference_downscale_factor ({scale})"
-            )
-        ref_height = height // scale
-        ref_width = width // scale
-
-        for video_path, strength in video_conditioning:
-            # Load video at scaled-down resolution (if scale > 1)
-            frame_gen = decode_video_by_frame(path=video_path, frame_cap=num_frames, device=self.device)
-            video = video_preprocess(frame_gen, ref_height, ref_width, self.dtype, self.device)
-            encoded_video = video_encoder(video)
-            reference_video_shape = VideoLatentShape.from_torch_shape(encoded_video.shape)
-
-            # Build attention_mask for ConditioningItemAttentionStrengthWrapper
-            if conditioning_attention_mask is not None:
-                # Downsample pixel-space mask to latent space, then scale by strength
-                latent_mask = self._downsample_mask_to_latent(
-                    mask=conditioning_attention_mask,
-                    target_latent_shape=reference_video_shape,
-                )
-                attn_mask = latent_mask * conditioning_attention_strength
-            elif conditioning_attention_strength < 1.0:
-                # Use scalar strength only
-                attn_mask = conditioning_attention_strength
-            else:
-                attn_mask = None
-
-            cond = VideoConditionByReferenceLatent(
-                latent=encoded_video,
-                downscale_factor=scale,
-                strength=strength,
-            )
-            if attn_mask is not None:
-                cond = ConditioningItemAttentionStrengthWrapper(cond, attention_mask=attn_mask)
-            conditionings.append(cond)
+        append_ic_lora_reference_video_conditionings(
+            conditionings,
+            video_conditioning,
+            height=height,
+            width=width,
+            num_frames=num_frames,
+            video_encoder=video_encoder,
+            dtype=self.dtype,
+            device=self.device,
+            reference_downscale_factor=self.reference_downscale_factor,
+            reference_temporal_scale_factor=self.reference_temporal_scale_factor,
+            conditioning_attention_strength=conditioning_attention_strength,
+            conditioning_attention_mask=conditioning_attention_mask,
+            tiling_config=None,
+        )
 
         if video_conditioning:
-            logging.info(f"[IC-LoRA] Added {len(video_conditioning)} video conditioning(s)")
+            logging.info("[IC-LoRA] Added %d video conditioning(s)", len(video_conditioning))
 
         return conditionings
 
-    @staticmethod
-    def _downsample_mask_to_latent(
-        mask: torch.Tensor,
-        target_latent_shape: VideoLatentShape,
-    ) -> torch.Tensor:
-        """
-        Downsample a pixel-space mask to latent space using VAE scale factors.
-        Handles causal temporal downsampling: the first frame is kept separately
-        (temporal scale factor = 1 for the first frame), while the remaining
-        frames are downsampled by the VAE's temporal scale factor.
-        Args:
-            mask: Pixel-space mask of shape (B, 1, F_pixel, H_pixel, W_pixel).
-                Values in [0, 1].
-            target_latent_shape: Expected latent shape after VAE encoding.
-                Used to determine the target (F_latent, H_latent, W_latent).
-        Returns:
-            Flattened latent-space mask of shape (B, F_lat * H_lat * W_lat),
-            matching the patchifier's token ordering (f, h, w).
-        """
-        b = mask.shape[0]
-        f_lat = target_latent_shape.frames
-        h_lat = target_latent_shape.height
-        w_lat = target_latent_shape.width
-
-        # Step 1: Spatial downsampling (area interpolation per frame)
-        f_pix = mask.shape[2]
-        spatial_down = torch.nn.functional.interpolate(
-            rearrange(mask, "b 1 f h w -> (b f) 1 h w"),
-            size=(h_lat, w_lat),
-            mode="area",
-        )
-        spatial_down = rearrange(spatial_down, "(b f) 1 h w -> b 1 f h w", b=b)
-
-        # Step 2: Causal temporal downsampling
-        # First frame: kept as-is (causal VAE encodes first frame independently)
-        first_frame = spatial_down[:, :, :1, :, :]  # (B, 1, 1, H_lat, W_lat)
-
-        if f_pix > 1 and f_lat > 1:
-            # Remaining frames: downsample by temporal factor via group-mean
-            t = (f_pix - 1) // (f_lat - 1)  # temporal downscale factor
-            assert (f_pix - 1) % (f_lat - 1) == 0, (
-                f"Pixel frames ({f_pix}) not compatible with latent frames ({f_lat}): "
-                f"(f_pix - 1) must be divisible by (f_lat - 1)"
-            )
-            rest = rearrange(spatial_down[:, :, 1:, :, :], "b 1 (f t) h w -> b 1 f t h w", t=t)
-            rest = rest.mean(dim=3)  # (B, 1, F_lat-1, H_lat, W_lat)
-            latent_mask = torch.cat([first_frame, rest], dim=2)  # (B, 1, F_lat, H_lat, W_lat)
-        else:
-            latent_mask = first_frame
-
-        # Flatten to (B, F_lat * H_lat * W_lat) matching patchifier token order (f, h, w)
-        return rearrange(latent_mask, "b 1 f h w -> b (f h w)")
-
 
 @torch.inference_mode()
 def main() -> None:
-    logging.getLogger().setLevel(logging.INFO)
+    logging.basicConfig(level=logging.INFO)
     checkpoint_path = detect_checkpoint_path(distilled=True)
     params = detect_params(checkpoint_path)
     parser = default_2_stage_distilled_arg_parser(params=params)
@@ -462,7 +399,8 @@ def main() -> None:
         gemma_root=args.gemma_root,
         loras=tuple(args.lora) if args.lora else (),
         quantization=args.quantization,
-        torch_compile=args.compile,
+        compilation_config=args.compile,
+        offload_mode=args.offload_mode,
     )
     tiling_config = TilingConfig.default()
     video_chunks_number = get_video_chunks_number(args.num_frames, tiling_config)
@@ -479,7 +417,6 @@ def main() -> None:
         conditioning_attention_strength=conditioning_attention_strength,
         skip_stage_2=args.skip_stage_2,
         conditioning_attention_mask=conditioning_attention_mask,
-        streaming_prefetch_count=args.streaming_prefetch_count,
     )
 
     encode_video(
@@ -519,26 +456,5 @@ def _load_mask_video(
     return mask.clamp(0.0, 1.0)
 
 
-def _read_lora_reference_downscale_factor(lora_path: str) -> int:
-    """Read reference_downscale_factor from LoRA safetensors metadata.
-    Some IC-LoRA models are trained with reference videos at lower resolution than
-    the target output. This allows for more efficient training and can improve
-    generalization. The downscale factor indicates the ratio between target and
-    reference resolutions (e.g., factor=2 means reference is half the resolution).
-    Args:
-        lora_path: Path to the LoRA .safetensors file
-    Returns:
-        The reference downscale factor (1 if not specified in metadata, meaning
-        reference and target have the same resolution)
-    """
-    try:
-        with safe_open(lora_path, framework="pt") as f:
-            metadata = f.metadata() or {}
-            return int(metadata.get("reference_downscale_factor", 1))
-    except Exception as e:
-        logging.warning(f"Failed to read metadata from LoRA file '{lora_path}': {e}")
-        return 1
-
-
 if __name__ == "__main__":
     main()

packages/ltx-pipelines/src/ltx_pipelines/iclora_utils.py

@@ -0,0 +1,141 @@
+"""Shared IC-LoRA helpers: LoRA metadata, mask downsampling, reference-video conditioning.
+Used by ``ic_lora`` and ``lipdub`` (video reference path only). LipDub audio helpers live in ``lipdub.py``.
+"""
+
+from __future__ import annotations
+
+import logging
+
+import torch
+from einops import rearrange
+from safetensors import safe_open
+
+from ltx_core.conditioning import (
+    ConditioningItem,
+    ConditioningItemAttentionStrengthWrapper,
+    VideoConditionByReferenceLatent,
+)
+from ltx_core.model.video_vae import TilingConfig, VideoEncoder
+from ltx_core.types import VideoLatentShape
+from ltx_pipelines.utils.media_io import decode_video_by_frame, video_preprocess
+
+
+def read_lora_reference_downscale_factor(lora_path: str) -> int:
+    """Read ``reference_downscale_factor`` from LoRA safetensors metadata (default 1)."""
+    try:
+        with safe_open(lora_path, framework="pt") as f:
+            metadata = f.metadata() or {}
+            return int(metadata.get("reference_downscale_factor", 1))
+    except Exception as e:
+        logging.warning("Failed to read metadata from LoRA file '%s': %s", lora_path, e)
+        return 1
+
+
+def read_lora_reference_temporal_scale_factor(lora_path: str) -> int:
+    """Read ``reference_temporal_scale_factor`` from LoRA safetensors metadata (default 1)."""
+    try:
+        with safe_open(lora_path, framework="pt") as f:
+            metadata = f.metadata() or {}
+            return int(metadata.get("reference_temporal_scale_factor", 1))
+    except Exception as e:
+        logging.warning("Failed to read metadata from LoRA file '%s': %s", lora_path, e)
+        return 1
+
+
+def downsample_mask_video_to_latent(
+    mask: torch.Tensor,
+    target_latent_shape: VideoLatentShape,
+) -> torch.Tensor:
+    """Downsample a pixel-space mask video to flattened latent token weights."""
+    b = mask.shape[0]
+    f_lat = target_latent_shape.frames
+    h_lat = target_latent_shape.height
+    w_lat = target_latent_shape.width
+
+    f_pix = mask.shape[2]
+    spatial_down = torch.nn.functional.interpolate(
+        rearrange(mask, "b 1 f h w -> (b f) 1 h w"),
+        size=(h_lat, w_lat),
+        mode="area",
+    )
+    spatial_down = rearrange(spatial_down, "(b f) 1 h w -> b 1 f h w", b=b)
+
+    first_frame = spatial_down[:, :, :1, :, :]
+
+    if f_pix > 1 and f_lat > 1:
+        t = (f_pix - 1) // (f_lat - 1)
+        assert (f_pix - 1) % (f_lat - 1) == 0, (
+            f"Pixel frames ({f_pix}) not compatible with latent frames ({f_lat}): "
+            f"(f_pix - 1) must be divisible by (f_lat - 1)"
+        )
+        rest = rearrange(spatial_down[:, :, 1:, :, :], "b 1 (f t) h w -> b 1 f t h w", t=t)
+        rest = rest.mean(dim=3)
+        latent_mask = torch.cat([first_frame, rest], dim=2)
+    else:
+        latent_mask = first_frame
+
+    return rearrange(latent_mask, "b 1 f h w -> b (f h w)")
+
+
+def temporal_subsample(video: torch.Tensor, temporal_scale_factor: int) -> torch.Tensor:
+    """VAE-aligned temporal subsampling: keep frame 0, then every Nth frame."""
+    indices = [0, *list(range(1, video.shape[2], temporal_scale_factor))]
+    return video[:, :, indices]
+
+
+def append_ic_lora_reference_video_conditionings(  # noqa: PLR0913
+    conditionings: list[ConditioningItem],
+    video_conditioning: list[tuple[str, float]],
+    *,
+    height: int,
+    width: int,
+    num_frames: int,
+    video_encoder: VideoEncoder,
+    dtype: torch.dtype,
+    device: torch.device,
+    reference_downscale_factor: int,
+    reference_temporal_scale_factor: int = 1,
+    conditioning_attention_strength: float,
+    conditioning_attention_mask: torch.Tensor | None,
+    tiling_config: TilingConfig | None = None,
+) -> None:
+    """Append :class:`VideoConditionByReferenceLatent` items for each reference path."""
+    scale = reference_downscale_factor
+    if scale != 1 and (height % scale != 0 or width % scale != 0):
+        raise ValueError(
+            f"Output dimensions ({height}x{width}) must be divisible by reference_downscale_factor ({scale})"
+        )
+    ref_height = height // scale
+    ref_width = width // scale
+
+    for video_path, strength in video_conditioning:
+        frame_gen = decode_video_by_frame(path=video_path, frame_cap=num_frames, device=device)
+        video = video_preprocess(frame_gen, ref_height, ref_width, dtype, device)
+        if reference_temporal_scale_factor > 1:
+            video = temporal_subsample(video, reference_temporal_scale_factor)
+        if tiling_config is not None:
+            encoded_video = video_encoder.tiled_encode(video, tiling_config)
+        else:
+            encoded_video = video_encoder(video)
+        reference_video_shape = VideoLatentShape.from_torch_shape(encoded_video.shape)
+
+        if conditioning_attention_mask is not None:
+            latent_mask = downsample_mask_video_to_latent(
+                mask=conditioning_attention_mask,
+                target_latent_shape=reference_video_shape,
+            )
+            attn_mask = latent_mask * conditioning_attention_strength
+        elif conditioning_attention_strength < 1.0:
+            attn_mask = conditioning_attention_strength
+        else:
+            attn_mask = None
+
+        cond = VideoConditionByReferenceLatent(
+            latent=encoded_video,
+            downscale_factor=scale,
+            temporal_scale_factor=reference_temporal_scale_factor,
+            strength=strength,
+        )
+        if attn_mask is not None:
+            cond = ConditioningItemAttentionStrengthWrapper(cond, attention_mask=attn_mask)
+        conditionings.append(cond)

packages/ltx-pipelines/src/ltx_pipelines/keyframe_interpolation.py

@@ -12,10 +12,15 @@ from ltx_core.components.noisers import GaussianNoiser
 from ltx_core.components.schedulers import LTX2Scheduler
 from ltx_core.loader import LoraPathStrengthAndSDOps
 from ltx_core.loader.registry import Registry
+from ltx_core.model.transformer.compiling import CompilationConfig
 from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
 from ltx_core.quantization import QuantizationPolicy
 from ltx_core.types import Audio, VideoPixelShape
-from ltx_pipelines.utils.args import ImageConditioningInput, default_2_stage_arg_parser, detect_checkpoint_path
+from ltx_pipelines.utils.args import (
+    ImageConditioningInput,
+    default_2_stage_arg_parser,
+    detect_checkpoint_path,
+)
 from ltx_pipelines.utils.blocks import (
     AudioDecoder,
     DiffusionStage,
@@ -35,7 +40,7 @@ from ltx_pipelines.utils.helpers import (
     image_conditionings_by_adding_guiding_latent,
 )
 from ltx_pipelines.utils.media_io import encode_video
-from ltx_pipelines.utils.types import ModalitySpec
+from ltx_pipelines.utils.types import ModalitySpec, OffloadMode
 
 
 class KeyframeInterpolationPipeline:
@@ -58,13 +63,16 @@ class KeyframeInterpolationPipeline:
         device: torch.device | None = None,
         quantization: QuantizationPolicy | None = None,
         registry: Registry | None = None,
-        torch_compile: bool = False,
+        compilation_config: CompilationConfig | None = None,
+        offload_mode: OffloadMode = OffloadMode.NONE,
     ):
         self.device = device or get_device()
         self.dtype = torch.bfloat16
         self._scheduler = LTX2Scheduler()
 
-        self.prompt_encoder = PromptEncoder(checkpoint_path, gemma_root, self.dtype, self.device, registry=registry)
+        self.prompt_encoder = PromptEncoder(
+            checkpoint_path, gemma_root, self.dtype, self.device, registry=registry, offload_mode=offload_mode
+        )
         self.image_conditioner = ImageConditioner(checkpoint_path, self.dtype, self.device, registry=registry)
         self.stage_1 = DiffusionStage(
             checkpoint_path,
@@ -73,7 +81,8 @@ class KeyframeInterpolationPipeline:
             loras=tuple(loras),
             quantization=quantization,
             registry=registry,
-            torch_compile=torch_compile,
+            compilation_config=compilation_config,
+            offload_mode=offload_mode,
         )
         stage_2_loras = (*tuple(loras), *tuple(distilled_lora))
         self.stage_2 = DiffusionStage(
@@ -83,7 +92,8 @@ class KeyframeInterpolationPipeline:
             loras=stage_2_loras,
             quantization=quantization,
             registry=registry,
-            torch_compile=torch_compile,
+            compilation_config=compilation_config,
+            offload_mode=offload_mode,
         )
         self.upsampler = VideoUpsampler(
             checkpoint_path, spatial_upsampler_path, self.dtype, self.device, registry=registry
@@ -106,7 +116,6 @@ class KeyframeInterpolationPipeline:
         images: list[ImageConditioningInput],
         tiling_config: TilingConfig | None = None,
         enhance_prompt: bool = False,
-        streaming_prefetch_count: int | None = None,
         max_batch_size: int = 1,
         stage_1_sigmas: torch.Tensor | None = None,
         stage_2_sigmas: torch.Tensor = STAGE_2_DISTILLED_SIGMAS,
@@ -122,7 +131,6 @@ class KeyframeInterpolationPipeline:
             enhance_first_prompt=enhance_prompt,
             enhance_prompt_image=images[0][0] if len(images) > 0 else None,
             enhance_prompt_seed=seed,
-            streaming_prefetch_count=streaming_prefetch_count,
         )
         v_context_p, a_context_p = ctx_p.video_encoding, ctx_p.audio_encoding
         v_context_n, a_context_n = ctx_n.video_encoding, ctx_n.audio_encoding
@@ -179,7 +187,6 @@ class KeyframeInterpolationPipeline:
             audio=ModalitySpec(
                 context=a_context_p,
             ),
-            streaming_prefetch_count=streaming_prefetch_count,
             max_batch_size=max_batch_size,
         )
 
@@ -218,7 +225,6 @@ class KeyframeInterpolationPipeline:
                 noise_scale=stage_2_sigmas[0].item(),
                 initial_latent=audio_state.latent,
             ),
-            streaming_prefetch_count=streaming_prefetch_count,
         )
 
         decoded_video = self.video_decoder(video_state.latent, tiling_config, generator)
@@ -228,7 +234,7 @@ class KeyframeInterpolationPipeline:
 
 @torch.inference_mode()
 def main() -> None:
-    logging.getLogger().setLevel(logging.INFO)
+    logging.basicConfig(level=logging.INFO)
     checkpoint_path = detect_checkpoint_path()
     params = detect_params(checkpoint_path)
     parser = default_2_stage_arg_parser(params=params)
@@ -240,7 +246,8 @@ def main() -> None:
         gemma_root=args.gemma_root,
         loras=tuple(args.lora) if args.lora else (),
         quantization=args.quantization,
-        torch_compile=args.compile,
+        compilation_config=args.compile,
+        offload_mode=args.offload_mode,
     )
     tiling_config = TilingConfig.default()
     video_chunks_number = get_video_chunks_number(args.num_frames, tiling_config)
@@ -271,7 +278,6 @@ def main() -> None:
         ),
         images=args.images,
         tiling_config=tiling_config,
-        streaming_prefetch_count=args.streaming_prefetch_count,
         max_batch_size=args.max_batch_size,
     )
 

packages/ltx-pipelines/src/ltx_pipelines/lipdub.py

@@ -0,0 +1,335 @@
+"""Two-stage lip-dubbing pipeline with IC-LoRA and appended audio reference conditioning."""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Iterator
+
+import torch
+
+from ltx_core.components.noisers import GaussianNoiser
+from ltx_core.components.patchifiers import AudioPatchifier
+from ltx_core.conditioning import AudioConditionByReferenceLatent
+from ltx_core.loader import LoraPathStrengthAndSDOps
+from ltx_core.loader.registry import Registry
+from ltx_core.model.audio_vae import encode_audio as vae_encode_audio
+from ltx_core.model.transformer.compiling import CompilationConfig
+from ltx_core.model.video_vae import TilingConfig, VideoEncoder, get_video_chunks_number
+from ltx_core.quantization import QuantizationPolicy
+from ltx_core.types import Audio, AudioLatentShape, SpatioTemporalScaleFactors, VideoPixelShape
+from ltx_pipelines.iclora_utils import (
+    append_ic_lora_reference_video_conditionings,
+    read_lora_reference_downscale_factor,
+)
+from ltx_pipelines.utils.args import (
+    ImageConditioningInput,
+    detect_checkpoint_path,
+    lipdub_arg_parser,
+)
+from ltx_pipelines.utils.blocks import (
+    AudioConditioner,
+    AudioDecoder,
+    DiffusionStage,
+    ImageConditioner,
+    PromptEncoder,
+    VideoDecoder,
+    VideoUpsampler,
+)
+from ltx_pipelines.utils.constants import DISTILLED_SIGMAS, STAGE_2_DISTILLED_SIGMAS, detect_params
+from ltx_pipelines.utils.denoisers import SimpleDenoiser
+from ltx_pipelines.utils.helpers import assert_resolution, combined_image_conditionings, get_device
+from ltx_pipelines.utils.media_io import decode_audio_from_file, encode_video, get_videostream_metadata
+from ltx_pipelines.utils.types import ModalitySpec, OffloadMode
+
+
+def _snap_frames_to_8k1(frames: int) -> int:
+    """Round ``frames`` down to the nearest ``8k+1`` (the model's required frame count)."""
+    time_scale = SpatioTemporalScaleFactors.default().time
+    return ((frames - 1) // time_scale) * time_scale + 1
+
+
+class LipDubPipeline:
+    """Two-stage lip-dubbing with IC-LoRA video reference and appended audio reference tokens."""
+
+    def __init__(
+        self,
+        distilled_checkpoint_path: str,
+        spatial_upsampler_path: str,
+        gemma_root: str,
+        ic_lora: LoraPathStrengthAndSDOps,
+        device: torch.device | None = None,
+        quantization: QuantizationPolicy | None = None,
+        registry: Registry | None = None,
+        compilation_config: CompilationConfig | None = None,
+        offload_mode: OffloadMode = OffloadMode.NONE,
+    ) -> None:
+        self.device = device or get_device()
+        self.dtype = torch.bfloat16
+        self.ic_lora = ic_lora
+        loras = (ic_lora,)
+
+        self.prompt_encoder = PromptEncoder(
+            distilled_checkpoint_path,
+            gemma_root,
+            self.dtype,
+            self.device,
+            registry=registry,
+            offload_mode=offload_mode,
+        )
+        self.image_conditioner = ImageConditioner(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
+        self.audio_conditioner = AudioConditioner(
+            distilled_checkpoint_path,
+            self.dtype,
+            self.device,
+            registry=registry,
+        )
+        self.stage = DiffusionStage(
+            distilled_checkpoint_path,
+            self.dtype,
+            self.device,
+            loras=loras,
+            quantization=quantization,
+            registry=registry,
+            compilation_config=compilation_config,
+            offload_mode=offload_mode,
+        )
+        self.upsampler = VideoUpsampler(
+            distilled_checkpoint_path, spatial_upsampler_path, self.dtype, self.device, registry=registry
+        )
+        self.video_decoder = VideoDecoder(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
+        self.audio_decoder = AudioDecoder(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
+        self.reference_downscale_factor = read_lora_reference_downscale_factor(ic_lora.path)
+
+    def _create_stage_conditionings(
+        self,
+        images: list[ImageConditioningInput],
+        reference_video_path: str,
+        reference_strength: float,
+        height: int,
+        width: int,
+        num_frames: int,
+        video_encoder: VideoEncoder,
+        encode_tiling: TilingConfig | None,
+    ) -> list:
+        conditionings = combined_image_conditionings(
+            images=images,
+            height=height,
+            width=width,
+            video_encoder=video_encoder,
+            dtype=self.dtype,
+            device=self.device,
+        )
+        append_ic_lora_reference_video_conditionings(
+            conditionings,
+            [(reference_video_path, reference_strength)],
+            height=height,
+            width=width,
+            num_frames=num_frames,
+            video_encoder=video_encoder,
+            dtype=self.dtype,
+            device=self.device,
+            reference_downscale_factor=self.reference_downscale_factor,
+            conditioning_attention_strength=1.0,
+            conditioning_attention_mask=None,
+            tiling_config=encode_tiling,
+        )
+        return conditionings
+
+    def _encode_reference_audio_vae_latent(self, video_path: str) -> torch.Tensor:
+        audio = decode_audio_from_file(video_path, self.device)
+        if audio is None:
+            msg = f"No audio stream found in {video_path}"
+            raise ValueError(msg)
+        return self.audio_conditioner(lambda enc: vae_encode_audio(audio, enc, None))
+
+    @torch.inference_mode()
+    def __call__(  # noqa: PLR0913
+        self,
+        prompt: str,
+        seed: int,
+        height: int,
+        width: int,
+        images: list[ImageConditioningInput],
+        reference_video_path: str,
+        reference_strength: float = 1.0,
+        enhance_prompt: bool = False,
+        tiling_config: TilingConfig | None = None,
+        stage_1_sigmas: torch.Tensor = DISTILLED_SIGMAS,
+        stage_2_sigmas: torch.Tensor = STAGE_2_DISTILLED_SIGMAS,
+    ) -> tuple[Iterator[torch.Tensor], Audio]:
+        assert_resolution(height=height, width=width, is_two_stage=True)
+
+        meta = get_videostream_metadata(reference_video_path)
+        num_frames = _snap_frames_to_8k1(meta.frames)
+        frame_rate = float(meta.fps)
+
+        generator = torch.Generator(device=self.device).manual_seed(seed)
+        noiser = GaussianNoiser(generator=generator)
+
+        (ctx_p,) = self.prompt_encoder(
+            [prompt],
+            enhance_first_prompt=enhance_prompt,
+            enhance_prompt_image=images[0][0] if len(images) > 0 else None,
+            enhance_prompt_seed=seed,
+        )
+        video_context, audio_context = ctx_p.video_encoding, ctx_p.audio_encoding
+
+        stage_1_output_shape = VideoPixelShape(
+            batch=1,
+            frames=num_frames,
+            width=width // 2,
+            height=height // 2,
+            fps=frame_rate,
+        )
+        encode_tiling = TilingConfig.default()
+
+        def build_image_conditionings(output_shape: VideoPixelShape) -> list:
+            return self.image_conditioner(
+                lambda enc: self._create_stage_conditionings(
+                    images=images,
+                    reference_video_path=reference_video_path,
+                    reference_strength=reference_strength,
+                    height=output_shape.height,
+                    width=output_shape.width,
+                    num_frames=num_frames,
+                    video_encoder=enc,
+                    encode_tiling=encode_tiling,
+                )
+            )
+
+        def build_audio_ref_conditioning(audio_latent: torch.Tensor) -> AudioConditionByReferenceLatent:
+            ref_patch, ref_pos = patchify_lipdub_audio_reference_latent(
+                audio_latent,
+                negative_positions=True,
+                device=self.device,
+            )
+            return AudioConditionByReferenceLatent(ref_patch, ref_pos, strength=1.0)
+
+        stage_1_conditionings = build_image_conditionings(stage_1_output_shape)
+
+        ref_vae = self._encode_reference_audio_vae_latent(reference_video_path)
+        audio_conditionings = [build_audio_ref_conditioning(ref_vae)]
+
+        stage_1_sigmas_tensor = stage_1_sigmas.to(dtype=torch.float32, device=self.device)
+        video_state, audio_state = self.stage(
+            denoiser=SimpleDenoiser(video_context, audio_context),
+            sigmas=stage_1_sigmas_tensor,
+            noiser=noiser,
+            width=stage_1_output_shape.width,
+            height=stage_1_output_shape.height,
+            frames=num_frames,
+            fps=frame_rate,
+            video=ModalitySpec(
+                context=video_context,
+                conditionings=stage_1_conditionings,
+            ),
+            audio=ModalitySpec(
+                context=audio_context,
+                conditionings=audio_conditionings,
+            ),
+        )
+
+        s1_audio_latent = audio_state.latent.clone()
+
+        upscaled_video_latent = self.upsampler(video_state.latent[:1])
+        stage_2_sigmas_tensor = stage_2_sigmas.to(dtype=torch.float32, device=self.device)
+        stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
+        stage_2_conditionings = build_image_conditionings(stage_2_output_shape)
+
+        stage_2_audio_conditionings = [build_audio_ref_conditioning(s1_audio_latent)]
+
+        video_state, _audio_unused = self.stage(
+            denoiser=SimpleDenoiser(video_context, audio_context),
+            sigmas=stage_2_sigmas_tensor,
+            noiser=noiser,
+            width=width,
+            height=height,
+            frames=num_frames,
+            fps=frame_rate,
+            video=ModalitySpec(
+                context=video_context,
+                conditionings=stage_2_conditionings,
+                noise_scale=stage_2_sigmas_tensor[0].item(),
+                initial_latent=upscaled_video_latent,
+            ),
+            audio=ModalitySpec(
+                context=audio_context,
+                conditionings=stage_2_audio_conditionings,
+                frozen=True,
+                noise_scale=0.0,
+                initial_latent=s1_audio_latent,
+            ),
+        )
+
+        decoded_video = self.video_decoder(video_state.latent, tiling_config, generator)
+        decoded_audio = self.audio_decoder(s1_audio_latent)
+        return decoded_video, decoded_audio
+
+
+def patchify_lipdub_audio_reference_latent(
+    vae_latents: torch.Tensor,
+    *,
+    negative_positions: bool,
+    device: torch.device,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Patchify audio VAE latents and build RoPE positions (optional negative shift for reference)."""
+    patchifier = AudioPatchifier(patch_size=1)
+    patchified = patchifier.patchify(vae_latents)
+    b, c, _t, mel_bins = vae_latents.shape
+    seq_len = patchified.shape[1]
+    latent_coords = patchifier.get_patch_grid_bounds(
+        output_shape=AudioLatentShape(batch=b, channels=c, frames=seq_len, mel_bins=mel_bins),
+        device=device,
+    )
+    positions = latent_coords.to(dtype=torch.float32)
+    if negative_positions:
+        aud_dur = positions[:, :, -1, 1].max().item()
+        positions = positions - aud_dur - 0.04
+    return patchified, positions
+
+
+@torch.inference_mode()
+def main() -> None:
+    logging.basicConfig(level=logging.INFO)
+    checkpoint_path = detect_checkpoint_path(distilled=True)
+    params = detect_params(checkpoint_path)
+    parser = lipdub_arg_parser(params=params)
+    args = parser.parse_args()
+
+    if not args.lora or len(args.lora) != 1:
+        raise ValueError("LipDub requires exactly one --lora (the lip-dub IC-LoRA).")
+
+    pipeline = LipDubPipeline(
+        distilled_checkpoint_path=args.distilled_checkpoint_path,
+        spatial_upsampler_path=args.spatial_upsampler_path,
+        gemma_root=args.gemma_root,
+        ic_lora=args.lora[0],
+        quantization=args.quantization,
+        compilation_config=args.compile,
+        offload_mode=args.offload_mode,
+    )
+    tiling_config = TilingConfig.default()
+    src = get_videostream_metadata(args.reference_video)
+    video_chunks_number = get_video_chunks_number(_snap_frames_to_8k1(src.frames), tiling_config)
+    video, audio = pipeline(
+        prompt=args.prompt,
+        seed=args.seed,
+        height=args.height,
+        width=args.width,
+        images=[],
+        reference_video_path=args.reference_video,
+        reference_strength=args.reference_strength,
+        tiling_config=tiling_config,
+        enhance_prompt=args.enhance_prompt,
+    )
+    encode_video(
+        video=video,
+        fps=int(src.fps),
+        audio=audio,
+        output_path=args.output_path,
+        video_chunks_number=video_chunks_number,
+    )
+
+
+if __name__ == "__main__":
+    main()

packages/ltx-pipelines/src/ltx_pipelines/retake.py

@@ -11,6 +11,7 @@ from ltx_core.components.schedulers import LTX2Scheduler
 from ltx_core.conditioning.types.noise_mask_cond import TemporalRegionMask
 from ltx_core.loader import LoraPathStrengthAndSDOps
 from ltx_core.loader.registry import Registry
+from ltx_core.model.transformer.compiling import CompilationConfig
 from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
 from ltx_core.quantization import QuantizationPolicy
 from ltx_core.types import (
@@ -36,7 +37,7 @@ from ltx_pipelines.utils.media_io import (
     encode_video,
     get_videostream_metadata,
 )
-from ltx_pipelines.utils.types import ModalitySpec
+from ltx_pipelines.utils.types import ModalitySpec, OffloadMode
 
 
 class RetakePipeline:
@@ -73,7 +74,8 @@ class RetakePipeline:
         quantization: QuantizationPolicy | None = None,
         registry: Registry | None = None,
         distilled: bool = True,
-        torch_compile: bool = False,
+        compilation_config: CompilationConfig | None = None,
+        offload_mode: OffloadMode = OffloadMode.NONE,
     ):
         self.device = device or get_device()
         self.dtype = torch.bfloat16
@@ -86,6 +88,7 @@ class RetakePipeline:
             dtype=self.dtype,
             device=self.device,
             registry=registry,
+            offload_mode=offload_mode,
         )
         self.image_conditioner = ImageConditioner(
             checkpoint_path=checkpoint_path,
@@ -106,7 +109,8 @@ class RetakePipeline:
             loras=tuple(loras),
             quantization=quantization,
             registry=registry,
-            torch_compile=torch_compile,
+            compilation_config=compilation_config,
+            offload_mode=offload_mode,
         )
         self.video_decoder = VideoDecoder(
             checkpoint_path=checkpoint_path,
@@ -141,7 +145,6 @@ class RetakePipeline:
         regenerate_audio: bool = True,
         enhance_prompt: bool = False,
         tiling_config: TilingConfig | None = None,
-        streaming_prefetch_count: int | None = None,
         max_batch_size: int = 1,
         sigmas: torch.Tensor | None = None,
     ) -> tuple[Iterator[torch.Tensor], torch.Tensor]:
@@ -210,7 +213,6 @@ class RetakePipeline:
             prompts_to_encode,
             enhance_first_prompt=enhance_prompt,
             enhance_prompt_seed=seed,
-            streaming_prefetch_count=streaming_prefetch_count,
         )
 
         v_context_p, a_context_p = contexts[0].video_encoding, contexts[0].audio_encoding
@@ -269,7 +271,6 @@ class RetakePipeline:
             fps=output_shape.fps,
             video=video_modality_spec,
             audio=audio_modality_spec,
-            streaming_prefetch_count=streaming_prefetch_count,
             max_batch_size=max_batch_size,
         )
 
@@ -283,7 +284,7 @@ class RetakePipeline:
 @torch.inference_mode()
 def main() -> None:
     """CLI entry point for retake (regenerate a time region)."""
-    logging.getLogger().setLevel(logging.INFO)
+    logging.basicConfig(level=logging.INFO)
     parser = video_editing_arg_parser(distilled=True)
     parser.description = "Retake: regenerate a time region of a video with LTX-2."
     args = parser.parse_args()
@@ -307,8 +308,9 @@ def main() -> None:
         gemma_root=args.gemma_root,
         loras=tuple(args.lora) if args.lora else (),
         quantization=args.quantization,
-        distilled=args.distilled,
-        torch_compile=args.compile,
+        distilled=True,
+        compilation_config=args.compile,
+        offload_mode=args.offload_mode,
     )
     params = detect_params(args.distilled_checkpoint_path)
     tiling_config = TilingConfig.default()
@@ -321,7 +323,6 @@ def main() -> None:
         video_guider_params=params.video_guider_params,
         audio_guider_params=params.audio_guider_params,
         tiling_config=tiling_config,
-        streaming_prefetch_count=args.streaming_prefetch_count,
         max_batch_size=args.max_batch_size,
     )
     video_chunks_number = get_video_chunks_number(src.frames, tiling_config)

packages/ltx-pipelines/src/ltx_pipelines/t2a_one_stage.py

@@ -0,0 +1,192 @@
+import logging
+
+import torch
+
+from ltx_core.components.guiders import (
+    MultiModalGuiderFactory,
+    MultiModalGuiderParams,
+    create_multimodal_guider_factory,
+)
+from ltx_core.components.noisers import GaussianNoiser
+from ltx_core.components.schedulers import LTX2Scheduler
+from ltx_core.loader import LoraPathStrengthAndSDOps
+from ltx_core.loader.registry import Registry
+from ltx_core.model.transformer import LTXV_AUDIO_ONLY_MODEL_COMFY_RENAMING_MAP, LTXAudioOnlyModelConfigurator
+from ltx_core.model.transformer.compiling import CompilationConfig
+from ltx_core.quantization import QuantizationPolicy
+from ltx_core.types import Audio
+from ltx_pipelines.utils import get_device
+from ltx_pipelines.utils.args import (
+    default_1_stage_t2a_arg_parser,
+    detect_checkpoint_path,
+)
+from ltx_pipelines.utils.blocks import (
+    AudioDecoder,
+    DiffusionStage,
+    PromptEncoder,
+)
+from ltx_pipelines.utils.constants import detect_params
+from ltx_pipelines.utils.denoisers import FactoryGuidedDenoiser
+from ltx_pipelines.utils.media_io import encode_audio
+from ltx_pipelines.utils.types import ModalitySpec, OffloadMode
+
+# Placeholder pixel dimensions used for ``VideoPixelShape`` construction.
+# Audio-only generation reads ``frames`` and ``fps`` from the pixel shape via
+# ``AudioLatentShape.from_video_pixel_shape`` (height/width are unused).
+_AUDIO_ONLY_PLACEHOLDER_RES = 512
+
+
+class T2AOneStagePipeline:
+    """
+    Single-stage text-to-audio generation pipeline.
+    Generates audio at the target duration in a single diffusion pass with
+    classifier-free guidance (CFG) on the audio modality only. The video
+    modality is fully absent — the transformer runs audio-only by passing
+    ``video=None`` to the ``DiffusionStage``.
+    Assumes full non distilled model is provided in the checkpoint_path.
+    """
+
+    def __init__(
+        self,
+        checkpoint_path: str,
+        gemma_root: str,
+        loras: list[LoraPathStrengthAndSDOps],
+        device: torch.device | None = None,
+        quantization: QuantizationPolicy | None = None,
+        registry: Registry | None = None,
+        compilation_config: CompilationConfig | None = None,
+        offload_mode: OffloadMode = OffloadMode.NONE,
+    ):
+        self.dtype = torch.bfloat16
+        self.device = device or get_device()
+        self._scheduler = LTX2Scheduler()
+        self.prompt_encoder = PromptEncoder(
+            checkpoint_path=checkpoint_path,
+            gemma_root=gemma_root,
+            dtype=self.dtype,
+            device=self.device,
+            registry=registry,
+            offload_mode=offload_mode,
+        )
+        # Audio-only: build an audio-only transformer (model_configurator) so the video
+        # weights are never instantiated, plus a use-case-specific SDOps that restricts
+        # checkpoint reads to the audio model's keys, so the video weights are never even
+        # read from disk (the loader skips any key the SDOps maps to None).
+        self.stage = DiffusionStage(
+            checkpoint_path=checkpoint_path,
+            dtype=self.dtype,
+            device=self.device,
+            loras=tuple(loras),
+            quantization=quantization,
+            registry=registry,
+            compilation_config=compilation_config,
+            offload_mode=offload_mode,
+            model_configurator=LTXAudioOnlyModelConfigurator,
+            model_sd_ops=LTXV_AUDIO_ONLY_MODEL_COMFY_RENAMING_MAP,
+        )
+        self.audio_decoder = AudioDecoder(
+            checkpoint_path=checkpoint_path,
+            dtype=self.dtype,
+            device=self.device,
+            registry=registry,
+        )
+
+    def __call__(
+        self,
+        prompt: str,
+        negative_prompt: str,
+        seed: int,
+        num_frames: int,
+        frame_rate: float,
+        num_inference_steps: int,
+        audio_guider_params: MultiModalGuiderParams | MultiModalGuiderFactory,
+        enhance_prompt: bool = False,
+        max_batch_size: int = 1,
+        sigmas: torch.Tensor | None = None,
+    ) -> Audio:
+        generator = torch.Generator(device=self.device).manual_seed(seed)
+        noiser = GaussianNoiser(generator=generator)
+
+        ctx_p, ctx_n = self.prompt_encoder(
+            [prompt, negative_prompt],
+            enhance_first_prompt=enhance_prompt,
+            enhance_prompt_image=None,
+            enhance_prompt_seed=seed,
+        )
+        a_context_p = ctx_p.audio_encoding
+        a_context_n = ctx_n.audio_encoding
+
+        sigmas = (sigmas if sigmas is not None else self._scheduler.execute(steps=num_inference_steps)).to(
+            dtype=torch.float32, device=self.device
+        )
+
+        # Normalize to a guider factory. Plain ``MultiModalGuiderParams`` (the default /
+        # CLI case) becomes a simple sigma-independent guider, but callers may also pass
+        # their own factory for sigma-dependent guidance; ``FactoryGuidedDenoiser`` always
+        # consumes a factory.
+        audio_guider_factory = create_multimodal_guider_factory(
+            params=audio_guider_params,
+            negative_context=a_context_n,
+        )
+
+        _, audio_state = self.stage(
+            denoiser=FactoryGuidedDenoiser(
+                v_context=None,
+                a_context=a_context_p,
+                video_guider_factory=None,
+                audio_guider_factory=audio_guider_factory,
+            ),
+            sigmas=sigmas,
+            noiser=noiser,
+            width=_AUDIO_ONLY_PLACEHOLDER_RES,
+            height=_AUDIO_ONLY_PLACEHOLDER_RES,
+            frames=num_frames,
+            fps=frame_rate,
+            video=None,
+            audio=ModalitySpec(context=a_context_p),
+            max_batch_size=max_batch_size,
+        )
+
+        return self.audio_decoder(audio_state.latent)
+
+
+@torch.inference_mode()
+def main() -> None:
+    logging.basicConfig(level=logging.INFO)
+    checkpoint_path = detect_checkpoint_path()
+    params = detect_params(checkpoint_path)
+    parser = default_1_stage_t2a_arg_parser(params=params)
+    args = parser.parse_args()
+    pipeline = T2AOneStagePipeline(
+        checkpoint_path=args.checkpoint_path,
+        gemma_root=args.gemma_root,
+        loras=tuple(args.lora) if args.lora else (),
+        quantization=args.quantization,
+        compilation_config=args.compile,
+        offload_mode=args.offload_mode,
+    )
+    audio = pipeline(
+        prompt=args.prompt,
+        negative_prompt=args.negative_prompt,
+        seed=args.seed,
+        num_frames=args.num_frames,
+        frame_rate=args.frame_rate,
+        num_inference_steps=args.num_inference_steps,
+        audio_guider_params=MultiModalGuiderParams(
+            cfg_scale=args.audio_cfg_guidance_scale,
+            stg_scale=args.audio_stg_guidance_scale,
+            rescale_scale=args.audio_rescale_scale,
+            # Audio-only generation has no video modality, so the video->audio
+            # (v2a) cross-modal guidance is meaningless here. 1.0 disables it.
+            modality_scale=1.0,
+            skip_step=args.audio_skip_step,
+            stg_blocks=args.audio_stg_blocks,
+        ),
+        max_batch_size=args.max_batch_size,
+    )
+
+    encode_audio(audio=audio, output_path=args.output_path)
+
+
+if __name__ == "__main__":
+    main()

packages/ltx-pipelines/src/ltx_pipelines/ti2vid_one_stage.py

@@ -12,6 +12,7 @@ from ltx_core.components.noisers import GaussianNoiser
 from ltx_core.components.schedulers import LTX2Scheduler
 from ltx_core.loader import LoraPathStrengthAndSDOps
 from ltx_core.loader.registry import Registry
+from ltx_core.model.transformer.compiling import CompilationConfig
 from ltx_core.model.video_vae.tiling import TilingConfig
 from ltx_core.quantization import QuantizationPolicy
 from ltx_core.types import Audio
@@ -20,7 +21,11 @@ from ltx_pipelines.utils import (
     combined_image_conditionings,
     get_device,
 )
-from ltx_pipelines.utils.args import ImageConditioningInput, default_1_stage_arg_parser, detect_checkpoint_path
+from ltx_pipelines.utils.args import (
+    ImageConditioningInput,
+    default_1_stage_arg_parser,
+    detect_checkpoint_path,
+)
 from ltx_pipelines.utils.blocks import (
     AudioDecoder,
     DiffusionStage,
@@ -31,7 +36,7 @@ from ltx_pipelines.utils.blocks import (
 from ltx_pipelines.utils.constants import detect_params
 from ltx_pipelines.utils.denoisers import FactoryGuidedDenoiser
 from ltx_pipelines.utils.media_io import encode_video
-from ltx_pipelines.utils.types import ModalitySpec
+from ltx_pipelines.utils.types import ModalitySpec, OffloadMode
 
 
 class TI2VidOneStagePipeline:
@@ -51,7 +56,8 @@ class TI2VidOneStagePipeline:
         device: torch.device | None = None,
         quantization: QuantizationPolicy | None = None,
         registry: Registry | None = None,
-        torch_compile: bool = False,
+        compilation_config: CompilationConfig | None = None,
+        offload_mode: OffloadMode = OffloadMode.NONE,
     ):
         self.dtype = torch.bfloat16
         self.device = device or get_device()
@@ -62,6 +68,7 @@ class TI2VidOneStagePipeline:
             dtype=self.dtype,
             device=self.device,
             registry=registry,
+            offload_mode=offload_mode,
         )
         self.image_conditioner = ImageConditioner(
             checkpoint_path=checkpoint_path,
@@ -76,7 +83,8 @@ class TI2VidOneStagePipeline:
             loras=tuple(loras),
             quantization=quantization,
             registry=registry,
-            torch_compile=torch_compile,
+            compilation_config=compilation_config,
+            offload_mode=offload_mode,
         )
         self.video_decoder = VideoDecoder(
             checkpoint_path=checkpoint_path,
@@ -105,7 +113,6 @@ class TI2VidOneStagePipeline:
         audio_guider_params: MultiModalGuiderParams | MultiModalGuiderFactory,
         images: list[ImageConditioningInput],
         enhance_prompt: bool = False,
-        streaming_prefetch_count: int | None = None,
         tiling_config: TilingConfig | None = None,
         max_batch_size: int = 1,
         sigmas: torch.Tensor | None = None,
@@ -121,7 +128,6 @@ class TI2VidOneStagePipeline:
             enhance_first_prompt=enhance_prompt,
             enhance_prompt_image=images[0][0] if len(images) > 0 else None,
             enhance_prompt_seed=seed,
-            streaming_prefetch_count=streaming_prefetch_count,
         )
         v_context_p, a_context_p = ctx_p.video_encoding, ctx_p.audio_encoding
         v_context_n, a_context_n = ctx_n.video_encoding, ctx_n.audio_encoding
@@ -170,7 +176,6 @@ class TI2VidOneStagePipeline:
             audio=ModalitySpec(
                 context=a_context_p,
             ),
-            streaming_prefetch_count=streaming_prefetch_count,
             max_batch_size=max_batch_size,
         )
 
@@ -181,7 +186,7 @@ class TI2VidOneStagePipeline:
 
 @torch.inference_mode()
 def main() -> None:
-    logging.getLogger().setLevel(logging.INFO)
+    logging.basicConfig(level=logging.INFO)
     checkpoint_path = detect_checkpoint_path()
     params = detect_params(checkpoint_path)
     parser = default_1_stage_arg_parser(params=params)
@@ -191,7 +196,8 @@ def main() -> None:
         gemma_root=args.gemma_root,
         loras=tuple(args.lora) if args.lora else (),
         quantization=args.quantization,
-        torch_compile=args.compile,
+        compilation_config=args.compile,
+        offload_mode=args.offload_mode,
     )
     video, audio = pipeline(
         prompt=args.prompt,
@@ -219,7 +225,6 @@ def main() -> None:
             stg_blocks=args.audio_stg_blocks,
         ),
         images=args.images,
-        streaming_prefetch_count=args.streaming_prefetch_count,
         max_batch_size=args.max_batch_size,
     )
 

packages/ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages.py

@@ -12,10 +12,15 @@ from ltx_core.components.noisers import GaussianNoiser
 from ltx_core.components.schedulers import LTX2Scheduler
 from ltx_core.loader import LoraPathStrengthAndSDOps
 from ltx_core.loader.registry import Registry
+from ltx_core.model.transformer.compiling import CompilationConfig
 from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
 from ltx_core.quantization import QuantizationPolicy
 from ltx_core.types import Audio, VideoPixelShape
-from ltx_pipelines.utils.args import ImageConditioningInput, default_2_stage_arg_parser, detect_checkpoint_path
+from ltx_pipelines.utils.args import (
+    ImageConditioningInput,
+    default_2_stage_arg_parser,
+    detect_checkpoint_path,
+)
 from ltx_pipelines.utils.blocks import (
     AudioDecoder,
     DiffusionStage,
@@ -35,7 +40,7 @@ from ltx_pipelines.utils.helpers import (
     get_device,
 )
 from ltx_pipelines.utils.media_io import encode_video
-from ltx_pipelines.utils.types import ModalitySpec
+from ltx_pipelines.utils.types import ModalitySpec, OffloadMode
 
 
 class TI2VidTwoStagesPipeline:
@@ -57,13 +62,16 @@ class TI2VidTwoStagesPipeline:
         device: torch.device | None = None,
         quantization: QuantizationPolicy | None = None,
         registry: Registry | None = None,
-        torch_compile: bool = False,
+        compilation_config: CompilationConfig | None = None,
+        offload_mode: OffloadMode = OffloadMode.NONE,
     ):
         self.device = device or get_device()
         self.dtype = torch.bfloat16
         self._scheduler = LTX2Scheduler()
 
-        self.prompt_encoder = PromptEncoder(checkpoint_path, gemma_root, self.dtype, self.device, registry=registry)
+        self.prompt_encoder = PromptEncoder(
+            checkpoint_path, gemma_root, self.dtype, self.device, registry=registry, offload_mode=offload_mode
+        )
         self.image_conditioner = ImageConditioner(checkpoint_path, self.dtype, self.device, registry=registry)
         self.upsampler = VideoUpsampler(
             checkpoint_path, spatial_upsampler_path, self.dtype, self.device, registry=registry
@@ -78,7 +86,8 @@ class TI2VidTwoStagesPipeline:
             loras=tuple(loras),
             quantization=quantization,
             registry=registry,
-            torch_compile=torch_compile,
+            compilation_config=compilation_config,
+            offload_mode=offload_mode,
         )
         self.stage_2 = DiffusionStage(
             checkpoint_path,
@@ -87,7 +96,8 @@ class TI2VidTwoStagesPipeline:
             loras=(*tuple(loras), *distilled_lora),
             quantization=quantization,
             registry=registry,
-            torch_compile=torch_compile,
+            compilation_config=compilation_config,
+            offload_mode=offload_mode,
         )
 
     def __call__(  # noqa: PLR0913
@@ -105,7 +115,6 @@ class TI2VidTwoStagesPipeline:
         images: list[ImageConditioningInput],
         tiling_config: TilingConfig | None = None,
         enhance_prompt: bool = False,
-        streaming_prefetch_count: int | None = None,
         max_batch_size: int = 1,
         stage_1_sigmas: torch.Tensor | None = None,
         stage_2_sigmas: torch.Tensor = STAGE_2_DISTILLED_SIGMAS,
@@ -121,7 +130,6 @@ class TI2VidTwoStagesPipeline:
             enhance_first_prompt=enhance_prompt,
             enhance_prompt_image=images[0][0] if len(images) > 0 else None,
             enhance_prompt_seed=seed,
-            streaming_prefetch_count=streaming_prefetch_count,
         )
         v_context_p, a_context_p = ctx_p.video_encoding, ctx_p.audio_encoding
         v_context_n, a_context_n = ctx_n.video_encoding, ctx_n.audio_encoding
@@ -170,7 +178,6 @@ class TI2VidTwoStagesPipeline:
             fps=frame_rate,
             video=ModalitySpec(context=v_context_p, conditionings=stage_1_conditionings),
             audio=ModalitySpec(context=a_context_p),
-            streaming_prefetch_count=streaming_prefetch_count,
             max_batch_size=max_batch_size,
         )
 
@@ -208,7 +215,6 @@ class TI2VidTwoStagesPipeline:
                 noise_scale=stage_2_sigmas[0].item(),
                 initial_latent=audio_state.latent,
             ),
-            streaming_prefetch_count=streaming_prefetch_count,
         )
 
         decoded_video = self.video_decoder(video_state.latent, tiling_config, generator)
@@ -218,7 +224,7 @@ class TI2VidTwoStagesPipeline:
 
 @torch.inference_mode()
 def main() -> None:
-    logging.getLogger().setLevel(logging.INFO)
+    logging.basicConfig(level=logging.INFO)
     checkpoint_path = detect_checkpoint_path()
     params = detect_params(checkpoint_path)
     parser = default_2_stage_arg_parser(params=params)
@@ -230,7 +236,8 @@ def main() -> None:
         gemma_root=args.gemma_root,
         loras=tuple(args.lora) if args.lora else (),
         quantization=args.quantization,
-        torch_compile=args.compile,
+        compilation_config=args.compile,
+        offload_mode=args.offload_mode,
     )
     tiling_config = TilingConfig.default()
     video_chunks_number = get_video_chunks_number(args.num_frames, tiling_config)
@@ -261,7 +268,6 @@ def main() -> None:
         ),
         images=args.images,
         tiling_config=tiling_config,
-        streaming_prefetch_count=args.streaming_prefetch_count,
         max_batch_size=args.max_batch_size,
     )
 

packages/ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages_hq.py

@@ -9,6 +9,7 @@ from ltx_core.components.noisers import GaussianNoiser
 from ltx_core.components.schedulers import LTX2Scheduler
 from ltx_core.loader import LoraPathStrengthAndSDOps
 from ltx_core.loader.registry import Registry
+from ltx_core.model.transformer.compiling import CompilationConfig
 from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
 from ltx_core.quantization import QuantizationPolicy
 from ltx_core.types import Audio, VideoLatentShape, VideoPixelShape
@@ -33,7 +34,7 @@ from ltx_pipelines.utils.helpers import (
 )
 from ltx_pipelines.utils.media_io import encode_video
 from ltx_pipelines.utils.samplers import res2s_audio_video_denoising_loop
-from ltx_pipelines.utils.types import ModalitySpec
+from ltx_pipelines.utils.types import ModalitySpec, OffloadMode
 
 
 class TI2VidTwoStagesHQPipeline:
@@ -60,7 +61,8 @@ class TI2VidTwoStagesHQPipeline:
         device: torch.device | None = None,
         quantization: QuantizationPolicy | None = None,
         registry: Registry | None = None,
-        torch_compile: bool = False,
+        compilation_config: CompilationConfig | None = None,
+        offload_mode: OffloadMode = OffloadMode.NONE,
     ):
         self.device = device or get_device()
         self.dtype = torch.bfloat16
@@ -77,7 +79,9 @@ class TI2VidTwoStagesHQPipeline:
             sd_ops=distilled_lora[0].sd_ops,
         )
 
-        self.prompt_encoder = PromptEncoder(checkpoint_path, gemma_root, self.dtype, self.device, registry=registry)
+        self.prompt_encoder = PromptEncoder(
+            checkpoint_path, gemma_root, self.dtype, self.device, registry=registry, offload_mode=offload_mode
+        )
         self.image_conditioner = ImageConditioner(checkpoint_path, self.dtype, self.device, registry=registry)
         self.upsampler = VideoUpsampler(
             checkpoint_path, spatial_upsampler_path, self.dtype, self.device, registry=registry
@@ -92,7 +96,8 @@ class TI2VidTwoStagesHQPipeline:
             loras=(*loras, distilled_lora_stage_1),
             quantization=quantization,
             registry=registry,
-            torch_compile=torch_compile,
+            compilation_config=compilation_config,
+            offload_mode=offload_mode,
         )
         self.stage_2 = DiffusionStage(
             checkpoint_path,
@@ -101,7 +106,8 @@ class TI2VidTwoStagesHQPipeline:
             loras=(*loras, distilled_lora_stage_2),
             quantization=quantization,
             registry=registry,
-            torch_compile=torch_compile,
+            compilation_config=compilation_config,
+            offload_mode=offload_mode,
         )
 
     @torch.inference_mode()
@@ -120,7 +126,6 @@ class TI2VidTwoStagesHQPipeline:
         images: list[ImageConditioningInput],
         tiling_config: TilingConfig | None = None,
         enhance_prompt: bool = False,
-        streaming_prefetch_count: int | None = None,
         max_batch_size: int = 1,
         stage_1_sigmas: torch.Tensor | None = None,
         stage_2_sigmas: torch.Tensor = STAGE_2_DISTILLED_SIGMAS,
@@ -136,7 +141,6 @@ class TI2VidTwoStagesHQPipeline:
             enhance_first_prompt=enhance_prompt,
             enhance_prompt_image=images[0][0] if len(images) > 0 else None,
             enhance_prompt_seed=seed,
-            streaming_prefetch_count=streaming_prefetch_count,
         )
         v_context_p, a_context_p = ctx_p.video_encoding, ctx_p.audio_encoding
         v_context_n, a_context_n = ctx_n.video_encoding, ctx_n.audio_encoding
@@ -190,7 +194,6 @@ class TI2VidTwoStagesHQPipeline:
             video=ModalitySpec(context=v_context_p, conditionings=stage_1_conditionings),
             audio=ModalitySpec(context=a_context_p),
             loop=res2s_audio_video_denoising_loop,
-            streaming_prefetch_count=streaming_prefetch_count,
             max_batch_size=max_batch_size,
         )
 
@@ -231,7 +234,6 @@ class TI2VidTwoStagesHQPipeline:
                 initial_latent=audio_state.latent,
             ),
             loop=res2s_audio_video_denoising_loop,
-            streaming_prefetch_count=streaming_prefetch_count,
         )
 
         decoded_video = self.video_decoder(video_state.latent, tiling_config, generator)
@@ -241,7 +243,7 @@ class TI2VidTwoStagesHQPipeline:
 
 @torch.inference_mode()
 def main() -> None:
-    logging.getLogger().setLevel(logging.INFO)
+    logging.basicConfig(level=logging.INFO)
     parser = hq_2_stage_arg_parser(params=LTX_2_3_HQ_PARAMS)
     args = parser.parse_args()
     pipeline = TI2VidTwoStagesHQPipeline(
@@ -253,7 +255,8 @@ def main() -> None:
         gemma_root=args.gemma_root,
         loras=tuple(args.lora) if args.lora else (),
         quantization=args.quantization,
-        torch_compile=args.compile,
+        compilation_config=args.compile,
+        offload_mode=args.offload_mode,
     )
     tiling_config = TilingConfig.default()
     video_chunks_number = get_video_chunks_number(args.num_frames, tiling_config)
@@ -284,7 +287,6 @@ def main() -> None:
         ),
         images=args.images,
         tiling_config=tiling_config,
-        streaming_prefetch_count=args.streaming_prefetch_count,
         max_batch_size=args.max_batch_size,
     )
 

packages/ltx-pipelines/src/ltx_pipelines/utils/__init__.py

@@ -16,15 +16,17 @@ from ltx_pipelines.utils.helpers import (
     image_conditionings_by_adding_guiding_latent,
 )
 from ltx_pipelines.utils.samplers import (
+    euler_cfg_pp_denoising_loop,
     euler_denoising_loop,
     gradient_estimating_euler_denoising_loop,
     res2s_audio_video_denoising_loop,
 )
-from ltx_pipelines.utils.types import Denoiser, ModalitySpec
+from ltx_pipelines.utils.types import DenoisedLatentResult, Denoiser, ModalitySpec
 
 __all__ = [
     "AudioConditioner",
     "AudioDecoder",
+    "DenoisedLatentResult",
     "Denoiser",
     "DiffusionStage",
     "FactoryGuidedDenoiser",
@@ -38,6 +40,7 @@ __all__ = [
     "assert_resolution",
     "cleanup_memory",
     "combined_image_conditionings",
+    "euler_cfg_pp_denoising_loop",
     "euler_denoising_loop",
     "get_device",
     "gradient_estimating_euler_denoising_loop",

packages/ltx-pipelines/src/ltx_pipelines/utils/args.py

@@ -1,8 +1,11 @@
 import argparse
+import json
+from collections.abc import Sequence
 from pathlib import Path
-from typing import NamedTuple
+from typing import Any, NamedTuple
 
 from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps
+from ltx_core.model.transformer.compiling import CompilationConfig
 from ltx_core.quantization import QuantizationPolicy
 from ltx_pipelines.utils.constants import (
     DEFAULT_IMAGE_CRF,
@@ -12,6 +15,8 @@ from ltx_pipelines.utils.constants import (
     LTX_2_3_PARAMS,
     PipelineParams,
 )
+from ltx_pipelines.utils.quantization_factory import QuantizationKind
+from ltx_pipelines.utils.types import OffloadMode
 
 
 class ImageConditioningInput(NamedTuple):
@@ -30,7 +35,7 @@ class VideoConditioningAction(argparse.Action):
         option_string: str | None = None,  # noqa: ARG002
     ) -> None:
         path, strength_str = values
-        resolved_path = resolve_path(path)
+        resolved_path = resolve_existing_path(path)
         strength = float(strength_str)
         current = getattr(namespace, self.dest) or []
         current.append((resolved_path, strength))
@@ -56,7 +61,7 @@ class VideoMaskConditioningAction(argparse.Action):
             msg = f"{option_string} requires exactly 2 arguments (MASK_PATH STRENGTH), got {len(values)}"
             raise argparse.ArgumentError(self, msg)
 
-        mask_path = resolve_path(values[0])
+        mask_path = resolve_existing_path(values[0])
         strength = float(values[1])
         setattr(namespace, self.dest, (mask_path, strength))
 
@@ -74,7 +79,7 @@ class ImageAction(argparse.Action):
             raise argparse.ArgumentError(self, msg)
 
         conditioning = ImageConditioningInput(
-            path=resolve_path(values[0]),
+            path=resolve_existing_path(values[0]),
             frame_idx=int(values[1]),
             strength=float(values[2]),
             crf=int(values[3]) if len(values) > 3 else DEFAULT_IMAGE_CRF,
@@ -99,7 +104,7 @@ class LoraAction(argparse.Action):
         path = values[0]
         strength_str = values[1] if len(values) > 1 else str(DEFAULT_LORA_STRENGTH)
 
-        resolved_path = resolve_path(path)
+        resolved_path = resolve_existing_path(path)
         strength = float(strength_str)
 
         current = getattr(namespace, self.dest) or []
@@ -107,49 +112,153 @@ class LoraAction(argparse.Action):
         setattr(namespace, self.dest, current)
 
 
-def resolve_path(path: str) -> str:
-    return str(Path(path).expanduser().resolve().as_posix())
+class CompileAction(argparse.Action):
+    """Parse ``--compile [KEY=VALUE ...]`` into a :class:`CompilationConfig`.
+    The flag is absent           -> ``args.compile`` stays at its default (``None``).
+    The flag is passed alone     -> ``CompilationConfig()`` (vanilla torch defaults).
+    The flag is passed with args -> ``CompilationConfig`` with the given fields overridden.
+    Errors (unknown key, malformed value, duplicate key, empty value) raise
+    :class:`argparse.ArgumentError` so argparse formats them as friendly CLI
+    messages rather than uncaught tracebacks.
+    """
 
+    _ALLOWED_KEYS = frozenset({"mode", "backend", "fullgraph", "dynamic", "inductor_config", "dynamo_config"})
 
-QUANTIZATION_POLICIES = ("fp8-cast", "fp8-scaled-mm")
-
-
-class QuantizationAction(argparse.Action):
     def __call__(
         self,
         parser: argparse.ArgumentParser,  # noqa: ARG002
         namespace: argparse.Namespace,
         values: list[str],
-        option_string: str | None = None,
+        option_string: str | None = None,  # noqa: ARG002
     ) -> None:
-        if len(values) > 2:
-            msg = (
-                f"{option_string} accepts at most 2 arguments (POLICY and optional AMAX_PATH), got {len(values)} values"
-            )
-            raise argparse.ArgumentError(self, msg)
+        overrides: dict[str, object] = {}
+        for item in values:
+            if "=" not in item:
+                raise argparse.ArgumentError(self, f"expects KEY=VALUE pairs, got: {item!r}")
+            key, _, raw = item.partition("=")
+            key = key.strip()
+            if key not in self._ALLOWED_KEYS:
+                raise argparse.ArgumentError(
+                    self,
+                    f"{key!r} is not a CompilationConfig field; valid keys: {sorted(self._ALLOWED_KEYS)}",
+                )
+            if key in overrides:
+                raise argparse.ArgumentError(self, f"{key} given more than once")
+            if key == "mode":
+                overrides[key] = self._parse_mode(raw)
+            elif key == "backend":
+                overrides[key] = self._parse_non_empty(key, raw)
+            elif key == "fullgraph":
+                overrides[key] = self._parse_bool(key, raw)
+            elif key == "dynamic":
+                overrides[key] = self._parse_dynamic(raw)
+            elif key in ("inductor_config", "dynamo_config"):
+                overrides[key] = self._parse_json_dict(key, raw)
+        setattr(namespace, self.dest, CompilationConfig(**overrides))
 
-        policy_name = values[0]
-        if policy_name not in QUANTIZATION_POLICIES:
-            msg = f"Unknown quantization policy '{policy_name}'. Choose from: {', '.join(QUANTIZATION_POLICIES)}"
-            raise argparse.ArgumentError(self, msg)
+    def _parse_mode(self, raw: str) -> str | None:
+        stripped = raw.strip()
+        if not stripped:
+            raise argparse.ArgumentError(self, "mode=... value cannot be empty (use mode=none to clear)")
+        if stripped.lower() == "none":
+            return None
+        return stripped
 
-        if policy_name == "fp8-cast":
-            if len(values) > 1:
-                msg = f"{option_string} fp8-cast does not accept additional arguments"
-                raise argparse.ArgumentError(self, msg)
-            policy = QuantizationPolicy.fp8_cast()
-        elif policy_name == "fp8-scaled-mm":
-            amax_path = resolve_path(values[1]) if len(values) > 1 else None
-            policy = QuantizationPolicy.fp8_scaled_mm(amax_path)
+    def _parse_non_empty(self, key: str, raw: str) -> str:
+        stripped = raw.strip()
+        if not stripped:
+            raise argparse.ArgumentError(self, f"{key}=... value cannot be empty")
+        return stripped
 
-        setattr(namespace, self.dest, policy)
+    def _parse_bool(self, key: str, raw: str) -> bool:
+        normalized = raw.strip().lower()
+        if normalized in ("true", "1"):
+            return True
+        if normalized in ("false", "0"):
+            return False
+        raise argparse.ArgumentError(self, f"{key}=... must be true or false; got {raw!r}")
+
+    def _parse_dynamic(self, raw: str) -> bool | None:
+        normalized = raw.strip().lower()
+        if normalized in ("auto", "none"):
+            return None
+        if normalized in ("true", "1"):
+            return True
+        if normalized in ("false", "0"):
+            return False
+        raise argparse.ArgumentError(self, f"dynamic=... must be auto/true/false; got {raw!r}")
+
+    def _parse_json_dict(self, key: str, raw: str) -> dict[str, Any]:
+        # Inline JSON object starts with '{'; otherwise treat the value as a path to a JSON file.
+        stripped = raw.strip()
+        if not stripped:
+            raise argparse.ArgumentError(self, f"{key}=... value cannot be empty")
+        if stripped.startswith("{"):
+            source = stripped
+        else:
+            path = Path(stripped).expanduser()
+            if not path.is_file():
+                raise argparse.ArgumentError(
+                    self, f"{key}=... must be a JSON object or a path to a JSON file; got {raw!r}"
+                )
+            source = path.read_text()
+        try:
+            value = json.loads(source)
+        except json.JSONDecodeError as e:
+            raise argparse.ArgumentError(self, f"{key}=... must be a JSON object; got {raw!r} ({e.msg})") from None
+        if not isinstance(value, dict):
+            raise argparse.ArgumentError(self, f"{key}=... must decode to a JSON object; got {type(value).__name__}")
+        return value
+
+
+def resolve_path(path: str) -> str:
+    return str(Path(path).expanduser().resolve().as_posix())
+
+
+def resolve_existing_path(path: str) -> str:
+    """Resolve *path* and verify it exists."""
+    resolved = resolve_path(path)
+    if not Path(resolved).exists():
+        raise argparse.ArgumentError(None, f"Path not found: {resolved}")
+    return resolved
+
+
+QUANTIZATION_POLICIES = tuple(k.value for k in QuantizationKind)
+
+
+def _resolve_quantization(namespace: argparse.Namespace) -> None:
+    # Resolution is deferred until after parse_args because fp8-scaled-mm needs the
+    # checkpoint path, which isn't on the namespace when the --quantization argument
+    # is parsed.
+    name = getattr(namespace, "quantization", None)
+    if name is None or isinstance(name, QuantizationPolicy):
+        return
+    try:
+        kind = QuantizationKind(name)
+    except ValueError:
+        return
+    ckpt = getattr(namespace, "checkpoint_path", None) or getattr(namespace, "distilled_checkpoint_path", None)
+    if ckpt is None:
+        raise SystemExit(f"--quantization {kind.value} requires --checkpoint-path (or --distilled-checkpoint-path).")
+    namespace.quantization = kind.to_policy(checkpoint_path=ckpt)
+
+
+class _PipelineArgumentParser(argparse.ArgumentParser):
+    def parse_args(  # type: ignore[override]
+        self,
+        args: Sequence[str] | None = None,
+        namespace: argparse.Namespace | None = None,
+    ) -> argparse.Namespace:
+        ns = super().parse_args(args, namespace)
+        _resolve_quantization(ns)
+        return ns
 
 
 def detect_checkpoint_path(distilled: bool = False) -> str:
     """Pre-parse argv to extract the checkpoint path before building the full parser."""
     pre = argparse.ArgumentParser(add_help=False)
     flag = "--distilled-checkpoint-path" if distilled else "--checkpoint-path"
-    pre.add_argument(flag, type=resolve_path, required=True)
+    pre.add_argument(flag, type=resolve_existing_path, required=True)
     known, _ = pre.parse_known_args()
     return known.distilled_checkpoint_path if distilled else known.checkpoint_path
 
@@ -158,18 +267,18 @@ def basic_arg_parser(
     params: PipelineParams = LTX_2_3_PARAMS,
     distilled: bool = False,
 ) -> argparse.ArgumentParser:
-    parser = argparse.ArgumentParser()
+    parser = _PipelineArgumentParser()
     if distilled:
         parser.add_argument(
             "--distilled-checkpoint-path",
-            type=resolve_path,
+            type=resolve_existing_path,
             required=True,
             help="Path to LTX-2 distilled model checkpoint (.safetensors file).",
         )
     else:
         parser.add_argument(
             "--checkpoint-path",
-            type=resolve_path,
+            type=resolve_existing_path,
             required=True,
             help="Path to LTX-2 model checkpoint (.safetensors file).",
         )
@@ -184,7 +293,7 @@ def basic_arg_parser(
         )
     parser.add_argument(
         "--gemma-root",
-        type=resolve_path,
+        type=resolve_existing_path,
         required=True,
         help="Path to the root directory containing the Gemma text encoder model files.",
     )
@@ -231,16 +340,19 @@ def basic_arg_parser(
         except ValueError as e:
             raise argparse.ArgumentTypeError(f"must be an integer, got {value}") from e
 
-    # Layer streaming
+    # Weight offloading
     parser.add_argument(
-        "--streaming-prefetch-count",
-        type=_positive_int,
-        default=None,
-        metavar="N",
+        "--offload",
+        dest="offload_mode",
+        type=OffloadMode,
+        default=OffloadMode.NONE,
+        choices=list(OffloadMode),
         help=(
-            "Enable layer streaming prefetching N layers ahead. "
-            "At most 1 + N layers reside on GPU at once. "
-            "Must be >= 1. Example: --streaming-prefetch-count 2"
+            "Weight offloading strategy. "
+            "'none' keeps all weights on GPU (default). "
+            "'cpu' pins weights in CPU RAM, streams to GPU per layer. "
+            "'disk' reads weights from disk on demand (lowest memory). "
+            "Example: --offload cpu"
         ),
     )
 
@@ -260,22 +372,32 @@ def basic_arg_parser(
 
     parser.add_argument(
         "--quantization",
-        dest="quantization",
-        action=QuantizationAction,
-        nargs="+",
-        metavar=("POLICY", "AMAX_PATH"),
+        choices=QUANTIZATION_POLICIES,
         default=None,
         help=(
             f"Quantization policy: {', '.join(QUANTIZATION_POLICIES)}. "
             "fp8-cast uses FP8 casting with upcasting during inference. "
-            "fp8-scaled-mm uses FP8 scaled matrix multiplication (optionally provide amax calibration file path). "
-            "Example: --quantization fp8-cast or --quantization fp8-scaled-mm /path/to/amax.json"
+            "fp8-scaled-mm uses FP8 scaled matrix multiplication; the layer set is auto-discovered "
+            "from the checkpoint's .weight_scale tensors. "
+            "Example: --quantization fp8-cast or --quantization fp8-scaled-mm"
         ),
     )
     parser.add_argument(
         "--compile",
-        action="store_true",
-        help="Enable torch.compile for transformer blocks to optimize performance.",
+        nargs="*",
+        action=CompileAction,
+        default=None,
+        metavar="KEY=VALUE",
+        help=(
+            "Enable torch.compile for transformer blocks. Pass alone for defaults, "
+            "or with KEY=VALUE overrides for any CompilationConfig field. "
+            "Keys: mode, backend, fullgraph, dynamic, inductor_config, dynamo_config. "
+            "inductor_config/dynamo_config take JSON objects (inline or a path to a .json file) "
+            "that fully replace the defaults. "
+            "Examples: --compile  or  --compile mode=reduce-overhead  or  "
+            "--compile mode=reduce-overhead fullgraph=true backend=eager  or  "
+            "--compile inductor_config='{\"max_autotune\": true}'"
+        ),
     )
     return parser
 
@@ -338,12 +460,59 @@ def video_editing_arg_parser(
     (no height/width/num-frames; resolution comes from input video). Default is distilled checkpoint only.
     """
     parser = basic_arg_parser(distilled=distilled)
-    parser.add_argument("--video-path", type=resolve_path, required=True, help="Path to the source video.")
+    parser.add_argument("--video-path", type=resolve_existing_path, required=True, help="Path to the source video.")
     parser.add_argument("--start-time", type=float, required=True, help="Start time of the region to regenerate (s).")
     parser.add_argument("--end-time", type=float, required=True, help="End time of the region to regenerate (s).")
     return parser
 
 
+def lipdub_arg_parser(
+    params: PipelineParams = LTX_2_3_PARAMS,
+) -> argparse.ArgumentParser:
+    """Argument parser for the lip-dub pipeline.
+    Frame count and frame rate are derived from the reference video at runtime (the frame count
+    is silently snapped down to the nearest 8k+1), so this parser intentionally omits
+    --num-frames, --frame-rate, and --image. Distilled checkpoint only.
+    """
+    parser = basic_arg_parser(params=params, distilled=True)
+    parser.add_argument(
+        "--height",
+        type=int,
+        default=params.stage_2_height,
+        help=(
+            f"Height of the generated video in pixels, should be divisible by 64 (default: {params.stage_2_height})."
+        ),
+    )
+    parser.add_argument(
+        "--width",
+        type=int,
+        default=params.stage_2_width,
+        help=f"Width of the generated video in pixels, should be divisible by 64 (default: {params.stage_2_width}).",
+    )
+    parser.add_argument(
+        "--spatial-upsampler-path",
+        type=resolve_path,
+        required=True,
+        help=(
+            "Path to the spatial upsampler model used to increase the resolution "
+            "of the generated video in the latent space."
+        ),
+    )
+    parser.add_argument(
+        "--reference-video",
+        type=resolve_path,
+        required=True,
+        help="Reference video file (video + audio track used for IC-LoRA and audio identity).",
+    )
+    parser.add_argument(
+        "--reference-strength",
+        type=float,
+        default=1.0,
+        help="Strength for IC-LoRA video reference conditioning (default: 1.0).",
+    )
+    return parser
+
+
 def default_1_stage_arg_parser(params: PipelineParams = LTX_2_3_PARAMS) -> argparse.ArgumentParser:
     video_guider = params.video_guider_params
     audio_guider = params.audio_guider_params
@@ -413,7 +582,7 @@ def default_1_stage_arg_parser(params: PipelineParams = LTX_2_3_PARAMS) -> argpa
         default=video_guider.skip_step,
         help=(
             "Video skip step N controls periodic skipping during the video diffusion process: "
-            "only steps where step_index % (N + 1) == 0 are processed, all others are skipped "
+            "only steps where step_index %% (N + 1) == 0 are processed, all others are skipped "
             f"(e.g., 0 = no skipping; 1 = skip every other step; 2 = skip 2 of every 3 steps; "
             f"default: {video_guider.skip_step})."
         ),
@@ -473,7 +642,7 @@ def default_1_stage_arg_parser(params: PipelineParams = LTX_2_3_PARAMS) -> argpa
         default=audio_guider.skip_step,
         help=(
             "Audio skip step N controls periodic skipping during the audio diffusion process: "
-            "only steps where step_index % (N + 1) == 0 are processed, all others are skipped "
+            "only steps where step_index %% (N + 1) == 0 are processed, all others are skipped "
             f"(e.g., 0 = no skipping; 1 = skip every other step; 2 = skip 2 of every 3 steps; "
             f"default: {audio_guider.skip_step})."
         ),
@@ -481,6 +650,62 @@ def default_1_stage_arg_parser(params: PipelineParams = LTX_2_3_PARAMS) -> argpa
     return parser
 
 
+def default_1_stage_t2a_arg_parser(params: PipelineParams = LTX_2_3_PARAMS) -> argparse.ArgumentParser:
+    """Argument parser for single-stage text-to-audio pipelines (audio-only)."""
+    audio_guider = params.audio_guider_params
+    parser = basic_arg_parser(params=params)
+    parser.add_argument(
+        "--num-frames",
+        type=int,
+        default=params.num_frames,
+        help="Number of frames used to derive audio duration (num-frames / frame-rate).",
+    )
+    parser.add_argument(
+        "--frame-rate",
+        type=float,
+        default=params.frame_rate,
+        help="Frame rate used with --num-frames to derive the audio duration.",
+    )
+    parser.add_argument(
+        "--negative-prompt",
+        type=str,
+        default=DEFAULT_NEGATIVE_PROMPT,
+        help="Negative prompt to steer audio generation away from artifacts.",
+    )
+    parser.add_argument(
+        "--audio-cfg-guidance-scale",
+        type=float,
+        default=audio_guider.cfg_scale,
+        help=f"Audio CFG scale (default: {audio_guider.cfg_scale}).",
+    )
+    parser.add_argument(
+        "--audio-stg-guidance-scale",
+        type=float,
+        default=audio_guider.stg_scale,
+        help=f"Audio STG scale (default: {audio_guider.stg_scale}).",
+    )
+    parser.add_argument(
+        "--audio-rescale-scale",
+        type=float,
+        default=audio_guider.rescale_scale,
+        help=f"Audio rescale scale (default: {audio_guider.rescale_scale}).",
+    )
+    parser.add_argument(
+        "--audio-stg-blocks",
+        type=int,
+        nargs="*",
+        default=audio_guider.stg_blocks,
+        help=f"Blocks to perturb for Audio STG (default: {audio_guider.stg_blocks}).",
+    )
+    parser.add_argument(
+        "--audio-skip-step",
+        type=int,
+        default=audio_guider.skip_step,
+        help=f"Audio skip step (default: {audio_guider.skip_step}).",
+    )
+    return parser
+
+
 def default_2_stage_arg_parser(params: PipelineParams = LTX_2_3_PARAMS) -> argparse.ArgumentParser:
     parser = default_1_stage_arg_parser(params=params)
     parser.set_defaults(height=params.stage_2_height, width=params.stage_2_width)
@@ -513,7 +738,7 @@ def default_2_stage_arg_parser(params: PipelineParams = LTX_2_3_PARAMS) -> argpa
     )
     parser.add_argument(
         "--spatial-upsampler-path",
-        type=resolve_path,
+        type=resolve_existing_path,
         required=True,
         help=(
             "Path to the spatial upsampler model used to increase the resolution "
@@ -556,7 +781,7 @@ def default_2_stage_distilled_arg_parser(params: PipelineParams = LTX_2_3_PARAMS
             )
     parser.add_argument(
         "--spatial-upsampler-path",
-        type=resolve_path,
+        type=resolve_existing_path,
         required=True,
         help=(
             "Path to the spatial upsampler model used to increase the resolution "

packages/ltx-pipelines/src/ltx_pipelines/utils/blocks.py

@@ -6,6 +6,7 @@ removes the need for :class:`ModelLedger`.
 
 from __future__ import annotations
 
+import copy
 import logging
 from collections.abc import Iterator
 from contextlib import AbstractContextManager, contextmanager
@@ -15,13 +16,16 @@ from typing import Callable, TypeVar
 import torch
 
 from ltx_core.batch_split import BatchSplitAdapter
+from ltx_core.block_streaming import DISK_CPU_SLOTS, StreamingModelBuilder
 from ltx_core.components.diffusion_steps import EulerDiffusionStep
 from ltx_core.components.noisers import Noiser
 from ltx_core.components.patchifiers import AudioPatchifier, VideoLatentPatchifier
 from ltx_core.components.protocols import DiffusionStepProtocol
-from ltx_core.layer_streaming import LayerStreamingWrapper
 from ltx_core.loader import SDOps
-from ltx_core.loader.primitives import LoraPathStrengthAndSDOps
+from ltx_core.loader.attention_ops import set_attention_module_op
+from ltx_core.loader.fuse_loras import bf16_fuse_rule
+from ltx_core.loader.module_ops import ModuleOps
+from ltx_core.loader.primitives import BuilderProtocol, LoraPathStrengthAndSDOps, ModelBuilderProtocol
 from ltx_core.loader.registry import DummyRegistry, Registry
 from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder as Builder
 from ltx_core.model.audio_vae import (
@@ -35,14 +39,24 @@ from ltx_core.model.audio_vae import (
 from ltx_core.model.audio_vae import (
     decode_audio as vae_decode_audio,
 )
+from ltx_core.model.model_protocol import LTXModelProtocol, ModelConfigurator
 from ltx_core.model.transformer import (
     LTXV_MODEL_COMFY_RENAMING_MAP,
     LTXModelConfigurator,
     X0Model,
 )
-from ltx_core.model.transformer.compiling import COMPILE_TRANSFORMER, modify_sd_ops_for_compilation
+from ltx_core.model.transformer.attention import (
+    AttentionCallable,
+    AttentionFunction,
+)
+from ltx_core.model.transformer.compiling import (
+    CompilationConfig,
+    build_compile_transformer_op,
+    modify_sd_ops_for_compilation,
+)
 from ltx_core.model.upsampler import LatentUpsamplerConfigurator, upsample_video
 from ltx_core.model.video_vae import (
+    MEMORY_EFFICIENT_DECODE,
     VAE_DECODER_COMFY_KEYS_FILTER,
     VAE_ENCODER_COMFY_KEYS_FILTER,
     TilingConfig,
@@ -50,7 +64,7 @@ from ltx_core.model.video_vae import (
     VideoEncoder,
     VideoEncoderConfigurator,
 )
-from ltx_core.quantization import QuantizationPolicy
+from ltx_core.quantization import QuantizationPolicy, fp8_cast_fuse_rule
 from ltx_core.text_encoders.gemma import (
     EMBEDDINGS_PROCESSOR_KEY_OPS,
     GEMMA_LLM_KEY_OPS,
@@ -59,7 +73,7 @@ from ltx_core.text_encoders.gemma import (
     GemmaTextEncoderConfigurator,
     module_ops_from_gemma_root,
 )
-from ltx_core.text_encoders.gemma.embeddings_processor import EmbeddingsProcessorOutput
+from ltx_core.text_encoders.gemma.embeddings_processor import EmbeddingsProcessor, EmbeddingsProcessorOutput
 from ltx_core.tools import AudioLatentTools, LatentTools, VideoLatentTools
 from ltx_core.types import Audio, AudioLatentShape, LatentState, VideoLatentShape, VideoPixelShape
 from ltx_core.utils import find_matching_file
@@ -70,7 +84,7 @@ from ltx_pipelines.utils.helpers import (
     generate_enhanced_prompt,
 )
 from ltx_pipelines.utils.samplers import euler_denoising_loop
-from ltx_pipelines.utils.types import Denoiser, ModalitySpec
+from ltx_pipelines.utils.types import Denoiser, ModalitySpec, OffloadMode
 
 logger = logging.getLogger(__name__)
 
@@ -83,36 +97,62 @@ _M = TypeVar("_M", bound=torch.nn.Module)
 # ---------------------------------------------------------------------------
 
 
+def _chain_quantization(
+    sd_ops: SDOps,
+    module_ops: tuple[ModuleOps, ...],
+    quantization: QuantizationPolicy,
+) -> tuple[SDOps, tuple[ModuleOps, ...]]:
+    chained_sd_ops = sd_ops
+    if quantization.sd_ops is not None:
+        chained_sd_ops = SDOps(
+            name=f"sd_ops_chain_{sd_ops.name}+{quantization.sd_ops.name}",
+            mapping=(*sd_ops.mapping, *quantization.sd_ops.mapping),
+        )
+    return chained_sd_ops, (*module_ops, *quantization.module_ops)
+
+
+def _apply_compile_ops(
+    sd_ops: SDOps,
+    module_ops: tuple[ModuleOps, ...],
+    loras: tuple[LoraPathStrengthAndSDOps, ...],
+    number_of_layers: int,
+    compilation_config: CompilationConfig,
+) -> tuple[SDOps, tuple[ModuleOps, ...], tuple[LoraPathStrengthAndSDOps, ...]]:
+    """Rewrite sd_ops/module_ops/LoRAs for compiled blocks (params land under ``_orig_mod``)."""
+    sd_ops = modify_sd_ops_for_compilation(sd_ops, number_of_layers)
+    compile_op = build_compile_transformer_op(compilation_config)
+    module_ops = (*module_ops, compile_op)
+    loras = tuple(
+        LoraPathStrengthAndSDOps(
+            lora.path,
+            lora.strength,
+            modify_sd_ops_for_compilation(lora.sd_ops, number_of_layers),
+        )
+        for lora in loras
+    )
+    return sd_ops, module_ops, loras
+
+
 @contextmanager
 def _streaming_model(
-    model: _M,
-    layers_attr: str,
+    builder: StreamingModelBuilder,
+    offload_mode: OffloadMode,
     target_device: torch.device,
-    prefetch_count: int,
-) -> Iterator[_M]:
-    """Wrap *model* with :class:`LayerStreamingWrapper`, yield it, then tear down."""
-    wrapped = LayerStreamingWrapper(
-        model,
-        layers_attr=layers_attr,
-        target_device=target_device,
-        prefetch_count=prefetch_count,
+    dtype: torch.dtype,
+) -> Iterator:
+    """Build a streaming wrapper, yield it, then tear down and free memory."""
+    cpu_slots_count = DISK_CPU_SLOTS if offload_mode == OffloadMode.DISK else None
+    wrapped = builder.build(
+        device=target_device,
+        dtype=dtype,
+        cpu_slots_count=cpu_slots_count,
     )
     try:
-        yield wrapped  # type: ignore[misc]
+        yield wrapped
     finally:
         wrapped.teardown()
         wrapped.to("meta")
         cleanup_memory()
-        # Flush the host (pinned) memory cache so that freed pinned pages are
-        # returned to the OS.  Without this, sequential streaming models
-        # (e.g. text encoder then transformer) exhaust host memory because the
-        # CachingHostAllocator keeps freed blocks cached indefinitely.
-        torch.cuda.synchronize(device=target_device)
-        try:
-            if hasattr(torch._C, "_host_emptyCache"):
-                torch._C._host_emptyCache()
-        except Exception:
-            logger.warning("Host empty cache cleanup failed; ignoring.", exc_info=True)
 
 
 def _build_state(
@@ -154,7 +194,7 @@ class DiffusionStage:
     pattern in every pipeline.
     """
 
-    def __init__(
+    def __init__(  # noqa: PLR0913
         self,
         checkpoint_path: str,
         dtype: torch.dtype,
@@ -162,65 +202,130 @@ class DiffusionStage:
         loras: tuple[LoraPathStrengthAndSDOps, ...] = (),
         quantization: QuantizationPolicy | None = None,
         registry: Registry | None = None,
-        torch_compile: bool = False,
+        compilation_config: CompilationConfig | None = None,
+        offload_mode: OffloadMode = OffloadMode.NONE,
+        transformer_builder: ModelBuilderProtocol[LTXModelProtocol] | None = None,
+        model_configurator: type[ModelConfigurator] = LTXModelConfigurator,
+        model_sd_ops: SDOps = LTXV_MODEL_COMFY_RENAMING_MAP,
     ) -> None:
+        self._checkpoint_path = checkpoint_path
         self._dtype = dtype
         self._device = device
         self._quantization = quantization
-        self._torch_compile = torch_compile
-        self._transformer_builder = Builder(
-            model_path=checkpoint_path,
-            model_class_configurator=LTXModelConfigurator,
-            model_sd_ops=LTXV_MODEL_COMFY_RENAMING_MAP,
-            loras=tuple(loras),
-            registry=registry or DummyRegistry(),
+        self._compilation_config = compilation_config
+        self._offload_mode = offload_mode
+        # A quantization policy may pin its own configurator; otherwise use the one
+        # provided by the caller (defaults to the audio-video LTXModelConfigurator).
+        configurator = (
+            quantization.model_configurator
+            if quantization is not None and quantization.model_configurator is not None
+            else model_configurator
         )
+        if transformer_builder is not None:
+            self._transformer_builder = transformer_builder
+        else:
+            self._transformer_builder = Builder(
+                model_path=checkpoint_path,
+                model_class_configurator=configurator,
+                model_sd_ops=model_sd_ops,
+                loras=tuple(loras),
+                registry=registry or DummyRegistry(),
+            )
+
+        if offload_mode != OffloadMode.NONE:
+            # WeightsProvider currently only supports plain bf16 + fp8_cast LoRA fusion
+            # (no companion-key emission). Quantization policies that emit
+            # companion keys (e.g. ``.weight_scale``) cannot be streamed yet.
+            if quantization is not None and quantization.fuse_rule is not fp8_cast_fuse_rule:
+                raise ValueError(
+                    "Block streaming is not supported with this quantization policy "
+                    "(only bf16 and fp8_cast are currently supported)."
+                )
+            streaming_sd_ops: SDOps = model_sd_ops
+            streaming_module_ops: tuple[ModuleOps, ...] = ()
+            streaming_loras = tuple(loras)
+
+            if compilation_config:
+                number_of_layers = self._transformer_builder.model_config()["transformer"]["num_layers"]
+                streaming_sd_ops, streaming_module_ops, streaming_loras = _apply_compile_ops(
+                    streaming_sd_ops, streaming_module_ops, streaming_loras, number_of_layers
+                )
+            if quantization is not None:
+                streaming_sd_ops, streaming_module_ops = _chain_quantization(
+                    streaming_sd_ops, streaming_module_ops, quantization
+                )
+            self._streaming_builder = StreamingModelBuilder(
+                model_class_configurator=configurator,
+                model_path=checkpoint_path,
+                model_sd_ops=streaming_sd_ops,
+                module_ops=streaming_module_ops,
+                loras=streaming_loras,
+                registry=registry or DummyRegistry(),
+                fuse_rule=quantization.fuse_rule if quantization is not None else bf16_fuse_rule,
+                blocks_attr="transformer_blocks",
+                blocks_prefix="transformer_blocks",
+            )
+
+    def with_attention(self, attention: AttentionFunction | AttentionCallable | None) -> "DiffusionStage":
+        """Return a new ``DiffusionStage`` that pins the transformer build to ``attention``.
+        Functional: never mutates ``self``. The returned stage shares all other
+        configuration with the original; only the underlying builders' ``module_ops``
+        gain a ``set_attention_module_op(attention)`` entry so subsequent transformer
+        builds use that kernel. ``attention=None`` is a no-op (returns ``self``).
+        """
+        if attention is None:
+            return self
+        op = set_attention_module_op(attention)
+        new = copy.copy(self)
+        new._transformer_builder = self._transformer_builder.with_module_ops(
+            (*self._transformer_builder.module_ops, op),
+        )
+        if self._offload_mode != OffloadMode.NONE:
+            new._streaming_builder = self._streaming_builder.with_module_ops(
+                (*self._streaming_builder.module_ops, op),
+            )
+        return new
 
     def _build_transformer(self, *, device: torch.device | None = None, **kwargs: object) -> X0Model:
         target = device or self._device
         sd_ops = self._transformer_builder.model_sd_ops
         module_ops = self._transformer_builder.module_ops
         loras = self._transformer_builder.loras
-        if self._torch_compile:
-            module_ops = (*module_ops, COMPILE_TRANSFORMER)
+        if self._compilation_config is not None:
             number_of_layers = self._transformer_builder.model_config()["transformer"]["num_layers"]
-            sd_ops = modify_sd_ops_for_compilation(sd_ops, number_of_layers)
-            loras = tuple(
-                LoraPathStrengthAndSDOps(
-                    lora.path,
-                    lora.strength,
-                    modify_sd_ops_for_compilation(
-                        lora.sd_ops if lora.sd_ops is not None else SDOps(name="identity"), number_of_layers
-                    ),
-                )
-                for lora in loras
+            sd_ops, module_ops, loras = _apply_compile_ops(
+                sd_ops, module_ops, loras, number_of_layers, self._compilation_config
             )
         if self._quantization is not None:
-            module_ops = (*module_ops, *self._quantization.module_ops)
-            sd_ops = SDOps(
-                name=f"sd_ops_chain_{sd_ops.name}+{self._quantization.sd_ops.name}",
-                mapping=(*sd_ops.mapping, *self._quantization.sd_ops.mapping),
-            )
+            sd_ops, module_ops = _chain_quantization(sd_ops, module_ops, self._quantization)
 
         builder = self._transformer_builder.with_module_ops(module_ops).with_sd_ops(sd_ops).with_loras(loras)
+        if self._quantization is not None:
+            builder = builder.with_fuse_rule(self._quantization.fuse_rule)
         return X0Model(builder.build(device=target, **kwargs)).to(target).eval()
 
-    def _transformer_ctx(
-        self,
-        streaming_prefetch_count: int | None,
-        **kwargs: object,
-    ) -> AbstractContextManager:
-        if streaming_prefetch_count is not None:
-            return _streaming_model(
-                self._build_transformer(device=torch.device("cpu"), **kwargs),
-                layers_attr="velocity_model.transformer_blocks",
-                target_device=self._device,
-                prefetch_count=streaming_prefetch_count,
-            )
+    @contextmanager
+    def _streaming_transformer_ctx(self) -> Iterator[X0Model]:
+        with _streaming_model(
+            self._streaming_builder, self._offload_mode, self._device, self._dtype
+        ) as streaming_wrapper:
+            yield X0Model(streaming_wrapper).eval()
+
+    def _transformer_ctx(self, **kwargs: object) -> AbstractContextManager:
+        if self._offload_mode != OffloadMode.NONE:
+            return self._streaming_transformer_ctx()
         return gpu_model(self._build_transformer(**kwargs))
 
-    def __call__(  # noqa: PLR0913
+    def model_context(self, **kwargs: object) -> AbstractContextManager:
+        """Build the transformer, yield it, then free its memory on exit.
+        Keyword arguments are forwarded to the underlying builder (e.g.
+        ``video_tools`` required by ``TiledDataParallelBuilder``).
+        """
+        return self._transformer_ctx(**kwargs)
+
+    def run(  # noqa: PLR0913
         self,
+        transformer: object,
         denoiser: Denoiser,
         sigmas: torch.Tensor,
         noiser: Noiser,
@@ -232,27 +337,14 @@ class DiffusionStage:
         audio: ModalitySpec | None = None,
         stepper: DiffusionStepProtocol | None = None,
         loop: Callable[..., tuple[LatentState | None, LatentState | None]] | None = None,
-        streaming_prefetch_count: int | None = None,
         max_batch_size: int = 1,
     ) -> tuple[LatentState | None, LatentState | None]:
-        """Build transformer → run denoising loop → free transformer.
-        Args:
-            width: Output width in pixels.
-            height: Output height in pixels.
-            frames: Number of output frames.
-            fps: Frame rate.
-            loop: Denoising loop function. Must accept
-                ``(sigmas, video_state, audio_state, stepper, transformer, denoiser)``
-                as the first six positional arguments. When ``None``, resolves to
-                :func:`euler_denoising_loop` at call time.
-            streaming_prefetch_count: When set, build the transformer on CPU and
-                wrap with :class:`LayerStreamingWrapper` for memory-efficient
-                inference, prefetching this many layers ahead.
-            max_batch_size: Maximum batch size per transformer forward pass.
-                Guided denoisers make up to 4 transformer calls per step.
-                When set to a value > 1, the transformer batches multiple
-                calls together, reducing layer-streaming PCIe transfers.
-                Default ``1`` preserves sequential behavior.
+        """Run denoising with a pre-built transformer.
+        Same semantics as ``__call__`` but accepts a pre-built transformer so
+        the model can be shared across multiple calls (e.g. tiled inference
+        inside a single ``model_context()`` block). Audio supports
+        ``ModalitySpec(frozen=True)`` to keep the latent unchanged throughout
+        denoising while still providing cross-modal context to the transformer.
         Returns ``(video_state | None, audio_state | None)`` with cleared
         conditionings and unpatchified latents for present modalities.
         """
@@ -261,7 +353,6 @@ class DiffusionStage:
 
         if loop is None:
             loop = euler_denoising_loop
-
         if stepper is None:
             stepper = EulerDiffusionStep()
 
@@ -281,28 +372,80 @@ class DiffusionStage:
             audio_tools = AudioLatentTools(AudioPatchifier(patch_size=1), a_shape)
             audio_state = _build_state(audio, audio_tools, noiser, self._dtype, self._device)
 
-        with self._transformer_ctx(streaming_prefetch_count, video_tools=video_tools) as base_transformer:
-            transformer = BatchSplitAdapter(base_transformer, max_batch_size=max_batch_size)
-            video_state, audio_state = loop(
-                sigmas=sigmas,
-                video_state=video_state,
-                audio_state=audio_state,
-                stepper=stepper,
-                transformer=transformer,
-                denoiser=denoiser,
-            )
+        wrapped = BatchSplitAdapter(transformer, max_batch_size=max_batch_size)  # type: ignore[arg-type]
+        video_state, audio_state = loop(
+            sigmas=sigmas,
+            video_state=video_state,
+            audio_state=audio_state,
+            stepper=stepper,
+            transformer=wrapped,
+            denoiser=denoiser,
+        )
 
-        # Post-process: clear conditionings and unpatchify
         if video_state is not None and video_tools is not None:
             video_state = video_tools.clear_conditioning(video_state)
             video_state = video_tools.unpatchify(video_state)
-
         if audio_state is not None and audio_tools is not None:
             audio_state = audio_tools.clear_conditioning(audio_state)
             audio_state = audio_tools.unpatchify(audio_state)
 
         return video_state, audio_state
 
+    def __call__(  # noqa: PLR0913
+        self,
+        denoiser: Denoiser,
+        sigmas: torch.Tensor,
+        noiser: Noiser,
+        width: int,
+        height: int,
+        frames: int,
+        fps: float,
+        video: ModalitySpec | None = None,
+        audio: ModalitySpec | None = None,
+        stepper: DiffusionStepProtocol | None = None,
+        loop: Callable[..., tuple[LatentState | None, LatentState | None]] | None = None,
+        max_batch_size: int = 1,
+    ) -> tuple[LatentState | None, LatentState | None]:
+        """Build transformer -> run denoising loop -> free transformer.
+        Returns ``(video_state | None, audio_state | None)`` with cleared
+        conditionings and unpatchified latents for present modalities.
+        """
+        # Build video_tools up front so it can be forwarded to the transformer
+        # context (required by TiledDataParallelBuilder in multi-GPU mode).
+        # `run()` rebuilds its own tools internally; the duplication is cheap.
+        video_tools: LatentTools | None = None
+        if video is not None:
+            pixel_shape = VideoPixelShape(batch=1, frames=frames, height=height, width=width, fps=fps)
+            v_shape = VideoLatentShape.from_pixel_shape(pixel_shape)
+            video_tools = VideoLatentTools(VideoLatentPatchifier(patch_size=1), v_shape, fps)
+
+        mode = "streaming" if self._offload_mode != OffloadMode.NONE else "standard"
+        logger.info("Building transformer (%s) from %s", mode, self._checkpoint_path)
+        with self._transformer_ctx(video_tools=video_tools) as transformer:
+            logger.info(
+                "Running denoising loop (%d steps, %dx%d %d frames @ %.1f fps)",
+                len(sigmas) - 1,
+                width,
+                height,
+                frames,
+                fps,
+            )
+            return self.run(
+                transformer,
+                denoiser,
+                sigmas,
+                noiser,
+                width,
+                height,
+                frames,
+                fps,
+                video,
+                audio,
+                stepper,
+                loop,
+                max_batch_size,
+            )
+
 
 # ---------------------------------------------------------------------------
 # PromptEncoder
@@ -322,21 +465,43 @@ class PromptEncoder:
         dtype: torch.dtype,
         device: torch.device,
         registry: Registry | None = None,
+        offload_mode: OffloadMode = OffloadMode.NONE,
+        text_encoder_builder: BuilderProtocol | None = None,
     ) -> None:
+        self._gemma_root = gemma_root
+        self._checkpoint_path = checkpoint_path
         self._dtype = dtype
         self._device = device
+        self._offload_mode = offload_mode
 
-        module_ops = module_ops_from_gemma_root(gemma_root)
-        model_folder = find_matching_file(gemma_root, "model*.safetensors").parent
-        weight_paths = [str(p) for p in model_folder.rglob("*.safetensors")]
-
-        self._text_encoder_builder = Builder(
-            model_path=tuple(weight_paths),
-            model_class_configurator=GemmaTextEncoderConfigurator,
-            model_sd_ops=GEMMA_LLM_KEY_OPS,
-            module_ops=(GEMMA_MODEL_OPS, *module_ops),
-            registry=registry or DummyRegistry(),
-        )
+        if text_encoder_builder is not None:
+            if offload_mode != OffloadMode.NONE:
+                raise ValueError(
+                    "text_encoder_builder cannot be used with offload_mode != OffloadMode.NONE "
+                    "because no streaming text encoder builder is available."
+                )
+            self._text_encoder_builder = text_encoder_builder
+            self._streaming_text_encoder_builder = None
+        else:
+            module_ops = module_ops_from_gemma_root(gemma_root)
+            model_folder = find_matching_file(gemma_root, "model*.safetensors").parent
+            weight_paths = [str(p) for p in model_folder.rglob("*.safetensors")]
+            self._text_encoder_builder = Builder(
+                model_path=tuple(weight_paths),
+                model_class_configurator=GemmaTextEncoderConfigurator,
+                model_sd_ops=GEMMA_LLM_KEY_OPS,
+                module_ops=(GEMMA_MODEL_OPS, *module_ops),
+                registry=registry or DummyRegistry(),
+            )
+            self._streaming_text_encoder_builder = StreamingModelBuilder(
+                model_path=tuple(weight_paths),
+                model_class_configurator=GemmaTextEncoderConfigurator,
+                model_sd_ops=GEMMA_LLM_KEY_OPS,
+                module_ops=(GEMMA_MODEL_OPS, *module_ops),
+                registry=registry or DummyRegistry(),
+                blocks_attr="model.model.language_model.layers",
+                blocks_prefix="model.model.language_model.layers",
+            )
         self._embeddings_processor_builder = Builder(
             model_path=checkpoint_path,
             model_class_configurator=EmbeddingsProcessorConfigurator,
@@ -344,18 +509,18 @@ class PromptEncoder:
             registry=registry or DummyRegistry(),
         )
 
-    def _text_encoder_ctx(
-        self,
-        streaming_prefetch_count: int | None,
-    ) -> AbstractContextManager:
-        if streaming_prefetch_count is not None:
-            return _streaming_model(
-                self._text_encoder_builder.build(device=torch.device("cpu"), dtype=self._dtype).eval(),
-                layers_attr="model.model.language_model.layers",
-                target_device=self._device,
-                prefetch_count=streaming_prefetch_count,
-            )
-        return gpu_model(self._text_encoder_builder.build(device=self._device, dtype=self._dtype).eval())
+    def _build_text_encoder(self) -> torch.nn.Module:
+        """Build the Gemma text encoder (non-streaming path)."""
+        return self._text_encoder_builder.build(device=self._device, dtype=self._dtype).eval()
+
+    def _build_embeddings_processor(self) -> EmbeddingsProcessor:
+        """Build the embeddings processor on the target device."""
+        return self._embeddings_processor_builder.build(device=self._device, dtype=self._dtype).eval()
+
+    def _text_encoder_ctx(self) -> AbstractContextManager:
+        if self._offload_mode != OffloadMode.NONE:
+            return _streaming_model(self._streaming_text_encoder_builder, self._offload_mode, self._device, self._dtype)
+        return gpu_model(self._build_text_encoder())
 
     def __call__(
         self,
@@ -364,21 +529,22 @@ class PromptEncoder:
         enhance_first_prompt: bool = False,
         enhance_prompt_image: str | None = None,
         enhance_prompt_seed: int = 42,
-        streaming_prefetch_count: int | None = None,
     ) -> list[EmbeddingsProcessorOutput]:
-        """Encode *prompts* through Gemma → embeddings processor, freeing each model after use."""
-        with self._text_encoder_ctx(streaming_prefetch_count) as text_encoder:
+        """Encode *prompts* through Gemma -> embeddings processor, freeing each model after use."""
+        logger.info("Building text encoder from %s", self._gemma_root)
+        with self._text_encoder_ctx() as text_encoder:
             if enhance_first_prompt:
                 prompts = list(prompts)
                 prompts[0] = generate_enhanced_prompt(
                     text_encoder, prompts[0], enhance_prompt_image, seed=enhance_prompt_seed
                 )
-            raw_outputs = [text_encoder.encode(p) for p in prompts]
+            raw_outputs = text_encoder.encode(prompts)
+        logger.info("Text encoder done, building embeddings processor from %s", self._checkpoint_path)
 
-        with gpu_model(
-            self._embeddings_processor_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
-        ) as embeddings_processor:
-            return [embeddings_processor.process_hidden_states(hs, mask) for hs, mask in raw_outputs]
+        with gpu_model(self._build_embeddings_processor()) as embeddings_processor:
+            result = [embeddings_processor.process_hidden_states(hs, mask) for hs, mask in raw_outputs]
+        logger.info("Prompt encoding complete")
+        return result
 
 
 # ---------------------------------------------------------------------------
@@ -408,7 +574,7 @@ class ImageConditioner:
         )
 
     def _build_encoder(self) -> VideoEncoder:
-        return self._encoder_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
+        return self._encoder_builder.build(device=self._device, dtype=self._dtype).eval()
 
     def __call__(self, fn: Callable[[VideoEncoder], T]) -> T:
         """Build video encoder → call *fn(encoder)* → free encoder."""
@@ -432,6 +598,7 @@ class VideoUpsampler:
         device: torch.device,
         registry: Registry | None = None,
     ) -> None:
+        self._upsampler_path = upsampler_path
         self._dtype = dtype
         self._device = device
         self._encoder_builder = Builder(
@@ -448,13 +615,10 @@ class VideoUpsampler:
 
     def __call__(self, latent: torch.Tensor) -> torch.Tensor:
         """Upsample *latent* using video encoder + spatial upsampler, then free both."""
+        logger.info("Building video encoder + spatial upsampler from %s", self._upsampler_path)
         with (
-            gpu_model(
-                self._encoder_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
-            ) as encoder,
-            gpu_model(
-                self._upsampler_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
-            ) as upsampler,
+            gpu_model(self._encoder_builder.build(device=self._device, dtype=self._dtype).eval()) as encoder,
+            gpu_model(self._upsampler_builder.build(device=self._device, dtype=self._dtype).eval()) as upsampler,
         ):
             return upsample_video(latent=latent, video_encoder=encoder, upsampler=upsampler)
 
@@ -475,15 +639,22 @@ class VideoDecoder:
         dtype: torch.dtype,
         device: torch.device,
         registry: Registry | None = None,
+        memory_efficient: bool = True,
+        decoder_builder: BuilderProtocol | None = None,
     ) -> None:
+        self._checkpoint_path = checkpoint_path
         self._dtype = dtype
         self._device = device
-        self._decoder_builder = Builder(
-            model_path=checkpoint_path,
-            model_class_configurator=VideoDecoderConfigurator,
-            model_sd_ops=VAE_DECODER_COMFY_KEYS_FILTER,
-            registry=registry or DummyRegistry(),
-        )
+        if decoder_builder is not None:
+            self._decoder_builder = decoder_builder
+        else:
+            self._decoder_builder = Builder(
+                model_path=checkpoint_path,
+                model_class_configurator=VideoDecoderConfigurator,
+                model_sd_ops=VAE_DECODER_COMFY_KEYS_FILTER,
+                registry=registry or DummyRegistry(),
+                module_ops=(MEMORY_EFFICIENT_DECODE,) if memory_efficient else (),
+            )
 
     def __call__(
         self,
@@ -492,7 +663,8 @@ class VideoDecoder:
         generator: torch.Generator | None = None,
     ) -> Iterator[torch.Tensor]:
         """Decode *latent* to pixel-space video chunks. Decoder freed after exhaustion."""
-        decoder = self._decoder_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
+        logger.info("Building video decoder from %s", self._checkpoint_path)
+        decoder = self._decoder_builder.build(device=self._device, dtype=self._dtype).eval()
         return _cleanup_iter(decoder.decode_video(latent, tiling_config, generator), decoder)
 
 
@@ -511,6 +683,7 @@ class AudioDecoder:
         device: torch.device,
         registry: Registry | None = None,
     ) -> None:
+        self._checkpoint_path = checkpoint_path
         self._dtype = dtype
         self._device = device
         self._decoder_builder = Builder(
@@ -528,13 +701,10 @@ class AudioDecoder:
 
     def __call__(self, latent: torch.Tensor) -> Audio:
         """Decode audio *latent* through VAE decoder + vocoder, then free both."""
+        logger.info("Building audio decoder + vocoder from %s", self._checkpoint_path)
         with (
-            gpu_model(
-                self._decoder_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
-            ) as decoder,
-            gpu_model(
-                self._vocoder_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
-            ) as vocoder,
+            gpu_model(self._decoder_builder.build(device=self._device, dtype=self._dtype).eval()) as decoder,
+            gpu_model(self._vocoder_builder.build(device=self._device, dtype=self._dtype).eval()) as vocoder,
         ):
             return vae_decode_audio(latent, decoder, vocoder)
 
@@ -568,7 +738,5 @@ class AudioConditioner:
 
     def __call__(self, fn: Callable[[torch.nn.Module], T]) -> T:
         """Build audio encoder → call *fn(encoder)* → free encoder."""
-        with gpu_model(
-            self._encoder_builder.build(device=self._device, dtype=self._dtype).to(self._device).eval()
-        ) as encoder:
+        with gpu_model(self._encoder_builder.build(device=self._device, dtype=self._dtype).eval()) as encoder:
             return fn(encoder)

packages/ltx-pipelines/src/ltx_pipelines/utils/color_conversion.py

@@ -0,0 +1,224 @@
+"""Color space conversion utilities for video encoding.
+Provides GPU-accelerated RGB to YUV420 conversion that runs between the
+VAE decoder (which yields float RGB chunks) and ``encode_video``, bypassing
+pyav's CPU-side libswscale conversion. The ``FrameConverter`` also carries
+the codec metadata (pixel format, colour space, colour range) that
+``encode_video`` needs to tag the output stream.
+"""
+
+from __future__ import annotations
+
+import enum
+from collections.abc import Callable
+from dataclasses import dataclass, field
+
+import torch
+
+
+class ColorSpace(enum.Enum):
+    """YUV color space standard."""
+
+    BT_709 = "bt709"
+    BT_2020_NCL = "bt2020ncl"
+
+    @property
+    def av_colorspace(self) -> int:
+        """FFmpeg ``AVCOL_SPC_*`` constant for ``codec_context.colorspace``."""
+        return _AV_COLORSPACE[self]
+
+
+class ColorRange(enum.Enum):
+    """YUV color range."""
+
+    MPEG = "mpeg"
+    JPEG = "jpeg"
+
+    @property
+    def av_color_range(self) -> int:
+        """FFmpeg ``AVCOL_RANGE_*`` constant for ``codec_context.color_range``."""
+        return _AV_COLOR_RANGE[self]
+
+
+class PixelFormat(enum.Enum):
+    """Pixel format for video frames."""
+
+    RGB24 = "rgb24"
+    YUV420P = "yuv420p"
+
+    @property
+    def av_format(self) -> str:
+        """PyAV format string for ``VideoFrame.from_ndarray``."""
+        return self.value
+
+
+_AV_COLORSPACE = {
+    ColorSpace.BT_709: 1,  # AVCOL_SPC_BT709
+    ColorSpace.BT_2020_NCL: 9,  # AVCOL_SPC_BT2020_NCL
+}
+
+_AV_COLOR_RANGE = {
+    ColorRange.MPEG: 1,  # AVCOL_RANGE_MPEG (limited)
+    ColorRange.JPEG: 2,  # AVCOL_RANGE_JPEG (full)
+}
+
+# BT.709 RGB->YUV matrix (row-major: each row produces one of Y, U, V)
+_BT709_MATRIX = torch.tensor(
+    [
+        [0.2126, 0.7152, 0.0722],
+        [-0.1146, -0.3854, 0.5],
+        [0.5, -0.4542, -0.0458],
+    ],
+    dtype=torch.float32,
+)
+
+# BT.2020 NCL RGB->YUV matrix
+_KR_2020 = 0.2627
+_KG_2020 = 0.6780
+_KB_2020 = 0.0593
+_BT2020_MATRIX = torch.tensor(
+    [
+        [_KR_2020, _KG_2020, _KB_2020],
+        [-_KR_2020 / 1.8814, -_KG_2020 / 1.8814, 0.5],
+        [0.5, -_KG_2020 / 1.4746, -_KB_2020 / 1.4746],
+    ],
+    dtype=torch.float32,
+)
+
+_COLOR_SPACE_MATRICES = {
+    ColorSpace.BT_709: _BT709_MATRIX,
+    ColorSpace.BT_2020_NCL: _BT2020_MATRIX,
+}
+
+
+@dataclass(frozen=True)
+class FrameConverter:
+    """Converts ``[*, C, H, W]`` float ``[0, 1]`` frames to uint8.
+    Carries encoding metadata so ``encode_video`` can derive pixel format,
+    color space, and color range from the converter itself.
+    The ``fn_`` callable **may mutate its input** (PyTorch trailing-underscore
+    convention).  Callers that need to keep the original ``frames`` afterwards
+    must pass ``frames.clone()``.  Inside ``encode_video``'s per-chunk
+    generator each chunk is consumed once, so direct passthrough is safe.
+    """
+
+    pixel_format: PixelFormat
+    fn_: Callable[[torch.Tensor], torch.Tensor] = field(repr=False)
+    color_space: ColorSpace | None = None
+    color_range: ColorRange | None = None
+
+    def __call__(self, frames: torch.Tensor) -> torch.Tensor:
+        return self.fn_(frames)
+
+
+def rgb_to_yuv(image: torch.Tensor, color_space: ColorSpace) -> torch.Tensor:
+    """Convert an RGB image to YUV.
+    The image data is assumed to be in the range of ``[0, 1]``.
+    Uses a single matrix multiply for better memory locality.
+    Args:
+        image: RGB image with shape ``(*, 3, H, W)``.
+        color_space: Color space standard for the conversion matrix.
+    Returns:
+        YUV image with shape ``(*, 3, H, W)``.
+    """
+    if len(image.shape) < 3 or image.shape[-3] != 3:
+        raise ValueError(f"Input size must have a shape of (*, 3, H, W). Got {image.shape}")
+
+    mat = _COLOR_SPACE_MATRICES[color_space].to(device=image.device, dtype=image.dtype)
+    # [*, 3, H, W] -> [*, H, W, 3] @ [3, 3]^T -> [*, H, W, 3] -> [*, 3, H, W]
+    pixels = image.movedim(-3, -1)  # [*, H, W, 3]
+    yuv = pixels @ mat.T  # [*, H, W, 3]
+    return yuv.movedim(-1, -3)  # [*, 3, H, W]
+
+
+def apply_color_range_(y: torch.Tensor, uv: torch.Tensor, color_range: ColorRange) -> tuple[torch.Tensor, torch.Tensor]:
+    """Scale Y and UV planes to the specified color range, in-place.
+    Args:
+        y: Luma plane in ``[0, 1]``.
+        uv: Chroma planes centered at 0.
+        color_range: Target color range.
+    Returns:
+        Scaled ``(Y, UV)`` tensors (modified in-place).
+    """
+    if color_range == ColorRange.MPEG:
+        y.mul_(219).add_(16)
+        uv.mul_(224).add_(128)
+    elif color_range == ColorRange.JPEG:
+        y.mul_(255)
+        uv.add_(0.5).mul_(255)
+    else:
+        raise ValueError(f"Unsupported color range: {color_range}")
+    return y, uv
+
+
+def rgb_to_yuv420(
+    image: torch.Tensor, color_space: ColorSpace, color_range: ColorRange
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Convert an RGB image to YUV 4:2:0 with chroma subsampling.
+    Chroma is subsampled by averaging 2x2 pixel blocks (chroma siting
+    ``(128, 128)``).
+    Args:
+        image: RGB image with shape ``(*, 3, H, W)`` in ``[0, 1]``.
+            H and W must be divisible by 2.
+        color_space: Color space standard.
+        color_range: Color range for the output.
+    Returns:
+        ``(Y, UV)`` where Y has shape ``(*, 1, H, W)`` and UV has shape
+        ``(*, 2, H//2, W//2)``.
+    """
+    if len(image.shape) < 3 or image.shape[-3] != 3:
+        raise ValueError(f"Input size must have a shape of (*, 3, H, W). Got {image.shape}")
+    if image.shape[-2] % 2 != 0 or image.shape[-1] % 2 != 0:
+        raise ValueError(f"Input H and W must be divisible by 2. Got {image.shape}")
+
+    yuv = rgb_to_yuv(image, color_space)
+    y = yuv[..., :1, :, :]
+    # Subsample chroma: average 2x2 blocks via avg_pool2d (contiguous, fused kernel)
+    uv_full = yuv[..., 1:3, :, :].contiguous()
+    # Flatten leading dims for avg_pool2d which expects [N, C, H, W]
+    lead = uv_full.shape[:-3]
+    uv_flat = uv_full.reshape(-1, 2, uv_full.shape[-2], uv_full.shape[-1])
+    uv = torch.nn.functional.avg_pool2d(uv_flat, kernel_size=2, stride=2)
+    uv = uv.reshape(*lead, 2, uv.shape[-2], uv.shape[-1])
+
+    return apply_color_range_(y, uv, color_range)
+
+
+def pack_i420(y: torch.Tensor, uv: torch.Tensor) -> torch.Tensor:
+    """Pack Y and UV planes into I420 layout for pyav.
+    I420 packs the three planes into a single 2D array of height ``H * 3 // 2``
+    and width ``W``.  The Y plane occupies the first ``H`` rows.  The UV tensor
+    ``(*, 2, H//2, W//2)`` is reshaped to ``(*, H//2, W)`` -- U rows packed
+    two-by-two followed by V rows packed two-by-two -- and appended below.
+    Args:
+        y: Luma with shape ``(*, 1, H, W)``.
+        uv: Chroma with shape ``(*, 2, H//2, W//2)``.
+    Returns:
+        Packed tensor with shape ``(*, H*3//2, W)`` uint8.
+    """
+    y_plane = y[..., 0, :, :]  # [*, H, W]
+    uv_packed = uv.reshape(*uv.shape[:-3], uv.shape[-2], uv.shape[-1] * 2)  # [*, H//2, W]
+    packed = torch.cat([y_plane, uv_packed], dim=-2)  # [*, H*3//2, W]
+    return packed.clamp_(0, 255).to(torch.uint8)
+
+
+def _rgb_uint8_fn_(frames: torch.Tensor) -> torch.Tensor:
+    """In-place: mutates ``frames`` via ``clamp_`` + ``mul_``, returns a uint8 view."""
+    return frames.clamp_(0.0, 1.0).mul_(255.0).to(torch.uint8).movedim(-3, -1)
+
+
+rgb_uint8_converter_ = FrameConverter(pixel_format=PixelFormat.RGB24, fn_=_rgb_uint8_fn_)
+"""``(*, 3, H, W)`` float ``[0, 1]`` to ``(*, H, W, 3)`` uint8.  Mutates input."""
+
+
+def _yuv420p_bt709_fn_(frames: torch.Tensor) -> torch.Tensor:
+    y, uv = rgb_to_yuv420(frames, ColorSpace.BT_709, ColorRange.MPEG)
+    return pack_i420(y, uv)
+
+
+yuv420p_bt709_converter_ = FrameConverter(
+    pixel_format=PixelFormat.YUV420P,
+    fn_=_yuv420p_bt709_fn_,
+    color_space=ColorSpace.BT_709,
+    color_range=ColorRange.MPEG,
+)
+"""``(*, 3, H, W)`` float ``[0, 1]`` to ``(*, H*3//2, W)`` uint8 YUV420p BT.709 MPEG."""

packages/ltx-pipelines/src/ltx_pipelines/utils/denoisers.py

@@ -20,6 +20,7 @@ from ltx_core.guidance.perturbations import (
 from ltx_core.model.transformer import X0Model
 from ltx_core.types import LatentState
 from ltx_pipelines.utils.helpers import modality_from_latent_state
+from ltx_pipelines.utils.types import DenoisedLatentResult
 
 _POSITIVE_ONLY_GUIDER = MultiModalGuider(
     params=MultiModalGuiderParams(cfg_scale=1.0, stg_scale=0.0, modality_scale=1.0),
@@ -53,7 +54,7 @@ def _repeat_state(state: LatentState, n: int) -> LatentState:
     )
 
 
-def _guided_denoise(  # noqa: PLR0913
+def _guided_denoise(  # noqa: PLR0913,PLR0915
     transformer: X0Model,
     video_state: LatentState | None,
     audio_state: LatentState | None,
@@ -66,7 +67,8 @@ def _guided_denoise(  # noqa: PLR0913
     last_denoised_video: torch.Tensor | None,
     last_denoised_audio: torch.Tensor | None,
     step_index: int,
-) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+    force_uncond_pass: bool = False,
+) -> tuple[DenoisedLatentResult | None, DenoisedLatentResult | None]:
     """Core guided denoising — batches all guidance passes into one transformer call.
     Collects per-pass contexts first, then builds a single batched Modality
     per present modality via :func:`modality_from_latent_state`.  When wrapped
@@ -80,7 +82,9 @@ def _guided_denoise(  # noqa: PLR0913
     a_skip = audio_guider.should_skip_step(step_index)
 
     if v_skip and a_skip:
-        return last_denoised_video, last_denoised_audio
+        video_result = DenoisedLatentResult.result_or_none(denoised=last_denoised_video)
+        audio_result = DenoisedLatentResult.result_or_none(denoised=last_denoised_audio)
+        return video_result, audio_result
 
     if video_state is not None and v_context is None:
         raise ValueError("v_context is required when video_state is provided")
@@ -91,10 +95,12 @@ def _guided_denoise(  # noqa: PLR0913
     _pass = tuple[str, torch.Tensor | None, torch.Tensor | None, PerturbationConfig]
     passes: list[_pass] = [("cond", v_context, a_context, PerturbationConfig.empty())]
 
-    if video_guider.do_unconditional_generation() or audio_guider.do_unconditional_generation():
-        if video_guider.do_unconditional_generation() and video_guider.negative_context is None:
+    v_needs_neg = video_guider.do_unconditional_generation() or (force_uncond_pass and video_state is not None)
+    a_needs_neg = audio_guider.do_unconditional_generation() or (force_uncond_pass and audio_state is not None)
+    if v_needs_neg or a_needs_neg:
+        if v_needs_neg and video_guider.negative_context is None:
             raise ValueError("Negative context is required for unconditioned denoising")
-        if audio_guider.do_unconditional_generation() and audio_guider.negative_context is None:
+        if a_needs_neg and audio_guider.negative_context is None:
             raise ValueError("Negative context is required for unconditioned denoising")
         v_neg = video_guider.negative_context if video_guider.negative_context is not None else v_context
         a_neg = audio_guider.negative_context if audio_guider.negative_context is not None else a_context
@@ -132,6 +138,8 @@ def _guided_denoise(  # noqa: PLR0913
     ptb_configs = [ptb for _, _, _, ptb in passes]
     n = len(passes)
 
+    orig_b = (video_state or audio_state).latent.shape[0]
+
     def _batched_sigma(state: LatentState) -> torch.Tensor:
         """Expand scalar sigma to (n * B,) matching the repeated state."""
         return sigma.expand(state.latent.shape[0] * n)
@@ -156,8 +164,16 @@ def _guided_denoise(  # noqa: PLR0913
             enabled=not a_skip,
         )
 
+    # Replicate each pass's PerturbationConfig to all `orig_b` samples it
+    # carries, so `BatchedPerturbationConfig.mask_like` returns a per-sample
+    # mask (length n*orig_b) instead of a per-pass mask (length n). Without
+    # this expansion the mask is broadcast against a (n*orig_b, T, D) tensor
+    # and the multiplication fails with a batch-dim mismatch whenever
+    # `orig_b > 1` (e.g. multi-prompt benchmark panels).
+    batched_ptb_configs = [ptb for ptb in ptb_configs for _ in range(orig_b)]
+
     all_v, all_a = transformer(
-        video=batched_video, audio=batched_audio, perturbations=BatchedPerturbationConfig(ptb_configs)
+        video=batched_video, audio=batched_audio, perturbations=BatchedPerturbationConfig(batched_ptb_configs)
     )
 
     # Split results back and combine via guiders.
@@ -166,13 +182,22 @@ def _guided_denoise(  # noqa: PLR0913
     r = dict(zip(pass_names, zip(splits_v, splits_a, strict=True), strict=True))
 
     cond_v, cond_a = r["cond"]
+    cond_v = cond_v if isinstance(cond_v, torch.Tensor) else torch.tensor(cond_v)
+    cond_a = cond_a if isinstance(cond_a, torch.Tensor) else torch.tensor(cond_a)
     uncond_v, uncond_a = r.get("uncond", (0.0, 0.0))
     ptb_v, ptb_a = r.get("ptb", (0.0, 0.0))
     mod_v, mod_a = r.get("mod", (0.0, 0.0))
 
     denoised_video = last_denoised_video if v_skip else video_guider.calculate(cond_v, uncond_v, ptb_v, mod_v)
     denoised_audio = last_denoised_audio if a_skip else audio_guider.calculate(cond_a, uncond_a, ptb_a, mod_a)
-    return denoised_video, denoised_audio
+    return (
+        DenoisedLatentResult.result_or_none(
+            denoised=denoised_video, uncond=uncond_v, cond=cond_v, ptb=ptb_v, mod=mod_v
+        ),
+        DenoisedLatentResult.result_or_none(
+            denoised=denoised_audio, uncond=uncond_a, cond=cond_a, ptb=ptb_a, mod=mod_a
+        ),
+    )
 
 
 class SimpleDenoiser:
@@ -195,11 +220,15 @@ class SimpleDenoiser:
         audio_state: LatentState | None,
         sigmas: torch.Tensor,
         step_index: int,
-    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+    ) -> tuple[DenoisedLatentResult | None, DenoisedLatentResult | None]:
         sigma = sigmas[step_index]
         pos_video = modality_from_latent_state(video_state, self.v_context, sigma) if video_state is not None else None
         pos_audio = modality_from_latent_state(audio_state, self.a_context, sigma) if audio_state is not None else None
-        return transformer(video=pos_video, audio=pos_audio, perturbations=None)
+        denoised_video, denoised_audio = transformer(video=pos_video, audio=pos_audio, perturbations=None)
+        return (
+            DenoisedLatentResult.result_or_none(denoised=denoised_video),
+            DenoisedLatentResult.result_or_none(denoised=denoised_audio),
+        )
 
 
 class GuidedDenoiser:
@@ -214,11 +243,13 @@ class GuidedDenoiser:
         a_context: torch.Tensor | None,
         video_guider: MultiModalGuider | None = None,
         audio_guider: MultiModalGuider | None = None,
+        force_uncond_pass: bool = False,
     ) -> None:
         self.v_context = v_context
         self.a_context = a_context
         self.video_guider = video_guider
         self.audio_guider = audio_guider
+        self.force_uncond_pass = force_uncond_pass
         self._last_denoised_video: torch.Tensor | None = None
         self._last_denoised_audio: torch.Tensor | None = None
 
@@ -229,8 +260,8 @@ class GuidedDenoiser:
         audio_state: LatentState | None,
         sigmas: torch.Tensor,
         step_index: int,
-    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
-        denoised_video, denoised_audio = _guided_denoise(
+    ) -> tuple[DenoisedLatentResult | None, DenoisedLatentResult | None]:
+        guided_denoise_result_v, guided_denoise_result_a = _guided_denoise(
             transformer=transformer,
             video_state=video_state,
             audio_state=audio_state,
@@ -242,10 +273,11 @@ class GuidedDenoiser:
             last_denoised_video=self._last_denoised_video,
             last_denoised_audio=self._last_denoised_audio,
             step_index=step_index,
+            force_uncond_pass=self.force_uncond_pass,
         )
-        self._last_denoised_video = denoised_video
-        self._last_denoised_audio = denoised_audio
-        return denoised_video, denoised_audio
+        self._last_denoised_video = guided_denoise_result_v.denoised
+        self._last_denoised_audio = guided_denoise_result_a.denoised
+        return guided_denoise_result_v, guided_denoise_result_a
 
 
 class FactoryGuidedDenoiser:
@@ -257,11 +289,13 @@ class FactoryGuidedDenoiser:
         a_context: torch.Tensor | None,
         video_guider_factory: MultiModalGuiderFactory | None = None,
         audio_guider_factory: MultiModalGuiderFactory | None = None,
+        force_uncond_pass: bool = False,
     ) -> None:
         self.v_context = v_context
         self.a_context = a_context
         self.video_guider_factory = video_guider_factory
         self.audio_guider_factory = audio_guider_factory
+        self.force_uncond_pass = force_uncond_pass
         self._last_denoised_video: torch.Tensor | None = None
         self._last_denoised_audio: torch.Tensor | None = None
         self._sigma_vals_cached: list[float] | None = None
@@ -273,7 +307,7 @@ class FactoryGuidedDenoiser:
         audio_state: LatentState | None,
         sigmas: torch.Tensor,
         step_index: int,
-    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+    ) -> tuple[DenoisedLatentResult | None, DenoisedLatentResult | None]:
         if self._sigma_vals_cached is None:
             self._sigma_vals_cached = sigmas.detach().cpu().tolist()
         sigma_val = self._sigma_vals_cached[step_index]
@@ -287,7 +321,7 @@ class FactoryGuidedDenoiser:
             else None
         )
 
-        denoised_video, denoised_audio = _guided_denoise(
+        guided_denoise_result_v, guided_denoise_result_a = _guided_denoise(
             transformer=transformer,
             video_state=video_state,
             audio_state=audio_state,
@@ -299,7 +333,8 @@ class FactoryGuidedDenoiser:
             last_denoised_video=self._last_denoised_video,
             last_denoised_audio=self._last_denoised_audio,
             step_index=step_index,
+            force_uncond_pass=self.force_uncond_pass,
         )
-        self._last_denoised_video = denoised_video
-        self._last_denoised_audio = denoised_audio
-        return denoised_video, denoised_audio
+        self._last_denoised_video = guided_denoise_result_v.denoised
+        self._last_denoised_audio = guided_denoise_result_a.denoised
+        return guided_denoise_result_v, guided_denoise_result_a

packages/ltx-pipelines/src/ltx_pipelines/utils/helpers.py

@@ -37,6 +37,11 @@ def cleanup_memory() -> None:
     gc.collect()
     torch.cuda.empty_cache()
     torch.cuda.synchronize()
+    try:
+        if hasattr(torch._C, "_host_emptyCache"):
+            torch._C._host_emptyCache()
+    except Exception:
+        logging.warning("Host empty cache cleanup failed; ignoring.", exc_info=True)
 
 
 def _conform_latent_length(latent: torch.Tensor, expected_frames_count: int) -> torch.Tensor:

packages/ltx-pipelines/src/ltx_pipelines/utils/media_io.py

@@ -1,23 +1,37 @@
+import enum
 import logging
 import math
+import threading
 from collections.abc import Generator, Iterator
 from fractions import Fraction
 from io import BytesIO
+from pathlib import Path
+from queue import Queue
 
 import av
 import numpy as np
+import OpenImageIO
 import torch
 from einops import rearrange
 from PIL import Image
 from torch._prims_common import DeviceLikeType
 from tqdm import tqdm
 
+from ltx_core.hdr import LogC3
 from ltx_core.types import Audio, VideoPixelShape
+from ltx_pipelines.utils.color_conversion import FrameConverter, PixelFormat, yuv420p_bt709_converter_
 from ltx_pipelines.utils.constants import DEFAULT_IMAGE_CRF
 
 logger = logging.getLogger(__name__)
 
 
+class ResizeMode(enum.Enum):
+    """How to fit a conditioning video to the target resolution."""
+
+    CENTER_CROP = "center_crop"
+    REFLECT_PAD = "reflect_pad"
+
+
 def resize_aspect_ratio_preserving(image: torch.Tensor, long_side: int) -> torch.Tensor:
     """
     Resize image preserving aspect ratio (filling target long side).
@@ -75,8 +89,18 @@ def resize_and_center_crop(tensor: torch.Tensor, height: int, width: int) -> tor
     return tensor
 
 
-def normalize_latent(latent: torch.Tensor, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
-    return (latent / 127.5 - 1.0).to(device=device, dtype=dtype)
+def normalize_images(images: torch.Tensor, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
+    return (images / 127.5 - 1.0).to(device=device, dtype=dtype)
+
+
+def to_vae_range(x: torch.Tensor) -> torch.Tensor:
+    """Map [0, 1] to [-1, 1] (VAE input convention)."""
+    return torch.clamp(x, 0.0, 1.0) * 2.0 - 1.0
+
+
+def from_vae_range(z: torch.Tensor) -> torch.Tensor:
+    """Map [-1, 1] (VAE output convention) to [0, 1]."""
+    return torch.clamp((z + 1.0) / 2.0, 0.0, 1.0)
 
 
 def load_image_and_preprocess(
@@ -95,7 +119,7 @@ def load_image_and_preprocess(
     image = preprocess(image=image, crf=crf)
     image = torch.tensor(image, dtype=torch.float32, device=device)
     image = resize_and_center_crop(image, height, width)
-    image = normalize_latent(image, device, dtype)
+    image = normalize_images(image, device, dtype)
     return image
 
 
@@ -116,14 +140,118 @@ def video_preprocess(
     Returns:
         Tensor of shape (1, C, F, height, width) with values in [-1, 1].
     """
-    result = None
+    result: torch.Tensor | None = None
     for f in frames:
         frame = resize_and_center_crop(f.to(torch.float32), height, width)
-        frame = normalize_latent(frame, device, dtype)
+        frame = normalize_images(frame, device, dtype)
         result = frame if result is None else torch.cat([result, frame], dim=2)
+    if result is None:
+        raise ValueError("video_preprocess received an empty frame generator; no frames were decoded from the source.")
     return result
 
 
+def align_resolution(
+    width: int,
+    height: int,
+    resize_mode: ResizeMode,
+    divisor: int = 64,
+) -> tuple[int, int, int, int]:
+    """Compute aligned generation dimensions and crop-back size.
+    Args:
+        width: Source video width (need not be aligned).
+        height: Source video height (need not be aligned).
+        resize_mode: CENTER_CROP rounds down; REFLECT_PAD rounds up.
+        divisor: Alignment divisor (default 64 for two-stage pipelines).
+    Returns:
+        ``(gen_width, gen_height, crop_width, crop_height)`` where
+        ``gen_*`` are multiples of *divisor* and ``crop_*`` are the
+        original dimensions to trim back to after decoding.  When no
+        cropping is needed ``crop_*`` equals ``gen_*``.
+    """
+    if resize_mode is ResizeMode.REFLECT_PAD:
+        gen_w = ((width + divisor - 1) // divisor) * divisor
+        gen_h = ((height + divisor - 1) // divisor) * divisor
+    else:
+        gen_w = (width // divisor) * divisor
+        gen_h = (height // divisor) * divisor
+
+    crop_w = width if gen_w != width else gen_w
+    crop_h = height if gen_h != height else gen_h
+    return gen_w, gen_h, crop_w, crop_h
+
+
+def resize_and_reflect_pad(tensor: torch.Tensor, height: int, width: int) -> torch.Tensor:
+    """Resize tensor to fit within target, then reflect-pad to exact dimensions.
+    Unlike resize_and_center_crop which stretches and crops, this preserves the
+    original aspect ratio and pads the shorter dimension with reflected pixels.
+    When the target is already >= the source in both dimensions, interpolation
+    is skipped entirely to preserve original pixels.
+    Args:
+        tensor: Input with shape (H, W, C) or (F, H, W, C)
+        height: Target height
+        width: Target width
+    Returns:
+        Tensor with shape (1, C, 1, height, width) for 3D or (1, C, F, height, width) for 4D
+    """
+    if tensor.ndim == 3:
+        tensor = rearrange(tensor, "h w c -> 1 c h w")
+    elif tensor.ndim == 4:
+        tensor = rearrange(tensor, "f h w c -> f c h w")
+    else:
+        raise ValueError(f"Expected input with 3 or 4 dimensions; got shape {tensor.shape}.")
+
+    _, _, src_h, src_w = tensor.shape
+
+    if height >= src_h and width >= src_w:
+        new_h, new_w = src_h, src_w
+    else:
+        scale = min(height / src_h, width / src_w)
+        new_h = round(src_h * scale)
+        new_w = round(src_w * scale)
+        tensor = torch.nn.functional.interpolate(tensor, size=(new_h, new_w), mode="bilinear", align_corners=False)
+
+    pad_bottom = height - new_h
+    pad_right = width - new_w
+    if pad_bottom > 0 or pad_right > 0:
+        pad_mode = "reflect" if pad_bottom < new_h and pad_right < new_w else "replicate"
+        tensor = torch.nn.functional.pad(tensor, (0, pad_right, 0, pad_bottom), mode=pad_mode)
+
+    tensor = rearrange(tensor, "f c h w -> 1 c f h w")
+    return tensor
+
+
+def load_video_conditioning_hdr(
+    video_path: str,
+    height: int,
+    width: int,
+    frame_cap: int,
+    dtype: torch.dtype,
+    device: torch.device,
+    hdr_transform: str = "logc3",
+    resize_mode: ResizeMode = ResizeMode.CENTER_CROP,
+) -> Iterator[torch.Tensor]:
+    """Load a video and yield preprocessed frames for HDR IC-LoRA conditioning.
+    Decodes through the standard path and applies the LDR compression that
+    matches training. Callers are responsible for providing Rec.709 SDR
+    input — the HDR IC-LoRA was trained on that color space.
+    Args:
+        hdr_transform: LDR-compression name (currently only ``logc3``).
+        resize_mode: How to fit the video to the target resolution.
+    Yields:
+        Per-frame tensors of shape ``(1, C, 1, height, width)``.
+    """
+    if hdr_transform != "logc3":
+        raise ValueError(f"Unsupported HDR transform: {hdr_transform}")
+
+    resize_fn = resize_and_reflect_pad if resize_mode is ResizeMode.REFLECT_PAD else resize_and_center_crop
+
+    for f in decode_video_by_frame(path=video_path, frame_cap=frame_cap, device=device):
+        frame = resize_fn(f.to(torch.float32), height, width)
+        ldr = (frame / 255.0).clamp(0.0, 1.0)
+        compressed = LogC3().compress_ldr(ldr)
+        yield to_vae_range(compressed).to(device=device, dtype=dtype)
+
+
 def decode_image(image_path: str) -> np.ndarray:
     image = Image.open(image_path)
     np_array = np.array(image)[..., :3]
@@ -202,47 +330,138 @@ def encode_video(
     audio: Audio | None,
     output_path: str,
     video_chunks_number: int,
+    frame_converter: FrameConverter = yuv420p_bt709_converter_,
+    crf: int = 19,
+    preset: str = "veryfast",
+    thread_count: int = 0,
 ) -> None:
     if isinstance(video, torch.Tensor):
         video = iter([video])
 
-    first_chunk = next(video)
+    def convert(chunk: torch.Tensor) -> torch.Tensor:
+        return frame_converter(chunk.movedim(-1, -3))
 
-    _, height, width, _ = first_chunk.shape
+    first_chunk = convert(next(video))
+
+    if frame_converter.pixel_format == PixelFormat.RGB24:
+        height, width = first_chunk.shape[-3], first_chunk.shape[-2]
+    else:
+        height = first_chunk.shape[-2] * 2 // 3
+        width = first_chunk.shape[-1]
 
     container = av.open(output_path, mode="w")
-    stream = container.add_stream("libx264", rate=int(fps))
-    stream.width = width
-    stream.height = height
-    stream.pix_fmt = "yuv420p"
+    success = False
+    try:
+        stream = container.add_stream("libx264", rate=int(fps), options={"crf": str(crf), "preset": preset})
+        stream.width = width
+        stream.height = height
+        stream.pix_fmt = "yuv420p"
+        stream.codec_context.thread_count = thread_count
+        stream.codec_context.thread_type = "FRAME"
+        if frame_converter.color_space is not None:
+            stream.codec_context.colorspace = frame_converter.color_space.av_colorspace
+        if frame_converter.color_range is not None:
+            stream.codec_context.color_range = frame_converter.color_range.av_color_range
 
-    if audio is not None:
-        audio_stream = _prepare_audio_stream(container, audio.sampling_rate)
+        if audio is not None:
+            audio_stream = _prepare_audio_stream(container, audio.sampling_rate)
 
-    def all_tiles(
-        first_chunk: torch.Tensor, tiles_generator: Generator[tuple[torch.Tensor, int], None, None]
-    ) -> Generator[tuple[torch.Tensor, int], None, None]:
-        yield first_chunk
-        yield from tiles_generator
+        av_format = frame_converter.pixel_format.av_format
 
-    for video_chunk in tqdm(all_tiles(first_chunk, video), total=video_chunks_number):
-        video_chunk_cpu = video_chunk.to("cpu").numpy()
-        for frame_array in video_chunk_cpu:
-            frame = av.VideoFrame.from_ndarray(frame_array, format="rgb24")
-            for packet in stream.encode(frame):
-                container.mux(packet)
+        def cpu_chunks() -> Generator[np.ndarray, None, None]:
+            yield first_chunk.to("cpu").numpy()
+            for chunk in video:
+                yield convert(chunk).to("cpu").numpy()
 
-    # Flush encoder
-    for packet in stream.encode():
-        container.mux(packet)
+        _encode_chunks_threaded(
+            container=container,
+            stream=stream,
+            av_format=av_format,
+            chunks=cpu_chunks(),
+            progress_total=video_chunks_number,
+        )
 
-    if audio is not None:
-        _write_audio(container, audio_stream, audio)
-
-    container.close()
+        if audio is not None:
+            _write_audio(container, audio_stream, audio)
+        success = True
+    finally:
+        container.close()
+        if not success:
+            Path(output_path).unlink(missing_ok=True)
     logger.info(f"Video saved to {output_path}")
 
 
+def encode_audio(audio: Audio, output_path: str) -> None:
+    """Save an audio waveform as a 16-bit PCM ``.wav`` file at the source sampling rate.
+    Reuses :func:`_write_audio` (the same muxing path used by :func:`encode_video`);
+    the only difference is a PCM (``pcm_s16le``) stream in a WAV container instead of
+    the AAC stream used for muxed video.
+    """
+    container = av.open(output_path, mode="w")
+    audio_stream = container.add_stream("pcm_s16le", rate=audio.sampling_rate)
+    audio_stream.codec_context.sample_rate = audio.sampling_rate
+    audio_stream.codec_context.layout = "stereo"
+    audio_stream.codec_context.time_base = Fraction(1, audio.sampling_rate)
+    try:
+        _write_audio(container, audio_stream, audio)
+    finally:
+        container.close()
+    logger.info(f"Audio saved to {output_path}")
+
+
+def _encode_chunks_threaded(
+    container: av.container.Container,
+    stream: av.video.stream.VideoStream,
+    av_format: str,
+    chunks: Iterator[np.ndarray],
+    progress_total: int,
+) -> None:
+    """Run libx264 frame.encode + container.mux on a background thread while
+    the caller produces numpy chunks on the current thread. The 1-slot queue
+    lets the producer get one chunk ahead (so the next VAE/gather chunk
+    overlaps with libx264 encoding the previous chunk) without buffering more
+    than one chunk in CPU memory.
+    """
+    chunk_queue: Queue[np.ndarray | None] = Queue(maxsize=1)
+    encoder_error: list[BaseException] = []
+
+    def encoder_worker() -> None:
+        error: BaseException | None = None
+        while True:
+            arr = chunk_queue.get()
+            if arr is None:
+                break
+            if error is not None:
+                continue
+            try:
+                for frame_array in arr:
+                    frame = av.VideoFrame.from_ndarray(frame_array, format=av_format)
+                    for packet in stream.encode(frame):
+                        container.mux(packet)
+            except Exception as e:
+                error = e
+        if error is None:
+            try:
+                for packet in stream.encode():
+                    container.mux(packet)
+            except Exception as e:
+                error = e
+        if error is not None:
+            encoder_error.append(error)
+
+    encoder_thread = threading.Thread(target=encoder_worker, name="h264-encoder")
+    encoder_thread.start()
+    try:
+        for arr in tqdm(chunks, total=progress_total):
+            chunk_queue.put(arr)
+    finally:
+        chunk_queue.put(None)
+        encoder_thread.join()
+
+    if encoder_error:
+        raise encoder_error[0]
+
+
 _INT_FORMAT_MAX: dict[str, float] = {
     "u8": 128.0,
     "u8p": 128.0,
@@ -481,3 +700,85 @@ def preprocess(image: np.array, crf: float = DEFAULT_IMAGE_CRF) -> np.array:
     with BytesIO(video_bytes) as video_file:
         image_array = decode_single_frame(video_file)
     return image_array
+
+
+def save_exr_tensor(tensor: torch.Tensor, file_path: str | Path, half: bool = False) -> None:
+    """Save a single tensor frame as EXR with linear sRGB colorspace metadata.
+    Args:
+        tensor: ``[H, W, C]`` or ``[C, H, W]`` float tensor.
+        file_path: Output path (e.g. ``frame_0000.exr``).
+        half: Force float16 output with ZIP compression.
+    """
+    if tensor.dim() == 3 and tensor.shape[0] == 3:
+        tensor = tensor.permute(1, 2, 0)
+    use_half = half or tensor.dtype in (torch.float16, torch.half)
+    img_np = np.ascontiguousarray(tensor.cpu().numpy().astype(np.float32))
+    file_path = str(file_path)
+
+    h, w = img_np.shape[:2]
+    fmt = OpenImageIO.HALF if use_half else OpenImageIO.FLOAT
+    spec = OpenImageIO.ImageSpec(w, h, 3, fmt)
+    spec.channelnames = ("R", "G", "B")
+    spec.attribute("compression", "zip")
+    spec.attribute("chromaticities", "float[8]", (0.64, 0.33, 0.30, 0.60, 0.15, 0.06, 0.3127, 0.3290))
+    spec.attribute("colorSpace", "sRGB")
+
+    out = OpenImageIO.ImageOutput.create(file_path)
+    if out is None:
+        raise RuntimeError(
+            f"Failed to create EXR writer for '{file_path}'. Ensure OpenImageIO is built with OpenEXR support."
+        )
+    try:
+        if not out.open(file_path, spec):
+            raise RuntimeError(f"Failed to open EXR file '{file_path}': {out.geterror()}")
+        if not out.write_image(img_np):
+            raise RuntimeError(f"Failed to write EXR image '{file_path}': {out.geterror()}")
+    finally:
+        out.close()
+
+
+def _linear_to_srgb(x: np.ndarray) -> np.ndarray:
+    """Linear -> sRGB OETF per IEC 61966-2-1. Input assumed in [0, 1]."""
+    x = np.clip(x, 0.0, 1.0)
+    return np.where(x <= 0.0031308, x * 12.92, 1.055 * np.power(x, 1.0 / 2.4) - 0.055)
+
+
+def encode_exr_sequence_to_mp4(exr_dir: Path, output_mp4: Path, frame_rate: float) -> None:
+    """Convert a linear EXR frame sequence to sRGB and encode to H.264 .mp4 via PyAV.
+    Exposure is fixed at EV=0 (no gain). Each EXR frame is clamped to [0, 1],
+    passed through the sRGB OETF, quantised to 8-bit BGR, and fed to a libx264
+    stream (crf 18, yuv420p). ``frame_rate`` is the original source video's
+    frame rate so playback matches the input timing.
+    """
+    import os  # noqa: PLC0415
+
+    os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"
+    import cv2  # noqa: PLC0415
+
+    exr_files = sorted(exr_dir.glob("frame_*.exr"))
+    if not exr_files:
+        raise FileNotFoundError(f"No EXR frames found in {exr_dir}")
+
+    container = av.open(str(output_mp4), mode="w")
+    stream = container.add_stream("libx264", rate=Fraction(frame_rate).limit_denominator(1000))
+    stream.pix_fmt = "yuv420p"
+    stream.options = {"crf": "18", "movflags": "+faststart"}
+
+    try:
+        for i, exr_path in enumerate(exr_files):
+            hdr = cv2.imread(str(exr_path), cv2.IMREAD_UNCHANGED).astype(np.float32)
+            sdr = _linear_to_srgb(np.maximum(hdr, 0.0))
+            bgr8 = (sdr * 255.0 + 0.5).astype(np.uint8)
+
+            if i == 0:
+                stream.height = bgr8.shape[0]
+                stream.width = bgr8.shape[1]
+
+            frame = av.VideoFrame.from_ndarray(bgr8, format="bgr24")
+            for packet in stream.encode(frame):
+                container.mux(packet)
+
+        for packet in stream.encode():
+            container.mux(packet)
+    finally:
+        container.close()

packages/ltx-pipelines/src/ltx_pipelines/utils/quantization_factory.py

@@ -0,0 +1,36 @@
+"""User-facing quantization-policy dispatch.
+``ltx-core`` exposes one ``build_policy`` factory per backend. This module
+provides the user-facing string-keyed dispatch used by CLI args and pipeline
+defaults — keeping the enum out of ``ltx-core`` so adding/removing backends is
+a single-file change here.
+"""
+
+from enum import Enum
+
+from typing_extensions import assert_never
+
+from ltx_core.quantization import QuantizationPolicy
+from ltx_core.quantization.fp8_cast import build_policy as _build_fp8_cast_policy
+from ltx_core.quantization.fp8_scaled_mm import build_policy as _build_fp8_scaled_mm_policy
+
+
+class QuantizationKind(str, Enum):
+    FP8_CAST = "fp8-cast"
+    FP8_SCALED_MM = "fp8-scaled-mm"
+
+    def to_policy(self, checkpoint_path: str | None = None) -> QuantizationPolicy:
+        """Build the :class:`QuantizationPolicy` for this kind.
+        ``checkpoint_path`` is required for both backends: ``FP8_SCALED_MM``
+        uses it to discover the layer set from ``.weight_scale`` tensors,
+        and ``FP8_CAST`` uses it to fold any prequant scales into the fp8
+        weight at load time.
+        """
+        if checkpoint_path is None:
+            raise ValueError(f"{self.value} quantization requires checkpoint_path.")
+        match self:
+            case QuantizationKind.FP8_CAST:
+                return _build_fp8_cast_policy(checkpoint_path)
+            case QuantizationKind.FP8_SCALED_MM:
+                return _build_fp8_scaled_mm_policy(checkpoint_path)
+            case _:
+                assert_never(self)