[Model] Add LTX-2.5 support (#6070)

Signed-off-by: lishunyang12 <lishunyang12@163.com>

.buildkite/cuda/test-nightly.yml

@@ -459,6 +459,12 @@ steps:
           - pytest -s -v tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_ltx.py -m "full_model and cuda" --run-level "full_model"
         mirror_hardwares: h100_2
 
+      - label: ":full_moon: Diffusion X2V · LTX-2.5 Online Test"
+        timeout_in_minutes: 60
+        commands:
+          - pytest -s -v tests/e2e/online_serving/test_ltx25.py -m "full_model and cuda" --run-level "full_model"
+        mirror_hardwares: h100_1
+
       - label: ":full_moon: Diffusion X2V · Wan2.2 Accuracy Test"
         timeout_in_minutes: 180
         commands:
@@ -487,6 +493,13 @@ steps:
           - pytest -s -v 'tests/e2e/accuracy/ltx/test_ltx_official_similarity.py::test_ltx_matches_official[ltx23_two_stage_layer_fused_i2v]' -m full_model --run-level full_model
         mirror_hardwares: h100_1
 
+      - label: ":full_moon: Diffusion X2V · LTX-2.5 Official Parity Test"
+        timeout_in_minutes: 180
+        commands:
+          - pytest -s -v tests/e2e/accuracy/ltx/test_ltx25_official_similarity.py::test_ltx25_distilled_two_stage_matches_official -m full_model --run-level full_model
+          - pytest -s -v tests/e2e/accuracy/ltx/test_ltx25_official_similarity.py::test_ltx25_full_matches_official -m full_model --run-level full_model
+        mirror_hardwares: h100_1
+
       - label: ":full_moon: Diffusion X2V · Perf Test"
         key: nightly-diffusion-x2v-performance
         timeout_in_minutes: 180

.gitignore

@@ -254,6 +254,7 @@ outputs/
 results/
 generated/
 output_*/
+/tests/e2e/accuracy/artifacts/
 
 # Configuration overrides
 configs/local.yaml

docs/models/supported_models.md

@@ -43,7 +43,12 @@ th {
 | `Wan22VACEPipeline` | Wan2.2-VACE | `Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers` | ✅︎ |   |   |   | — |
 | `LTX2Pipeline` | LTX-2 / LTX-2.3 one-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | | — |
 | `LTX2TwoStagePipeline` | LTX-2 / LTX-2.3 ordinary two-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` + matching Lightricks LoRA and upsampler | ✅︎ | ✅︎ | | | — |
-| `LTX2DistilledPipeline` | LTX-2 / LTX-2.3 full-distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled`, `diffusers/LTX-2.3-Distilled-Diffusers` + `Lightricks/LTX-2.3` upsampler | ✅︎ | ✅︎ | | | — |
+| `LTX2DistilledOneStagePipeline` | LTX-2 / LTX-2.3 merged-distilled one-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled`, `diffusers/LTX-2.3-Distilled-Diffusers` | ✅︎ | ✅︎ | | | — |
+| `LTX2DistilledTwoStagePipeline` | LTX-2 / LTX-2.3 merged-distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled`, `diffusers/LTX-2.3-Distilled-Diffusers` + matching Lightricks upsampler | ✅︎ | ✅︎ | | | — |
+| `LTX2Pipeline` | LTX-2.5 Full/SFT one-stage T2V and I2V | `Lightricks/LTX-2.5-Diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/LTX/LTX-2.5.md) |
+| `LTX2TwoStagePipeline` | LTX-2.5 Full/SFT two-stage T2V and I2V | `Lightricks/LTX-2.5-Diffusers` + `Lightricks/LTX-2.5` LoRA | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/LTX/LTX-2.5.md) |
+| `LTX2DistilledOneStagePipeline` | LTX-2.5 merged-distilled one-stage T2V and I2V | `Lightricks/LTX-2.5-Diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/LTX/LTX-2.5.md) |
+| `LTX2DistilledTwoStagePipeline` | LTX-2.5 merged-distilled two-stage T2V and I2V | `Lightricks/LTX-2.5-Diffusers` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/LTX/LTX-2.5.md) |
 | `LingBotVideoPipeline` | LingBot-Video dense and MoE T2I, T2V, TI2V | `robbyant/lingbot-video-dense-1.3b`, `robbyant/lingbot-video-moe-30b-a3b` | ✅︎ | | | | — |
 | `MiniMaxH3Pipeline` | MiniMax H3 T2VA, first/last-keyframe FL2VA, and mixed-reference Ref2VA | `MiniMaxAI/MiniMax-H3` | ✅︎ | ✅︎ | | | [Published](https://recipes.vllm.ai/MiniMaxAI/MiniMax-H3) |
 | `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |

docs/user_guide/diffusion_features.md

@@ -161,6 +161,7 @@ The following tables show which models support each feature:
 | **Wan2.1-VACE**              |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **LTX-2**                    |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **LTX-2.3**                  |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
+| **LTX-2.5**                  |     ❌     | ❓ (one-stage) | ❓ (Ulysses only) | ✅ (Full only) |         ❓         |         ❌         |   ✅    |             ✅             |      ✅ (decode)      |    ❓ (FP8)     |        ❌         |
 | **Helios**                   |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ✅*        |
 | **HunyuanVideo-1.5 T2V I2V** |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |  ✅ (encode/decode)   |       ✅        |        ❌         |
 | **DreamID-Omni**             |     ❌     |     ❌      |           ❌           |       ✅        |         ❌         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
@@ -174,7 +175,7 @@ The following tables show which models support each feature:
 **Frame Interpolation Support**
 
 - **Supported**: Wan2.2 text-to-video, image-to-video, and TI2V pipelines
-- **Not supported**: Wan2.1-VACE, LTX-2, LTX-2.3, Helios, HunyuanVideo-1.5, DreamID-Omni
+- **Not supported**: Wan2.1-VACE, LTX-2, LTX-2.3, LTX-2.5, Helios, HunyuanVideo-1.5, DreamID-Omni
 
 ### AudioGen
 

examples/offline_inference/image_to_video/image_to_video.py

@@ -68,6 +68,7 @@ from vllm_omni.model_extras import (
 from vllm_omni.model_extras import (
     get_extra_body_params,
     get_model_class_name,
+    should_preserve_reference_image_size,
 )
 from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.platforms import current_omni_platform
@@ -454,7 +455,13 @@ def main():
 
     media_inputs: dict[str, Any] = {}
     if image is not None:
-        media_inputs["image"] = image.resize((width, height), PIL.Image.Resampling.LANCZOS)
+        preserve_image_size = should_preserve_reference_image_size(
+            model_class_name,
+            model=args.model,
+        )
+        media_inputs["image"] = (
+            image if preserve_image_size else image.resize((width, height), PIL.Image.Resampling.LANCZOS)
+        )
     if last_image is not None:
         media_inputs["last_image"] = last_image.resize((width, height), PIL.Image.Resampling.LANCZOS)
     if mask_image is not None:

examples/offline_inference/text_to_video/text_to_video.py

@@ -570,6 +570,8 @@ def main():
         num_frames=args.num_frames,
         extra_args=extra_args,
     )
+    if args.frame_rate is not None:
+        sampling_kwargs["frame_rate"] = args.frame_rate
     if args.guidance_scale_high is not None:
         sampling_kwargs["guidance_scale_2"] = args.guidance_scale_high
 

recipes/LTX/LTX-2.5.md

@@ -0,0 +1,175 @@
+# LTX-2.5
+
+> Text-to-video and first-frame image-to-video generation with synchronized audio
+
+## Overview
+
+vLLM-Omni supports the gated
+[`Lightricks/LTX-2.5-Diffusers`](https://huggingface.co/Lightricks/LTX-2.5-Diffusers)
+checkpoint for T2V and first-frame I2V. Every output includes synchronized
+48 kHz stereo audio.
+
+The raw [`Lightricks/LTX-2.5`](https://huggingface.co/Lightricks/LTX-2.5)
+repository is not directly loadable with `--model`. It supplies the official
+upsampler and LoRA sidecars used by the Full/SFT two-stage pipeline, so accept
+both model licenses and authenticate before first use.
+
+## Choose a pipeline
+
+| Pipeline | Mode | Default size | Schedule |
+|---|---|---:|---:|
+| `LTX2Pipeline` | Full/SFT one-stage | 960x544 | 30 steps |
+| `LTX2TwoStagePipeline` | Full/SFT two-stage | 1920x1088 | 30 + 3 steps |
+| `LTX2DistilledOneStagePipeline` | Distilled one-stage | 960x544 | 8 steps |
+| `LTX2DistilledTwoStagePipeline` | Distilled two-stage | 1920x1088 | 8 + 3 steps |
+
+Both two-stage pipelines generate at half resolution, apply the official x2
+latent upsampler, and run a three-step refinement stage. Select the class with
+`--model-class-name`; no `--task-type` flag is required. Supplying one initial
+image selects I2V, while omitting it selects T2V.
+
+## Prerequisites
+
+```bash
+hf auth login
+export MODEL=Lightricks/LTX-2.5-Diffusers
+```
+
+Install matching vLLM and vLLM-Omni versions, and ensure `ffmpeg` and
+`ffprobe` are on `PATH`. I2V requires PyAV backed by an FFmpeg build with the
+`libx264` encoder.
+
+## Hardware
+
+| GPU | Status | Recommended scope |
+|---|---|---|
+| NVIDIA B300 | Verified | All four canonical pipelines |
+| NVIDIA B200 or H200 | Capacity-based recommendation; not yet verified | All four canonical pipelines |
+| NVIDIA GB200 or GB300 | Capacity-based recommendation; not yet verified | All four canonical pipelines |
+| NVIDIA H100 80 GB | FP8 recipe | Distilled one-stage at 960x544 |
+
+The 1920x1088 two-stage examples require about 114 GB of peak GPU memory.
+80 GB GPUs do not have enough safety margin for the canonical two-stage
+configuration; use a larger GPU or reduce the output configuration.
+
+### H100 80 GB FP8
+
+Use the distilled one-stage pipeline with FP8 and cuDNN attention:
+
+```bash
+vllm serve Lightricks/LTX-2.5-Diffusers \
+  --omni \
+  --model-class-name LTX2DistilledOneStagePipeline \
+  --quantization fp8 \
+  --diffusion-attention-backend CUDNN_ATTN \
+  --host 0.0.0.0 \
+  --port 8000
+```
+
+Use the T2V or I2V request below after the server is ready.
+
+## Offline inference
+
+Choose values from the pipeline table. For example, the distilled two-stage
+path uses:
+
+```bash
+export PIPELINE=LTX2DistilledTwoStagePipeline
+export WIDTH=1920
+export HEIGHT=1088
+export STEPS=8
+```
+
+### Text-to-video
+
+```bash
+python examples/offline_inference/text_to_video/text_to_video.py \
+  --model "${MODEL}" \
+  --model-class-name "${PIPELINE}" \
+  --prompt "A cinematic shot of a red fox walking through a snowy forest at dawn, the camera tracking alongside, snow crunching underfoot." \
+  --width "${WIDTH}" \
+  --height "${HEIGHT}" \
+  --num-frames 121 \
+  --num-inference-steps "${STEPS}" \
+  --frame-rate 24 \
+  --fps 24 \
+  --seed 42 \
+  --output ltx25-t2v.mp4
+```
+
+### First-frame image-to-video
+
+```bash
+python examples/offline_inference/image_to_video/image_to_video.py \
+  --model "${MODEL}" \
+  --model-class-name "${PIPELINE}" \
+  --image /absolute/path/to/first-frame.png \
+  --prompt "The red fox walks forward while the camera tracks alongside." \
+  --width "${WIDTH}" \
+  --height "${HEIGHT}" \
+  --num-frames 121 \
+  --num-inference-steps "${STEPS}" \
+  --frame-rate 24 \
+  --fps 24 \
+  --seed 42 \
+  --output ltx25-i2v.mp4
+```
+
+LTX-2.5 uses the official CRF-18 first-frame conditioning path by default.
+
+## Online serving
+
+Start one server for the selected pipeline:
+
+```bash
+vllm serve "${MODEL}" \
+  --omni \
+  --model-class-name "${PIPELINE}" \
+  --host 0.0.0.0 \
+  --port 8000 \
+  --stage-init-timeout 900
+```
+
+T2V request:
+
+```bash
+curl -sS --fail-with-body \
+  -X POST http://127.0.0.1:8000/v1/videos/sync \
+  -F 'prompt=A cinematic shot of a red fox walking through a snowy forest at dawn, the camera tracking alongside, snow crunching underfoot.' \
+  -o ltx25-online-t2v.mp4
+```
+
+For I2V, add exactly one first frame:
+
+```bash
+export FIRST_FRAME=/absolute/path/to/first-frame.png
+
+curl -sS --fail-with-body \
+  -X POST http://127.0.0.1:8000/v1/videos/sync \
+  -F "input_reference=@${FIRST_FRAME};type=image/png" \
+  -F 'prompt=The red fox walks forward while the camera tracks alongside.' \
+  -o ltx25-online-i2v.mp4
+```
+
+Restart the server after changing `PIPELINE` because each class selects a
+different weight layout and execution topology.
+
+## Constraints
+
+- Online `width`, `height`, `num_frames`, `fps`, and `num_inference_steps`
+  are optional and use the selected pipeline defaults when omitted. `seed` is
+  optional and random when omitted.
+- `num_frames` must be `8k+1` when overridden.
+- One-stage width and height must be divisible by 32; two-stage final
+  dimensions must be divisible by 64.
+- Full/SFT pipelines accept negative prompts. Distilled pipelines are
+  positive-only and reject negative prompts.
+- I2V accepts exactly one initial image per prompt.
+- The supported `--model` value is `Lightricks/LTX-2.5-Diffusers`.
+
+## References
+
+- [LTX-2.5-Diffusers model card](https://huggingface.co/Lightricks/LTX-2.5-Diffusers)
+- [Official LTX-2 repository](https://github.com/Lightricks/LTX-2)
+- [Generic T2V example](../../examples/offline_inference/text_to_video/text_to_video.py)
+- [Generic I2V example](../../examples/offline_inference/image_to_video/image_to_video.py)

recipes/LTX/LTX-2.md

@@ -8,10 +8,12 @@
 |---|---|---|
 | `LTX2Pipeline` | LTX-2 one-stage T2V/I2V | `Lightricks/LTX-2` |
 | `LTX2TwoStagePipeline` | LTX-2 ordinary two-stage T2V/I2V | `Lightricks/LTX-2` |
-| `LTX2DistilledPipeline` | LTX-2 full-distilled two-stage T2V/I2V | `rootonchair/LTX-2-19b-distilled` |
+| `LTX2DistilledOneStagePipeline` | LTX-2 merged-distilled one-stage T2V/I2V | `rootonchair/LTX-2-19b-distilled` |
+| `LTX2DistilledTwoStagePipeline` | LTX-2 merged-distilled two-stage T2V/I2V | `rootonchair/LTX-2-19b-distilled` |
 | `LTX2Pipeline` | LTX-2.3 one-stage T2V/I2V | `diffusers/LTX-2.3-Diffusers` |
 | `LTX2TwoStagePipeline` | LTX-2.3 ordinary two-stage T2V/I2V | `diffusers/LTX-2.3-Diffusers`<br>`Lightricks/LTX-2.3` |
-| `LTX2DistilledPipeline` | LTX-2.3 full-distilled two-stage T2V/I2V | `diffusers/LTX-2.3-Distilled-Diffusers`<br>`Lightricks/LTX-2.3` |
+| `LTX2DistilledOneStagePipeline` | LTX-2.3 merged-distilled one-stage T2V/I2V | `diffusers/LTX-2.3-Distilled-Diffusers` |
+| `LTX2DistilledTwoStagePipeline` | LTX-2.3 merged-distilled two-stage T2V/I2V | `diffusers/LTX-2.3-Distilled-Diffusers`<br>`Lightricks/LTX-2.3` |
 
 Repositories in the table are download units. A full pipeline repository
 contains the Transformer, text encoder, connectors, VAEs, vocoder, scheduler,
@@ -26,9 +28,12 @@ image selects I2V. Both one-stage repositories declare this class, so
 the raw `Lightricks/LTX-2.3` safetensors repository is not directly loadable.
 
 `LTX2TwoStagePipeline` samples the regular model at half resolution, upsamples,
-then refines with the distilled LoRA. `LTX2DistilledPipeline` uses a fully
-merged distilled Transformer in both stages. Both entries support T2V and I2V;
-select their class explicitly.
+then refines with the distilled LoRA. `LTX2DistilledOneStagePipeline` uses a
+merged distilled Transformer without upsampling, while
+`LTX2DistilledTwoStagePipeline` uses it in both stages. All entries support
+T2V and I2V; select their class explicitly. The deprecated
+`LTX2DistilledPipeline` name remains an alias for
+`LTX2DistilledTwoStagePipeline`.
 
 ## API Migration
 
@@ -52,8 +57,8 @@ The consolidation also removes these registry names without aliases:
 | `LTX23Pipeline` | `LTX2Pipeline`; checkpoint metadata selects LTX-2.3 |
 | `LTX2ImageToVideoPipeline` | `LTX2Pipeline` with `image=` |
 | `LTX23ImageToVideoPipeline` | `LTX2Pipeline` with `image=`; checkpoint metadata selects LTX-2.3 |
-| `LTX2TwoStagesPipeline` | `LTX2DistilledPipeline` |
-| `LTX2ImageToVideoTwoStagesPipeline` | `LTX2DistilledPipeline` with `image=` |
+| `LTX2TwoStagesPipeline` | `LTX2DistilledTwoStagePipeline` |
+| `LTX2ImageToVideoTwoStagesPipeline` | `LTX2DistilledTwoStagePipeline` with `image=` |
 
 Passing any second positional argument now raises `TypeError`. These changes
 affect direct Python callers and explicit `--model-class-name` overrides;
@@ -112,10 +117,10 @@ Start a two-stage checkpoint with its explicit class:
 
 ```bash
 vllm serve rootonchair/LTX-2-19b-distilled --omni \
-  --model-class-name LTX2DistilledPipeline --stage-init-timeout 600
+  --model-class-name LTX2DistilledTwoStagePipeline --stage-init-timeout 600
 # LTX-2.3 full-distilled; the v1.1 x2 upsampler is resolved when absent
 vllm serve diffusers/LTX-2.3-Distilled-Diffusers --omni \
-  --model-class-name LTX2DistilledPipeline --stage-init-timeout 600
+  --model-class-name LTX2DistilledTwoStagePipeline --stage-init-timeout 600
 # Ordinary LTX-2.3 two-stage
 vllm serve diffusers/LTX-2.3-Diffusers --omni \
   --model-class-name LTX2TwoStagePipeline \

recipes/README.md

@@ -55,6 +55,7 @@ recipes/
 | [`MiniMaxAI/MiniMax-H3-MUSA.md`](./MiniMaxAI/MiniMax-H3-MUSA.md) | MiniMax H3 T2VA, FL2VA, and Ref2VA on MUSA | 4x MTT S5000 (TP4 Ref2VA profile) |
 | [`krea/Krea-2.md`](./krea/Krea-2.md) | Text-to-image (Turbo + Raw), offline + online, with LoRA | 1x H100 80GB |
 | [`LTX/LTX-2.md`](./LTX/LTX-2.md) | LTX-2/LTX-2.3 text-to-video and image-to-video with synchronized audio | H200 141GB / 96GB-class GPU |
+| [`LTX/LTX-2.5.md`](./LTX/LTX-2.5.md) | LTX-2.5-Diffusers: Full/SFT one-stage and distilled two-stage T2V/I2V with synchronized audio | NVIDIA B300; cuDNN-qualified |
 | [`MammothModa2/MammothModa2.md`](./MammothModa2/MammothModa2.md) | Preview and Dev text-to-image (AR → DiT); Dev text/image understanding | Preview: 1x L40S 48GB / 1x ≥40GB GPU; Dev: 1x NVIDIA GPU with sufficient cache headroom |
 | [`meituan-longcat/LongCat-Video-Avatar-1.5.md`](./meituan-longcat/LongCat-Video-Avatar-1.5.md) | Audio-driven avatar video generation (AT2V / AI2V, single- and multi-speaker, AVC continuation) | 1x H100 80GB |
 | [`mistralai/Voxtral-TTS.md`](./mistralai/Voxtral-TTS.md) | Online serving for TTS | 1x RTX 4090 24GB |

requirements/common.txt

@@ -1,15 +1,15 @@
 # Common dependencies for all platforms
-# transformers: kept exactly aligned with the upstream vLLM constraint.
-# vLLM pins ">= 4.56.0, != 5.0.*, != 5.1.*, != 5.2.*, != 5.3.*, != 5.4.*, != 5.5.0".
-# The previous extra "< 5.9.0" exclusion was removed when re-aligning with
-# upstream: transformers 5.9 changed _adjust_missing_and_unexpected_keys from
+# transformers: required unified multimodal model classes are unavailable
+# before 5.10.1; 5.10.0 is yanked, and 5.15 has a known construction
+# regression. transformers 5.9 also changed
+# _adjust_missing_and_unexpected_keys from
 # list+list to set|set, breaking models that declare
 # _keys_to_ignore_on_load_unexpected as a list (e.g. OpenMOSS-Team/MOSS-TTS-Nano)
 # with "TypeError: unsupported operand type(s) for |: 'list' and 'set'".
 # That breakage is now handled in code instead of by pinning:
 # MossTTSNanoForGeneration wraps weight loading in
 # transformers_keys_to_ignore_compat (vllm_omni/model_executor/models/utils.py).
-transformers >= 5.5.3
+transformers >= 5.10.1, < 5.15
 av>=14.0.0
 omegaconf>=2.3.0
 diffusers==0.38.0

tests/diffusion/attention/test_cudnn_attn.py

@@ -5,6 +5,7 @@ from contextlib import nullcontext
 
 import pytest
 import torch
+from torch.nn.attention import SDPBackend
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.backends.cudnn_attn import CuDNNAttentionImpl
@@ -12,6 +13,48 @@ from vllm_omni.diffusion.attention.backends.cudnn_attn import CuDNNAttentionImpl
 pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
 
 
+@pytest.mark.parametrize(
+    ("query_sequence_length", "key_sequence_length", "expected_backend"),
+    [
+        (1, 1, SDPBackend.MATH),
+        (8, 1, SDPBackend.MATH),
+        (1, 8, SDPBackend.MATH),
+        (8, 8, SDPBackend.CUDNN_ATTENTION),
+    ],
+)
+def test_cudnn_routes_single_token_attention_to_math(
+    monkeypatch, query_sequence_length, key_sequence_length, expected_backend
+):
+    selected_backends = []
+
+    def fake_sdpa_kernel(backends):
+        selected_backends.append(backends)
+        return nullcontext()
+
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.attention.backends.cudnn_attn.sdpa_kernel",
+        fake_sdpa_kernel,
+    )
+    monkeypatch.setattr(
+        torch.nn.functional,
+        "scaled_dot_product_attention",
+        lambda query, _key, _value, **_kwargs: query,
+    )
+    impl = CuDNNAttentionImpl(
+        num_heads=2,
+        head_size=4,
+        softmax_scale=0.5,
+    )
+    query = torch.randn(1, query_sequence_length, 2, 4)
+    key = torch.randn(1, key_sequence_length, 2, 4)
+    value = torch.randn_like(key)
+
+    output = impl.forward_cuda(query, key, value)
+
+    assert output.shape == query.shape
+    assert selected_backends == [[expected_backend]]
+
+
 def test_cudnn_slices_valid_kv_prefix_without_padding_mask(monkeypatch):
     observed = {}
 

tests/diffusion/cache/test_cache_dit.py

