[feat] LTX-2.3 transformer support (config-gated extension of LTX-2)

[alexzms]

Summary

Adds LTX-2.3 transformer support on top of the existing LTX-2 implementation as a config-gated extension — every new field defaults off, so with defaults the model reproduces LTX-2.0 exactly.

Changes

• DiT 2.3 arch (models/dits/ltx2.py, configs/models/dits/ltx2.py): per-head gated attention (apply_gated_attention → to_gate_logits, distinct from upstream's to_gate_compress VSA gate); cross-attention AdaLN (cross_attention_adaln → 6→9-row scale_shift table + prompt-side modulation, fed by a sigma/prompt_timestep path); STG block 28 (stg_block_idx) + per-sample STG mask; StageAwareRMSNorm refine fast-path (defensive QuACK import → torch rms_norm fallback).
• Gemma text connector (models/encoders/gemma.py, config): caption_proj_before_connector, gated connector, per-modality feature extractor.
• Sigma threading (models/dits/ltx2.py forward + pipelines/basic/ltx2/stages/ltx2_denoising.py): video_sigma/audio_sigma populate Modality.sigma to drive the prompt AdaLN. No-op for 2.0 (sigma unused when prompt_adaln is None).
• VAE decoder (models/vaes/ltx2vae.py, +8 lines): apply the existing compress_all channel-multiplier scheme to compress_space/compress_time too (the 2.3 base VAE grows decoder channels 128→1024 via these). No-op when multiplier=1 (2.0).
• Registration (presets.py, registry.py): LTX2_3_BASE preset (30 steps, CFG 3.0 + negative prompt, STG [28]) + 2.3 model entry + distilled detector aliases.

Backward compatibility

Every 2.3 path is gated by a config flag (default off) or no-ops for 2.0 (sigma/VAE multipliers). Existing LTX2 tests pass on this branch.

Test plan — what IS validated

• py_compile + import all touched modules
• 35 existing LTX2 tests pass (LTX-2.0 unaffected): test_ltx2_stage_overrides (12), test_ltx2_continuation (17), test_nvfp4_ltx2_wiring (6)
• E2E — single-GPU, eager: generates coherent LTX-2.3 base t2v (832×1280, 30 steps, CFG 3, num_gpus=1, enable_torch_compile=False), e2e ≈ 103s, output matches prompt

NOT yet validated (reviewer attention)

• Sequence-parallel / multi-GPU: this PR keeps upstream's original_seq_len SP attention-mask path and does not port internal's diverged scheme (create_attention_mask_for_padding, skip_*_self_attn_blocks lists). The new gated-attn / cross-attn-AdaLN / STG-mask paths are untested under SP — needs a multi-GPU run.
• torch.compile (fullgraph=True): compile-enable showed no setup-time graph break, but a full warmup+measured compiled run was not completed — graph-break behavior of the new 2.3 paths is unverified.
• Functional 2.3 generation against the official 2.3 weights in CI.

Review notes

1. to_gate_logits (2.3 per-head gate on x) is distinct from to_gate_compress (VSA gate on context) — kept separate, not cross-wired.
2. STG: the hard if not skip_video_self_attn branch became an always-compute multiply-mask (skip→×0 equivalent), so attn1 always runs on the perturbed pass; confirm acceptable.
3. VAE compress_space/compress_time now mirror the existing compress_all channel-multiplier convention.

Excludes: i2v conditioning + LoRA fix + NVFP4 per-stage profile (already upstream), training/distillation, demo UI. Audio BWE is the companion PR (#1398).

[mergify]

Merge Protections

Your pull request matches the following merge protections and will not be merged until they are valid.

🔴 PR merge requirements

Waiting for

• check-success=fastcheck-passed
• check-success=full-suite-passed
• check-success~=pre-commit

This rule is failing.

• check-success=fastcheck-passed
• check-success=full-suite-passed
• check-success~=pre-commit
• #approved-reviews-by>=1
• title~=(?i)^\[(feat|feature|bugfix|fix|refactor|perf|ci|doc|docs|misc|chore|kernel|new.?model|skill|skills|infra)\]

[gemini-code-assist]

Code Review

This pull request introduces support for LTX-2.3 models, adding features such as gated self-attention, cross-attention AdaLN, and separate video/audio feature extractors in the Gemma text encoder, while maintaining backward compatibility with LTX-2.0. The code review identified several critical and high-severity issues, including a bug where modality.sigma is never populated (which will cause runtime crashes when cross_attention_adaln is enabled), a backward compatibility break for LTX-2.0 VSA checkpoints due to unconditional parameter renaming, a potential crash with single-element 1D STG masks, and potential runtime errors from calling .view() on non-contiguous tensors.

[gemini-code-assist]

Critical Bug: modality.sigma is never populated

In LTX2Transformer3DModel.forward (lines 2962 and 3007), video_modality and audio_modality are instantiated without passing the sigma parameter. As a result, modality.sigma always defaults to None.

This means prompt_timestep will always be None here, which will cause apply_cross_attention_adaln to raise a ValueError at runtime whenever cross_attention_adaln is enabled.

Fix: Update LTX2Transformer3DModel.forward to accept sigma (e.g., from kwargs or as an explicit argument) and pass it when instantiating the Modality objects.

[gemini-code-assist]

High Severity: Backward Compatibility Break for LTX-2.0 VSA Checkpoints

Renaming to_gate_compress to to_gate_logits unconditionally in param_names_mapping will break backward compatibility for LTX-2.0 checkpoints that use Video Sparse Attention (VSA).

In LTX-2.0 VSA, to_gate_compress is a valid parameter representing the VSA gate. Unconditional renaming will attempt to load these weights into to_gate_logits, causing shape mismatches and runtime errors.

Recommendation: Perform this renaming conditionally during weight loading (e.g., in load_weights by checking the model config or parameter shapes) rather than using a global static regex mapping.

[gemini-code-assist]

Medium Severity: Potential Crash with Single-Element 1D STG Mask

If skip_flag is a 1D tensor of shape [1], skip_flag.shape[0] != bsz will raise a ValueError when bsz > 1. In PyTorch, single-element 1D tensors are often used as broadcastable scalars.

It is safer to check skip_flag.numel() == 1 and expand it to bsz to handle this common edge case.

Suggested change

	if skip_flag.ndim == 0:
	perturb = skip_flag.reshape(1).expand(bsz)
	else:
	if skip_flag.shape[0] != bsz:
	raise ValueError(
	"Per-sample STG mask batch size mismatch: "
	f"got {skip_flag.shape[0]}, expected {bsz}")
	perturb = skip_flag.reshape(bsz)
	if skip_flag.numel() == 1:
	perturb = skip_flag.reshape(1).expand(bsz)
	else:
	if skip_flag.shape[0] != bsz:
	raise ValueError(
	"Per-sample STG mask batch size mismatch: "
	f"got {skip_flag.shape[0]}, expected {bsz}")
	perturb = skip_flag.reshape(bsz)

[gemini-code-assist]

Medium Severity: Potential View Crash on Non-Contiguous Tensor

If out is not contiguous (which can happen depending on the attention backend or sequence parallel operations), out.view will raise a RuntimeError. Using reshape is safer as it automatically handles non-contiguous tensors by copying them if necessary.

Suggested change

	out = out.view(b, q_len, self.heads, self.dim_head)
	out = out.reshape(b, q_len, self.heads, self.dim_head)

[gemini-code-assist]

Medium Severity: Potential View Crash on Non-Contiguous Tensor

If out is not contiguous, out.view will raise a RuntimeError. Using reshape is safer as it automatically handles non-contiguous tensors.

Suggested change

	out = out.view(b, q_len, self.heads, self.dim_head)
	out = out.reshape(b, q_len, self.heads, self.dim_head)

[SolitaryThinker]

Hi @alexzms — automated review from Gob, one of @SolitaryThinker's AI reviewers. Findings aren't all human-verified; ping @SolitaryThinker if anything looks off.

TL;DR

Solid config-gated extension — default-off everywhere, the LTX-2.0 path is empirically preserved by the existing 35-test suite, and the new arch flags / preset / registry / pipeline plumbing all line up cleanly. ship-with-fixes: 1 × S1, 3 × S2, 2 × S3.

The dominant concern is the new top-of-dict to_gate_compress→to_gate_logits param_names_mapping. It matches any checkpoint key containing to_gate_compress, which is also a legitimate LTX-2.0 VSA module and a default LoRA target module — so LTX-2.0 LoRAs trained with the default lora_target_modules list will silently retarget to a differently-shaped param after this PR. Detail and fix options below.

The rest is test gaps (zero new tests for ~575 LoC of new model code) and reviewer-attention items you already flagged (SP / compile / official-2.3-weights).

---

S1-1 — to_gate_compress rename has wider blast radius than intended

fastvideo/configs/models/dits/ltx2.py:14-19 adds four rules at the top of param_names_mapping:

r"^(.*)\.to_gate_compress\.(.*)$": r"model.\1.to_gate_logits.\2",
# (+ three more-specific anchored variants)

But to_gate_compress already exists in LTX-2.0:

1. fastvideo/models/dits/ltx2.py:1381-1388 / 1544-1551 on origin/main — created when the attention backend is VIDEO_SPARSE_ATTN. Weight shape: Linear(context_dim, inner_dim) = [4096, 4096] for LTX-2.
2. fastvideo/train/utils/lora.py:36 and fastvideo/pipelines/lora_pipeline.py:171 — listed as a default LoRA target module.

The new to_gate_logits is Linear(query_dim, heads) = [32, 4096]. Completely different shape.

Concrete failure modes after this PR:

• An LTX-2.0 VSA-trained checkpoint loads through the new mapping → *.to_gate_compress.weight gets retargeted to a [32, 4096] slot. Best case: strict loader hard-fails on shape mismatch. Worst case: permissive loader silently drops it, losing the VSA gate weight.
• Higher-probability case: any LTX-2.0 LoRA trained with the default lora_target_modules list ships *.to_gate_compress.lora_A/lora_B keys. The mapping rewrites those to to_gate_logits.lora_A/B, breaking every such LoRA in the wild.

Mitigating context (why S1 not S0): the FastVideo-shipped LTX-2.0 base checkpoint at FastVideo/LTX2-Diffusers is not VSA-trained, so the key doesn't appear in its state-dict — the PR's empirical "35 LTX-2.0 tests pass" is consistent with this.

Suggested fix (pick one):

• Option A — explicit gate: move the four to_gate_compress rules into a separate dict that LTX2VideoArchConfig merges in only when apply_gated_attention=True. LTX-2.0 (default) stays provably untouched.
• Option B — narrow the regex: anchor on the new-only block prefixes (\.blocks\.\d+\.(attn1|attn2|audio_attn1|audio_attn2|audio_to_video_attn|video_to_audio_attn)\.to_gate_compress\.), explicitly excluding attn_masked.to_gate_compress (the VSA path).

Either way, please add a regression test: feed a synthetic state-dict with an *.attn_masked.to_gate_compress.weight key through get_param_names_mapping(LTX2VideoConfig().param_names_mapping) and assert the name is unchanged.

---

S2-1 — Zero new tests for ~575 LoC of new model code

git diff origin/main...HEAD --stat -- fastvideo/tests/ tests/ is empty.

The existing LTX-2.0 tests cover the unchanged-default axis — they verify the new defaults don't break 2.0, which is the main backward-compat guarantee. But the new LTX-2.3 paths themselves have no component-level coverage. Minimum set worth adding:

• Config matrix smoketest: build LTX2Transformer3DModel with each new flag toggled, run a random-tensor forward, assert grad flow reaches to_gate_logits / prompt_scale_shift_table / audio_prompt_scale_shift_table / prompt_adaln_single.*.
• Param-mapping unit test for S1-1 (covered above).
• _build_attn_keep_mask test: bool False (LTX-2.0 all-ones), bool True off-block, scalar tensor, [B] tensor, plus the size-mismatch error path.
• apply_cross_attention_adaln shape test: pin the [B, T, 2, -1] reshape contract on prompt_timestep.
• Gemma v2 extractor parity: with use_v2_feature_extractor=False, _run_feature_extractor now returns (shared, shared) — assert object identity and that _run_connectors numerics match the LTX-2.0 path.

---

S2-2 — STG self-attention path now always computes attn1, multiplies by 0 to "skip"

fastvideo/models/dits/ltx2.py:2017-2106. You flagged this in the PR description ("the hard if not skip_video_self_attn branch became an always-compute multiply-mask … confirm acceptable"), so this is by design.

For the LTX-2.0 default unperturbed pass (skip_video_self_attn=False) the mask is all-ones and math is identical. The two concrete things worth surfacing:

1. Wall-clock cost on STG perturbed pass: LTX-2.0 previously short-circuited attn1 entirely on the configured block; now it always computes and zeros out. That's ~one extra self-attention per STG-active block per step on the negative pass. Worth a one-liner quantifying it (or stating it's trivial) so users know.
2. SSIM stability: please confirm the "35 existing LTX-2.0 tests pass" includes test_ltx2_similarity.py with STG enabled — bitwise-identical output isn't guaranteed across non-deterministic SDPA/flash backends even on the all-ones path, since the LTX-2.0 SSIM reference was generated when attn1 was not invoked at all on STG blocks.

Suggested fix (preserves both numerics and wall-clock on LTX-2.0 STG while keeping the new per-sample mask path):

if torch.is_tensor(skip_video_self_attn) or (self.cross_attention_adaln and self.idx == self.stg_block_idx):
    # 2.3-style per-sample mask: always compute, multiply by mask
    ...
else:
    # 2.0-style bool skip: original short-circuit
    if not skip_video_self_attn:
        vx = vx + self.attn1(norm_vx, ...) * vgate_msa

---

S2-3 — preprocess_text_embeddings discards the audio branch under v2 extractor

fastvideo/models/encoders/gemma.py:604-612:

video_features, _ = self._run_feature_extractor(
    outputs.hidden_states, attention_mask, padding_side=target_padding_side,
)
return video_features, attention_mask

When use_v2_feature_extractor=True (LTX-2.3), _run_feature_extractor now returns distinct (video, audio) tensors with different output dims. preprocess_text_embeddings keeps only video_features and the LTX-2 preprocess pipeline (fastvideo/pipelines/preprocess/ltx2/ltx2_preprocess_pipelines.py:77) writes only that. Downstream training then feeds the same tensor through both video and audio connectors — exactly the LTX-2.0 collapse that the v2 extractor was added to avoid.

Inference is unaffected (the runtime forward keeps both branches). But this is a landmine for the audio-BWE follow-up #1398 and any future LTX-2.3 training work — silent quality regression rather than loud failure.

Suggested fix: when use_v2_feature_extractor, either extend the preprocess return tuple to (video_embeds, audio_embeds, attention_mask) and have the preprocess pipeline persist both, or raise a NotImplementedError on this code path so the failure is loud.

---

S3-1 — StageAwareRMSNorm docstring undersells the backward-compat reasoning

fastvideo/models/dits/ltx2.py:49-66: the docstring says "outside the refine stage this behaves identically to torch.nn.RMSNorm" but doesn't mention that the trainable weight parameter is inherited from super().__init__, so LTX-2.0 q_norm/k_norm checkpoint weights load unchanged. That's the load-bearing fact for backward-compat; one extra sentence saves the next reviewer five minutes.

---

S3-2 — LTX-2.3 detector and distilled-alias intent could use comments

fastvideo/registry.py:282-289: the detector matches the substring "2.3" in the LTX-2.0 exclusion lambda. Works for listed repos; cheap to harden by anchoring on "-2.3" / "2.3-". Cosmetic.

fastvideo/registry.py:230-241 (the existing distilled detector) is unchanged but now also matches the new LTX-2.3 distilled aliases. If the intent is "2.3 distilled inference uses the same preset as 2.0 distilled" (which the PR description implies), a one-line comment next to the alias clarifies it; if a 2.3-specific distilled preset is intended later, this is silently wrong now.

---

Files reviewed

registry.py, configs/models/dits/ltx2.py, configs/models/encoders/gemma.py, models/dits/ltx2.py, models/encoders/gemma.py, models/vaes/ltx2vae.py, pipelines/basic/ltx2/presets.py, pipelines/basic/ltx2/stages/ltx2_denoising.py, plus read-only verification of fastvideo/tests/ LTX-2 suite, train/utils/lora.py, pipelines/lora_pipeline.py, and training/ltx2_training_pipeline.py callers.

---

— Gob (@SolitaryThinker's AI reviewer). Full review (including S3 items) is archived locally.

[mergify]

Pre-commit checks failed

Hi @alexzms, the pre-commit checks have failed. To fix them locally:

# Install pre-commit if you haven't already
uv pip install pre-commit
pre-commit install

# Run all checks and auto-fix what's possible
pre-commit run --all-files

Common fixes:

• yapf: yapf -i <file> (formatting)
• ruff: ruff check --fix <file> (linting)
• codespell: codespell --write-changes <file> (spelling)

After fixing, commit and push the changes. The checks will re-run automatically.

For future commits, pre-commit will run automatically on changed files before each commit.

[alexzms]

Pushed 06906ff9 to address Gob's S1.

Diagnosis verified, with one path-shape correction: the LTX-2.0 to_gate_compress parameter is a sibling of attn_masked on the attention module (fastvideo/models/dits/ltx2.py:1497-1504 on this branch, :1381-1388 on origin/main), not a child — so the path is e.g. attn1.to_gate_compress, not attn1.attn_masked.to_gate_compress. The blast radius is exactly as Gob flagged; a path-based exclusion of attn_masked.* would not have worked because both LTX-2.0 (VSA) and LTX-2.3 (gated) place their to_gate_compress at the same depth.

Going with Option A: moved the 4 rename rules out of the static default_factory and inject them at the front of param_names_mapping in __post_init__ only when apply_gated_attention=True. With the flag off (LTX-2.0 default), the rules never run and load behaviour is byte-identical to origin/main. With it on (LTX-2.3), they fire before the generic prefix-strip rules (first-match-wins iteration in get_param_names_mapping).

Regression test: fastvideo/tests/api/test_ltx2_param_mapping.py covers the LTX-2.0 VSA + default-target LoRA preservation, the LTX-2.3 rename path across all 4 prefix variants, and the rule-ordering invariant. 14 new tests + 48 existing LTX-2 api tests all pass.

[SolitaryThinker]

/merge

[mergify]

Pre-commit checks failed

Hi @alexzms, the pre-commit checks have failed. To fix them locally:

# Install pre-commit if you haven't already
uv pip install pre-commit
pre-commit install

# Run all checks and auto-fix what's possible
pre-commit run --all-files

Common fixes:

• yapf: yapf -i <file> (formatting)
• ruff: ruff check --fix <file> (linting)
• codespell: codespell --write-changes <file> (spelling)

After fixing, commit and push the changes. The checks will re-run automatically.

For future commits, pre-commit will run automatically on changed files before each commit.