From 83d1f11449a0eaf2910a0a8e34c055df20e4f169 Mon Sep 17 00:00:00 2001
From: Mu GuanLin <1203789601@qq.com>
Date: Fri, 14 Aug 2026 00:13:14 +0800
Subject: [PATCH] [Feat] LTX Standard Two-Stage Pipeline (#5500)
Signed-off-by: mglyn <1203789601@qq.com>
---
.buildkite/cuda/test-nightly.yml | 10 +-
docs/models/supported_models.md | 3 +-
docs/user_guide/quantization/fp8.md | 4 +-
recipes/LTX/LTX-2.md | 118 +++--
.../models/ltx2/test_ltx2_latents.py | 43 +-
.../models/ltx2/test_ltx2_parallel.py | 14 +
.../models/ltx2/test_ltx2_phase_adapter.py | 483 ++++++++++++++++++
.../models/ltx2/test_ltx2_pipeline.py | 372 +++++++++++++-
.../models/ltx2/test_ltx2_transformer.py | 93 +++-
tests/e2e/accuracy/ltx/run_ltx_reference.py | 155 ++++--
.../ltx/test_ltx_official_similarity.py | 444 ++++++++++++----
tests/e2e/online_serving/test_ltx.py | 109 ++++
.../e2e/online_serving/test_ltx2_expansion.py | 91 ----
tests/model_extras/test_model_extras.py | 11 +-
tests/model_tests/diffusion/test_alignment.py | 1 +
vllm_omni/diffusion/models/ltx2/__init__.py | 6 +-
.../models/ltx2/ltx2_adapter_parser.py | 156 ++++++
.../diffusion/models/ltx2/ltx2_components.py | 133 ++++-
.../diffusion/models/ltx2/ltx2_denoise.py | 23 +-
.../diffusion/models/ltx2/ltx2_guidance.py | 32 +-
.../diffusion/models/ltx2/ltx2_latents.py | 50 +-
.../models/ltx2/ltx2_phase_adapter.py | 430 ++++++++++++++++
.../diffusion/models/ltx2/ltx2_recipes.py | 50 ++
.../diffusion/models/ltx2/ltx2_request.py | 31 +-
.../diffusion/models/ltx2/ltx2_runtime.py | 38 +-
.../diffusion/models/ltx2/ltx2_transformer.py | 63 ++-
.../models/ltx2/pipeline_ltx2_two_stage.py | 45 +-
vllm_omni/diffusion/registry.py | 6 +
vllm_omni/model_extras/registry.py | 1 +
29 files changed, 2644 insertions(+), 371 deletions(-)
create mode 100644 tests/diffusion/models/ltx2/test_ltx2_phase_adapter.py
create mode 100644 tests/e2e/online_serving/test_ltx.py
delete mode 100644 tests/e2e/online_serving/test_ltx2_expansion.py
create mode 100644 vllm_omni/diffusion/models/ltx2/ltx2_adapter_parser.py
create mode 100644 vllm_omni/diffusion/models/ltx2/ltx2_phase_adapter.py
diff --git a/.buildkite/cuda/test-nightly.yml b/.buildkite/cuda/test-nightly.yml
index d3246a98f..ea74847b9 100644
--- a/.buildkite/cuda/test-nightly.yml
+++ b/.buildkite/cuda/test-nightly.yml
@@ -456,7 +456,7 @@ steps:
timeout_in_minutes: 90
commands:
- pytest -s -v tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py -m "full_model and cuda" --run-level "full_model"
- - pytest -s -v tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_ltx2_expansion.py -m "full_model and cuda" --run-level "full_model"
+ - 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 · Wan2.2 Accuracy Test"
@@ -477,10 +477,16 @@ steps:
commands:
- pytest -s -v tests/e2e/accuracy/hunyuanvideo15_t2v/test_hunyuanvideo15_t2v_video_similarity.py -m full_model --run-level full_model
- pytest -s -v tests/e2e/accuracy/hunyuanvideo15_i2v/test_hunyuanvideo15_i2v_video_similarity.py -m full_model --run-level full_model
- - pytest -s -v 'tests/e2e/accuracy/ltx/test_ltx_official_similarity.py::test_ltx_one_stage_matches_official[ltx2_3_i2v]' -m full_model --run-level full_model
- pytest -s -v tests/e2e/accuracy/test_video_streaming_output_similarity.py -m full_model --run-level full_model
mirror_hardwares: h100_2
+ - label: ":full_moon: Diffusion X2V · LTX Accuracy Test"
+ timeout_in_minutes: 180
+ commands:
+ - pytest -s -v 'tests/e2e/accuracy/ltx/test_ltx_official_similarity.py::test_ltx_matches_official[ltx2_3_i2v]' -m full_model --run-level full_model
+ - 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 · Perf Test"
key: nightly-diffusion-x2v-performance
timeout_in_minutes: 180
diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md
index b9686047a..73fdc55e4 100644
--- a/docs/models/supported_models.md
+++ b/docs/models/supported_models.md
@@ -42,7 +42,8 @@ th {
| `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
| `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` | ✅︎ | ✅︎ | | | — |
-| `LTX2DistilledPipeline` | LTX-2 distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | | — |
+| `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 | ✅︎ | ✅︎ | | | — |
| `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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
diff --git a/docs/user_guide/quantization/fp8.md b/docs/user_guide/quantization/fp8.md
index ddca3ea19..05473b15e 100644
--- a/docs/user_guide/quantization/fp8.md
+++ b/docs/user_guide/quantization/fp8.md
@@ -80,12 +80,14 @@ warmup_quack_fp8([(14040, 2048, 6144), (14040, 2048, 2048)])
## Model Type Support
-### Diffusion Model (Qwen-Image, Wan2.2)
+### Diffusion Models
| Model | HF models | Online | Pre-calibrated | Recommendation | `ignored_layers` | Text-Encoder quantization |
|-------|-----------|:-------:|:------:|----------------|------------------|------------------|
| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Yes | Yes | Skip sensitive image-stream MLPs when quality regresses | `img_mlp` | |
| Wan2.2 | Wan2.2 diffusion pipelines | Not validated | Not validated | Validate against BF16 before documenting as supported | TBD | |
+| LTX-2 | `Lightricks/LTX-2`, `rootonchair/LTX-2-19b-distilled` | Yes | Not validated | Transformer only; use dynamic phase LoRA for ordinary two-stage | None | |
+| LTX-2.3 | `diffusers/LTX-2.3-Diffusers`, `diffusers/LTX-2.3-Distilled-Diffusers` | Yes | Not validated | Transformer only; use dynamic phase LoRA for ordinary two-stage | None | |
| Z-Image | `Tongyi-MAI/Z-Image-Turbo` | Yes | Yes | All layers | None | ✅︎ |
| FLUX.1 | `black-forest-labs/FLUX.1-dev`, `black-forest-labs/FLUX.1-schnell` | Yes | Yes | All layers | None | |
| FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | Yes | Not validated | All layers | None | ✅︎ |
diff --git a/recipes/LTX/LTX-2.md b/recipes/LTX/LTX-2.md
index 3142f9fd1..123880d09 100644
--- a/recipes/LTX/LTX-2.md
+++ b/recipes/LTX/LTX-2.md
@@ -4,11 +4,20 @@
## Pipelines
-| Model | Checkpoint | Task | `--model-class-name` | Batching |
-|---|---|---|---|---|
-| LTX-2 | `Lightricks/LTX-2` | One-stage T2V/I2V | `LTX2Pipeline` | Yes |
-| LTX-2 distilled | `rootonchair/LTX-2-19b-distilled` | Two-stage T2V/I2V | `LTX2DistilledPipeline` | No |
-| LTX-2.3 | `diffusers/LTX-2.3-Diffusers` | One-stage T2V/I2V | `LTX2Pipeline` | Yes |
+| `--model-class-name` | Task | Required checkpoint repositories |
+|---|---|---|
+| `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` |
+| `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`
`Lightricks/LTX-2.3` |
+| `LTX2DistilledPipeline` | LTX-2.3 full-distilled two-stage T2V/I2V | `diffusers/LTX-2.3-Distilled-Diffusers`
`Lightricks/LTX-2.3` |
+
+Repositories in the table are download units. A full pipeline repository
+contains the Transformer, text encoder, connectors, VAEs, vocoder, scheduler,
+and tokenizer; an additional repository supplies LoRA or upsampler sidecars.
+The runtime first searches the model root for sidecars under their official
+filenames, then downloads them from the matching Lightricks Hub repository.
`LTX2Pipeline` is the unified one-stage entry. Checkpoint metadata selects the
LTX-2 or LTX-2.3 profile; omitting an image selects T2V, while one initial
@@ -16,10 +25,10 @@ image selects I2V. Both one-stage repositories declare this class, so
`--model-class-name` is optional. LTX-2.3 requires the Diffusers checkpoint;
the raw `Lightricks/LTX-2.3` safetensors repository is not directly loadable.
-`LTX2DistilledPipeline` is the unified distilled entry: omit `image` for T2V
-or provide one initial image for I2V. As in the official LTX pipelines,
-distilled mode is selected explicitly rather than inferred from the checkpoint
-name or metadata; always pass `--model-class-name LTX2DistilledPipeline`.
+`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.
## API Migration
@@ -68,24 +77,28 @@ The model recipe also supplies the default negative prompt. Top-level
API defaults `num_frames` to `1`, so set it explicitly; the offline LTX scripts
default to `121`.
-## Distilled Two-Stage Defaults
+## Two-Stage Defaults
-| Parameter | Value |
-|---|---:|
-| Final width × height | 1536 × 1024 |
-| Stage 1 width × height | 768 × 512 |
-| Frames / frame rate | 121 / 24 |
-| Stage 1 / Stage 2 denoise steps | 8 / 3 |
-| Guidance | Fixed positive-only |
+| Parameter | Ordinary | Full-distilled |
+|---|---:|---:|
+| Final width × height | 1536 × 1024 | 1536 × 1024 |
+| Stage 1 width × height | 768 × 512 | 768 × 512 |
+| Frames / frame rate | 121 / 24 | 121 / 24 |
+| Stage 1 / Stage 2 steps | 40 (LTX-2) or 30 (LTX-2.3) / 3 | 8 / 3 |
+| Guidance | Stage 1 guided; Stage 2 positive-only | Fixed positive-only |
-The dimensions passed to the API are the final output dimensions and must be
-divisible by 64. Both one- and two-stage requests require `num_frames = 8k+1`.
-The distilled recipe fixes both sigma schedules; `num_inference_steps` must be
-`8` when supplied.
+API dimensions are final dimensions and must be divisible by 64. All LTX
+requests require `num_frames = 8k+1`. Ordinary Stage 1 uses the LTX-2 or
+LTX-2.3 one-stage defaults shown above. Distilled schedules are fixed, so
+`num_inference_steps`, when supplied, must be `8`. Both entries reject custom
+sigmas and input latents.
+
+Ordinary two-stage uses layer-fused LoRA for an unquantized BF16 Transformer
+and automatically switches to dynamic LoRA when quantization is enabled.
## Serving
-Start either checkpoint:
+Start either one-stage checkpoint:
```bash
vllm serve Lightricks/LTX-2 --omni --stage-init-timeout 600
@@ -95,11 +108,19 @@ vllm serve Lightricks/LTX-2 --omni --stage-init-timeout 600
vllm serve diffusers/LTX-2.3-Diffusers --omni --stage-init-timeout 600
```
-Start the distilled checkpoint with its two-stage entry:
+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
+# 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
+# Ordinary LTX-2.3 two-stage
+vllm serve diffusers/LTX-2.3-Diffusers --omni \
+ --model-class-name LTX2TwoStagePipeline \
+ --enable-layerwise-offload \
+ --stage-init-timeout 600
```
The same server handles T2V and I2V. A T2V request using the selected recipe's
@@ -135,8 +156,9 @@ together with `input_reference`.
## Guidance
-LTX supports independent video/audio CFG, spatio-temporal guidance (STG),
-cross-modality guidance, and rescaling:
+One-stage and ordinary Stage 1 support independent video/audio CFG,
+spatio-temporal guidance (STG), cross-modality guidance, and rescaling.
+Distilled stages and ordinary Stage 2 are fixed positive-only.
| Parameter | Default | Effect | Alias |
|---|---:|---|---|
@@ -172,9 +194,9 @@ complete guidance plan in parallel. Despite the legacy `cfg` name, it covers
text CFG, STG, and cross-modality guidance passes; guidance rescaling is
applied after all pass predictions are gathered.
-The default one-stage recipe has four Transformer passes per denoise step:
-`cond`, `uncond`, `ptb` (STG), and `mod` (cross-modality). The useful balanced
-configurations are therefore:
+The default one-stage recipe and ordinary Stage 1 have four Transformer passes
+per denoise step: `cond`, `uncond`, `ptb` (STG), and `mod`
+(cross-modality). The useful balanced configurations are therefore:
| `--cfg-parallel-size` | Passes per rank | Guidance-slot utilization | Notes |
|---:|---:|---:|---|
@@ -195,9 +217,8 @@ vllm serve diffusers/LTX-2.3-Diffusers --omni \
```
The total device count is the product of `cfg_parallel_size` and the other
-configured parallel dimensions. Distilled pipelines use positive-only
-guidance, so `cfg_parallel_size > 1` does not provide guidance speedup for
-them.
+configured parallel dimensions. Positive-only distilled phases do not benefit
+from `cfg_parallel_size > 1`.
### Python API
@@ -242,22 +263,22 @@ noted below.
| `image` | image or batch, `None` | Direct value wins over request images; no image selects T2V. I2V accepts one image per prompt, and a batch cannot mix T2V/I2V. |
| `prompt` | string or list, `None` | Positive-text fallback; request prompts win. Mutually exclusive with `prompt_embeds`. |
| `negative_prompt` | string or list, `None` | Fallback after request values and before the recipe default; mutually exclusive with negative embeddings. |
-| `height` | `int`, `None` | Request → direct value → recipe default; divisible by 32 one-stage or 64 distilled. |
+| `height` | `int`, `None` | Request → direct value → recipe default; divisible by 32 one-stage or 64 two-stage. |
| `width` | `int`, `None` | Same precedence and alignment as `height`. |
| `num_frames` | `int`, `None` | Request → direct value → recipe default; must be `8k+1` and also determines audio duration with `frame_rate`. |
| `frame_rate` | `float`, `None` | Request `frame_rate` → request `fps` → direct value → recipe default. |
-| `num_inference_steps` | `int`, `None` | Request → direct value → recipe default; minimum 2 one-stage, fixed at 8 for distilled Stage 1. Custom one-stage `sigmas` determine actual steps. |
+| `num_inference_steps` | `int`, `None` | Request → direct value → recipe default; minimum 2 one-stage and ordinary Stage 1, fixed at 8 for distilled Stage 1. Custom one-stage `sigmas` determine actual steps. |
| `sigmas` | list of float, `None` | One-stage only. Request values win; every request in a fused batch must use the same schedule. |
| `timesteps` | list of int, `None` | Compatibility slot; LTX accepts only `None`. Use `sigmas`. |
-| `guidance_scale` | `float`, `None` | One-stage common video/audio CFG fallback; an explicit request value wins. Distilled two-stage uses fixed positive-only guidance. |
+| `guidance_scale` | `float`, `None` | One-stage and ordinary Stage 1 common video/audio CFG fallback; an explicit request value wins. Distilled phases use fixed positive-only guidance. |
| `guidance_rescale` | `float`, `None` | Accepts only `None` or `0.0`; use the modality rescale fields. |
| `noise_scale` | `float`, `0.0` | Compatibility slot; LTX accepts only `0.0`. |
| `num_videos_per_prompt` | `int`, `1` | Output-count fallback; positive request `num_outputs_per_prompt` wins. |
| `generator` | generator or list, `None` | Explicit RNG; otherwise request generators/seeds are collated. Lists must match the effective output batch. |
-| `latents` | tensor, `None` | Request tensors win. One-stage accepts packed `[B, S, C]` and validated unpacked video layouts. |
-| `audio_latents` | tensor, `None` | One-stage initial audio latents; request tensors win and are collated. |
+| `latents` | tensor, `None` | One-stage only. Request tensors win; packed `[B, S, C]` and validated unpacked video layouts are accepted. |
+| `audio_latents` | tensor, `None` | One-stage only. Request tensors win and are collated. |
| `prompt_embeds` | tensor, `None` | Precomputed positive conditioning; requires `prompt_attention_mask` and cannot accompany `prompt`. |
-| `negative_prompt_embeds` | tensor, `None` | One-stage precomputed negative conditioning; requires its mask and cannot accompany `negative_prompt`. |
+| `negative_prompt_embeds` | tensor, `None` | One-stage and ordinary Stage 1 conditioning; requires its mask and cannot accompany `negative_prompt`. |
| `prompt_attention_mask` | tensor, `None` | Mask for positive embeddings; request `prompt_attention_mask` or `attention_mask` wins. |
| `negative_prompt_attention_mask` | tensor, `None` | Mask for negative embeddings; request negative mask fields win. |
| `decode_timestep` | float or list, `0.0` | Video-VAE decode timestep; request value wins. Lists may match 1, prompt batch, or output batch. |
@@ -273,13 +294,13 @@ request prompt payload; LTX guidance fields live in sampling `extra_args`.
### Recipe-Specific Request Capabilities
-| Override | One-stage | Distilled two-stage |
-|---|---|---|
-| Guidance | Supported | Fixed positive-only |
-| Negative prompt/embeddings | Supported | Rejected |
-| `num_inference_steps` | Supported | Fixed at 8 for Stage 1; Stage 2 uses 3 |
-| Custom `sigmas` | Supported | Rejected; both phases use fixed schedules |
-| Video/audio latents | Supported | Rejected |
+| Override | One-stage | Ordinary two-stage | Distilled two-stage |
+|---|---|---|---|
+| Guidance | Supported | Stage 1 only; Stage 2 is positive-only | Fixed positive-only |
+| Negative prompt/embeddings | Supported | Supported by Stage 1 | Rejected |
+| `num_inference_steps` | Supported | Controls Stage 1; Stage 2 uses 3 | Fixed at 8 for Stage 1; Stage 2 uses 3 |
+| Custom `sigmas` | Supported | Rejected | Rejected; both phases use fixed schedules |
+| Video/audio latents | Supported | Rejected | Rejected |
These capability checks apply equally to direct `forward` keywords and
values in `OmniDiffusionSamplingParams`; unsupported values fail instead of
@@ -309,6 +330,10 @@ bundled offline CLI do not currently expose `sigmas`.
`--max-num-seqs 1`.
- `--cfg-parallel-size` shards the complete LTX guidance plan, including STG,
modality guidance, and rescale-compatible prediction gathering.
+- Sequence parallelism may pad audio latents. Pure Ulysses masks the padding;
+ Ring cannot, so audio length must be SP-divisible. Use `ring_degree=1` or a
+ divisible request shape.
+- Cache-DiT is one-stage only; multi-stage configurations are rejected.
## Operational Notes
@@ -323,14 +348,15 @@ bundled offline CLI do not currently expose `sigmas`.
request parameter.
- For benchmarks, use `tests/dfx/perf/tests/test_ltx2_vllm_omni.json` with
`tests/dfx/perf/scripts/run_diffusion_benchmark.py`.
-- Current two-stage support is limited to the distilled checkpoint; official
- LTX two-stage/HQ execution remains out of scope.
+- Ordinary and full-distilled two-stage T2V/I2V are supported; HQ execution
+ remains out of scope.
## References
-
-
-
+-
- [Online video generation](../../docs/user_guide/examples/online_serving/text_to_video.md)
- [Diffusion execution modes](../../docs/user_guide/diffusion/execution_modes.md)
- [T2V offline example](../../examples/offline_inference/text_to_video/text_to_video.md)
diff --git a/tests/diffusion/models/ltx2/test_ltx2_latents.py b/tests/diffusion/models/ltx2/test_ltx2_latents.py
index 36f7e9687..cad5c300d 100644
--- a/tests/diffusion/models/ltx2/test_ltx2_latents.py
+++ b/tests/diffusion/models/ltx2/test_ltx2_latents.py
@@ -23,7 +23,7 @@ def _make_pipeline(pipeline_cls, sequence_parallel_size: int = 1):
return pipeline
-def test_prepare_video_latents_samples_directly_in_packed_token_space():
+def test_prepare_video_latents_matches_official_values_and_token_major_layout():
pipeline = _make_pipeline(LTX2Pipeline)
pipeline.vae_spatial_compression_ratio = 8
pipeline.vae_temporal_compression_ratio = 8
@@ -44,6 +44,7 @@ def test_prepare_video_latents_samples_directly_in_packed_token_space():
)
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+ assert actual.stride()[1:] == (1, actual.shape[1])
def test_prepare_audio_latents_samples_directly_in_packed_token_space():
@@ -81,6 +82,42 @@ def test_prepare_audio_latents_pads_generated_dummy_length_for_sp():
assert original_num_frames == 1
assert padded_num_frames == 2
assert latents.shape == (1, 2, 128)
+ torch.testing.assert_close(latents[:, 1:], torch.zeros_like(latents[:, 1:]))
+
+
+def test_prepare_audio_latents_request_rng_is_invariant_to_sp_padding():
+ pipeline_sp1 = _make_pipeline(LTX2Pipeline, sequence_parallel_size=1)
+ pipeline_sp4 = _make_pipeline(LTX2Pipeline, sequence_parallel_size=4)
+ generator_sp1 = torch.Generator().manual_seed(42)
+ generator_sp4 = torch.Generator().manual_seed(42)
+
+ latents_sp1, _, _ = pipeline_sp1.prepare_audio_latents(
+ batch_size=1,
+ num_channels_latents=2,
+ num_mel_bins=8,
+ audio_latent_length=3,
+ dtype=torch.float32,
+ device=torch.device("cpu"),
+ generator=generator_sp1,
+ )
+ latents_sp4, _, _ = pipeline_sp4.prepare_audio_latents(
+ batch_size=1,
+ num_channels_latents=2,
+ num_mel_bins=8,
+ audio_latent_length=3,
+ dtype=torch.float32,
+ device=torch.device("cpu"),
+ generator=generator_sp4,
+ )
+
+ torch.testing.assert_close(latents_sp4[:, :3], latents_sp1, rtol=0, atol=0)
+ torch.testing.assert_close(latents_sp4[:, 3:], torch.zeros_like(latents_sp4[:, 3:]))
+ torch.testing.assert_close(
+ torch.randn(8, generator=generator_sp4),
+ torch.randn(8, generator=generator_sp1),
+ rtol=0,
+ atol=0,
+ )
def test_prepare_audio_latents_pads_packed_sequence_dim_for_provided_latents():
@@ -137,7 +174,9 @@ def test_prepare_audio_latents_accepts_already_padded_4d_latents_for_sp():
assert original_num_frames == 10
assert padded_num_frames == 12
assert padded.shape == (1, 12, 8)
- torch.testing.assert_close(padded, latent_ops.pack_audio_latents(latents))
+ packed = latent_ops.pack_audio_latents(latents)
+ torch.testing.assert_close(padded[:, :10], packed[:, :10])
+ torch.testing.assert_close(padded[:, 10:], torch.zeros_like(padded[:, 10:]))
def test_resolve_audio_latent_length_preserves_legacy_4d_shape_inference():
diff --git a/tests/diffusion/models/ltx2/test_ltx2_parallel.py b/tests/diffusion/models/ltx2/test_ltx2_parallel.py
index 8d8c22e79..95ca6e868 100644
--- a/tests/diffusion/models/ltx2/test_ltx2_parallel.py
+++ b/tests/diffusion/models/ltx2/test_ltx2_parallel.py
@@ -12,6 +12,19 @@ import torch
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+def test_velocity_from_x0_uses_official_host_scalar_sigma(mocker):
+ from vllm_omni.diffusion.models.ltx2.ltx2_guidance import velocity_from_x0
+
+ sigma = mocker.MagicMock()
+ sigma.to.return_value.item.return_value = 0.5
+
+ actual = velocity_from_x0(torch.tensor([1.0]), torch.tensor([0.0]), sigma)
+
+ sigma.to.assert_called_once_with(torch.float32)
+ sigma.to.return_value.item.assert_called_once_with()
+ torch.testing.assert_close(actual, torch.tensor([2.0]))
+
+
class TestCFGParallelHelpers:
"""Test LTX-2.3 CFG helper math without loading model weights."""
@@ -199,6 +212,7 @@ class TestCFGParallelHelpers:
),
attention_kwargs=None,
audio_scheduler=SimpleNamespace(sigmas=torch.stack([audio_sigma])),
+ original_audio_num_frames=state.audio.shape[1],
)
actual_video, actual_audio = executor.predict_parallel_guidance(
diff --git a/tests/diffusion/models/ltx2/test_ltx2_phase_adapter.py b/tests/diffusion/models/ltx2/test_ltx2_phase_adapter.py
new file mode 100644
index 000000000..4b1d18cb5
--- /dev/null
+++ b/tests/diffusion/models/ltx2/test_ltx2_phase_adapter.py
@@ -0,0 +1,483 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Unit tests for the fixed LTX refinement-phase adapter."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from safetensors.torch import save_file
+
+from vllm_omni.diffusion.models.ltx2 import ltx2_phase_adapter
+from vllm_omni.diffusion.models.ltx2.ltx2_adapter_parser import (
+ AdapterTarget,
+ parse_ltx_adapter,
+)
+from vllm_omni.diffusion.models.ltx2.ltx2_components import LTX23_TWO_STAGE_COMPONENT_PROFILE
+from vllm_omni.diffusion.models.ltx2.ltx2_phase_adapter import LTXPhaseAdapterRuntime
+from vllm_omni.diffusion.models.ltx2.ltx2_recipes import (
+ LTX2_TWO_STAGE_RECIPE,
+ LTX_DISTILLED_ADAPTER_SLOT,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _target(
+ lora_a,
+ lora_b,
+ *,
+ source_module="proj",
+ module="proj",
+ target_slice=None,
+):
+ return AdapterTarget(
+ source_module=source_module,
+ module=module,
+ slice=target_slice,
+ a_key="a",
+ b_key="b",
+ a_shape=tuple(lora_a.shape),
+ b_shape=tuple(lora_b.shape),
+ )
+
+
+def _wrap(
+ layer,
+ lora_a,
+ lora_b,
+ *,
+ source_module="proj",
+ module="proj",
+ target_slice=None,
+ layer_fused=False,
+):
+ target = _target(
+ lora_a,
+ lora_b,
+ source_module=source_module,
+ module=module,
+ target_slice=target_slice,
+ )
+ wrapper = ltx2_phase_adapter._PhaseAdapterLinear(
+ layer,
+ [target],
+ adapter_dtype=lora_a.dtype,
+ layer_fused=layer_fused,
+ )
+ wrapper.load_target(target, lora_a, lora_b)
+ wrapper.set_enabled(True)
+ return wrapper
+
+
+class _LoRAAttention(torch.nn.Module):
+ def __init__(self) -> None:
+ super().__init__()
+ self.to_qkv = torch.nn.Identity()
+
+
+class _LoRATransformer(torch.nn.Module):
+ stacked_params_mapping = ((".attn1.to_qkv", ".attn1.to_q", "q"),)
+
+ def __init__(self) -> None:
+ super().__init__()
+ self.proj_in = torch.nn.Linear(4, 6, bias=False)
+ block = torch.nn.Module()
+ block.attn1 = _LoRAAttention()
+ self.transformer_blocks = torch.nn.ModuleList([block])
+
+
+def test_ltx_adapter_parser_maps_official_and_packed_module_names(tmp_path):
+ adapter_path = tmp_path / "distilled.safetensors"
+ save_file(
+ {
+ "diffusion_model.text_embedding_projection.aggregate_embed.lora_A.weight": torch.randn(2, 4),
+ "diffusion_model.patchify_proj.lora_A.weight": torch.randn(2, 4),
+ "diffusion_model.patchify_proj.lora_B.weight": torch.randn(6, 2),
+ "diffusion_model.transformer_blocks.0.attn1.to_q.lora_A.weight": torch.randn(2, 4),
+ "diffusion_model.transformer_blocks.0.attn1.to_q.lora_B.weight": torch.randn(4, 2),
+ },
+ adapter_path,
+ )
+
+ manifest = parse_ltx_adapter(_LoRATransformer(), str(adapter_path))
+ assert {(target.source_module, target.module, target.slice) for target in manifest.targets} == {
+ ("proj_in", "proj_in", None),
+ ("transformer_blocks.0.attn1.to_q", "transformer_blocks.0.attn1.to_qkv", "q"),
+ }
+
+
+def test_ltx_adapter_parser_rejects_unpaired_b_tensor(tmp_path):
+ adapter_path = tmp_path / "distilled.safetensors"
+ save_file(
+ {"diffusion_model.patchify_proj.lora_B.weight": torch.randn(6, 2)},
+ adapter_path,
+ )
+
+ with pytest.raises(ValueError, match="Missing paired LoRA A tensors"):
+ parse_ltx_adapter(_LoRATransformer(), str(adapter_path))
+
+
+def _pipeline(*, dtype, quantization_config):
+ return SimpleNamespace(
+ pipeline_recipe=LTX2_TWO_STAGE_RECIPE,
+ component_profile=LTX23_TWO_STAGE_COMPONENT_PROFILE,
+ transformer=object(),
+ od_config=SimpleNamespace(
+ model="unused",
+ lora_path=None,
+ dtype=dtype,
+ quantization_config=quantization_config,
+ ),
+ )
+
+
+@pytest.mark.parametrize(
+ ("dtype", "quantization_config", "layer_fused"),
+ [
+ (torch.float32, object(), False),
+ (torch.bfloat16, None, True),
+ ],
+)
+def test_ltx_phase_adapter_selects_execution_from_quantization(
+ monkeypatch,
+ dtype,
+ quantization_config,
+ layer_fused,
+):
+ pipeline = _pipeline(dtype=dtype, quantization_config=quantization_config)
+ manifest = object()
+ captured = {}
+
+ def resolve_artifact(*_args, **kwargs):
+ assert not kwargs
+ return "adapter"
+
+ def parse(transformer, path):
+ assert transformer is pipeline.transformer
+ assert path == "adapter"
+ return manifest
+
+ class Runtime:
+ def __init__(self, transformer, parsed_manifest, **kwargs):
+ captured.update(transformer=transformer, manifest=parsed_manifest, **kwargs)
+
+ monkeypatch.setattr(ltx2_phase_adapter, "resolve_ltx_artifact", resolve_artifact)
+ monkeypatch.setattr(ltx2_phase_adapter, "parse_ltx_adapter", parse)
+ monkeypatch.setattr(ltx2_phase_adapter, "LTXPhaseAdapterRuntime", Runtime)
+
+ runtime = ltx2_phase_adapter.build_ltx_phase_adapter(pipeline)
+
+ assert isinstance(runtime, Runtime)
+ assert captured == {
+ "transformer": pipeline.transformer,
+ "manifest": manifest,
+ "dtype": dtype,
+ "layer_fused": layer_fused,
+ }
+
+
+def test_ltx_phase_adapter_rejects_non_bfloat16_layer_fused(monkeypatch):
+ pipeline = _pipeline(dtype=torch.float16, quantization_config=None)
+ monkeypatch.setattr(ltx2_phase_adapter, "resolve_ltx_artifact", lambda *_args, **_kwargs: "adapter")
+
+ with pytest.raises(ValueError, match="bfloat16"):
+ ltx2_phase_adapter.build_ltx_phase_adapter(pipeline)
+
+
+def test_ltx_phase_adapter_keeps_one_transformer_and_switches_a_fixed_slot(tmp_path):
+ class TinyTransformer(torch.nn.Module):
+ stacked_params_mapping = ()
+
+ def __init__(self) -> None:
+ super().__init__()
+ self.proj = torch.nn.Linear(2, 2, bias=False)
+
+ adapter_path = tmp_path / "distilled.safetensors"
+ save_file(
+ {
+ "diffusion_model.proj.lora_A.weight": torch.tensor([[1.0, 2.0]]),
+ "diffusion_model.proj.lora_B.weight": torch.tensor([[3.0], [4.0]]),
+ },
+ adapter_path,
+ )
+ transformer = TinyTransformer()
+ runtime = LTXPhaseAdapterRuntime(
+ transformer,
+ parse_ltx_adapter(transformer, str(adapter_path)),
+ dtype=torch.float32,
+ )
+
+ assert runtime.base_parameter_name("proj.weight") == "proj.base_layer.weight"
+ assert {name for name, _ in transformer.named_parameters()} == {"proj.base_layer.weight"}
+ with torch.no_grad():
+ transformer.proj.base_layer.weight.copy_(torch.eye(2))
+ runtime.finalize()
+
+ x = torch.tensor([[1.0, 1.0]])
+ runtime.activate(None)
+ torch.testing.assert_close(transformer.proj(x), torch.tensor([[1.0, 1.0]]))
+ runtime.activate(LTX_DISTILLED_ADAPTER_SLOT)
+ torch.testing.assert_close(transformer.proj(x), torch.tensor([[10.0, 13.0]]))
+
+
+def test_ltx_layer_fused_adapter_matches_official_bf16_weight_fusion():
+ layer = torch.nn.Linear(4, 3, bias=True, dtype=torch.bfloat16)
+ base_weight = torch.tensor(
+ [
+ [0.5, -0.25, 0.125, 0.75],
+ [-0.375, 0.625, 0.25, -0.5],
+ [0.875, -0.125, -0.75, 0.375],
+ ],
+ dtype=torch.bfloat16,
+ )
+ lora_a = torch.tensor(
+ [
+ [0.1015625, -0.203125, 0.3046875, -0.40625],
+ [0.5078125, -0.609375, 0.7109375, -0.8125],
+ ],
+ dtype=torch.bfloat16,
+ )
+ lora_b = torch.tensor(
+ [[0.9140625, -0.1171875], [-0.21875, 0.3203125], [0.421875, -0.5234375]],
+ dtype=torch.bfloat16,
+ )
+ with torch.no_grad():
+ layer.weight.copy_(base_weight)
+ layer.bias.copy_(torch.tensor([0.125, -0.25, 0.5], dtype=torch.bfloat16))
+ wrapper = _wrap(layer, lora_a, lora_b, layer_fused=True)
+
+ expected_weight = torch.matmul(lora_b, lora_a)
+ expected_weight.add_(base_weight)
+ x = torch.tensor([[0.6015625, -0.703125, 0.8046875, -0.90625]], dtype=torch.bfloat16)
+ expected = torch.nn.functional.linear(x, expected_weight, layer.bias)
+
+ assert torch.equal(wrapper(x), expected)
+ assert torch.equal(layer.weight, base_weight)
+
+
+def test_ltx_layer_fused_row_parallel_uses_one_collective(monkeypatch):
+ from vllm.model_executor.layers.linear import UnquantizedLinearMethod
+
+ class FakeRow(torch.nn.Module):
+ input_size = 4
+ output_size = 2
+ input_size_per_partition = 2
+ input_is_parallel = True
+ tp_rank = 0
+ tp_size = 2
+ reduce_results = True
+ skip_bias_add = False
+ return_bias = False
+
+ def __init__(self):
+ super().__init__()
+ self.weight = torch.nn.Parameter(torch.eye(2, dtype=torch.bfloat16))
+ self.bias = torch.nn.Parameter(torch.tensor([0.25, -0.5], dtype=torch.bfloat16))
+ self.quant_method = UnquantizedLinearMethod()
+
+ monkeypatch.setattr(ltx2_phase_adapter, "RowParallelLinear", FakeRow)
+ reductions = []
+
+ def all_reduce(value):
+ reductions.append(value.clone())
+ return value + 1
+
+ monkeypatch.setattr(ltx2_phase_adapter, "tensor_model_parallel_all_reduce", all_reduce)
+ layer = FakeRow()
+ lora_a = torch.tensor([[0.5, -0.25, 0.75, -0.125]], dtype=torch.bfloat16)
+ lora_b = torch.tensor([[0.25], [-0.5]], dtype=torch.bfloat16)
+ wrapper = _wrap(layer, lora_a, lora_b, layer_fused=True)
+
+ expected_weight = torch.matmul(lora_b, lora_a[:, :2])
+ expected_weight.add_(layer.weight)
+ x = torch.tensor([[1.0, 2.0]], dtype=torch.bfloat16)
+ expected_local = torch.nn.functional.linear(x, expected_weight, layer.bias)
+
+ actual = wrapper(x)
+ assert len(reductions) == 1
+ assert torch.equal(reductions[0], expected_local)
+ assert torch.equal(actual, expected_local + 1)
+
+
+def test_ltx_layer_fused_column_parallel_uses_local_weight_shard(monkeypatch):
+ from vllm.model_executor.layers.linear import UnquantizedLinearMethod
+
+ class FakeColumn(torch.nn.Module):
+ input_size = 3
+ output_size = 4
+ output_size_per_partition = 2
+ tp_rank = 1
+ tp_size = 2
+ gather_output = False
+ skip_bias_add = False
+ return_bias = False
+
+ def __init__(self):
+ super().__init__()
+ self.weight = torch.nn.Parameter(
+ torch.tensor([[0.5, -0.25, 0.125], [-0.375, 0.625, 0.25]], dtype=torch.bfloat16)
+ )
+ self.bias = torch.nn.Parameter(torch.tensor([0.25, -0.5], dtype=torch.bfloat16))
+ self.quant_method = UnquantizedLinearMethod()
+
+ monkeypatch.setattr(ltx2_phase_adapter, "ColumnParallelLinear", FakeColumn)
+ layer = FakeColumn()
+ base_weight = layer.weight.detach().clone()
+ lora_a = torch.tensor([[0.5, -0.25, 0.75]], dtype=torch.bfloat16)
+ lora_b = torch.tensor([[0.25], [-0.5], [0.75], [-1.0]], dtype=torch.bfloat16)
+ wrapper = _wrap(layer, lora_a, lora_b, layer_fused=True)
+
+ expected_weight = torch.matmul(lora_b[2:], lora_a)
+ expected_weight.add_(base_weight)
+ x = torch.tensor([[1.0, -0.5, 0.25]], dtype=torch.bfloat16)
+ expected = torch.nn.functional.linear(x, expected_weight, layer.bias)
+
+ assert torch.equal(wrapper(x), expected)
+ assert torch.equal(layer.weight, base_weight)
+
+
+def test_ltx_layer_fused_packed_qkv_updates_only_the_target_slice(monkeypatch):
+ from vllm.model_executor.layers.linear import UnquantizedLinearMethod
+
+ class FakeColumn(torch.nn.Module):
+ pass
+
+ class FakeQKV(FakeColumn):
+ input_size = 4
+ head_size = 2
+ v_head_size = 2
+ total_num_heads = 2
+ total_num_kv_heads = 2
+ num_heads = 1
+ num_kv_heads = 1
+ num_kv_head_replicas = 1
+ tp_rank = 1
+ tp_size = 2
+ gather_output = False
+ skip_bias_add = True
+ return_bias = True
+
+ def __init__(self):
+ super().__init__()
+ self.weight = torch.nn.Parameter(
+ torch.arange(24, dtype=torch.float32).reshape(6, 4).to(torch.bfloat16) / 16
+ )
+ self.bias = torch.nn.Parameter(torch.arange(6, dtype=torch.float32).to(torch.bfloat16) / 8)
+ self.quant_method = UnquantizedLinearMethod()
+
+ monkeypatch.setattr(ltx2_phase_adapter, "ColumnParallelLinear", FakeColumn)
+ monkeypatch.setattr(ltx2_phase_adapter, "QKVParallelLinear", FakeQKV)
+ layer = FakeQKV()
+ base_weight = layer.weight.detach().clone()
+ lora_a = torch.tensor([[0.5, -0.25, 0.75, -0.125]], dtype=torch.bfloat16)
+ lora_b = torch.tensor([[0.25], [-0.5], [0.75], [-1.0]], dtype=torch.bfloat16)
+ wrapper = _wrap(
+ layer,
+ lora_a,
+ lora_b,
+ source_module="to_k",
+ module="to_qkv",
+ target_slice="k",
+ layer_fused=True,
+ )
+
+ expected_weight = base_weight.clone()
+ expected_k = torch.matmul(lora_b[2:], lora_a)
+ expected_k.add_(base_weight[2:4])
+ expected_weight[2:4].copy_(expected_k)
+ x = torch.tensor([[1.0, -0.5, 0.25, -0.125]], dtype=torch.bfloat16)
+ expected_output = torch.nn.functional.linear(x, expected_weight)
+
+ actual_output, returned_bias = wrapper(x)
+ assert torch.equal(actual_output, expected_output)
+ assert returned_bias is layer.bias
+ assert torch.equal(layer.weight, base_weight)
+
+
+@pytest.mark.parametrize(
+ ("layer_factory", "target_slice", "expected_a", "expected_b", "expected_output"),
+ [
+ (
+ lambda: type(
+ "FakeColumn",
+ (torch.nn.Module,),
+ {
+ "input_size": 4,
+ "output_size": 4,
+ "output_size_per_partition": 2,
+ "tp_rank": 1,
+ "gather_output": False,
+ "forward": lambda self, x: x.new_zeros((*x.shape[:-1], 2)),
+ },
+ )(),
+ None,
+ torch.tensor([[1.0, 10.0, 100.0, 1000.0]]),
+ torch.tensor([[1.0], [2.0], [3.0], [4.0]]),
+ torch.tensor([[12963.0, 17284.0]]),
+ ),
+ (
+ lambda: type(
+ "FakeRow",
+ (torch.nn.Module,),
+ {
+ "input_size": 4,
+ "output_size": 2,
+ "input_size_per_partition": 2,
+ "tp_rank": 1,
+ "reduce_results": False,
+ "forward": lambda self, x: x.new_zeros((*x.shape[:-1], 2)),
+ },
+ )(),
+ None,
+ torch.tensor([[1.0, 10.0, 100.0, 1000.0]]),
+ torch.tensor([[2.0], [3.0]]),
+ torch.tensor([[8600.0, 12900.0]]),
+ ),
+ (
+ lambda: type(
+ "FakeQKV",
+ (torch.nn.Module,),
+ {
+ "input_size": 4,
+ "head_size": 2,
+ "v_head_size": 2,
+ "total_num_heads": 2,
+ "total_num_kv_heads": 2,
+ "num_heads": 1,
+ "num_kv_heads": 1,
+ "num_kv_head_replicas": 1,
+ "tp_rank": 1,
+ "forward": lambda self, x: x.new_zeros((*x.shape[:-1], 6)),
+ },
+ )(),
+ "k",
+ torch.tensor([[1.0, 10.0, 100.0, 1000.0]]),
+ torch.tensor([[1.0], [2.0], [3.0], [4.0]]),
+ torch.tensor([[0.0, 0.0, 12963.0, 17284.0, 0.0, 0.0]]),
+ ),
+ ],
+)
+def test_ltx_phase_adapter_uses_local_tp_lora_slices(
+ monkeypatch,
+ layer_factory,
+ target_slice,
+ expected_a,
+ expected_b,
+ expected_output,
+):
+ layer = layer_factory()
+ layer.register_buffer("device_anchor", torch.empty(0))
+ layer_type = type(layer)
+ if target_slice == "k":
+ monkeypatch.setattr(ltx2_phase_adapter, "QKVParallelLinear", layer_type)
+ elif "Column" in layer_type.__name__:
+ monkeypatch.setattr(ltx2_phase_adapter, "ColumnParallelLinear", layer_type)
+ else:
+ monkeypatch.setattr(ltx2_phase_adapter, "RowParallelLinear", layer_type)
+
+ wrapper = _wrap(layer, expected_a, expected_b, target_slice=target_slice)
+ torch.testing.assert_close(wrapper(torch.tensor([[1.0, 2.0, 3.0, 4.0]])), expected_output)
+ expected_local_a_width = getattr(layer, "input_size_per_partition", layer.input_size)
+ assert wrapper.pieces[0].lora_a.shape[-1] == expected_local_a_width
diff --git a/tests/diffusion/models/ltx2/test_ltx2_pipeline.py b/tests/diffusion/models/ltx2/test_ltx2_pipeline.py
index d70820f41..2f3ae522e 100644
--- a/tests/diffusion/models/ltx2/test_ltx2_pipeline.py
+++ b/tests/diffusion/models/ltx2/test_ltx2_pipeline.py
@@ -13,13 +13,21 @@ import pytest
import torch
from vllm_omni.diffusion.data import DiffusionOutput
-from vllm_omni.diffusion.models.ltx2 import ltx2_components
+from vllm_omni.diffusion.models.ltx2 import (
+ ltx2_components,
+ ltx2_transformer,
+)
from vllm_omni.diffusion.models.ltx2.ltx2_components import (
LTX2_COMPONENT_PROFILE,
LTX2_DISTILLED_COMPONENT_PROFILE,
+ LTX2_TWO_STAGE_COMPONENT_PROFILE,
LTX23_COMPONENT_PROFILE,
+ LTX23_DISTILLED_COMPONENT_PROFILE,
+ LTX23_TWO_STAGE_COMPONENT_PROFILE,
_install_connector_attention,
detect_ltx_model_version,
+ resolve_ltx_artifact,
+ resolve_ltx_component_profile,
)
from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import LTXI2VConditioningMixin
from vllm_omni.diffusion.models.ltx2.ltx2_denoise import (
@@ -33,6 +41,7 @@ from vllm_omni.diffusion.models.ltx2.ltx2_guidance import (
LTXGuidancePlan,
LTXGuidanceSpec,
LTXModalityGuidance,
+ _repeat_batch,
build_perturbation_kwargs,
combine_guided_x0,
)
@@ -40,18 +49,29 @@ 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,
+ LTX2_TWO_STAGE_RECIPE,
+ LTX23_DISTILLED_TWO_STAGE_RECIPE,
LTX23_ONE_STAGE_RECIPE,
+ LTX23_TWO_STAGE_RECIPE,
LTX_DEFAULT_NEGATIVE_PROMPT,
LTX_POSITIVE_ONLY_RECIPE,
+ resolve_ltx_pipeline_recipe,
+)
+from vllm_omni.diffusion.models.ltx2.ltx2_request import (
+ LTXRequestInputs,
+ validate_ltx_checkpoint,
+ validate_pipeline_request,
)
-from vllm_omni.diffusion.models.ltx2.ltx2_request import LTXRequestInputs, validate_pipeline_request
from vllm_omni.diffusion.models.ltx2.ltx2_runtime import LTXRuntime
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import (
LTX2I2VDMD2Pipeline,
LTX2Pipeline,
LTX2T2VDMD2Pipeline,
)
-from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import LTX2DistilledPipeline
+from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import (
+ LTX2DistilledPipeline,
+ LTX2TwoStagePipeline,
+)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -104,10 +124,22 @@ def test_ltx_public_entries_share_runtime_and_keep_recipe_boundaries():
)
assert issubclass(LTX2Pipeline, LTXRuntime)
+ assert issubclass(LTX2TwoStagePipeline, LTXRuntime)
assert issubclass(LTX2DistilledPipeline, LTXRuntime)
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 LTX2DistilledPipeline.component_profile is LTX2_DISTILLED_COMPONENT_PROFILE
+ for pipeline_cls, profile in (
+ (LTX2Pipeline, LTX2_COMPONENT_PROFILE),
+ (LTX2TwoStagePipeline, LTX2_TWO_STAGE_COMPONENT_PROFILE),
+ (LTX2DistilledPipeline, LTX2_DISTILLED_COMPONENT_PROFILE),
+ ):
+ assert pipeline_cls._dit_modules == list(profile.dit_modules)
+ assert pipeline_cls._encoder_modules == list(profile.encoder_modules)
+ assert pipeline_cls._vae_modules == list(profile.vae_modules)
+ assert pipeline_cls._resident_modules == list(profile.resident_modules)
assert LTX2_COMPONENT_PROFILE.video_vae_cls is DistributedAutoencoderKLLTX2Video
assert LTX23_COMPONENT_PROFILE.video_vae_cls is DistributedAutoencoderKLLTX2Video
@@ -123,6 +155,72 @@ def test_ltx_checkpoint_version_detection_uses_metadata(tmp_path):
assert detect_ltx_model_version(str(tmp_path)) == "2"
+def test_ltx_artifact_falls_back_to_unpinned_hub_download(tmp_path, monkeypatch):
+ filename = "ltx-sidecar.safetensors"
+ calls = []
+
+ def fake_download(**kwargs):
+ calls.append(kwargs)
+ return "/cache/ltx-sidecar.safetensors"
+
+ monkeypatch.setattr(ltx2_components, "hf_hub_download", fake_download)
+
+ assert (
+ resolve_ltx_artifact(str(tmp_path / "model"), "Lightricks/test", filename) == "/cache/ltx-sidecar.safetensors"
+ )
+ assert calls == [{"repo_id": "Lightricks/test", "filename": filename}]
+
+
+def test_ltx_artifact_prefers_model_root(tmp_path, monkeypatch):
+ filename = "ltx-sidecar.safetensors"
+ expected = tmp_path / filename
+ expected.write_bytes(b"sidecar")
+ monkeypatch.setattr(ltx2_components, "hf_hub_download", lambda **_kwargs: pytest.fail("unexpected Hub lookup"))
+
+ assert resolve_ltx_artifact(str(tmp_path), "Lightricks/test", filename) == str(expected)
+
+
+@pytest.mark.parametrize(
+ ("expected_kind", "scheduler_config"),
+ [
+ ("regular", {"use_dynamic_shifting": True, "shift_terminal": 0.1}),
+ ("distilled", {"use_dynamic_shifting": False, "shift_terminal": None}),
+ ],
+)
+def test_ltx_checkpoint_validation_accepts_matching_scheduler_metadata(expected_kind, scheduler_config):
+ validate_ltx_checkpoint(
+ scheduler_config,
+ expected_kind=expected_kind,
+ pipeline_name="TestLTXPipeline",
+ )
+
+
+@pytest.mark.parametrize(
+ ("expected_kind", "scheduler_config", "error"),
+ [
+ ("regular", {"use_dynamic_shifting": False, "shift_terminal": None}, "regular non-distilled"),
+ ("distilled", {"use_dynamic_shifting": True, "shift_terminal": 0.1}, "merged distilled"),
+ ],
+)
+def test_ltx_checkpoint_validation_rejects_mismatched_scheduler_metadata(expected_kind, scheduler_config, error):
+ with pytest.raises(ValueError, match=error):
+ validate_ltx_checkpoint(
+ scheduler_config,
+ expected_kind=expected_kind,
+ pipeline_name="TestLTXPipeline",
+ )
+
+
+def test_ltx_rejects_advanced_uaa_before_component_initialization():
+ od_config = SimpleNamespace(
+ model="unused",
+ parallel_config=SimpleNamespace(ulysses_mode="advanced_uaa"),
+ )
+
+ with pytest.raises(ValueError, match="does not support ulysses_mode='advanced_uaa'"):
+ LTX2Pipeline(od_config=od_config)
+
+
def test_ltx23_checkpoint_selects_version_specific_one_stage_profile(tmp_path, monkeypatch):
from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
@@ -144,16 +242,104 @@ def test_ltx23_checkpoint_selects_version_specific_one_stage_profile(tmp_path, m
assert pipe.reports_stage_durations
-def test_ltx_one_stage_entry_exposes_batch_image_and_distributed_decode_contracts():
+def test_ltx_multistage_rejects_cache_dit_before_component_initialization(tmp_path, monkeypatch):
+ from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
+
+ (tmp_path / "model_index.json").write_text(json.dumps({"vocoder": ["ltx2", "LTX2Vocoder"]}))
+ initialized = False
+
+ def stub_components(*_args, **_kwargs):
+ nonlocal initialized
+ initialized = True
+
+ monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)
+
+ with pytest.raises(ValueError, match="only one-stage LTX recipes"):
+ LTX2TwoStagePipeline(
+ od_config=SimpleNamespace(
+ model=str(tmp_path),
+ cache_backend="cache_dit",
+ )
+ )
+
+ assert not initialized
+
+
+def test_ltx23_checkpoint_selects_version_specific_two_stage_profiles(tmp_path, monkeypatch):
+ from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
+
+ (tmp_path / "model_index.json").write_text(json.dumps({"vocoder": ["ltx2", "LTX2VocoderWithBWE"]}))
+
+ 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)
+ monkeypatch.setattr(ltx2_runtime, "build_ltx_phase_adapter", lambda *_args: SimpleNamespace())
+ pipe = LTX2TwoStagePipeline(
+ od_config=SimpleNamespace(
+ model=str(tmp_path),
+ enable_diffusion_pipeline_profiler=False,
+ lora_path=None,
+ model_config={},
+ )
+ )
+
+ assert pipe.model_version == "2.3"
+ assert pipe.component_profile is LTX23_TWO_STAGE_COMPONENT_PROFILE
+ assert pipe.pipeline_recipe is LTX23_TWO_STAGE_RECIPE
+
+
+def test_ltx23_checkpoint_selects_full_distilled_profile_and_recipe(tmp_path, monkeypatch):
+ from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
+
+ (tmp_path / "model_index.json").write_text(json.dumps({"vocoder": ["ltx2", "LTX2VocoderWithBWE"]}))
+
+ 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 = LTX2DistilledPipeline(
+ od_config=SimpleNamespace(
+ model=str(tmp_path),
+ enable_diffusion_pipeline_profiler=False,
+ )
+ )
+
+ assert pipe.model_version == "2.3"
+ assert pipe.component_profile is LTX23_DISTILLED_COMPONENT_PROFILE
+ assert pipe.pipeline_recipe is LTX23_DISTILLED_TWO_STAGE_RECIPE
+ assert pipe._phase_adapter is None
+ assert pipe.preserve_sp_padded_audio_duration
+ assert pipe.reports_stage_durations
+
+
+def test_ltx_entries_expose_batch_image_and_distributed_decode_contracts():
assert LTX2Pipeline.supports_request_batch
assert LTX2Pipeline.support_image_input
assert LTX2Pipeline.unified_text_image_entry
assert LTX2Pipeline.distributed_video_decode
+ assert not LTX2TwoStagePipeline.supports_request_batch
+ assert LTX2TwoStagePipeline.support_image_input
+ assert LTX2TwoStagePipeline.unified_text_image_entry
assert not LTX2DistilledPipeline.supports_request_batch
assert LTX2DistilledPipeline.support_image_input
assert LTX2DistilledPipeline.unified_text_image_entry
+def test_ltx_two_stage_resolves_request_embedded_images():
+ pipe = object.__new__(LTX2TwoStagePipeline)
+ request_inputs = SimpleNamespace(latents=None)
+ image = torch.zeros(3, 8, 8)
+ req = SimpleNamespace(prompts=[{"prompt": "image conditioned", "multi_modal_data": {"image": image}}])
+
+ assert pipe._resolve_request_image(req, None, request_inputs) is image
+
+
def test_ltx_dmd2_entries_retain_their_task_boundaries():
assert not LTX2T2VDMD2Pipeline.support_image_input
assert LTX2I2VDMD2Pipeline.support_image_input
@@ -330,6 +516,15 @@ def test_ltx_guidance_uses_official_close_to_disabled_semantics():
assert LTXGuidancePlan.build(LTXGuidanceSpec(video=guidance)).names == ("cond",)
+def test_ltx_positive_only_guidance_preserves_token_major_layout():
+ tensor = torch.arange(24).view(1, 3, 8).transpose(1, 2).contiguous().transpose(1, 2)
+
+ repeated = _repeat_batch(tensor, 1)
+
+ assert repeated is tensor
+ assert repeated.stride() == tensor.stride()
+
+
def test_ltx_guidance_rescale_is_invariant_to_request_batching():
cond = torch.tensor(
[
@@ -380,6 +575,32 @@ def test_ltx_guidance_rescale_handles_zero_variance_prediction():
torch.testing.assert_close(actual, torch.zeros_like(cond))
+def test_ltx_audio_guidance_rescale_ignores_sp_padding_tokens():
+ logical_cond = torch.tensor([[[1.0, 2.0], [3.0, 5.0]]])
+ logical_uncond = torch.tensor([[[0.5, -1.0], [2.0, 1.0]]])
+ guidance = LTXModalityGuidance(cfg_scale=3.0, rescale_scale=1.0)
+ expected = combine_guided_x0(
+ cond=logical_cond,
+ uncond_text=logical_uncond,
+ uncond_perturbed=0.0,
+ uncond_modality=0.0,
+ guidance=guidance,
+ )
+ padded_cond = torch.cat([logical_cond, torch.tensor([[[1000.0, -1000.0]]])], dim=1)
+ padded_uncond = torch.cat([logical_uncond, torch.tensor([[[-500.0, 500.0]]])], dim=1)
+
+ actual = combine_guided_x0(
+ cond=padded_cond,
+ uncond_text=padded_uncond,
+ uncond_perturbed=0.0,
+ uncond_modality=0.0,
+ guidance=guidance,
+ rescale_token_count=2,
+ )
+
+ torch.testing.assert_close(actual[:, :2], expected)
+
+
def test_ltx_guidance_builds_per_sample_stg_and_modality_masks():
plan = LTXGuidancePlan.build(LTX23_ONE_STAGE_RECIPE.request_guidance)
perturbations = build_perturbation_kwargs(plan, batch_size=2, reference=torch.ones(8, 1, 1))
@@ -547,11 +768,19 @@ def test_ltx_connector_attention_dispatches_through_omni_kernel(has_learned_regi
def test_ltx2_distilled_pipeline_shares_recipe_runtime():
assert issubclass(LTX2DistilledPipeline, LTXI2VConditioningMixin)
- assert issubclass(LTX2DistilledPipeline, LTX2Pipeline)
+ assert not issubclass(LTX2DistilledPipeline, LTX2Pipeline)
assert LTX2DistilledPipeline._run_recipe is LTXRuntime._run_recipe
assert "_forward_request" not in LTX2DistilledPipeline.__dict__
+def test_ltx_two_stage_entry_is_a_thin_pipeline_variant():
+ assert "_run_recipe" not in LTX2TwoStagePipeline.__dict__
+ assert "_forward_request" not in LTX2TwoStagePipeline.__dict__
+ assert "load_weights" not in LTX2TwoStagePipeline.__dict__
+ assert issubclass(LTX2TwoStagePipeline, LTXI2VConditioningMixin)
+ assert not issubclass(LTX2TwoStagePipeline, LTX2Pipeline)
+
+
def test_ltx_variants_share_denoise_loop_and_i2v_conditioning():
assert "_denoise_loop" not in LTX2Pipeline.__dict__
assert issubclass(LTX2Pipeline, LTXI2VConditioningMixin)
@@ -594,7 +823,8 @@ def test_denoise_executor_owns_progress_and_interrupt():
torch.testing.assert_close(state.audio, torch.tensor(11.0))
-def test_ltx2_distilled_two_stage_executes_declarative_phase_plan():
+@pytest.mark.parametrize("pipeline_cls", [LTX2TwoStagePipeline, LTX2DistilledPipeline])
+def test_ltx_two_stage_executes_declarative_i2v_phase_plan(pipeline_cls):
request_inputs = LTXRequestInputs(
prompt="prompt",
negative_prompt="",
@@ -671,12 +901,15 @@ def test_ltx2_distilled_two_stage_executes_declarative_phase_plan():
return latents.repeat_interleave(2, dim=-2).repeat_interleave(2, dim=-1)
prompt_context_sentinel = prompt_context
- pipeline = object.__new__(LTX2DistilledPipeline)
+ pipeline = object.__new__(pipeline_cls)
torch.nn.Module.__init__(pipeline)
pipeline.device = torch.device("cpu")
pipeline.vae_spatial_compression_ratio = 32
pipeline.vae_temporal_compression_ratio = 8
pipeline.latent_upsampler = FakeUpsampler()
+ activated_slots = []
+ if pipeline_cls is LTX2TwoStagePipeline:
+ pipeline._phase_adapter = SimpleNamespace(activate=activated_slots.append)
object.__setattr__(pipeline, "_resolve_request_inputs", resolve_request_inputs)
object.__setattr__(pipeline, "run_phase", run_phase)
object.__setattr__(pipeline, "decode_phase", decode_phase)
@@ -686,14 +919,18 @@ def test_ltx2_distilled_two_stage_executes_declarative_phase_plan():
assert len(phase_calls) == 2
assert phase_calls[1][2] is prompt_context_sentinel
+ if pipeline_cls is LTX2TwoStagePipeline:
+ assert activated_slots == [None, "ltx_distilled"]
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))
+ expected_audio = 2.0 if pipeline_cls is LTX2TwoStagePipeline else 4.0
+ torch.testing.assert_close(output.output[1], torch.full((1, 8, 1, 2), expected_audio))
-def test_ltx2_distilled_stage2_reapplies_final_resolution_i2v_conditioning():
+@pytest.mark.parametrize("pipeline_cls", [LTX2TwoStagePipeline, LTX2DistilledPipeline])
+def test_ltx_two_stage_stage2_reapplies_final_resolution_i2v_conditioning(pipeline_cls):
from vllm_omni.diffusion.models.ltx2 import ltx2_latents
- pipeline = object.__new__(LTX2DistilledPipeline)
+ pipeline = object.__new__(pipeline_cls)
torch.nn.Module.__init__(pipeline)
pipeline.vae_spatial_compression_ratio = 32
pipeline.vae_temporal_compression_ratio = 8
@@ -750,6 +987,18 @@ def test_ltx2_distilled_two_stage_recipe_is_fixed_positive_only():
assert (LTX2_DISTILLED_TWO_STAGE_RECIPE.height, LTX2_DISTILLED_TWO_STAGE_RECIPE.width) == (1024, 1536)
+def test_ltx23_full_distilled_route_uses_merged_weights_without_lora():
+ assert resolve_ltx_component_profile("distilled_two_stage", "2.3") is LTX23_DISTILLED_COMPONENT_PROFILE
+ assert resolve_ltx_pipeline_recipe("distilled_two_stage", "2.3") is LTX23_DISTILLED_TWO_STAGE_RECIPE
+ assert LTX23_DISTILLED_TWO_STAGE_RECIPE is LTX2_DISTILLED_TWO_STAGE_RECIPE
+ assert all(phase.adapter_slot is None for phase in LTX23_DISTILLED_TWO_STAGE_RECIPE.phases)
+ assert LTX23_DISTILLED_COMPONENT_PROFILE.distilled_lora_filename is None
+ assert LTX23_DISTILLED_COMPONENT_PROFILE.vocoder_cls is LTX23_COMPONENT_PROFILE.vocoder_cls
+ assert LTX23_DISTILLED_COMPONENT_PROFILE.latent_upsampler_filename == (
+ "ltx-2.3-spatial-upscaler-x2-1.1.safetensors"
+ )
+
+
@pytest.mark.parametrize(
("request_inputs", "error"),
[
@@ -958,6 +1207,98 @@ def test_ltx_custom_sigmas_bypass_scheduler_shifting():
torch.testing.assert_close(timesteps, expected[:-1] * scheduler.config.num_train_timesteps)
+@pytest.mark.parametrize(
+ ("recipe", "steps", "stg_block"),
+ [
+ (LTX2_TWO_STAGE_RECIPE, 40, 29),
+ (LTX23_TWO_STAGE_RECIPE, 30, 28),
+ ],
+)
+def test_ltx_official_two_stage_recipes_only_vary_by_model_defaults(recipe, steps, stg_block):
+ stage1, stage2 = recipe.phases
+ assert (recipe.width, recipe.height) == (1536, 1024)
+ assert recipe.num_inference_steps == steps
+ assert stage1.spatial_downscale == 2
+ assert stage1.adapter_slot is None
+ assert stage1.guidance.video.stg_blocks == (stg_block,)
+ assert stage1.guidance.audio.stg_blocks == (stg_block,)
+ assert stage2.input_transform == "spatial_upsample"
+ assert stage2.adapter_slot == "ltx_distilled"
+ assert stage2.guidance == LTXGuidanceSpec.positive_only()
+ assert stage2.num_inference_steps == 3
+ assert stage2.sigmas == (0.909375, 0.725, 0.421875, 0.0)
+ # Official Stage 2 refines video only; its audio result is discarded.
+ assert (recipe.video_output_phase, recipe.audio_output_phase) == (1, 0)
+
+
+def test_ltx_two_stage_entries_select_the_official_adapter_slots():
+ phase_switches = []
+ ordinary_pipeline = object.__new__(LTX2TwoStagePipeline)
+ ordinary_pipeline._phase_adapter = SimpleNamespace(activate=phase_switches.append)
+ for phase in LTX2_TWO_STAGE_RECIPE.phases:
+ ordinary_pipeline._enter_phase(phase)
+ assert phase_switches == [None, "ltx_distilled"]
+
+ distilled_pipeline = object.__new__(LTX2DistilledPipeline)
+ for phase in LTX2_DISTILLED_TWO_STAGE_RECIPE.phases:
+ distilled_pipeline._enter_phase(phase)
+
+
+def test_ltx23_gate_projection_declares_its_global_column_parallel_layout(monkeypatch):
+ class FakeColumn(torch.nn.Module):
+ def __init__(
+ self,
+ input_size,
+ output_size,
+ *,
+ bias,
+ gather_output,
+ return_bias,
+ quant_config,
+ prefix,
+ ):
+ super().__init__()
+ self.input_size = input_size
+ self.output_size = output_size
+ self.bias_enabled = bias
+ self.gather_output = gather_output
+ self.return_bias = return_bias
+ self.quant_config = quant_config
+ self.prefix = prefix
+
+ class FakeRow(torch.nn.Module):
+ def __init__(self, *args, **kwargs):
+ super().__init__()
+
+ class FakeAttention(torch.nn.Module):
+ def __init__(self, **kwargs):
+ super().__init__()
+
+ monkeypatch.setattr(ltx2_transformer, "ColumnParallelLinear", FakeColumn)
+ monkeypatch.setattr(ltx2_transformer, "RowParallelLinear", FakeRow)
+ monkeypatch.setattr(ltx2_transformer, "_LTX2ParallelAttention", FakeAttention)
+ monkeypatch.setattr(ltx2_transformer, "get_tensor_model_parallel_world_size", lambda: 2)
+
+ attention = ltx2_transformer.LTX2Attention(
+ query_dim=16,
+ heads=4,
+ dim_head=4,
+ qk_norm="rms_norm",
+ apply_gated_attention=True,
+ pack_qkv=False,
+ prefix="transformer_blocks.0.attn1",
+ )
+
+ gate = attention.to_gate_logits
+ assert isinstance(gate, FakeColumn)
+ assert gate.input_size == 16
+ assert gate.output_size == 4
+ assert gate.gather_output is False
+ assert gate.return_bias is False
+ assert gate.quant_config is None
+ assert gate.prefix == "transformer_blocks.0.attn1.to_gate_logits"
+
+
class TestLTXRequestParsing:
def test_request_resolves_official_and_explicit_negative_prompts_per_sample(self):
from vllm_omni.diffusion.request import OmniDiffusionRequest
@@ -1432,6 +1773,7 @@ class TestPipelineComponents:
def test_distilled_profile_keeps_upsampler_in_managed_resident_components(self):
assert LTX2_DISTILLED_COMPONENT_PROFILE.resident_modules == ("vocoder", "latent_upsampler")
+ assert LTX23_DISTILLED_COMPONENT_PROFILE.resident_modules == ("vocoder", "latent_upsampler")
class TestLTX23DecodeConditioning:
@@ -1477,6 +1819,11 @@ class TestRegistryIntegration:
from vllm_omni.diffusion.registry import _DIFFUSION_MODELS
assert _DIFFUSION_MODELS["LTX2Pipeline"] == ("ltx2", "pipeline_ltx2", "LTX2Pipeline")
+ assert _DIFFUSION_MODELS["LTX2TwoStagePipeline"] == (
+ "ltx2",
+ "pipeline_ltx2_two_stage",
+ "LTX2TwoStagePipeline",
+ )
assert _DIFFUSION_MODELS["LTX2DistilledPipeline"] == (
"ltx2",
"pipeline_ltx2_two_stage",
@@ -1488,6 +1835,9 @@ class TestRegistryIntegration:
"LTX23ImageToVideoPipeline",
"LTX2TwoStagesPipeline",
"LTX2ImageToVideoTwoStagesPipeline",
+ "LTX2ImageToVideoTwoStagePipeline",
+ "LTX23TwoStagePipeline",
+ "LTX23ImageToVideoTwoStagePipeline",
}
assert removed_entries.isdisjoint(_DIFFUSION_MODELS)
@@ -1497,6 +1847,7 @@ class TestRegistryIntegration:
expected = [
"LTX2Pipeline",
+ "LTX2TwoStagePipeline",
"LTX2DistilledPipeline",
"LTX2T2VDMD2Pipeline",
"LTX2I2VDMD2Pipeline",
@@ -1509,6 +1860,7 @@ class TestRegistryIntegration:
"model_class_name",
[
"LTX2Pipeline",
+ "LTX2TwoStagePipeline",
"LTX2DistilledPipeline",
"LTX2T2VDMD2Pipeline",
"LTX2I2VDMD2Pipeline",
diff --git a/tests/diffusion/models/ltx2/test_ltx2_transformer.py b/tests/diffusion/models/ltx2/test_ltx2_transformer.py
index 73e10463e..f74b2991d 100644
--- a/tests/diffusion/models/ltx2/test_ltx2_transformer.py
+++ b/tests/diffusion/models/ltx2/test_ltx2_transformer.py
@@ -1,13 +1,20 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from types import SimpleNamespace
+
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.ltx2_transformer import LTX2VideoTransformer3DModel, _make_rms_norm
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+from vllm_omni.diffusion.models.ltx2.ltx2_transformer import (
+ LTX2AudioVideoAttnProcessor,
+ LTX2VideoTransformer3DModel,
+ _LTX2ParallelAttention,
+ _make_rms_norm,
+)
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@@ -66,3 +73,83 @@ def test_ltx_sp_plan_shards_video_and_audio_timesteps_together(rope_type):
assert video_timestep.expected_dims == audio_timestep.expected_dims == 2
assert not video_timestep.split_output
assert not audio_timestep.split_output
+
+
+def test_ltx_sp_keeps_cross_attention_key_padding_mask(monkeypatch):
+ monkeypatch.setattr(LTX2AudioVideoAttnProcessor, "_is_sp_enabled", staticmethod(lambda: True))
+ processor = LTX2AudioVideoAttnProcessor()
+ hidden_states = torch.zeros(1, 3, 4)
+ encoder_hidden_states = torch.zeros(1, 4, 4)
+ key_padding_mask = torch.tensor([[True, True, True, False]])
+
+ actual = processor._prepare_attention_mask(
+ SimpleNamespace(),
+ hidden_states,
+ encoder_hidden_states,
+ key_padding_mask,
+ batch_size=1,
+ sequence_length=4,
+ )
+
+ torch.testing.assert_close(actual, key_padding_mask[:, None, None, :])
+
+
+def _fake_ltx_parallel_attention():
+ calls = []
+
+ class Backend:
+ @staticmethod
+ def get_name():
+ return "FLASH_ATTN"
+
+ @staticmethod
+ def supports_attention_mask():
+ return True
+
+ def run(name):
+ def forward(query, _key, _value, _metadata):
+ calls.append(name)
+ return torch.zeros_like(query)
+
+ return forward
+
+ layer = object.__new__(_LTX2ParallelAttention)
+ nn.Module.__init__(layer)
+ layer.attn_backend = Backend
+ layer.attention = SimpleNamespace(forward=run("native"))
+ layer.sdpa_fallback = SimpleNamespace(forward=run("sdpa"))
+ layer.backend_pref = "FLASH_ATTN"
+ return layer, calls
+
+
+@pytest.mark.parametrize(("query_length", "key_length"), [(4, 4), (3, 4)])
+def test_ltx_masked_cross_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, :])
+
+ 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():
+ 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"]
diff --git a/tests/e2e/accuracy/ltx/run_ltx_reference.py b/tests/e2e/accuracy/ltx/run_ltx_reference.py
index a4a180047..3ab794662 100644
--- a/tests/e2e/accuracy/ltx/run_ltx_reference.py
+++ b/tests/e2e/accuracy/ltx/run_ltx_reference.py
@@ -20,6 +20,11 @@ from PIL import Image
def _parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser()
parser.add_argument("--backend", choices=("official", "omni"), required=True)
+ parser.add_argument(
+ "--pipeline-kind",
+ choices=("one_stage", "distilled", "two_stage"),
+ default="one_stage",
+ )
parser.add_argument("--request", type=Path, required=True)
parser.add_argument("--output-dir", type=Path, required=True)
parser.add_argument("--model")
@@ -27,6 +32,8 @@ def _parse_args() -> argparse.Namespace:
parser.add_argument("--official-root", type=Path)
parser.add_argument("--checkpoint", type=Path)
parser.add_argument("--gemma-root", type=Path)
+ parser.add_argument("--spatial-upsampler", type=Path)
+ parser.add_argument("--distilled-lora", type=Path)
parser.add_argument("--enable-layerwise-offload", action="store_true")
return parser.parse_args()
@@ -91,9 +98,13 @@ def _configure_official_sdpa(pipeline: Any) -> None:
attention=attention,
masked_attention=attention,
)
- owners_and_attributes = (
- (pipeline.stage, "_transformer_builder"),
+ owners_and_attributes = [
(pipeline.prompt_encoder, "_embeddings_processor_builder"),
+ ]
+ owners_and_attributes.extend(
+ (stage, "_transformer_builder")
+ for name in ("stage", "stage_1", "stage_2")
+ if (stage := getattr(pipeline, name, None)) is not None
)
for owner, attribute in owners_and_attributes:
builder = getattr(owner, attribute)
@@ -104,6 +115,12 @@ def _configure_official_sdpa(pipeline: Any) -> None:
)
+def _require_path(path: Path | None, description: str) -> str:
+ if path is None or not path.is_file():
+ raise ValueError(f"Official {description} is required and must exist: {path}")
+ return str(path)
+
+
@torch.inference_mode()
def _run_official(args: argparse.Namespace, request: dict[str, Any]) -> None:
if args.official_root is None or args.checkpoint is None or args.gemma_root is None:
@@ -113,16 +130,54 @@ def _run_official(args: argparse.Namespace, request: dict[str, Any]) -> None:
torch.backends.cuda.enable_cudnn_sdp(False)
from ltx_core.components.guiders import MultiModalGuiderParams
+ from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps
+ from ltx_pipelines.distilled import DistilledPipeline
from ltx_pipelines.ti2vid_one_stage import TI2VidOneStagePipeline
+ from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline
from ltx_pipelines.utils.args import ImageConditioningInput
from ltx_pipelines.utils.types import OffloadMode
- pipeline = TI2VidOneStagePipeline(
- checkpoint_path=str(args.checkpoint),
- gemma_root=str(args.gemma_root),
- loras=(),
- offload_mode=OffloadMode.CPU if args.enable_layerwise_offload else OffloadMode.NONE,
- )
+ if request.get("pipeline_kind", "one_stage") != args.pipeline_kind:
+ raise ValueError(
+ f"Request pipeline kind {request.get('pipeline_kind')!r} does not match {args.pipeline_kind!r}"
+ )
+ checkpoint = _require_path(args.checkpoint, "checkpoint")
+ gemma_root = str(args.gemma_root)
+ if not Path(gemma_root).is_dir():
+ raise ValueError(f"Official Gemma root is required and must be a directory: {gemma_root}")
+ offload_mode = OffloadMode.CPU if args.enable_layerwise_offload else OffloadMode.NONE
+ pipeline: Any
+ if args.pipeline_kind == "one_stage":
+ pipeline = TI2VidOneStagePipeline(
+ checkpoint_path=checkpoint,
+ gemma_root=gemma_root,
+ loras=(),
+ offload_mode=offload_mode,
+ )
+ elif args.pipeline_kind == "distilled":
+ pipeline = DistilledPipeline(
+ distilled_checkpoint_path=checkpoint,
+ spatial_upsampler_path=_require_path(args.spatial_upsampler, "spatial upsampler"),
+ gemma_root=gemma_root,
+ loras=(),
+ offload_mode=offload_mode,
+ )
+ else:
+ distilled_lora = [
+ LoraPathStrengthAndSDOps(
+ path=_require_path(args.distilled_lora, "distilled LoRA"),
+ strength=1.0,
+ sd_ops=LTXV_LORA_COMFY_RENAMING_MAP,
+ )
+ ]
+ pipeline = TI2VidTwoStagesPipeline(
+ checkpoint_path=checkpoint,
+ distilled_lora=distilled_lora,
+ spatial_upsampler_path=_require_path(args.spatial_upsampler, "spatial upsampler"),
+ gemma_root=gemma_root,
+ loras=(),
+ offload_mode=offload_mode,
+ )
_configure_official_sdpa(pipeline)
image_path = request.get("image")
images = (
@@ -140,34 +195,45 @@ def _run_official(args: argparse.Namespace, request: dict[str, Any]) -> None:
)
]
)
- 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,
- )
+ if args.pipeline_kind == "distilled":
+ 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,
+ )
+ else:
+ 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,
+ )
video_tensor = torch.cat([chunk.detach().cpu() for chunk in video], dim=0)
_save_outputs(
args.output_dir,
@@ -179,6 +245,7 @@ def _run_official(args: argparse.Namespace, request: dict[str, Any]) -> None:
"attention_backend": "torch_sdpa",
"official_revision": os.environ.get("VLLM_TEST_LTX_OFFICIAL_REVISION"),
"checkpoint": str(args.checkpoint),
+ "pipeline_kind": args.pipeline_kind,
},
)
@@ -253,6 +320,10 @@ def _canonical_video(video: Any) -> torch.Tensor:
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")
+ if request.get("pipeline_kind", "one_stage") != args.pipeline_kind:
+ raise ValueError(
+ f"Request pipeline kind {request.get('pipeline_kind')!r} does not match {args.pipeline_kind!r}"
+ )
attention_config = {"default": {"backend": "TORCH_SDPA"}}
@@ -273,7 +344,8 @@ def _run_omni(args: argparse.Namespace, request: dict[str, Any]) -> None:
parallel_config=DiffusionParallelConfig(),
)
try:
- model_class_name = get_model_class_name(omni)
+ detected_model_class_name = get_model_class_name(omni)
+ model_class_name = args.model_class_name
sampling_params = OmniDiffusionSamplingParams(
height=request["height"],
width=request["width"],
@@ -289,10 +361,9 @@ def _run_omni(args: argparse.Namespace, request: dict[str, Any]) -> None:
key: value for key, value in request.items() if key.startswith("video_") or key.startswith("audio_")
}
apply_declared_extra_args(sampling_params, get_extra_body_params(model_class_name), guidance)
- prompt: dict[str, Any] = {
- "prompt": request["prompt"],
- "negative_prompt": request["negative_prompt"],
- }
+ prompt: dict[str, Any] = {"prompt": request["prompt"]}
+ if request.get("negative_prompt"):
+ 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:
@@ -312,6 +383,8 @@ def _run_omni(args: argparse.Namespace, request: dict[str, Any]) -> None:
"attention_backend": "torch_sdpa",
"model": args.model,
"model_class_name": model_class_name,
+ "model_config_class_name": detected_model_class_name,
+ "pipeline_kind": args.pipeline_kind,
},
)
finally:
diff --git a/tests/e2e/accuracy/ltx/test_ltx_official_similarity.py b/tests/e2e/accuracy/ltx/test_ltx_official_similarity.py
index 647ece7a7..8a3be4a99 100644
--- a/tests/e2e/accuracy/ltx/test_ltx_official_similarity.py
+++ b/tests/e2e/accuracy/ltx/test_ltx_official_similarity.py
@@ -3,10 +3,9 @@
"""E2E accuracy guard against a pinned Lightricks LTX pipeline revision.
-The 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.
+The original reduced one-stage guards remain unchanged. Four complementary
+default-shape cases cover both model versions, both two-stage weight families,
+and T2V/I2V without expanding every combination into a duplicate golden.
"""
from __future__ import annotations
@@ -17,8 +16,9 @@ import shutil
import subprocess
import sys
import time
-from dataclasses import dataclass
+from dataclasses import asdict, dataclass, replace
from pathlib import Path
+from typing import Literal
import numpy as np
import pytest
@@ -53,74 +53,260 @@ NEGATIVE_PROMPT = (
# Both runtimes use PyTorch SDPA with the current Torch dispatch defaults.
ATTENTION_BACKEND = "torch_sdpa"
-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
+
+
+@dataclass(frozen=True)
+class LTXArtifact:
+ repo_id: str
+ filename: str
+ revision: str
+ env: str
+ repo_type: Literal["model", "dataset"] = "model"
+
+
+@dataclass(frozen=True)
+class LTXAccuracyThresholds:
+ video_ssim_mean: float
+ video_ssim_min: float
+ video_psnr_mean_db: float
+ # Zero disables gating for waveform-sensitive audio that is not expected
+ # to align with the official output.
+ audio_relative_l2: float
+ audio_cosine_similarity: float
+
+
+STRICT_THRESHOLDS = LTXAccuracyThresholds(
+ video_ssim_mean=0.95,
+ video_ssim_min=0.90,
+ video_psnr_mean_db=30.0,
+ audio_relative_l2=0.20,
+ audio_cosine_similarity=0.95,
+)
@dataclass(frozen=True)
class LTXAccuracyCase:
name: str
+ pipeline_kind: Literal["one_stage", "distilled", "two_stage"]
model_id: str
model_revision: str
model_env: str
model_class_name: str
- checkpoint_repo: str
- checkpoint_filename: str
- checkpoint_revision: str
- checkpoint_env: str
- stg_block: int
+ checkpoint: LTXArtifact
+ width: int
+ height: int
+ num_frames: int
+ num_inference_steps: int
+ seed: int
+ stg_block: int | None
+ thresholds: LTXAccuracyThresholds
prompt: str = PROMPT
- image_repo: str | None = None
- image_filename: str | None = None
- image_revision: str | None = None
+ gemma_model_id: str | None = None
+ gemma_model_revision: str | None = None
+ gemma_model_env: str | None = None
+ image: LTXArtifact | None = None
+ spatial_upsampler: LTXArtifact | None = None
+ distilled_lora: LTXArtifact | None = None
+ # Official uses block streaming; Omni uses layerwise offload. Both preserve
+ # bf16 arithmetic while keeping full-resolution two-stage cases on one H100.
+ enable_layerwise_offload: bool = False
-CASES = (
+LTX2_REVISION = "47da56e2ad66ce4125a9922b4a8826bf407f9d0a"
+LTX23_REVISION = "4229404625088d21c4f112eb640fb04a0900ee25"
+LTX2_CHECKPOINT = LTXArtifact(
+ repo_id="Lightricks/LTX-2",
+ filename="ltx-2-19b-dev.safetensors",
+ revision=LTX2_REVISION,
+ env="VLLM_TEST_LTX2_OFFICIAL_CHECKPOINT",
+)
+LTX2_DISTILLED_CHECKPOINT = replace(
+ LTX2_CHECKPOINT,
+ filename="ltx-2-19b-distilled.safetensors",
+ env="VLLM_TEST_LTX2_DISTILLED_OFFICIAL_CHECKPOINT",
+)
+LTX2_UPSAMPLER = replace(
+ LTX2_CHECKPOINT,
+ filename="ltx-2-spatial-upscaler-x2-1.0.safetensors",
+ env="VLLM_TEST_LTX2_UPSAMPLER",
+)
+LTX2_DISTILLED_LORA = replace(
+ LTX2_CHECKPOINT,
+ filename="ltx-2-19b-distilled-lora-384.safetensors",
+ env="VLLM_TEST_LTX2_DISTILLED_LORA",
+)
+LTX23_CHECKPOINT = LTXArtifact(
+ repo_id="Lightricks/LTX-2.3",
+ filename="ltx-2.3-22b-dev.safetensors",
+ revision=LTX23_REVISION,
+ env="VLLM_TEST_LTX23_OFFICIAL_CHECKPOINT",
+)
+LTX23_DISTILLED_CHECKPOINT = replace(
+ LTX23_CHECKPOINT,
+ # The pinned Diffusers checkpoint contains the original merged distilled
+ # Transformer. Keep the reference Transformer on the same version; the
+ # spatial upsampler is versioned independently and remains on 1.1 below.
+ filename="ltx-2.3-22b-distilled.safetensors",
+ env="VLLM_TEST_LTX23_DISTILLED_OFFICIAL_CHECKPOINT",
+)
+LTX23_UPSAMPLER = replace(
+ LTX23_CHECKPOINT,
+ filename="ltx-2.3-spatial-upscaler-x2-1.1.safetensors",
+ env="VLLM_TEST_LTX23_UPSAMPLER",
+)
+LTX23_DISTILLED_LORA = replace(
+ LTX23_CHECKPOINT,
+ filename="ltx-2.3-22b-distilled-lora-384-1.1.safetensors",
+ env="VLLM_TEST_LTX23_DISTILLED_LORA",
+)
+I2V_IMAGE = LTXArtifact(
+ repo_id="huggingface/documentation-images",
+ filename="diffusers/svd/rocket.png",
+ revision="645d8364f0c7a101180b364811b5a11a362e4010",
+ env="VLLM_TEST_LTX_I2V_IMAGE",
+ repo_type="dataset",
+)
+
+
+LEGACY_CASES = (
LTXAccuracyCase(
name="ltx2",
+ pipeline_kind="one_stage",
model_id="Lightricks/LTX-2",
- model_revision="47da56e2ad66ce4125a9922b4a8826bf407f9d0a",
+ model_revision=LTX2_REVISION,
model_env="VLLM_TEST_LTX2_MODEL",
model_class_name="LTX2Pipeline",
- checkpoint_repo="Lightricks/LTX-2",
- checkpoint_filename="ltx-2-19b-dev.safetensors",
- checkpoint_revision="47da56e2ad66ce4125a9922b4a8826bf407f9d0a",
- checkpoint_env="VLLM_TEST_LTX2_OFFICIAL_CHECKPOINT",
+ checkpoint=LTX2_CHECKPOINT,
+ width=512,
+ height=384,
+ num_frames=25,
+ num_inference_steps=20,
+ seed=42,
stg_block=29,
+ thresholds=STRICT_THRESHOLDS,
),
LTXAccuracyCase(
name="ltx2_3",
+ pipeline_kind="one_stage",
model_id="diffusers/LTX-2.3-Diffusers",
model_revision="8eee8edcf067e838b843f926ec4d4cc9b2be1aaf",
model_env="VLLM_TEST_LTX23_MODEL",
model_class_name="LTX2Pipeline",
- checkpoint_repo="Lightricks/LTX-2.3",
- checkpoint_filename="ltx-2.3-22b-dev.safetensors",
- checkpoint_revision="4229404625088d21c4f112eb640fb04a0900ee25",
- checkpoint_env="VLLM_TEST_LTX23_OFFICIAL_CHECKPOINT",
+ checkpoint=LTX23_CHECKPOINT,
+ width=512,
+ height=384,
+ num_frames=25,
+ num_inference_steps=20,
+ seed=42,
stg_block=28,
+ thresholds=STRICT_THRESHOLDS,
),
LTXAccuracyCase(
name="ltx2_3_i2v",
+ pipeline_kind="one_stage",
model_id="diffusers/LTX-2.3-Diffusers",
model_revision="8eee8edcf067e838b843f926ec4d4cc9b2be1aaf",
model_env="VLLM_TEST_LTX23_MODEL",
model_class_name="LTX2Pipeline",
- checkpoint_repo="Lightricks/LTX-2.3",
- checkpoint_filename="ltx-2.3-22b-dev.safetensors",
- checkpoint_revision="4229404625088d21c4f112eb640fb04a0900ee25",
- checkpoint_env="VLLM_TEST_LTX23_OFFICIAL_CHECKPOINT",
+ checkpoint=LTX23_CHECKPOINT,
+ width=512,
+ height=384,
+ num_frames=25,
+ num_inference_steps=20,
+ seed=42,
stg_block=28,
- image_repo="huggingface/documentation-images",
- image_filename="diffusers/svd/rocket.png",
- image_revision="645d8364f0c7a101180b364811b5a11a362e4010",
+ thresholds=STRICT_THRESHOLDS,
+ image=I2V_IMAGE,
),
)
+TWO_STAGE_CASES = (
+ LTXAccuracyCase(
+ name="ltx2_distilled_t2v",
+ pipeline_kind="distilled",
+ model_id="rootonchair/LTX-2-19b-distilled",
+ model_revision="388e2846f54aae51687498ffb6b27c7c2c9ce9e5",
+ model_env="VLLM_TEST_LTX2_DISTILLED_MODEL",
+ model_class_name="LTX2DistilledPipeline",
+ checkpoint=LTX2_DISTILLED_CHECKPOINT,
+ spatial_upsampler=LTX2_UPSAMPLER,
+ width=1536,
+ height=1024,
+ num_frames=121,
+ num_inference_steps=8,
+ seed=10,
+ stg_block=None,
+ gemma_model_id="Lightricks/LTX-2",
+ gemma_model_revision=LTX2_REVISION,
+ gemma_model_env="VLLM_TEST_LTX2_MODEL",
+ thresholds=STRICT_THRESHOLDS,
+ ),
+ LTXAccuracyCase(
+ name="ltx23_distilled_i2v",
+ pipeline_kind="distilled",
+ model_id="diffusers/LTX-2.3-Distilled-Diffusers",
+ model_revision="432e0d3c2d1769aaa4d295f9243f7062bf6b47ee",
+ model_env="VLLM_TEST_LTX23_DISTILLED_MODEL",
+ model_class_name="LTX2DistilledPipeline",
+ checkpoint=LTX23_DISTILLED_CHECKPOINT,
+ spatial_upsampler=LTX23_UPSAMPLER,
+ width=1536,
+ height=1024,
+ num_frames=121,
+ num_inference_steps=8,
+ seed=10,
+ stg_block=None,
+ gemma_model_id="diffusers/LTX-2.3-Diffusers",
+ gemma_model_revision="8eee8edcf067e838b843f926ec4d4cc9b2be1aaf",
+ gemma_model_env="VLLM_TEST_LTX23_MODEL",
+ thresholds=STRICT_THRESHOLDS,
+ image=I2V_IMAGE,
+ ),
+ LTXAccuracyCase(
+ name="ltx2_two_stage_layer_fused_t2v",
+ pipeline_kind="two_stage",
+ model_id="Lightricks/LTX-2",
+ model_revision=LTX2_REVISION,
+ model_env="VLLM_TEST_LTX2_MODEL",
+ model_class_name="LTX2TwoStagePipeline",
+ checkpoint=LTX2_CHECKPOINT,
+ spatial_upsampler=LTX2_UPSAMPLER,
+ distilled_lora=LTX2_DISTILLED_LORA,
+ width=1536,
+ height=1024,
+ num_frames=121,
+ num_inference_steps=40,
+ seed=10,
+ stg_block=29,
+ enable_layerwise_offload=True,
+ thresholds=STRICT_THRESHOLDS,
+ ),
+ LTXAccuracyCase(
+ name="ltx23_two_stage_layer_fused_i2v",
+ pipeline_kind="two_stage",
+ model_id="diffusers/LTX-2.3-Diffusers",
+ model_revision="8eee8edcf067e838b843f926ec4d4cc9b2be1aaf",
+ model_env="VLLM_TEST_LTX23_MODEL",
+ model_class_name="LTX2TwoStagePipeline",
+ checkpoint=LTX23_CHECKPOINT,
+ spatial_upsampler=LTX23_UPSAMPLER,
+ distilled_lora=LTX23_DISTILLED_LORA,
+ width=1536,
+ height=1024,
+ num_frames=121,
+ num_inference_steps=30,
+ seed=10,
+ stg_block=28,
+ enable_layerwise_offload=True,
+ thresholds=STRICT_THRESHOLDS,
+ image=I2V_IMAGE,
+ ),
+)
+
+CASES = (*LEGACY_CASES, *TWO_STAGE_CASES)
+
+
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)
@@ -209,6 +395,7 @@ def _resolve_model(case: LTXAccuracyCase) -> Path:
"model_index.json",
"audio_vae/*",
"connectors/*",
+ "latent_upsampler/*",
"processor/*",
"scheduler/*",
"text_encoder/config.json",
@@ -223,74 +410,82 @@ def _resolve_model(case: LTXAccuracyCase) -> Path:
)
-def _resolve_gemma_root(model: Path) -> Path:
+def _resolve_gemma_root(case: LTXAccuracyCase, 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
+ if case.gemma_model_id is not None:
+ configured_model = os.environ.get(case.gemma_model_env or "")
+ if configured_model and Path(configured_model).is_dir():
+ return Path(configured_model)
+ return Path(
+ snapshot_download(
+ repo_id=case.gemma_model_id,
+ revision=case.gemma_model_revision,
+ allow_patterns=[
+ "processor/*",
+ "text_encoder/*",
+ "tokenizer/*",
+ ],
+ )
+ )
return model
-def _resolve_checkpoint(case: LTXAccuracyCase, model: Path) -> Path:
- configured_checkpoint = os.environ.get(case.checkpoint_env)
- if configured_checkpoint:
- checkpoint = Path(configured_checkpoint)
- assert checkpoint.is_file(), f"Official checkpoint not found: {checkpoint}"
- return checkpoint
- model_checkpoint = model / case.checkpoint_filename
- if model_checkpoint.is_file():
- return model_checkpoint
+def _resolve_artifact(artifact: LTXArtifact, model: Path | None = None) -> Path:
+ configured_path = os.environ.get(artifact.env)
+ if configured_path:
+ path = Path(configured_path)
+ assert path.is_file(), f"Configured LTX artifact not found: {path}"
+ return path
+ if model is not None:
+ model_path = model / artifact.filename
+ if model_path.is_file():
+ return model_path
return Path(
hf_hub_download(
- repo_id=case.checkpoint_repo,
- filename=case.checkpoint_filename,
- revision=case.checkpoint_revision,
+ repo_id=artifact.repo_id,
+ repo_type=None if artifact.repo_type == "model" else artifact.repo_type,
+ filename=artifact.filename,
+ revision=artifact.revision,
)
)
def _resolve_image(case: LTXAccuracyCase) -> Path | None:
- if case.image_filename is None:
+ if case.image is None:
return None
- if case.image_repo is None or case.image_revision is None:
- raise ValueError(f"Incomplete image source for LTX accuracy case {case.name!r}.")
- configured_image = os.environ.get("VLLM_TEST_LTX_I2V_IMAGE")
- if configured_image:
- image = Path(configured_image)
- assert image.is_file(), f"LTX I2V conditioning image not found: {image}"
- return image
- return Path(
- hf_hub_download(
- repo_id=case.image_repo,
- repo_type="dataset",
- filename=case.image_filename,
- revision=case.image_revision,
- )
- )
+ return _resolve_artifact(case.image)
def _request(case: LTXAccuracyCase, image: Path | None) -> dict[str, object]:
request: dict[str, object] = {
+ "pipeline_kind": case.pipeline_kind,
"prompt": case.prompt,
- "negative_prompt": NEGATIVE_PROMPT,
- "width": 512,
- "height": 384,
- "num_frames": 25,
+ "negative_prompt": "" if case.pipeline_kind == "distilled" else NEGATIVE_PROMPT,
+ "width": case.width,
+ "height": case.height,
+ "num_frames": case.num_frames,
"fps": 24,
- "num_inference_steps": 20,
- "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": [case.stg_block],
- "audio_stg_blocks": [case.stg_block],
+ "num_inference_steps": case.num_inference_steps,
+ "seed": case.seed,
}
+ if case.pipeline_kind != "distilled":
+ assert case.stg_block is not None
+ request.update(
+ 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=[case.stg_block],
+ audio_stg_blocks=[case.stg_block],
+ )
if image is not None:
request["image"] = str(image.resolve())
return request
@@ -299,20 +494,30 @@ def _request(case: LTXAccuracyCase, image: Path | None) -> dict[str, object]:
def _video_metrics(reference: np.ndarray, prediction: np.ndarray) -> dict[str, float]:
assert reference.shape == prediction.shape
assert reference.ndim == 4 and reference.shape[-1] == 3
+ device = torch.device("cpu")
+ ssim = StructuralSimilarityIndexMeasure(data_range=1.0).to(device)
+ psnr = PeakSignalNoiseRatio(data_range=1.0).to(device)
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))
+ max_abs = 0.0
+ absolute_error_sum = 0.0
+ with torch.inference_mode():
+ for reference_frame, prediction_frame in zip(reference, prediction, strict=True):
+ reference_tensor = torch.from_numpy(reference_frame).permute(2, 0, 1).unsqueeze(0).to(device)
+ prediction_tensor = torch.from_numpy(prediction_frame).permute(2, 0, 1).unsqueeze(0).to(device)
+ ssim_scores.append(float(ssim(prediction_tensor, reference_tensor)))
+ psnr_scores.append(float(psnr(prediction_tensor, reference_tensor)))
+ ssim.reset()
+ psnr.reset()
+ difference = np.abs(reference_frame.astype(np.float64) - prediction_frame.astype(np.float64))
+ max_abs = max(max_abs, float(difference.max()))
+ absolute_error_sum += float(difference.sum())
return {
"ssim_mean": float(np.mean(ssim_scores)),
"ssim_min": float(np.min(ssim_scores)),
"psnr_mean_db": float(np.mean(psnr_scores)),
- "max_abs": float(difference.max()),
- "mean_abs": float(difference.mean()),
+ "max_abs": max_abs,
+ "mean_abs": absolute_error_sum / reference.size,
}
@@ -343,13 +548,19 @@ def _audio_metrics(reference: np.ndarray, prediction: np.ndarray) -> dict[str, f
@pytest.mark.benchmark
@pytest.mark.diffusion
@hardware_test(res={"cuda": "H100"}, num_cards=1)
-def test_ltx_one_stage_matches_official(case: LTXAccuracyCase, accuracy_artifact_root: Path) -> None:
- """Compare official and Omni raw AV outputs from the same E2E request."""
+def test_ltx_matches_official(case: LTXAccuracyCase, accuracy_artifact_root: Path) -> None:
+ """Compare official and Omni raw AV outputs from the same single-H100 request."""
+ configured_artifact_root = os.environ.get("VLLM_TEST_LTX_ARTIFACT_ROOT")
+ if configured_artifact_root:
+ accuracy_artifact_root = Path(configured_artifact_root)
+ accuracy_artifact_root.mkdir(parents=True, exist_ok=True)
output_root = reset_artifact_dir(accuracy_artifact_root / "ltx_official" / case.name)
official_root, official_revision = _official_source(accuracy_artifact_root / "ltx_official")
model = _resolve_model(case)
- gemma_root = _resolve_gemma_root(model)
- checkpoint = _resolve_checkpoint(case, model)
+ gemma_root = _resolve_gemma_root(case, model)
+ checkpoint = _resolve_artifact(case.checkpoint, model)
+ spatial_upsampler = _resolve_artifact(case.spatial_upsampler, model) if case.spatial_upsampler is not None else None
+ distilled_lora = _resolve_artifact(case.distilled_lora, model) if case.distilled_lora is not None else None
image = _resolve_image(case)
request_path = output_root / "request.json"
request_path.write_text(json.dumps(_request(case, image), indent=2) + "\n")
@@ -359,8 +570,13 @@ def test_ltx_one_stage_matches_official(case: LTXAccuracyCase, accuracy_artifact
str(runner),
"--request",
str(request_path),
+ "--pipeline-kind",
+ case.pipeline_kind,
]
- if os.environ.get("VLLM_TEST_LTX_ENABLE_LAYERWISE_OFFLOAD", "").lower() in {"1", "true", "yes", "on"}:
+ enable_layerwise_offload = case.enable_layerwise_offload or os.environ.get(
+ "VLLM_TEST_LTX_ENABLE_LAYERWISE_OFFLOAD", ""
+ ).lower() in {"1", "true", "yes", "on"}
+ if enable_layerwise_offload:
runner_args.append("--enable-layerwise-offload")
env = os.environ.copy()
env["VLLM_TEST_LTX_OFFICIAL_REVISION"] = official_revision
@@ -372,10 +588,9 @@ def test_ltx_one_stage_matches_official(case: LTXAccuracyCase, accuracy_artifact
)
official_output = output_root / "official"
- _run(
- _official_runner_prefix()
- + runner_args
- + [
+ official_args = _official_runner_prefix() + runner_args
+ official_args.extend(
+ [
"--backend",
"official",
"--output-dir",
@@ -386,12 +601,16 @@ def test_ltx_one_stage_matches_official(case: LTXAccuracyCase, accuracy_artifact
str(checkpoint),
"--gemma-root",
str(gemma_root),
- ],
- env=env,
+ ]
)
+ if spatial_upsampler is not None:
+ official_args.extend(["--spatial-upsampler", str(spatial_upsampler)])
+ if distilled_lora is not None:
+ official_args.extend(["--distilled-lora", str(distilled_lora)])
+ _run(official_args, env=env)
omni_output = output_root / "omni"
- _run(
+ omni_args = (
[sys.executable]
+ runner_args
+ [
@@ -403,9 +622,9 @@ def test_ltx_one_stage_matches_official(case: LTXAccuracyCase, accuracy_artifact
str(model),
"--model-class-name",
case.model_class_name,
- ],
- env=env,
+ ]
)
+ _run(omni_args, env=env)
official_metadata = json.loads((official_output / "metadata.json").read_text())
omni_metadata = json.loads((omni_output / "metadata.json").read_text())
@@ -423,18 +642,27 @@ def test_ltx_one_stage_matches_official(case: LTXAccuracyCase, accuracy_artifact
result = {
"case": case.name,
"task": "i2v" if image is not None else "t2v",
+ "pipeline_kind": case.pipeline_kind,
+ "model_class_name": case.model_class_name,
"attention_backend": ATTENTION_BACKEND,
"official_revision": official_revision,
"model_revision": case.model_revision,
- "checkpoint_revision": case.checkpoint_revision,
+ "checkpoint_revision": case.checkpoint.revision,
+ "spatial_upsampler_revision": (case.spatial_upsampler.revision if case.spatial_upsampler is not None else None),
+ "distilled_lora_revision": case.distilled_lora.revision if case.distilled_lora is not None else None,
+ "enable_layerwise_offload": enable_layerwise_offload,
+ "thresholds": asdict(case.thresholds),
+ "request": _request(case, image),
"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"] >= VIDEO_SSIM_MEAN_THRESHOLD
- assert video_metrics["ssim_min"] >= VIDEO_SSIM_MIN_THRESHOLD
- assert video_metrics["psnr_mean_db"] >= VIDEO_PSNR_MEAN_THRESHOLD
- assert audio_metrics["relative_l2"] <= AUDIO_RELATIVE_L2_THRESHOLD
- assert audio_metrics["cosine_similarity"] >= AUDIO_COSINE_THRESHOLD
+ assert video_metrics["ssim_mean"] >= case.thresholds.video_ssim_mean
+ assert video_metrics["ssim_min"] >= case.thresholds.video_ssim_min
+ assert video_metrics["psnr_mean_db"] >= case.thresholds.video_psnr_mean_db
+ if case.thresholds.audio_relative_l2 > 0:
+ assert audio_metrics["relative_l2"] <= case.thresholds.audio_relative_l2
+ if case.thresholds.audio_cosine_similarity > 0:
+ assert audio_metrics["cosine_similarity"] >= case.thresholds.audio_cosine_similarity
diff --git a/tests/e2e/online_serving/test_ltx.py b/tests/e2e/online_serving/test_ltx.py
new file mode 100644
index 000000000..43df5df2d
--- /dev/null
+++ b/tests/e2e/online_serving/test_ltx.py
@@ -0,0 +1,109 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Online-serving smoke tests for the public LTX-2 pipeline entries."""
+
+import os
+
+import pytest
+
+from tests.helpers.mark import hardware_marks
+from tests.helpers.media import generate_synthetic_image
+from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
+
+DISTILLED_MODEL = os.getenv("VLLM_TEST_LTX23_DISTILLED_MODEL", "diffusers/LTX-2.3-Distilled-Diffusers")
+BASE_MODEL = os.getenv("VLLM_TEST_LTX23_MODEL", "diffusers/LTX-2.3-Diffusers")
+PROMPT = "A serene lake at sunset with mountains in the background."
+SINGLE_CARD_MARKS = hardware_marks(res={"cuda": "H100"})
+HSDP_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
+
+
+def _server(
+ model: str,
+ model_class_name: str,
+ *extra_args: str,
+) -> OmniServerParams:
+ return OmniServerParams(
+ model=model,
+ server_args=[
+ "--enforce-eager",
+ *extra_args,
+ "--model-class-name",
+ model_class_name,
+ ],
+ )
+
+
+def _cases():
+ return [
+ pytest.param(
+ _server(BASE_MODEL, "LTX2Pipeline"),
+ 3,
+ id="one_stage",
+ marks=SINGLE_CARD_MARKS,
+ ),
+ pytest.param(
+ _server(DISTILLED_MODEL, "LTX2DistilledPipeline"),
+ 8,
+ id="distilled",
+ marks=SINGLE_CARD_MARKS,
+ ),
+ pytest.param(
+ _server(
+ DISTILLED_MODEL,
+ "LTX2DistilledPipeline",
+ "--use-hsdp",
+ "--hsdp-shard-size",
+ "2",
+ ),
+ 8,
+ id="distilled_hsdp2",
+ marks=HSDP_MARKS,
+ ),
+ pytest.param(
+ _server(
+ BASE_MODEL,
+ "LTX2TwoStagePipeline",
+ "--enable-layerwise-offload",
+ ),
+ 3,
+ id="two_stage",
+ marks=SINGLE_CARD_MARKS,
+ ),
+ ]
+
+
+@pytest.mark.parametrize(
+ ("omni_server", "num_inference_steps"),
+ _cases(),
+ indirect=["omni_server"],
+)
+def test_ltx_pipeline_entries(
+ omni_server: OmniServer,
+ num_inference_steps: int,
+ openai_client: OpenAIClientHandler,
+ subtests,
+):
+ for task in ("t2v", "i2v"):
+ with subtests.test(task=task):
+ form_data = {
+ "prompt": PROMPT,
+ "model": omni_server.model,
+ "height": 128,
+ "width": 128,
+ "num_frames": 9,
+ "fps": 8,
+ "num_inference_steps": num_inference_steps,
+ "seed": 42,
+ }
+ request_config = {
+ "model": omni_server.model,
+ "form_data": form_data,
+ }
+ 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)
diff --git a/tests/e2e/online_serving/test_ltx2_expansion.py b/tests/e2e/online_serving/test_ltx2_expansion.py
deleted file mode 100644
index 4d0eb6c90..000000000
--- a/tests/e2e/online_serving/test_ltx2_expansion.py
+++ /dev/null
@@ -1,91 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Expansion tests for LTX-2 two-stage pipelines.
-
-Coverage:
-- HSDP with LTX2DistilledPipeline (text-to-video and image-to-video)
-"""
-
-import os
-
-import pytest
-
-from tests.helpers.mark import hardware_marks
-from tests.helpers.media import generate_synthetic_image
-from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
-
-pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
-
-MODEL = os.getenv("VLLM_TEST_LTX2_MODEL", "rootonchair/LTX-2-19b-distilled")
-PROMPT = "A serene lake at sunset with mountains in the background."
-PARALLEL_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
-HSDP_ARGS = ["--use-hsdp", "--hsdp-shard-size", "2"]
-
-
-def _cases():
- cases = []
-
- cases.append(
- pytest.param(
- OmniServerParams(
- model=MODEL,
- server_args=[
- *HSDP_ARGS,
- "--model-class-name",
- "LTX2DistilledPipeline",
- ],
- ),
- False,
- id="t2v_hsdp",
- marks=PARALLEL_MARKS,
- )
- )
-
- cases.append(
- pytest.param(
- OmniServerParams(
- model=MODEL,
- server_args=[
- *HSDP_ARGS,
- "--model-class-name",
- "LTX2DistilledPipeline",
- ],
- ),
- True,
- id="i2v_hsdp",
- marks=PARALLEL_MARKS,
- )
- )
-
- return cases
-
-
-@pytest.mark.parametrize(("omni_server", "is_i2v"), _cases(), indirect=["omni_server"])
-def test_ltx2_two_stage_hsdp(
- omni_server: OmniServer,
- is_i2v: bool,
- openai_client: OpenAIClientHandler,
-):
- # Keep CI small while exercising the fixed 8+3-step distilled recipe.
- # Height and width describe the final output; Stage 1 runs at half size.
- form_data = {
- "prompt": PROMPT,
- "model": omni_server.model,
- "height": 128,
- "width": 128,
- "num_frames": 9,
- "fps": 8,
- "num_inference_steps": 8,
- "seed": 42,
- }
-
- request_config = {
- "model": omni_server.model,
- "form_data": form_data,
- }
-
- if is_i2v:
- request_config["image_reference"] = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
-
- openai_client.send_video_diffusion_request(request_config)
diff --git a/tests/model_extras/test_model_extras.py b/tests/model_extras/test_model_extras.py
index d57671adc..557fb5f6c 100644
--- a/tests/model_extras/test_model_extras.py
+++ b/tests/model_extras/test_model_extras.py
@@ -233,10 +233,13 @@ def test_ltx_extra_registry_declares_official_guidance_params() -> None:
}
)
- assert get_extra_body_params("LTX2Pipeline") == expected
- assert get_extra_output_params("LTX2Pipeline") == frozenset()
- assert get_extra_body_params("LTX2DistilledPipeline") == expected
- assert get_extra_output_params("LTX2DistilledPipeline") == frozenset()
+ for pipeline_name in (
+ "LTX2Pipeline",
+ "LTX2TwoStagePipeline",
+ "LTX2DistilledPipeline",
+ ):
+ assert get_extra_body_params(pipeline_name) == expected
+ assert get_extra_output_params(pipeline_name) == frozenset()
@pytest.mark.core_model
diff --git a/tests/model_tests/diffusion/test_alignment.py b/tests/model_tests/diffusion/test_alignment.py
index be064f603..f5d6c18c4 100644
--- a/tests/model_tests/diffusion/test_alignment.py
+++ b/tests/model_tests/diffusion/test_alignment.py
@@ -27,6 +27,7 @@ EXCLUDED_MODELS = [
"OvisImagePipeline",
"WanPipeline",
"WanVACEPipeline",
+ "LTX2TwoStagePipeline",
"LTX2DistilledPipeline",
"LTX2T2VDMD2Pipeline",
"LTX2I2VDMD2Pipeline",
diff --git a/vllm_omni/diffusion/models/ltx2/__init__.py b/vllm_omni/diffusion/models/ltx2/__init__.py
index 300295bd8..20ada6adc 100644
--- a/vllm_omni/diffusion/models/ltx2/__init__.py
+++ b/vllm_omni/diffusion/models/ltx2/__init__.py
@@ -12,12 +12,16 @@ from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import (
LTX2Pipeline,
LTX2T2VDMD2Pipeline,
)
-from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import LTX2DistilledPipeline
+from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import (
+ LTX2DistilledPipeline,
+ LTX2TwoStagePipeline,
+)
__all__ = [
"LTX2Pipeline",
"LTX2T2VDMD2Pipeline",
"LTX2I2VDMD2Pipeline",
+ "LTX2TwoStagePipeline",
"LTX2DistilledPipeline",
"get_ltx2_post_process_func",
"load_transformer_config",
diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_adapter_parser.py b/vllm_omni/diffusion/models/ltx2/ltx2_adapter_parser.py
new file mode 100644
index 000000000..d8228370e
--- /dev/null
+++ b/vllm_omni/diffusion/models/ltx2/ltx2_adapter_parser.py
@@ -0,0 +1,156 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Parse official LTX distilled adapters without materializing their tensors.
+
+The official LTX adapters are raw safetensors files rather than PEFT
+checkpoints. This module is deliberately LTX-specific: it translates the
+official names into the locally constructed Diffusers transformer names and
+describes the resulting adapter targets. Applying those targets is the job
+of :mod:`ltx2_phase_adapter`.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+from dataclasses import dataclass
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+from safetensors import safe_open
+
+_IGNORED_OFFICIAL_LORA_MODULES = {"text_embedding_projection.aggregate_embed"}
+_OFFICIAL_TO_DIFFUSERS_PREFIXES = (
+ ("audio_prompt_adaln_single.", "audio_prompt_adaln."),
+ ("prompt_adaln_single.", "prompt_adaln."),
+ ("audio_adaln_single.", "audio_time_embed."),
+ ("adaln_single.", "time_embed."),
+ ("audio_patchify_proj", "audio_proj_in"),
+ ("patchify_proj", "proj_in"),
+ ("av_ca_video_scale_shift_adaln_single.", "av_cross_attn_video_scale_shift."),
+ ("av_ca_audio_scale_shift_adaln_single.", "av_cross_attn_audio_scale_shift."),
+ ("av_ca_a2v_gate_adaln_single.", "av_cross_attn_video_a2v_gate."),
+ ("av_ca_v2a_gate_adaln_single.", "av_cross_attn_audio_v2a_gate."),
+)
+
+
+@dataclass(frozen=True)
+class AdapterTarget:
+ """One A/B pair mapped onto a local transformer linear layer."""
+
+ # Keep the original Diffusers module name for diagnostics after LTX packs
+ # Q/K/V into the runtime module path.
+ source_module: str
+ module: str
+ slice: str | None
+ a_key: str
+ b_key: str
+ a_shape: tuple[int, ...]
+ b_shape: tuple[int, ...]
+
+
+@dataclass(frozen=True)
+class AdapterManifest:
+ """Tensor-free description of the fixed refinement adapter."""
+
+ path: str
+ targets: tuple[AdapterTarget, ...]
+
+
+def _to_diffusers_module_name(name: str) -> str:
+ for official_prefix, diffusers_prefix in _OFFICIAL_TO_DIFFUSERS_PREFIXES:
+ if name.startswith(official_prefix):
+ return diffusers_prefix + name[len(official_prefix) :]
+ return name
+
+
+def _resolve_lora_target(transformer: nn.Module, module_name: str) -> tuple[str, str | int | None] | None:
+ modules = dict(transformer.named_modules())
+ if module_name in modules:
+ return module_name, None
+ for packed_suffix, sub_suffix, shard_id in getattr(transformer, "stacked_params_mapping", ()):
+ if sub_suffix in module_name:
+ packed_name = module_name.replace(sub_suffix, packed_suffix)
+ if packed_name in modules:
+ return packed_name, shard_id
+ return None
+
+
+def parse_ltx_adapter(transformer: nn.Module, path: str) -> AdapterManifest:
+ """Translate one official raw LTX adapter into local Transformer targets."""
+ path = str(Path(path))
+ targets: list[AdapterTarget] = []
+ unresolved: list[str] = []
+ seen: set[tuple[str, str | int | None]] = set()
+
+ with safe_open(path, framework="pt", device="cpu") as handle:
+ tensor_names = set(handle.keys())
+ official_name_set = {
+ tensor_name.removeprefix("diffusion_model.").rsplit(".lora_", 1)[0]
+ for tensor_name in tensor_names
+ if tensor_name.endswith(".lora_A.weight")
+ }
+ official_b_names = {
+ tensor_name.removeprefix("diffusion_model.").rsplit(".lora_", 1)[0]
+ for tensor_name in tensor_names
+ if tensor_name.endswith(".lora_B.weight")
+ }
+ missing_a = sorted(official_b_names - official_name_set - _IGNORED_OFFICIAL_LORA_MODULES)
+ if missing_a:
+ raise ValueError(f"Missing paired LoRA A tensors for {missing_a} in {path}.")
+
+ official_names = sorted(official_name_set)
+ for official_name in official_names:
+ if official_name in _IGNORED_OFFICIAL_LORA_MODULES:
+ continue
+ key_prefix = f"diffusion_model.{official_name}"
+ key_a = f"{key_prefix}.lora_A.weight"
+ key_b = f"{key_prefix}.lora_B.weight"
+ if key_b not in tensor_names:
+ raise ValueError(f"Missing paired LoRA tensor {key_b!r} in {path}.")
+
+ a_shape = tuple(handle.get_slice(key_a).get_shape())
+ b_shape = tuple(handle.get_slice(key_b).get_shape())
+ if len(a_shape) != 2 or len(b_shape) != 2 or b_shape[1] != a_shape[0]:
+ raise ValueError(f"Invalid LoRA pair for {official_name}: A={a_shape}, B={b_shape}.")
+
+ source_module = _to_diffusers_module_name(official_name)
+ resolved = _resolve_lora_target(transformer, source_module)
+ if resolved is None:
+ unresolved.append(official_name)
+ continue
+ module, shard_id = resolved
+ target_id = (module, shard_id)
+ if target_id in seen:
+ raise ValueError(f"Official LTX LoRA contains duplicate target {module!r}, slice {shard_id!r}.")
+ seen.add(target_id)
+ targets.append(
+ AdapterTarget(
+ source_module=source_module,
+ module=module,
+ slice=shard_id,
+ a_key=key_a,
+ b_key=key_b,
+ a_shape=a_shape,
+ b_shape=b_shape,
+ )
+ )
+
+ if unresolved:
+ raise ValueError(f"Official LTX LoRA contains unmapped modules: {unresolved}.")
+ if not targets:
+ raise ValueError(f"No Transformer LoRA tensors were found in {path}.")
+ return AdapterManifest(path=path, targets=tuple(targets))
+
+
+def iter_adapter_tensors(manifest: AdapterManifest) -> Iterator[tuple[AdapterTarget, torch.Tensor, torch.Tensor]]:
+ """Stream adapter pairs in manifest order without a full checkpoint copy."""
+
+ with safe_open(manifest.path, framework="pt", device="cpu") as handle:
+ for target in manifest.targets:
+ yield (
+ target,
+ handle.get_tensor(target.a_key),
+ handle.get_tensor(target.b_key),
+ )
diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_components.py b/vllm_omni/diffusion/models/ltx2/ltx2_components.py
index 54c866eb7..7f2f31a60 100644
--- a/vllm_omni/diffusion/models/ltx2/ltx2_components.py
+++ b/vllm_omni/diffusion/models/ltx2/ltx2_components.py
@@ -9,7 +9,8 @@ import inspect
import json
import logging
import os
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
+from pathlib import Path
from typing import TYPE_CHECKING, Any
import torch
@@ -19,6 +20,8 @@ from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel
from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder
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 vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
@@ -32,6 +35,7 @@ from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
if TYPE_CHECKING:
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+from .ltx2_request import LTXCheckpointKind, validate_ltx_checkpoint
from .ltx2_transformer import (
LTX2VideoTransformer3DModel,
apply_interleaved_rotary_emb,
@@ -70,6 +74,9 @@ class LTXComponentProfile:
video_vae_cls: type = AutoencoderKLLTX2Video
vocoder_cls: type = LTX2Vocoder
vocoder_fallback_cls: type | None = None
+ artifact_repo_id: str | None = None
+ latent_upsampler_filename: str | None = None
+ distilled_lora_filename: str | None = None
LTX2_COMPONENT_PROFILE = LTXComponentProfile(
@@ -101,16 +108,96 @@ LTX2_DISTILLED_COMPONENT_PROFILE = LTXComponentProfile(
video_vae_cls=DistributedAutoencoderKLLTX2Video,
)
+LTX23_DISTILLED_COMPONENT_PROFILE = replace(
+ LTX23_COMPONENT_PROFILE,
+ name="ltx2_3_distilled",
+ resident_modules=(*LTX23_COMPONENT_PROFILE.resident_modules, "latent_upsampler"),
+ artifact_repo_id="Lightricks/LTX-2.3",
+ latent_upsampler_filename="ltx-2.3-spatial-upscaler-x2-1.1.safetensors",
+)
+
+LTX2_TWO_STAGE_COMPONENT_PROFILE = replace(
+ LTX2_COMPONENT_PROFILE,
+ name="ltx2_two_stage",
+ resident_modules=(*LTX2_COMPONENT_PROFILE.resident_modules, "latent_upsampler"),
+ artifact_repo_id="Lightricks/LTX-2",
+ latent_upsampler_filename="ltx-2-spatial-upscaler-x2-1.0.safetensors",
+ distilled_lora_filename="ltx-2-19b-distilled-lora-384.safetensors",
+)
+
+LTX23_TWO_STAGE_COMPONENT_PROFILE = replace(
+ LTX23_COMPONENT_PROFILE,
+ name="ltx2_3_two_stage",
+ resident_modules=(*LTX23_COMPONENT_PROFILE.resident_modules, "latent_upsampler"),
+ artifact_repo_id="Lightricks/LTX-2.3",
+ latent_upsampler_filename="ltx-2.3-spatial-upscaler-x2-1.1.safetensors",
+ distilled_lora_filename="ltx-2.3-22b-distilled-lora-384-1.1.safetensors",
+)
+
_COMPONENT_PROFILES: dict[tuple[str, str], LTXComponentProfile] = {
("one_stage", "2"): LTX2_COMPONENT_PROFILE,
("one_stage", "2.3"): LTX23_COMPONENT_PROFILE,
+ ("two_stage", "2"): LTX2_TWO_STAGE_COMPONENT_PROFILE,
+ ("two_stage", "2.3"): LTX23_TWO_STAGE_COMPONENT_PROFILE,
("distilled_two_stage", "2"): LTX2_DISTILLED_COMPONENT_PROFILE,
+ ("distilled_two_stage", "2.3"): LTX23_DISTILLED_COMPONENT_PROFILE,
("dmd2", "2"): LTX2_COMPONENT_PROFILE,
("dmd2", "2.3"): LTX23_COMPONENT_PROFILE,
}
+def resolve_ltx_artifact(
+ model: str,
+ repo_id: str,
+ filename: str,
+) -> str:
+ """Resolve an official LTX sidecar from the model root or its Hub repository."""
+ candidate = Path(model) / filename
+ if candidate.is_file():
+ return str(candidate)
+
+ try:
+ return hf_hub_download(repo_id=repo_id, filename=filename)
+ except Exception as exc:
+ raise FileNotFoundError(
+ f"Unable to resolve LTX artifact {filename!r}. Searched {candidate}; "
+ f"place the file in the model root or make {repo_id} available."
+ ) from exc
+
+
+def _load_ltx_latent_upsampler_single_file(path: str, dtype: torch.dtype) -> LTX2LatentUpsamplerModel:
+ """Load an official single-file upsampler into the Diffusers module."""
+ with safe_open(path, framework="pt", device="cpu") as handle:
+ metadata = handle.metadata() or {}
+ raw_config = json.loads(metadata.get("config", "{}"))
+ config = {
+ "in_channels": raw_config.get("in_channels", 128),
+ "mid_channels": raw_config.get("mid_channels", 1024),
+ "num_blocks_per_stage": raw_config.get("num_blocks_per_stage", 4),
+ "dims": raw_config.get("dims", 3),
+ "spatial_upsample": raw_config.get("spatial_upsample", True),
+ "temporal_upsample": raw_config.get("temporal_upsample", False),
+ "rational_spatial_scale": raw_config.get("spatial_scale", 2.0),
+ "use_rational_resampler": raw_config.get("rational_resampler", False),
+ }
+ with torch.device("cpu"):
+ upsampler = LTX2LatentUpsamplerModel(**config).to(dtype=dtype)
+
+ state_dict = load_file(path, device="cpu")
+ if "upsampler.0.weight" in state_dict and hasattr(upsampler.upsampler, "conv"):
+ state_dict["upsampler.conv.weight"] = state_dict.pop("upsampler.0.weight")
+ state_dict["upsampler.conv.bias"] = state_dict.pop("upsampler.0.bias")
+ missing, unexpected = upsampler.load_state_dict(state_dict, strict=False)
+ unresolved_missing = set(missing) - {"upsampler.blur_down.kernel"}
+ if unresolved_missing or unexpected:
+ raise ValueError(
+ f"Invalid LTX latent upsampler {path}: missing={sorted(unresolved_missing)}, "
+ f"unexpected={sorted(unexpected)}."
+ )
+ return upsampler
+
+
def resolve_ltx_component_profile(pipeline_kind: str, model_version: str) -> LTXComponentProfile:
"""Resolve component construction independently from execution recipes."""
try:
@@ -392,16 +479,34 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None:
)
if "latent_upsampler" in profile.resident_modules:
- # BlurDownsample constructs an integer kernel that must be initialized
- # on CPU; component placement is handled uniformly after construction.
- with torch.device("cpu"):
- pipeline.latent_upsampler = _load_component(
- LTX2LatentUpsamplerModel,
+ upsampler_config = os.path.join(model, "latent_upsampler", "config.json")
+ if os.path.isfile(upsampler_config) or not local_files_only:
+ try:
+ pipeline.latent_upsampler = _load_component(
+ LTX2LatentUpsamplerModel,
+ model,
+ "latent_upsampler",
+ local_files_only=local_files_only,
+ dtype=dtype,
+ )
+ except (OSError, ValueError):
+ if profile.latent_upsampler_filename is None or profile.artifact_repo_id is None:
+ raise
+ upsampler_path = resolve_ltx_artifact(
+ model,
+ profile.artifact_repo_id,
+ profile.latent_upsampler_filename,
+ )
+ pipeline.latent_upsampler = _load_ltx_latent_upsampler_single_file(upsampler_path, dtype)
+ else:
+ if profile.latent_upsampler_filename is None or profile.artifact_repo_id is None:
+ raise FileNotFoundError(f"LTX latent upsampler component not found under {model}.")
+ upsampler_path = resolve_ltx_artifact(
model,
- "latent_upsampler",
- local_files_only=local_files_only,
- dtype=dtype,
+ profile.artifact_repo_id,
+ profile.latent_upsampler_filename,
)
+ pipeline.latent_upsampler = _load_ltx_latent_upsampler_single_file(upsampler_path, dtype)
transformer_config = load_transformer_config(model, "transformer", local_files_only)
quant_config = getattr(od_config, "quantization_config", None)
@@ -412,6 +517,16 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None:
subfolder="scheduler",
local_files_only=local_files_only,
)
+ 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"
+ validate_ltx_checkpoint(
+ pipeline.scheduler.config,
+ expected_kind=expected_checkpoint_kind,
+ pipeline_name=type(pipeline).__name__,
+ )
pipeline.vae_spatial_compression_ratio = pipeline.vae.spatial_compression_ratio
pipeline.vae_temporal_compression_ratio = pipeline.vae.temporal_compression_ratio
diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_denoise.py b/vllm_omni/diffusion/models/ltx2/ltx2_denoise.py
index 4a149be85..dd55f3953 100644
--- a/vllm_omni/diffusion/models/ltx2/ltx2_denoise.py
+++ b/vllm_omni/diffusion/models/ltx2/ltx2_denoise.py
@@ -62,13 +62,14 @@ class LTXDenoiseContext:
audio_latents: torch.Tensor
video_coords: torch.Tensor
audio_coords: torch.Tensor
+ audio_attention_mask: torch.Tensor | None = None
conditioning_mask: torch.Tensor | None = None
conditioning_mask_for_model: torch.Tensor | None = None
@dataclass
class LTXPhaseResult:
- """Denoised, unpacked AV latents and the context used to produce them."""
+ """Denoised AV latents and the context used to produce them."""
forward_context: LTXForwardContext
video: torch.Tensor
@@ -284,7 +285,8 @@ def prepare_rope_coords_stage(
latents: torch.Tensor,
audio_latents: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
- video_coords = pipeline.transformer.rope.prepare_video_coords(
+ transformer = pipeline.transformer
+ video_coords = transformer.rope.prepare_video_coords(
latents.shape[0],
forward_ctx.latent_num_frames,
forward_ctx.latent_height,
@@ -292,7 +294,7 @@ def prepare_rope_coords_stage(
latents.device,
fps=forward_ctx.request_inputs.frame_rate,
)
- audio_coords = pipeline.transformer.audio_rope.prepare_audio_coords(
+ audio_coords = transformer.audio_rope.prepare_audio_coords(
audio_latents.shape[0],
forward_ctx.padded_audio_num_frames,
audio_latents.device,
@@ -331,6 +333,7 @@ def build_transformer_kwargs(
# with a learned register, making all output context tokens valid.
"encoder_attention_mask": None,
"audio_encoder_attention_mask": None,
+ "audio_attention_mask": denoise_ctx.audio_attention_mask,
"num_frames": forward_ctx.latent_num_frames,
"height": forward_ctx.latent_height,
"width": forward_ctx.latent_width,
@@ -442,11 +445,25 @@ class LTXPhaseExecutor:
video_audio_step_adapter=video_audio_step_adapter,
)
video_coords, audio_coords = prepare_rope_coords_stage(pipeline, forward_ctx, latents, audio_latents)
+ ring_degree = getattr(pipeline.od_config.parallel_config, "ring_degree", 1) or 1
+ if padded_audio_num_frames > original_audio_num_frames and ring_degree > 1:
+ raise ValueError(
+ "LTX audio padding requires an attention mask, which Ring sequence parallelism does not support. "
+ "Use Ulysses-only SP or choose a request whose audio latent length is divisible by the SP size."
+ )
denoise_ctx = LTXDenoiseContext(
latents=latents,
audio_latents=audio_latents,
video_coords=video_coords,
audio_coords=audio_coords,
+ audio_attention_mask=(
+ torch.arange(padded_audio_num_frames, device=audio_latents.device)
+ .lt(original_audio_num_frames)
+ .unsqueeze(0)
+ .expand(audio_latents.shape[0], -1)
+ if padded_audio_num_frames > original_audio_num_frames
+ else None
+ ),
conditioning_mask=conditioning_mask,
)
denoise_ctx = pipeline._prepare_denoise_context_for_guidance(forward_ctx, denoise_ctx)
diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_guidance.py b/vllm_omni/diffusion/models/ltx2/ltx2_guidance.py
index f35f5b998..bc84e7c32 100644
--- a/vllm_omni/diffusion/models/ltx2/ltx2_guidance.py
+++ b/vllm_omni/diffusion/models/ltx2/ltx2_guidance.py
@@ -126,7 +126,10 @@ 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."""
- sigma = sigma.to(torch.float32)
+ # 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.
+ sigma = sigma.to(torch.float32).item()
return ((sample.to(torch.float32) - x0.to(torch.float32)) / sigma).to(sample.dtype)
@@ -148,6 +151,7 @@ def combine_guided_x0(
uncond_perturbed: torch.Tensor | float,
uncond_modality: torch.Tensor | float,
guidance: LTXModalityGuidance,
+ rescale_token_count: int | None = None,
) -> torch.Tensor:
dtype = cond.dtype
cond = cond.float()
@@ -163,9 +167,18 @@ def combine_guided_x0(
# uses its full-tensor standard deviation. In Omni, dimension 0 may
# instead contain independently scheduled requests. Reduce each item
# separately so its result cannot depend on unrelated co-batched work.
- reduce_dims = tuple(range(1, pred.ndim))
- cond_std = cond.std(dim=reduce_dims, keepdim=True)
- pred_std = pred.std(dim=reduce_dims, keepdim=True)
+ stats_cond = cond
+ stats_pred = pred
+ if rescale_token_count is not None:
+ if cond.ndim < 2 or not 0 < rescale_token_count <= cond.shape[1]:
+ raise ValueError(
+ f"Guidance rescale token count must be in [1, {cond.shape[1]}], got {rescale_token_count}."
+ )
+ stats_cond = cond[:, :rescale_token_count]
+ stats_pred = pred[:, :rescale_token_count]
+ reduce_dims = tuple(range(1, stats_pred.ndim))
+ cond_std = stats_cond.std(dim=reduce_dims, keepdim=True)
+ pred_std = stats_pred.std(dim=reduce_dims, keepdim=True)
eps = torch.finfo(pred.dtype).eps
factor = cond_std / pred_std.clamp_min(eps)
factor = torch.where(pred_std > eps, factor, torch.ones_like(factor))
@@ -175,6 +188,8 @@ def combine_guided_x0(
def _repeat_batch(tensor: torch.Tensor, repeats: int) -> torch.Tensor:
+ if repeats == 1:
+ return tensor
return tensor.repeat((repeats,) + (1,) * (tensor.ndim - 1))
@@ -278,6 +293,11 @@ class LTXGuidanceExecutor:
if model_pass_count > 1:
denoise_ctx.video_coords = _repeat_batch(denoise_ctx.video_coords, model_pass_count)
denoise_ctx.audio_coords = _repeat_batch(denoise_ctx.audio_coords, model_pass_count)
+ if denoise_ctx.audio_attention_mask is not None:
+ denoise_ctx.audio_attention_mask = _repeat_batch(
+ denoise_ctx.audio_attention_mask,
+ model_pass_count,
+ )
return denoise_ctx
@staticmethod
@@ -322,6 +342,7 @@ class LTXGuidanceExecutor:
guidance: LTXModalityGuidance,
*,
model_sigma: torch.Tensor | None = None,
+ rescale_token_count: int | None = None,
) -> torch.Tensor:
model_sigma = sigma if model_sigma is None else model_sigma
# Official guidance reduces contiguous BSC tensors. The fused
@@ -334,6 +355,7 @@ class LTXGuidanceExecutor:
uncond_perturbed=x0.get("ptb", 0.0),
uncond_modality=x0.get("mod", 0.0),
guidance=guidance,
+ rescale_token_count=rescale_token_count,
)
return velocity_from_x0(sample, guided, sigma)
@@ -434,6 +456,7 @@ class LTXGuidanceExecutor:
audio_splits,
forward_ctx.audio_scheduler.sigmas[index],
plan.spec.audio,
+ rescale_token_count=forward_ctx.original_audio_num_frames,
),
)
@@ -509,6 +532,7 @@ class LTXGuidanceExecutor:
audio_splits,
forward_ctx.audio_scheduler.sigmas[index],
plan.spec.audio,
+ rescale_token_count=forward_ctx.original_audio_num_frames,
),
)
diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_latents.py b/vllm_omni/diffusion/models/ltx2/ltx2_latents.py
index bb0155e79..11929d331 100644
--- a/vllm_omni/diffusion/models/ltx2/ltx2_latents.py
+++ b/vllm_omni/diffusion/models/ltx2/ltx2_latents.py
@@ -117,7 +117,9 @@ def create_noised_state(
device=latents.device,
dtype=latents.dtype,
)
- return noise_scale * noise + (1 - noise_scale) * latents
+ if isinstance(noise_scale, torch.Tensor):
+ noise_scale = noise_scale.float()
+ return torch.lerp(latents.float(), noise.float(), noise_scale).to(latents.dtype)
def pack_audio_latents(latents: torch.Tensor) -> torch.Tensor:
@@ -135,6 +137,16 @@ def unpad_audio_latents(latents: torch.Tensor, num_frames: int) -> torch.Tensor:
return latents[:, :num_frames]
+def clear_audio_padding(latents: torch.Tensor, num_frames: int) -> torch.Tensor:
+ """Keep Omni's SP-only audio padding outside the 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]:
+ return latents
+ padding = latents.new_zeros(latents.shape[0], latents.shape[1] - num_frames, latents.shape[2])
+ return torch.cat([latents[:, :num_frames], padding], dim=1)
+
+
def get_sp_padded_audio_latent_length(audio_latent_length: int, sp_size: int) -> int:
if sp_size > 1:
audio_latent_length += (sp_size - (audio_latent_length % sp_size)) % sp_size
@@ -205,7 +217,12 @@ def prepare_video_latents(
f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
f" size of {batch_size}. Make sure the batch size matches the length of the generators."
)
- return randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+ 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)
def prepare_audio_latents(
@@ -226,6 +243,16 @@ def prepare_audio_latents(
sp_size = getattr(pipeline.od_config.parallel_config, "sequence_parallel_size", 1) or 1
padded_latent_length = get_sp_padded_audio_latent_length(original_latent_length, int(sp_size))
+ def pad_logical_latents(logical_latents: torch.Tensor) -> torch.Tensor:
+ if padded_latent_length == original_latent_length:
+ return logical_latents
+ padding = logical_latents.new_zeros(
+ logical_latents.shape[0],
+ padded_latent_length - original_latent_length,
+ logical_latents.shape[2],
+ )
+ return torch.cat([logical_latents, padding], dim=1)
+
if latents is not None:
if latents.ndim == 4:
latents = pack_audio_latents(latents)
@@ -237,29 +264,24 @@ def prepare_audio_latents(
pipeline.audio_vae.latents_mean,
pipeline.audio_vae.latents_std,
)
- latents = create_noised_state(latents, noise_scale, generator)
if latents.shape[1] not in {original_latent_length, padded_latent_length}:
raise ValueError(
"Provided `audio_latents` has incompatible audio frame count "
f"{latents.shape[1]}; expected {original_latent_length} or {padded_latent_length}."
)
- if latents.shape[1] == original_latent_length and padded_latent_length > original_latent_length:
- padding = torch.zeros(
- latents.shape[0],
- padded_latent_length - original_latent_length,
- latents.shape[2],
- dtype=latents.dtype,
- device=latents.device,
- )
- latents = torch.cat([latents, padding], dim=1)
+ # Padding is an Omni implementation detail, not part of the official
+ # audio state. Draw request RNG only for logical tokens so later phases
+ # remain seed-invariant across SP degrees, then append deterministic 0s.
+ latents = create_noised_state(latents[:, :original_latent_length], noise_scale, generator)
+ latents = pad_logical_latents(latents)
return latents.to(device=device, dtype=dtype), original_latent_length, padded_latent_length
- shape = (batch_size, padded_latent_length, num_channels_latents * latent_mel_bins)
+ shape = (batch_size, original_latent_length, num_channels_latents * latent_mel_bins)
if isinstance(generator, list) and len(generator) != batch_size:
raise ValueError(
f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
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)
+ latents = pad_logical_latents(randn_tensor(shape, generator=generator, device=device, dtype=dtype))
return latents, original_latent_length, padded_latent_length
diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_phase_adapter.py b/vllm_omni/diffusion/models/ltx2/ltx2_phase_adapter.py
new file mode 100644
index 000000000..a205fce6b
--- /dev/null
+++ b/vllm_omni/diffusion/models/ltx2/ltx2_phase_adapter.py
@@ -0,0 +1,430 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Construction and execution of a fixed LTX refinement-phase adapter.
+
+Checkpoint key parsing remains in :mod:`ltx2_adapter_parser`; this module
+resolves the declared artifact and makes the existing Transformer linears
+adapter-capable while retaining their quantization and tensor-parallel paths.
+"""
+
+from __future__ import annotations
+
+from collections import defaultdict
+from collections.abc import Iterable
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import vllm.envs as envs
+from vllm.distributed import (
+ split_tensor_along_last_dim,
+ tensor_model_parallel_all_reduce,
+)
+from vllm.logger import init_logger
+from vllm.model_executor.layers.batch_invariant import linear_batch_invariant
+from vllm.model_executor.layers.linear import (
+ ColumnParallelLinear,
+ QKVParallelLinear,
+ RowParallelLinear,
+ UnquantizedLinearMethod,
+)
+from vllm.model_executor.layers.utils import dispatch_unquantized_gemm
+from vllm.platforms import current_platform
+
+from .ltx2_adapter_parser import AdapterManifest, AdapterTarget, iter_adapter_tensors, parse_ltx_adapter
+from .ltx2_components import resolve_ltx_artifact
+from .ltx2_recipes import LTX_DISTILLED_ADAPTER_SLOT
+
+logger = init_logger(__name__)
+
+
+@dataclass(frozen=True)
+class _AdapterLayout:
+ """Rank-local storage and execution layout for one LoRA pair."""
+
+ kind: str
+ input_size: int
+ output_size: int
+ a_input_start: int
+ a_input_size: int
+ b_output_start: int
+ b_output_size: int
+ output_start: int
+ reduce_results: bool = False
+
+
+def _module_device(module: nn.Module) -> torch.device:
+ for tensor in (*module.parameters(), *module.buffers()):
+ if tensor.device.type != "meta":
+ return tensor.device
+ raise RuntimeError(f"Cannot materialize an LTX phase adapter for meta module {module._get_name()}.")
+
+
+def _build_layout(layer: nn.Module, target: AdapterTarget) -> _AdapterLayout:
+ if isinstance(layer, QKVParallelLinear):
+ if target.slice not in {"q", "k", "v"}:
+ raise ValueError(
+ f"LTX packed QKV target {target.module!r} requires a q/k/v adapter slice, got {target.slice!r}."
+ )
+ q_local = layer.num_heads * layer.head_size
+ k_local = layer.num_kv_heads * layer.head_size
+ v_local = layer.num_kv_heads * layer.v_head_size
+ if target.slice == "q":
+ output_size = layer.total_num_heads * layer.head_size
+ b_output_start = layer.tp_rank * q_local
+ b_output_size = q_local
+ output_start = 0
+ elif target.slice == "k":
+ output_size = layer.total_num_kv_heads * layer.head_size
+ shard_rank = layer.tp_rank // layer.num_kv_head_replicas
+ b_output_start = shard_rank * k_local
+ b_output_size = k_local
+ output_start = q_local
+ else:
+ output_size = layer.total_num_kv_heads * layer.v_head_size
+ shard_rank = layer.tp_rank // layer.num_kv_head_replicas
+ b_output_start = shard_rank * v_local
+ b_output_size = v_local
+ output_start = q_local + k_local
+ return _AdapterLayout(
+ kind="qkv",
+ input_size=layer.input_size,
+ output_size=output_size,
+ a_input_start=0,
+ a_input_size=layer.input_size,
+ b_output_start=b_output_start,
+ b_output_size=b_output_size,
+ output_start=output_start,
+ )
+
+ if isinstance(layer, ColumnParallelLinear):
+ if layer.gather_output:
+ raise ValueError(
+ f"LTX phase adapter does not support gathered ColumnParallelLinear target {target.module!r}."
+ )
+ return _AdapterLayout(
+ kind="column",
+ input_size=layer.input_size,
+ output_size=layer.output_size,
+ a_input_start=0,
+ a_input_size=layer.input_size,
+ b_output_start=layer.tp_rank * layer.output_size_per_partition,
+ b_output_size=layer.output_size_per_partition,
+ output_start=0,
+ )
+
+ if isinstance(layer, RowParallelLinear):
+ return _AdapterLayout(
+ kind="row",
+ input_size=layer.input_size,
+ output_size=layer.output_size,
+ a_input_start=layer.tp_rank * layer.input_size_per_partition,
+ a_input_size=layer.input_size_per_partition,
+ b_output_start=0,
+ b_output_size=layer.output_size,
+ output_start=0,
+ reduce_results=layer.reduce_results,
+ )
+
+ if isinstance(layer, nn.Linear):
+ return _AdapterLayout(
+ kind="replicated",
+ input_size=layer.in_features,
+ output_size=layer.out_features,
+ a_input_start=0,
+ a_input_size=layer.in_features,
+ b_output_start=0,
+ b_output_size=layer.out_features,
+ output_start=0,
+ )
+
+ raise TypeError(f"LTX phase adapter target {target.module!r} has unsupported layer type {type(layer).__name__}.")
+
+
+class _AdapterPiece(nn.Module):
+ """One rank-local A/B pair stored as non-persistent buffers."""
+
+ def __init__(self, target: AdapterTarget, layout: _AdapterLayout) -> None:
+ super().__init__()
+ self.target = target
+ self.layout = layout
+ # The fixed graph is installed before base checkpoint loading. The
+ # empty buffers are replaced only after the base layer's quant/TP layout
+ # has settled and before offload/compile snapshots are taken.
+ self.register_buffer("lora_a", torch.empty(0), persistent=False)
+ self.register_buffer("lora_b", torch.empty(0), persistent=False)
+
+ def load(self, lora_a: torch.Tensor, lora_b: torch.Tensor, *, device: torch.device, dtype: torch.dtype) -> None:
+ if tuple(lora_a.shape) != self.target.a_shape or tuple(lora_b.shape) != self.target.b_shape:
+ raise ValueError(
+ f"LTX adapter source shape changed for {self.target.source_module!r}: "
+ f"A={tuple(lora_a.shape)}, B={tuple(lora_b.shape)}."
+ )
+ if lora_a.ndim != 2 or lora_b.ndim != 2:
+ raise ValueError(f"LTX adapter target {self.target.source_module!r} must use rank-2 A/B tensors.")
+ rank = self.target.a_shape[0]
+ if lora_a.shape[1] != self.layout.input_size or lora_b.shape[0] != self.layout.output_size:
+ raise ValueError(
+ f"LTX adapter shape does not match {self.target.module!r}: "
+ f"expected A=({rank}, {self.layout.input_size}), "
+ f"B=({self.layout.output_size}, {rank}); "
+ f"got A={tuple(lora_a.shape)}, B={tuple(lora_b.shape)}."
+ )
+
+ local_a = lora_a.narrow(1, self.layout.a_input_start, self.layout.a_input_size)
+ local_b = lora_b.narrow(0, self.layout.b_output_start, self.layout.b_output_size)
+ self.lora_a = local_a.to(device=device, dtype=dtype)
+ self.lora_b = local_b.to(device=device, dtype=dtype)
+
+ def delta(self, input_: torch.Tensor) -> torch.Tensor:
+ if self.lora_a.numel() == 0 or self.lora_b.numel() == 0:
+ raise RuntimeError(f"LTX phase adapter target {self.target.module!r} was not materialized.")
+ if self.layout.kind == "row" and input_.shape[-1] != self.layout.a_input_size:
+ input_ = input_.narrow(-1, self.layout.a_input_start, self.layout.a_input_size)
+ input_ = input_.to(dtype=self.lora_a.dtype)
+ delta = F.linear(F.linear(input_, self.lora_a), self.lora_b)
+ if self.layout.kind == "row" and self.layout.reduce_results:
+ delta = tensor_model_parallel_all_reduce(delta)
+ return delta
+
+ def weight_delta(self) -> torch.Tensor:
+ """Materialize the rank-local BF16 weight delta used by official fusion."""
+ if self.lora_a.numel() == 0 or self.lora_b.numel() == 0:
+ raise RuntimeError(f"LTX phase adapter target {self.target.module!r} was not materialized.")
+ return torch.matmul(self.lora_b, self.lora_a)
+
+
+def _apply_unquantized_gemm_with_weight(
+ layer: nn.Module,
+ input_: torch.Tensor,
+ weight: torch.Tensor,
+ bias: torch.Tensor | None,
+) -> torch.Tensor:
+ if envs.VLLM_BATCH_INVARIANT and current_platform.is_cuda_alike():
+ return linear_batch_invariant(input_, weight, bias)
+ return dispatch_unquantized_gemm()(layer, input_, weight, bias)
+
+
+def _forward_with_weight(
+ layer: nn.Module,
+ input_: torch.Tensor,
+ weight: torch.Tensor,
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor | None]:
+ """Run an unquantized base layer with an ephemeral fused weight."""
+ if isinstance(layer, nn.Linear):
+ return F.linear(input_, weight, layer.bias)
+
+ quant_method = getattr(layer, "quant_method", None)
+ if not isinstance(quant_method, UnquantizedLinearMethod):
+ raise ValueError("LTX layer-fused phase LoRA requires unquantized Transformer linears.")
+
+ if isinstance(layer, RowParallelLinear):
+ if layer.input_is_parallel:
+ input_parallel = input_
+ else:
+ input_parallel = split_tensor_along_last_dim(input_, num_partitions=layer.tp_size)[
+ layer.tp_rank
+ ].contiguous()
+ bias = None if layer.tp_rank > 0 or layer.skip_bias_add else layer.bias
+ output_parallel = _apply_unquantized_gemm_with_weight(layer, input_parallel, weight, bias)
+ if layer.reduce_results and layer.tp_size > 1:
+ output = tensor_model_parallel_all_reduce(output_parallel)
+ else:
+ output = output_parallel
+ elif isinstance(layer, ColumnParallelLinear):
+ bias = layer.bias if not layer.skip_bias_add else None
+ output = _apply_unquantized_gemm_with_weight(layer, input_, weight, bias)
+ else:
+ raise TypeError(f"Unsupported LTX layer-fused linear type {type(layer).__name__}.")
+
+ if not layer.return_bias:
+ return output
+ return output, layer.bias if layer.skip_bias_add else None
+
+
+class _PhaseAdapterLinear(nn.Module):
+ """Delegate the base linear unchanged and add an enabled adapter slot."""
+
+ def __init__(
+ self,
+ base_layer: nn.Module,
+ targets: Iterable[AdapterTarget],
+ *,
+ adapter_dtype: torch.dtype,
+ layer_fused: bool = False,
+ ) -> None:
+ super().__init__()
+ self.base_layer = base_layer
+ self.adapter_dtype = adapter_dtype
+ if layer_fused and not (
+ isinstance(base_layer, nn.Linear)
+ or isinstance(getattr(base_layer, "quant_method", None), UnquantizedLinearMethod)
+ ):
+ raise ValueError("LTX layer-fused phase LoRA requires unquantized Transformer linears.")
+ self.layer_fused = layer_fused
+ pieces = [_AdapterPiece(target, _build_layout(base_layer, target)) for target in targets]
+ self.pieces = nn.ModuleList(pieces)
+ self._pieces_by_target = {piece.target: piece for piece in pieces}
+ self._enabled = False
+
+ def load_target(self, target: AdapterTarget, lora_a: torch.Tensor, lora_b: torch.Tensor) -> None:
+ try:
+ piece = self._pieces_by_target[target]
+ except KeyError as exc:
+ raise ValueError(f"Adapter target {target.module!r} is not installed on this layer.") from exc
+ piece.load(lora_a, lora_b, device=_module_device(self.base_layer), dtype=self.adapter_dtype)
+
+ def set_enabled(self, enabled: bool) -> None:
+ self._enabled = enabled
+
+ def _add_adapter_delta(self, input_: torch.Tensor, output: torch.Tensor) -> torch.Tensor:
+ if not self._enabled:
+ return output
+ if output.requires_grad:
+ output = output.clone()
+ for piece in self.pieces:
+ delta = piece.delta(input_).to(dtype=output.dtype)
+ start = piece.layout.output_start
+ output[..., start : start + piece.layout.b_output_size].add_(delta)
+ return output
+
+ @torch.no_grad()
+ def _fused_weight(self) -> torch.Tensor:
+ base_weight = self.base_layer.weight
+ fused_weight = base_weight.clone()
+ for piece in self.pieces:
+ start = piece.layout.output_start
+ base_slice = base_weight.narrow(0, start, piece.layout.b_output_size)
+ delta = piece.weight_delta()
+ if delta.shape != base_slice.shape:
+ raise ValueError(
+ f"LTX layer-fused delta shape {tuple(delta.shape)} does not match base weight slice "
+ f"{tuple(base_slice.shape)} for {piece.target.module!r}."
+ )
+ # Match official fusion exactly: materialize B@A in the adapter
+ # dtype, then add the base tensor in-place before the main GEMM.
+ delta.add_(base_slice)
+ fused_weight.narrow(0, start, piece.layout.b_output_size).copy_(delta)
+ return fused_weight
+
+ def forward(self, input_: torch.Tensor):
+ if not self._enabled:
+ return self.base_layer(input_)
+ if self.layer_fused:
+ return _forward_with_weight(self.base_layer, input_, self._fused_weight())
+ output = self.base_layer(input_)
+ if isinstance(output, tuple):
+ return (self._add_adapter_delta(input_, output[0]), *output[1:])
+ return self._add_adapter_delta(input_, output)
+
+
+class LTXPhaseAdapterRuntime:
+ """Install and operate one fixed phase adapter on an existing transformer."""
+
+ def __init__(
+ self,
+ transformer: nn.Module,
+ manifest: AdapterManifest,
+ *,
+ dtype: torch.dtype,
+ layer_fused: bool = False,
+ ) -> None:
+ self.transformer = transformer
+ self.manifest = manifest
+ self._wrappers: dict[str, _PhaseAdapterLinear] = {}
+ self._target_names: tuple[str, ...] = ()
+ self._materialized = False
+ grouped: dict[str, list[AdapterTarget]] = defaultdict(list)
+ for target in manifest.targets:
+ grouped[target.module].append(target)
+
+ original_layers = {name: self.transformer.get_submodule(name) for name in grouped}
+ for name, targets in grouped.items():
+ wrapper = _PhaseAdapterLinear(
+ original_layers[name],
+ targets,
+ adapter_dtype=dtype,
+ layer_fused=layer_fused,
+ )
+ parent_name, _, child_name = name.rpartition(".")
+ parent = self.transformer.get_submodule(parent_name) if parent_name else self.transformer
+ if child_name not in parent._modules:
+ raise ValueError(f"LTX phase adapter target {name!r} is no longer a direct module child.")
+ parent._modules[child_name] = wrapper
+ self._wrappers[name] = wrapper
+
+ self._target_names = tuple(sorted(self._wrappers, key=len, reverse=True))
+ # LTX's custom loader receives the original checkpoint names. Keep the
+ # adaptation local to that loader rather than changing AutoWeightsLoader.
+ self.transformer._phase_adapter_parameter_name = self.base_parameter_name
+ logger.info(
+ "Installed %d LTX %s phase-adapter linears",
+ len(self._wrappers),
+ "layer-fused" if layer_fused else "dynamic",
+ )
+
+ def base_parameter_name(self, name: str) -> str:
+ for target_name in self._target_names:
+ prefix = f"{target_name}."
+ if name.startswith(prefix):
+ return f"{target_name}.base_layer.{name[len(prefix) :]}"
+ return name
+
+ def finalize(self) -> None:
+ """Stream rank-local A/B buffers after base weight processing completes."""
+ if self._materialized:
+ return
+ for target, lora_a, lora_b in iter_adapter_tensors(self.manifest):
+ self._wrappers[target.module].load_target(target, lora_a, lora_b)
+ self._materialized = True
+ logger.info("Materialized %d LTX phase-adapter tensor pairs", len(self.manifest.targets))
+
+ def activate(self, adapter_slot: str | None) -> None:
+ if adapter_slot is not None:
+ if adapter_slot != LTX_DISTILLED_ADAPTER_SLOT:
+ raise ValueError(f"Unknown LTX phase adapter slot {adapter_slot!r}.")
+ if not self._materialized:
+ raise RuntimeError("LTX phase adapter data must be materialized before activation.")
+ for wrapper in self._wrappers.values():
+ wrapper.set_enabled(adapter_slot is not None)
+
+
+def build_ltx_phase_adapter(pipeline: Any) -> LTXPhaseAdapterRuntime | None:
+ """Build the fixed adapter required by a multi-phase LTX recipe."""
+ adapter_slots = {phase.adapter_slot for phase in pipeline.pipeline_recipe.phases if phase.adapter_slot is not None}
+ if not adapter_slots:
+ return None
+ if adapter_slots != {LTX_DISTILLED_ADAPTER_SLOT}:
+ raise ValueError(f"Unsupported LTX phase adapter slots: {sorted(adapter_slots)}.")
+
+ profile = pipeline.component_profile
+ if profile.distilled_lora_filename is None or profile.artifact_repo_id is None:
+ raise ValueError(f"{profile.name} does not declare the required distilled adapter artifact.")
+ if getattr(pipeline.od_config, "lora_path", None) is not None:
+ raise ValueError(
+ f"{pipeline.__class__.__name__} reserves LoRA execution for its phase adapter; "
+ "request or static LoRA composition is not supported yet."
+ )
+
+ quantization_config = getattr(pipeline.od_config, "quantization_config", None)
+ layer_fused = quantization_config is None
+ adapter_path = resolve_ltx_artifact(
+ pipeline.od_config.model,
+ profile.artifact_repo_id,
+ profile.distilled_lora_filename,
+ )
+
+ if layer_fused and pipeline.od_config.dtype is not torch.bfloat16:
+ raise ValueError("LTX layer-fused phase LoRA requires a bfloat16 pipeline dtype.")
+
+ manifest = parse_ltx_adapter(pipeline.transformer, adapter_path)
+ runtime = LTXPhaseAdapterRuntime(
+ pipeline.transformer,
+ manifest,
+ dtype=pipeline.od_config.dtype,
+ layer_fused=layer_fused,
+ )
+ return runtime
diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_recipes.py b/vllm_omni/diffusion/models/ltx2/ltx2_recipes.py
index a70f349ac..bb42bc0e9 100644
--- a/vllm_omni/diffusion/models/ltx2/ltx2_recipes.py
+++ b/vllm_omni/diffusion/models/ltx2/ltx2_recipes.py
@@ -24,6 +24,7 @@ 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"
@dataclass(frozen=True)
@@ -36,6 +37,7 @@ class LTXPhaseRecipe:
sigmas: tuple[float, ...] | None = None
noise_scale: float = 0.0
input_transform: Literal["initial", "spatial_upsample"] = "initial"
+ adapter_slot: str | None = None
allow_guidance_override: bool = True
use_official_sigma_schedule: bool = True
@@ -157,11 +159,59 @@ LTX2_DISTILLED_TWO_STAGE_RECIPE = LTXPipelineRecipe(
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
+# distilled LoRA is loaded for either phase.
+LTX23_DISTILLED_TWO_STAGE_RECIPE = LTX2_DISTILLED_TWO_STAGE_RECIPE
+
+
+def _official_two_stage_recipe(one_stage_recipe: LTXPipelineRecipe) -> LTXPipelineRecipe:
+ return LTXPipelineRecipe(
+ height=one_stage_recipe.height * 2,
+ width=one_stage_recipe.width * 2,
+ num_frames=one_stage_recipe.num_frames,
+ frame_rate=one_stage_recipe.frame_rate,
+ num_inference_steps=one_stage_recipe.num_inference_steps,
+ negative_prompt=one_stage_recipe.negative_prompt,
+ phases=(
+ LTXPhaseRecipe(
+ name="generate_lowres",
+ guidance=one_stage_recipe.request_guidance,
+ spatial_downscale=2,
+ noise_scale=1.0,
+ ),
+ 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",
+ adapter_slot=LTX_DISTILLED_ADAPTER_SLOT,
+ allow_guidance_override=False,
+ use_official_sigma_schedule=False,
+ ),
+ ),
+ video_output_phase=1,
+ # The official second stage refines video only and deliberately
+ # discards its audio result. Decode the full-context Stage-1 audio.
+ audio_output_phase=0,
+ allow_request_sigmas=False,
+ 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)
+
_PIPELINE_RECIPES: dict[tuple[str, str], LTXPipelineRecipe] = {
("one_stage", "2"): LTX2_ONE_STAGE_RECIPE,
("one_stage", "2.3"): LTX23_ONE_STAGE_RECIPE,
+ ("two_stage", "2"): LTX2_TWO_STAGE_RECIPE,
+ ("two_stage", "2.3"): LTX23_TWO_STAGE_RECIPE,
("distilled_two_stage", "2"): LTX2_DISTILLED_TWO_STAGE_RECIPE,
+ ("distilled_two_stage", "2.3"): LTX23_DISTILLED_TWO_STAGE_RECIPE,
("dmd2", "2"): LTX_POSITIVE_ONLY_RECIPE,
("dmd2", "2.3"): LTX_POSITIVE_ONLY_RECIPE,
}
diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_request.py b/vllm_omni/diffusion/models/ltx2/ltx2_request.py
index e2eef19df..91d62943a 100644
--- a/vllm_omni/diffusion/models/ltx2/ltx2_request.py
+++ b/vllm_omni/diffusion/models/ltx2/ltx2_request.py
@@ -5,8 +5,9 @@
from __future__ import annotations
+from collections.abc import Mapping
from dataclasses import dataclass, replace
-from typing import Any
+from typing import Any, Literal
import torch
@@ -49,6 +50,34 @@ class LTXRequestInputs:
return self.guidance.video.cfg_scale
+LTXCheckpointKind = Literal["distilled", "regular"]
+
+
+def validate_ltx_checkpoint(
+ scheduler_config: Mapping[str, Any],
+ *,
+ expected_kind: LTXCheckpointKind | None,
+ pipeline_name: str,
+) -> None:
+ """Ensure scheduler metadata matches the checkpoint required by a profile."""
+ if expected_kind is None:
+ return
+
+ if expected_kind == "distilled":
+ if scheduler_config.get("use_dynamic_shifting", True) or scheduler_config.get("shift_terminal") is not None:
+ raise ValueError(
+ f"{pipeline_name} requires a merged distilled LTX2 checkpoint; "
+ "the loaded scheduler metadata describes a non-distilled checkpoint."
+ )
+ return
+
+ if not scheduler_config.get("use_dynamic_shifting", False) or scheduler_config.get("shift_terminal") is None:
+ raise ValueError(
+ f"{pipeline_name} requires a regular non-distilled LTX checkpoint; "
+ "the loaded scheduler metadata describes a distilled checkpoint."
+ )
+
+
def validate_pipeline_request(
request_inputs: LTXRequestInputs,
*,
diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py b/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py
index 4b74ebcb8..cb4130259 100644
--- a/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py
+++ b/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py
@@ -47,6 +47,7 @@ from .ltx2_guidance import (
LTXGuidanceExecutor,
LTXGuidancePlan,
)
+from .ltx2_phase_adapter import LTXPhaseAdapterRuntime, build_ltx_phase_adapter
from .ltx2_recipes import (
LTXPhaseRecipe,
LTXPipelineRecipe,
@@ -137,9 +138,21 @@ class LTXRuntime(
def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = "") -> None:
del prefix
+ parallel_config = getattr(od_config, "parallel_config", None)
+ if getattr(parallel_config, "ulysses_mode", "strict") == "advanced_uaa":
+ raise ValueError(
+ 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.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:
+ 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."
+ )
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)
@@ -150,6 +163,7 @@ class LTXRuntime(
super().__init__()
self._guidance_plan = LTXGuidancePlan.build(self.pipeline_recipe.request_guidance)
initialize_pipeline_components(self, od_config)
+ self._phase_adapter: LTXPhaseAdapterRuntime | None = build_ltx_phase_adapter(self)
self.setup_diffusion_pipeline_profiler(
enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler
)
@@ -205,6 +219,8 @@ class LTXRuntime(
max_sequence_length=max_sequence_length,
)
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)
validate_pipeline_request(
request_inputs,
@@ -214,6 +230,13 @@ class LTXRuntime(
pipeline_name=self.__class__.__name__,
request_sigmas=request_sigmas,
)
+ phase_adapter = getattr(self, "_phase_adapter", None)
+ if phase_adapter is not None and any(
+ getattr(sampling, "lora_request", None) is not None for sampling in req.sampling_params_list
+ ):
+ 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)
def _run_recipe(
@@ -258,8 +281,20 @@ class LTXRuntime(
return self.decode_phase(output_phase)
def _enter_phase(self, phase: LTXPhaseRecipe) -> None:
- """Hook for a future phase-weight strategy."""
self._active_phase_name = phase.name
+ phase_adapter = getattr(self, "_phase_adapter", None)
+ if phase_adapter is None:
+ if phase.adapter_slot is not None:
+ raise RuntimeError(f"LTX phase {phase.name!r} requires adapter slot {phase.adapter_slot!r}.")
+ return
+ phase_adapter.activate(phase.adapter_slot)
+
+ def eval(self):
+ result = super().eval()
+ phase_adapter = getattr(self, "_phase_adapter", None)
+ if phase_adapter is not None:
+ phase_adapter.finalize()
+ return result
def prepare_latents(
self,
@@ -692,6 +727,7 @@ class LTXRuntime(
noise_pred_audio,
timestep,
)
+ audio = latent_ops.clear_audio_padding(audio, forward_ctx.original_audio_num_frames)
return latent_ops.LTXAVState(video=video, audio=audio)
def _unpack_and_denormalize_stage(
diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py b/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py
index 95724bd19..a59655cd3 100644
--- a/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py
+++ b/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py
@@ -355,6 +355,19 @@ def to_ltx_padding_mask(attention_mask: torch.Tensor) -> torch.Tensor:
return attention_mask
+class _LTX2ParallelAttention(Attention):
+ """Preserve LTX SP collectives while applying model-specific padding masks."""
+
+ def _run_local_attention(self, query, key, value, attn_metadata):
+ attention_mask = attn_metadata.attn_mask if attn_metadata is not None else None
+ if attention_mask is not None:
+ # LTX masks describe K/V padding. Keep this semantic local to LTX:
+ # Flash varlen applies a 2D mask to both Q and K, while SP uses a
+ # broadcast 4D key mask that only SDPA consumes correctly.
+ return self.sdpa_fallback.forward(query, key, value, attn_metadata)
+ return super()._run_local_attention(query, key, value, attn_metadata)
+
+
class LTX2AudioVideoAttnProcessor:
r"""
Processor for implementing attention (SDPA is used by default if you're using PyTorch 2.0) for the LTX-2.0 model.
@@ -396,11 +409,10 @@ class LTX2AudioVideoAttnProcessor:
return None
if self._is_sp_enabled():
- # In SP, Ulysses expects a 2D padding mask that matches query length.
- # For cross-attention, encoder sequence length != query length, so drop the mask.
- if encoder_hidden_states is not None and encoder_hidden_states.shape[1] != hidden_states.shape[1]:
- return None
- return to_ltx_padding_mask(attention_mask)
+ # Keep the full key mask replicated while strict Ulysses reshards
+ # Q/K/V. The broadcast Q dimension also works when cross-attention
+ # has different Q and K lengths, without changing generic Ulysses.
+ return to_ltx_padding_mask(attention_mask)[:, None, None, :]
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
@@ -696,7 +708,7 @@ class LTX2Attention(torch.nn.Module):
torch.nn.Dropout(dropout) if dropout > 0 else torch.nn.Identity(),
]
)
- self.attn = Attention(
+ self.attn = _LTX2ParallelAttention(
num_heads=self.query_num_heads,
head_size=dim_head,
num_kv_heads=self.kv_num_heads,
@@ -709,7 +721,15 @@ class LTX2Attention(torch.nn.Module):
# LTX-2.3: per-head gated attention
# leave unquantized for this linear
if apply_gated_attention:
- self.to_gate_logits = nn.Linear(query_dim, self.query_num_heads, bias=True)
+ self.to_gate_logits = ColumnParallelLinear(
+ query_dim,
+ heads,
+ bias=True,
+ gather_output=False,
+ return_bias=False,
+ quant_config=None,
+ prefix=f"{prefix}.to_gate_logits" if prefix else "to_gate_logits",
+ )
else:
self.to_gate_logits = None
@@ -1498,6 +1518,14 @@ class LTX2VideoTransformer3DModel(nn.Module):
_layerwise_offload_blocks_attrs = ["transformer_blocks"]
_hsdp_shard_conditions = [is_transformer_block_module]
_sp_plan: dict[str, Any] | None = None
+ stacked_params_mapping = (
+ (".attn1.to_qkv", ".attn1.to_q", "q"),
+ (".attn1.to_qkv", ".attn1.to_k", "k"),
+ (".attn1.to_qkv", ".attn1.to_v", "v"),
+ (".audio_attn1.to_qkv", ".audio_attn1.to_q", "q"),
+ (".audio_attn1.to_qkv", ".audio_attn1.to_k", "k"),
+ (".audio_attn1.to_qkv", ".audio_attn1.to_v", "v"),
+ )
packed_modules_mapping = {
"to_qkv": ["to_q", "to_k", "to_v"],
}
@@ -1828,6 +1856,7 @@ class LTX2VideoTransformer3DModel(nn.Module):
audio_sigma: torch.Tensor | None = None,
encoder_attention_mask: torch.Tensor | None = None,
audio_encoder_attention_mask: torch.Tensor | None = None,
+ audio_attention_mask: torch.Tensor | None = None,
num_frames: int | None = None,
height: int | None = None,
width: int | None = None,
@@ -1861,6 +1890,8 @@ class LTX2VideoTransformer3DModel(nn.Module):
Optional multiplicative text attention mask of shape `(batch_size, text_seq_len)`.
audio_encoder_attention_mask (`torch.Tensor`, *optional*):
Optional multiplicative text attention mask of shape `(batch_size, text_seq_len)` for audio modeling.
+ audio_attention_mask (`torch.Tensor`, *optional*):
+ Optional audio-token key padding mask shared by audio self-attention and audio-to-video attention.
num_frames (`int`, *optional*):
The number of latent video frames. Used if calculating the video coordinates for RoPE.
height (`int`, *optional*):
@@ -2044,6 +2075,8 @@ class LTX2VideoTransformer3DModel(nn.Module):
"ca_audio_rotary_emb": audio_cross_attn_rotary_emb,
"encoder_attention_mask": encoder_attention_mask,
"audio_encoder_attention_mask": audio_encoder_attention_mask,
+ "audio_self_attention_mask": audio_attention_mask,
+ "a2v_cross_attention_mask": audio_attention_mask,
"video_self_attention_perturbation_mask": perturbation_mask_for("video_self_attention", block_idx),
"audio_self_attention_perturbation_mask": perturbation_mask_for("audio_self_attention", block_idx),
"a2v_cross_attention_perturbation_mask": perturbation_mask_for("a2v_cross_attention", block_idx),
@@ -2097,16 +2130,8 @@ class LTX2VideoTransformer3DModel(nn.Module):
Returns:
Set of parameter names that were successfully loaded.
"""
- stacked_params_mapping = [
- (".attn1.to_qkv", ".attn1.to_q", "q"),
- (".attn1.to_qkv", ".attn1.to_k", "k"),
- (".attn1.to_qkv", ".attn1.to_v", "v"),
- (".audio_attn1.to_qkv", ".audio_attn1.to_q", "q"),
- (".audio_attn1.to_qkv", ".audio_attn1.to_k", "k"),
- (".audio_attn1.to_qkv", ".audio_attn1.to_v", "v"),
- ]
-
params_dict = dict(self.named_parameters())
+ phase_adapter_parameter_name = getattr(self, "_phase_adapter_parameter_name", None)
tp_size = get_tensor_model_parallel_world_size()
tp_rank = get_tensor_model_parallel_rank() if tp_size > 1 else 0
loaded_params: set[str] = set()
@@ -2127,10 +2152,12 @@ class LTX2VideoTransformer3DModel(nn.Module):
return weight
for name, loaded_weight in weights:
- for param_name, weight_name, shard_id in stacked_params_mapping:
+ for param_name, weight_name, shard_id in self.stacked_params_mapping:
if weight_name not in name:
continue
packed_name = name.replace(weight_name, param_name)
+ if callable(phase_adapter_parameter_name):
+ packed_name = phase_adapter_parameter_name(packed_name)
if packed_name not in params_dict:
continue
name = packed_name
@@ -2139,6 +2166,8 @@ class LTX2VideoTransformer3DModel(nn.Module):
weight_loader(param, loaded_weight, shard_id)
break
else:
+ if callable(phase_adapter_parameter_name):
+ name = phase_adapter_parameter_name(name)
if name not in params_dict:
logger.warning(
"Skipping transformer weight %s -- not found in model "
diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_two_stage.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_two_stage.py
index 0e5bbc558..3cb6b84f1 100644
--- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_two_stage.py
+++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_two_stage.py
@@ -1,27 +1,48 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Distilled two-stage entry points for the LTX model family."""
+"""Two-stage entry points for the LTX model family."""
-from vllm_omni.diffusion.data import OmniDiffusionConfig
+from typing import ClassVar
-from .ltx2_components import LTX2_DISTILLED_COMPONENT_PROFILE
+from .ltx2_components import (
+ LTX2_DISTILLED_COMPONENT_PROFILE,
+ LTX2_TWO_STAGE_COMPONENT_PROFILE,
+)
from .ltx2_components import (
get_ltx2_post_process_func as get_ltx2_post_process_func, # noqa: F401
)
-from .ltx2_recipes import LTX2_DISTILLED_TWO_STAGE_RECIPE
-from .pipeline_ltx2 import LTX2Pipeline
+from .ltx2_conditioning import LTXI2VConditioningMixin
+from .ltx2_recipes import (
+ LTX2_DISTILLED_TWO_STAGE_RECIPE,
+ LTX2_TWO_STAGE_RECIPE,
+)
+from .ltx2_runtime import LTXRuntime
-class LTX2DistilledPipeline(LTX2Pipeline):
- """Unified LTX-2 distilled two-stage T2V/I2V entry."""
+class _LTX2TwoStageBase(LTXI2VConditioningMixin, LTXRuntime):
+ """Static discovery contract shared by the two-stage public entries."""
- pipeline_kind = "distilled_two_stage"
- component_profile = LTX2_DISTILLED_COMPONENT_PROFILE
- pipeline_recipe = LTX2_DISTILLED_TWO_STAGE_RECIPE
+ _dit_modules: ClassVar[list[str]] = list(LTX2_TWO_STAGE_COMPONENT_PROFILE.dit_modules)
+ _encoder_modules: ClassVar[list[str]] = list(LTX2_TWO_STAGE_COMPONENT_PROFILE.encoder_modules)
+ _vae_modules: ClassVar[list[str]] = list(LTX2_TWO_STAGE_COMPONENT_PROFILE.vae_modules)
+ _resident_modules: ClassVar[list[str]] = list(LTX2_TWO_STAGE_COMPONENT_PROFILE.resident_modules)
supports_request_batch = False
support_image_input = True
unified_text_image_entry = True
- def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
- super().__init__(od_config=od_config, prefix=prefix)
+
+class LTX2TwoStagePipeline(_LTX2TwoStageBase):
+ """Regular checkpoint with low-resolution generation and LoRA refinement."""
+
+ pipeline_kind = "two_stage"
+ component_profile = LTX2_TWO_STAGE_COMPONENT_PROFILE
+ pipeline_recipe = LTX2_TWO_STAGE_RECIPE
+
+
+class LTX2DistilledPipeline(_LTX2TwoStageBase):
+ """Unified LTX-2/LTX-2.3 full-distilled two-stage T2V/I2V entry."""
+
+ pipeline_kind = "distilled_two_stage"
+ component_profile = LTX2_DISTILLED_COMPONENT_PROFILE
+ pipeline_recipe = LTX2_DISTILLED_TWO_STAGE_RECIPE
diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py
index 75ef7391a..0e814886b 100644
--- a/vllm_omni/diffusion/registry.py
+++ b/vllm_omni/diffusion/registry.py
@@ -71,6 +71,11 @@ _DIFFUSION_MODELS = {
"pipeline_ltx2",
"LTX2Pipeline",
),
+ "LTX2TwoStagePipeline": (
+ "ltx2",
+ "pipeline_ltx2_two_stage",
+ "LTX2TwoStagePipeline",
+ ),
"LTX2DistilledPipeline": (
"ltx2",
"pipeline_ltx2_two_stage",
@@ -516,6 +521,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
"WanPipeline": "get_wan22_post_process_func",
"WanVACEPipeline": "get_wan22_vace_post_process_func",
"LTX2Pipeline": "get_ltx2_post_process_func",
+ "LTX2TwoStagePipeline": "get_ltx2_post_process_func",
"LTX2DistilledPipeline": "get_ltx2_post_process_func",
"LTX2T2VDMD2Pipeline": "get_ltx2_post_process_func",
"LTX2I2VDMD2Pipeline": "get_ltx2_post_process_func",
diff --git a/vllm_omni/model_extras/registry.py b/vllm_omni/model_extras/registry.py
index 8d7ee1e29..bc923ea86 100644
--- a/vllm_omni/model_extras/registry.py
+++ b/vllm_omni/model_extras/registry.py
@@ -205,6 +205,7 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = {
}
for model_class_name in (
"LTX2Pipeline",
+ "LTX2TwoStagePipeline",
"LTX2DistilledPipeline",
)
},