@@ -36,7 +36,6 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 SEPARATE_CFG_TRANSFORMERS = [
     DreamIdOmniModel,
-    LTX2VideoTransformer3DModel,
     HeliosTransformer3DModel,
     LongCatImageTransformer2DModel,
     Cosmos3VFMTransformer,
@@ -124,6 +123,13 @@ def test_cache_dit_configs_have_separate_cfg(transformer_model):
     assert transformer_model._cache_dit_adapter_config.has_separate_cfg is True
 
 
+def test_ltx2_cache_dit_uses_one_forward_per_denoise_step():
+    """LTX batches all guidance passes into one Transformer invocation."""
+    adapter_config = LTX2VideoTransformer3DModel._cache_dit_adapter_config
+
+    assert adapter_config.has_separate_cfg is False
+
+
 @patch("vllm_omni.diffusion.cache.cachedit.model_specific.BlockAdapter")
 @patch("vllm_omni.diffusion.cache.cachedit.model_specific.cache_dit")
 def test_separate_wan22_custom_enabler_has_separate_cfg(mock_cache_dit, mock_block_adapter):

tests/diffusion/model_loader/test_hub_prefetch.py

@@ -0,0 +1,72 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import contextlib
+
+import huggingface_hub
+import pytest
+
+from vllm_omni.diffusion.model_loader import hub_prefetch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def test_prefetch_subfolders_propagates_revision(monkeypatch):
+    calls = []
+
+    def fake_snapshot_download(**kwargs):
+        calls.append(kwargs)
+
+    monkeypatch.setattr(huggingface_hub, "snapshot_download", fake_snapshot_download)
+    monkeypatch.setattr(hub_prefetch, "_repo_prefetch_lock", lambda _model: contextlib.nullcontext())
+
+    hub_prefetch.prefetch_subfolders(
+        "org/model",
+        ["transformer"],
+        revision="pinned-revision",
+    )
+
+    assert calls == [
+        {
+            "repo_id": "org/model",
+            "revision": "pinned-revision",
+            "allow_patterns": ["transformer/*", "transformer/**", "*.json", "*.txt"],
+        }
+    ]
+
+
+def test_from_pretrained_retry_preserves_revision(monkeypatch):
+    factory_calls = []
+    prefetch_calls = []
+
+    def fake_factory(model, **kwargs):
+        factory_calls.append((model, kwargs))
+        if len(factory_calls) == 1:
+            raise OSError("partial cache")
+        return "loaded"
+
+    def fake_prefetch(model, subfolders, **kwargs):
+        prefetch_calls.append((model, tuple(subfolders), kwargs))
+
+    monkeypatch.setattr(hub_prefetch, "prefetch_subfolders", fake_prefetch)
+    monkeypatch.setattr(hub_prefetch.time, "sleep", lambda _seconds: None)
+
+    result = hub_prefetch.from_pretrained_with_prefetch(
+        fake_factory,
+        "org/model",
+        subfolder="transformer",
+        prefetch_list=["transformer", "vae"],
+        revision="pinned-revision",
+        max_attempts=2,
+    )
+
+    assert result == "loaded"
+    assert len(factory_calls) == 2
+    assert all(call[1]["revision"] == "pinned-revision" for call in factory_calls)
+    assert prefetch_calls == [
+        (
+            "org/model",
+            ("transformer", "vae"),
+            {"local_files_only": False, "revision": "pinned-revision"},
+        )
+    ]

tests/diffusion/models/ltx2/test_ltx25_pipeline.py

@@ -0,0 +1,881 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Focused LTX-2.5 pipeline correctness tests."""
+
+import json
+from dataclasses import replace
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.data import DiffusionOutput
+from vllm_omni.diffusion.models.ltx2 import ltx2_components, ltx2_latents, ltx2_runtime
+from vllm_omni.diffusion.models.ltx2.ltx2_components import (
+    LTX25_DISTILLED_COMPONENT_PROFILE,
+    LTX25_DISTILLED_ONE_STAGE_COMPONENT_PROFILE,
+    LTX25_FULL_COMPONENT_PROFILE,
+    LTX25_TWO_STAGE_COMPONENT_PROFILE,
+    detect_ltx_model_version,
+    resolve_ltx_checkpoint_kind,
+    resolve_ltx_component_profile,
+)
+from vllm_omni.diffusion.models.ltx2.ltx2_denoise import (
+    LTXPhaseResult,
+    _first_frame_keyframes_mask,
+    _official_ltx_sigmas,
+)
+from vllm_omni.diffusion.models.ltx2.ltx2_guidance import (
+    LTX_GUIDANCE_EXECUTOR,
+    LTXGuidanceSpec,
+    LTXModalityGuidance,
+    velocity_from_x0,
+    x0_from_velocity,
+)
+from vllm_omni.diffusion.models.ltx2.ltx2_latents import LTXAVState
+from vllm_omni.diffusion.models.ltx2.ltx2_recipes import (
+    LTX2_DISTILLED_TWO_STAGE_RECIPE,
+    LTX2_ONE_STAGE_RECIPE,
+    LTX23_ONE_STAGE_RECIPE,
+    LTX25_DISTILLED_ONE_STAGE_RECIPE,
+    LTX25_DISTILLED_TWO_STAGE_RECIPE,
+    LTX25_FULL_RECIPE,
+    LTX25_TWO_STAGE_RECIPE,
+    LTX_DISTILLED_SIGMAS,
+    LTX_POSITIVE_ONLY_RECIPE,
+    resolve_ltx_pipeline_recipe,
+)
+from vllm_omni.diffusion.models.ltx2.ltx2_request import LTXRequestInputs
+from vllm_omni.diffusion.models.ltx2.ltx2_runtime import LTXRuntime
+from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import (
+    LTX2DistilledOneStagePipeline,
+    LTX2Pipeline,
+)
+from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import (
+    LTX2DistilledTwoStagePipeline,
+    LTX2TwoStagePipeline,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.mark.parametrize("model_version", ["2", "2.3", "2.5"])
+def test_ltx_timestep_adaln_expands_per_token_like_official_without_sp(model_version):
+    pipe = object.__new__(LTX2Pipeline)
+    pipe.model_version = model_version
+    pipe.od_config = SimpleNamespace(parallel_config=SimpleNamespace(sequence_parallel_size=1))
+    ts = torch.tensor([750.0, 500.0])
+
+    kwargs = LTXRuntime._denoise_timestep_kwargs(
+        pipe,
+        ts,
+        SimpleNamespace(),
+        SimpleNamespace(),
+        video_token_count=4,
+        audio_token_count=3,
+    )
+
+    torch.testing.assert_close(kwargs["timestep"], ts[:, None].expand(-1, 4))
+    torch.testing.assert_close(kwargs["audio_timestep"], ts[:, None].expand(-1, 3))
+    assert kwargs["sigma"] is ts
+    assert kwargs["audio_sigma"] is ts
+
+
+def test_ltx_ancestral_step_clears_sp_audio_padding_after_scheduler_update(monkeypatch):
+    state = LTXAVState(
+        video=torch.zeros(1, 2, 2),
+        audio=torch.zeros(1, 4, 2),
+    )
+    forward_ctx = SimpleNamespace(
+        original_audio_num_frames=3,
+        sampler="euler_ancestral",
+    )
+    denoise_ctx = SimpleNamespace(latents=None, audio_latents=None)
+    pipe = SimpleNamespace(
+        _predict_noise_for_step=lambda *_args: (torch.zeros_like(state.video), torch.zeros_like(state.audio))
+    )
+
+    def fake_step(_pipeline, actual_forward_ctx, _denoise_ctx, *_args):
+        assert actual_forward_ctx.sampler == "euler_ancestral"
+        updated_audio = torch.ones_like(state.audio)
+        updated_audio[:, 3:] = 99.0
+        return torch.ones_like(state.video), updated_audio
+
+    monkeypatch.setattr(ltx2_runtime, "step_denoised_latents", fake_step)
+
+    actual = LTXRuntime._denoise_step(
+        pipe,
+        0,
+        torch.tensor(1.0),
+        state,
+        forward_ctx,
+        denoise_ctx,
+    )
+
+    torch.testing.assert_close(actual.audio[:, :3], torch.ones_like(actual.audio[:, :3]))
+    torch.testing.assert_close(actual.audio[:, 3:], torch.zeros_like(actual.audio[:, 3:]))
+
+
+def test_ltx25_first_causal_latent_frame_is_marked_for_keyframe_embedding():
+    reference = torch.zeros(2, 12, 8)
+
+    mask = _first_frame_keyframes_mask(reference, latent_num_frames=3)
+
+    assert mask.shape == (2, 12, 1)
+    torch.testing.assert_close(mask[:, :4], torch.ones(2, 4, 1))
+    torch.testing.assert_close(mask[:, 4:], torch.zeros(2, 8, 1))
+
+
+def _make_ltx_request_pipe(cls):
+    pipe = object.__new__(cls)
+    torch.nn.Module.__init__(pipe)
+    pipe.device = torch.device("cpu")
+    pipe.tokenizer_max_length = 99
+    pipe.vae_spatial_compression_ratio = 32
+    pipe.vae_temporal_compression_ratio = 8
+    return pipe
+
+
+def _resolve_request_inputs_for_test(
+    pipe,
+    req,
+    *,
+    guidance_scale=4.0,
+    negative_prompt=None,
+):
+    return pipe._resolve_request_inputs(
+        req,
+        prompt=None,
+        negative_prompt=negative_prompt,
+        height=None,
+        width=None,
+        num_frames=None,
+        frame_rate=None,
+        num_inference_steps=None,
+        guidance_scale=guidance_scale,
+        num_videos_per_prompt=1,
+        generator=None,
+        latents=None,
+        audio_latents=None,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+        prompt_attention_mask=None,
+        negative_prompt_attention_mask=None,
+        decode_timestep=0.0,
+        decode_noise_scale=None,
+        output_type="np",
+        max_sequence_length=None,
+    )
+
+
+def test_ltx25_missing_gemma4_recommends_supported_transformers_range(monkeypatch):
+    pipe = SimpleNamespace(component_profile=replace(LTX25_FULL_COMPONENT_PROFILE, text_encoder_cls=None))
+    od_config = SimpleNamespace(model="Lightricks/LTX-2.5-Diffusers", dtype=torch.bfloat16)
+
+    monkeypatch.setattr(ltx2_components, "get_local_device", lambda: torch.device("cpu"))
+    monkeypatch.setattr(ltx2_components, "prefetch_subfolders", lambda *_args, **_kwargs: None)
+    monkeypatch.setattr(ltx2_components.AutoTokenizer, "from_pretrained", lambda *_args, **_kwargs: object())
+
+    with pytest.raises(ImportError) as exc_info:
+        ltx2_components.initialize_pipeline_components(pipe, od_config)
+
+    assert str(exc_info.value) == (
+        "LTX-2.5 requires Gemma4UnifiedForConditionalGeneration; install transformers>=5.10.1,<5.15."
+    )
+
+
+def test_ltx_converted_component_loading_propagates_revision(monkeypatch):
+    revision = "pinned-revision"
+    calls = {"components": []}
+    profile = replace(
+        LTX25_DISTILLED_COMPONENT_PROFILE,
+        text_encoder_cls=object,
+        video_vae_cls=object,
+        vocoder_cls=object,
+        vocoder_fallback_cls=None,
+        scheduler_use_dynamic_shifting=False,
+    )
+    pipeline = SimpleNamespace(component_profile=profile, pipeline_kind="distilled_two_stage")
+    od_config = SimpleNamespace(
+        model="org/converted-ltx25",
+        revision=revision,
+        dtype=torch.bfloat16,
+        quantization_config=None,
+    )
+
+    monkeypatch.setattr(ltx2_components, "get_local_device", lambda: torch.device("cpu"))
+    monkeypatch.setattr(ltx2_components, "_install_connector_attention", lambda _connectors, **_kwargs: None)
+    monkeypatch.setattr(ltx2_components, "_place_aux_components", lambda _pipeline: None)
+
+    def fake_prefetch(model, subfolders, **kwargs):
+        calls["prefetch"] = (model, tuple(subfolders), kwargs)
+
+    def fake_tokenizer(*args, **kwargs):
+        calls["tokenizer"] = (args, kwargs)
+        return SimpleNamespace(model_max_length=1024)
+
+    def fake_component(component_cls, model, subfolder, **kwargs):
+        calls["components"].append((component_cls, model, subfolder, kwargs))
+        if subfolder == "vae":
+            return SimpleNamespace(spatial_compression_ratio=32, temporal_compression_ratio=8)
+        if subfolder == "audio_vae":
+            return SimpleNamespace(
+                mel_compression_ratio=4,
+                temporal_compression_ratio=4,
+                config=SimpleNamespace(sample_rate=16_000, mel_hop_length=160),
+            )
+        return object()
+
+    def fake_transformer_config(model, subfolder, local_files_only, *, revision):
+        calls["transformer_config"] = (model, subfolder, local_files_only, revision)
+        return {"component": "transformer"}
+
+    monkeypatch.setattr(ltx2_components, "prefetch_subfolders", fake_prefetch)
+    monkeypatch.setattr(ltx2_components.AutoTokenizer, "from_pretrained", fake_tokenizer)
+    monkeypatch.setattr(ltx2_components, "_load_component", fake_component)
+    monkeypatch.setattr(ltx2_components, "load_transformer_config", fake_transformer_config)
+    monkeypatch.setattr(
+        ltx2_components,
+        "create_transformer_from_config",
+        lambda *_args, **_kwargs: SimpleNamespace(config=SimpleNamespace(patch_size=1, patch_size_t=1)),
+    )
+
+    def fake_scheduler(*args, **kwargs):
+        calls["scheduler"] = (args, kwargs)
+        return SimpleNamespace(config={"use_dynamic_shifting": False, "shift_terminal": None})
+
+    monkeypatch.setattr(ltx2_components.FlowMatchEulerDiscreteScheduler, "from_pretrained", fake_scheduler)
+
+    ltx2_components.initialize_pipeline_components(pipeline, od_config)
+
+    assert pipeline.weights_sources[0].revision == revision
+    assert calls["prefetch"][2]["revision"] == revision
+    assert calls["tokenizer"][1]["revision"] == revision
+    assert {call[2] for call in calls["components"]} == {
+        "text_encoder",
+        "connectors",
+        "vae",
+        "audio_vae",
+        "vocoder",
+        "latent_upsampler",
+    }
+    assert all(call[3]["revision"] == revision for call in calls["components"])
+    assert calls["transformer_config"] == (
+        od_config.model,
+        profile.transformer_subfolder,
+        False,
+        revision,
+    )
+    assert calls["scheduler"][1]["revision"] == revision
+
+
+def test_ltx_checkpoint_explicit_version_precedes_structural_heuristics(tmp_path):
+    (tmp_path / "model_index.json").write_text(json.dumps({"model_version": "2.3.1"}))
+    (tmp_path / "transformer").mkdir()
+    (tmp_path / "transformer" / "config.json").write_text(json.dumps({"ff_bias": False}))
+
+    assert detect_ltx_model_version(str(tmp_path)) == "2.3"
+
+
+def test_ltx25_checkpoint_selects_full_one_stage_profile(tmp_path, monkeypatch):
+    from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
+
+    (tmp_path / "model_index.json").write_text(
+        json.dumps({"text_encoder": ["transformers", "Gemma4UnifiedForConditionalGeneration"]})
+    )
+
+    def stub_components(pipe, od_config):
+        pipe.od_config = od_config
+        pipe.vae_spatial_compression_ratio = 32
+
+    monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)
+    monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", lambda *_args, **_kwargs: None)
+
+    pipe = LTX2Pipeline(od_config=SimpleNamespace(model=str(tmp_path), enable_diffusion_pipeline_profiler=False))
+
+    assert pipe.model_version == "2.5"
+    assert pipe.component_profile is LTX25_FULL_COMPONENT_PROFILE
+    assert pipe.pipeline_recipe is LTX25_FULL_RECIPE
+    assert pipe.preserve_sp_padded_audio_duration
+    assert pipe.reports_stage_durations
+
+
+def test_ltx25_full_recipe_matches_official_sft_defaults():
+    (phase,) = LTX25_FULL_RECIPE.phases
+
+    assert (LTX25_FULL_RECIPE.height, LTX25_FULL_RECIPE.width) == (544, 960)
+    assert LTX25_FULL_RECIPE.num_frames == 121
+    assert LTX25_FULL_RECIPE.frame_rate == 24.0
+    assert LTX25_FULL_RECIPE.num_inference_steps == 30
+    assert phase.guidance.video == LTXModalityGuidance(
+        cfg_scale=3.0,
+        stg_scale=1.0,
+        modality_scale=3.0,
+        rescale_scale=0.7,
+        stg_blocks=(28,),
+    )
+    assert phase.guidance.audio == LTXModalityGuidance(
+        cfg_scale=7.0,
+        stg_scale=1.0,
+        modality_scale=3.0,
+        rescale_scale=0.7,
+        stg_blocks=(28,),
+    )
+    assert phase.sigmas is None
+    assert phase.noise_scale == 1.0
+    assert phase.use_official_sigma_schedule
+    assert LTX25_FULL_COMPONENT_PROFILE.transformer_subfolder == "transformer_full"
+    assert LTX25_FULL_COMPONENT_PROFILE.scheduler_use_dynamic_shifting
+    assert LTX25_FULL_COMPONENT_PROFILE.scheduler_shift_terminal == 0.1
+
+
+def test_ltx25_full_sigma_schedule_uses_official_default_sequence_anchor():
+    scheduler = SimpleNamespace(
+        config={
+            "base_image_seq_len": 1024,
+            "max_image_seq_len": 4096,
+            "base_shift": 0.95,
+            "max_shift": 2.05,
+            "shift_terminal": 0.1,
+        }
+    )
+
+    sigmas = _official_ltx_sigmas(
+        scheduler,
+        steps=30,
+        device=torch.device("cpu"),
+    )
+
+    torch.testing.assert_close(
+        sigmas[[0, 1, -2, -1]],
+        torch.tensor([1.0, 0.9949570, 0.1, 0.0]),
+        rtol=1e-5,
+        atol=1e-6,
+    )
+
+
+def test_ltx25_checkpoint_selects_distilled_two_stage_profile(tmp_path, monkeypatch):
+    from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
+
+    (tmp_path / "model_index.json").write_text(
+        json.dumps({"text_encoder": ["transformers", "Gemma4UnifiedForConditionalGeneration"]})
+    )
+
+    def stub_components(pipe, od_config):
+        pipe.od_config = od_config
+        pipe.vae_spatial_compression_ratio = 32
+
+    monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)
+    monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", lambda *_args, **_kwargs: None)
+
+    pipe = LTX2DistilledTwoStagePipeline(
+        od_config=SimpleNamespace(
+            model=str(tmp_path),
+            enable_diffusion_pipeline_profiler=False,
+        )
+    )
+
+    assert pipe.component_profile is LTX25_DISTILLED_COMPONENT_PROFILE
+    assert pipe.pipeline_recipe is LTX25_DISTILLED_TWO_STAGE_RECIPE
+
+
+def test_ltx25_full_config_enrichment_uses_selected_transformer_subfolder(tmp_path):
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+    (tmp_path / "model_index.json").write_text(json.dumps({"_class_name": "LTX2Pipeline", "model_version": "2.5"}))
+    for subfolder, marker in (("transformer", "distilled"), ("transformer_full", "full")):
+        path = tmp_path / subfolder
+        path.mkdir()
+        (path / "config.json").write_text(json.dumps({"marker": marker}))
+
+    config = OmniDiffusionConfig(model=str(tmp_path), model_class_name="LTX2Pipeline")
+    config.enrich_config()
+
+    assert config.tf_model_config.get("marker") == "full"
+
+
+def test_ltx25_four_public_pipeline_semantics_are_disjoint():
+    assert resolve_ltx_component_profile("one_stage", "2.5") is LTX25_FULL_COMPONENT_PROFILE
+    assert resolve_ltx_pipeline_recipe("one_stage", "2.5") is LTX25_FULL_RECIPE
+    assert resolve_ltx_component_profile("two_stage", "2.5") is LTX25_TWO_STAGE_COMPONENT_PROFILE
+    assert resolve_ltx_pipeline_recipe("two_stage", "2.5") is LTX25_TWO_STAGE_RECIPE
+    assert resolve_ltx_component_profile("distilled_one_stage", "2.5") is LTX25_DISTILLED_ONE_STAGE_COMPONENT_PROFILE
+    assert resolve_ltx_pipeline_recipe("distilled_one_stage", "2.5") is LTX25_DISTILLED_ONE_STAGE_RECIPE
+    assert resolve_ltx_component_profile("distilled_two_stage", "2.5") is LTX25_DISTILLED_COMPONENT_PROFILE
+    assert resolve_ltx_pipeline_recipe("distilled_two_stage", "2.5") is LTX25_DISTILLED_TWO_STAGE_RECIPE
+
+    assert LTX2Pipeline.pipeline_kind == "one_stage"
+    assert LTX2TwoStagePipeline.pipeline_kind == "two_stage"
+    assert LTX2DistilledOneStagePipeline.pipeline_kind == "distilled_one_stage"
+    assert LTX2DistilledTwoStagePipeline.pipeline_kind == "distilled_two_stage"
+
+
+@pytest.mark.parametrize(
+    ("pipeline_kind", "expected"),
+    [
+        ("one_stage", "regular"),
+        ("two_stage", "regular"),
+        ("distilled_one_stage", "distilled"),
+        ("distilled_two_stage", "distilled"),
+        ("dmd2", None),
+    ],
+)
+def test_ltx_checkpoint_kind_is_derived_from_pipeline_kind(pipeline_kind, expected):
+    assert resolve_ltx_checkpoint_kind(pipeline_kind) == expected
+
+
+def test_ltx_checkpoint_kind_rejects_unknown_pipeline():
+    with pytest.raises(ValueError, match="Unsupported LTX pipeline kind"):
+        resolve_ltx_checkpoint_kind("unknown")
+
+
+def test_ltx25_distilled_two_stage_recipe_matches_model_card_resolution():
+    stage1, stage2 = LTX25_DISTILLED_TWO_STAGE_RECIPE.phases
+
+    assert (LTX25_DISTILLED_TWO_STAGE_RECIPE.height, LTX25_DISTILLED_TWO_STAGE_RECIPE.width) == (1088, 1920)
+    assert stage1.spatial_downscale == 2
+    assert stage1.sigmas == LTX_DISTILLED_SIGMAS
+    assert stage1.sampler == "euler_ancestral"
+    assert stage2.sampler == "euler"
+    assert stage2.sigmas == (0.909375, 0.725, 0.421875, 0.0)
+    assert stage2.input_transform == "spatial_upsample"
+
+    assert LTX25_DISTILLED_TWO_STAGE_RECIPE.negative_prompt == ""
+    assert not LTX25_DISTILLED_TWO_STAGE_RECIPE.allow_negative_prompt
+
+
+def test_ltx25_distilled_one_stage_uses_native_low_resolution():
+    (phase,) = LTX25_DISTILLED_ONE_STAGE_RECIPE.phases
+
+    assert (LTX25_DISTILLED_ONE_STAGE_RECIPE.height, LTX25_DISTILLED_ONE_STAGE_RECIPE.width) == (544, 960)
+    assert phase.spatial_downscale == 1
+    assert phase.sigmas == LTX_DISTILLED_SIGMAS
+    assert phase.sampler == "euler_ancestral"
+    assert not LTX25_DISTILLED_ONE_STAGE_RECIPE.supports_cache_dit
+
+
+def test_ltx25_regular_two_stage_uses_full_transformer_then_distilled_lora():
+    stage1, stage2 = LTX25_TWO_STAGE_RECIPE.phases
+
+    assert (LTX25_TWO_STAGE_RECIPE.height, LTX25_TWO_STAGE_RECIPE.width) == (1088, 1920)
+    assert LTX25_TWO_STAGE_RECIPE.num_inference_steps == 30
+    assert stage1.spatial_downscale == 2
+    assert stage1.guidance == LTX25_FULL_RECIPE.request_guidance
+    assert stage2.adapter_slot == "ltx_distilled"
+    assert LTX25_TWO_STAGE_RECIPE.allow_request_phase_sigmas
+    assert LTX25_TWO_STAGE_COMPONENT_PROFILE.transformer_subfolder == "transformer_full"
+    assert LTX25_TWO_STAGE_COMPONENT_PROFILE.distilled_lora_filename is not None
+
+
+@pytest.mark.parametrize(
+    "recipe",
+    [
+        LTX2_ONE_STAGE_RECIPE,
+        LTX23_ONE_STAGE_RECIPE,
+        LTX_POSITIVE_ONLY_RECIPE,
+    ],
+)
+def test_ltx_one_stage_recipes_declare_cache_dit_supported(recipe):
+    assert recipe.supports_cache_dit
+
+
+@pytest.mark.parametrize(
+    "recipe",
+    [
+        LTX25_FULL_RECIPE,
+        LTX25_DISTILLED_ONE_STAGE_RECIPE,
+        LTX2_DISTILLED_TWO_STAGE_RECIPE,
+        LTX25_DISTILLED_TWO_STAGE_RECIPE,
+    ],
+)
+def test_ltx25_and_multistage_recipes_declare_cache_dit_unsupported(recipe):
+    assert not recipe.supports_cache_dit
+
+
+def test_ltx_positive_only_guidance_preserves_official_x0_roundtrip():
+    sample = torch.full((1, 1), 1e8, dtype=torch.float32)
+    velocity = torch.ones((1, 1), dtype=torch.bfloat16)
+    sigma = torch.tensor(0.5)
+
+    actual = LTX_GUIDANCE_EXECUTOR._guide_modality(
+        sample,
+        {"cond": velocity},
+        sigma,
+        LTXModalityGuidance(),
+    )
+    expected = velocity_from_x0(sample, x0_from_velocity(sample, velocity, sigma), sigma)
+
+    assert actual.dtype == torch.float32
+    torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0)
+
+
+def test_ltx_velocity_from_x0_materializes_official_scalar_sigma():
+    item_calls = []
+
+    class TrackingTensor(torch.Tensor):
+        @staticmethod
+        def __new__(cls):
+            return torch.Tensor._make_subclass(cls, torch.tensor(0.725), False)
+
+        def item(self):
+            item_calls.append(True)
+            return super().item()
+
+    sample = torch.tensor([[1.0]], dtype=torch.bfloat16)
+    x0 = torch.tensor([[0.5]], dtype=torch.bfloat16)
+
+    velocity_from_x0(sample, x0, TrackingTensor())
+
+    assert item_calls == [True]
+
+
+def test_ltx_ancestral_positive_only_guidance_preserves_raw_velocity():
+    sample = torch.full((1, 1), 1e8, dtype=torch.float32)
+    velocity = torch.ones((1, 1), dtype=torch.bfloat16)
+
+    actual = LTX_GUIDANCE_EXECUTOR._guide_modality(
+        sample,
+        {"cond": velocity},
+        torch.tensor(0.5),
+        LTXModalityGuidance(),
+        preserve_positive_velocity=True,
+    )
+
+    assert actual is velocity
+
+
+@pytest.mark.parametrize("model_version", ["2.5", "2.3"])
+def test_ltx_t2v_denoise_state_uses_official_bfloat16_dtype(model_version):
+    expected_dtype = torch.bfloat16
+    captured = {}
+
+    def prepare_latents(**kwargs):
+        captured["video_dtype"] = kwargs["dtype"]
+        return torch.empty(1, 1, 1, dtype=kwargs["dtype"])
+
+    def prepare_audio_latents(*_args, **kwargs):
+        captured["audio_dtype"] = kwargs["dtype"]
+        return torch.empty(1, 1, 1, dtype=kwargs["dtype"]), 1, 1
+
+    pipeline = SimpleNamespace(
+        model_version=model_version,
+        transformer=SimpleNamespace(config=SimpleNamespace(in_channels=128)),
+        prepare_latents=prepare_latents,
+        prepare_audio_latents=prepare_audio_latents,
+        audio_sampling_rate=48_000,
+        audio_hop_length=160,
+        audio_vae_temporal_compression_ratio=4,
+        audio_vae_mel_compression_ratio=4,
+        audio_vae=SimpleNamespace(config=SimpleNamespace(mel_bins=64, latent_channels=8)),
+        _resolve_audio_latent_length=lambda length, _latents: length,
+    )
+    request = SimpleNamespace(
+        num_videos_per_prompt=1,
+        height=64,
+        width=64,
+        num_frames=1,
+        frame_rate=24.0,
+        generator=None,
+        latents=None,
+        audio_latents=None,
+        audio_latents_normalized=False,
+    )
+    prompt = SimpleNamespace(
+        batch_size=1,
+        positive_connector_prompt_embeds=torch.empty(1, dtype=torch.bfloat16),
+        positive_connector_audio_prompt_embeds=torch.empty(1, dtype=torch.bfloat16),
+    )
+
+    LTXRuntime._prepare_video_latents_stage(
+        pipeline,
+        request,
+        prompt,
+        device=torch.device("cpu"),
+        noise_scale=1.0,
+    )
+    LTXRuntime._prepare_audio_latents_stage(
+        pipeline,
+        request,
+        prompt,
+        device=torch.device("cpu"),
+        noise_scale=1.0,
+    )
+
+    assert captured == {"video_dtype": expected_dtype, "audio_dtype": expected_dtype}
+
+
+def test_ltx_seeded_latents_match_official_packed_rng_layout():
+    pipeline = SimpleNamespace(
+        model_version="2.5",
+        vae_spatial_compression_ratio=2,
+        vae_temporal_compression_ratio=2,
+        transformer_spatial_patch_size=1,
+        transformer_temporal_patch_size=1,
+        audio_vae_mel_compression_ratio=2,
+        od_config=SimpleNamespace(parallel_config=SimpleNamespace(sequence_parallel_size=1)),
+    )
+    device = torch.device("cpu")
+    actual_generator = torch.Generator(device=device).manual_seed(123)
+    expected_generator = torch.Generator(device=device).manual_seed(123)
+
+    actual_video = ltx2_latents.prepare_video_latents(
+        pipeline,
+        batch_size=1,
+        num_channels_latents=2,
+        height=4,
+        width=6,
+        num_frames=3,
+        dtype=torch.float32,
+        device=device,
+        generator=actual_generator,
+    )
+    actual_audio, original_length, padded_length = ltx2_latents.prepare_audio_latents(
+        pipeline,
+        batch_size=1,
+        num_channels_latents=2,
+        audio_latent_length=3,
+        num_mel_bins=4,
+        dtype=torch.float32,
+        device=device,
+        generator=actual_generator,
+    )
+
+    expected_video = torch.randn((1, 12, 2), generator=expected_generator)
+    expected_audio = torch.randn((1, 3, 4), generator=expected_generator)
+
+    torch.testing.assert_close(actual_video, expected_video, rtol=0.0, atol=0.0)
+    torch.testing.assert_close(actual_audio, expected_audio, rtol=0.0, atol=0.0)
+    assert (original_length, padded_length) == (3, 3)
+
+
+def test_ltx25_distilled_two_stage_executes_custom_phase_schedules():
+    request_inputs = LTXRequestInputs(
+        prompt="prompt",
+        negative_prompt="",
+        height=64,
+        width=64,
+        num_frames=1,
+        frame_rate=24.0,
+        num_inference_steps=8,
+        guidance=LTXGuidanceSpec.positive_only(),
+        num_videos_per_prompt=1,
+        generator=None,
+        latents=None,
+        audio_latents=None,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+        prompt_attention_mask=None,
+        negative_prompt_attention_mask=None,
+        decode_timestep=0.25,
+        decode_noise_scale=0.5,
+        output_type="np",
+        max_sequence_length=16,
+    )
+    prompt_context = object()
+    phase_calls = []
+
+    def resolve_request_inputs(req, **kwargs):
+        return request_inputs
+
+    def run_phase(req, inputs, *, prompt_context=None, **kwargs):
+        phase_recipe = kwargs["phase_recipe"]
+        assert kwargs["image"] is source_image
+        phase_calls.append((phase_recipe, inputs, prompt_context))
+        if len(phase_calls) == 1:
+            assert prompt_context is None
+            assert inputs.decode_timestep == 0.25
+            assert inputs.decode_noise_scale == 0.5
+            context = prompt_context_sentinel
+            assert phase_recipe.name == "generate_lowres"
+            assert (inputs.height, inputs.width) == (32, 32)
+            assert inputs.num_inference_steps == 2
+            assert kwargs["sigmas"] == stage_1_sigmas
+            assert kwargs["noise_scale"] == 1.0
+            video = torch.ones(1, 128, 1, 1, 1)
+            audio = torch.full((1, 8, 1, 2), 2.0)
+            audio_for_next_phase = torch.full((1, 8, 1, 2), 5.0)
+        else:
+            assert prompt_context is prompt_context_sentinel
+            assert phase_recipe.name == "refine"
+            assert (inputs.height, inputs.width) == (64, 64)
+            assert inputs.latents.shape == (1, 128, 1, 2, 2)
+            torch.testing.assert_close(inputs.audio_latents, torch.full((1, 8, 1, 2), 5.0))
+            assert inputs.audio_latents_normalized
+            assert inputs.guidance_scale == 1.0
+            assert inputs.num_inference_steps == 2
+            assert kwargs["sigmas"] == stage_2_sigmas
+            assert kwargs["noise_scale"] == stage_2_sigmas[0]
+            assert inputs.decode_timestep == 0.0
+            assert inputs.decode_noise_scale is None
+            context = prompt_context
+            video = torch.full((1, 128, 1, 2, 2), 3.0)
+            audio = torch.full((1, 8, 1, 2), 4.0)
+            audio_for_next_phase = None
+        return LTXPhaseResult(
+            forward_context=SimpleNamespace(prompt_context=context),
+            video=video,
+            audio=audio,
+            audio_for_next_phase=audio_for_next_phase,
+        )
+
+    def decode_phase(phase):
+        return DiffusionOutput(output=(phase.video, phase.audio))
+
+    class FakeUpsampler(torch.nn.Module):
+        dtype = torch.float32
+
+        def forward(self, latents):
+            return latents.repeat_interleave(2, dim=-2).repeat_interleave(2, dim=-1)
+
+    prompt_context_sentinel = prompt_context
+    stage_1_sigmas = [1.0, 0.5, 0.0]
+    stage_2_sigmas = [0.8, 0.2, 0.0]
+    pipeline = object.__new__(LTX2TwoStagePipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.pipeline_recipe = LTX25_DISTILLED_TWO_STAGE_RECIPE
+    pipeline.device = torch.device("cpu")
+    pipeline.vae_spatial_compression_ratio = 32
+    pipeline.vae_temporal_compression_ratio = 8
+    pipeline.latent_upsampler = FakeUpsampler()
+    object.__setattr__(pipeline, "_resolve_request_inputs", resolve_request_inputs)
+    object.__setattr__(pipeline, "run_phase", run_phase)
+    object.__setattr__(pipeline, "decode_phase", decode_phase)
+
+    source_image = object()
+    output = pipeline.forward(
+        SimpleNamespace(sampling_params_list=[], prompts=[]),
+        image=source_image,
+        stage_1_sigmas=stage_1_sigmas,
+        stage_2_sigmas=stage_2_sigmas,
+    )
+
+    assert len(phase_calls) == 2
+    assert phase_calls[1][2] is prompt_context_sentinel
+    torch.testing.assert_close(output.output[0], torch.full((1, 128, 1, 2, 2), 3.0))
+    torch.testing.assert_close(output.output[1], torch.full((1, 8, 1, 2), 4.0))
+
+
+class TestLTXRequestParsing:
+    def test_request_resolves_custom_sigmas_from_extra_args(self):
+        from vllm_omni.diffusion.request import OmniDiffusionRequest
+        from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+        from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+        sigmas = [1.0, 0.5, 0.0]
+        req = DiffusionRequestBatch(
+            [
+                OmniDiffusionRequest(
+                    prompt="prompt",
+                    sampling_params=OmniDiffusionSamplingParams(extra_args={"sigmas": sigmas}),
+                    request_id="ltx-custom-sigmas",
+                )
+            ]
+        )
+
+        assert _make_ltx_request_pipe(LTX2Pipeline)._resolve_request_sigmas(req, None) == sigmas
+
+    def test_request_resolves_per_stage_sigmas_from_extra_args(self):
+        from vllm_omni.diffusion.request import OmniDiffusionRequest
+        from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+        from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+        stage_1_sigmas = [1.0, 0.4, 0.0]
+        req = DiffusionRequestBatch(
+            [
+                OmniDiffusionRequest(
+                    prompt="prompt",
+                    sampling_params=OmniDiffusionSamplingParams(extra_args={"stage_1_sigmas": stage_1_sigmas}),
+                    request_id="ltx-custom-phase-sigmas",
+                )
+            ]
+        )
+
+        resolved = _make_ltx_request_pipe(LTX2TwoStagePipeline)._resolve_request_phase_sigmas(
+            req, [1.0, 0.0], [0.8, 0.0]
+        )
+
+        assert resolved == (stage_1_sigmas, [0.8, 0.0])
+
+    @pytest.mark.parametrize("value", [-1, 52, 1.5, True])
+    def test_request_rejects_invalid_image_crf(self, value):
+        from vllm_omni.diffusion.request import OmniDiffusionRequest
+        from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+        from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+        req = DiffusionRequestBatch(
+            [
+                OmniDiffusionRequest(
+                    prompt="prompt",
+                    sampling_params=OmniDiffusionSamplingParams(extra_args={"image_crf": value}),
+                    request_id="ltx-invalid-image-crf",
+                )
+            ]
+        )
+
+        with pytest.raises(ValueError, match="image_crf"):
+            _resolve_request_inputs_for_test(_make_ltx_request_pipe(LTX2Pipeline), req)
+
+
+class TestLTXForwardStages:
+    def test_gemma_tokenization_matches_official_bos_and_left_padding(self):
+        pipe = _make_ltx_request_pipe(LTX2Pipeline)
+
+        class FakeTokenizer:
+            bos_token_id = 2
+            eos_token_id = 1
+            eos_token = "<eos>"
+            pad_token_id = 0
+            pad_token = "<pad>"
+            padding_side = "right"
+
+            def __init__(self):
+                self.calls = []
+
+            def __call__(self, text, **kwargs):
+                self.calls.append((text, kwargs))
+                ids = [7, 8, 9, self.eos_token_id] if text == "long" else [6, self.eos_token_id]
+                return SimpleNamespace(input_ids=ids)
+
+            def pad(self, encoded, *, padding, max_length, return_tensors, return_attention_mask):
+                assert padding == "max_length"
+                assert return_tensors == "pt"
+                assert return_attention_mask is True
+                rows = []
+                masks = []
+                for ids in encoded["input_ids"]:
+                    pad = max_length - len(ids)
+                    rows.append([self.pad_token_id] * pad + ids)
+                    masks.append([0] * pad + [1] * len(ids))
+                return SimpleNamespace(input_ids=torch.tensor(rows), attention_mask=torch.tensor(masks))
+
+        class FakeTextEncoder(torch.nn.Module):
+            dtype = torch.float32
+
+            def __init__(self):
+                super().__init__()
+                self.seen_input_ids = None
+                self.seen_attention_mask = None
+
+            def forward(self, *, input_ids, attention_mask, output_hidden_states):
+                assert output_hidden_states is True
+                self.seen_input_ids = input_ids.clone()
+                self.seen_attention_mask = attention_mask.clone()
+                return SimpleNamespace(hidden_states=(input_ids.unsqueeze(-1).float(),))
+
+        tokenizer = FakeTokenizer()
+        text_encoder = FakeTextEncoder()
+        object.__setattr__(pipe, "tokenizer", tokenizer)
+        object.__setattr__(pipe, "text_encoder", text_encoder)
+
+        pipe._get_gemma_prompt_embeds(
+            [" long ", "short"],
+            max_sequence_length=4,
+            device=torch.device("cpu"),
+            dtype=torch.float32,
+        )
+
+        torch.testing.assert_close(text_encoder.seen_input_ids, torch.tensor([[2, 7, 8, 9], [0, 2, 6, 1]]))
+        torch.testing.assert_close(text_encoder.seen_attention_mask, torch.tensor([[1, 1, 1, 1], [0, 1, 1, 1]]))
+        assert tokenizer.padding_side == "left"
+        assert [call[0] for call in tokenizer.calls] == ["long", "short"]

tests/diffusion/models/ltx2/test_ltx2_conditioning.py

@@ -3,6 +3,7 @@
 
 """Unit tests for LTX image-to-video input and conditioning behavior."""
 
+import sys
 from types import SimpleNamespace
 
 import pytest
@@ -23,7 +24,7 @@ def _make_ltx23_request_pipe(cls):
 
 
 class TestLTXImageToVideoForwardStages:
-    def test_i2v_pil_preprocessing_matches_official_center_crop(self):
+    def test_legacy_i2v_pil_preprocessing_preserves_aspect_ratio_and_center_crops(self):
         from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import _preprocess_i2v_pil_images
 
         pixels = torch.zeros(4, 8, 3, dtype=torch.uint8)
@@ -35,6 +36,136 @@ class TestLTXImageToVideoForwardStages:
 
         torch.testing.assert_close(actual, expected)
 
+    def test_ltx25_i2v_pil_preprocessing_preserves_aspect_ratio_after_crf(self, monkeypatch):
+        import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as conditioning
+
+        pixels = torch.zeros(4, 8, 3, dtype=torch.uint8)
+        pixels[:, :, 0] = torch.arange(8, dtype=torch.uint8)
+        image = Image.fromarray(pixels.numpy())
+        monkeypatch.setattr(conditioning, "_apply_image_conditioning_crf", lambda image_array, _crf: image_array)
+
+        actual = conditioning._preprocess_i2v_pil_images(image, height=4, width=4, crf=18)
+        expected = pixels[:, 2:6].permute(2, 0, 1).unsqueeze(0).float() / 127.5 - 1.0
+
+        torch.testing.assert_close(actual, expected)
+
+    def test_ltx25_i2v_applies_crf18_before_resize(self, monkeypatch):
+        import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as conditioning
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
+
+        captured_crfs = []
+        captured_prepare_kwargs = []
+
+        def fake_preprocess(image, *, height, width, crf, device, dtype):
+            del image, device
+            captured_crfs.append(crf)
+            return torch.zeros(1, 3, height, width, dtype=dtype)
+
+        def fake_prepare_latents(**kwargs):
+            captured_prepare_kwargs.append(kwargs)
+            return kwargs["image"], None
+
+        monkeypatch.setattr(conditioning, "_preprocess_i2v_pil_images", fake_preprocess)
+        pipe = object.__new__(LTX2Pipeline)
+        torch.nn.Module.__init__(pipe)
+        pipe.model_version = "2.5"
+        object.__setattr__(pipe, "transformer", SimpleNamespace(config=SimpleNamespace(in_channels=4)))
+        object.__setattr__(pipe, "prepare_latents", fake_prepare_latents)
+        request_inputs = SimpleNamespace(
+            height=4,
+            width=4,
+            num_frames=3,
+            num_videos_per_prompt=1,
+            generator=None,
+            latents=None,
+        )
+        prompt_context = SimpleNamespace(
+            batch_size=1,
+            positive_connector_prompt_embeds=torch.zeros(1, 1, 1, dtype=torch.bfloat16),
+        )
+
+        pipe._prepare_video_latents_stage(
+            request_inputs,
+            prompt_context,
+            device=torch.device("cpu"),
+            noise_scale=0.0,
+            image=Image.new("RGB", (8, 4)),
+        )
+
+        assert captured_crfs == [18]
+        assert captured_prepare_kwargs[0]["dtype"] is torch.bfloat16
+        assert captured_prepare_kwargs[0]["image"].dtype is torch.bfloat16
+
+    def test_ltx25_i2v_rejects_inputs_that_cannot_receive_crf18(self):
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
+
+        pipe = object.__new__(LTX2Pipeline)
+        torch.nn.Module.__init__(pipe)
+        pipe.model_version = "2.5"
+
+        with pytest.raises(ValueError, match="requires PIL images"):
+            pipe._prepare_video_latents_stage(
+                SimpleNamespace(),
+                SimpleNamespace(),
+                device=torch.device("cpu"),
+                noise_scale=0.0,
+                image=torch.zeros(3, 4, 4),
+            )
+
+    def test_ltx25_i2v_tensor_input_can_opt_out_of_crf18(self):
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
+
+        captured_prepare_kwargs = []
+
+        def fake_prepare_latents(**kwargs):
+            captured_prepare_kwargs.append(kwargs)
+            return kwargs["image"], None
+
+        pipe = object.__new__(LTX2Pipeline)
+        torch.nn.Module.__init__(pipe)
+        pipe.model_version = "2.5"
+        object.__setattr__(pipe, "transformer", SimpleNamespace(config=SimpleNamespace(in_channels=4)))
+        object.__setattr__(pipe, "prepare_latents", fake_prepare_latents)
+        request_inputs = SimpleNamespace(
+            height=4,
+            width=4,
+            num_frames=3,
+            num_videos_per_prompt=1,
+            generator=None,
+            latents=None,
+            image_crf=0,
+        )
+        prompt_context = SimpleNamespace(
+            batch_size=1,
+            positive_connector_prompt_embeds=torch.zeros(1, 1, 1, dtype=torch.bfloat16),
+        )
+
+        pipe._prepare_video_latents_stage(
+            request_inputs,
+            prompt_context,
+            device=torch.device("cpu"),
+            noise_scale=0.0,
+            image=torch.zeros(3, 4, 4),
+        )
+
+        assert captured_prepare_kwargs[0]["image"].shape == (1, 3, 4, 4)
+
+    def test_ltx25_i2v_crf_preserves_single_pixel_dimension(self, monkeypatch):
+        from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import _apply_image_conditioning_crf
+
+        image = torch.zeros(1, 4, 3, dtype=torch.uint8).numpy()
+        monkeypatch.setitem(sys.modules, "av", None)
+
+        assert _apply_image_conditioning_crf(image, 18) is image
+
+    def test_ltx25_i2v_reports_missing_pyav(self, monkeypatch):
+        from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import _apply_image_conditioning_crf
+
+        monkeypatch.setitem(sys.modules, "av", None)
+
+        with pytest.raises(ImportError, match="PyAV with a libx264 encoder"):
+            _apply_image_conditioning_crf(torch.zeros(4, 4, 3, dtype=torch.uint8).numpy(), 18)
+
     def test_forward_resolves_request_image_and_delegates_to_shared_recipe_runtime(self):
         from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
         from vllm_omni.diffusion.request import OmniDiffusionRequest
@@ -98,7 +229,7 @@ class TestLTXImageToVideoForwardStages:
         with pytest.raises(ValueError, match="cannot mix text-to-video and image-to-video"):
             pipe._resolve_request_image(mixed_req, None, request_inputs)
 
-    def test_denoise_timestep_kwargs_masks_video_only(self):
+    def test_denoise_timestep_kwargs_masks_video_and_keeps_audio_per_token(self):
         from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
 
         pipe = object.__new__(LTX2Pipeline)
@@ -189,13 +320,18 @@ class TestLTXImageToVideoConditioning:
         torch.testing.assert_close(conditioning_mask, torch.tensor([[1.0, 0.0, 0.0]]))
         torch.testing.assert_close(out, torch.tensor([[[40.0, 41.0], [20.0, 21.0], [30.0, 31.0]]]))
 
-    def test_ltx23_i2v_5d_latents_noise_preserves_conditioning_frame(self, monkeypatch):
+    @pytest.mark.parametrize(
+        ("model_version", "sampled_shape"),
+        [("2.3", (1, 3, 2)), ("2.5", (1, 3, 2))],
+    )
+    def test_i2v_5d_latents_noise_uses_packed_rng_layout(self, monkeypatch, model_version, sampled_shape):
         import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as ltx2_conditioning
         import vllm_omni.diffusion.models.ltx2.ltx2_latents as ltx2_latents
         from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
 
         pipe = object.__new__(LTX2Pipeline)
         torch.nn.Module.__init__(pipe)
+        pipe.model_version = model_version
         pipe.vae_spatial_compression_ratio = 1
         pipe.vae_temporal_compression_ratio = 1
         pipe.transformer_spatial_patch_size = 1
@@ -237,14 +373,20 @@ class TestLTXImageToVideoConditioning:
 
         torch.testing.assert_close(conditioning_mask, torch.tensor([[1.0, 0.0, 0.0]]))
         torch.testing.assert_close(out, torch.tensor([[[40.0, 41.0], [1.0, 1.0], [1.0, 1.0]]]))
-        assert sampled_shapes == [(1, 3, 2)]
+        assert sampled_shapes == [sampled_shape]
 
-    def test_ltx23_i2v_image_noise_is_sampled_after_packing(self, monkeypatch):
+    @pytest.mark.parametrize(
+        ("model_version", "sampled_shape"),
+        [("2.3", (1, 3, 2)), ("2.5", (1, 3, 2))],
+    )
+    def test_i2v_image_noise_uses_packed_rng_layout(self, monkeypatch, model_version, sampled_shape):
         import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as ltx2_conditioning
+        import vllm_omni.diffusion.models.ltx2.ltx2_latents as ltx2_latents
         from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
 
         pipe = object.__new__(LTX2Pipeline)
         torch.nn.Module.__init__(pipe)
+        pipe.model_version = model_version
         pipe.vae_spatial_compression_ratio = 1
         pipe.vae_temporal_compression_ratio = 1
         pipe.transformer_spatial_patch_size = 1
@@ -266,6 +408,7 @@ class TestLTXImageToVideoConditioning:
             return torch.ones(shape, device=device, dtype=dtype)
 
         monkeypatch.setattr(ltx2_conditioning, "randn_tensor", fake_randn_tensor)
+        monkeypatch.setattr(ltx2_latents, "randn_tensor", fake_randn_tensor)
 
         out, conditioning_mask = pipe.prepare_latents(
             image=torch.zeros(1, 3, 1, 1),
@@ -279,10 +422,54 @@ class TestLTXImageToVideoConditioning:
             device=torch.device("cpu"),
         )
 
-        assert sampled_shapes == [(1, 3, 2)]
+        assert sampled_shapes == [sampled_shape]
         torch.testing.assert_close(conditioning_mask, torch.tensor([[1.0, 0.0, 0.0]]))
         torch.testing.assert_close(out, torch.tensor([[[10.0, 11.0], [1.0, 1.0], [1.0, 1.0]]]))
 
+    def test_ltx25_re_noise_uses_official_scalar_noise_scale(self, monkeypatch):
+        import vllm_omni.diffusion.models.ltx2.ltx2_latents as ltx2_latents
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
+
+        pipe = object.__new__(LTX2Pipeline)
+        torch.nn.Module.__init__(pipe)
+        pipe.model_version = "2.5"
+        pipe.vae_spatial_compression_ratio = 1
+        pipe.vae_temporal_compression_ratio = 1
+        pipe.transformer_spatial_patch_size = 1
+        pipe.transformer_temporal_patch_size = 1
+        pipe.vae = SimpleNamespace(
+            latents_mean=torch.zeros(2),
+            latents_std=torch.ones(2),
+            config=SimpleNamespace(scaling_factor=1.0),
+        )
+        object.__setattr__(
+            pipe,
+            "_encode_i2v_image_latents",
+            lambda *_args, **_kwargs: torch.zeros(1, 2, 1, 1, 1, dtype=torch.bfloat16),
+        )
+        captured = {}
+
+        def fake_create_noised_state(latents, noise_scale, generator=None):
+            del generator
+            captured["noise_scale"] = noise_scale
+            return latents
+
+        monkeypatch.setattr(ltx2_latents, "create_noised_state", fake_create_noised_state)
+        pipe.prepare_latents(
+            image=torch.zeros(1, 3, 1, 1),
+            batch_size=1,
+            num_channels_latents=2,
+            height=1,
+            width=1,
+            num_frames=3,
+            noise_scale=0.909375,
+            dtype=torch.bfloat16,
+            device=torch.device("cpu"),
+            latents=torch.zeros(1, 2, 3, 1, 1, dtype=torch.bfloat16),
+        )
+
+        assert captured["noise_scale"] == 0.909375
+
     def test_ltx23_i2v_video_step_preserves_conditioning_frame(self):
         from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
 
@@ -385,6 +572,99 @@ class TestLTXImageToVideoConditioning:
         torch.testing.assert_close(audio, torch.tensor([[[-2.0]]]))
         assert pipeline.step_indices == [0, 1]
 
+    def test_ltx25_ancestral_adapter_matches_official_seeded_step_and_preserves_i2v_mask(self):
+        from vllm_omni.diffusion.models.ltx2.ltx2_denoise import LTXVideoAudioStepAdapter
+
+        device = torch.device("cpu")
+        sigmas = torch.tensor([1.0, 0.5, 0.0])
+        pipeline = SimpleNamespace(device=device, scheduler=SimpleNamespace(sigmas=sigmas))
+        mask = torch.tensor([[1.0, 0.0]])
+        adapter = LTXVideoAudioStepAdapter(
+            pipeline,
+            SimpleNamespace(sigmas=sigmas.clone()),
+            latent_num_frames=1,
+            latent_height=1,
+            latent_width=1,
+            image_conditioned=True,
+            sampler="euler_ancestral",
+            generator=torch.Generator(device=device).manual_seed(42),
+            conditioning_mask=mask,
+        )
+        video = torch.tensor([[[3.0], [2.0]]])
+        audio = torch.tensor([[[4.0]]])
+        video_velocity = torch.tensor([[[7.0], [0.25]]])
+        audio_velocity = torch.tensor([[[0.5]]])
+
+        reference_generator = torch.Generator(device=device).manual_seed(10042)
+        video_noise = torch.randn(video.shape, generator=reference_generator)
+        audio_noise = torch.randn(audio.shape, generator=reference_generator)
+
+        def official_step(sample, velocity, noise):
+            sigma = sigmas[0]
+            sigma_next = sigmas[1]
+            denoised = sample.float() - velocity.float() * sigma
+            sigma_down = sigma_next * (sigma_next / sigma)
+            ratio = sigma_down / sigma
+            deterministic = ratio * sample.float() + (1.0 - ratio) * denoised
+            alpha_next = 1.0 - sigma_next
+            alpha_down = 1.0 - sigma_down
+            coeff = (sigma_next**2 - sigma_down**2 * alpha_next**2 / alpha_down**2).sqrt()
+            return (alpha_next / alpha_down) * deterministic + noise * coeff
+
+        expected_video = official_step(video, video_velocity, video_noise)
+        expected_video[:, :1] = video[:, :1]
+        expected_audio = official_step(audio, audio_velocity, audio_noise)
+
+        ((actual_video, actual_audio),) = adapter.step(
+            (video_velocity, audio_velocity),
+            (sigmas[0], sigmas[0]),
+            (video, audio),
+        )
+        torch.testing.assert_close(actual_video, expected_video)
+        torch.testing.assert_close(actual_audio, expected_audio)
+
+        terminal_video_velocity = torch.full_like(actual_video, 0.125)
+        terminal_audio_velocity = torch.full_like(actual_audio, 0.25)
+        ((terminal_video, terminal_audio),) = adapter.step(
+            (terminal_video_velocity, terminal_audio_velocity),
+            (sigmas[1], sigmas[1]),
+            (actual_video, actual_audio),
+        )
+        expected_terminal_video = actual_video - terminal_video_velocity * sigmas[1]
+        expected_terminal_video[:, :1] = actual_video[:, :1]
+        expected_terminal_audio = actual_audio - terminal_audio_velocity * sigmas[1]
+        torch.testing.assert_close(terminal_video, expected_terminal_video)
+        torch.testing.assert_close(terminal_audio, expected_terminal_audio)
+
+    def test_ltx25_ancestral_quantizes_x0_to_model_dtype_before_fp32_step(self):
+        from vllm_omni.diffusion.models.ltx2.ltx2_denoise import _ancestral_euler_step_from_velocity
+
+        sample = torch.tensor([[[3.140625, -1.234375]]], dtype=torch.bfloat16)
+        velocity = torch.tensor([[[0.333984375, 2.71875]]], dtype=torch.bfloat16)
+        sigmas = torch.tensor([0.73, 0.41, 0.0], dtype=torch.float32)
+        generator = torch.Generator(device="cpu").manual_seed(10042)
+        reference_generator = torch.Generator(device="cpu").manual_seed(10042)
+
+        sigma, sigma_next = sigmas[:2]
+        denoised = (sample.float() - velocity.float() * sigma).to(sample.dtype).float()
+        sigma_down = sigma_next * (sigma_next / sigma)
+        sigma_down_ratio = sigma_down / sigma
+        expected = sigma_down_ratio * sample.float() + (1.0 - sigma_down_ratio) * denoised
+        alpha_next = 1.0 - sigma_next
+        alpha_down = 1.0 - sigma_down
+        coefficient = (sigma_next**2 - sigma_down**2 * alpha_next**2 / alpha_down**2).sqrt()
+        noise = torch.randn(
+            sample.shape,
+            generator=reference_generator,
+            dtype=sample.dtype,
+            device=sample.device,
+        )
+        expected = ((alpha_next / alpha_down) * expected + noise.float() * coefficient).to(sample.dtype)
+
+        actual = _ancestral_euler_step_from_velocity(sample, velocity, sigmas, 0, generator)
+
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
     def test_i2v_guidance_uses_zero_sigma_for_conditioned_tokens(self):
         from vllm_omni.diffusion.models.ltx2.ltx2_denoise import LTXDenoiseContext
         from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline

tests/diffusion/models/ltx2/test_ltx2_latents.py

@@ -202,3 +202,68 @@ def test_prepare_audio_latents_rejects_incompatible_provided_length():
             device=torch.device("cpu"),
             latents=latents,
         )
+
+
+def test_create_noised_state_matches_official_fp32_lerp():
+    latents = torch.linspace(-2, 2, 4096, dtype=torch.bfloat16).reshape(1, 32, 128)
+    expected_generator = torch.Generator().manual_seed(42)
+    noise = torch.randn(latents.shape, generator=expected_generator, dtype=latents.dtype)
+    expected = torch.lerp(latents.float(), noise.float(), torch.tensor(0.15, dtype=torch.bfloat16).float()).to(
+        latents.dtype
+    )
+
+    actual = latent_ops.create_noised_state(
+        latents,
+        torch.tensor(0.15, dtype=torch.bfloat16),
+        torch.Generator().manual_seed(42),
+    )
+
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
+def test_create_conditioned_noised_state_matches_official_two_lerps():
+    latents = torch.linspace(-2, 2, 24, dtype=torch.bfloat16).reshape(1, 3, 8)
+    clean_latents = latents.clone()
+    clean_latents[:, 0] = 3.0
+    denoise_mask = torch.tensor([[[0.0], [1.0], [1.0]]])
+    expected_generator = torch.Generator().manual_seed(42)
+    noise = torch.randn(latents.shape, generator=expected_generator, dtype=latents.dtype)
+    noised = torch.lerp(latents.float(), noise.float(), 0.909375)
+    expected = torch.lerp(clean_latents.float(), noised, denoise_mask).to(latents.dtype)
+
+    actual = latent_ops.create_conditioned_noised_state(
+        latents,
+        clean_latents,
+        denoise_mask,
+        0.909375,
+        torch.Generator().manual_seed(42),
+    )
+
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
+def test_prepare_supplied_video_latents_uses_official_token_major_layout():
+    pipeline = _make_pipeline(LTX2Pipeline)
+    supplied = torch.arange(512, dtype=torch.float32).reshape(1, 32, 16)
+    actual_video = pipeline.prepare_latents(
+        batch_size=1,
+        num_channels_latents=4,
+        height=64,
+        width=64,
+        num_frames=9,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
+        latents=supplied,
+    )
+
+    torch.testing.assert_close(actual_video, supplied, rtol=0, atol=0)
+    assert actual_video.stride()[1:] == (1, actual_video.shape[1])
+
+
+def test_clear_audio_padding_keeps_padding_outside_sampler_state():
+    updated = torch.tensor([[[1.0, 2.0], [3.0, 4.0], [99.0, -99.0]]])
+
+    actual = latent_ops.clear_audio_padding(updated, 2)
+
+    torch.testing.assert_close(actual[:, :2], updated[:, :2])
+    torch.testing.assert_close(actual[:, 2:], torch.zeros_like(actual[:, 2:]))

tests/diffusion/models/ltx2/test_ltx2_parallel.py

@@ -28,6 +28,15 @@ def test_velocity_from_x0_uses_official_host_scalar_sigma(mocker):
 class TestCFGParallelHelpers:
     """Test LTX-2.3 CFG helper math without loading model weights."""
 
+    def test_single_guidance_pass_preserves_official_token_layout(self):
+        from vllm_omni.diffusion.models.ltx2.ltx2_guidance import _repeat_batch
+
+        tokens = torch.randn(1, 16, 32).transpose(1, 2).contiguous().transpose(1, 2)
+        repeated = _repeat_batch(tokens, 1)
+
+        assert repeated is tokens
+        assert repeated.stride() == tokens.stride()
+
     def test_combine_cfg_noise_matches_x0_space_formula(self):
         from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
 

tests/diffusion/models/ltx2/test_ltx2_pipeline.py

@@ -20,6 +20,7 @@ from vllm_omni.diffusion.models.ltx2 import (
 from vllm_omni.diffusion.models.ltx2.ltx2_components import (
     LTX2_COMPONENT_PROFILE,
     LTX2_DISTILLED_COMPONENT_PROFILE,
+    LTX2_DISTILLED_ONE_STAGE_COMPONENT_PROFILE,
     LTX2_TWO_STAGE_COMPONENT_PROFILE,
     LTX23_COMPONENT_PROFILE,
     LTX23_DISTILLED_COMPONENT_PROFILE,
@@ -47,6 +48,7 @@ from vllm_omni.diffusion.models.ltx2.ltx2_guidance import (
 )
 from vllm_omni.diffusion.models.ltx2.ltx2_latents import LTXAVState
 from vllm_omni.diffusion.models.ltx2.ltx2_recipes import (
+    LTX2_DISTILLED_ONE_STAGE_RECIPE,
     LTX2_DISTILLED_TWO_STAGE_RECIPE,
     LTX2_ONE_STAGE_RECIPE,
     LTX2_TWO_STAGE_RECIPE,
@@ -64,12 +66,14 @@ from vllm_omni.diffusion.models.ltx2.ltx2_request import (
 )
 from vllm_omni.diffusion.models.ltx2.ltx2_runtime import LTXRuntime
 from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import (
+    LTX2DistilledOneStagePipeline,
     LTX2I2VDMD2Pipeline,
     LTX2Pipeline,
     LTX2T2VDMD2Pipeline,
 )
 from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import (
     LTX2DistilledPipeline,
+    LTX2DistilledTwoStagePipeline,
     LTX2TwoStagePipeline,
 )
 
@@ -125,15 +129,21 @@ def test_ltx_public_entries_share_runtime_and_keep_recipe_boundaries():
 
     assert issubclass(LTX2Pipeline, LTXRuntime)
     assert issubclass(LTX2TwoStagePipeline, LTXRuntime)
+    assert issubclass(LTX2DistilledOneStagePipeline, LTXRuntime)
+    assert issubclass(LTX2DistilledTwoStagePipeline, LTXRuntime)
     assert issubclass(LTX2DistilledPipeline, LTXRuntime)
+    assert LTX2DistilledPipeline is LTX2DistilledTwoStagePipeline
     assert LTX2Pipeline.component_profile is LTX2_COMPONENT_PROFILE
     assert LTX2Pipeline.pipeline_recipe is LTX2_ONE_STAGE_RECIPE
     assert LTX2TwoStagePipeline.component_profile is LTX2_TWO_STAGE_COMPONENT_PROFILE
     assert LTX2TwoStagePipeline.pipeline_recipe is LTX2_TWO_STAGE_RECIPE
+    assert LTX2DistilledOneStagePipeline.component_profile is LTX2_DISTILLED_ONE_STAGE_COMPONENT_PROFILE
+    assert LTX2DistilledOneStagePipeline.pipeline_recipe is LTX2_DISTILLED_ONE_STAGE_RECIPE
     assert LTX2DistilledPipeline.component_profile is LTX2_DISTILLED_COMPONENT_PROFILE
     for pipeline_cls, profile in (
         (LTX2Pipeline, LTX2_COMPONENT_PROFILE),
         (LTX2TwoStagePipeline, LTX2_TWO_STAGE_COMPONENT_PROFILE),
+        (LTX2DistilledOneStagePipeline, LTX2_DISTILLED_ONE_STAGE_COMPONENT_PROFILE),
         (LTX2DistilledPipeline, LTX2_DISTILLED_COMPONENT_PROFILE),
     ):
         assert pipeline_cls._dit_modules == list(profile.dit_modules)
@@ -254,7 +264,7 @@ def test_ltx_multistage_rejects_cache_dit_before_component_initialization(tmp_pa
 
     monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)
 
-    with pytest.raises(ValueError, match="only one-stage LTX recipes"):
+    with pytest.raises(ValueError, match="not qualified for this LTX recipe"):
         LTX2TwoStagePipeline(
             od_config=SimpleNamespace(
                 model=str(tmp_path),
@@ -709,15 +719,48 @@ def test_ltx_component_loader_installs_omni_connector_attention(monkeypatch):
 
     assert len(installed) == 2
     assert all(processor.__class__.__name__ == "_LTXConnectorAttnProcessor" for processor in installed)
-    assert all(processor.has_learned_registers for processor in installed)
     assert all(kernel["role"] == "ltx2.connector" for kernel in kernels)
     assert all(kernel["role_category"] == "self" for kernel in kernels)
     assert all(kernel["skip_sequence_parallel"] for kernel in kernels)
     assert all(kernel["disable_kv_quant"] for kernel in kernels)
+    assert all(processor.has_learned_registers for processor in installed)
+    assert not any(processor.preserve_learned_register_mask for processor in installed)
 
 
-@pytest.mark.parametrize("has_learned_registers", [False, True])
-def test_ltx_connector_attention_dispatches_through_omni_kernel(has_learned_registers):
+def test_ltx25_connector_attention_preserves_learned_register_mask(monkeypatch):
+    installed = []
+
+    class OmniAttention:
+        def __init__(self, **_kwargs):
+            pass
+
+    class Attention:
+        heads = 32
+        head_dim = 128
+        inner_kv_dim = 4096
+
+        def set_processor(self, processor):
+            installed.append(processor)
+
+    monkeypatch.setattr(ltx2_components, "OmniAttention", OmniAttention)
+    connectors = SimpleNamespace(
+        video_connector=SimpleNamespace(
+            learnable_registers=object(),
+            transformer_blocks=[SimpleNamespace(attn1=Attention())],
+        ),
+        audio_connector=SimpleNamespace(
+            learnable_registers=object(),
+            transformer_blocks=[SimpleNamespace(attn1=Attention())],
+        ),
+    )
+
+    _install_connector_attention(connectors, preserve_learned_register_mask=True)
+
+    assert all(processor.has_learned_registers for processor in installed)
+    assert all(processor.preserve_learned_register_mask for processor in installed)
+
+
+def test_ltx_connector_attention_dispatches_through_omni_kernel():
     calls = []
 
     class Backend:
@@ -732,12 +775,16 @@ def test_ltx_connector_attention_dispatches_through_omni_kernel(has_learned_regi
             calls.append((query, key, value, metadata))
             return query
 
+    class UpcastingNorm(torch.nn.Module):
+        def forward(self, tensor):
+            return tensor.float()
+
     attention = SimpleNamespace(
         heads=2,
         head_dim=4,
         inner_kv_dim=8,
-        norm_q=torch.nn.Identity(),
-        norm_k=torch.nn.Identity(),
+        norm_q=UpcastingNorm(),
+        norm_k=UpcastingNorm(),
         to_q=torch.nn.Identity(),
         to_k=torch.nn.Identity(),
         to_v=torch.nn.Identity(),
@@ -746,24 +793,29 @@ def test_ltx_connector_attention_dispatches_through_omni_kernel(has_learned_regi
         rope_type="split",
         omni_attention=OmniAttention(),
     )
-    hidden_states = torch.randn(2, 3, 8)
+    hidden_states = torch.randn(2, 3, 8, dtype=torch.bfloat16)
     additive_mask = torch.zeros(2, 1, 3, 3)
+    rotary = (
+        torch.randn(2, 2, 3, 2, dtype=torch.float32),
+        torch.randn(2, 2, 3, 2, dtype=torch.float32),
+    )
 
-    output = ltx2_components._LTXConnectorAttnProcessor(has_learned_registers=has_learned_registers)(
+    output = ltx2_components._LTXConnectorAttnProcessor()(
         attention,
         hidden_states,
         attention_mask=additive_mask,
+        query_rotary_emb=rotary,
     )
 
     query, key, value, metadata = calls[0]
     assert query.shape == key.shape == value.shape == (2, 3, 2, 4)
-    if has_learned_registers:
-        assert metadata is None
-    else:
-        assert metadata.attn_mask.shape == (2, 3)
-        assert metadata.attn_mask.dtype == torch.bool
-        assert metadata.attn_mask.all()
-    torch.testing.assert_close(output, hidden_states)
+    assert query.dtype == key.dtype == value.dtype == torch.bfloat16
+    assert metadata.attn_mask.shape == (2, 3)
+    assert metadata.attn_mask.dtype == torch.bool
+    assert metadata.attn_mask.all()
+    expected_rotary = tuple(component.to(torch.bfloat16) for component in rotary)
+    expected = ltx2_components.apply_split_rotary_emb(hidden_states, expected_rotary, head_dim=4)
+    torch.testing.assert_close(output, expected, rtol=0, atol=0)
 
 
 def test_ltx2_distilled_pipeline_shares_recipe_runtime():
@@ -1707,6 +1759,7 @@ class TestLTXForwardStages:
         assert seen["request_inputs"].output_type == "latent"
         assert seen["kwargs"] == {
             "request_sigmas": request_sigmas,
+            "request_phase_sigmas": None,
             "image": None,
         }
 
@@ -1824,6 +1877,16 @@ class TestRegistryIntegration:
             "pipeline_ltx2_two_stage",
             "LTX2TwoStagePipeline",
         )
+        assert _DIFFUSION_MODELS["LTX2DistilledOneStagePipeline"] == (
+            "ltx2",
+            "pipeline_ltx2",
+            "LTX2DistilledOneStagePipeline",
+        )
+        assert _DIFFUSION_MODELS["LTX2DistilledTwoStagePipeline"] == (
+            "ltx2",
+            "pipeline_ltx2_two_stage",
+            "LTX2DistilledTwoStagePipeline",
+        )
         assert _DIFFUSION_MODELS["LTX2DistilledPipeline"] == (
             "ltx2",
             "pipeline_ltx2_two_stage",
@@ -1848,6 +1911,8 @@ class TestRegistryIntegration:
         expected = [
             "LTX2Pipeline",
             "LTX2TwoStagePipeline",
+            "LTX2DistilledOneStagePipeline",
+            "LTX2DistilledTwoStagePipeline",
             "LTX2DistilledPipeline",
             "LTX2T2VDMD2Pipeline",
             "LTX2I2VDMD2Pipeline",
@@ -1861,6 +1926,8 @@ class TestRegistryIntegration:
         [
             "LTX2Pipeline",
             "LTX2TwoStagePipeline",
+            "LTX2DistilledOneStagePipeline",
+            "LTX2DistilledTwoStagePipeline",
             "LTX2DistilledPipeline",
             "LTX2T2VDMD2Pipeline",
             "LTX2I2VDMD2Pipeline",

tests/diffusion/models/ltx2/test_ltx2_transformer.py

@@ -7,18 +7,54 @@ import pytest
 import torch
 from torch import nn
 
-from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig
+from vllm_omni.diffusion.models.ltx2 import ltx2_transformer
 from vllm_omni.diffusion.models.ltx2.ltx2_transformer import (
-    LTX2AudioVideoAttnProcessor,
+    LTX2FeedForward,
     LTX2VideoTransformer3DModel,
-    _LTX2ParallelAttention,
+    LTX2VideoTransformerBlock,
     _make_rms_norm,
+    apply_interleaved_rotary_emb,
+    apply_keyframes_absolute_embedding,
+    apply_split_rotary_emb,
 )
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 
 
+def test_ltx_interleaved_rope_matches_official_model_dtype_arithmetic():
+    torch.manual_seed(0)
+    x = torch.randn(2, 3, 8, dtype=torch.bfloat16)
+    cos = torch.randn(2, 3, 8, dtype=x.dtype)
+    sin = torch.randn(2, 3, 8, dtype=x.dtype)
+    x_real, x_imag = x.unflatten(2, (4, 2)).unbind(-1)
+    rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(2)
+    expected = x * cos + rotated * sin
+
+    actual = apply_interleaved_rotary_emb(x, (cos, sin))
+
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
+def test_ltx_split_rope_matches_official_model_dtype_arithmetic():
+    torch.manual_seed(1)
+    x = torch.randn(2, 3, 8, dtype=torch.bfloat16)
+    cos = torch.randn(2, 3, 4, dtype=x.dtype)
+    sin = torch.randn(2, 3, 4, dtype=x.dtype)
+    split_x = x.reshape(2, 3, 2, 4)
+    first_x, second_x = split_x.chunk(2, dim=-2)
+    cos_u = cos.unsqueeze(-2)
+    sin_u = sin.unsqueeze(-2)
+    expected = split_x * cos_u
+    expected[..., :1, :].addcmul_(-sin_u, second_x)
+    expected[..., 1:, :].addcmul_(sin_u, first_x)
+    expected = expected.reshape_as(x)
+
+    actual = apply_split_rotary_emb(x, (cos, sin), head_dim=8)
+
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
 def test_ltx_rms_norm_no_affine_matches_official_torch_module():
     norm = _make_rms_norm(8, eps=1e-6, elementwise_affine=False)
 
@@ -36,11 +72,53 @@ def test_ltx_rms_norm_affine_weight_remains_parameter():
     assert "weight" in norm.state_dict()
 
 
-def test_ltx_transformer_has_separate_cfg_cache_dit_config():
+def test_ltx_attention_restores_projection_dtype_after_upcasting_qk_norms():
+    calls = []
+
+    class UpcastingNorm(nn.Module):
+        def forward(self, tensor):
+            return tensor.float()
+
+    class RecordingAttention:
+        def __call__(self, query, key, value, metadata):
+            calls.append((query, key, value, metadata))
+            return query
+
+    attention = SimpleNamespace(
+        heads=2,
+        head_dim=4,
+        to_qkv=None,
+        to_q=nn.Identity(),
+        to_k=nn.Identity(),
+        to_v=nn.Identity(),
+        norm_q=UpcastingNorm(),
+        norm_k=UpcastingNorm(),
+        to_gate_logits=None,
+        attn=RecordingAttention(),
+        to_out=(nn.Identity(), nn.Identity()),
+        rope_type="split",
+    )
+    hidden_states = torch.randn(2, 3, 8, dtype=torch.bfloat16)
+    rotary = (
+        torch.ones(2, 2, 3, 2, dtype=torch.float32),
+        torch.zeros(2, 2, 3, 2, dtype=torch.float32),
+    )
+
+    processor = ltx2_transformer.LTX2AudioVideoAttnProcessor()
+    processor._slice_rope_for_tp = lambda rope, _: rope
+    output = processor(attention, hidden_states, query_rotary_emb=rotary)
+
+    query, key, value, metadata = calls[0]
+    assert query.dtype == key.dtype == value.dtype == torch.bfloat16
+    assert metadata is None
+    torch.testing.assert_close(output, hidden_states)
+
+
+def test_ltx_transformer_has_fused_cfg_cache_dit_config():
     adapter_config = getattr(LTX2VideoTransformer3DModel, "_cache_dit_adapter_config")
 
     assert isinstance(adapter_config, CacheDiTAdapterConfig)
-    assert adapter_config.has_separate_cfg
+    assert not adapter_config.has_separate_cfg
 
 
 def test_ltx_transformer_exposes_hsdp_shard_conditions_for_blocks():
@@ -76,8 +154,9 @@ def test_ltx_sp_plan_shards_video_and_audio_timesteps_together(rope_type):
 
 
 def test_ltx_sp_keeps_cross_attention_key_padding_mask(monkeypatch):
-    monkeypatch.setattr(LTX2AudioVideoAttnProcessor, "_is_sp_enabled", staticmethod(lambda: True))
-    processor = LTX2AudioVideoAttnProcessor()
+    processor_cls = ltx2_transformer.LTX2AudioVideoAttnProcessor
+    monkeypatch.setattr(processor_cls, "_is_sp_enabled", staticmethod(lambda: True))
+    processor = processor_cls()
     hidden_states = torch.zeros(1, 3, 4)
     encoder_hidden_states = torch.zeros(1, 4, 4)
     key_padding_mask = torch.tensor([[True, True, True, False]])
@@ -113,7 +192,7 @@ def _fake_ltx_parallel_attention():
 
         return forward
 
-    layer = object.__new__(_LTX2ParallelAttention)
+    layer = object.__new__(ltx2_transformer._LTX2ParallelAttention)
     nn.Module.__init__(layer)
     layer.attn_backend = Backend
     layer.attention = SimpleNamespace(forward=run("native"))
@@ -123,33 +202,281 @@ def _fake_ltx_parallel_attention():
 
 
 @pytest.mark.parametrize(("query_length", "key_length"), [(4, 4), (3, 4)])
-def test_ltx_masked_cross_attention_uses_sdpa(query_length, key_length):
+def test_ltx_masked_attention_uses_sdpa(query_length, key_length):
     layer, calls = _fake_ltx_parallel_attention()
     query = torch.zeros(1, query_length, 2, 4, dtype=torch.bfloat16)
     key = value = torch.zeros(1, key_length, 2, 4, dtype=torch.bfloat16)
     key_padding_mask = torch.tensor([[True] * (key_length - 1) + [False]])
-    metadata = AttentionMetadata(attn_mask=key_padding_mask[:, None, None, :])
+    metadata = ltx2_transformer.AttentionMetadata(attn_mask=key_padding_mask[:, None, None, :])
 
     layer._run_local_attention(query, key, value, metadata)
 
     assert calls == ["sdpa"]
 
 
-def test_ltx_masked_self_attention_uses_sdpa():
-    layer, calls = _fake_ltx_parallel_attention()
-    query = key = value = torch.zeros(1, 4, 2, 4, dtype=torch.bfloat16)
-    key_padding_mask = torch.tensor([[True, True, True, False]])
-    metadata = AttentionMetadata(attn_mask=key_padding_mask[:, None, None, :])
-
-    layer._run_local_attention(query, key, value, metadata)
-
-    assert calls == ["sdpa"]
-
-
-def test_ltx_unmasked_cross_attention_keeps_native_backend():
+def test_ltx_unmasked_attention_keeps_native_backend():
     layer, calls = _fake_ltx_parallel_attention()
     query = key = value = torch.zeros(1, 4, 2, 4, dtype=torch.bfloat16)
 
     layer._run_local_attention(query, key, value, None)
 
     assert calls == ["native"]
+
+
+def test_ltx25_feed_forward_bias_flag_controls_both_linears(monkeypatch):
+    class FakeParallelLinear(nn.Linear):
+        def __init__(self, dim_in, dim_out, *, bias=True, **kwargs):
+            super().__init__(dim_in, dim_out, bias=bias)
+
+    monkeypatch.setattr(ltx2_transformer, "ColumnParallelLinear", FakeParallelLinear)
+    monkeypatch.setattr(ltx2_transformer, "RowParallelLinear", FakeParallelLinear)
+
+    without_bias = LTX2FeedForward(8, inner_dim=16, bias=False)
+    with_bias = LTX2FeedForward(8, inner_dim=16, bias=True)
+
+    assert without_bias.net[0].proj.bias is None
+    assert without_bias.net[2].bias is None
+    assert with_bias.net[0].proj.bias is not None
+    assert with_bias.net[2].bias is not None
+
+
+def _build_tiny_ltx_transformer(monkeypatch, **overrides):
+    class FakeBlock(nn.Module):
+        def __init__(self, **kwargs):
+            super().__init__()
+            self.init_kwargs = kwargs
+
+        def forward(self, hidden_states, audio_hidden_states, *_args, **_kwargs):
+            return hidden_states, audio_hidden_states
+
+    monkeypatch.setattr(ltx2_transformer, "LTX2VideoTransformerBlock", FakeBlock)
+    kwargs = {
+        "in_channels": 4,
+        "out_channels": 4,
+        "num_attention_heads": 2,
+        "attention_head_dim": 4,
+        "cross_attention_dim": 8,
+        "vae_scale_factors": (1, 1, 1),
+        "pos_embed_max_pos": 4,
+        "base_height": 8,
+        "base_width": 8,
+        "audio_in_channels": 4,
+        "audio_out_channels": 4,
+        "audio_num_attention_heads": 2,
+        "audio_attention_head_dim": 4,
+        "audio_cross_attention_dim": 8,
+        "audio_scale_factor": 1,
+        "audio_pos_embed_max_pos": 4,
+        "audio_sampling_rate": 16,
+        "audio_hop_length": 4,
+        "num_layers": 1,
+        "caption_channels": 8,
+        "use_prompt_embeddings": False,
+        "cross_attn_mod": True,
+        "audio_cross_attn_mod": True,
+    }
+    kwargs.update(overrides)
+    return LTX2VideoTransformer3DModel(**kwargs)
+
+
+def test_ltx25_keyframe_absolute_embedding_matches_official():
+    hidden_states = torch.arange(32, dtype=torch.bfloat16).view(1, 4, 8)
+    keyframes_mask = torch.tensor([[[1.0], [1.0], [0.0], [0.0]]])
+    embedding = torch.linspace(-0.5, 0.5, 8, dtype=torch.bfloat16).view(1, 8)
+
+    actual = apply_keyframes_absolute_embedding(hidden_states, keyframes_mask, embedding)
+    expected = hidden_states + (keyframes_mask > 0).to(torch.bfloat16) * embedding
+
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+    assert apply_keyframes_absolute_embedding(hidden_states, None, embedding) is hidden_states
+
+
+def test_ltx25_constructor_flags_and_native_weight_loader(monkeypatch):
+    legacy_model = _build_tiny_ltx_transformer(monkeypatch)
+
+    assert legacy_model.config.ff_bias is True
+    assert legacy_model.config.audio_ff_bias is True
+    assert legacy_model.config.use_prompt_adaln_single is True
+    assert legacy_model.config.use_keyframes_abs_pos_embedding is False
+    assert hasattr(legacy_model, "prompt_adaln")
+    assert hasattr(legacy_model, "audio_prompt_adaln")
+    assert not hasattr(legacy_model, "keyframes_abs_pos_embedding")
+    assert legacy_model.transformer_blocks[0].init_kwargs["ff_bias"] is True
+    assert legacy_model.transformer_blocks[0].init_kwargs["audio_ff_bias"] is True
+
+    ltx25_model = _build_tiny_ltx_transformer(
+        monkeypatch,
+        ff_bias=False,
+        audio_ff_bias=True,
+        use_prompt_adaln_single=False,
+        use_keyframes_abs_pos_embedding=True,
+    )
+
+    assert ltx25_model.keyframes_abs_pos_embedding.shape == (1, 8)
+    torch.testing.assert_close(ltx25_model.keyframes_abs_pos_embedding, torch.zeros(1, 8))
+    assert not hasattr(ltx25_model, "prompt_adaln")
+    assert not hasattr(ltx25_model, "audio_prompt_adaln")
+    assert ltx25_model.transformer_blocks[0].init_kwargs["ff_bias"] is False
+    assert ltx25_model.transformer_blocks[0].init_kwargs["audio_ff_bias"] is True
+
+    monkeypatch.setattr(ltx2_transformer, "get_tensor_model_parallel_world_size", lambda: 1)
+    expected = torch.arange(8, dtype=torch.float32).view(1, 8)
+    loaded = ltx25_model.load_weights([("keyframes_abs_pos_embedding", expected)])
+
+    assert loaded == {"keyframes_abs_pos_embedding"}
+    torch.testing.assert_close(ltx25_model.keyframes_abs_pos_embedding, expected)
+
+
+def test_ltx_cross_attention_modulations_use_official_timesteps(monkeypatch):
+    class CaptureAda(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.calls = []
+
+        def forward(self, timestep, *, batch_size, hidden_dtype):
+            self.calls.append(timestep.detach().clone())
+            embedding = torch.zeros(timestep.numel(), 8, dtype=hidden_dtype)
+            return embedding, embedding
+
+    class CaptureRope(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.output_dtypes = []
+
+        def forward(self, _coords, *, device, out_dtype=None):
+            self.output_dtypes.append(out_dtype)
+            empty = torch.empty(0, device=device, dtype=out_dtype)
+            return empty, empty
+
+    model = _build_tiny_ltx_transformer(monkeypatch)
+    rope = CaptureRope()
+    model.rope = model.audio_rope = model.cross_attn_rope = model.cross_attn_audio_rope = rope
+    model.prompt_modulation = False
+    video_scale_shift = CaptureAda()
+    video_gate = CaptureAda()
+    audio_scale_shift = CaptureAda()
+    audio_gate = CaptureAda()
+    model.av_cross_attn_video_scale_shift = video_scale_shift
+    model.av_cross_attn_video_a2v_gate = video_gate
+    model.av_cross_attn_audio_scale_shift = audio_scale_shift
+    model.av_cross_attn_audio_v2a_gate = audio_gate
+
+    model(
+        hidden_states=torch.zeros(1, 2, 4),
+        audio_hidden_states=torch.zeros(1, 1, 4),
+        encoder_hidden_states=torch.zeros(1, 1, 8),
+        audio_encoder_hidden_states=torch.zeros(1, 1, 8),
+        timestep=torch.tensor([[0.0, 2.0]]),
+        audio_timestep=torch.tensor([[5.0]]),
+        sigma=torch.tensor([2.0]),
+        audio_sigma=torch.tensor([5.0]),
+        video_coords=torch.zeros(1, 3, 2, 2),
+        audio_coords=torch.zeros(1, 1, 1, 2),
+    )
+
+    torch.testing.assert_close(video_scale_shift.calls[0], torch.tensor([0.0, 2.0]))
+    torch.testing.assert_close(video_gate.calls[0], torch.tensor([5.0]))
+    torch.testing.assert_close(audio_scale_shift.calls[0], torch.tensor([5.0]))
+    torch.testing.assert_close(audio_gate.calls[0], torch.tensor([2.0]))
+    assert len(rope.output_dtypes) == 4
+    assert all(dtype is torch.float32 for dtype in rope.output_dtypes)
+
+
+def test_ltx25_static_prompt_table_is_used_without_prompt_adaln():
+    class CaptureAttention(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.encoder_hidden_states = None
+
+        def forward(self, hidden_states, encoder_hidden_states=None, **kwargs):
+            if encoder_hidden_states is not None:
+                self.encoder_hidden_states = encoder_hidden_states.detach().clone()
+            return torch.zeros_like(hidden_states)
+
+    class ZeroOutput(nn.Module):
+        def forward(self, hidden_states):
+            return torch.zeros_like(hidden_states)
+
+    block = object.__new__(LTX2VideoTransformerBlock)
+    nn.Module.__init__(block)
+    block.video_cross_attn_adaln = True
+    block.audio_cross_attn_adaln = True
+    block.cross_attn_adaln = True
+    block.scale_shift_table = nn.Parameter(torch.zeros(9, 2))
+    block.audio_scale_shift_table = nn.Parameter(torch.zeros(9, 2))
+    block.prompt_scale_shift_table = nn.Parameter(torch.tensor([[1.0, 2.0], [3.0, 4.0]]))
+    block.audio_prompt_scale_shift_table = nn.Parameter(torch.tensor([[0.5, 1.0], [1.5, 2.0]]))
+    block.norm1 = nn.Identity()
+    block.audio_norm1 = nn.Identity()
+    block.norm2 = nn.Identity()
+    block.audio_norm2 = nn.Identity()
+    block.norm3 = nn.Identity()
+    block.audio_norm3 = nn.Identity()
+    block.attn1 = CaptureAttention()
+    block.audio_attn1 = CaptureAttention()
+    block.attn2 = CaptureAttention()
+    block.audio_attn2 = CaptureAttention()
+    block.ff = ZeroOutput()
+    block.audio_ff = ZeroOutput()
+
+    hidden_states = torch.ones(1, 1, 2)
+    audio_hidden_states = torch.ones(1, 1, 2)
+    encoder_hidden_states = torch.ones(1, 1, 2)
+    audio_encoder_hidden_states = torch.ones(1, 1, 2)
+    block(
+        hidden_states,
+        audio_hidden_states,
+        encoder_hidden_states,
+        audio_encoder_hidden_states,
+        temb=torch.zeros(1, 1, 18),
+        temb_audio=torch.zeros(1, 1, 18),
+        temb_ca_scale_shift=torch.empty(0),
+        temb_ca_audio_scale_shift=torch.empty(0),
+        temb_ca_gate=torch.empty(0),
+        temb_ca_audio_gate=torch.empty(0),
+        temb_prompt=None,
+        temb_prompt_audio=None,
+        use_a2v_cross_attention=False,
+        use_v2a_cross_attention=False,
+    )
+
+    torch.testing.assert_close(block.attn2.encoder_hidden_states, torch.tensor([[[5.0, 7.0]]]))
+    torch.testing.assert_close(block.audio_attn2.encoder_hidden_states, torch.tensor([[[3.0, 4.0]]]))
+
+
+def test_ltx_gated_attention_uses_tp_sharded_projection(monkeypatch):
+    column_calls = []
+
+    class FakeColumnParallelLinear(nn.Linear):
+        def __init__(self, input_size, output_size, *, bias=True, **kwargs):
+            column_calls.append((output_size, kwargs))
+            super().__init__(input_size, output_size, bias=bias)
+
+    class FakeRowParallelLinear(nn.Linear):
+        def __init__(self, input_size, output_size, *, bias=True, **_kwargs):
+            super().__init__(input_size, output_size, bias=bias)
+
+    class FakeParallelAttention(nn.Module):
+        def __init__(self, **_kwargs):
+            super().__init__()
+
+    monkeypatch.setattr(ltx2_transformer, "ColumnParallelLinear", FakeColumnParallelLinear)
+    monkeypatch.setattr(ltx2_transformer, "RowParallelLinear", FakeRowParallelLinear)
+    monkeypatch.setattr(ltx2_transformer, "_LTX2ParallelAttention", FakeParallelAttention)
+    monkeypatch.setattr(ltx2_transformer, "get_tensor_model_parallel_world_size", lambda: 2)
+
+    attention = ltx2_transformer.LTX2Attention(
+        query_dim=8,
+        heads=4,
+        kv_heads=4,
+        dim_head=2,
+        qk_norm="rms_norm",
+        apply_gated_attention=True,
+        pack_qkv=False,
+    )
+
+    gate_output_size, gate_kwargs = column_calls[-1]
+    assert gate_output_size == 4
+    assert gate_kwargs["gather_output"] is False
+    assert gate_kwargs["return_bias"] is False
+    assert attention.query_num_heads == 2

tests/e2e/accuracy/ltx/run_ltx25_reference.py

@@ -0,0 +1,591 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Generate raw LTX video and audio outputs with the official or Omni runtime."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+from PIL import Image
+
+
+def _parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--backend", choices=("official", "omni"), required=True)
+    parser.add_argument("--request", type=Path, required=True)
+    parser.add_argument("--output-dir", type=Path, required=True)
+    parser.add_argument("--model")
+    parser.add_argument("--model-class-name")
+    parser.add_argument("--official-root", type=Path)
+    parser.add_argument("--checkpoint", type=Path)
+    parser.add_argument("--gemma-root", type=Path)
+    parser.add_argument("--transformer-path", type=Path)
+    parser.add_argument("--text-encoder-path", type=Path)
+    parser.add_argument("--video-vae-path", type=Path)
+    parser.add_argument("--audio-vae-path", type=Path)
+    parser.add_argument("--spatial-upsampler-path", type=Path)
+    parser.add_argument("--distilled-lora-path", type=Path)
+    parser.add_argument("--connector-model", type=Path)
+    parser.add_argument(
+        "--official-pipeline",
+        choices=("distilled_two_stage", "full_one_stage", "full_two_stage"),
+        default="distilled_two_stage",
+    )
+    parser.add_argument("--enable-layerwise-offload", action="store_true")
+    return parser.parse_args()
+
+
+def _save_outputs(
+    output_dir: Path,
+    *,
+    video: torch.Tensor,
+    audio: torch.Tensor,
+    audio_sample_rate: int,
+    fps: float,
+    metadata: dict[str, Any],
+) -> None:
+    output_dir.mkdir(parents=True, exist_ok=True)
+    video = video.detach().float().cpu().clamp(0.0, 1.0)
+    audio = audio.detach().float().cpu()
+    np.save(output_dir / "video.npy", video.numpy())
+    np.save(output_dir / "audio.npy", audio.numpy())
+
+    from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
+
+    frames_u8 = video.mul(255.0).round().to(torch.uint8).numpy()
+    audio_array = audio.numpy()
+    while audio_array.ndim > 2 and audio_array.shape[0] == 1:
+        audio_array = audio_array[0]
+    (output_dir / "output.mp4").write_bytes(
+        mux_video_audio_bytes(
+            frames_u8,
+            audio_array,
+            fps=float(fps),
+            audio_sample_rate=int(audio_sample_rate),
+            video_codec_options={"preset": "ultrafast"},
+        )
+    )
+
+    frame_indices = sorted({0, video.shape[0] // 2, video.shape[0] - 1})
+    for index in frame_indices:
+        frame = video[index].mul(255.0).round().to(torch.uint8).numpy()
+        Image.fromarray(frame).save(output_dir / f"frame_{index:04d}.png")
+
+    metadata.update(
+        {
+            "video_shape": list(video.shape),
+            "audio_shape": list(audio.shape),
+            "audio_sample_rate": int(audio_sample_rate),
+            "frame_indices": frame_indices,
+            "mp4": "output.mp4",
+        }
+    )
+    (output_dir / "metadata.json").write_text(json.dumps(metadata, indent=2) + "\n")
+
+
+def _reset_accelerator_peak_memory_stats() -> None:
+    if torch.accelerator.is_available():
+        torch.accelerator.reset_peak_memory_stats()
+
+
+def _accelerator_peak_memory_stats() -> dict[str, float]:
+    if not torch.accelerator.is_available():
+        return {}
+    torch.accelerator.synchronize()
+    mib = 1024**2
+    return {
+        "peak_memory_allocated_mb": torch.accelerator.max_memory_allocated() / mib,
+        "peak_memory_reserved_mb": torch.accelerator.max_memory_reserved() / mib,
+    }
+
+
+def _insert_official_paths(official_root: Path) -> None:
+    for relative_path in ("packages/ltx-core/src", "packages/ltx-pipelines/src"):
+        path = str((official_root / relative_path).resolve())
+        if path not in sys.path:
+            sys.path.insert(0, path)
+
+
+def _configure_official_sdpa(pipeline: Any) -> None:
+    """Pin official connector and denoiser attention to cuDNN, with MATH fallback."""
+    from ltx_core.loader.attention_ops import set_attention_module_op
+    from ltx_core.model.transformer.attention import PytorchAttention
+    from torch.nn.attention import SDPBackend
+
+    attention = PytorchAttention(priority=[SDPBackend.CUDNN_ATTENTION, SDPBackend.MATH])
+    module_op = set_attention_module_op(
+        attention=attention,
+        masked_attention=attention,
+    )
+    transformer_owners = [
+        owner
+        for attribute in ("stage", "stage_1", "stage_2")
+        if (owner := getattr(pipeline, attribute, None)) is not None
+    ]
+    owners_and_attributes = [
+        *((owner, "_transformer_builder") for owner in transformer_owners),
+        (pipeline.prompt_encoder, "_embeddings_processor_builder"),
+    ]
+    for owner, attribute in owners_and_attributes:
+        builder = getattr(owner, attribute)
+        setattr(
+            owner,
+            attribute,
+            builder.with_module_ops((*builder.module_ops, module_op)),
+        )
+
+
+def _use_diffusers_connector_weights(prompt_encoder: Any, model: Path) -> None:
+    """Run official connector code with the checkpoint under test."""
+    from safetensors import safe_open
+
+    connector_root = model / "connectors"
+    shards = sorted(connector_root.glob("*.safetensors"))
+    if not shards:
+        raise ValueError(f"No connector safetensors found under {connector_root}")
+
+    original_build_embeddings_processor = prompt_encoder._build_embeddings_processor
+
+    def build_embeddings_processor():
+        processor = original_build_embeddings_processor()
+        parameters = dict(processor.named_parameters())
+        expected = {name for name in parameters if name.startswith(("video_connector.", "audio_connector."))}
+        loaded: set[str] = set()
+        with torch.no_grad():
+            for shard in shards:
+                with safe_open(str(shard), framework="pt", device="cpu") as weights:
+                    for source_name in weights.keys():
+                        if not source_name.startswith(("video_connector.", "audio_connector.")):
+                            continue
+                        target_name = (
+                            source_name.replace(".transformer_blocks.", ".transformer_1d_blocks.")
+                            .replace(".attn1.norm_q.", ".attn1.q_norm.")
+                            .replace(".attn1.norm_k.", ".attn1.k_norm.")
+                        )
+                        parameter = parameters[target_name]
+                        parameter.copy_(weights.get_tensor(source_name).to(parameter.device, parameter.dtype))
+                        loaded.add(target_name)
+        if loaded != expected:
+            missing = sorted(expected - loaded)
+            unexpected = sorted(loaded - expected)
+            raise ValueError(f"Connector weight mismatch: missing={missing}, unexpected={unexpected}")
+        return processor
+
+    prompt_encoder._build_embeddings_processor = build_embeddings_processor
+
+
+@torch.inference_mode()
+def _run_official(args: argparse.Namespace, request: dict[str, Any]) -> None:
+    if args.official_root is None:
+        raise ValueError("Official backend requires --official-root")
+    _insert_official_paths(args.official_root)
+    # Keep Gemma on the same SDPA path in both subprocesses. DiT attention is
+    # still pinned independently to the explicit cuDNN backend below.
+    torch.backends.cuda.enable_cudnn_sdp(False)
+    _reset_accelerator_peak_memory_stats()
+
+    from ltx_pipelines.utils.args import ImageConditioningInput
+    from ltx_pipelines.utils.model_paths import ModelPaths
+    from ltx_pipelines.utils.types import OffloadMode
+
+    offload_mode = OffloadMode.CPU if args.enable_layerwise_offload else OffloadMode.NONE
+    image_path = request.get("image")
+    images = (
+        []
+        if image_path is None
+        else [
+            ImageConditioningInput(
+                path=str(image_path),
+                frame_idx=0,
+                strength=1.0,
+                crf=request.get("image_crf", 18),
+            )
+        ]
+    )
+
+    if args.transformer_path is not None:
+        required = {
+            "--text-encoder-path": args.text_encoder_path,
+            "--video-vae-path": args.video_vae_path,
+            "--audio-vae-path": args.audio_vae_path,
+        }
+        if args.official_pipeline in {"distilled_two_stage", "full_two_stage"}:
+            required["--spatial-upsampler-path"] = args.spatial_upsampler_path
+        if args.official_pipeline == "full_two_stage":
+            required["--distilled-lora-path"] = args.distilled_lora_path
+        missing = [name for name, value in required.items() if value is None]
+        if missing:
+            raise ValueError(f"Official LTX-2.5 split backend is missing: {', '.join(missing)}")
+
+        model_paths = ModelPaths.from_split(
+            transformer_path=str(args.transformer_path),
+            text_encoder_path=str(args.text_encoder_path),
+            video_vae_path=str(args.video_vae_path),
+            audio_vae_path=str(args.audio_vae_path),
+        )
+        if args.official_pipeline == "distilled_two_stage":
+            from ltx_pipelines.distilled import DistilledPipeline
+
+            pipeline = DistilledPipeline(
+                model_paths=model_paths,
+                spatial_upsampler_path=str(args.spatial_upsampler_path),
+                loras=(),
+                offload_mode=offload_mode,
+            )
+            if args.connector_model is not None:
+                _use_diffusers_connector_weights(pipeline.prompt_encoder, args.connector_model)
+            _configure_official_sdpa(pipeline)
+            video, audio, _, _ = pipeline(
+                prompt=request["prompt"],
+                seed=request["seed"],
+                height=request["height"],
+                width=request["width"],
+                num_frames=request["num_frames"],
+                frame_rate=request["fps"],
+                images=images,
+            )
+            pipeline_name = "DistilledPipeline"
+        elif args.official_pipeline == "full_two_stage":
+            from ltx_core.components.guiders import MultiModalGuiderParams
+            from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps
+            from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline
+
+            pipeline = TI2VidTwoStagesPipeline(
+                model_paths=model_paths,
+                distilled_lora=[
+                    LoraPathStrengthAndSDOps(
+                        path=str(args.distilled_lora_path),
+                        strength=1.0,
+                        sd_ops=LTXV_LORA_COMFY_RENAMING_MAP,
+                    )
+                ],
+                spatial_upsampler_path=str(args.spatial_upsampler_path),
+                loras=[],
+                offload_mode=offload_mode,
+            )
+            if args.connector_model is not None:
+                _use_diffusers_connector_weights(pipeline.prompt_encoder, args.connector_model)
+            _configure_official_sdpa(pipeline)
+            video, audio, _, _ = pipeline(
+                prompt=request["prompt"],
+                negative_prompt=request["negative_prompt"],
+                seed=request["seed"],
+                height=request["height"],
+                width=request["width"],
+                num_frames=request["num_frames"],
+                frame_rate=request["fps"],
+                num_inference_steps=request["num_inference_steps"],
+                video_guider_params=MultiModalGuiderParams(
+                    cfg_scale=request["video_cfg_scale"],
+                    stg_scale=request["video_stg_scale"],
+                    rescale_scale=request["video_rescale_scale"],
+                    modality_scale=request["video_modality_scale"],
+                    skip_step=0,
+                    stg_blocks=request["video_stg_blocks"],
+                ),
+                audio_guider_params=MultiModalGuiderParams(
+                    cfg_scale=request["audio_cfg_scale"],
+                    stg_scale=request["audio_stg_scale"],
+                    rescale_scale=request["audio_rescale_scale"],
+                    modality_scale=request["audio_modality_scale"],
+                    skip_step=0,
+                    stg_blocks=request["audio_stg_blocks"],
+                ),
+                images=images,
+                max_batch_size=4,
+                stage_1_sigmas=torch.tensor(request["stage_1_sigmas"], dtype=torch.float32),
+                stage_2_sigmas=torch.tensor(request["stage_2_sigmas"], dtype=torch.float32),
+            )
+            pipeline_name = "TI2VidTwoStagesPipeline"
+        else:
+            from ltx_core.components.guiders import MultiModalGuiderParams
+            from ltx_pipelines.ti2vid_one_stage import TI2VidOneStagePipeline
+
+            pipeline = TI2VidOneStagePipeline(
+                model_paths=model_paths,
+                loras=(),
+                offload_mode=offload_mode,
+            )
+            if args.connector_model is not None:
+                _use_diffusers_connector_weights(pipeline.prompt_encoder, args.connector_model)
+            _configure_official_sdpa(pipeline)
+            request_sigmas = request.get("sigmas")
+            video, audio, _ = pipeline(
+                prompt=request["prompt"],
+                negative_prompt=request["negative_prompt"],
+                seed=request["seed"],
+                height=request["height"],
+                width=request["width"],
+                num_frames=request["num_frames"],
+                frame_rate=request["fps"],
+                num_inference_steps=request["num_inference_steps"],
+                video_guider_params=MultiModalGuiderParams(
+                    cfg_scale=request["video_cfg_scale"],
+                    stg_scale=request["video_stg_scale"],
+                    rescale_scale=request["video_rescale_scale"],
+                    modality_scale=request["video_modality_scale"],
+                    skip_step=0,
+                    stg_blocks=request["video_stg_blocks"],
+                ),
+                audio_guider_params=MultiModalGuiderParams(
+                    cfg_scale=request["audio_cfg_scale"],
+                    stg_scale=request["audio_stg_scale"],
+                    rescale_scale=request["audio_rescale_scale"],
+                    modality_scale=request["audio_modality_scale"],
+                    skip_step=0,
+                    stg_blocks=request["audio_stg_blocks"],
+                ),
+                images=images,
+                max_batch_size=4,
+                sigmas=None if request_sigmas is None else torch.tensor(request_sigmas, dtype=torch.float32),
+            )
+            pipeline_name = "TI2VidOneStagePipeline"
+        checkpoint_metadata = {
+            "transformer_path": str(args.transformer_path),
+            "video_vae_path": str(args.video_vae_path),
+            "connector_model": None if args.connector_model is None else str(args.connector_model),
+            "pipeline": pipeline_name,
+        }
+    else:
+        if args.checkpoint is None or args.gemma_root is None:
+            raise ValueError("Official monolith backend requires --checkpoint and --gemma-root")
+
+        from ltx_core.components.guiders import MultiModalGuiderParams
+        from ltx_pipelines.ti2vid_one_stage import TI2VidOneStagePipeline
+
+        pipeline = TI2VidOneStagePipeline(
+            model_paths=ModelPaths.from_monolith(str(args.checkpoint), str(args.gemma_root)),
+            loras=(),
+            offload_mode=offload_mode,
+        )
+        _configure_official_sdpa(pipeline)
+        video, audio, _ = pipeline(
+            prompt=request["prompt"],
+            negative_prompt=request["negative_prompt"],
+            seed=request["seed"],
+            height=request["height"],
+            width=request["width"],
+            num_frames=request["num_frames"],
+            frame_rate=request["fps"],
+            num_inference_steps=request["num_inference_steps"],
+            video_guider_params=MultiModalGuiderParams(
+                cfg_scale=request["video_cfg_scale"],
+                stg_scale=request["video_stg_scale"],
+                rescale_scale=request["video_rescale_scale"],
+                modality_scale=request["video_modality_scale"],
+                skip_step=0,
+                stg_blocks=request["video_stg_blocks"],
+            ),
+            audio_guider_params=MultiModalGuiderParams(
+                cfg_scale=request["audio_cfg_scale"],
+                stg_scale=request["audio_stg_scale"],
+                rescale_scale=request["audio_rescale_scale"],
+                modality_scale=request["audio_modality_scale"],
+                skip_step=0,
+                stg_blocks=request["audio_stg_blocks"],
+            ),
+            images=images,
+            max_batch_size=4,
+        )
+        checkpoint_metadata = {
+            "checkpoint": str(args.checkpoint),
+            "pipeline": "TI2VidOneStagePipeline",
+        }
+
+    video_tensor = torch.cat([chunk.detach().cpu() for chunk in video], dim=0)
+    _save_outputs(
+        args.output_dir,
+        video=_canonical_video(video_tensor),
+        audio=audio.waveform,
+        audio_sample_rate=audio.sampling_rate,
+        fps=float(request["fps"]),
+        metadata={
+            "backend": "official",
+            "attention_backend": "torch_sdpa_cudnn_bf16",
+            "official_revision": os.environ.get("VLLM_TEST_LTX_OFFICIAL_REVISION"),
+            "seed": request["seed"],
+            "sigmas": request.get("sigmas"),
+            "stage_1_sigmas": request.get("stage_1_sigmas"),
+            "stage_2_sigmas": request.get("stage_2_sigmas"),
+            **_accelerator_peak_memory_stats(),
+            **checkpoint_metadata,
+        },
+    )
+
+
+def _omni_worker_peak_memory_mb(output: Any) -> float:
+    result = output[0] if isinstance(output, list) and output else output
+    return float(getattr(result, "peak_memory_mb", 0.0) or 0.0)
+
+
+def _unwrap_omni_output(output: Any) -> tuple[Any, Any, int]:
+    from vllm_omni.outputs import OmniRequestOutput
+
+    audio = None
+    audio_sample_rate = None
+    frames = output[0] if isinstance(output, list) and output else output
+    if isinstance(frames, OmniRequestOutput):
+        multimodal_output = frames.multimodal_output or {}
+        audio = multimodal_output.get("audio")
+        audio_sample_rate = multimodal_output.get("audio_sample_rate")
+        nested_output = getattr(frames, "request_output", None)
+        if frames.is_pipeline_output and isinstance(nested_output, OmniRequestOutput):
+            frames = nested_output
+            multimodal_output = frames.multimodal_output or {}
+            audio = multimodal_output.get("audio", audio)
+            audio_sample_rate = multimodal_output.get("audio_sample_rate", audio_sample_rate)
+        if isinstance(frames, OmniRequestOutput):
+            if not frames.images:
+                raise ValueError("No video frames found in OmniRequestOutput")
+            frames = frames.images
+
+    if isinstance(frames, list) and len(frames) == 1:
+        frames = frames[0]
+    if isinstance(frames, tuple) and len(frames) == 2:
+        frames, audio = frames
+    elif isinstance(frames, dict):
+        audio = frames.get("audio", audio)
+        audio_sample_rate = frames.get("audio_sample_rate", audio_sample_rate)
+        frames = frames.get("frames", frames.get("video"))
+
+    if frames is None or audio is None or audio_sample_rate is None:
+        raise ValueError("Omni output did not contain video, audio, and audio_sample_rate")
+    return frames, audio, int(audio_sample_rate)
+
+
+def _canonical_video(video: Any) -> torch.Tensor:
+    if isinstance(video, list):
+        if len(video) == 1:
+            return _canonical_video(video[0])
+        tensors = [torch.as_tensor(np.asarray(item) if not isinstance(item, torch.Tensor) else item) for item in video]
+        if all(item.ndim == 3 for item in tensors):
+            video = torch.stack(tensors)
+        elif all(item.ndim == 4 for item in tensors):
+            video = torch.cat(tensors)
+        else:
+            raise ValueError(f"Cannot combine video items with dimensions {[item.ndim for item in tensors]}")
+    tensor = torch.as_tensor(np.asarray(video) if not isinstance(video, torch.Tensor) else video).detach().cpu()
+    if tensor.ndim == 5:
+        tensor = tensor[0]
+    if tensor.ndim != 4:
+        raise ValueError(f"Expected a 4D video tensor, got {tuple(tensor.shape)}")
+    if tensor.shape[-1] in (3, 4):
+        tensor = tensor[..., :3]
+    elif tensor.shape[1] in (3, 4):
+        tensor = tensor[:, :3].permute(0, 2, 3, 1)
+    elif tensor.shape[0] in (3, 4):
+        tensor = tensor[:3].permute(1, 2, 3, 0)
+    else:
+        raise ValueError(f"Cannot infer video channel dimension from {tuple(tensor.shape)}")
+    tensor = tensor.float()
+    if tensor.numel() and tensor.max() > 1:
+        tensor = tensor / 255.0
+    elif tensor.numel() and tensor.min() < 0:
+        tensor = tensor.clamp(-1.0, 1.0).add(1.0).mul(0.5)
+    return tensor.clamp(0.0, 1.0)
+
+
+@torch.inference_mode()
+def _run_omni(args: argparse.Namespace, request: dict[str, Any]) -> None:
+    if args.model is None or args.model_class_name is None:
+        raise ValueError("Omni backend requires --model and --model-class-name")
+
+    # vLLM imports can otherwise change the process-wide Gemma dispatch.
+    torch.backends.cuda.enable_cudnn_sdp(False)
+    _reset_accelerator_peak_memory_stats()
+
+    attention_config = {"default": {"backend": "CUDNN_ATTN"}}
+
+    from vllm_omni.diffusion.data import DiffusionParallelConfig
+    from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
+    from vllm_omni.entrypoints.omni import Omni
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+    from vllm_omni.model_extras import get_extra_body_params
+    from vllm_omni.platforms import current_omni_platform
+
+    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(request["seed"])
+    omni = Omni(
+        model=args.model,
+        model_class_name=args.model_class_name,
+        enforce_eager=True,
+        enable_layerwise_offload=args.enable_layerwise_offload,
+        diffusion_attention_config=attention_config,
+        parallel_config=DiffusionParallelConfig(),
+    )
+    try:
+        # The explicit public class argument selects the worker pipeline. The
+        # parent Omni proxy can only expose the checkpoint's base class, which
+        # is insufficient to distinguish LTX execution modes.
+        model_class_name = args.model_class_name
+        sampling_params = OmniDiffusionSamplingParams(
+            height=request["height"],
+            width=request["width"],
+            generator=generator,
+            guidance_scale=None,
+            num_inference_steps=request["num_inference_steps"],
+            num_frames=request["num_frames"],
+            fps=request["fps"],
+            frame_rate=float(request["fps"]),
+            output_type="np",
+        )
+        extra_args = {
+            key: value for key, value in request.items() if key.startswith("video_") or key.startswith("audio_")
+        }
+        for key in ("sigmas", "stage_1_sigmas", "stage_2_sigmas", "image_crf"):
+            if key in request:
+                extra_args[key] = request[key]
+        apply_declared_extra_args(sampling_params, get_extra_body_params(model_class_name), extra_args)
+        prompt: dict[str, Any] = {"prompt": request["prompt"]}
+        if "negative_prompt" in request:
+            prompt["negative_prompt"] = request["negative_prompt"]
+        image_path = request.get("image")
+        if image_path is not None:
+            with Image.open(str(image_path)) as source_image:
+                image = source_image.convert("RGB")
+                image.load()
+            prompt["multi_modal_data"] = {"image": image}
+        output = omni.generate(prompt, sampling_params)
+        video, audio, audio_sample_rate = _unwrap_omni_output(output)
+        audio_tensor = torch.as_tensor(np.asarray(audio) if not isinstance(audio, torch.Tensor) else audio)
+        _save_outputs(
+            args.output_dir,
+            video=_canonical_video(video),
+            audio=audio_tensor,
+            audio_sample_rate=audio_sample_rate,
+            fps=float(request["fps"]),
+            metadata={
+                "backend": "omni",
+                "attention_backend": "torch_sdpa_cudnn_bf16",
+                "seed": request["seed"],
+                "sigmas": request.get("sigmas"),
+                "model": args.model,
+                "stage_1_sigmas": request.get("stage_1_sigmas"),
+                "stage_2_sigmas": request.get("stage_2_sigmas"),
+                "model_class_name": model_class_name,
+                "worker_peak_memory_mb": _omni_worker_peak_memory_mb(output),
+                **_accelerator_peak_memory_stats(),
+            },
+        )
+    finally:
+        omni.shutdown()
+
+
+def main() -> None:
+    args = _parse_args()
+    request = json.loads(args.request.read_text())
+    if args.backend == "official":
+        _run_official(args, request)
+    else:
+        _run_omni(args, request)
+
+
+if __name__ == "__main__":
+    main()

tests/e2e/accuracy/ltx/test_ltx25_official_similarity.py

@@ -0,0 +1,753 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""E2E accuracy guards against pinned official Lightricks LTX references.
+
+The LTX-2/2.3 comparison runs both runtimes through PyTorch SDPA and uses
+``max_batch_size=4`` in the official reference to match Omni's fused guidance
+batch. Video and audio guidance use the official non-HQ one-stage defaults;
+only the generation shape and step count are reduced for CI runtime.
+LTX-2.5 runs the official split-artifact pipelines with connector weights from
+the same Diffusers checkpoint under test. The distilled case covers the fixed
+two-stage schedule; the Full/SFT case covers raw dev weights and an explicit
+shared one-stage schedule.
+"""
+
+from __future__ import annotations
+
+import json
+import math
+import os
+import shutil
+import subprocess
+import sys
+import time
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+from huggingface_hub import hf_hub_download, snapshot_download
+from torchmetrics.image import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure
+
+from tests.e2e.accuracy.helpers import reset_artifact_dir
+from tests.helpers import skip_if_gated_repo_inaccessible
+from tests.helpers.mark import hardware_test
+
+OFFICIAL_REPOSITORY = "https://github.com/Lightricks/LTX-2.git"
+LTX25_OFFICIAL_REVISION = "7954dcb0d986bdc36ef272564a9789ade07fcc65"
+LTX25_OMNI_MODEL_ID = "Lightricks/LTX-2.5-Diffusers"
+LTX25_OMNI_MODEL_REVISION = "a6de4b5354f078db24d9cf4778c14846788aea3d"
+LTX25_OFFICIAL_MODEL_ID = "Lightricks/LTX-2.5"
+LTX25_OFFICIAL_MODEL_REVISION = "8a4ff96f581e72bedc1b44367581c49d544a05f1"
+LTX25_OFFICIAL_COMMON_FILES = {
+    "text_encoder": "text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors",
+    "video_vae": "vae/ltx-2.5-video-vae-conv-bf16.safetensors",
+    "audio_vae": "vae/ltx-2.5-audio-vae-bf16.safetensors",
+}
+LTX25_OFFICIAL_VARIANT_FILES = {
+    "distilled": {
+        "transformer": "diffusion_models/ltx-2.5-22b-distilled-transformer-bf16.safetensors",
+        "spatial_upsampler": "latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors",
+    },
+    "full": {
+        "transformer": "diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors",
+    },
+    "full_two_stage": {
+        "transformer": "diffusion_models/ltx-2.5-22b-dev-transformer-bf16.safetensors",
+        "spatial_upsampler": "latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors",
+        "distilled_lora": "loras/ltx-2.5-22b-distilled-lora-450-bf16.safetensors",
+    },
+}
+# Version selected by the pinned official source's uv.lock. Keep it isolated
+# from Omni's runtime and development dependencies.
+OFFICIAL_OPENIMAGEIO_VERSION = "3.1.11.0"
+LTX25_PROMPT = "A cinematic shot of a red fox walking through a snowy forest at dawn, the camera tracking alongside."
+LTX25_I2V_IMAGE_REPO = "huggingface/documentation-images"
+LTX25_I2V_IMAGE_FILENAME = "diffusers/svd/rocket.png"
+LTX25_I2V_IMAGE_REVISION = "645d8364f0c7a101180b364811b5a11a362e4010"
+NEGATIVE_PROMPT = (
+    "blurry, out of focus, overexposed, underexposed, low contrast, washed out colors, excessive noise, "
+    "grainy texture, poor lighting, flickering, motion blur, distorted proportions, unnatural skin tones, "
+    "deformed facial features, asymmetrical face, missing facial features, extra limbs, disfigured hands, "
+    "wrong hand count, artifacts around text, inconsistent perspective, camera shake, incorrect depth of field, "
+    "background too sharp, background clutter, distracting reflections, harsh shadows, inconsistent lighting "
+    "direction, color banding, cartoonish rendering, 3D CGI look, unrealistic materials, uncanny valley effect, "
+    "incorrect ethnicity, wrong gender, exaggerated expressions, wrong gaze direction, mismatched lip sync, "
+    "silent or muted audio, distorted voice, robotic voice, echo, background noise, off-sync audio, incorrect "
+    "dialogue, added dialogue, repetitive speech, jittery movement, awkward pauses, incorrect timing, unnatural "
+    "transitions, inconsistent framing, tilted camera, flat lighting, inconsistent tone, cinematic oversaturation, "
+    "stylized filters, or AI artifacts."
+)
+
+# Release parity exercises the accelerated BF16 cuDNN path in both runtimes.
+# Degenerate one-token attention falls back to PyTorch SDPA MATH.
+ATTENTION_BACKEND = "torch_sdpa_cudnn_bf16"
+VIDEO_SSIM_MEAN_THRESHOLD = 0.95
+VIDEO_SSIM_MIN_THRESHOLD = 0.90
+VIDEO_PSNR_MEAN_THRESHOLD = 30.0
+AUDIO_RELATIVE_L2_THRESHOLD = 0.2
+AUDIO_COSINE_THRESHOLD = 0.95
+
+# Decoded-output gates for the pinned official/Omni BF16 cuDNN comparison.
+# The metrics artifact preserves the exact observed values.
+LTX25_VIDEO_SSIM_MEAN_THRESHOLD = 0.99
+LTX25_I2V_VIDEO_SSIM_MEAN_THRESHOLD = 0.99
+LTX25_VIDEO_SSIM_MIN_THRESHOLD = 0.99
+LTX25_VIDEO_PSNR_MEAN_THRESHOLD = 40.0
+LTX25_VIDEO_LPIPS_MEAN_THRESHOLD = 0.01
+LTX25_AUDIO_RELATIVE_L2_THRESHOLD = 0.32
+LTX25_AUDIO_COSINE_THRESHOLD = 0.95
+
+# Full/SFT uses independently converted dev weights, but the release contract
+# still requires decoded-video parity at the same 0.99 SSIM floor as the
+# distilled paths when both runtimes use BF16 cuDNN attention.
+LTX25_FULL_VIDEO_SSIM_MEAN_THRESHOLD = 0.99
+LTX25_FULL_VIDEO_SSIM_MIN_THRESHOLD = 0.99
+LTX25_FULL_VIDEO_PSNR_MEAN_THRESHOLD = 40.0
+LTX25_FULL_AUDIO_RELATIVE_L2_THRESHOLD = 0.3
+LTX25_FULL_VIDEO_LPIPS_MEAN_THRESHOLD = 0.01
+LTX25_FULL_I2V_VIDEO_LPIPS_MEAN_THRESHOLD = 0.02
+LTX25_FULL_AUDIO_COSINE_THRESHOLD = AUDIO_COSINE_THRESHOLD
+LTX25_STAGE_2_SIGMAS = [0.909375, 0.725, 0.421875, 0.0]
+
+
+def test_ltx_reference_runner_unwraps_flattened_pipeline_output() -> None:
+    from vllm_omni.outputs import OmniRequestOutput
+
+    from .run_ltx25_reference import _omni_worker_peak_memory_mb, _unwrap_omni_output
+
+    frame = object()
+    audio = object()
+    output = OmniRequestOutput(
+        stage_id=0,
+        images=[frame],
+        _multimodal_output={"audio": audio, "audio_sample_rate": 48_000},
+    )
+
+    frames, actual_audio, sample_rate = _unwrap_omni_output(output)
+
+    assert frames is frame
+    assert actual_audio is audio
+    assert sample_rate == 48_000
+    output.peak_memory_mb = 12_345.0
+    assert _omni_worker_peak_memory_mb([output]) == 12_345.0
+
+
+def _run(command: list[str], *, env: dict[str, str], timeout: int = 1800) -> None:
+    start = time.perf_counter()
+    subprocess.run(command, env=env, timeout=timeout, check=True)
+    print(f"{' '.join(command[:3])} finished in {time.perf_counter() - start:.1f}s")
+
+
+def _clone_pinned_source(root: Path, repository: str, revision: str) -> None:
+    root.mkdir(parents=True)
+    subprocess.run(["git", "init", "-q", str(root)], check=True)
+    subprocess.run(["git", "-C", str(root), "remote", "add", "origin", repository], check=True)
+    last_error: subprocess.CalledProcessError | None = None
+    for attempt in range(3):
+        try:
+            subprocess.run(
+                ["git", "-C", str(root), "fetch", "--depth", "1", "origin", revision],
+                check=True,
+            )
+            last_error = None
+            break
+        except subprocess.CalledProcessError as error:
+            last_error = error
+            if attempt < 2:
+                time.sleep(5 * (attempt + 1))
+    if last_error is not None:
+        raise last_error
+    subprocess.run(["git", "-C", str(root), "checkout", "-q", "--detach", "FETCH_HEAD"], check=True)
+
+
+def _git_revision(root: Path) -> str | None:
+    try:
+        return subprocess.check_output(
+            ["git", "-C", str(root), "rev-parse", "HEAD"],
+            stderr=subprocess.DEVNULL,
+            text=True,
+        ).strip()
+    except subprocess.CalledProcessError:
+        return None
+
+
+def _ltx25_official_source(artifact_root: Path) -> tuple[Path, str]:
+    repository = os.environ.get("VLLM_TEST_LTX25_OFFICIAL_REPOSITORY", OFFICIAL_REPOSITORY)
+    revision = os.environ.get("VLLM_TEST_LTX25_OFFICIAL_REVISION", LTX25_OFFICIAL_REVISION)
+    configured_root = os.environ.get("VLLM_TEST_LTX25_OFFICIAL_ROOT")
+    root = Path(configured_root) if configured_root else artifact_root / f"official-source-{revision[:12]}"
+    actual_revision = _git_revision(root) if root.exists() else None
+    if actual_revision != revision and configured_root:
+        raise AssertionError(f"Official LTX-2.5 source revision mismatch: {actual_revision} != {revision}")
+    if actual_revision != revision:
+        if root.exists():
+            shutil.rmtree(root)
+        _clone_pinned_source(root, repository, revision)
+        actual_revision = _git_revision(root)
+    assert actual_revision == revision, f"Official LTX-2.5 source revision mismatch: {actual_revision} != {revision}"
+    return root, revision
+
+
+def _official_runner_prefix() -> list[str]:
+    """Run the reference with its missing binary dependency isolated from CI."""
+    uv = shutil.which("uv")
+    assert uv is not None, "uv is required to run the pinned official LTX reference"
+    return [
+        uv,
+        "run",
+        "--no-project",
+        "--with",
+        f"openimageio=={OFFICIAL_OPENIMAGEIO_VERSION}",
+        "--python",
+        sys.executable,
+        "python",
+    ]
+
+
+def _resolve_ltx25_omni_model(*, transformer_subfolder: str = "transformer") -> tuple[Path, str, str | None]:
+    if transformer_subfolder not in {"transformer", "transformer_full"}:
+        raise ValueError(f"Unsupported LTX-2.5 transformer subfolder: {transformer_subfolder!r}")
+    configured_model = os.environ.get("VLLM_TEST_LTX25_MODEL")
+    configured_revision = os.environ.get("VLLM_TEST_LTX25_MODEL_REVISION")
+    if configured_model and Path(configured_model).exists():
+        model = Path(configured_model).resolve()
+        snapshot_revision = model.name if model.parent.name == "snapshots" else None
+        assert (model / transformer_subfolder).is_dir(), f"LTX-2.5 component not found: {model / transformer_subfolder}"
+        return model, configured_model, configured_revision or snapshot_revision
+    model_id = configured_model or LTX25_OMNI_MODEL_ID
+    revision = configured_revision
+    if revision is None and model_id == LTX25_OMNI_MODEL_ID:
+        revision = LTX25_OMNI_MODEL_REVISION
+    model = Path(
+        snapshot_download(
+            repo_id=model_id,
+            revision=revision,
+            allow_patterns=[
+                "model_index.json",
+                "audio_vae/*",
+                "connectors/config.json",
+                "connectors/diffusion_pytorch_model.safetensors.index.json",
+                "connectors/diffusion_pytorch_model-*.safetensors",
+                "scheduler/*",
+                "text_encoder/config.json",
+                "text_encoder/generation_config.json",
+                "text_encoder/model*",
+                "tokenizer/*",
+                f"{transformer_subfolder}/*",
+                "vae/*",
+                "vocoder/*",
+            ],
+        )
+    ).resolve()
+    snapshot_revision = model.name if model.parent.name == "snapshots" else None
+    return model, model_id, snapshot_revision or revision
+
+
+def _resolve_ltx25_official_artifacts(
+    *, checkpoint_variant: str = "distilled"
+) -> tuple[dict[str, Path], str, str | None]:
+    try:
+        variant_files = LTX25_OFFICIAL_VARIANT_FILES[checkpoint_variant]
+    except KeyError as exc:
+        raise ValueError(f"Unsupported official LTX-2.5 checkpoint variant: {checkpoint_variant!r}") from exc
+    files = {**LTX25_OFFICIAL_COMMON_FILES, **variant_files}
+    configured_model = os.environ.get("VLLM_TEST_LTX25_OFFICIAL_MODEL")
+    configured_revision = os.environ.get("VLLM_TEST_LTX25_OFFICIAL_MODEL_REVISION")
+    if configured_model and Path(configured_model).exists():
+        root = Path(configured_model).resolve()
+        model_source = configured_model
+        snapshot_revision = root.name if root.parent.name == "snapshots" else None
+        resolved_revision = configured_revision or snapshot_revision
+    else:
+        model_source = configured_model or LTX25_OFFICIAL_MODEL_ID
+        revision = configured_revision
+        if revision is None and model_source == LTX25_OFFICIAL_MODEL_ID:
+            revision = LTX25_OFFICIAL_MODEL_REVISION
+        root = Path(
+            snapshot_download(
+                repo_id=model_source,
+                revision=revision,
+                allow_patterns=list(files.values()),
+            )
+        ).resolve()
+        snapshot_revision = root.name if root.parent.name == "snapshots" else None
+        resolved_revision = snapshot_revision or revision
+
+    artifacts = {name: root / relative_path for name, relative_path in files.items()}
+    missing = [f"{name}: {path}" for name, path in artifacts.items() if not path.is_file()]
+    assert not missing, f"Official LTX-2.5 artifacts are missing: {', '.join(missing)}"
+    return artifacts, model_source, resolved_revision
+
+
+def _resolve_gemma_root(model: Path) -> Path:
+    configured_root = os.environ.get("VLLM_TEST_LTX_GEMMA_ROOT")
+    if configured_root:
+        root = Path(configured_root)
+        assert root.is_dir(), f"Gemma root not found: {root}"
+        return root
+    return model
+
+
+def _resolve_ltx25_image() -> Path:
+    configured_image = os.environ.get("VLLM_TEST_LTX25_I2V_IMAGE")
+    if configured_image:
+        image = Path(configured_image)
+        assert image.is_file(), f"LTX-2.5 I2V conditioning image not found: {image}"
+        return image
+    return Path(
+        hf_hub_download(
+            repo_id=LTX25_I2V_IMAGE_REPO,
+            repo_type="dataset",
+            filename=LTX25_I2V_IMAGE_FILENAME,
+            revision=LTX25_I2V_IMAGE_REVISION,
+        )
+    )
+
+
+def _require_ltx25_model_access() -> None:
+    configured_omni_model = os.environ.get("VLLM_TEST_LTX25_MODEL")
+    if not (configured_omni_model and Path(configured_omni_model).exists()):
+        omni_model = configured_omni_model or LTX25_OMNI_MODEL_ID
+        omni_revision = os.environ.get("VLLM_TEST_LTX25_MODEL_REVISION")
+        if omni_revision is None and omni_model == LTX25_OMNI_MODEL_ID:
+            omni_revision = LTX25_OMNI_MODEL_REVISION
+        skip_if_gated_repo_inaccessible(
+            omni_model,
+            revision=omni_revision,
+            filename="model_index.json",
+        )
+
+    configured_official_model = os.environ.get("VLLM_TEST_LTX25_OFFICIAL_MODEL")
+    if not (configured_official_model and Path(configured_official_model).exists()):
+        official_model = configured_official_model or LTX25_OFFICIAL_MODEL_ID
+        official_revision = os.environ.get("VLLM_TEST_LTX25_OFFICIAL_MODEL_REVISION")
+        if official_revision is None and official_model == LTX25_OFFICIAL_MODEL_ID:
+            official_revision = LTX25_OFFICIAL_MODEL_REVISION
+        skip_if_gated_repo_inaccessible(
+            official_model,
+            revision=official_revision,
+            filename=LTX25_OFFICIAL_COMMON_FILES["text_encoder"],
+        )
+
+
+def _ltx25_request(image: Path | None = None) -> dict[str, object]:
+    request: dict[str, object] = {
+        "prompt": LTX25_PROMPT,
+        # Keep the nightly parity test small enough to run the pinned official
+        # pipeline and Omni sequentially on one H100. The public recipe covers
+        # the release-resolution 1920x1088x121 configuration.
+        "width": 1024,
+        "height": 576,
+        "num_frames": 25,
+        "fps": 24,
+        "num_inference_steps": 8,
+        "seed": 42,
+    }
+    if image is not None:
+        request.update(image=str(image.resolve()), image_crf=18)
+    return request
+
+
+def _ltx25_full_sigmas(num_inference_steps: int) -> list[float]:
+    """Materialize the official LTX-2.5 one-stage schedule in FP32."""
+    sigmas = torch.linspace(1.0, 0.0, num_inference_steps + 1)
+    base_anchor = 1024
+    max_anchor = 4096
+    base_shift = 0.95
+    max_shift = 2.05
+    slope = (max_shift - base_shift) / (max_anchor - base_anchor)
+    sigma_shift = max_anchor * slope + (base_shift - slope * base_anchor)
+    exp_shift = math.exp(sigma_shift)
+    sigmas = torch.where(sigmas != 0, exp_shift / (exp_shift + (1 / sigmas - 1)), 0)
+
+    non_zero = sigmas != 0
+    one_minus_sigmas = 1.0 - sigmas[non_zero]
+    scale = one_minus_sigmas[-1] / (1.0 - 0.1)
+    sigmas[non_zero] = 1.0 - one_minus_sigmas / scale
+    return [float(sigma) for sigma in sigmas]
+
+
+def _ltx25_full_request(image: Path | None = None) -> dict[str, object]:
+    num_inference_steps = 30
+    request: dict[str, object] = {
+        "prompt": LTX25_PROMPT,
+        "negative_prompt": NEGATIVE_PROMPT,
+        # Full guidance expands to a four-way batch. Keep the decoded output
+        # small enough for the pinned official and Omni runs to share one H100.
+        "width": 512,
+        "height": 384,
+        "num_frames": 25,
+        "fps": 24,
+        "num_inference_steps": num_inference_steps,
+        "seed": 42,
+        "video_cfg_scale": 3.0,
+        "audio_cfg_scale": 7.0,
+        "video_stg_scale": 1.0,
+        "audio_stg_scale": 1.0,
+        "video_modality_scale": 3.0,
+        "audio_modality_scale": 3.0,
+        "video_rescale_scale": 0.7,
+        "audio_rescale_scale": 0.7,
+        "video_stg_blocks": [28],
+        "audio_stg_blocks": [28],
+        "sigmas": _ltx25_full_sigmas(num_inference_steps),
+    }
+    if image is not None:
+        request.update(image=str(image.resolve()), image_crf=18)
+    return request
+
+
+def test_ltx25_full_request_pins_official_schedule() -> None:
+    request = _ltx25_full_request()
+    sigmas = request["sigmas"]
+
+    assert isinstance(sigmas, list)
+    assert isinstance(request["num_inference_steps"], int)
+    assert len(sigmas) == request["num_inference_steps"] + 1
+    assert math.isclose(sigmas[0], 1.0, rel_tol=0.0, abs_tol=1e-6)
+    assert sigmas[-1] == 0.0
+    assert all(sigmas[index] > sigmas[index + 1] for index in range(len(sigmas) - 1))
+
+
+def test_ltx25_i2v_requests_pin_official_crf() -> None:
+    image = Path("conditioning.png")
+    for request in (_ltx25_request(image), _ltx25_full_request(image)):
+        assert request["image"] == str(image.resolve())
+        assert request["image_crf"] == 18
+
+
+def _video_metrics(reference: np.ndarray, prediction: np.ndarray) -> dict[str, float]:
+    from benchmarks.diffusion.quantization_quality import compute_lpips_video
+
+    assert reference.shape == prediction.shape
+    assert reference.ndim == 4 and reference.shape[-1] == 3
+    ssim_scores: list[float] = []
+    psnr_scores: list[float] = []
+    for reference_frame, prediction_frame in zip(reference, prediction, strict=True):
+        reference_tensor = torch.from_numpy(reference_frame).permute(2, 0, 1).unsqueeze(0)
+        prediction_tensor = torch.from_numpy(prediction_frame).permute(2, 0, 1).unsqueeze(0)
+        ssim_scores.append(float(StructuralSimilarityIndexMeasure(data_range=1.0)(prediction_tensor, reference_tensor)))
+        psnr_scores.append(float(PeakSignalNoiseRatio(data_range=1.0)(prediction_tensor, reference_tensor)))
+    difference = np.abs(reference.astype(np.float64) - prediction.astype(np.float64))
+    return {
+        "ssim_mean": float(np.mean(ssim_scores)),
+        "ssim_min": float(np.min(ssim_scores)),
+        "psnr_mean_db": float(np.mean(psnr_scores)),
+        "lpips_alex_mean": compute_lpips_video(reference, prediction),
+        "max_abs": float(difference.max()),
+        "mean_abs": float(difference.mean()),
+    }
+
+
+def _canonical_audio(audio: np.ndarray) -> np.ndarray:
+    while audio.ndim > 2 and audio.shape[0] == 1:
+        audio = audio[0]
+    return audio.astype(np.float64)
+
+
+def _audio_metrics(reference: np.ndarray, prediction: np.ndarray) -> dict[str, float | bool]:
+    reference = _canonical_audio(reference)
+    prediction = _canonical_audio(prediction)
+    assert reference.shape == prediction.shape
+    difference = reference - prediction
+    reference_norm = max(float(np.linalg.norm(reference)), 1e-12)
+    prediction_norm = max(float(np.linalg.norm(prediction)), 1e-12)
+    return {
+        "bitwise_equal": bool(np.array_equal(reference, prediction)),
+        "max_abs": float(np.abs(difference).max()),
+        "mean_abs": float(np.abs(difference).mean()),
+        "relative_l2": float(np.linalg.norm(difference) / reference_norm),
+        "cosine_similarity": float(np.vdot(reference.ravel(), prediction.ravel()) / (reference_norm * prediction_norm)),
+    }
+
+
+@pytest.mark.full_model
+@pytest.mark.benchmark
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("task", ["t2v", "i2v"])
+def test_ltx25_distilled_two_stage_matches_official(accuracy_artifact_root: Path, task: str) -> None:
+    """Compare Omni with the official runtime using the same checkpoint connector weights."""
+    _require_ltx25_model_access()
+    artifact_parent = accuracy_artifact_root / "ltx_official"
+    output_root = reset_artifact_dir(artifact_parent / f"ltx2_5_distilled_{task}")
+    official_root, official_revision = _ltx25_official_source(artifact_parent)
+    official_artifacts, official_model_source, official_model_revision = _resolve_ltx25_official_artifacts()
+    omni_model, omni_model_source, omni_model_revision = _resolve_ltx25_omni_model()
+    image = _resolve_ltx25_image() if task == "i2v" else None
+    request_path = output_root / "request.json"
+    request_path.write_text(json.dumps(_ltx25_request(image), indent=2) + "\n")
+
+    runner = Path(__file__).with_name("run_ltx25_reference.py")
+    runner_args = [
+        str(runner),
+        "--request",
+        str(request_path),
+        "--enable-layerwise-offload",
+    ]
+    env = os.environ.copy()
+    env["VLLM_TEST_LTX_OFFICIAL_REVISION"] = official_revision
+    env["PYTHONUNBUFFERED"] = "1"
+    repository_root = Path(__file__).resolve().parents[4]
+    existing_pythonpath = env.get("PYTHONPATH")
+    env["PYTHONPATH"] = (
+        str(repository_root) if not existing_pythonpath else f"{repository_root}{os.pathsep}{existing_pythonpath}"
+    )
+
+    official_output = output_root / "official"
+    _run(
+        _official_runner_prefix()
+        + runner_args
+        + [
+            "--backend",
+            "official",
+            "--output-dir",
+            str(official_output),
+            "--official-root",
+            str(official_root),
+            "--transformer-path",
+            str(official_artifacts["transformer"]),
+            "--text-encoder-path",
+            str(official_artifacts["text_encoder"]),
+            "--video-vae-path",
+            str(official_artifacts["video_vae"]),
+            "--audio-vae-path",
+            str(official_artifacts["audio_vae"]),
+            "--spatial-upsampler-path",
+            str(official_artifacts["spatial_upsampler"]),
+            "--connector-model",
+            str(omni_model),
+        ],
+        env=env,
+    )
+
+    omni_output = output_root / "omni"
+    _run(
+        [sys.executable]
+        + runner_args
+        + [
+            "--backend",
+            "omni",
+            "--output-dir",
+            str(omni_output),
+            "--model",
+            str(omni_model),
+            "--model-class-name",
+            "LTX2DistilledTwoStagePipeline",
+        ],
+        env=env,
+    )
+
+    official_metadata = json.loads((official_output / "metadata.json").read_text())
+    omni_metadata = json.loads((omni_output / "metadata.json").read_text())
+    assert official_metadata["official_revision"] == official_revision
+    assert official_metadata["pipeline"] == "DistilledPipeline"
+    assert official_metadata["attention_backend"] == ATTENTION_BACKEND
+    assert official_metadata["connector_model"] == str(omni_model)
+    assert omni_metadata["model_class_name"] == "LTX2DistilledTwoStagePipeline"
+    assert omni_metadata["attention_backend"] == ATTENTION_BACKEND
+    assert official_metadata["audio_sample_rate"] == omni_metadata["audio_sample_rate"]
+    video_metrics = _video_metrics(
+        np.load(official_output / "video.npy"),
+        np.load(omni_output / "video.npy"),
+    )
+    audio_metrics = _audio_metrics(
+        np.load(official_output / "audio.npy"),
+        np.load(omni_output / "audio.npy"),
+    )
+    result = {
+        "case": "ltx2_5_distilled",
+        "task": task,
+        "attention_backend": ATTENTION_BACKEND,
+        "official_revision": official_revision,
+        "official_model": official_model_source,
+        "official_model_revision": official_model_revision,
+        "official_connector_model": omni_model_source,
+        "official_connector_revision": omni_model_revision,
+        "omni_model": omni_model_source,
+        "omni_model_revision": omni_model_revision,
+        "resolved_omni_model_path": str(omni_model),
+        "peak_memory_mb": {
+            "official_allocated": official_metadata["peak_memory_allocated_mb"],
+            "official_reserved": official_metadata["peak_memory_reserved_mb"],
+            "omni_worker_reported": omni_metadata["worker_peak_memory_mb"],
+            "omni_parent_allocated": omni_metadata["peak_memory_allocated_mb"],
+            "omni_parent_reserved": omni_metadata["peak_memory_reserved_mb"],
+        },
+        "video": video_metrics,
+        "audio": audio_metrics,
+    }
+    (output_root / "metrics.json").write_text(json.dumps(result, indent=2) + "\n")
+    print(json.dumps(result, indent=2))
+
+    ssim_mean_threshold = LTX25_VIDEO_SSIM_MEAN_THRESHOLD if task == "t2v" else LTX25_I2V_VIDEO_SSIM_MEAN_THRESHOLD
+    assert video_metrics["ssim_mean"] >= ssim_mean_threshold
+    assert video_metrics["ssim_min"] >= LTX25_VIDEO_SSIM_MIN_THRESHOLD
+    assert video_metrics["psnr_mean_db"] >= LTX25_VIDEO_PSNR_MEAN_THRESHOLD
+    assert video_metrics["lpips_alex_mean"] <= LTX25_VIDEO_LPIPS_MEAN_THRESHOLD
+    assert audio_metrics["cosine_similarity"] >= LTX25_AUDIO_COSINE_THRESHOLD
+    assert audio_metrics["relative_l2"] <= LTX25_AUDIO_RELATIVE_L2_THRESHOLD
+
+
+@pytest.mark.full_model
+@pytest.mark.benchmark
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("task", ["t2v", "i2v"])
+@pytest.mark.parametrize("pipeline_mode", ["full_one_stage", "full_two_stage"])
+def test_ltx25_full_matches_official(accuracy_artifact_root: Path, task: str, pipeline_mode: str) -> None:
+    """Compare raw official Full/SFT weights with Omni's converted Full pipeline."""
+    _require_ltx25_model_access()
+    artifact_parent = accuracy_artifact_root / "ltx_official"
+    output_root = reset_artifact_dir(artifact_parent / f"ltx2_5_{pipeline_mode}_{task}")
+    model_class_name = {
+        "full_one_stage": "LTX2Pipeline",
+        "full_two_stage": "LTX2TwoStagePipeline",
+    }[pipeline_mode]
+    official_root, official_revision = _ltx25_official_source(artifact_parent)
+    official_artifacts, official_model_source, official_model_revision = _resolve_ltx25_official_artifacts(
+        checkpoint_variant="full" if pipeline_mode == "full_one_stage" else "full_two_stage"
+    )
+    omni_model, omni_model_source, omni_model_revision = _resolve_ltx25_omni_model(
+        transformer_subfolder="transformer_full"
+    )
+    image = _resolve_ltx25_image() if task == "i2v" else None
+    request = _ltx25_full_request(image)
+    if pipeline_mode == "full_two_stage":
+        request["stage_1_sigmas"] = request.pop("sigmas")
+        request["stage_2_sigmas"] = LTX25_STAGE_2_SIGMAS
+    request_path = output_root / "request.json"
+    request_path.write_text(json.dumps(request, indent=2) + "\n")
+
+    runner = Path(__file__).with_name("run_ltx25_reference.py")
+    runner_args = [
+        str(runner),
+        "--request",
+        str(request_path),
+        "--enable-layerwise-offload",
+    ]
+    env = os.environ.copy()
+    env["VLLM_TEST_LTX_OFFICIAL_REVISION"] = official_revision
+    env["PYTHONUNBUFFERED"] = "1"
+    repository_root = Path(__file__).resolve().parents[4]
+    existing_pythonpath = env.get("PYTHONPATH")
+    env["PYTHONPATH"] = (
+        str(repository_root) if not existing_pythonpath else f"{repository_root}{os.pathsep}{existing_pythonpath}"
+    )
+
+    official_output = output_root / "official"
+    official_sidecar_args: list[str] = []
+    if pipeline_mode == "full_two_stage":
+        official_sidecar_args = [
+            "--spatial-upsampler-path",
+            str(official_artifacts["spatial_upsampler"]),
+            "--distilled-lora-path",
+            str(official_artifacts["distilled_lora"]),
+        ]
+
+    _run(
+        _official_runner_prefix()
+        + runner_args
+        + [
+            "--backend",
+            "official",
+            "--output-dir",
+            str(official_output),
+            "--official-root",
+            str(official_root),
+            "--official-pipeline",
+            pipeline_mode,
+            "--transformer-path",
+            str(official_artifacts["transformer"]),
+            "--text-encoder-path",
+            str(official_artifacts["text_encoder"]),
+            "--video-vae-path",
+            str(official_artifacts["video_vae"]),
+            "--audio-vae-path",
+            str(official_artifacts["audio_vae"]),
+            "--connector-model",
+            str(omni_model),
+        ]
+        + official_sidecar_args,
+        env=env,
+    )
+
+    omni_output = output_root / "omni"
+    _run(
+        [sys.executable]
+        + runner_args
+        + [
+            "--backend",
+            "omni",
+            "--output-dir",
+            str(omni_output),
+            "--model",
+            str(omni_model),
+            "--model-class-name",
+            model_class_name,
+        ],
+        env=env,
+    )
+
+    official_metadata = json.loads((official_output / "metadata.json").read_text())
+    omni_metadata = json.loads((omni_output / "metadata.json").read_text())
+    assert official_metadata["official_revision"] == official_revision
+    assert (
+        official_metadata["pipeline"]
+        == {
+            "full_one_stage": "TI2VidOneStagePipeline",
+            "full_two_stage": "TI2VidTwoStagesPipeline",
+        }[pipeline_mode]
+    )
+    assert official_metadata["attention_backend"] == ATTENTION_BACKEND
+    assert official_metadata["connector_model"] == str(omni_model)
+    assert omni_metadata["model_class_name"] == model_class_name
+    assert omni_metadata["attention_backend"] == ATTENTION_BACKEND
+    assert official_metadata["seed"] == omni_metadata["seed"]
+    for schedule_name in ("sigmas", "stage_1_sigmas", "stage_2_sigmas"):
+        assert official_metadata[schedule_name] == omni_metadata[schedule_name]
+    assert official_metadata["audio_sample_rate"] == omni_metadata["audio_sample_rate"]
+    video_metrics = _video_metrics(
+        np.load(official_output / "video.npy"),
+        np.load(omni_output / "video.npy"),
+    )
+    audio_metrics = _audio_metrics(
+        np.load(official_output / "audio.npy"),
+        np.load(omni_output / "audio.npy"),
+    )
+    result = {
+        "case": f"ltx2_5_{pipeline_mode}",
+        "task": task,
+        "attention_backend": ATTENTION_BACKEND,
+        "official_revision": official_revision,
+        "official_model": official_model_source,
+        "official_model_revision": official_model_revision,
+        "official_connector_model": omni_model_source,
+        "official_connector_revision": omni_model_revision,
+        "omni_model": omni_model_source,
+        "omni_model_revision": omni_model_revision,
+        "resolved_omni_model_path": str(omni_model),
+        "peak_memory_mb": {
+            "official_allocated": official_metadata["peak_memory_allocated_mb"],
+            "official_reserved": official_metadata["peak_memory_reserved_mb"],
+            "omni_worker_reported": omni_metadata["worker_peak_memory_mb"],
+            "omni_parent_allocated": omni_metadata["peak_memory_allocated_mb"],
+            "omni_parent_reserved": omni_metadata["peak_memory_reserved_mb"],
+        },
+        "video": video_metrics,
+        "audio": audio_metrics,
+    }
+    (output_root / "metrics.json").write_text(json.dumps(result, indent=2) + "\n")
+    print(json.dumps(result, indent=2))
+
+    assert video_metrics["ssim_mean"] >= LTX25_FULL_VIDEO_SSIM_MEAN_THRESHOLD
+    assert video_metrics["ssim_min"] >= LTX25_FULL_VIDEO_SSIM_MIN_THRESHOLD
+    assert video_metrics["psnr_mean_db"] >= LTX25_FULL_VIDEO_PSNR_MEAN_THRESHOLD
+    lpips_threshold = (
+        LTX25_FULL_VIDEO_LPIPS_MEAN_THRESHOLD if task == "t2v" else LTX25_FULL_I2V_VIDEO_LPIPS_MEAN_THRESHOLD
+    )
+    assert audio_metrics["relative_l2"] <= LTX25_FULL_AUDIO_RELATIVE_L2_THRESHOLD
+    assert video_metrics["lpips_alex_mean"] <= lpips_threshold
+    assert audio_metrics["cosine_similarity"] >= LTX25_FULL_AUDIO_COSINE_THRESHOLD

tests/e2e/online_serving/test_ltx25.py

@@ -0,0 +1,89 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Online serving smokes for all four canonical LTX-2.5 pipelines."""
+
+import os
+
+import pytest
+
+from tests.helpers import skip_if_gated_repo_inaccessible
+from tests.helpers.mark import hardware_marks
+from tests.helpers.media import generate_synthetic_image
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+DEFAULT_MODEL = "Lightricks/LTX-2.5-Diffusers"
+DEFAULT_REVISION = "a6de4b5354f078db24d9cf4778c14846788aea3d"
+MODEL = os.getenv("VLLM_TEST_LTX25_MODEL", DEFAULT_MODEL)
+MODEL_REVISION = os.getenv("VLLM_TEST_LTX25_MODEL_REVISION", DEFAULT_REVISION if MODEL == DEFAULT_MODEL else "")
+PROMPT = "A red fox walks through a snowy forest while the camera remains fixed."
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
+SINGLE_CARD_MARKS = hardware_marks(res={"cuda": "H100"})
+
+
+@pytest.fixture(scope="module", autouse=True)
+def require_ltx25_model_access() -> None:
+    if not os.path.isdir(MODEL):
+        skip_if_gated_repo_inaccessible(
+            MODEL,
+            revision=MODEL_REVISION or None,
+            filename="model_index.json",
+        )
+
+
+def _server(model_class_name: str) -> OmniServerParams:
+    return OmniServerParams(
+        model=MODEL,
+        server_args=[
+            *(["--revision", MODEL_REVISION] if MODEL_REVISION else []),
+            "--model-class-name",
+            model_class_name,
+            "--enforce-eager",
+            "--enable-layerwise-offload",
+            "--diffusion-attention-backend",
+            "CUDNN_ATTN",
+        ],
+    )
+
+
+def _cases():
+    return [
+        pytest.param(_server("LTX2Pipeline"), 30, id="full_one_stage", marks=SINGLE_CARD_MARKS),
+        pytest.param(_server("LTX2TwoStagePipeline"), 30, id="full_two_stage", marks=SINGLE_CARD_MARKS),
+        pytest.param(_server("LTX2DistilledOneStagePipeline"), 8, id="distilled_one_stage", marks=SINGLE_CARD_MARKS),
+        pytest.param(_server("LTX2DistilledTwoStagePipeline"), 8, id="distilled_two_stage", marks=SINGLE_CARD_MARKS),
+    ]
+
+
+@pytest.mark.parametrize(("omni_server", "num_inference_steps"), _cases(), indirect=["omni_server"])
+def test_ltx25_pipeline_entries(
+    omni_server: OmniServer,
+    num_inference_steps: int,
+    openai_client: OpenAIClientHandler,
+    subtests,
+) -> None:
+    """Generate T2V and first-frame I2V through the synchronous video API."""
+    for task in ("t2v", "i2v"):
+        with subtests.test(task=task):
+            request_config = {
+                "model": omni_server.model,
+                "form_data": {
+                    "model": omni_server.model,
+                    "prompt": PROMPT,
+                    "height": 128,
+                    "width": 128,
+                    "num_frames": 9,
+                    "fps": 24,
+                    "num_inference_steps": num_inference_steps,
+                    "seed": 42,
+                },
+            }
+            if task == "i2v":
+                request_config["image_reference"] = (
+                    f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
+                )
+
+            openai_client.send_video_diffusion_request(request_config)

tests/entrypoints/openai_api/test_video_server.py

@@ -32,7 +32,7 @@ from vllm_omni.entrypoints.openai.protocol.videos import (
     VideoParams,
     VideoResponse,
 )
-from vllm_omni.entrypoints.openai.serving_video import OmniOpenAIServingVideo
+from vllm_omni.entrypoints.openai.serving_video import OmniOpenAIServingVideo, ReferenceImage
 from vllm_omni.entrypoints.openai.storage import LocalStorageManager
 from vllm_omni.entrypoints.openai.stores import AsyncDictStore, TaskRegistry
 from vllm_omni.errors import GuardrailViolationError
@@ -461,6 +461,51 @@ def test_i2v_video_generation_resizes_input_to_requested_dimensions(test_client,
     assert input_image.size == (96, 64)
 
 
+def test_i2v_resize_policy_can_defer_to_pipeline(monkeypatch):
+    engine = FakeAsyncOmni()
+    engine.get_diffusion_od_config = lambda: SimpleNamespace(
+        model="org/model",
+        model_class_name="ExamplePipeline",
+        revision="pinned-revision",
+    )
+    captured = {}
+
+    def fake_policy(model_class_name, *, model, revision=None):
+        captured.update(
+            model_class_name=model_class_name,
+            model=model,
+            revision=revision,
+        )
+        return True
+
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.openai.serving_video.should_preserve_reference_image_size",
+        fake_policy,
+    )
+    handler = OmniOpenAIServingVideo.for_diffusion(
+        diffusion_engine=engine,
+        model_name="fallback/model",
+    )
+    image = Image.new("RGB", (48, 32))
+
+    asyncio.run(
+        handler._run_and_extract(
+            VideoGenerationRequest(prompt="A bear playing with yarn.", width=96, height=64),
+            "pipeline-owned-resize",
+            reference_image=ReferenceImage(image),
+        )
+    )
+
+    input_image = engine.captured_prompt["multi_modal_data"]["image"]
+    assert isinstance(input_image, Image.Image)
+    assert input_image.size == (48, 32)
+    assert captured == {
+        "model_class_name": "ExamplePipeline",
+        "model": "org/model",
+        "revision": "pinned-revision",
+    }
+
+
 def test_i2v_extra_params_dimensions_preserve_input_image_geometry(test_client, mocker: MockerFixture):
     image_bytes = _make_test_image_bytes((48, 48))
     mocker.patch(

tests/entrypoints/test_resolve_dreamzero_config.py

@@ -49,7 +49,9 @@ def test_dreamzero_config_sets_model_class_and_policy_config(monkeypatch):
 def test_dreamzero_enrich_config_preserves_explicit_model_class_name(monkeypatch):
     monkeypatch.setattr(
         "vllm.transformers_utils.config.get_hf_file_to_dict",
-        lambda path, _model: None if path == "model_index.json" else {"model_type": "vla", "architectures": ["VLA"]},
+        lambda path, _model, **_kwargs: (
+            None if path == "model_index.json" else {"model_type": "vla", "architectures": ["VLA"]}
+        ),
     )
 
     od_config = OmniDiffusionConfig(

tests/helpers/__init__.py

@@ -11,11 +11,16 @@ affect in-process Omni (``OmniRunner`` / offline e2e) vs subprocess-based
 import pytest
 
 
-def skip_if_gated_repo_inaccessible(repo_id: str) -> None:
+def skip_if_gated_repo_inaccessible(
+    repo_id: str,
+    *,
+    revision: str | None = None,
+    filename: str = "config.json",
+) -> None:
     """Skip the test if a gated HuggingFace repo is not accessible.
 
-    Tries to download the model's config.json via ``hf_hub_download``,
-    which performs an actual file-access check (unlike ``HfApi().model_info()``
+    Tries to download ``filename`` via ``hf_hub_download``, which performs an
+    actual file-access check (unlike ``HfApi().model_info()``
     that only checks metadata).  If the token has metadata access but not
     file-download access, ``hf_hub_download`` will raise ``GatedRepoError``
     and we skip cleanly.
@@ -26,7 +31,7 @@ def skip_if_gated_repo_inaccessible(repo_id: str) -> None:
     except Exception:
         return
     try:
-        hf_hub_download(repo_id=repo_id, filename="config.json")
+        hf_hub_download(repo_id=repo_id, filename=filename, revision=revision)
     except GatedRepoError as exc:
         pytest.skip(
             f"Skipping: gated HF repo {repo_id!r} inaccessible to the current "

tests/model_extras/test_model_extras.py

@@ -20,6 +20,7 @@ from vllm_omni.model_extras import (
     get_output_tensor_range,
     get_x_to_text_model_family,
     should_init_extra_args_for_non_diffusion_stages,
+    should_preserve_reference_image_size,
 )
 
 
@@ -230,18 +231,62 @@ def test_ltx_extra_registry_declares_official_guidance_params() -> None:
             "audio_rescale_scale",
             "video_stg_blocks",
             "audio_stg_blocks",
+            "sigmas",
+            "stage_1_sigmas",
+            "stage_2_sigmas",
+            "image_crf",
         }
     )
 
     for pipeline_name in (
         "LTX2Pipeline",
         "LTX2TwoStagePipeline",
+        "LTX2DistilledOneStagePipeline",
+        "LTX2DistilledTwoStagePipeline",
         "LTX2DistilledPipeline",
     ):
         assert get_extra_body_params(pipeline_name) == expected
         assert get_extra_output_params(pipeline_name) == frozenset()
 
 
+@pytest.mark.parametrize(
+    ("model_version", "expected"),
+    [("2", False), ("2.3", False), ("2.5", True)],
+)
+def test_ltx_reference_image_size_policy(tmp_path, model_version: str, expected: bool) -> None:
+    (tmp_path / "model_index.json").write_text(
+        '{"_class_name": "LTX2Pipeline", "model_version": "' + model_version + '"}'
+    )
+
+    assert (
+        should_preserve_reference_image_size(
+            None,
+            model=str(tmp_path),
+        )
+        is expected
+    )
+
+
+def test_reference_image_size_policy_threads_revision(monkeypatch: pytest.MonkeyPatch) -> None:
+    captured = {}
+
+    def fake_policy(*, model, revision=None):
+        captured.update(model=model, revision=revision)
+        return True
+
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.models.ltx2.ltx2_components.preserves_reference_image_size",
+        fake_policy,
+    )
+
+    assert should_preserve_reference_image_size(
+        "LTX2Pipeline",
+        model="org/model",
+        revision="pinned-revision",
+    )
+    assert captured == {"model": "org/model", "revision": "pinned-revision"}
+
+
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_ming_flash_omni_extra_registry_declares_request_and_response_params() -> None:
@@ -424,6 +469,13 @@ def test_vace_extra_registry_has_no_pipeline_params() -> None:
 def test_unknown_pipeline_has_empty_extra_registry() -> None:
     assert get_extra_body_params("UnknownPipeline") == frozenset()
     assert get_extra_output_params("UnknownPipeline") == frozenset()
+    assert (
+        should_preserve_reference_image_size(
+            "UnknownPipeline",
+            model="org/model",
+        )
+        is False
+    )
     assert should_init_extra_args_for_non_diffusion_stages("UnknownPipeline") is False
 
 

tests/model_tests/diffusion/test_alignment.py

@@ -28,6 +28,8 @@ EXCLUDED_MODELS = [
     "WanPipeline",
     "WanVACEPipeline",
     "LTX2TwoStagePipeline",
+    "LTX2DistilledOneStagePipeline",
+    "LTX2DistilledTwoStagePipeline",
     "LTX2DistilledPipeline",
     "LTX2T2VDMD2Pipeline",
     "LTX2I2VDMD2Pipeline",

vllm_omni/diffusion/attention/backends/cudnn_attn.py

@@ -85,6 +85,24 @@ class CuDNNAttentionImpl(AttentionImpl):
 
         enable_gqa = query.shape[2] != key.shape[2]
         query, key, value = (x.permute(0, 2, 1, 3) for x in (query, key, value))
+        # cuDNN has no kernel for degenerate attention shapes when either
+        # sequence has one token. Route those degenerate shapes before entering the
+        # cuDNN-only context: under torch.compile, FakeTensor evaluation raises
+        # before the eager RuntimeError fallback below can run.
+        if query.shape[-2] == 1 or key.shape[-2] == 1:
+            with sdpa_kernel([SDPBackend.MATH]):
+                output = torch.nn.functional.scaled_dot_product_attention(
+                    query,
+                    key,
+                    value,
+                    attn_mask=attention_mask,
+                    dropout_p=0.0,
+                    is_causal=self.causal,
+                    scale=self.softmax_scale,
+                    enable_gqa=enable_gqa,
+                )
+            return output.permute(0, 2, 1, 3)
+
         # Pin cuDNN exclusively. A priority list like [CUDNN, FLASH, MATH] hits a
         # PyTorch SDPA dispatch quirk: when FLASH rejects a non-None attn_mask,
         # cuDNN gets runtime-disabled in the same call and the dispatcher falls
@@ -93,7 +111,7 @@ class CuDNNAttentionImpl(AttentionImpl):
         #
         # Fall back to the default SDPA dispatcher if cuDNN rejects the shape,
         # e.g. under torch.compile where Dynamo sees a symbolic head_dim and
-        # cuDNN's kernel selection fails (observed in LTX-2 audio attention).
+        # cuDNN's kernel selection fails for some symbolic attention shapes.
         # The unpinned dispatcher then picks EFFICIENT/MATH instead of raising.
         try:
             with sdpa_kernel([SDPBackend.CUDNN_ATTENTION]):

vllm_omni/diffusion/data.py

@@ -538,7 +538,11 @@ class DiffusionCacheConfig:
         raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{item}'")
 
 
-def resolve_model_class_name(model: str | None, diffusion_load_format: str = "default") -> str | None:
+def resolve_model_class_name(
+    model: str | None,
+    diffusion_load_format: str = "default",
+    revision: str | None = None,
+) -> str | None:
     """Resolve the diffusion pipeline class name from the model config.
 
     Read-only counterpart of ``OmniDiffusionConfig.enrich_config``, safe to call
@@ -550,11 +554,10 @@ def resolve_model_class_name(model: str | None, diffusion_load_format: str = "de
 
     if not model:
         return None
-
     is_lance_subfolder = os.path.basename(str(model).rstrip("/")) in {"Lance_3B", "Lance_3B_Video"}
 
     # Diffusers models: read _class_name from the pipeline index.
-    model_index = get_diffusion_model_index(model)
+    model_index = get_diffusion_model_index(model, revision=revision)
     if model_index is not None:
         return model_index.get("_class_name")
     if diffusion_load_format == "diffusers":
@@ -562,7 +565,7 @@ def resolve_model_class_name(model: str | None, diffusion_load_format: str = "de
 
     # Other models: map model_type / architecture from config.json.
     try:
-        cfg = get_hf_file_to_dict("config.json", model) or {}
+        cfg = get_hf_file_to_dict("config.json", model, revision=revision) or {}
     except Exception:
         cfg = {}
     model_type = cfg.get("model_type")
@@ -1210,9 +1213,20 @@ class OmniDiffusionConfig:
                             exc,
                         )
                 else:
-                    tf_config_dict = get_hf_file_to_dict("transformer/config.json", self.model)
+                    from vllm_omni.model_extras import get_transformer_config_subfolder
+
+                    transformer_subfolder = get_transformer_config_subfolder(
+                        self.model_class_name,
+                        model=self.model,
+                        revision=self.revision,
+                    )
+                    tf_config_dict = get_hf_file_to_dict(
+                        f"{transformer_subfolder}/config.json",
+                        self.model,
+                        revision=self.revision,
+                    )
                     if tf_config_dict is None:
-                        tf_config_dict = get_hf_file_to_dict("unet/config.json", self.model)
+                        tf_config_dict = get_hf_file_to_dict("unet/config.json", self.model, revision=self.revision)
                     if tf_config_dict is not None:
                         self.set_tf_model_config(TransformerConfig.from_dict(tf_config_dict))
                     else:
@@ -1232,7 +1246,7 @@ class OmniDiffusionConfig:
                     "that require additional inputs."
                 )
             else:
-                cfg = get_hf_file_to_dict("config.json", self.model)
+                cfg = get_hf_file_to_dict("config.json", self.model, revision=self.revision)
                 if cfg is None:
                     # Lance ships its top-level config.json one directory above
                     # the per-checkpoint subfolders (``Lance_3B/`` or

vllm_omni/diffusion/model_loader/hub_prefetch.py

@@ -265,6 +265,7 @@ def prefetch_subfolders(
     subfolders: Iterable[str],
     *,
     local_files_only: bool | None = None,
+    revision: str | None = None,
     include_root_metadata: bool = True,
 ) -> None:
     """Materialise ``model``'s ``subfolders`` in the HF cache before loading.
@@ -278,6 +279,7 @@ def prefetch_subfolders(
         local_files_only: When ``True``, skip the prefetch entirely.
             When ``None`` (default), auto-detect: skip if *model* is a
             local directory, run otherwise.
+        revision: Optional Hub revision shared with component loaders.
         include_root_metadata: When True, also pull ``*.json`` at the repo
             root so ``model_index.json`` / ``config.json`` resolution during
             ``from_pretrained`` also hits a warm cache.
@@ -334,6 +336,7 @@ def prefetch_subfolders(
             with _repo_prefetch_lock(model):
                 snapshot_download(
                     repo_id=model,
+                    revision=revision,
                     allow_patterns=allow_patterns,
                 )
             logger.info("Prefetch complete for %s", model)
@@ -450,6 +453,7 @@ def from_pretrained_with_prefetch(
     on the first failure exactly as before.
     """
     prefetch_list = list(prefetch_list)
+    revision = from_pretrained_kwargs.get("revision")
     can_heal = not local_files_only and bool(model) and not os.path.isdir(model)
     last_exc: BaseException | None = None
 
@@ -481,7 +485,7 @@ def from_pretrained_with_prefetch(
             # Force a fresh, verified snapshot of every component this pipeline
             # needs - not just ``subfolder`` - so a sibling component that was
             # also half-written gets repaired in the same pass.
-            prefetch_subfolders(model, prefetch_list, local_files_only=False)
+            prefetch_subfolders(model, prefetch_list, local_files_only=False, revision=revision)
 
     assert last_exc is not None  # loop only exits via return or a caught exc
     raise last_exc

vllm_omni/diffusion/models/ltx2/__init__.py

@@ -8,21 +8,25 @@ from vllm_omni.diffusion.models.ltx2.ltx2_components import (
 )
 from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel
 from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import (
+    LTX2DistilledOneStagePipeline,
     LTX2I2VDMD2Pipeline,
     LTX2Pipeline,
     LTX2T2VDMD2Pipeline,
 )
 from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import (
     LTX2DistilledPipeline,
+    LTX2DistilledTwoStagePipeline,
     LTX2TwoStagePipeline,
 )
 
 __all__ = [
     "LTX2Pipeline",
+    "LTX2DistilledOneStagePipeline",
     "LTX2T2VDMD2Pipeline",
     "LTX2I2VDMD2Pipeline",
     "LTX2TwoStagePipeline",
     "LTX2DistilledPipeline",
+    "LTX2DistilledTwoStagePipeline",
     "get_ltx2_post_process_func",
     "load_transformer_config",
     "create_transformer_from_config",

vllm_omni/diffusion/models/ltx2/ltx2_components.py

@@ -22,7 +22,7 @@ from diffusers.video_processor import VideoProcessor
 from huggingface_hub import hf_hub_download
 from safetensors import safe_open
 from safetensors.torch import load_file
-from transformers import AutoTokenizer, Gemma3ForConditionalGeneration
+from transformers import AutoModelForImageTextToText, AutoTokenizer, Gemma3ForConditionalGeneration
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.layer import Attention as OmniAttention
@@ -48,6 +48,14 @@ try:
 except ImportError:
     LTX2VocoderWithBWE = None
 
+try:
+    from transformers import Gemma4UnifiedForConditionalGeneration as _Gemma4UnifiedForConditionalGeneration
+except ImportError:
+    _Gemma4UnifiedForConditionalGeneration = None
+
+
+_LTX25_TEXT_ENCODER_CLS = AutoModelForImageTextToText if _Gemma4UnifiedForConditionalGeneration is not None else None
+
 
 _LTX_COMPONENT_SUBFOLDERS = (
     "tokenizer",
@@ -73,10 +81,15 @@ class LTXComponentProfile:
     resident_modules: tuple[str, ...] = ()
     video_vae_cls: type = AutoencoderKLLTX2Video
     vocoder_cls: type = LTX2Vocoder
+    text_encoder_cls: type | None = Gemma3ForConditionalGeneration
     vocoder_fallback_cls: type | None = None
     artifact_repo_id: str | None = None
     latent_upsampler_filename: str | None = None
     distilled_lora_filename: str | None = None
+    transformer_subfolder: str = "transformer"
+    scheduler_use_dynamic_shifting: bool = False
+    scheduler_shift_terminal: float | None = None
+    preserve_connector_attention_mask: bool = False
 
 
 LTX2_COMPONENT_PROFILE = LTXComponentProfile(
@@ -99,6 +112,23 @@ LTX23_COMPONENT_PROFILE = LTXComponentProfile(
     vocoder_fallback_cls=LTX2Vocoder,
 )
 
+LTX25_FULL_COMPONENT_PROFILE = LTXComponentProfile(
+    name="ltx2_5_full",
+    dit_modules=("transformer",),
+    encoder_modules=("text_encoder", "connectors"),
+    vae_modules=("vae", "audio_vae"),
+    resident_modules=("vocoder",),
+    video_vae_cls=DistributedAutoencoderKLLTX2Video,
+    vocoder_cls=LTX2VocoderWithBWE or LTX2Vocoder,
+    vocoder_fallback_cls=LTX2Vocoder,
+    text_encoder_cls=_LTX25_TEXT_ENCODER_CLS,
+    transformer_subfolder="transformer_full",
+    scheduler_use_dynamic_shifting=True,
+    scheduler_shift_terminal=0.1,
+    preserve_connector_attention_mask=True,
+)
+
+
 LTX2_DISTILLED_COMPONENT_PROFILE = LTXComponentProfile(
     name="ltx2_distilled",
     dit_modules=("transformer",),
@@ -108,6 +138,32 @@ LTX2_DISTILLED_COMPONENT_PROFILE = LTXComponentProfile(
     video_vae_cls=DistributedAutoencoderKLLTX2Video,
 )
 
+LTX25_DISTILLED_COMPONENT_PROFILE = LTXComponentProfile(
+    name="ltx2_5_distilled",
+    dit_modules=("transformer",),
+    encoder_modules=("text_encoder", "connectors"),
+    vae_modules=("vae", "audio_vae"),
+    resident_modules=("vocoder", "latent_upsampler"),
+    video_vae_cls=DistributedAutoencoderKLLTX2Video,
+    vocoder_cls=LTX2VocoderWithBWE or LTX2Vocoder,
+    vocoder_fallback_cls=LTX2Vocoder,
+    text_encoder_cls=_LTX25_TEXT_ENCODER_CLS,
+    preserve_connector_attention_mask=True,
+)
+
+LTX2_DISTILLED_ONE_STAGE_COMPONENT_PROFILE = replace(
+    LTX2_DISTILLED_COMPONENT_PROFILE,
+    name="ltx2_distilled_one_stage",
+    resident_modules=LTX2_COMPONENT_PROFILE.resident_modules,
+)
+
+LTX25_DISTILLED_ONE_STAGE_COMPONENT_PROFILE = replace(
+    LTX25_DISTILLED_COMPONENT_PROFILE,
+    name="ltx2_5_distilled_one_stage",
+    resident_modules=LTX25_FULL_COMPONENT_PROFILE.resident_modules,
+)
+
+
 LTX23_DISTILLED_COMPONENT_PROFILE = replace(
     LTX23_COMPONENT_PROFILE,
     name="ltx2_3_distilled",
@@ -116,6 +172,12 @@ LTX23_DISTILLED_COMPONENT_PROFILE = replace(
     latent_upsampler_filename="ltx-2.3-spatial-upscaler-x2-1.1.safetensors",
 )
 
+LTX23_DISTILLED_ONE_STAGE_COMPONENT_PROFILE = replace(
+    LTX23_DISTILLED_COMPONENT_PROFILE,
+    name="ltx2_3_distilled_one_stage",
+    resident_modules=LTX23_COMPONENT_PROFILE.resident_modules,
+)
+
 LTX2_TWO_STAGE_COMPONENT_PROFILE = replace(
     LTX2_COMPONENT_PROFILE,
     name="ltx2_two_stage",
@@ -135,18 +197,44 @@ LTX23_TWO_STAGE_COMPONENT_PROFILE = replace(
 )
 
 
+LTX25_TWO_STAGE_COMPONENT_PROFILE = replace(
+    LTX25_FULL_COMPONENT_PROFILE,
+    name="ltx2_5_two_stage",
+    resident_modules=(*LTX25_FULL_COMPONENT_PROFILE.resident_modules, "latent_upsampler"),
+    artifact_repo_id="Lightricks/LTX-2.5",
+    latent_upsampler_filename=("latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors"),
+    distilled_lora_filename="loras/ltx-2.5-22b-distilled-lora-450-bf16.safetensors",
+)
+
 _COMPONENT_PROFILES: dict[tuple[str, str], LTXComponentProfile] = {
     ("one_stage", "2"): LTX2_COMPONENT_PROFILE,
     ("one_stage", "2.3"): LTX23_COMPONENT_PROFILE,
+    ("one_stage", "2.5"): LTX25_FULL_COMPONENT_PROFILE,
     ("two_stage", "2"): LTX2_TWO_STAGE_COMPONENT_PROFILE,
     ("two_stage", "2.3"): LTX23_TWO_STAGE_COMPONENT_PROFILE,
+    ("two_stage", "2.5"): LTX25_TWO_STAGE_COMPONENT_PROFILE,
+    ("distilled_one_stage", "2"): LTX2_DISTILLED_ONE_STAGE_COMPONENT_PROFILE,
+    ("distilled_one_stage", "2.3"): LTX23_DISTILLED_ONE_STAGE_COMPONENT_PROFILE,
+    ("distilled_one_stage", "2.5"): LTX25_DISTILLED_ONE_STAGE_COMPONENT_PROFILE,
     ("distilled_two_stage", "2"): LTX2_DISTILLED_COMPONENT_PROFILE,
     ("distilled_two_stage", "2.3"): LTX23_DISTILLED_COMPONENT_PROFILE,
+    ("distilled_two_stage", "2.5"): LTX25_DISTILLED_COMPONENT_PROFILE,
     ("dmd2", "2"): LTX2_COMPONENT_PROFILE,
     ("dmd2", "2.3"): LTX23_COMPONENT_PROFILE,
 }
 
 
+def resolve_ltx_checkpoint_kind(pipeline_kind: str) -> LTXCheckpointKind | None:
+    """Derive checkpoint requirements from the execution contract."""
+    if pipeline_kind in {"one_stage", "two_stage"}:
+        return "regular"
+    if pipeline_kind in {"distilled_one_stage", "distilled_two_stage"}:
+        return "distilled"
+    if pipeline_kind == "dmd2":
+        return None
+    raise ValueError(f"Unsupported LTX pipeline kind: {pipeline_kind!r}.")
+
+
 def resolve_ltx_artifact(
     model: str,
     repo_id: str,
@@ -206,7 +294,7 @@ def resolve_ltx_component_profile(pipeline_kind: str, model_version: str) -> LTX
         raise ValueError(f"Unsupported LTX component kind/version: {pipeline_kind!r}/{model_version!r}.") from exc
 
 
-def _load_ltx_metadata_json(model: str, filename: str) -> dict[str, Any]:
+def _load_ltx_metadata_json(model: str, filename: str, revision: str | None = None) -> dict[str, Any]:
     """Load small checkpoint metadata without relying on repository names."""
     if os.path.isdir(model):
         path = os.path.join(model, filename)
@@ -214,7 +302,7 @@ def _load_ltx_metadata_json(model: str, filename: str) -> dict[str, Any]:
             return {}
     else:
         try:
-            path = hf_hub_download(repo_id=model, filename=filename)
+            path = hf_hub_download(repo_id=model, filename=filename, revision=revision)
         except Exception:
             return {}
     try:
@@ -225,17 +313,34 @@ def _load_ltx_metadata_json(model: str, filename: str) -> dict[str, Any]:
     return value if isinstance(value, dict) else {}
 
 
-def detect_ltx_model_version(model: str) -> str:
-    """Detect LTX-2 versus LTX-2.3 from checkpoint component metadata.
+def detect_ltx_model_version(model: str, revision: str | None = None) -> str:
+    """Detect the LTX model version from checkpoint component metadata."""
+    model_index = _load_ltx_metadata_json(model, "model_index.json", revision)
+    model_version = str(model_index.get("model_version", ""))
+    if model_version.startswith("2.5"):
+        return "2.5"
 
-    Official checkpoints use ``model_version`` metadata. Diffusers repositories
-    expose the equivalent distinction through the BWE vocoder introduced by
-    LTX-2.3. Unknown conversions retain the official LTX-2 fallback.
-    """
-    model_index = _load_ltx_metadata_json(model, "model_index.json")
-    if str(model_index.get("model_version", "")).startswith("2.3"):
+    text_encoder_entry = model_index.get("text_encoder")
+    if isinstance(text_encoder_entry, (list, tuple)) and text_encoder_entry:
+        text_encoder_class = str(text_encoder_entry[-1])
+    elif isinstance(text_encoder_entry, dict):
+        text_encoder_class = str(text_encoder_entry.get("_class_name", ""))
+    else:
+        text_encoder_class = ""
+    if text_encoder_class == "Gemma4UnifiedForConditionalGeneration":
+        return "2.5"
+
+    text_encoder_config = _load_ltx_metadata_json(model, "text_encoder/config.json", revision)
+    if text_encoder_config.get("model_type") in ("gemma4_unified", "gemma4"):
+        return "2.5"
+
+    if model_version.startswith("2.3"):
         return "2.3"
 
+    transformer_config = _load_ltx_metadata_json(model, "transformer/config.json", revision)
+    if transformer_config.get("ff_bias") is False:
+        return "2.5"
+
     vocoder_entry = model_index.get("vocoder")
     if isinstance(vocoder_entry, (list, tuple)) and vocoder_entry:
         vocoder_class = str(vocoder_entry[-1])
@@ -246,21 +351,32 @@ def detect_ltx_model_version(model: str) -> str:
     if vocoder_class == "LTX2VocoderWithBWE":
         return "2.3"
 
-    vocoder_config = _load_ltx_metadata_json(model, "vocoder/config.json")
+    vocoder_config = _load_ltx_metadata_json(model, "vocoder/config.json", revision)
     if str(vocoder_config.get("model_version", "")).startswith("2.3"):
         return "2.3"
     if vocoder_config.get("_class_name") == "LTX2VocoderWithBWE":
         return "2.3"
-
-    logger.info("Using LTX-2 defaults for checkpoint %s", model)
+    if "bwe" in vocoder_config or "bwe_config" in vocoder_config:
+        return "2.3"
     return "2"
 
 
+def preserves_reference_image_size(*, model: str | None, revision: str | None = None) -> bool:
+    """Preserve source geometry only for the LTX-2.5 CRF-18 I2V path."""
+    return model is not None and detect_ltx_model_version(model, revision=revision) == "2.5"
+
+
 class _LTXConnectorAttnProcessor:
     """Preserve official connector math around Omni attention dispatch."""
 
-    def __init__(self, *, has_learned_registers: bool) -> None:
+    def __init__(
+        self,
+        *,
+        has_learned_registers: bool = False,
+        preserve_learned_register_mask: bool = False,
+    ) -> None:
         self.has_learned_registers = has_learned_registers
+        self.preserve_learned_register_mask = preserve_learned_register_mask
 
     def __call__(
         self,
@@ -274,12 +390,20 @@ class _LTXConnectorAttnProcessor:
         encoder_hidden_states = encoder_hidden_states if encoder_hidden_states is not None else hidden_states
         gate_logits = attn.to_gate_logits(hidden_states) if attn.to_gate_logits is not None else None
 
-        query = attn.norm_q(attn.to_q(hidden_states))
-        key = attn.norm_k(attn.to_k(encoder_hidden_states))
+        query = attn.to_q(hidden_states)
+        key = attn.to_k(encoder_hidden_states)
         value = attn.to_v(encoder_hidden_states)
+        # Offload hooks may execute affine-free Q/K norms in FP32. Restore the
+        # projection dtype before attention, matching the fully resident path.
+        query = attn.norm_q(query).to(dtype=value.dtype)
+        key = attn.norm_k(key).to(dtype=value.dtype)
 
         if query_rotary_emb is not None:
+            # Diffusers builds connector RoPE in FP32, while the official
+            # connector materializes it in the hidden-state dtype.
+            query_rotary_emb = tuple(component.to(value.dtype) for component in query_rotary_emb)
             key_rotary_emb = key_rotary_emb if key_rotary_emb is not None else query_rotary_emb
+            key_rotary_emb = tuple(component.to(value.dtype) for component in key_rotary_emb)
             if attn.rope_type == "interleaved":
                 query = apply_interleaved_rotary_emb(query, query_rotary_emb)
                 key = apply_interleaved_rotary_emb(key, key_rotary_emb)
@@ -289,6 +413,10 @@ class _LTXConnectorAttnProcessor:
             else:
                 raise ValueError(f"Unsupported LTX connector RoPE type: {attn.rope_type}")
 
+        # Keep Q/K in the projection dtype expected by the attention backend.
+        query = query.to(dtype=value.dtype)
+        key = key.to(dtype=value.dtype)
+
         batch_size, _, inner_dim = query.shape
         head_dim = inner_dim // attn.heads
         kv_heads = attn.inner_kv_dim // attn.head_dim
@@ -296,10 +424,10 @@ class _LTXConnectorAttnProcessor:
         key = key.view(batch_size, -1, kv_heads, head_dim)
         value = value.view(batch_size, -1, kv_heads, head_dim)
 
-        # The connector replaces padding tokens with learned registers before
-        # entering its blocks, so every key is valid and the old padding mask
-        # becomes an all-keep no-op.
-        if self.has_learned_registers:
+        # LTX-2/2.3 replace padding tokens with learned registers, so their
+        # old padding mask becomes a no-op. LTX-2.5 preserves that all-zero
+        # additive mask to match the official masked-SDPA dispatch.
+        if self.has_learned_registers and not self.preserve_learned_register_mask:
             attention_mask = None
         elif attention_mask is not None and attn.omni_attention.attn_backend.get_name().upper() == "FLASH_ATTN":
             attention_mask = to_ltx_padding_mask(attention_mask)
@@ -316,7 +444,11 @@ class _LTXConnectorAttnProcessor:
         return attn.to_out[1](hidden_states)
 
 
-def _install_connector_attention(connectors: LTX2TextConnectors) -> None:
+def _install_connector_attention(
+    connectors: LTX2TextConnectors,
+    *,
+    preserve_learned_register_mask: bool = False,
+) -> None:
     for connector_name in ("video_connector", "audio_connector"):
         connector = getattr(connectors, connector_name, None)
         has_learned_registers = getattr(connector, "learnable_registers", None) is not None
@@ -335,15 +467,20 @@ def _install_connector_attention(connectors: LTX2TextConnectors) -> None:
                     skip_sequence_parallel=True,
                     disable_kv_quant=True,
                 )
-                attention.set_processor(_LTXConnectorAttnProcessor(has_learned_registers=has_learned_registers))
+                attention.set_processor(
+                    _LTXConnectorAttnProcessor(
+                        has_learned_registers=has_learned_registers,
+                        preserve_learned_register_mask=preserve_learned_register_mask,
+                    )
+                )
 
 
-def _detect_vocoder_output_sample_rate(model: str) -> int | None:
+def _detect_vocoder_output_sample_rate(model: str, revision: str | None = None) -> int | None:
     """Read the generated waveform sample rate from the vocoder config."""
     vocoder_config_path = os.path.join(model, "vocoder", "config.json")
     if not os.path.exists(vocoder_config_path):
         try:
-            vocoder_config_path = hf_hub_download(model, "vocoder/config.json")
+            vocoder_config_path = hf_hub_download(model, "vocoder/config.json", revision=revision)
         except Exception:
             return None
     try:
@@ -355,7 +492,10 @@ def _detect_vocoder_output_sample_rate(model: str) -> int | None:
 
 def get_ltx2_post_process_func(od_config: Any):
     """Build the common LTX engine-output adapter."""
-    output_sample_rate = _detect_vocoder_output_sample_rate(od_config.model)
+    output_sample_rate = _detect_vocoder_output_sample_rate(
+        od_config.model,
+        revision=getattr(od_config, "revision", None),
+    )
 
     def post_process_func(output: tuple[torch.Tensor, torch.Tensor] | torch.Tensor):
         if not (isinstance(output, tuple) and len(output) == 2):
@@ -378,6 +518,7 @@ def _load_component(
     *,
     local_files_only: bool,
     dtype: torch.dtype,
+    revision: str | None,
 ) -> Any:
     return from_pretrained_with_prefetch(
         component_cls.from_pretrained,
@@ -385,6 +526,7 @@ def _load_component(
         subfolder=subfolder,
         prefetch_list=_LTX_COMPONENT_SUBFOLDERS,
         local_files_only=local_files_only,
+        revision=revision,
         torch_dtype=dtype,
     )
 
@@ -411,31 +553,36 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None:
     pipeline.device = get_local_device()
     dtype = getattr(od_config, "dtype", torch.bfloat16)
     model = od_config.model
+    revision = getattr(od_config, "revision", None)
     local_files_only = os.path.exists(model)
 
     pipeline.weights_sources = [
         DiffusersPipelineLoader.ComponentSource(
             model_or_path=model,
-            subfolder="transformer",
-            revision=None,
+            subfolder=profile.transformer_subfolder,
+            revision=revision,
             prefix="transformer.",
             fall_back_to_pt=True,
         ),
     ]
-    prefetch_subfolders(model, _LTX_COMPONENT_SUBFOLDERS, local_files_only=local_files_only)
+    prefetch_subfolders(model, _LTX_COMPONENT_SUBFOLDERS, local_files_only=local_files_only, revision=revision)
 
     pipeline.tokenizer = AutoTokenizer.from_pretrained(
         model,
         subfolder="tokenizer",
         local_files_only=local_files_only,
+        revision=revision,
     )
+    if profile.text_encoder_cls is None:
+        raise ImportError("LTX-2.5 requires Gemma4UnifiedForConditionalGeneration; install transformers>=5.10.1,<5.15.")
     with torch.device("cpu"):
         pipeline.text_encoder = _load_component(
-            Gemma3ForConditionalGeneration,
+            profile.text_encoder_cls,
             model,
             "text_encoder",
             local_files_only=local_files_only,
             dtype=dtype,
+            revision=revision,
         )
     pipeline.connectors = _load_component(
         LTX2TextConnectors,
@@ -443,14 +590,19 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None:
         "connectors",
         local_files_only=local_files_only,
         dtype=dtype,
+        revision=revision,
+    )
+    _install_connector_attention(
+        pipeline.connectors,
+        preserve_learned_register_mask=profile.preserve_connector_attention_mask,
     )
-    _install_connector_attention(pipeline.connectors)
     pipeline.vae = _load_component(
         profile.video_vae_cls,
         model,
         "vae",
         local_files_only=local_files_only,
         dtype=dtype,
+        revision=revision,
     )
     pipeline.audio_vae = _load_component(
         AutoencoderKLLTX2Audio,
@@ -458,6 +610,7 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None:
         "audio_vae",
         local_files_only=local_files_only,
         dtype=dtype,
+        revision=revision,
     )
     try:
         pipeline.vocoder = _load_component(
@@ -466,6 +619,7 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None:
             "vocoder",
             local_files_only=local_files_only,
             dtype=dtype,
+            revision=revision,
         )
     except (TypeError, OSError, ValueError):
         if profile.vocoder_fallback_cls is None or profile.vocoder_fallback_cls is profile.vocoder_cls:
@@ -476,6 +630,7 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None:
             "vocoder",
             local_files_only=local_files_only,
             dtype=dtype,
+            revision=revision,
         )
 
     if "latent_upsampler" in profile.resident_modules:
@@ -488,6 +643,7 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None:
                     "latent_upsampler",
                     local_files_only=local_files_only,
                     dtype=dtype,
+                    revision=revision,
                 )
             except (OSError, ValueError):
                 if profile.latent_upsampler_filename is None or profile.artifact_repo_id is None:
@@ -508,7 +664,9 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None:
             )
             pipeline.latent_upsampler = _load_ltx_latent_upsampler_single_file(upsampler_path, dtype)
 
-    transformer_config = load_transformer_config(model, "transformer", local_files_only)
+    transformer_config = load_transformer_config(
+        model, profile.transformer_subfolder, local_files_only, revision=revision
+    )
     quant_config = getattr(od_config, "quantization_config", None)
     pipeline.transformer = create_transformer_from_config(transformer_config, quant_config=quant_config)
     _place_aux_components(pipeline)
@@ -516,15 +674,17 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None:
         model,
         subfolder="scheduler",
         local_files_only=local_files_only,
+        revision=revision,
     )
-    expected_checkpoint_kind: LTXCheckpointKind | None = None
-    if pipeline.pipeline_kind == "two_stage":
-        expected_checkpoint_kind = "regular"
-    elif pipeline.pipeline_kind == "distilled_two_stage":
-        expected_checkpoint_kind = "distilled"
+    if profile.scheduler_use_dynamic_shifting:
+        pipeline.scheduler = FlowMatchEulerDiscreteScheduler.from_config(
+            pipeline.scheduler.config,
+            use_dynamic_shifting=True,
+            shift_terminal=profile.scheduler_shift_terminal,
+        )
     validate_ltx_checkpoint(
         pipeline.scheduler.config,
-        expected_kind=expected_checkpoint_kind,
+        expected_kind=resolve_ltx_checkpoint_kind(pipeline.pipeline_kind),
         pipeline_name=type(pipeline).__name__,
     )
 
@@ -553,6 +713,8 @@ def load_transformer_config(
     model_path: str,
     subfolder: str = "transformer",
     local_files_only: bool = True,
+    *,
+    revision: str | None = None,
 ) -> dict:
     """Load an LTX transformer config from a local model or the HF Hub."""
     if local_files_only:
@@ -563,6 +725,7 @@ def load_transformer_config(
         config_path = hf_hub_download(
             repo_id=model_path,
             filename=f"{subfolder}/config.json",
+            revision=revision,
         )
     with open(config_path) as config_file:
         return json.load(config_file)

vllm_omni/diffusion/models/ltx2/ltx2_conditioning.py

@@ -7,6 +7,7 @@ from __future__ import annotations
 
 import math
 from dataclasses import dataclass
+from io import BytesIO
 from typing import TYPE_CHECKING, Any
 
 import numpy as np
@@ -52,19 +53,18 @@ def _preprocess_i2v_pil_images(
     *,
     height: int,
     width: int,
+    crf: int = 0,
     device: torch.device | None = None,
     dtype: torch.dtype | None = None,
 ) -> torch.Tensor:
-    """Match official LTX aspect-preserving resize and center crop."""
+    """Match the version-specific official LTX PIL preprocessing."""
     image_list = images if isinstance(images, list) else [images]
     processed = []
     for image in image_list:
-        pixels = (
-            torch.from_numpy(np.array(image.convert("RGB"), copy=True))
-            .permute(2, 0, 1)
-            .unsqueeze(0)
-            .to(device=device, dtype=torch.float32)
-        )
+        image_array = np.array(image.convert("RGB"), copy=True)
+        if crf:
+            image_array = _apply_image_conditioning_crf(image_array, crf)
+        pixels = torch.from_numpy(image_array).permute(2, 0, 1).unsqueeze(0).to(device=device, dtype=torch.float32)
         source_height, source_width = pixels.shape[-2:]
         scale = max(height / source_height, width / source_width)
         resized_height = math.ceil(source_height * scale)
@@ -82,6 +82,47 @@ def _preprocess_i2v_pil_images(
     return torch.cat(processed, dim=0)
 
 
+def _apply_image_conditioning_crf(image: np.ndarray, crf: int) -> np.ndarray:
+    """Round-trip one RGB frame through H.264 to match LTX-2.5 training."""
+    if image.dtype != np.uint8 or image.ndim != 3 or image.shape[-1] != 3:
+        raise ValueError("LTX image-conditioning CRF expects a uint8 RGB image.")
+    if min(image.shape[:2]) < 2:
+        return image
+    try:
+        import av
+    except ImportError as exc:
+        raise ImportError("PyAV with a libx264 encoder is required for LTX-2.5 image conditioning (CRF 18).") from exc
+
+    with BytesIO() as output_file:
+        container = av.open(output_file, "w", format="mp4")
+        try:
+            stream = container.add_stream("libx264", rate=1, options={"crf": str(crf), "preset": "veryfast"})
+            height = image.shape[0] // 2 * 2
+            width = image.shape[1] // 2 * 2
+            stream.height = height
+            stream.width = width
+            frame = av.VideoFrame.from_ndarray(image[:height, :width], format="rgb24").reformat(format="yuv420p")
+            container.mux(stream.encode(frame))
+            container.mux(stream.encode())
+        except Exception as exc:
+            raise RuntimeError(
+                "The installed PyAV/FFmpeg build cannot encode libx264 video required "
+                "for LTX-2.5 image conditioning (CRF 18)."
+            ) from exc
+        finally:
+            container.close()
+        video_bytes = output_file.getvalue()
+
+    with BytesIO(video_bytes) as video_file:
+        container = av.open(video_file)
+        try:
+            stream = next(stream for stream in container.streams if stream.type == "video")
+            frame = next(container.decode(stream))
+        finally:
+            container.close()
+    return frame.to_ndarray(format="rgb24")
+
+
 class LTXTextConditioningMixin:
     """Shared Gemma encoding and text connector orchestration."""
 
@@ -106,13 +147,30 @@ class LTXTextConditioningMixin:
         if self.tokenizer.pad_token is None:
             self.tokenizer.pad_token = self.tokenizer.eos_token
 
-        text_inputs = self.tokenizer(
-            [text.strip() for text in prompt],
+        bos_token_id = self.tokenizer.bos_token_id
+        if bos_token_id is None:
+            raise ValueError("Tokenizer is missing bos_token_id; LTX prompt encoding requires a leading BOS.")
+
+        input_ids = []
+        for text in prompt:
+            encoded = self.tokenizer(
+                text.strip(),
+                padding=False,
+                truncation=True,
+                max_length=max_sequence_length,
+                return_tensors=None,
+            )
+            ids = encoded.input_ids
+            if not ids or ids[0] != bos_token_id:
+                ids = [bos_token_id, *ids][:max_sequence_length]
+            input_ids.append(ids)
+
+        text_inputs = self.tokenizer.pad(
+            {"input_ids": input_ids},
             padding="max_length",
             max_length=max_sequence_length,
-            truncation=True,
-            add_special_tokens=True,
             return_tensors="pt",
+            return_attention_mask=True,
         )
         text_input_ids = text_inputs.input_ids.to(device)
         prompt_attention_mask = text_inputs.attention_mask.to(device)
@@ -121,7 +179,6 @@ class LTXTextConditioningMixin:
             attention_mask=prompt_attention_mask,
             output_hidden_states=True,
         ).hidden_states
-
         prompt_embeds = torch.stack(hidden_states, dim=-1).flatten(2, 3).to(dtype=dtype)
         prompt_embeds = _repeat_prompt_tensor_for_outputs(prompt_embeds, num_videos_per_prompt)
         prompt_attention_mask = prompt_attention_mask.view(batch_size, -1)
@@ -501,6 +558,7 @@ class LTXI2VConditioningMixin:
             vae_temporal_compression_ratio=self.vae_temporal_compression_ratio,
         )
         mask_shape = (batch_size, 1, num_frames, height, width)
+        is_ltx25 = getattr(self, "model_version", None) == "2.5"
 
         if latents is not None:
             unpacked_latents = latents.ndim == 5
@@ -518,6 +576,7 @@ class LTXI2VConditioningMixin:
                     self.vae.latents_std,
                     self.vae.config.scaling_factor,
                 )
+                clean_latents = latents
                 if image is not None:
                     image_latents = self._encode_i2v_image_latents(
                         image,
@@ -525,14 +584,19 @@ class LTXI2VConditioningMixin:
                         generator=generator,
                         dtype=dtype,
                     )
-                    # Official distilled Stage 2 starts from the upsampled Stage 1
-                    # state, then reapplies the source image at final resolution.
-                    latents = self._replace_first_latent_frame(latents, image_latents)
+                    # Official Stage 2 starts from the upsampled Stage 1 state,
+                    # then reapplies the source image at final resolution.
+                    clean_latents = self._replace_first_latent_frame(clean_latents, image_latents)
                 latents = latent_ops.pack_latents(
                     latents,
                     self.transformer_spatial_patch_size,
                     self.transformer_temporal_patch_size,
                 )
+                clean_latents = latent_ops.pack_latents(
+                    clean_latents,
+                    self.transformer_spatial_patch_size,
+                    self.transformer_temporal_patch_size,
+                )
             else:
                 conditioning_mask = latents.new_zeros(mask_shape)
                 conditioning_mask[:, :, 0] = 1.0
@@ -548,14 +612,14 @@ class LTXI2VConditioningMixin:
                     f" {latents.shape}, but the expected shape is {conditioning_mask.shape + (num_channels_latents,)}."
                 )
             if unpacked_latents:
-                # Official LTX patchifies the latent state before drawing noise.
-                # Sampling the packed shape preserves its seed-to-output contract.
-                latents = latent_ops.create_noised_state(
+                latents = latent_ops.create_conditioned_noised_state(
                     latents,
-                    noise_scale * (1 - conditioning_mask).unsqueeze(-1),
+                    clean_latents,
+                    (1 - conditioning_mask.float()).unsqueeze(-1),
+                    noise_scale,
                     generator,
                 )
-            elif image is not None:
+            elif not unpacked_latents and image is not None:
                 image_latents = self._encode_i2v_image_latents(
                     image,
                     batch_size=batch_size,
@@ -595,13 +659,27 @@ class LTXI2VConditioningMixin:
             self.transformer_spatial_patch_size,
             self.transformer_temporal_patch_size,
         )
-        conditioning_mask = latent_ops.pack_latents(
+        packed_conditioning_mask = latent_ops.pack_latents(
             conditioning_mask,
             self.transformer_spatial_patch_size,
             self.transformer_temporal_patch_size,
         ).squeeze(-1)
-        noise = randn_tensor(init_latents.shape, generator=generator, device=device, dtype=dtype)
-        latents = init_latents * conditioning_mask.unsqueeze(-1) + noise * (1 - conditioning_mask).unsqueeze(-1)
+        if is_ltx25:
+            clean_latents = init_latents
+            init_latents = torch.zeros_like(init_latents)
+            latents = latent_ops.create_conditioned_noised_state(
+                init_latents,
+                clean_latents,
+                (1 - packed_conditioning_mask.float()).unsqueeze(-1),
+                noise_scale,
+                generator=generator,
+            )
+        else:
+            noise = randn_tensor(init_latents.shape, generator=generator, device=device, dtype=dtype)
+            latents = init_latents * packed_conditioning_mask.unsqueeze(-1) + noise * (
+                1 - packed_conditioning_mask
+            ).unsqueeze(-1)
+        conditioning_mask = packed_conditioning_mask
         return latents, conditioning_mask
 
     def _step_video_latents_i2v(
@@ -665,6 +743,17 @@ class LTXI2VConditioningMixin:
             )
 
         if image is not None:
+            default_crf = 18 if getattr(self, "model_version", "2") == "2.5" else 0
+            image_crf = getattr(request_inputs, "image_crf", None)
+            image_crf = default_crf if image_crf is None else int(image_crf)
+            is_pil_input = isinstance(image, PIL.Image.Image) or (
+                isinstance(image, list) and image and all(isinstance(item, PIL.Image.Image) for item in image)
+            )
+            if image_crf != 0 and not is_pil_input:
+                raise ValueError(
+                    "`image_crf` re-compression requires PIL images. Pass PIL images, or set `image_crf=0` "
+                    "to skip re-compression."
+                )
             if isinstance(image, torch.Tensor):
                 if image.ndim == 3:
                     image = image.unsqueeze(0)
@@ -677,6 +766,7 @@ class LTXI2VConditioningMixin:
                     image,
                     height=request_inputs.height,
                     width=request_inputs.width,
+                    crf=image_crf,
                     device=device,
                     dtype=prompt_context.positive_connector_prompt_embeds.dtype,
                 )

vllm_omni/diffusion/models/ltx2/ltx2_denoise.py

@@ -42,6 +42,7 @@ class LTXForwardContext:
     original_audio_num_frames: int
     padded_audio_num_frames: int
     timesteps: torch.Tensor
+    sampler: str
     audio_scheduler: Any
     video_audio_step_adapter: Any
 
@@ -88,6 +89,11 @@ class LTXDenoisePipeline(Protocol):
 
 LTXDenoiseStep = Callable[[int, torch.Tensor, LTXAVState], LTXAVState]
 
+LTX_ANCESTRAL_ETA = 1.0
+LTX_ANCESTRAL_S_NOISE = 1.0
+LTX_ANCESTRAL_NOISE_SEED_OFFSET = 10000
+LTX_OFFICIAL_DEFAULT_SEED = 10
+
 
 class LTXDenoiseExecutor:
     """Run the one shared LTX denoise loop.
@@ -138,6 +144,9 @@ class LTXVideoAudioStepAdapter:
         latent_width: int,
         *,
         image_conditioned: bool,
+        sampler: str = "euler",
+        generator: torch.Generator | list[torch.Generator] | None = None,
+        conditioning_mask: torch.Tensor | None = None,
     ) -> None:
         self._pipeline = pipeline
         self._audio_scheduler = audio_scheduler
@@ -145,11 +154,27 @@ class LTXVideoAudioStepAdapter:
         self._latent_height = latent_height
         self._latent_width = latent_width
         self._image_conditioned = image_conditioned
+        self._sampler = sampler
+        self._conditioning_mask = conditioning_mask
+        self._ancestral_generators = (
+            _make_ancestral_generators(generator, pipeline.device) if sampler == "euler_ancestral" else None
+        )
         self._step_index = 0
 
     def step(self, noise_pred, t, latents, return_dict=False, generator=None):
         del t, return_dict, generator
-        if self._image_conditioned:
+        if self._sampler == "euler_ancestral":
+            video_out = _ancestral_euler_step_from_velocity(
+                latents[0],
+                noise_pred[0],
+                self._pipeline.scheduler.sigmas,
+                self._step_index,
+                self._ancestral_generators,
+            )
+            if self._conditioning_mask is not None:
+                mask = self._conditioning_mask.unsqueeze(-1).to(video_out.dtype)
+                video_out = torch.lerp(video_out, latents[0], mask)
+        elif self._image_conditioned:
             video_out = self._pipeline._step_video_latents_i2v(
                 noise_pred[0],
                 latents[0],
@@ -165,17 +190,96 @@ class LTXVideoAudioStepAdapter:
                 self._pipeline.scheduler.sigmas,
                 self._step_index,
             )
-        audio_out = euler_step_from_velocity(
-            latents[1],
-            noise_pred[1],
-            self._audio_scheduler.sigmas,
-            self._step_index,
-        )
+        if self._sampler == "euler_ancestral":
+            audio_out = _ancestral_euler_step_from_velocity(
+                latents[1],
+                noise_pred[1],
+                self._audio_scheduler.sigmas,
+                self._step_index,
+                self._ancestral_generators,
+            )
+        else:
+            audio_out = euler_step_from_velocity(
+                latents[1],
+                noise_pred[1],
+                self._audio_scheduler.sigmas,
+                self._step_index,
+            )
         self._step_index += 1
         return ((video_out, audio_out),)
 
 
-def _official_ltx_sigmas(scheduler: Any, steps: int, device: torch.device) -> torch.Tensor:
+def _make_ancestral_generators(
+    generator: torch.Generator | list[torch.Generator] | None,
+    device: torch.device,
+) -> torch.Generator | list[torch.Generator]:
+    """Create the official independent ``seed + 10000`` noise stream."""
+
+    def _offset(source: torch.Generator | None) -> torch.Generator:
+        seed = LTX_OFFICIAL_DEFAULT_SEED if source is None else source.initial_seed()
+        return torch.Generator(device=device).manual_seed(seed + LTX_ANCESTRAL_NOISE_SEED_OFFSET)
+
+    if isinstance(generator, list):
+        return [_offset(item) for item in generator]
+    return _offset(generator)
+
+
+def _randn_like_with_generators(
+    sample: torch.Tensor,
+    generators: torch.Generator | list[torch.Generator],
+) -> torch.Tensor:
+    if not isinstance(generators, list):
+        return torch.randn(sample.shape, generator=generators, dtype=sample.dtype, device=sample.device)
+    if len(generators) != sample.shape[0]:
+        raise ValueError(
+            f"LTX ancestral sampling received {len(generators)} generators for batch size {sample.shape[0]}."
+        )
+    return torch.cat(
+        [
+            torch.randn((1, *sample.shape[1:]), generator=item, dtype=sample.dtype, device=sample.device)
+            for item in generators
+        ]
+    )
+
+
+def _ancestral_euler_step_from_velocity(
+    sample: torch.Tensor,
+    velocity: torch.Tensor,
+    sigmas: torch.Tensor,
+    step_index: int,
+    generators: torch.Generator | list[torch.Generator] | None,
+) -> torch.Tensor:
+    """Apply the official LTX-2.5 rectified-flow ancestral Euler step."""
+    sigma = sigmas[step_index].to(torch.float32)
+    sigma_next = sigmas[step_index + 1].to(torch.float32)
+    # Official X0Model materializes the denoised prediction in the model
+    # dtype before the ancestral step promotes it back to float32.
+    denoised_model_dtype = (sample.float() - velocity.float() * sigma).to(sample.dtype)
+    denoised = denoised_model_dtype.float()
+    if sigma_next == 0:
+        return denoised.to(sample.dtype)
+    if generators is None:
+        raise ValueError("LTX ancestral Euler sampling requires a noise generator.")
+
+    downstep_ratio = 1.0 + (sigma_next / sigma - 1.0) * LTX_ANCESTRAL_ETA
+    sigma_down = sigma_next * downstep_ratio
+    sigma_down_ratio = sigma_down / sigma
+    x_next = sigma_down_ratio * sample.float() + (1.0 - sigma_down_ratio) * denoised
+
+    alpha_next = 1.0 - sigma_next
+    alpha_down = 1.0 - sigma_down
+    renoise_coeff = (sigma_next**2 - sigma_down**2 * alpha_next**2 / alpha_down**2).clamp(min=0).sqrt()
+    noise = _randn_like_with_generators(sample, generators)
+    x_next = (alpha_next / alpha_down) * x_next + noise.float() * LTX_ANCESTRAL_S_NOISE * renoise_coeff
+    result = x_next.to(sample.dtype)
+    return result
+
+
+def _official_ltx_sigmas(
+    scheduler: Any,
+    steps: int,
+    device: torch.device,
+) -> torch.Tensor:
     """Build the official LTX one-stage sigma schedule in torch fp32."""
     config = scheduler.config
     base_anchor = config.get("base_image_seq_len", 1024)
@@ -185,7 +289,8 @@ def _official_ltx_sigmas(scheduler: Any, steps: int, device: torch.device) -> to
 
     sigmas = torch.linspace(1.0, 0.0, steps + 1)
     slope = (max_shift - base_shift) / (max_anchor - base_anchor)
-    # Official LTX2 one-stage intentionally uses the max sequence anchor, so the shift stays at max_shift.
+    # Official LTX one-stage pipelines omit the latent when constructing
+    # this schedule, so the shift stays at the max sequence anchor.
     sigma_shift = max_anchor * slope + (base_shift - slope * base_anchor)
     exp_shift = math.exp(sigma_shift)
     sigmas = torch.where(sigmas != 0, exp_shift / (exp_shift + (1 / sigmas - 1)), 0)
@@ -224,6 +329,9 @@ def prepare_scheduler_stage(
     latent_width: int,
     use_official_sigma_schedule: bool,
     image_conditioned: bool = False,
+    sampler: str = "euler",
+    generator: torch.Generator | list[torch.Generator] | None = None,
+    conditioning_mask: torch.Tensor | None = None,
 ) -> tuple[Any, Any, torch.Tensor]:
     if sigmas is not None and timesteps is not None:
         raise ValueError("Only one of `sigmas` or `timesteps` may be provided.")
@@ -236,6 +344,9 @@ def prepare_scheduler_stage(
         latent_height,
         latent_width,
         image_conditioned=image_conditioned,
+        sampler=sampler,
+        generator=generator,
+        conditioning_mask=conditioning_mask,
     )
     if sigmas is not None:
         scheduler_sigmas = torch.as_tensor(sigmas, dtype=torch.float32, device=device)
@@ -302,6 +413,20 @@ def prepare_rope_coords_stage(
     return video_coords, audio_coords
 
 
+def _first_frame_keyframes_mask(reference: torch.Tensor, latent_num_frames: int) -> torch.Tensor:
+    """Mark the first causal latent frame, matching official LTX-2.5."""
+    if latent_num_frames <= 0:
+        raise ValueError(f"LTX latent frame count must be positive, got {latent_num_frames}.")
+    tokens_per_frame, remainder = divmod(reference.shape[1], latent_num_frames)
+    if remainder:
+        raise ValueError(
+            f"LTX video token count {reference.shape[1]} is not divisible by {latent_num_frames} latent frames."
+        )
+    mask = reference.new_zeros((reference.shape[0], reference.shape[1], 1))
+    mask[:, :tokens_per_frame] = 1
+    return mask
+
+
 def build_transformer_kwargs(
     pipeline: Any,
     forward_ctx: LTXForwardContext,
@@ -317,9 +442,13 @@ def build_transformer_kwargs(
     attention_kwargs: dict[str, Any] | None = None,
 ) -> dict[str, Any]:
     del encoder_attention_mask, audio_encoder_attention_mask
+    keyframes_mask = None
+    if getattr(pipeline.transformer.config, "use_keyframes_abs_pos_embedding", False):
+        keyframes_mask = _first_frame_keyframes_mask(hidden_states, forward_ctx.latent_num_frames)
     return {
         "hidden_states": hidden_states,
         "audio_hidden_states": audio_hidden_states,
+        "keyframes_mask": keyframes_mask,
         "encoder_hidden_states": encoder_hidden_states,
         "audio_encoder_hidden_states": audio_encoder_hidden_states,
         **pipeline._denoise_timestep_kwargs(
@@ -426,6 +555,9 @@ class LTXPhaseExecutor:
             latent_width=latent_width,
             use_official_sigma_schedule=phase_recipe.use_official_sigma_schedule,
             image_conditioned=conditioning_mask is not None,
+            sampler=phase_recipe.sampler,
+            generator=request_inputs.generator,
+            conditioning_mask=conditioning_mask,
         )
         forward_ctx = LTXForwardContext(
             req=req,
@@ -441,6 +573,7 @@ class LTXPhaseExecutor:
             original_audio_num_frames=original_audio_num_frames,
             padded_audio_num_frames=padded_audio_num_frames,
             timesteps=timesteps_tensor,
+            sampler=phase_recipe.sampler,
             audio_scheduler=audio_scheduler,
             video_audio_step_adapter=video_audio_step_adapter,
         )

vllm_omni/diffusion/models/ltx2/ltx2_guidance.py

@@ -126,9 +126,9 @@ def x0_from_velocity(sample: torch.Tensor, velocity: torch.Tensor, sigma: torch.
 
 def velocity_from_x0(sample: torch.Tensor, x0: torch.Tensor, sigma: torch.Tensor) -> torch.Tensor:
     """Convert x0 back to velocity using the official fp32 arithmetic order."""
-    # Official LTX converts the scalar scheduler sigma to a host value before
-    # division. A CUDA scalar tensor selects different kernel arithmetic and
-    # can cross bf16 rounding boundaries late in the denoise trajectory.
+    # Official LTX materializes the scalar schedule value on the host before
+    # dividing. CUDA tensor-scalar division can cross bfloat16 rounding bounds
+    # even when the scalar has the same float32 value.
     sigma = sigma.to(torch.float32).item()
     return ((sample.to(torch.float32) - x0.to(torch.float32)) / sigma).to(sample.dtype)
 
@@ -305,10 +305,17 @@ class LTXGuidanceExecutor:
         ts: torch.Tensor,
         video_token_count: int,
         audio_token_count: int,
+        *,
+        expand_for_sequence_parallel: bool = False,
     ) -> dict[str, torch.Tensor]:
+        timestep = ts
+        audio_timestep = ts
+        if expand_for_sequence_parallel:
+            timestep = ts.reshape(-1, 1).expand(-1, video_token_count)
+            audio_timestep = ts.reshape(-1, 1).expand(-1, audio_token_count)
         return {
-            "timestep": ts.reshape(-1, 1).expand(-1, video_token_count),
-            "audio_timestep": ts.reshape(-1, 1).expand(-1, audio_token_count),
+            "timestep": timestep,
+            "audio_timestep": audio_timestep,
             "sigma": ts,
             "audio_sigma": ts,
         }
@@ -342,8 +349,11 @@ class LTXGuidanceExecutor:
         guidance: LTXModalityGuidance,
         *,
         model_sigma: torch.Tensor | None = None,
+        preserve_positive_velocity: bool = False,
         rescale_token_count: int | None = None,
     ) -> torch.Tensor:
+        if preserve_positive_velocity and tuple(splits) == ("cond",):
+            return splits["cond"]
         model_sigma = sigma if model_sigma is None else model_sigma
         # Official guidance reduces contiguous BSC tensors. The fused
         # transformer output may be channel-major after splitting, which
@@ -368,6 +378,7 @@ class LTXGuidanceExecutor:
         state: LTXAVState,
         forward_ctx: LTXForwardContext,
         denoise_ctx: LTXDenoiseContext,
+        preserve_positive_velocity: bool = False,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         guidance_world_size = get_guidance_parallel_world_size()
         self.validate_guidance_world_size(plan, guidance_world_size)
@@ -450,12 +461,14 @@ class LTXGuidanceExecutor:
                 video_sigma,
                 plan.spec.video,
                 model_sigma=pipeline._video_guidance_model_sigma(video_sigma, denoise_ctx),
+                preserve_positive_velocity=preserve_positive_velocity,
             ),
             self._guide_modality(
                 state.audio,
                 audio_splits,
                 forward_ctx.audio_scheduler.sigmas[index],
                 plan.spec.audio,
+                preserve_positive_velocity=preserve_positive_velocity,
                 rescale_token_count=forward_ctx.original_audio_num_frames,
             ),
         )
@@ -469,9 +482,19 @@ class LTXGuidanceExecutor:
         state: LTXAVState,
         forward_ctx: LTXForwardContext,
         denoise_ctx: LTXDenoiseContext,
+        preserve_positive_velocity: bool = False,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         if forward_ctx.guidance_parallel_ready:
-            return self.predict_parallel_guidance(pipeline, plan, index, timestep, state, forward_ctx, denoise_ctx)
+            return self.predict_parallel_guidance(
+                pipeline,
+                plan,
+                index,
+                timestep,
+                state,
+                forward_ctx,
+                denoise_ctx,
+                preserve_positive_velocity=preserve_positive_velocity,
+            )
 
         prompt = forward_ctx.prompt_context
         video_contexts: list[torch.Tensor] = []
@@ -526,12 +549,14 @@ class LTXGuidanceExecutor:
                     pipeline.scheduler.sigmas[index],
                     denoise_ctx,
                 ),
+                preserve_positive_velocity=preserve_positive_velocity,
             ),
             self._guide_modality(
                 state.audio,
                 audio_splits,
                 forward_ctx.audio_scheduler.sigmas[index],
                 plan.spec.audio,
+                preserve_positive_velocity=preserve_positive_velocity,
                 rescale_token_count=forward_ctx.original_audio_num_frames,
             ),
         )

vllm_omni/diffusion/models/ltx2/ltx2_latents.py

@@ -20,6 +20,11 @@ class LTXAVState:
     audio: torch.Tensor
 
 
+def official_video_token_layout(latents: torch.Tensor) -> torch.Tensor:
+    """Match the token-major view produced by the official LTX patchifier."""
+    return latents.transpose(1, 2).contiguous().transpose(1, 2)
+
+
 def pack_latents(
     latents: torch.Tensor,
     patch_size: int = 1,
@@ -117,9 +122,23 @@ def create_noised_state(
         device=latents.device,
         dtype=latents.dtype,
     )
-    if isinstance(noise_scale, torch.Tensor):
-        noise_scale = noise_scale.float()
-    return torch.lerp(latents.float(), noise.float(), noise_scale).to(latents.dtype)
+    # Official GaussianNoisier blends in fp32, then materializes the next
+    # latent state in the model dtype. Batched request weights may arrive in
+    # the model dtype, so promote tensor weights alongside both endpoints.
+    lerp_weight = noise_scale.float() if isinstance(noise_scale, torch.Tensor) else noise_scale
+    return torch.lerp(latents.float(), noise.float(), lerp_weight).to(latents.dtype)
+
+
+def create_conditioned_noised_state(
+    latents: torch.Tensor,
+    clean_latents: torch.Tensor,
+    denoise_mask: torch.Tensor,
+    noise_scale: float | torch.Tensor,
+    generator: torch.Generator | list[torch.Generator] | None = None,
+) -> torch.Tensor:
+    """Match the official conditioned Gaussian noisier operation order."""
+    noised_latents = create_noised_state(latents, noise_scale, generator)
+    return torch.lerp(clean_latents.float(), noised_latents.float(), denoise_mask.float()).to(latents.dtype)
 
 
 def pack_audio_latents(latents: torch.Tensor) -> torch.Tensor:
@@ -138,7 +157,7 @@ def unpad_audio_latents(latents: torch.Tensor, num_frames: int) -> torch.Tensor:
 
 
 def clear_audio_padding(latents: torch.Tensor, num_frames: int) -> torch.Tensor:
-    """Keep Omni's SP-only audio padding outside the sampler state."""
+    """Keep SP-only audio padding outside the logical sampler state."""
     if not 0 < num_frames <= latents.shape[1]:
         raise ValueError(f"Audio frame count must be in [1, {latents.shape[1]}], got {num_frames}.")
     if num_frames == latents.shape[1]:
@@ -196,7 +215,8 @@ def prepare_video_latents(
             )
         if latents.ndim != 3:
             raise ValueError(f"Provided `latents` has shape {latents.shape}, expected [batch, seq, features].")
-        return create_noised_state(latents, noise_scale, generator).to(device=device, dtype=dtype)
+        latents = create_noised_state(latents, noise_scale, generator).to(device=device, dtype=dtype)
+        return official_video_token_layout(latents)
 
     num_frames, height, width = resolve_video_latent_shape(
         height,
@@ -218,11 +238,7 @@ def prepare_video_latents(
             f" size of {batch_size}. Make sure the batch size matches the length of the generators."
         )
     latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
-    # Official LTX samples the 5D video latent first and then patchifies it,
-    # leaving token-major storage behind the [B, tokens, channels] view.  The
-    # logical values are unchanged, but preserving that layout is important:
-    # CUDA selects a different bf16 GEMM path for a contiguous token tensor.
-    return latents.transpose(1, 2).contiguous().transpose(1, 2)
+    return official_video_token_layout(latents)
 
 
 def prepare_audio_latents(

vllm_omni/diffusion/models/ltx2/ltx2_recipes.py

@@ -3,7 +3,7 @@
 
 """Declarative execution recipes for the LTX model family."""
 
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 from typing import Literal
 
 from .ltx2_guidance import LTXGuidanceSpec, LTXModalityGuidance
@@ -22,6 +22,11 @@ LTX_DEFAULT_NEGATIVE_PROMPT = (
     "inconsistent tone, cinematic oversaturation, stylized filters, or AI artifacts."
 )
 
+LTX25_DEFAULT_NEGATIVE_PROMPT = (
+    "has_subtitles, has_blurbox, transition from black, transition to black, speech_ending_short, "
+    + LTX_DEFAULT_NEGATIVE_PROMPT
+)
+
 LTX_DISTILLED_SIGMAS = (1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875, 0.0)
 LTX_STAGE_2_DISTILLED_SIGMAS = (0.909375, 0.725, 0.421875, 0.0)
 LTX_DISTILLED_ADAPTER_SLOT = "ltx_distilled"
@@ -38,6 +43,7 @@ class LTXPhaseRecipe:
     noise_scale: float = 0.0
     input_transform: Literal["initial", "spatial_upsample"] = "initial"
     adapter_slot: str | None = None
+    sampler: Literal["euler", "euler_ancestral"] = "euler"
     allow_guidance_override: bool = True
     use_official_sigma_schedule: bool = True
 
@@ -46,6 +52,8 @@ class LTXPhaseRecipe:
             raise ValueError("LTX phase spatial_downscale must be positive.")
         if self.sigmas is not None and len(self.sigmas) < 2:
             raise ValueError("An explicit LTX sigma schedule must contain at least two denoise sigmas.")
+        if self.sampler not in ("euler", "euler_ancestral"):
+            raise ValueError(f"Unsupported LTX sampler: {self.sampler!r}.")
 
     @property
     def num_inference_steps(self) -> int | None:
@@ -66,9 +74,11 @@ class LTXPipelineRecipe:
     video_output_phase: int = -1
     audio_output_phase: int = -1
     allow_request_sigmas: bool = True
+    allow_request_phase_sigmas: bool = False
     allow_request_latents: bool = True
     allow_negative_prompt: bool = True
     fixed_num_inference_steps: bool = False
+    supports_cache_dit: bool = False
 
     def __post_init__(self) -> None:
         if not self.phases:
@@ -109,13 +119,29 @@ def _official_guidance(stg_block: int) -> LTXGuidanceSpec:
 
 
 LTX2_ONE_STAGE_RECIPE = LTXPipelineRecipe(
+    supports_cache_dit=True,
     phases=(LTXPhaseRecipe(name="generate", guidance=_official_guidance(29)),),
 )
 LTX23_ONE_STAGE_RECIPE = LTXPipelineRecipe(
+    supports_cache_dit=True,
     num_inference_steps=30,
     phases=(LTXPhaseRecipe(name="generate", guidance=_official_guidance(28)),),
 )
+LTX25_FULL_RECIPE = LTXPipelineRecipe(
+    height=544,
+    width=960,
+    num_inference_steps=30,
+    negative_prompt=LTX25_DEFAULT_NEGATIVE_PROMPT,
+    phases=(
+        LTXPhaseRecipe(
+            name="generate",
+            guidance=_official_guidance(28),
+            noise_scale=1.0,
+        ),
+    ),
+)
 LTX_POSITIVE_ONLY_RECIPE = LTXPipelineRecipe(
+    supports_cache_dit=True,
     phases=(
         LTXPhaseRecipe(
             name="generate",
@@ -159,6 +185,41 @@ LTX2_DISTILLED_TWO_STAGE_RECIPE = LTXPipelineRecipe(
     fixed_num_inference_steps=True,
 )
 
+LTX25_DISTILLED_TWO_STAGE_RECIPE = LTXPipelineRecipe(
+    height=1088,
+    width=1920,
+    num_inference_steps=len(LTX_DISTILLED_SIGMAS) - 1,
+    negative_prompt="",
+    phases=(
+        LTXPhaseRecipe(
+            name="generate_lowres",
+            guidance=LTXGuidanceSpec.positive_only(),
+            spatial_downscale=2,
+            sigmas=LTX_DISTILLED_SIGMAS,
+            noise_scale=1.0,
+            allow_guidance_override=False,
+            use_official_sigma_schedule=False,
+            sampler="euler_ancestral",
+        ),
+        LTXPhaseRecipe(
+            name="refine",
+            guidance=LTXGuidanceSpec.positive_only(),
+            sigmas=LTX_STAGE_2_DISTILLED_SIGMAS,
+            noise_scale=LTX_STAGE_2_DISTILLED_SIGMAS[0],
+            input_transform="spatial_upsample",
+            allow_guidance_override=False,
+            use_official_sigma_schedule=False,
+        ),
+    ),
+    video_output_phase=1,
+    audio_output_phase=1,
+    allow_request_sigmas=False,
+    allow_request_phase_sigmas=True,
+    allow_request_latents=False,
+    allow_negative_prompt=False,
+    fixed_num_inference_steps=True,
+)
+
 # LTX-2.3 full-distilled uses the same official fixed 8 + 3 sigma schedules
 # and request contract. Its distinct component profile selects the 22B merged
 # distilled Transformer, BWE vocoder, and matching spatial upsampler; no
@@ -166,6 +227,37 @@ LTX2_DISTILLED_TWO_STAGE_RECIPE = LTXPipelineRecipe(
 LTX23_DISTILLED_TWO_STAGE_RECIPE = LTX2_DISTILLED_TWO_STAGE_RECIPE
 
 
+def _distilled_one_stage_recipe(
+    two_stage_recipe: LTXPipelineRecipe,
+    *,
+    supports_cache_dit: bool = True,
+) -> LTXPipelineRecipe:
+    """Run the merged-distilled checkpoint at its native Stage-1 resolution."""
+    generate_phase = replace(
+        two_stage_recipe.phases[0],
+        name="generate",
+        spatial_downscale=1,
+    )
+    return replace(
+        two_stage_recipe,
+        height=two_stage_recipe.height // 2,
+        width=two_stage_recipe.width // 2,
+        phases=(generate_phase,),
+        video_output_phase=0,
+        audio_output_phase=0,
+        allow_request_phase_sigmas=False,
+        supports_cache_dit=supports_cache_dit,
+    )
+
+
+LTX2_DISTILLED_ONE_STAGE_RECIPE = _distilled_one_stage_recipe(LTX2_DISTILLED_TWO_STAGE_RECIPE)
+LTX23_DISTILLED_ONE_STAGE_RECIPE = _distilled_one_stage_recipe(LTX23_DISTILLED_TWO_STAGE_RECIPE)
+LTX25_DISTILLED_ONE_STAGE_RECIPE = _distilled_one_stage_recipe(
+    LTX25_DISTILLED_TWO_STAGE_RECIPE,
+    supports_cache_dit=False,
+)
+
+
 def _official_two_stage_recipe(one_stage_recipe: LTXPipelineRecipe) -> LTXPipelineRecipe:
     return LTXPipelineRecipe(
         height=one_stage_recipe.height * 2,
@@ -197,21 +289,29 @@ def _official_two_stage_recipe(one_stage_recipe: LTXPipelineRecipe) -> LTXPipeli
         # discards its audio result. Decode the full-context Stage-1 audio.
         audio_output_phase=0,
         allow_request_sigmas=False,
+        allow_request_phase_sigmas=True,
         allow_request_latents=False,
     )
 
 
 LTX2_TWO_STAGE_RECIPE = _official_two_stage_recipe(LTX2_ONE_STAGE_RECIPE)
 LTX23_TWO_STAGE_RECIPE = _official_two_stage_recipe(LTX23_ONE_STAGE_RECIPE)
+LTX25_TWO_STAGE_RECIPE = _official_two_stage_recipe(LTX25_FULL_RECIPE)
 
 
 _PIPELINE_RECIPES: dict[tuple[str, str], LTXPipelineRecipe] = {
     ("one_stage", "2"): LTX2_ONE_STAGE_RECIPE,
     ("one_stage", "2.3"): LTX23_ONE_STAGE_RECIPE,
+    ("one_stage", "2.5"): LTX25_FULL_RECIPE,
     ("two_stage", "2"): LTX2_TWO_STAGE_RECIPE,
     ("two_stage", "2.3"): LTX23_TWO_STAGE_RECIPE,
+    ("two_stage", "2.5"): LTX25_TWO_STAGE_RECIPE,
+    ("distilled_one_stage", "2"): LTX2_DISTILLED_ONE_STAGE_RECIPE,
+    ("distilled_one_stage", "2.3"): LTX23_DISTILLED_ONE_STAGE_RECIPE,
+    ("distilled_one_stage", "2.5"): LTX25_DISTILLED_ONE_STAGE_RECIPE,
     ("distilled_two_stage", "2"): LTX2_DISTILLED_TWO_STAGE_RECIPE,
     ("distilled_two_stage", "2.3"): LTX23_DISTILLED_TWO_STAGE_RECIPE,
+    ("distilled_two_stage", "2.5"): LTX25_DISTILLED_TWO_STAGE_RECIPE,
     ("dmd2", "2"): LTX_POSITIVE_ONLY_RECIPE,
     ("dmd2", "2.3"): LTX_POSITIVE_ONLY_RECIPE,
 }

vllm_omni/diffusion/models/ltx2/ltx2_request.py

@@ -43,6 +43,7 @@ class LTXRequestInputs:
     output_type: str
     max_sequence_length: int
     audio_latents_normalized: bool = False
+    image_crf: int | None = None
 
     @property
     def guidance_scale(self) -> float:
@@ -86,6 +87,7 @@ def validate_pipeline_request(
     vae_temporal_compression_ratio: int,
     pipeline_name: str,
     request_sigmas: list[float] | None = None,
+    request_phase_sigmas: tuple[list[float] | None, ...] | None = None,
 ) -> None:
     """Validate the request capabilities declared by an execution recipe."""
     alignment = vae_spatial_compression_ratio * pipeline_recipe.max_spatial_downscale
@@ -108,9 +110,28 @@ def validate_pipeline_request(
 
     if not pipeline_recipe.allow_request_sigmas and request_sigmas is not None:
         raise ValueError(f"{pipeline_name} uses fixed phase sigma schedules.")
+    if request_sigmas is not None and len(request_sigmas) < 2:
+        raise ValueError("A custom LTX sigma schedule must contain at least two values.")
 
-    if pipeline_recipe.fixed_num_inference_steps and (
-        request_inputs.num_inference_steps != pipeline_recipe.num_inference_steps
+    has_phase_sigmas = request_phase_sigmas is not None and any(item is not None for item in request_phase_sigmas)
+    if request_sigmas is not None and has_phase_sigmas:
+        raise ValueError("Use either `sigmas` or per-stage sigma schedules, not both.")
+    if has_phase_sigmas and not pipeline_recipe.allow_request_phase_sigmas:
+        raise ValueError(f"{pipeline_name} does not accept per-stage sigma schedules.")
+    if request_phase_sigmas is not None and len(request_phase_sigmas) != len(pipeline_recipe.phases):
+        raise ValueError(
+            f"{pipeline_name} has {len(pipeline_recipe.phases)} phases but received "
+            f"{len(request_phase_sigmas)} per-stage sigma schedules."
+        )
+    for index, phase_sigmas in enumerate(request_phase_sigmas or ()):
+        if phase_sigmas is not None and len(phase_sigmas) < 2:
+            raise ValueError(f"`stage_{index + 1}_sigmas` must contain at least two values.")
+
+    if (
+        request_sigmas is None
+        and (request_phase_sigmas is None or request_phase_sigmas[0] is None)
+        and pipeline_recipe.fixed_num_inference_steps
+        and (request_inputs.num_inference_steps != pipeline_recipe.num_inference_steps)
     ):
         raise ValueError(
             f"{pipeline_name} uses {pipeline_recipe.num_inference_steps} fixed Stage 1 denoise steps; "
@@ -165,6 +186,16 @@ def _get_extra_arg(sampling: Any, *names: str) -> Any:
     return None
 
 
+def _normalize_image_crf(value: Any) -> int | None:
+    if value is None:
+        return None
+    if isinstance(value, bool) or not isinstance(value, int):
+        raise ValueError("`image_crf` must be an integer from 0 through 51.")
+    if not 0 <= value <= 51:
+        raise ValueError("`image_crf` must be an integer from 0 through 51.")
+    return value
+
+
 def _normalize_stg_blocks(value: Any, default: tuple[int, ...]) -> tuple[int, ...]:
     if value is None:
         return default
@@ -274,6 +305,7 @@ class LTXRequestMixin:
         req: DiffusionRequestBatch,
         *,
         image: Any | None = None,
+        image_crf: int | None = None,
         prompt: str | list[str] | None = None,
         negative_prompt: str | list[str] | None = None,
         height: int | None = None,
@@ -282,6 +314,8 @@ class LTXRequestMixin:
         frame_rate: float | None = None,
         num_inference_steps: int | None = None,
         sigmas: list[float] | None = None,
+        stage_1_sigmas: list[float] | None = None,
+        stage_2_sigmas: list[float] | None = None,
         timesteps: list[int] | None = None,
         guidance_scale: float | None = None,
         guidance_rescale: float | None = None,
@@ -315,6 +349,7 @@ class LTXRequestMixin:
         return self._forward_request(
             req,
             image=image,
+            image_crf=image_crf,
             prompt=prompt,
             negative_prompt=negative_prompt,
             height=height,
@@ -323,6 +358,8 @@ class LTXRequestMixin:
             frame_rate=frame_rate,
             num_inference_steps=num_inference_steps,
             sigmas=sigmas,
+            stage_1_sigmas=stage_1_sigmas,
+            stage_2_sigmas=stage_2_sigmas,
             guidance_scale=guidance_scale,
             num_videos_per_prompt=num_videos_per_prompt,
             generator=generator,
@@ -343,12 +380,30 @@ class LTXRequestMixin:
         req: DiffusionRequestBatch,
         fallback: list[float] | None,
     ) -> list[float] | None:
-        request_sigmas = [sampling.sigmas for sampling in req.sampling_params_list]
+        request_sigmas = [
+            sampling.sigmas if sampling.sigmas is not None else _get_extra_arg(sampling, "sigmas")
+            for sampling in req.sampling_params_list
+        ]
         first = request_sigmas[0] if request_sigmas else None
         if any(item != first for item in request_sigmas[1:]):
             raise ValueError("Batched LTX requests must use identical custom sigmas.")
         return first if first is not None else fallback
 
+    @staticmethod
+    def _resolve_request_phase_sigmas(
+        req: DiffusionRequestBatch,
+        stage_1_fallback: list[float] | None,
+        stage_2_fallback: list[float] | None,
+    ) -> tuple[list[float] | None, list[float] | None] | None:
+        resolved: list[list[float] | None] = []
+        for name, fallback in (("stage_1_sigmas", stage_1_fallback), ("stage_2_sigmas", stage_2_fallback)):
+            values = [_get_extra_arg(sampling, name) for sampling in req.sampling_params_list]
+            first = values[0] if values else None
+            if any(item != first for item in values[1:]):
+                raise ValueError(f"Batched LTX requests must use identical `{name}` values.")
+            resolved.append(first if first is not None else fallback)
+        return tuple(resolved) if any(item is not None for item in resolved) else None
+
     def check_inputs(
         self,
         prompt: str | list[str] | None,
@@ -408,6 +463,7 @@ class LTXRequestMixin:
         decode_noise_scale: float | list[float] | None,
         output_type: str,
         max_sequence_length: int | None,
+        image_crf: int | None = None,
     ) -> LTXRequestInputs:
         sampling_params_list = req.sampling_params_list
         for sampling in sampling_params_list:
@@ -421,6 +477,12 @@ class LTXRequestMixin:
             if _get_extra_arg(sampling, "flow_shift") is not None:
                 raise ValueError("LTX does not support `flow_shift`; use `sigmas` for a custom schedule.")
 
+        request_image_crfs = [_get_extra_arg(item, "image_crf") for item in sampling_params_list]
+        first_image_crf = request_image_crfs[0] if request_image_crfs else None
+        if any(item != first_image_crf for item in request_image_crfs[1:]):
+            raise ValueError("Batched LTX requests must use identical `image_crf` values.")
+        image_crf = _normalize_image_crf(first_image_crf if first_image_crf is not None else image_crf)
+
         sampling = sampling_params_list[0]
         is_dummy_run = req.is_dummy_run()
         prompt = [item if isinstance(item, str) else (item.get("prompt") or "") for item in req.prompts] or prompt
@@ -551,4 +613,5 @@ class LTXRequestMixin:
             decode_noise_scale=decode_noise_scale,
             output_type=output_type,
             max_sequence_length=int(max_sequence_length),
+            image_crf=image_crf,
         )

vllm_omni/diffusion/models/ltx2/ltx2_runtime.py

@@ -144,20 +144,19 @@ class LTXRuntime(
                 f"{self.__class__.__name__} does not support ulysses_mode='advanced_uaa'. "
                 "Use the default ulysses_mode='strict' for LTX sequence parallelism."
             )
-        self.model_version = detect_ltx_model_version(od_config.model)
+        self.model_version = detect_ltx_model_version(od_config.model, revision=getattr(od_config, "revision", None))
         self.component_profile = resolve_ltx_component_profile(self.pipeline_kind, self.model_version)
         self.pipeline_recipe = resolve_ltx_pipeline_recipe(self.pipeline_kind, self.model_version)
-        if getattr(od_config, "cache_backend", "none") == "cache_dit" and len(self.pipeline_recipe.phases) > 1:
+        if getattr(od_config, "cache_backend", "none") == "cache_dit" and not self.pipeline_recipe.supports_cache_dit:
             raise ValueError(
                 f"{self.__class__.__name__} does not support cache_backend='cache_dit'. "
-                "Cache-DiT currently supports only one-stage LTX recipes; multi-stage recipes require "
-                "phase-aware cache enable and refresh."
+                "Cache-DiT is not qualified for this LTX recipe."
             )
         self._dit_modules = list(self.component_profile.dit_modules)
         self._encoder_modules = list(self.component_profile.encoder_modules)
         self._vae_modules = list(self.component_profile.vae_modules)
         self._resident_modules = list(self.component_profile.resident_modules)
-        if self.model_version == "2.3":
+        if self.model_version in ("2.3", "2.5"):
             self.preserve_sp_padded_audio_duration = True
             self.reports_stage_durations = True
         super().__init__()
@@ -173,6 +172,7 @@ class LTXRuntime(
         req: DiffusionRequestBatch,
         *,
         image: Any | None = None,
+        image_crf: int | None = None,
         prompt: str | list[str] | None = None,
         negative_prompt: str | list[str] | None = None,
         height: int | None = None,
@@ -181,6 +181,8 @@ class LTXRuntime(
         frame_rate: float | None = None,
         num_inference_steps: int | None = None,
         sigmas: list[float] | None = None,
+        stage_1_sigmas: list[float] | None = None,
+        stage_2_sigmas: list[float] | None = None,
         guidance_scale: float | None = None,
         num_videos_per_prompt: int | None = 1,
         generator: torch.Generator | list[torch.Generator] | None = None,
@@ -217,11 +219,13 @@ class LTXRuntime(
             decode_noise_scale=decode_noise_scale,
             output_type=output_type,
             max_sequence_length=max_sequence_length,
+            image_crf=image_crf,
         )
         image = self._resolve_request_image(req, image, request_inputs)
         if image is not None and not self.support_image_input:
             raise ValueError(f"{self.__class__.__name__} does not support `image` input.")
         request_sigmas = self._resolve_request_sigmas(req, sigmas)
+        request_phase_sigmas = self._resolve_request_phase_sigmas(req, stage_1_sigmas, stage_2_sigmas)
         validate_pipeline_request(
             request_inputs,
             pipeline_recipe=self.pipeline_recipe,
@@ -229,6 +233,7 @@ class LTXRuntime(
             vae_temporal_compression_ratio=self.vae_temporal_compression_ratio,
             pipeline_name=self.__class__.__name__,
             request_sigmas=request_sigmas,
+            request_phase_sigmas=request_phase_sigmas,
         )
         phase_adapter = getattr(self, "_phase_adapter", None)
         if phase_adapter is not None and any(
@@ -237,7 +242,13 @@ class LTXRuntime(
             raise ValueError(
                 f"{self.__class__.__name__} cannot compose a request LoRA with its internal phase adapter."
             )
-        return self._run_recipe(req, request_inputs, request_sigmas=request_sigmas, image=image)
+        return self._run_recipe(
+            req,
+            request_inputs,
+            request_sigmas=request_sigmas,
+            request_phase_sigmas=request_phase_sigmas,
+            image=image,
+        )
 
     def _run_recipe(
         self,
@@ -245,24 +256,39 @@ class LTXRuntime(
         request_inputs: LTXRequestInputs,
         *,
         request_sigmas: list[float] | None,
+        request_phase_sigmas: tuple[list[float] | None, ...] | None = None,
         image: Any | None = None,
     ) -> DiffusionOutput | list[DiffusionOutput]:
         """Execute one- and multi-phase recipes through the same control flow."""
         phase_results: list[LTXPhaseResult] = []
         prompt_context = None
-
-        for phase_recipe in self.pipeline_recipe.phases:
+        for phase_index, phase_recipe in enumerate(self.pipeline_recipe.phases):
+            override_sigmas = None if request_phase_sigmas is None else request_phase_sigmas[phase_index]
+            phase_sigmas = (
+                override_sigmas
+                if override_sigmas is not None
+                else (
+                    request_sigmas
+                    if request_sigmas is not None
+                    else (list(phase_recipe.sigmas) if phase_recipe.sigmas is not None else None)
+                )
+            )
             self._enter_phase(phase_recipe)
             phase_inputs = self._build_phase_inputs(
                 request_inputs,
                 phase_recipe,
                 phase_results[-1] if phase_results else None,
             )
+            if phase_sigmas is not None:
+                phase_inputs = replace(phase_inputs, num_inference_steps=len(phase_sigmas) - 1)
+            noise_scale = phase_recipe.noise_scale
+            if override_sigmas is not None and phase_recipe.input_transform == "spatial_upsample":
+                noise_scale = float(override_sigmas[0])
             phase_result = self.run_phase(
                 req,
                 phase_inputs,
-                noise_scale=phase_recipe.noise_scale,
-                sigmas=list(phase_recipe.sigmas) if phase_recipe.sigmas is not None else request_sigmas,
+                noise_scale=noise_scale,
+                sigmas=phase_sigmas,
                 timesteps=None,
                 attention_kwargs=None,
                 phase_recipe=phase_recipe,
@@ -282,6 +308,7 @@ class LTXRuntime(
 
     def _enter_phase(self, phase: LTXPhaseRecipe) -> None:
         self._active_phase_name = phase.name
+        self._guidance_plan = LTXGuidancePlan.build(phase.guidance)
         phase_adapter = getattr(self, "_phase_adapter", None)
         if phase_adapter is None:
             if phase.adapter_slot is not None:
@@ -645,7 +672,12 @@ class LTXRuntime(
         audio_token_count: int,
     ) -> dict[str, torch.Tensor]:
         del forward_ctx, denoise_ctx
-        return self.guidance_executor.timestep_kwargs(ts, video_token_count, audio_token_count)
+        return self.guidance_executor.timestep_kwargs(
+            ts,
+            video_token_count,
+            audio_token_count,
+            expand_for_sequence_parallel=True,
+        )
 
     def _video_guidance_model_sigma(
         self,
@@ -700,6 +732,7 @@ class LTXRuntime(
             state,
             forward_ctx,
             denoise_ctx,
+            preserve_positive_velocity=forward_ctx.sampler == "euler_ancestral",
         )
 
     def _denoise_step(

vllm_omni/diffusion/models/ltx2/ltx2_transformer.py

@@ -55,6 +55,18 @@ logger = init_logger(__name__)
 _RMSNORM_INIT_PARAMS = inspect.signature(RMSNorm.__init__).parameters
 
 
+def apply_keyframes_absolute_embedding(
+    hidden_states: torch.Tensor,
+    keyframes_mask: torch.Tensor | None,
+    embedding: torch.Tensor | None,
+) -> torch.Tensor:
+    """Add the official learned marker to first-frame/keyframe tokens."""
+    if keyframes_mask is None or embedding is None:
+        return hidden_states
+    mask = (keyframes_mask > 0).to(dtype=hidden_states.dtype)
+    return hidden_states + mask * embedding.to(device=hidden_states.device, dtype=hidden_states.dtype)
+
+
 def _make_rms_norm(hidden_size: int, *, eps: float, elementwise_affine: bool) -> nn.Module:
     """Bridge diffusers' RMSNorm API onto vLLM's `has_weight` variant."""
     if not elementwise_affine:
@@ -82,15 +94,10 @@ def _make_rms_norm(hidden_size: int, *, eps: float, elementwise_affine: bool) ->
 
 def apply_interleaved_rotary_emb(x: torch.Tensor, freqs: tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor:
     cos, sin = freqs
-    if cos.dtype != x.dtype:
-        cos = cos.to(dtype=x.dtype)
-    if sin.dtype != x.dtype:
-        sin = sin.to(dtype=x.dtype)
     # Concrete pair count instead of -1 keeps SDPA shape static under torch.compile.
     x_real, x_imag = x.unflatten(2, (x.shape[2] // 2, 2)).unbind(-1)  # [B, S, C // 2]
     x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(2)
-    out = x * cos + x_rotated * sin
-    return out
+    return x * cos + x_rotated * sin
 
 
 def apply_split_rotary_emb(
@@ -101,12 +108,6 @@ def apply_split_rotary_emb(
 ) -> torch.Tensor:
     # `head_dim` is plumbed in (not inferred via `-1`) so SDPA shape stays static under torch.compile.
     cos, sin = freqs
-    if cos.dtype != x.dtype:
-        cos = cos.to(dtype=x.dtype)
-    if sin.dtype != x.dtype:
-        sin = sin.to(dtype=x.dtype)
-
-    x_dtype = x.dtype
     needs_reshape = False
     if x.ndim != 4 and cos.ndim == 4:
         # cos is (#b, h, t, r) -> reshape x to (b, h, t, dim_per_head)
@@ -142,7 +143,6 @@ def apply_split_rotary_emb(
     if needs_reshape:
         out = out.swapaxes(1, 2).reshape(b, t, h * head_dim)
 
-    out = out.to(dtype=x_dtype)
     return out
 
 
@@ -269,6 +269,7 @@ class LTX2FeedForward(nn.Module):
             RowParallelLinear(
                 inner_dim,
                 dim_out,
+                bias=bias,
                 input_is_parallel=True,
                 return_bias=False,
                 quant_config=quant_config,
@@ -517,8 +518,8 @@ class LTX2AudioVideoAttnProcessor:
             is_self_attention=is_self_attention,
         )
 
-        query = attn.norm_q(query)
-        key = attn.norm_k(key)
+        query = attn.norm_q(query).to(dtype=value.dtype)
+        key = attn.norm_k(key).to(dtype=value.dtype)
 
         if query_rotary_emb is not None:
             query_rotary_emb = self._slice_rope_for_tp(query_rotary_emb, attn)
@@ -537,6 +538,10 @@ class LTX2AudioVideoAttnProcessor:
                     head_dim=attn.head_dim,
                 )
 
+        # Keep FP32 RoPE computation while satisfying attention dtype contracts.
+        query = query.to(dtype=value.dtype)
+        key = key.to(dtype=value.dtype)
+
         query = query.unflatten(2, (attn.heads, attn.head_dim))
         key = key.unflatten(2, (attn.heads, attn.head_dim))
         value_heads = value.unflatten(2, (attn.heads, attn.head_dim))
@@ -846,6 +851,8 @@ class LTX2VideoTransformerBlock(nn.Module):
         elementwise_affine: bool = False,
         rope_type: str = "interleaved",
         perturbed_attn: bool = False,
+        ff_bias: bool = True,
+        audio_ff_bias: bool = True,
         quant_config: "QuantizationConfig | None" = None,
         prefix: str = "",
     ):
@@ -965,6 +972,7 @@ class LTX2VideoTransformerBlock(nn.Module):
         self.ff = LTX2FeedForward(
             dim,
             activation_fn=activation_fn,
+            bias=ff_bias,
             quant_config=quant_config,
             prefix=f"{prefix}.ff" if prefix else "ff",
         )
@@ -973,6 +981,7 @@ class LTX2VideoTransformerBlock(nn.Module):
         self.audio_ff = LTX2FeedForward(
             audio_dim,
             activation_fn=activation_fn,
+            bias=audio_ff_bias,
             quant_config=quant_config,
             prefix=f"{prefix}.audio_ff" if prefix else "audio_ff",
         )
@@ -1078,21 +1087,34 @@ class LTX2VideoTransformerBlock(nn.Module):
         audio_hidden_states = audio_hidden_states + attn_audio_hidden_states * audio_gate_msa
 
         # 2. Video and Audio Cross-Attention with text embeddings (Q: Video/Audio; K,V: Text)
-        # LTX-2.3: compute prompt modulation params for K/V
-        if self.cross_attn_adaln and temb_prompt is not None:
-            video_prompt_ada_params = self.get_mod_params(self.prompt_scale_shift_table, temb_prompt, batch_size)
-            shift_text_kv, scale_text_kv = video_prompt_ada_params
+        if self.cross_attn_adaln:
+            if temb_prompt is not None:
+                shift_text_kv, scale_text_kv = self.get_mod_params(
+                    self.prompt_scale_shift_table, temb_prompt, batch_size
+                )
+            else:
+                shift_text_kv, scale_text_kv = (
+                    self.prompt_scale_shift_table[None, None]
+                    .to(device=hidden_states.device, dtype=hidden_states.dtype)
+                    .unbind(dim=2)
+                )
 
-            audio_prompt_ada_params = self.get_mod_params(
-                self.audio_prompt_scale_shift_table, temb_prompt_audio, batch_size
-            )
-            audio_shift_text_kv, audio_scale_text_kv = audio_prompt_ada_params
+            if temb_prompt_audio is not None:
+                audio_shift_text_kv, audio_scale_text_kv = self.get_mod_params(
+                    self.audio_prompt_scale_shift_table, temb_prompt_audio, batch_size
+                )
+            else:
+                audio_shift_text_kv, audio_scale_text_kv = (
+                    self.audio_prompt_scale_shift_table[None, None]
+                    .to(device=audio_hidden_states.device, dtype=audio_hidden_states.dtype)
+                    .unbind(dim=2)
+                )
 
         # 2.1. Video-Text Cross-Attention
         norm_hidden_states = self.norm2(hidden_states)
         if self.video_cross_attn_adaln:
             norm_hidden_states = norm_hidden_states * (1 + scale_text_q) + shift_text_q
-        if self.cross_attn_adaln and temb_prompt is not None:
+        if self.cross_attn_adaln:
             encoder_hidden_states = encoder_hidden_states * (1 + scale_text_kv) + shift_text_kv
 
         attn_hidden_states = self.attn2(
@@ -1109,7 +1131,7 @@ class LTX2VideoTransformerBlock(nn.Module):
         norm_audio_hidden_states = self.audio_norm2(audio_hidden_states)
         if self.audio_cross_attn_adaln:
             norm_audio_hidden_states = norm_audio_hidden_states * (1 + audio_scale_text_q) + audio_shift_text_q
-        if self.cross_attn_adaln and temb_prompt is not None:
+        if self.cross_attn_adaln:
             audio_encoder_hidden_states = audio_encoder_hidden_states * (1 + audio_scale_text_kv) + audio_shift_text_kv
 
         attn_audio_hidden_states = self.audio_attn2(
@@ -1508,7 +1530,10 @@ class LTX2VideoTransformer3DModel(nn.Module):
         block_forward_patterns={
             "transformer_blocks": ForwardPattern.Pattern_0,
         },
-        has_separate_cfg=True,
+        # LTX fuses all guidance passes into the batch dimension and invokes
+        # the transformer once per denoise step. Cache-DiT must therefore
+        # advance its step counter for every transformer forward.
+        has_separate_cfg=False,
         check_forward_pattern=False,
     )
 
@@ -1546,6 +1571,7 @@ class LTX2VideoTransformer3DModel(nn.Module):
                 # Shard video/audio latents across sequence
                 "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, split_output=False),
                 "audio_hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, split_output=False),
+                "keyframes_mask": SequenceParallelInput(split_dim=1, expected_dims=3, split_output=False),
                 # Shard prompt embeds across sequence
                 "encoder_hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, split_output=False),
                 "audio_encoder_hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, split_output=False),
@@ -1620,6 +1646,10 @@ class LTX2VideoTransformer3DModel(nn.Module):
         cross_attn_mod: bool = False,
         audio_gated_attn: bool = False,
         audio_cross_attn_mod: bool = False,
+        ff_bias: bool = True,
+        audio_ff_bias: bool = True,
+        use_prompt_adaln_single: bool = True,
+        use_keyframes_abs_pos_embedding: bool = False,
         quant_config: "QuantizationConfig | None" = None,
     ) -> None:
         super().__init__()
@@ -1666,12 +1696,19 @@ class LTX2VideoTransformer3DModel(nn.Module):
             timestep_scale_multiplier=timestep_scale_multiplier,
             cross_attn_timestep_scale_multiplier=cross_attn_timestep_scale_multiplier,
             rope_type=rope_type,
+            ff_bias=ff_bias,
+            audio_ff_bias=audio_ff_bias,
+            use_prompt_adaln_single=use_prompt_adaln_single,
+            use_keyframes_abs_pos_embedding=use_keyframes_abs_pos_embedding,
         )
 
         # 1. Patchification input projections
         self.proj_in = nn.Linear(in_channels, inner_dim)
         self.audio_proj_in = nn.Linear(audio_in_channels, audio_inner_dim)
 
+        if use_keyframes_abs_pos_embedding:
+            self.keyframes_abs_pos_embedding = nn.Parameter(torch.zeros(1, inner_dim))
+
         # 2. Prompt embeddings
         # LTX-2 (use_prompt_embeddings=True): caption projection in the transformer
         # LTX-2.3 (use_prompt_embeddings=False): caption projection in the connectors
@@ -1695,7 +1732,7 @@ class LTX2VideoTransformer3DModel(nn.Module):
 
         # 3.3. LTX-2.3: Prompt modulation from sigma
         self.prompt_modulation = cross_attn_mod or audio_cross_attn_mod
-        if self.prompt_modulation:
+        if self.prompt_modulation and use_prompt_adaln_single:
             self.prompt_adaln = LTX2AdaLayerNormSingle(inner_dim, num_mod_params=2, use_additional_conditions=False)
             self.audio_prompt_adaln = LTX2AdaLayerNormSingle(
                 audio_inner_dim, num_mod_params=2, use_additional_conditions=False
@@ -1815,6 +1852,8 @@ class LTX2VideoTransformer3DModel(nn.Module):
                     elementwise_affine=norm_elementwise_affine,
                     rope_type=rope_type,
                     perturbed_attn=perturbed_attn,
+                    ff_bias=ff_bias,
+                    audio_ff_bias=audio_ff_bias,
                     quant_config=quant_config,
                     prefix=f"transformer_blocks.{layer_idx}",
                 )
@@ -1852,6 +1891,7 @@ class LTX2VideoTransformer3DModel(nn.Module):
         audio_encoder_hidden_states: torch.Tensor,
         timestep: torch.LongTensor,
         audio_timestep: torch.LongTensor | None = None,
+        keyframes_mask: torch.Tensor | None = None,
         sigma: torch.Tensor | None = None,
         audio_sigma: torch.Tensor | None = None,
         encoder_attention_mask: torch.Tensor | None = None,
@@ -1952,9 +1992,7 @@ class LTX2VideoTransformer3DModel(nn.Module):
 
         video_rotary_emb = self.rope(video_coords, device=hidden_states.device, out_dtype=hidden_states.dtype)
         audio_rotary_emb = self.audio_rope(
-            audio_coords,
-            device=audio_hidden_states.device,
-            out_dtype=audio_hidden_states.dtype,
+            audio_coords, device=audio_hidden_states.device, out_dtype=audio_hidden_states.dtype
         )
 
         video_cross_attn_rotary_emb = self.cross_attn_rope(
@@ -1970,6 +2008,11 @@ class LTX2VideoTransformer3DModel(nn.Module):
 
         # 2. Patchify input projections
         hidden_states = self.proj_in(hidden_states)
+        hidden_states = apply_keyframes_absolute_embedding(
+            hidden_states,
+            keyframes_mask,
+            getattr(self, "keyframes_abs_pos_embedding", None),
+        )
         audio_hidden_states = self.audio_proj_in(audio_hidden_states)
 
         # 3. Prepare timestep embeddings and modulation parameters
@@ -2028,7 +2071,7 @@ class LTX2VideoTransformer3DModel(nn.Module):
         audio_cross_attn_v2a_gate = audio_cross_attn_v2a_gate.view(batch_size, -1, audio_cross_attn_v2a_gate.shape[-1])
 
         # 3.3. LTX-2.3: Compute prompt modulation from sigma
-        if self.prompt_modulation and sigma is not None:
+        if self.prompt_modulation and self.config.use_prompt_adaln_single and sigma is not None:
             temb_prompt, _ = self.prompt_adaln(sigma.flatten(), batch_size=batch_size, hidden_dtype=hidden_states.dtype)
             temb_prompt_audio, _ = self.audio_prompt_adaln(
                 audio_sigma.flatten(), batch_size=batch_size, hidden_dtype=audio_hidden_states.dtype

vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py

@@ -10,15 +10,12 @@ from typing import ClassVar
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin
 
-from .ltx2_components import LTX2_COMPONENT_PROFILE
+from .ltx2_components import LTX2_COMPONENT_PROFILE, LTX2_DISTILLED_ONE_STAGE_COMPONENT_PROFILE
 from .ltx2_components import (
     get_ltx2_post_process_func as get_ltx2_post_process_func,  # noqa: F401
 )
 from .ltx2_conditioning import LTXI2VConditioningMixin
-from .ltx2_recipes import (
-    LTX2_ONE_STAGE_RECIPE,
-    LTX_POSITIVE_ONLY_RECIPE,
-)
+from .ltx2_recipes import LTX2_DISTILLED_ONE_STAGE_RECIPE, LTX2_ONE_STAGE_RECIPE, LTX_POSITIVE_ONLY_RECIPE
 from .ltx2_runtime import LTXRuntime
 
 
@@ -36,6 +33,14 @@ class LTX2Pipeline(LTXI2VConditioningMixin, LTXRuntime):
     supports_request_batch = True
 
 
+class LTX2DistilledOneStagePipeline(LTX2Pipeline):
+    """Merged-distilled checkpoint at its native one-stage resolution."""
+
+    pipeline_kind = "distilled_one_stage"
+    component_profile = LTX2_DISTILLED_ONE_STAGE_COMPONENT_PROFILE
+    pipeline_recipe = LTX2_DISTILLED_ONE_STAGE_RECIPE
+
+
 class LTX2T2VDMD2Pipeline(DMD2PipelineMixin, LTX2Pipeline):
     """LTX2 T2V entry for FastGen DMD2-distilled models."""
 

vllm_omni/diffusion/models/ltx2/pipeline_ltx2_two_stage.py

@@ -40,9 +40,13 @@ class LTX2TwoStagePipeline(_LTX2TwoStageBase):
     pipeline_recipe = LTX2_TWO_STAGE_RECIPE
 
 
-class LTX2DistilledPipeline(_LTX2TwoStageBase):
-    """Unified LTX-2/LTX-2.3 full-distilled two-stage T2V/I2V entry."""
+class LTX2DistilledTwoStagePipeline(_LTX2TwoStageBase):
+    """Merged-distilled checkpoint with latent upsampling and refinement."""
 
     pipeline_kind = "distilled_two_stage"
     component_profile = LTX2_DISTILLED_COMPONENT_PROFILE
     pipeline_recipe = LTX2_DISTILLED_TWO_STAGE_RECIPE
+
+
+# Backwards-compatible name retained for existing LTX-2/LTX-2.3 recipes.
+LTX2DistilledPipeline = LTX2DistilledTwoStagePipeline

vllm_omni/diffusion/registry.py

@@ -76,6 +76,16 @@ _DIFFUSION_MODELS = {
         "pipeline_ltx2_two_stage",
         "LTX2TwoStagePipeline",
     ),
+    "LTX2DistilledOneStagePipeline": (
+        "ltx2",
+        "pipeline_ltx2",
+        "LTX2DistilledOneStagePipeline",
+    ),
+    "LTX2DistilledTwoStagePipeline": (
+        "ltx2",
+        "pipeline_ltx2_two_stage",
+        "LTX2DistilledTwoStagePipeline",
+    ),
     "LTX2DistilledPipeline": (
         "ltx2",
         "pipeline_ltx2_two_stage",
@@ -522,6 +532,8 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
     "WanVACEPipeline": "get_wan22_vace_post_process_func",
     "LTX2Pipeline": "get_ltx2_post_process_func",
     "LTX2TwoStagePipeline": "get_ltx2_post_process_func",
+    "LTX2DistilledOneStagePipeline": "get_ltx2_post_process_func",
+    "LTX2DistilledTwoStagePipeline": "get_ltx2_post_process_func",
     "LTX2DistilledPipeline": "get_ltx2_post_process_func",
     "LTX2T2VDMD2Pipeline": "get_ltx2_post_process_func",
     "LTX2I2VDMD2Pipeline": "get_ltx2_post_process_func",

vllm_omni/entrypoints/openai/serving_video.py

@@ -7,6 +7,7 @@ import copy
 import time
 from collections.abc import Mapping
 from dataclasses import dataclass
+from functools import cached_property
 from http import HTTPStatus
 from typing import Any, cast
 
@@ -30,6 +31,7 @@ from vllm_omni.entrypoints.openai.stage_params import (
 from vllm_omni.entrypoints.openai.utils import get_stage_type, parse_lora_request
 from vllm_omni.entrypoints.openai.video_api_utils import _encode_video_bytes, encode_video_base64
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt
+from vllm_omni.model_extras import should_preserve_reference_image_size
 from vllm_omni.outputs.output_metadata import (
     DiffusionMetadataMapping,
     DiffusionMultimodalOutput,
@@ -100,6 +102,21 @@ class OmniOpenAIServingVideo:
         if self._stage_configs is None and stage_configs is not None:
             self._stage_configs = stage_configs
 
+    @cached_property
+    def preserves_reference_image_size(self) -> bool:
+        """Return whether the active pipeline owns reference-image resizing."""
+        get_od_config = getattr(self._engine_client, "get_diffusion_od_config", None)
+        od_config = get_od_config() if callable(get_od_config) else getattr(self._engine_client, "od_config", None)
+        model_class_name = None if od_config is None else getattr(od_config, "model_class_name", None)
+        model = getattr(od_config, "model", None) if od_config is not None else None
+        model = model or self.model_name
+        revision = getattr(od_config, "revision", None) if od_config is not None else None
+        return should_preserve_reference_image_size(
+            model_class_name,
+            model=None if model is None else str(model),
+            revision=revision,
+        )
+
     @property
     def supports_mixed_reference_inputs(self) -> bool:
         """Return whether the configured diffusion model accepts mixed refs."""
@@ -154,7 +171,12 @@ class OmniOpenAIServingVideo:
         provided_fields = request.model_fields_set
         fps_provided = self._request_fps_provided(request)
         vp = request.resolve_video_params()
-        if input_image is not None and vp.width is not None and vp.height is not None:
+        if (
+            input_image is not None
+            and vp.width is not None
+            and vp.height is not None
+            and not self.preserves_reference_image_size
+        ):
             target_size = (vp.width, vp.height)
             image_items = input_image if isinstance(input_image, list) else [input_image]
             resized_images = [

vllm_omni/model_extras/__init__.py

@@ -10,8 +10,10 @@ from vllm_omni.model_extras.registry import (
     get_extra_output_params,
     get_model_class_name,
     get_output_tensor_range,
+    get_transformer_config_subfolder,
     get_x_to_text_model_family,
     should_init_extra_args_for_non_diffusion_stages,
+    should_preserve_reference_image_size,
 )
 
 __all__ = [
@@ -23,6 +25,8 @@ __all__ = [
     "get_extra_output_params",
     "get_model_class_name",
     "get_output_tensor_range",
+    "get_transformer_config_subfolder",
     "get_x_to_text_model_family",
     "should_init_extra_args_for_non_diffusion_stages",
+    "should_preserve_reference_image_size",
 ]

vllm_omni/model_extras/ltx2.py

@@ -21,7 +21,35 @@ LTX_EXTRA_BODY_PARAMS = frozenset(
         "audio_rescale_scale",
         "video_stg_blocks",
         "audio_stg_blocks",
+        "sigmas",
+        "stage_1_sigmas",
+        "stage_2_sigmas",
+        "image_crf",
     }
 )
 
 LTX_EXTRA_OUTPUT_PARAMS = frozenset()
+
+
+def ltx_preserves_reference_image_size(*, model: str | None, revision: str | None = None) -> bool:
+    """Keep source geometry until LTX-2.5 applies conditioning compression."""
+    if model is None:
+        return False
+    from vllm_omni.diffusion.models.ltx2.ltx2_components import (
+        preserves_reference_image_size as ltx_preserves_reference_image_size,
+    )
+
+    return ltx_preserves_reference_image_size(model=model, revision=revision)
+
+
+def ltx_transformer_config_subfolder(
+    *,
+    model: str | None,
+    revision: str | None = None,
+) -> str:
+    """Select the transformer config used by the requested LTX checkpoint."""
+    if model is None:
+        return "transformer"
+    from vllm_omni.diffusion.models.ltx2.ltx2_components import detect_ltx_model_version
+
+    return "transformer_full" if detect_ltx_model_version(model, revision=revision) == "2.5" else "transformer"

vllm_omni/model_extras/registry.py

@@ -4,7 +4,7 @@
 from __future__ import annotations
 
 from collections.abc import Callable, Mapping
-from typing import Any, Literal
+from typing import Any, Literal, Protocol
 
 from PIL import Image
 
@@ -34,7 +34,12 @@ from vllm_omni.model_extras.helios import (
 )
 from vllm_omni.model_extras.hunyuan_image3 import build_x_to_text_prompt as build_hunyuan_x_to_text_prompt
 from vllm_omni.model_extras.lingbot_video import LINGBOT_VIDEO_EXTRA_BODY_PARAMS
-from vllm_omni.model_extras.ltx2 import LTX_EXTRA_BODY_PARAMS, LTX_EXTRA_OUTPUT_PARAMS
+from vllm_omni.model_extras.ltx2 import (
+    LTX_EXTRA_BODY_PARAMS,
+    LTX_EXTRA_OUTPUT_PARAMS,
+    ltx_preserves_reference_image_size,
+    ltx_transformer_config_subfolder,
+)
 from vllm_omni.model_extras.magi_human import (
     MAGI_HUMAN_EXTRA_BODY_PARAMS,
     MAGI_HUMAN_EXTRA_OUTPUT_PARAMS,
@@ -96,6 +101,24 @@ XToTextPromptBuilder = Callable[[str, str, bool], tuple[dict[str, Any], list[int
 OutputTensorRange = Literal["negative_one_to_one", "zero_to_one"]
 
 
+class ReferenceImageSizeResolver(Protocol):
+    def __call__(
+        self,
+        *,
+        model: str | None,
+        revision: str | None = None,
+    ) -> bool: ...
+
+
+class TransformerConfigSubfolderResolver(Protocol):
+    def __call__(
+        self,
+        *,
+        model: str | None,
+        revision: str | None = None,
+    ) -> str: ...
+
+
 def default_x_to_text_prompt(
     model: str,
     prompt: str,
@@ -202,11 +225,14 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = {
         model_class_name: {
             "extra_body_params": LTX_EXTRA_BODY_PARAMS,
             "extra_output_params": LTX_EXTRA_OUTPUT_PARAMS,
+            "reference_image_size_resolver": ltx_preserves_reference_image_size,
         }
         for model_class_name in (
             "LTX2Pipeline",
             "LTX2TwoStagePipeline",
+            "LTX2DistilledOneStagePipeline",
             "LTX2DistilledPipeline",
+            "LTX2DistilledTwoStagePipeline",
         )
     },
     "WanVACEPipeline": {
@@ -229,6 +255,10 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = {
     },
 }
 
+for model_class_name in ("LTX2Pipeline", "LTX2TwoStagePipeline"):
+    _EXTRA_SPECS[model_class_name]["transformer_config_subfolder_resolver"] = ltx_transformer_config_subfolder
+
+
 # Multi-stage discovery reports the top-level wrapper rather than its DiT
 # submodule, so both names must resolve to the same request builders.
 _EXTRA_SPECS["MammothModa2ForConditionalGeneration"] = _EXTRA_SPECS["MammothModa2DiTPipeline"]
@@ -277,6 +307,36 @@ def get_output_tensor_range(model_class_name: str | None) -> OutputTensorRange:
     return spec.get("output_tensor_range", "negative_one_to_one")
 
 
+def get_transformer_config_subfolder(
+    model_class_name: str | None,
+    *,
+    model: str | None,
+    revision: str | None = None,
+) -> str:
+    """Return the model-declared DiT config subfolder, or the standard default."""
+    spec = _get_spec(model_class_name)
+    resolver: TransformerConfigSubfolderResolver | None = (
+        spec.get("transformer_config_subfolder_resolver") if spec else None
+    )
+    return resolver(model=model, revision=revision) if resolver else "transformer"
+
+
+def should_preserve_reference_image_size(
+    model_class_name: str | None,
+    *,
+    model: str | None,
+    revision: str | None = None,
+) -> bool:
+    """Return whether the selected pipeline owns reference-image resizing."""
+    if model_class_name is None and model is not None:
+        from vllm_omni.diffusion.data import resolve_model_class_name
+
+        model_class_name = resolve_model_class_name(model, revision=revision)
+    spec = _get_spec(model_class_name)
+    resolver: ReferenceImageSizeResolver | None = spec.get("reference_image_size_resolver") if spec else None
+    return bool(resolver and resolver(model=model, revision=revision))
+
+
 def should_init_extra_args_for_non_diffusion_stages(model_class_name: str | None) -> bool:
     spec = _get_spec(model_class_name)
     return bool(spec and spec.get("init_extra_args_for_non_diffusion_stages", False))