[Diffusion] Add LTX-2.5 Diffusion VAE decoder support (#6189)

Signed-off-by: mglyn <1203789601@qq.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

pyproject.toml

@@ -71,7 +71,6 @@ dev = [
     "voxcpm==2.0.3", # for VoxCPM2 path-locating in tests
     "torchdiffeq==0.2.5", # for covo_audio Test
     "FlagEmbedding==1.4.0", # for hunyuan_image3 tests
-    "kernels==0.14.1", # for accuracy tests comparison with diffusers output (0.15.0 incompatible with diffusers 0.38.0)
     "s3tokenizer==0.3.0", # for tts test
     # Provides top-level cosyvoice2; do not co-install stepaudio2-minicpmo,
     # whose top-level stepaudio2 package conflicts with this distribution.

recipes/LTX/LTX-2.5.md

@@ -10,14 +10,22 @@ checkpoint for T2V and first-frame I2V. Every output includes synchronized
 48 kHz stereo audio.
 
 The raw [`Lightricks/LTX-2.5`](https://huggingface.co/Lightricks/LTX-2.5)
-repository is not directly loadable with `--model`. It supplies the official
-upsampler and LoRA sidecars used by the Full/SFT two-stage pipeline, so accept
-both model licenses and authenticate before first use.
+repository is not directly loadable with `--model`. It supplies the canonical
+Diffusion VAE checkpoint and the official upsampler and LoRA sidecars used by
+the Full/SFT two-stage pipeline, so accept both model licenses and authenticate
+before first use.
+
+The vLLM-Omni integration code is provided under the project's
+[Apache-2.0 license](../../LICENSE). The LTX-2.5 model weights are separate
+artifacts covered by the
+[LTX-2.x Community License](https://github.com/Lightricks/LTX-2/blob/main/LICENSE.md),
+not the vLLM-Omni license. Users are responsible for reviewing and complying
+with the model license before downloading or using the weights.
 
 ## Choose a pipeline
 
 | Pipeline | Mode | Default size | Schedule |
-|---|---|---:|---:|
+| --- | --- | ---: | ---: |
 | `LTX2Pipeline` | Full/SFT one-stage | 960x544 | 30 steps |
 | `LTX2TwoStagePipeline` | Full/SFT two-stage | 1920x1088 | 30 + 3 steps |
 | `LTX2DistilledOneStagePipeline` | Distilled one-stage | 960x544 | 8 steps |
@@ -28,6 +36,40 @@ latent upsampler, and run a three-step refinement stage. Select the class with
 `--model-class-name`; no `--task-type` flag is required. Supplying one initial
 image selects I2V, while omitting it selects T2V.
 
+## Video decoder
+
+LTX-2.5 uses the canonical Native Diffusion VAE decoder by default. To opt in
+to the legacy convolutional VAE, set the startup-only stage-0 model extra
+`ltx2_use_conv_vae: true`. The choice applies to all requests and all four
+pipeline classes above. DiffVAE loads its NATTEN kernel from Hugging Face Hub
+during startup. `decode_timestep` and `decode_noise_scale` condition only the
+legacy ConvVAE and have no effect when DiffVAE is selected.
+
+For offline Python usage:
+
+```python
+omni = Omni(
+    model="Lightricks/LTX-2.5-Diffusers",
+    model_class_name="LTX2Pipeline",
+    stage_overrides='{"0":{"extras":{"ltx2_use_conv_vae":true}}}',
+)
+```
+
+For online serving:
+
+```bash
+vllm serve Lightricks/LTX-2.5-Diffusers \
+  --omni \
+  --model-class-name LTX2Pipeline \
+  --stage-overrides '{"0":{"extras":{"ltx2_use_conv_vae":true}}}'
+```
+
+Both decoders are untiled by default. Set `vae_use_tiling` for memory-saving
+serial tiling; DiffVAE tiles only above 80 frames or 768 pixels in either spatial
+dimension. DiffVAE also supports distributed VAE decode; see the
+[VAE Parallelism Guide](../../docs/user_guide/diffusion/parallelism/vae_parallelism.md).
+DiffVAE is decoder-only, so I2V still uses the convolutional VAE encoder.
+
 ## Prerequisites
 
 ```bash
@@ -42,7 +84,7 @@ Install matching vLLM and vLLM-Omni versions, and ensure `ffmpeg` and
 ## Hardware
 
 | GPU | Status | Recommended scope |
-|---|---|---|
+| --- | --- | --- |
 | NVIDIA B300 | Verified | All four canonical pipelines |
 | NVIDIA B200 or H200 | Capacity-based recommendation; not yet verified | All four canonical pipelines |
 | NVIDIA GB200 or GB300 | Capacity-based recommendation; not yet verified | All four canonical pipelines |

requirements/common.txt

@@ -13,6 +13,8 @@ transformers >= 5.10.1, < 5.15
 av>=14.0.0
 omegaconf>=2.3.0
 diffusers==0.40.0
+# Hub kernels, including LTX-2.5 NATTEN
+kernels==0.14.1
 safetensors>=0.8.0
 accelerate==1.12.0
 soundfile>=0.13.1

tests/config/test_omni_config.py

@@ -940,6 +940,35 @@ def test_structured_diffusion_config_rejects_non_boolean_compile_dynamic():
         omni_config_module._DiffusionConfigProjection(diffusion_compile_dynamic="false")
 
 
+def test_from_pipeline_config_routes_ltx2_conv_vae_extra(tmp_path):
+    deploy_path = tmp_path / "dreamzero_ltx2_extras.yaml"
+    deploy_path.write_text(
+        "\n".join(
+            [
+                "pipeline: dreamzero",
+                "async_chunk: false",
+                "stages:",
+                "  - stage_id: 0",
+                "    extras:",
+                "      ltx2_use_conv_vae: true",
+            ]
+        )
+    )
+
+    stage = _from_pipeline_key("dreamzero", deploy_config_path=str(deploy_path)).stage_by_id(0)
+
+    assert stage.diffusion_config.extras["ltx2_use_conv_vae"] is True
+
+
+def test_stage_override_routes_ltx2_conv_vae_extra():
+    stage = _from_pipeline_key(
+        "dreamzero",
+        cli_overrides={"stage_0_extras": {"ltx2_use_conv_vae": True}},
+    ).stage_by_id(0)
+
+    assert stage.diffusion_config.extras["ltx2_use_conv_vae"] is True
+
+
 def test_structured_diffusion_config_rejects_invalid_compile_granularity():
     with pytest.raises(ValidationError, match="diffusion_compile_granularity"):
         omni_config_module._DiffusionConfigProjection(diffusion_compile_granularity="block")

tests/diffusion/models/ltx2/test_ltx25_pipeline.py

@@ -1,11 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """Focused LTX-2.5 pipeline correctness tests."""
 
 import json
 from dataclasses import replace
 from types import SimpleNamespace
+from typing import Any
 
 import pytest
 import torch
@@ -170,7 +171,10 @@ def _resolve_request_inputs_for_test(
 
 
 def test_ltx25_missing_gemma4_recommends_supported_transformers_range(monkeypatch):
-    pipe = SimpleNamespace(component_profile=replace(LTX25_FULL_COMPONENT_PROFILE, text_encoder_cls=None))
+    pipe = SimpleNamespace(
+        component_profile=replace(LTX25_FULL_COMPONENT_PROFILE, text_encoder_cls=None),
+        use_diffusion_decoder=True,
+    )
     od_config = SimpleNamespace(model="Lightricks/LTX-2.5-Diffusers", dtype=torch.bfloat16)
 
     monkeypatch.setattr(ltx2_components, "get_local_device", lambda: torch.device("cpu"))
@@ -187,7 +191,19 @@ def test_ltx25_missing_gemma4_recommends_supported_transformers_range(monkeypatc
 
 def test_ltx_converted_component_loading_propagates_revision(monkeypatch):
     revision = "pinned-revision"
-    calls = {"components": []}
+    calls: dict[str, Any] = {"components": []}
+    natten_processor = object()
+
+    class FakeDiffusionDecoder:
+        def set_attn_processor(self, processor):
+            calls["diffusion_decoder_processor"] = processor
+
+        def enable_tiling(self):
+            calls["diffusion_decoder_tiling"] = True
+
+        def set_parallel_size(self, parallel_size, mode):
+            calls["diffusion_decoder_parallel"] = (parallel_size, mode)
+
     profile = replace(
         LTX25_DISTILLED_COMPONENT_PROFILE,
         text_encoder_cls=object,
@@ -196,12 +212,19 @@ def test_ltx_converted_component_loading_propagates_revision(monkeypatch):
         vocoder_fallback_cls=None,
         scheduler_use_dynamic_shifting=False,
     )
-    pipeline = SimpleNamespace(component_profile=profile, pipeline_kind="distilled_two_stage")
+    pipeline = SimpleNamespace(
+        component_profile=profile,
+        pipeline_kind="distilled_two_stage",
+        use_diffusion_decoder=True,
+    )
     od_config = SimpleNamespace(
         model="org/converted-ltx25",
         revision=revision,
         dtype=torch.bfloat16,
         quantization_config=None,
+        extras={},
+        vae_use_tiling=False,
+        parallel_config=SimpleNamespace(vae_patch_parallel_size=2, vae_parallel_mode="tile"),
     )
 
     monkeypatch.setattr(ltx2_components, "get_local_device", lambda: torch.device("cpu"))
@@ -227,6 +250,10 @@ def test_ltx_converted_component_loading_propagates_revision(monkeypatch):
             )
         return object()
 
+    def fake_native_diffusion_decoder(model, **kwargs):
+        calls["native_diffusion_decoder"] = (model, kwargs)
+        return FakeDiffusionDecoder()
+
     def fake_transformer_config(model, subfolder, local_files_only, *, revision):
         calls["transformer_config"] = (model, subfolder, local_files_only, revision)
         return {"component": "transformer"}
@@ -234,7 +261,13 @@ def test_ltx_converted_component_loading_propagates_revision(monkeypatch):
     monkeypatch.setattr(ltx2_components, "prefetch_subfolders", fake_prefetch)
     monkeypatch.setattr(ltx2_components.AutoTokenizer, "from_pretrained", fake_tokenizer)
     monkeypatch.setattr(ltx2_components, "_load_component", fake_component)
+    monkeypatch.setattr(ltx2_components, "_load_ltx25_native_diffusion_decoder", fake_native_diffusion_decoder)
     monkeypatch.setattr(ltx2_components, "load_transformer_config", fake_transformer_config)
+    monkeypatch.setattr(
+        ltx2_components,
+        "LTX2VideoVaeNeighborhoodNattenProcessor",
+        lambda: natten_processor,
+    )
     monkeypatch.setattr(
         ltx2_components,
         "create_transformer_from_config",
@@ -261,6 +294,14 @@ def test_ltx_converted_component_loading_propagates_revision(monkeypatch):
         "latent_upsampler",
     }
     assert all(call[3]["revision"] == revision for call in calls["components"])
+    assert "diffusion_decoder" not in calls["prefetch"][1]
+    assert calls["native_diffusion_decoder"] == (
+        od_config.model,
+        {"local_files_only": False, "dtype": torch.bfloat16, "revision": revision},
+    )
+    assert calls["diffusion_decoder_processor"] is natten_processor
+    assert calls["diffusion_decoder_tiling"] is True
+    assert calls["diffusion_decoder_parallel"] == (2, "tile")
     assert calls["transformer_config"] == (
         od_config.model,
         profile.transformer_subfolder,
@@ -270,6 +311,85 @@ def test_ltx_converted_component_loading_propagates_revision(monkeypatch):
     assert calls["scheduler"][1]["revision"] == revision
 
 
+def test_ltx25_native_diffusion_decoder_uses_diffusers_config_and_canonical_artifact(monkeypatch):
+    revision = "diffusers-revision"
+    config = {"decoder_stage_channels": [8, 8, 8, 8, 8]}
+    calls: dict[str, Any] = {}
+
+    class FakeDecoder:
+        def init_distributed(self):
+            calls["init_distributed"] = True
+
+    expected = FakeDecoder()
+
+    def fake_load_config(model, **kwargs):
+        calls["config"] = (model, kwargs)
+        return config
+
+    def fake_resolve(model, repo_id, filename, **kwargs):
+        calls["artifact"] = (model, repo_id, filename, kwargs)
+        return "/models/native-diffvae.safetensors"
+
+    def fake_load(cls, path, passed_config, dtype):
+        calls["load"] = (path, passed_config, dtype)
+        return expected
+
+    monkeypatch.setattr(
+        ltx2_components.DistributedLTX2VideoDiffusionDecoderModel,
+        "load_config",
+        staticmethod(fake_load_config),
+    )
+    monkeypatch.setattr(ltx2_components, "resolve_ltx_artifact", fake_resolve)
+    monkeypatch.setattr(
+        ltx2_components.DistributedLTX2VideoDiffusionDecoderModel,
+        "from_ltx25_native_checkpoint",
+        classmethod(fake_load),
+    )
+
+    actual = ltx2_components._load_ltx25_native_diffusion_decoder(
+        "Lightricks/LTX-2.5-Diffusers",
+        local_files_only=False,
+        dtype=torch.bfloat16,
+        revision=revision,
+    )
+
+    assert actual is expected
+    assert calls["config"] == (
+        "Lightricks/LTX-2.5-Diffusers",
+        {
+            "subfolder": "diffusion_decoder",
+            "local_files_only": False,
+            "revision": revision,
+        },
+    )
+    assert calls["artifact"] == (
+        "Lightricks/LTX-2.5-Diffusers",
+        "Lightricks/LTX-2.5",
+        "vae/ltx-2.5-video-vae-bf16.safetensors",
+        {
+            "model_revision": revision,
+            "artifact_revision": ltx2_components.LTX25_NATIVE_ARTIFACT_REVISION,
+        },
+    )
+    assert calls["load"] == ("/models/native-diffvae.safetensors", config, torch.bfloat16)
+    assert calls["init_distributed"] is True
+
+
+def test_ltx25_natten_failure_has_actionable_omni_remedy(monkeypatch):
+    class MissingNattenProcessor:
+        def __init__(self):
+            raise ImportError("kernels is unavailable")
+
+    monkeypatch.setattr(
+        ltx2_components,
+        "LTX2VideoVaeNeighborhoodNattenProcessor",
+        MissingNattenProcessor,
+    )
+
+    with pytest.raises(RuntimeError, match=r"kernels==0\.14\.1.*supported GPU.*allow Hub access"):
+        ltx2_components._create_ltx25_natten_processor()
+
+
 def test_ltx_checkpoint_explicit_version_precedes_structural_heuristics(tmp_path):
     (tmp_path / "model_index.json").write_text(json.dumps({"model_version": "2.3.1"}))
     (tmp_path / "transformer").mkdir()
@@ -411,6 +531,46 @@ def test_ltx25_four_public_pipeline_semantics_are_disjoint():
     assert LTX2DistilledTwoStagePipeline.pipeline_kind == "distilled_two_stage"
 
 
+@pytest.mark.parametrize(
+    "pipeline_cls",
+    [
+        LTX2Pipeline,
+        LTX2TwoStagePipeline,
+        LTX2DistilledOneStagePipeline,
+        LTX2DistilledTwoStagePipeline,
+    ],
+)
+@pytest.mark.parametrize(
+    ("extras", "use_diffusion_decoder"),
+    [({}, True), ({"ltx2_use_conv_vae": True}, False)],
+)
+def test_ltx25_all_public_pipelines_default_to_diffvae_with_conv_opt_in(
+    tmp_path, monkeypatch, pipeline_cls, extras, use_diffusion_decoder
+):
+    (tmp_path / "model_index.json").write_text(
+        json.dumps({"text_encoder": ["transformers", "Gemma4UnifiedForConditionalGeneration"]})
+    )
+
+    def stub_components(pipe, od_config):
+        pipe.od_config = od_config
+        pipe.vae_spatial_compression_ratio = 32
+
+    monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)
+    monkeypatch.setattr(ltx2_runtime, "build_ltx_phase_adapter", lambda _pipe: None)
+    monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", lambda *_args, **_kwargs: None)
+
+    pipe = pipeline_cls(
+        od_config=SimpleNamespace(
+            model=str(tmp_path),
+            extras=extras,
+            enable_diffusion_pipeline_profiler=False,
+        )
+    )
+
+    assert pipe.use_diffusion_decoder is use_diffusion_decoder
+    assert pipe._vae_modules.count("diffusion_decoder") == int(use_diffusion_decoder)
+
+
 @pytest.mark.parametrize(
     ("pipeline_kind", "expected"),
     [
@@ -673,7 +833,7 @@ def test_ltx25_distilled_two_stage_executes_custom_phase_schedules():
         max_sequence_length=16,
     )
     prompt_context = object()
-    phase_calls = []
+    phase_calls: list[Any] = []
 
     def resolve_request_inputs(req, **kwargs):
         return request_inputs

tests/diffusion/models/ltx2/test_ltx2_conditioning.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """Unit tests for LTX image-to-video input and conditioning behavior."""
 
@@ -258,6 +258,39 @@ class TestLTXImageToVideoConditioning:
 
         assert LTX2Pipeline.support_image_input is True
 
+    def test_ltx25_i2v_encode_uses_diffusion_decoder_statistics(self, monkeypatch):
+        import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as ltx2_conditioning
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
+
+        pipe = object.__new__(LTX2Pipeline)
+        torch.nn.Module.__init__(pipe)
+        pipe.use_diffusion_decoder = True
+        pipe.vae = SimpleNamespace(
+            encode=lambda image: image,
+            latents_mean=torch.full((2,), 100.0),
+            latents_std=torch.full((2,), 10.0),
+            config=SimpleNamespace(scaling_factor=1.0),
+        )
+        pipe.diffusion_decoder = SimpleNamespace(
+            latents_mean=torch.tensor([1.0, 3.0]),
+            latents_std=torch.tensor([2.0, 4.0]),
+            config=SimpleNamespace(scaling_factor=2.0),
+        )
+        monkeypatch.setattr(
+            ltx2_conditioning,
+            "retrieve_latents",
+            lambda *_args, **_kwargs: torch.tensor([[[[[3.0]]], [[[7.0]]]]]),
+        )
+
+        actual = pipe._encode_i2v_image_latents(
+            torch.zeros(1, 3, 1, 1),
+            batch_size=1,
+            generator=None,
+            dtype=torch.float32,
+        )
+
+        torch.testing.assert_close(actual, torch.full((1, 2, 1, 1, 1), 2.0))
+
     def test_ltx23_i2v_rejects_multi_image_prompt_list(self):
         from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
 

tests/diffusion/models/ltx2/test_ltx2_latents.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from types import SimpleNamespace
 
 import pytest
@@ -47,6 +50,39 @@ def test_prepare_video_latents_matches_official_values_and_token_major_layout():
     assert actual.stride()[1:] == (1, actual.shape[1])
 
 
+def test_prepare_video_latents_uses_diffusion_decoder_statistics():
+    pipeline = _make_pipeline(LTX2Pipeline)
+    pipeline.use_diffusion_decoder = True
+    pipeline.vae_spatial_compression_ratio = 1
+    pipeline.vae_temporal_compression_ratio = 1
+    pipeline.transformer_spatial_patch_size = 1
+    pipeline.transformer_temporal_patch_size = 1
+    pipeline.vae = SimpleNamespace(
+        latents_mean=torch.full((2,), 100.0),
+        latents_std=torch.full((2,), 10.0),
+        config=SimpleNamespace(scaling_factor=1.0),
+    )
+    pipeline.diffusion_decoder = SimpleNamespace(
+        latents_mean=torch.tensor([1.0, 3.0]),
+        latents_std=torch.tensor([2.0, 4.0]),
+        config=SimpleNamespace(scaling_factor=2.0),
+    )
+
+    actual = pipeline.prepare_latents(
+        batch_size=1,
+        num_channels_latents=2,
+        height=1,
+        width=1,
+        num_frames=1,
+        noise_scale=0.0,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
+        latents=torch.tensor([[[[[3.0]]], [[[7.0]]]]]),
+    )
+
+    torch.testing.assert_close(actual, torch.full((1, 1, 2), 2.0))
+
+
 def test_prepare_audio_latents_samples_directly_in_packed_token_space():
     pipeline = _make_pipeline(LTX2Pipeline)
 

tests/diffusion/models/ltx2/test_ltx2_phase_adapter.py

@@ -1,9 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """Unit tests for the fixed LTX refinement-phase adapter."""
 
 from types import SimpleNamespace
+from typing import Any
 
 import pytest
 import torch
@@ -148,10 +149,13 @@ def test_ltx_phase_adapter_selects_execution_from_quantization(
 ):
     pipeline = _pipeline(dtype=dtype, quantization_config=quantization_config)
     manifest = object()
-    captured = {}
+    captured: dict[str, Any] = {}
 
     def resolve_artifact(*_args, **kwargs):
-        assert not kwargs
+        assert kwargs == {
+            "model_revision": None,
+            "artifact_revision": None,
+        }
         return "adapter"
 
     def parse(transformer, path):

tests/diffusion/models/ltx2/test_ltx2_pipeline.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """Unit tests for the shared LTX pipeline runtime and public contracts."""
 
@@ -8,6 +8,7 @@ import os
 import tempfile
 from dataclasses import replace
 from types import SimpleNamespace
+from typing import Any
 
 import pytest
 import torch
@@ -165,7 +166,19 @@ 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):
+@pytest.mark.parametrize(
+    ("model", "repo_id", "expected_revision"),
+    [
+        ("Lightricks/test", "Lightricks/test", "model-revision"),
+        ("Lightricks/test-Diffusers", "Lightricks/test", "artifact-revision"),
+    ],
+)
+def test_ltx_artifact_uses_source_revision_and_hub_fallback(
+    monkeypatch,
+    model,
+    repo_id,
+    expected_revision,
+):
     filename = "ltx-sidecar.safetensors"
     calls = []
 
@@ -176,9 +189,22 @@ def test_ltx_artifact_falls_back_to_unpinned_hub_download(tmp_path, monkeypatch)
     monkeypatch.setattr(ltx2_components, "hf_hub_download", fake_download)
 
     assert (
-        resolve_ltx_artifact(str(tmp_path / "model"), "Lightricks/test", filename) == "/cache/ltx-sidecar.safetensors"
+        resolve_ltx_artifact(
+            model,
+            repo_id,
+            filename,
+            model_revision="model-revision",
+            artifact_revision="artifact-revision",
+        )
+        == "/cache/ltx-sidecar.safetensors"
     )
-    assert calls == [{"repo_id": "Lightricks/test", "filename": filename}]
+    assert calls == [
+        {
+            "repo_id": repo_id,
+            "filename": filename,
+            "revision": expected_revision,
+        }
+    ]
 
 
 def test_ltx_artifact_prefers_model_root(tmp_path, monkeypatch):
@@ -187,7 +213,41 @@ def test_ltx_artifact_prefers_model_root(tmp_path, monkeypatch):
     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)
+    assert resolve_ltx_artifact(
+        str(tmp_path),
+        "Lightricks/test",
+        filename,
+        model_revision="unused-model-revision",
+        artifact_revision="unused-artifact-revision",
+    ) == str(expected)
+
+
+def test_ltx_artifact_local_model_missing_sidecar_falls_back_to_hub(tmp_path, monkeypatch):
+    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),
+            "Lightricks/test",
+            "ltx-sidecar.safetensors",
+            model_revision="local-model-revision",
+            artifact_revision="artifact-revision",
+        )
+        == "/cache/ltx-sidecar.safetensors"
+    )
+    assert calls == [
+        {
+            "repo_id": "Lightricks/test",
+            "filename": "ltx-sidecar.safetensors",
+            "revision": "artifact-revision",
+        }
+    ]
 
 
 @pytest.mark.parametrize(
@@ -900,7 +960,7 @@ def test_ltx_two_stage_executes_declarative_i2v_phase_plan(pipeline_cls):
         max_sequence_length=16,
     )
     prompt_context = object()
-    phase_calls = []
+    phase_calls: list[Any] = []
 
     def resolve_request_inputs(req, **kwargs):
         return request_inputs
@@ -959,7 +1019,7 @@ def test_ltx_two_stage_executes_declarative_i2v_phase_plan(pipeline_cls):
     pipeline.vae_spatial_compression_ratio = 32
     pipeline.vae_temporal_compression_ratio = 8
     pipeline.latent_upsampler = FakeUpsampler()
-    activated_slots = []
+    activated_slots: list[Any] = []
     if pipeline_cls is LTX2TwoStagePipeline:
         pipeline._phase_adapter = SimpleNamespace(activate=activated_slots.append)
     object.__setattr__(pipeline, "_resolve_request_inputs", resolve_request_inputs)
@@ -1133,6 +1193,59 @@ def test_ltx_two_stage_recipe_rejects_sigmas():
         )
 
 
+def test_ltx_request_applies_diffvae_minimum_spatial_size_only_when_requested():
+    request = LTXRequestInputs(
+        prompt="prompt",
+        negative_prompt="",
+        height=128,
+        width=128,
+        num_frames=9,
+        frame_rate=24.0,
+        num_inference_steps=30,
+        guidance=LTX2_ONE_STAGE_RECIPE.request_guidance,
+        num_videos_per_prompt=1,
+        generator=None,
+        latents=None,
+        audio_latents=None,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+        prompt_attention_mask=None,
+        negative_prompt_attention_mask=None,
+        decode_timestep=0.0,
+        decode_noise_scale=None,
+        output_type="np",
+        max_sequence_length=16,
+    )
+
+    with pytest.raises(ValueError, match="DiffVAE requires video latent spatial dimensions of at least 7x7"):
+        validate_pipeline_request(
+            request,
+            pipeline_recipe=LTX2_ONE_STAGE_RECIPE,
+            vae_spatial_compression_ratio=32,
+            vae_temporal_compression_ratio=8,
+            pipeline_name="LTX2Pipeline",
+            min_video_latent_spatial_size=(7, 7),
+        )
+
+    # ConvVAE does not use NATTEN, so its existing 128x128 boundary remains valid.
+    validate_pipeline_request(
+        request,
+        pipeline_recipe=LTX2_ONE_STAGE_RECIPE,
+        vae_spatial_compression_ratio=32,
+        vae_temporal_compression_ratio=8,
+        pipeline_name="LTX2Pipeline",
+    )
+
+    validate_pipeline_request(
+        replace(request, height=256, width=256),
+        pipeline_recipe=LTX2_ONE_STAGE_RECIPE,
+        vae_spatial_compression_ratio=32,
+        vae_temporal_compression_ratio=8,
+        pipeline_name="LTX2Pipeline",
+        min_video_latent_spatial_size=(7, 7),
+    )
+
+
 @pytest.mark.parametrize(
     ("direct_kwargs", "sampling_kwargs", "error"),
     [
@@ -1284,7 +1397,7 @@ def test_ltx_official_two_stage_recipes_only_vary_by_model_defaults(recipe, step
 
 
 def test_ltx_two_stage_entries_select_the_official_adapter_slots():
-    phase_switches = []
+    phase_switches: list[Any] = []
     ordinary_pipeline = object.__new__(LTX2TwoStagePipeline)
     ordinary_pipeline._phase_adapter = SimpleNamespace(activate=phase_switches.append)
     for phase in LTX2_TWO_STAGE_RECIPE.phases:

tests/diffusion/models/ltx2/test_ltx2_vae.py

@@ -1,9 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
-"""Unit tests for LTX-2.3 VAE tiling and distributed decode behavior."""
+"""Unit tests for LTX video VAE tiling and distributed decode behavior."""
 
 from types import SimpleNamespace
+from typing import Any
 
 import pytest
 import torch
@@ -36,6 +37,441 @@ def test_ltx_base_vocoder_keeps_native_dtype(monkeypatch):
     assert output.dtype == torch.bfloat16
 
 
+class TestLTXDiffusionDecoder:
+    @pytest.mark.parametrize("mode", ["spatial_shard_height", "spatial_shard_width"])
+    def test_distributed_diffusion_decoder_rejects_non_tile_parallel_modes(self, mode):
+        from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder_distributed import (
+            DistributedLTX2VideoDiffusionDecoderModel,
+        )
+
+        model = object.__new__(DistributedLTX2VideoDiffusionDecoderModel)
+        torch.nn.Module.__init__(model)
+        model.distributed_executor = SimpleNamespace(set_parallel_size=lambda *_args, **_kwargs: None)
+
+        with pytest.raises(ValueError, match="only supports vae_parallel_mode='tile'"):
+            model.set_parallel_size(2, mode=mode)
+
+    def test_distributed_diffusion_decoder_accepts_tile_parallel_mode(self):
+        from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder_distributed import (
+            DistributedLTX2VideoDiffusionDecoderModel,
+        )
+
+        calls = []
+        model = object.__new__(DistributedLTX2VideoDiffusionDecoderModel)
+        torch.nn.Module.__init__(model)
+        model.distributed_executor = SimpleNamespace(
+            set_parallel_size=lambda parallel_size, mode: calls.append((parallel_size, mode))
+        )
+
+        model.set_parallel_size(4, mode="tile")
+
+        assert calls == [(4, "tile")]
+
+    def test_short_clip_keeps_stage5_temporal_context_then_crops_output(self):
+        from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder import (
+            LTX2VideoDiffusionDecoder3d,
+            LTX2VideoDiffusionDecoderModel,
+        )
+
+        class IdentityAttention:
+            def build_block_mask(self, _hidden_states):
+                return None
+
+        class IdentityBlock(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.attn = IdentityAttention()
+
+            def forward(self, hidden_states, _block_mask):
+                return hidden_states
+
+        class IdentityUpsample(torch.nn.Module):
+            def forward(self, hidden_states, *, drop_leading_frame):
+                assert drop_leading_frame
+                return hidden_states
+
+        decoder = object.__new__(LTX2VideoDiffusionDecoder3d)
+        torch.nn.Module.__init__(decoder)
+        decoder.det_stages = torch.nn.ModuleList([torch.nn.ModuleList([IdentityBlock()])])
+        decoder.upsamples = torch.nn.ModuleList([IdentityUpsample()])
+        decoder.trailing_pad_latent_frames = 2
+        decoder.temporal_compression_ratio = 8
+        decoder.stage5_kernel = (11, 3, 3)
+
+        # 17 frames contain one real frame followed by 16 propagated ghost
+        # frames. The reference keeps 11 for stage-5 NATTEN instead of
+        # cropping the tensor immediately back to one.
+        context = decoder.forward_stage_4(torch.zeros(1, 17, 3, 3, 4))
+        assert context.shape == (1, 11, 3, 3, 4)
+
+        class FakeDecoder(torch.nn.Module):
+            def forward(self, z, *, generator, num_inference_steps):
+                assert generator is None
+                assert num_inference_steps is None
+                return torch.zeros(z.shape[0], 3, 11, z.shape[3] * 2, z.shape[4] * 2)
+
+        model = object.__new__(LTX2VideoDiffusionDecoderModel)
+        torch.nn.Module.__init__(model)
+        model.decoder = FakeDecoder()
+        model.use_tiling = False
+        model.tile_sample_min_num_frames = 80
+        model.tile_sample_min_height = 768
+        model.tile_sample_min_width = 768
+        model.temporal_compression_ratio = 8
+        model.spatial_compression_ratio = 2
+
+        output = model.decode(torch.zeros(1, 4, 1, 2, 3), generator=None, return_dict=False)[0]
+        assert output.shape == (1, 3, 1, 4, 6)
+
+    @pytest.mark.parametrize(
+        ("model_version", "extras", "expected"),
+        [
+            ("2.5", {}, True),
+            ("2.5", {"ltx2_use_conv_vae": True}, False),
+            ("2.3", {}, False),
+        ],
+    )
+    def test_ltx25_diffusion_decoder_is_default_with_conv_vae_opt_in(self, model_version, extras, expected):
+        from vllm_omni.diffusion.models.ltx2.ltx2_components import _ltx2_use_diffusion_decoder
+
+        assert _ltx2_use_diffusion_decoder(SimpleNamespace(extras=extras), model_version) is expected
+
+    def test_conv_vae_opt_in_rejects_non_boolean_model_extra(self):
+        from vllm_omni.diffusion.models.ltx2.ltx2_components import _ltx2_use_diffusion_decoder
+
+        with pytest.raises(TypeError, match="ltx2_use_conv_vae"):
+            _ltx2_use_diffusion_decoder(SimpleNamespace(extras={"ltx2_use_conv_vae": "true"}), "2.5")
+
+    def test_native_diffusion_decoder_conversion_splits_qkv_and_folds_gates(self):
+        from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder import (
+            convert_ltx25_native_diffusion_decoder_state_dict,
+        )
+
+        qkv_weight = torch.arange(24, dtype=torch.float32).reshape(6, 4)
+        qkv_bias = torch.arange(6, dtype=torch.float32)
+        projection_weight = torch.arange(4, dtype=torch.float32).reshape(2, 2)
+        projection_bias = torch.ones(2)
+        gate = torch.tensor([2.0, 3.0])
+        mean = torch.tensor([0.1, 0.2])
+        std = torch.tensor([0.3, 0.4])
+        native = {
+            "encoder.ignored": torch.ones(1),
+            "per_channel_statistics.mean-of-means": mean,
+            "vae.per_channel_statistics.std-of-means": std,
+            "decoder.type_emb": torch.ones(2),
+            "decoder.coarse_head.weight": torch.ones(1),
+            "decoder.det_stages.0.0.attn.qkv.weight": qkv_weight,
+            "decoder.det_stages.0.0.attn.qkv.bias": qkv_bias,
+            "decoder.det_stages.0.0.attn.q_norm.weight": torch.tensor([5.0, 6.0]),
+            "decoder.det_stages.0.0.attn.k_norm.weight": torch.tensor([7.0, 8.0]),
+            "decoder.diff_blocks.0.attn.proj.weight": projection_weight,
+            "decoder.diff_blocks.0.attn.proj.bias": projection_bias,
+            "decoder.diff_blocks.0.gate_msa": gate,
+            "vae.decoder.t_embedder.mlp.0.weight": torch.ones(2, 2),
+        }
+
+        converted = convert_ltx25_native_diffusion_decoder_state_dict(native)
+
+        torch.testing.assert_close(converted["latents_mean"], mean)
+        torch.testing.assert_close(converted["latents_std"], std)
+        torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.to_q.weight"], qkv_weight[:2])
+        torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.to_k.weight"], qkv_weight[2:4])
+        torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.to_v.weight"], qkv_weight[4:])
+        torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.to_q.bias"], qkv_bias[:2])
+        torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.to_k.bias"], qkv_bias[2:4])
+        torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.to_v.bias"], qkv_bias[4:])
+        torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.norm_q.weight"], torch.tensor([5.0, 6.0]))
+        torch.testing.assert_close(converted["decoder.det_stages.0.0.attn.norm_k.weight"], torch.tensor([7.0, 8.0]))
+        torch.testing.assert_close(
+            converted["decoder.diff_blocks.0.attn.to_out.0.weight"], gate.unsqueeze(1) * projection_weight
+        )
+        torch.testing.assert_close(converted["decoder.diff_blocks.0.attn.to_out.0.bias"], gate * projection_bias)
+        assert "decoder.t_embedder.timestep_embedder.linear_1.weight" in converted
+        assert not any("type_emb" in key or "coarse" in key or "gate_msa" in key for key in converted)
+
+    def test_native_diffusion_decoder_conversion_rejects_invalid_fused_qkv(self):
+        from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder import (
+            convert_ltx25_native_diffusion_decoder_state_dict,
+        )
+
+        with pytest.raises(ValueError, match="not divisible by 3"):
+            convert_ltx25_native_diffusion_decoder_state_dict(
+                {"decoder.det_stages.0.0.attn.qkv.weight": torch.ones(5, 2)}
+            )
+
+    def test_decode_uses_diffusion_decoder_without_conv_vae_conditioning(self):
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
+
+        seen = {}
+
+        class ConvVae:
+            dtype = torch.float32
+            config = SimpleNamespace(timestep_conditioning=True)
+
+            def decode(self, *_args, **_kwargs):
+                raise AssertionError("the convolutional VAE decoder must not run")
+
+        class DiffusionDecoder:
+            dtype = torch.float32
+
+            def decode(self, latents, *, generator, return_dict):
+                seen["latents"] = latents
+                seen["generator"] = generator
+                seen["return_dict"] = return_dict
+                return (latents + 1,)
+
+        class AudioVae:
+            dtype = torch.float32
+
+            def decode(self, audio_latents, *, return_dict):
+                assert return_dict is False
+                return (audio_latents + 2,)
+
+        class VideoProcessor:
+            def postprocess_video(self, video, *, output_type):
+                seen["output_type"] = output_type
+                return video
+
+        pipe = object.__new__(LTX2Pipeline)
+        torch.nn.Module.__init__(pipe)
+        pipe.use_diffusion_decoder = True
+        pipe.vae = ConvVae()
+        pipe.diffusion_decoder = DiffusionDecoder()
+        pipe.audio_vae = AudioVae()
+        pipe.vocoder = torch.nn.Identity()
+        pipe.video_processor = VideoProcessor()
+
+        generator = torch.Generator().manual_seed(123)
+        latents = torch.ones(1, 1)
+        output = pipe._decode_output(
+            latents=latents,
+            audio_latents=torch.ones(1, 1),
+            output_type="pt",
+            connector_prompt_embeds=torch.ones(1, 1),
+            generator=generator,
+            device=torch.device("cpu"),
+            decode_timestep=0.7,
+            decode_noise_scale=0.5,
+            prompt_batch_size=1,
+        )
+
+        torch.testing.assert_close(seen["latents"], latents)
+        assert seen["generator"] is generator
+        assert seen["return_dict"] is False
+        assert seen["output_type"] == "pt"
+        torch.testing.assert_close(output.output[0], latents + 1)
+        torch.testing.assert_close(output.output[1], torch.full((1, 1), 3.0))
+
+    def test_diffusion_decode_denormalizes_with_decoder_statistics(self):
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
+
+        pipe = object.__new__(LTX2Pipeline)
+        torch.nn.Module.__init__(pipe)
+        pipe.use_diffusion_decoder = True
+        pipe.transformer_spatial_patch_size = 1
+        pipe.transformer_temporal_patch_size = 1
+        pipe.vae = SimpleNamespace(
+            latents_mean=torch.full((2,), 100.0),
+            latents_std=torch.full((2,), 10.0),
+            config=SimpleNamespace(scaling_factor=1.0),
+        )
+        pipe.diffusion_decoder = SimpleNamespace(
+            latents_mean=torch.tensor([10.0, 20.0]),
+            latents_std=torch.tensor([2.0, 4.0]),
+            config=SimpleNamespace(scaling_factor=2.0),
+        )
+        pipe.audio_vae = SimpleNamespace(
+            latents_mean=torch.tensor(0.0),
+            latents_std=torch.tensor(1.0),
+        )
+        forward_ctx = SimpleNamespace(
+            latent_num_frames=1,
+            latent_height=1,
+            latent_width=1,
+            original_audio_num_frames=1,
+            latent_mel_bins=2,
+        )
+
+        video, audio = pipe._unpack_and_denormalize_stage(
+            forward_ctx,
+            torch.tensor([[[1.0, 2.0]]]),
+            torch.zeros(1, 1, 2),
+        )
+
+        torch.testing.assert_close(video.flatten(), torch.tensor([11.0, 24.0]))
+        torch.testing.assert_close(audio, torch.zeros(1, 1, 1, 2))
+
+    def test_non_output_rank_does_not_enter_diffusion_decoder(self, monkeypatch):
+        import vllm_omni.diffusion.models.ltx2.ltx2_runtime as ltx_runtime
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
+
+        class ConvVae:
+            config = SimpleNamespace(timestep_conditioning=False)
+
+            def is_distributed_enabled(self):
+                raise AssertionError("DiffVAE decode must not inspect ConvVAE collectives")
+
+        class DiffusionDecoder:
+            dtype = torch.float32
+
+            def decode(self, *_args, **_kwargs):
+                raise AssertionError("only the output rank may run DiffVAE decode")
+
+        monkeypatch.setattr(ltx_runtime.torch.distributed, "is_initialized", lambda: True)
+        monkeypatch.setattr(ltx_runtime.torch.distributed, "get_rank", lambda: 1)
+        pipe = object.__new__(LTX2Pipeline)
+        torch.nn.Module.__init__(pipe)
+        pipe.use_diffusion_decoder = True
+        pipe.vae = ConvVae()
+        pipe.diffusion_decoder = DiffusionDecoder()
+
+        output = pipe._decode_output(
+            latents=torch.ones(1, 1),
+            audio_latents=torch.ones(1, 1),
+            output_type="pt",
+            connector_prompt_embeds=torch.ones(1, 1),
+            generator=None,
+            device=torch.device("cpu"),
+            decode_timestep=0.0,
+            decode_noise_scale=None,
+            prompt_batch_size=1,
+        )
+
+        assert output.output[0].numel() == 0
+        assert output.output[1].numel() == 0
+
+    def test_non_output_rank_enters_distributed_diffusion_decoder(self, monkeypatch):
+        import vllm_omni.diffusion.models.ltx2.ltx2_runtime as ltx_runtime
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
+
+        class ConvVae:
+            config = SimpleNamespace(timestep_conditioning=False)
+
+            def is_distributed_enabled(self):
+                raise AssertionError("DiffVAE decode must query the active decoder")
+
+        class DiffusionDecoder:
+            dtype = torch.float32
+
+            def __init__(self):
+                self.decode_calls = 0
+
+            def is_distributed_enabled(self):
+                return True
+
+            def decode(self, *_args, **_kwargs):
+                self.decode_calls += 1
+                return (torch.empty(0, 3, 0, 0, 0),)
+
+        monkeypatch.setattr(ltx_runtime.torch.distributed, "is_initialized", lambda: True)
+        monkeypatch.setattr(ltx_runtime.torch.distributed, "get_rank", lambda: 1)
+        pipe = object.__new__(LTX2Pipeline)
+        torch.nn.Module.__init__(pipe)
+        pipe.use_diffusion_decoder = True
+        pipe.vae = ConvVae()
+        pipe.diffusion_decoder = DiffusionDecoder()
+
+        output = pipe._decode_output(
+            latents=torch.ones(1, 1),
+            audio_latents=torch.ones(1, 1),
+            output_type="pt",
+            connector_prompt_embeds=torch.ones(1, 1),
+            generator=None,
+            device=torch.device("cpu"),
+            decode_timestep=0.0,
+            decode_noise_scale=None,
+            prompt_batch_size=1,
+        )
+
+        assert pipe.diffusion_decoder.decode_calls == 1
+        assert output.output[0].numel() == 0
+        assert output.output[1].numel() == 0
+
+    def test_diffusion_decoder_patch_parallel_size_one_uses_native_tiling(self, monkeypatch):
+        from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder import LTX2VideoDiffusionDecoderModel
+        from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder_distributed import (
+            DistributedLTX2VideoDiffusionDecoderModel,
+        )
+
+        expected = torch.ones(1, 3, 1, 2, 2)
+        seen = {}
+
+        def native_tiled_decode(self, z, generator=None, num_inference_steps=None):
+            seen["z"] = z
+            seen["generator"] = generator
+            seen["num_inference_steps"] = num_inference_steps
+            return expected
+
+        monkeypatch.setattr(LTX2VideoDiffusionDecoderModel, "tiled_decode", native_tiled_decode)
+        model = object.__new__(DistributedLTX2VideoDiffusionDecoderModel)
+        torch.nn.Module.__init__(model)
+        model.is_distributed_enabled = lambda: False
+        generator = torch.Generator().manual_seed(17)
+        z = torch.zeros(1, 1, 1, 1, 1)
+
+        result = DistributedLTX2VideoDiffusionDecoderModel.tiled_decode(
+            model,
+            z,
+            generator=generator,
+            num_inference_steps=1,
+        )
+
+        assert result is expected
+        assert seen["z"] is z
+        assert seen["generator"] is generator
+        assert seen["num_inference_steps"] == 1
+
+    def test_distributed_diffusion_tiles_preserve_serial_noise_order(self, monkeypatch):
+        from diffusers.utils.torch_utils import randn_tensor
+
+        from vllm_omni.diffusion.models.ltx2.ltx2_diffusion_decoder_distributed import (
+            DistributedLTX2VideoDiffusionDecoderModel,
+        )
+
+        model = DistributedLTX2VideoDiffusionDecoderModel(
+            out_channels=1,
+            latent_channels=1,
+            patch_size=1,
+            decoder_head_dim=8,
+            decoder_stage_channels=(8, 8, 8, 8, 8),
+            decoder_stage_depths=(1, 1, 1, 1, 1),
+            decoder_stage_kernels=((1, 1, 1),) * 4,
+            decoder_upsample_strides=((1, 1, 1),) * 4,
+            decoder_upsample_channel_reductions=(1, 1, 1, 1),
+            decoder_stage5_kernel=(1, 1, 1),
+            spatial_compression_ratio=1,
+            temporal_compression_ratio=1,
+        )
+        model.tile_sample_min_num_frames = 2
+        model.tile_sample_stride_num_frames = 1
+        model.tile_sample_min_height = 2
+        model.tile_sample_stride_height = 1
+        model.tile_sample_min_width = 2
+        model.tile_sample_stride_width = 1
+        model.distributed_executor = SimpleNamespace(group=None)
+        model.decoder.forward_stages_1_to_3 = lambda z: z.permute(0, 2, 3, 4, 1)
+        monkeypatch.setattr(torch.distributed, "broadcast", lambda *_args, **_kwargs: None)
+
+        z = torch.zeros(1, 1, 2, 3, 3)
+        distributed_generator = torch.Generator().manual_seed(123)
+        reference_generator = torch.Generator().manual_seed(123)
+        tasks, grid_spec = model._distributed_tile_split(z, distributed_generator, num_inference_steps=1)
+
+        assert grid_spec.grid_shape == (1, 3, 3)
+        assert grid_spec.split_dims == (2, 3, 4)
+        for task in tasks:
+            tile_shape = model._tiled_pixel_shape_from_features(
+                task.tensor,
+                drop_leading_frame=task.drop_leading_frame,
+                crop_trailing_ghost=task.crop_trailing_ghost,
+            )
+            expected = randn_tensor(tile_shape, generator=reference_generator, device=z.device, dtype=z.dtype)
+            actual = randn_tensor(tile_shape, generator=task.noise_generator, device=z.device, dtype=z.dtype)
+            torch.testing.assert_close(actual, expected)
+
+        assert torch.equal(distributed_generator.get_state(), reference_generator.get_state())
+
+
 class TestLTXOutputRank:
     @pytest.mark.parametrize(
         ("distributed_vae_state", "expected_decode_calls"),
@@ -138,7 +574,7 @@ class TestLTX23VaeDistributedDecode:
             },
             output_dtype=torch.float32,
         )
-        seen = {}
+        seen: dict[str, Any] = {}
 
         def exec_tile(task):
             return torch.full(tile_output_shapes[task.tile_id], float(task.tile_id + 1))

tests/e2e/accuracy/ltx/test_ltx25_official_similarity.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """E2E accuracy guards against pinned official Lightricks LTX references.
 
@@ -43,7 +43,7 @@ LTX25_OFFICIAL_MODEL_ID = "Lightricks/LTX-2.5"
 LTX25_OFFICIAL_MODEL_REVISION = "8a4ff96f581e72bedc1b44367581c49d544a05f1"
 LTX25_OFFICIAL_COMMON_FILES = {
     "text_encoder": "text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors",
-    "video_vae": "vae/ltx-2.5-video-vae-conv-bf16.safetensors",
+    "video_vae": "vae/ltx-2.5-video-vae-bf16.safetensors",
     "audio_vae": "vae/ltx-2.5-audio-vae-bf16.safetensors",
 }
 LTX25_OFFICIAL_VARIANT_FILES = {
@@ -208,6 +208,9 @@ def _resolve_ltx25_omni_model(*, transformer_subfolder: str = "transformer") ->
         model = Path(configured_model).resolve()
         snapshot_revision = model.name if model.parent.name == "snapshots" else None
         assert (model / transformer_subfolder).is_dir(), f"LTX-2.5 component not found: {model / transformer_subfolder}"
+        assert (model / "diffusion_decoder" / "config.json").is_file(), (
+            f"LTX-2.5 DiffVAE config not found: {model / 'diffusion_decoder' / 'config.json'}"
+        )
         return model, configured_model, configured_revision or snapshot_revision
     model_id = configured_model or LTX25_OMNI_MODEL_ID
     revision = configured_revision
@@ -223,6 +226,7 @@ def _resolve_ltx25_omni_model(*, transformer_subfolder: str = "transformer") ->
                 "connectors/config.json",
                 "connectors/diffusion_pytorch_model.safetensors.index.json",
                 "connectors/diffusion_pytorch_model-*.safetensors",
+                "diffusion_decoder/config.json",
                 "scheduler/*",
                 "text_encoder/config.json",
                 "text_encoder/generation_config.json",
@@ -470,7 +474,7 @@ def _assert_strict_similarity(
 @hardware_test(res={"cuda": "H100"}, num_cards=1)
 @pytest.mark.parametrize("task", ["t2v", "i2v"])
 def test_ltx25_distilled_two_stage_matches_official(accuracy_artifact_root: Path, task: str) -> None:
-    """Compare Omni with the official runtime using the same checkpoint connector weights."""
+    """Compare Omni with official canonical DiffVAE and matching connector weights."""
     _require_ltx25_model_access()
     artifact_parent = accuracy_artifact_root / "ltx_official"
     output_root = reset_artifact_dir(artifact_parent / f"ltx2_5_distilled_{task}")
@@ -545,6 +549,7 @@ def test_ltx25_distilled_two_stage_matches_official(accuracy_artifact_root: Path
     omni_metadata = json.loads((omni_output / "metadata.json").read_text())
     assert official_metadata["official_revision"] == official_revision
     assert official_metadata["pipeline"] == "DistilledPipeline"
+    assert official_metadata["video_vae_path"] == str(official_artifacts["video_vae"])
     assert official_metadata["attention_backend"] == ATTENTION_BACKEND
     assert official_metadata["connector_model"] == str(omni_model)
     assert omni_metadata["model_class_name"] == "LTX2DistilledTwoStagePipeline"
@@ -594,7 +599,7 @@ def test_ltx25_distilled_two_stage_matches_official(accuracy_artifact_root: Path
 @pytest.mark.parametrize("task", ["t2v", "i2v"])
 @pytest.mark.parametrize("pipeline_mode", ["full_one_stage", "full_two_stage"])
 def test_ltx25_full_matches_official(accuracy_artifact_root: Path, task: str, pipeline_mode: str) -> None:
-    """Compare raw official Full/SFT weights with Omni's converted Full pipeline."""
+    """Compare raw official Full/SFT and DiffVAE weights with Omni."""
     _require_ltx25_model_access()
     artifact_parent = accuracy_artifact_root / "ltx_official"
     output_root = reset_artifact_dir(artifact_parent / f"ltx2_5_{pipeline_mode}_{task}")
@@ -697,6 +702,7 @@ def test_ltx25_full_matches_official(accuracy_artifact_root: Path, task: str, pi
             "full_two_stage": "TI2VidTwoStagesPipeline",
         }[pipeline_mode]
     )
+    assert official_metadata["video_vae_path"] == str(official_artifacts["video_vae"])
     assert official_metadata["attention_backend"] == ATTENTION_BACKEND
     assert official_metadata["connector_model"] == str(omni_model)
     assert omni_metadata["model_class_name"] == model_class_name

tests/e2e/online_serving/test_ltx25.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """Online serving smokes for all four canonical LTX-2.5 pipelines."""
 
@@ -73,8 +73,8 @@ def test_ltx25_pipeline_entries(
                 "form_data": {
                     "model": omni_server.model,
                     "prompt": PROMPT,
-                    "height": 128,
-                    "width": 128,
+                    "height": 256,
+                    "width": 256,
                     "num_frames": 9,
                     "fps": 24,
                     "num_inference_steps": num_inference_steps,

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -33,6 +33,64 @@ def test_default_stage_config_includes_cache_backend():
     assert engine_args["model_stage"] == "diffusion"
 
 
+def test_default_stage_config_preserves_model_extras():
+    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg({"extras": {"ltx2_use_conv_vae": True}})[0]
+
+    assert stage_cfg["engine_args"]["extras"]["ltx2_use_conv_vae"] is True
+
+
+def test_default_stage_config_preserves_and_overrides_promoted_extras():
+    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+        {
+            "extras": {
+                "auxiliary_text_encoder": "/models/extras-llama",
+                "default_llama_model_id": "extras/default-llama",
+            },
+            "auxiliary_text_encoder": None,
+        }
+    )[0]
+
+    extras = stage_cfg["engine_args"]["extras"]
+    assert extras["auxiliary_text_encoder"] == "/models/extras-llama"
+    assert extras["default_llama_model_id"] == "extras/default-llama"
+
+    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+        {
+            "extras": {
+                "auxiliary_text_encoder": "/models/extras-llama",
+                "default_llama_model_id": "extras/default-llama",
+            },
+            "auxiliary_text_encoder": "/models/top-level-llama",
+            "default_llama_model_id": "top-level/default-llama",
+        }
+    )[0]
+
+    extras = stage_cfg["engine_args"]["extras"]
+    assert extras["auxiliary_text_encoder"] == "/models/top-level-llama"
+    assert extras["default_llama_model_id"] == "top-level/default-llama"
+
+
+def test_stage_override_preserves_model_extras_for_default_diffusion_stage(mocker):
+    """Local/unregistered Diffusers checkpoints still honor stage-0 extras."""
+
+    def resolve_with_default(*_args, default_stage_cfg_factory, **_kwargs):
+        return None, default_stage_cfg_factory(), None
+
+    mocker.patch(
+        "vllm_omni.engine.async_omni_engine.load_and_resolve_stage_configs",
+        side_effect=resolve_with_default,
+    )
+    engine = AsyncOmniEngine.__new__(AsyncOmniEngine)
+
+    _, stage_configs = engine._resolve_stage_configs(
+        "/models/LTX-2.5-Diffusers",
+        {"stage_overrides": '{"0":{"extras":{"ltx2_use_conv_vae":true}}}'},
+        trust_remote_code=False,
+    )
+
+    assert stage_configs[0]["engine_args"]["extras"]["ltx2_use_conv_vae"] is True
+
+
 def test_default_cache_config_used_when_missing():
     """Ensure default cache_config is synthesized when only backend is given."""
     stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(

vllm_omni/diffusion/models/ltx2/ltx2_components.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """Shared component construction helpers for the LTX model family."""
 
@@ -15,6 +15,7 @@ from typing import TYPE_CHECKING, Any
 
 import torch
 from diffusers import AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video, FlowMatchEulerDiscreteScheduler
+from diffusers.models.autoencoders.ltx2_diffusion_decoder import LTX2VideoVaeNeighborhoodNattenProcessor
 from diffusers.pipelines.ltx2 import LTX2TextConnectors
 from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel
 from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder
@@ -35,6 +36,12 @@ from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
 if TYPE_CHECKING:
     from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
 
+from .ltx2_diffusion_decoder import (
+    LTX25_NATIVE_ARTIFACT_REVISION,
+    LTX25_NATIVE_DIFFUSION_DECODER_FILENAME,
+    LTX25_NATIVE_DIFFUSION_DECODER_REPO_ID,
+)
+from .ltx2_diffusion_decoder_distributed import DistributedLTX2VideoDiffusionDecoderModel
 from .ltx2_request import LTXCheckpointKind, validate_ltx_checkpoint
 from .ltx2_transformer import (
     LTX2VideoTransformer3DModel,
@@ -69,6 +76,18 @@ _LTX_COMPONENT_SUBFOLDERS = (
 )
 logger = logging.getLogger(__name__)
 
+_LTX2_CONV_VAE_EXTRA = "ltx2_use_conv_vae"
+_LTX2_DIFFUSION_DECODER_SUBFOLDER = "diffusion_decoder"
+
+
+def _ltx2_use_diffusion_decoder(od_config: Any, model_version: str) -> bool:
+    """Select DiffVAE by default for LTX-2.5, with an explicit ConvVAE opt-in."""
+    extras = getattr(od_config, "extras", {}) or {}
+    use_conv_vae = extras.get(_LTX2_CONV_VAE_EXTRA, False)
+    if not isinstance(use_conv_vae, bool):
+        raise TypeError(f"{_LTX2_CONV_VAE_EXTRA} must be a bool, got {type(use_conv_vae)!r}")
+    return model_version == "2.5" and not use_conv_vae
+
 
 @dataclass(frozen=True)
 class LTXComponentProfile:
@@ -84,6 +103,7 @@ class LTXComponentProfile:
     text_encoder_cls: type | None = Gemma3ForConditionalGeneration
     vocoder_fallback_cls: type | None = None
     artifact_repo_id: str | None = None
+    artifact_revision: str | None = None
     latent_upsampler_filename: str | None = None
     distilled_lora_filename: str | None = None
     transformer_subfolder: str = "transformer"
@@ -202,6 +222,7 @@ LTX25_TWO_STAGE_COMPONENT_PROFILE = replace(
     name="ltx2_5_two_stage",
     resident_modules=(*LTX25_FULL_COMPONENT_PROFILE.resident_modules, "latent_upsampler"),
     artifact_repo_id="Lightricks/LTX-2.5",
+    artifact_revision=LTX25_NATIVE_ARTIFACT_REVISION,
     latent_upsampler_filename=("latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors"),
     distilled_lora_filename="loras/ltx-2.5-22b-distilled-lora-450-bf16.safetensors",
 )
@@ -239,14 +260,30 @@ def resolve_ltx_artifact(
     model: str,
     repo_id: str,
     filename: str,
+    *,
+    model_revision: str | None,
+    artifact_revision: str | None,
 ) -> str:
-    """Resolve an official LTX sidecar from the model root or its Hub repository."""
+    """Resolve an official LTX sidecar without crossing repository revisions.
+
+    A local model path selects where the primary components are loaded from;
+    it does not imply that independently hosted sidecars must be offline. Hub
+    offline behavior remains controlled by huggingface_hub itself.
+    """
     candidate = Path(model) / filename
     if candidate.is_file():
         return str(candidate)
 
+    # Hub revisions are repository-scoped. Reuse the model revision only when
+    # the model and artifact are in the same repository; otherwise use the
+    # independently pinned artifact revision.
+    revision = model_revision if model == repo_id else artifact_revision
     try:
-        return hf_hub_download(repo_id=repo_id, filename=filename)
+        return hf_hub_download(
+            repo_id=repo_id,
+            filename=filename,
+            revision=revision,
+        )
     except Exception as exc:
         raise FileNotFoundError(
             f"Unable to resolve LTX artifact {filename!r}. Searched {candidate}; "
@@ -254,6 +291,18 @@ def resolve_ltx_artifact(
         ) from exc
 
 
+def _create_ltx25_natten_processor() -> LTX2VideoVaeNeighborhoodNattenProcessor:
+    """Create the required production attention backend with an actionable error."""
+    try:
+        return LTX2VideoVaeNeighborhoodNattenProcessor()
+    except (ImportError, OSError, RuntimeError, ValueError) as exc:
+        raise RuntimeError(
+            "LTX-2.5 DiffVAE requires the shi-labs/natten Hub kernel. "
+            "Install kernels==0.14.1, use a supported GPU, leave "
+            "DIFFUSERS_DISABLE_REMOTE_CODE unset, and allow Hub access during kernel initialization."
+        ) 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:
@@ -519,18 +568,49 @@ def _load_component(
     local_files_only: bool,
     dtype: torch.dtype,
     revision: str | None,
+    prefetch_list: tuple[str, ...] = _LTX_COMPONENT_SUBFOLDERS,
 ) -> Any:
     return from_pretrained_with_prefetch(
         component_cls.from_pretrained,
         model,
         subfolder=subfolder,
-        prefetch_list=_LTX_COMPONENT_SUBFOLDERS,
+        prefetch_list=prefetch_list,
         local_files_only=local_files_only,
         revision=revision,
         torch_dtype=dtype,
     )
 
 
+def _load_ltx25_native_diffusion_decoder(
+    model: str,
+    *,
+    local_files_only: bool,
+    dtype: torch.dtype,
+    revision: str | None,
+) -> DistributedLTX2VideoDiffusionDecoderModel:
+    """Load canonical Native weights into the local Diffusers-compatible class."""
+    config = DistributedLTX2VideoDiffusionDecoderModel.load_config(
+        model,
+        subfolder=_LTX2_DIFFUSION_DECODER_SUBFOLDER,
+        local_files_only=local_files_only,
+        revision=revision,
+    )
+    checkpoint_path = resolve_ltx_artifact(
+        model,
+        LTX25_NATIVE_DIFFUSION_DECODER_REPO_ID,
+        LTX25_NATIVE_DIFFUSION_DECODER_FILENAME,
+        model_revision=revision,
+        artifact_revision=LTX25_NATIVE_ARTIFACT_REVISION,
+    )
+    decoder = DistributedLTX2VideoDiffusionDecoderModel.from_ltx25_native_checkpoint(
+        checkpoint_path,
+        config,
+        dtype,
+    )
+    decoder.init_distributed()
+    return decoder
+
+
 def _place_aux_components(pipeline: Any) -> None:
     parallel_config = getattr(pipeline.od_config, "parallel_config", None)
     use_managed_placement = bool(
@@ -555,6 +635,7 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None:
     model = od_config.model
     revision = getattr(od_config, "revision", None)
     local_files_only = os.path.exists(model)
+    use_diffusion_decoder = pipeline.use_diffusion_decoder
 
     pipeline.weights_sources = [
         DiffusersPipelineLoader.ComponentSource(
@@ -604,6 +685,27 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None:
         dtype=dtype,
         revision=revision,
     )
+    if use_diffusion_decoder:
+        pipeline.diffusion_decoder = _load_ltx25_native_diffusion_decoder(
+            model,
+            local_files_only=local_files_only,
+            dtype=dtype,
+            revision=revision,
+        )
+        # The canonical decoder is trained and validated with NATTEN. Its
+        # portable FlexAttention fallback materializes an impractically large
+        # block mask at production video sizes, so fail early if the matching
+        # Hub kernel cannot be loaded instead of failing later during decode.
+        pipeline.diffusion_decoder.set_attn_processor(_create_ltx25_natten_processor())
+        vae_patch_parallel_size = int(
+            getattr(getattr(od_config, "parallel_config", None), "vae_patch_parallel_size", 1)
+        )
+        if vae_patch_parallel_size > 1 or getattr(od_config, "vae_use_tiling", False):
+            pipeline.diffusion_decoder.enable_tiling()
+        pipeline.diffusion_decoder.set_parallel_size(
+            vae_patch_parallel_size,
+            mode=getattr(getattr(od_config, "parallel_config", None), "vae_parallel_mode", "tile"),
+        )
     pipeline.audio_vae = _load_component(
         AutoencoderKLLTX2Audio,
         model,
@@ -656,6 +758,8 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None:
                     model,
                     profile.artifact_repo_id,
                     profile.latent_upsampler_filename,
+                    model_revision=revision,
+                    artifact_revision=profile.artifact_revision,
                 )
                 pipeline.latent_upsampler = _load_ltx_latent_upsampler_single_file(upsampler_path, dtype)
         else:
@@ -665,6 +769,8 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None:
                 model,
                 profile.artifact_repo_id,
                 profile.latent_upsampler_filename,
+                model_revision=revision,
+                artifact_revision=profile.artifact_revision,
             )
             pipeline.latent_upsampler = _load_ltx_latent_upsampler_single_file(upsampler_path, dtype)
 

vllm_omni/diffusion/models/ltx2/ltx2_conditioning.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """Task-specific conditioning shared by LTX model versions."""
 
@@ -502,10 +502,12 @@ class LTXI2VConditioningMixin:
         init_latents = torch.cat(init_latents, dim=0).to(dtype=dtype)
         if num_videos_per_prompt > 1:
             init_latents = init_latents.repeat_interleave(num_videos_per_prompt, dim=0)
+        latents_mean, latents_std, scaling_factor = latent_ops.resolve_video_latent_statistics(self)
         return latent_ops.normalize_latents(
             init_latents,
-            self.vae.latents_mean,
-            self.vae.latents_std,
+            latents_mean,
+            latents_std,
+            scaling_factor,
         )
 
     @staticmethod
@@ -570,11 +572,12 @@ class LTXI2VConditioningMixin:
                     )
                 conditioning_mask = latents.new_zeros(mask_shape)
                 conditioning_mask[:, :, 0] = 1.0
+                latents_mean, latents_std, scaling_factor = latent_ops.resolve_video_latent_statistics(self)
                 latents = latent_ops.normalize_latents(
                     latents,
-                    self.vae.latents_mean,
-                    self.vae.latents_std,
-                    self.vae.config.scaling_factor,
+                    latents_mean,
+                    latents_std,
+                    scaling_factor,
                 )
                 clean_latents = latents
                 if image is not None:

vllm_omni/diffusion/models/ltx2/ltx2_diffusion_decoder.py

@@ -0,0 +1,245 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# SPDX-FileCopyrightText: Copyright 2026 Lightricks and The HuggingFace Team. All rights reserved.
+#
+# Native conversion rules and decoder behavior copied and modified from
+# Diffusers commit d035dcd7cc7c88e0a154609b62887d50bba9fdc2 (Apache-2.0).
+
+"""LTX-2.5 Native-weight adapter for Diffusers' diffusion VAE decoder."""
+
+from collections.abc import Mapping
+from typing import Any
+
+import torch
+from diffusers.models.autoencoders import (
+    LTX2VideoDiffusionDecoderModel as DiffusersLTX2VideoDiffusionDecoderModel,
+)
+from diffusers.models.autoencoders.ltx2_diffusion_decoder import (
+    LTX2VideoDiffusionDecoder3d as DiffusersLTX2VideoDiffusionDecoder3d,
+)
+from diffusers.models.autoencoders.vae import DecoderOutput
+from safetensors import safe_open
+
+LTX25_NATIVE_DIFFUSION_DECODER_REPO_ID = "Lightricks/LTX-2.5"
+LTX25_NATIVE_ARTIFACT_REVISION = "8a4ff96f581e72bedc1b44367581c49d544a05f1"
+LTX25_NATIVE_DIFFUSION_DECODER_FILENAME = "vae/ltx-2.5-video-vae-bf16.safetensors"
+
+_NATIVE_DECODER_PREFIXES = ("vae.decoder.", "decoder.")
+_NATIVE_STATISTICS_KEYS = {
+    "per_channel_statistics.mean-of-means": "latents_mean",
+    "per_channel_statistics.std-of-means": "latents_std",
+}
+_NATIVE_KEY_REPLACEMENTS = (
+    ("t_embedder.mlp.0.", "t_embedder.timestep_embedder.linear_1."),
+    ("t_embedder.mlp.2.", "t_embedder.timestep_embedder.linear_2."),
+    (".attn.proj.", ".attn.to_out.0."),
+    (".attn.q_norm.", ".attn.norm_q."),
+    (".attn.k_norm.", ".attn.norm_k."),
+)
+_GATE_FOLD_TARGETS = {
+    ".attn.to_out.0.weight": ".gate_msa",
+    ".attn.to_out.0.bias": ".gate_msa",
+    ".mlp.w_down.weight": ".gate_mlp",
+    ".context_proj.weight": ".gate_ctx",
+    ".context_proj.bias": ".gate_ctx",
+}
+_GATE_SUFFIXES = tuple(_GATE_FOLD_TARGETS.values())
+
+
+def _strip_native_decoder_prefix(key: str) -> str | None:
+    for prefix in _NATIVE_DECODER_PREFIXES:
+        if key.startswith(prefix):
+            return key.removeprefix(prefix)
+    return None
+
+
+def _native_statistics_target(key: str) -> str | None:
+    return _NATIVE_STATISTICS_KEYS.get(key.removeprefix("vae."))
+
+
+def _rename_native_decoder_key(key: str) -> str:
+    for source, target in _NATIVE_KEY_REPLACEMENTS:
+        key = key.replace(source, target)
+    return key
+
+
+def _fold_native_gate(key: str, value: torch.Tensor, gates: Mapping[str, torch.Tensor]) -> torch.Tensor:
+    target_suffix = next((suffix for suffix in _GATE_FOLD_TARGETS if key.endswith(suffix)), None)
+    if target_suffix is None:
+        return value
+    gate_key = key[: -len(target_suffix)] + _GATE_FOLD_TARGETS[target_suffix]
+    gate = gates.get(gate_key)
+    if gate is None:
+        return value
+    gate = gate.to(device=value.device, dtype=torch.float32)
+    value_float = value.to(dtype=torch.float32)
+    folded = gate.unsqueeze(1) * value_float if value.ndim == 2 else gate * value_float
+    return folded.to(dtype=value.dtype)
+
+
+def convert_ltx25_native_diffusion_decoder_state_dict(
+    native_state_dict: Mapping[str, torch.Tensor],
+) -> dict[str, torch.Tensor]:
+    """Convert the canonical Native LTX-2.5 DiffVAE state dict to Diffusers."""
+    decoder_state_dict: dict[str, torch.Tensor] = {}
+    converted: dict[str, torch.Tensor] = {}
+    for key, value in native_state_dict.items():
+        statistics_target = _native_statistics_target(key)
+        if statistics_target is not None:
+            converted[statistics_target] = value
+            continue
+        decoder_key = _strip_native_decoder_prefix(key)
+        if decoder_key is not None:
+            decoder_state_dict[decoder_key] = value
+
+    gates = {key: value for key, value in decoder_state_dict.items() if key.endswith(_GATE_SUFFIXES)}
+    for key, value in decoder_state_dict.items():
+        # ``type_emb`` is an unused shipping-checkpoint residue. Preview heads
+        # and legacy static gates are not parameters of the runtime decoder.
+        if key == "type_emb" or key.startswith("coarse_") or ".coarse_" in key or key in gates:
+            continue
+
+        converted_key = _rename_native_decoder_key(key)
+        value = _fold_native_gate(converted_key, value, gates)
+        if converted_key.endswith((".qkv.weight", ".qkv.bias")):
+            leaf = "weight" if converted_key.endswith(".weight") else "bias"
+            prefix = converted_key[: -len(f"qkv.{leaf}")]
+            if value.shape[0] % 3 != 0:
+                raise ValueError(
+                    f"Fused LTX-2.5 DiffVAE parameter {key!r} has leading dimension "
+                    f"{value.shape[0]}, which is not divisible by 3."
+                )
+            chunk = value.shape[0] // 3
+            converted[f"decoder.{prefix}to_q.{leaf}"] = value[:chunk].clone()
+            converted[f"decoder.{prefix}to_k.{leaf}"] = value[chunk : 2 * chunk].clone()
+            converted[f"decoder.{prefix}to_v.{leaf}"] = value[2 * chunk :].clone()
+            continue
+
+        converted[f"decoder.{converted_key}"] = value
+
+    return converted
+
+
+def load_ltx25_native_diffusion_decoder_state_dict(path: str) -> dict[str, torch.Tensor]:
+    """Read and convert only DiffVAE tensors from the canonical Native VAE file."""
+    with safe_open(path, framework="pt", device="cpu") as handle:
+        native_state_dict = {
+            key: handle.get_tensor(key)
+            for key in handle.keys()
+            if _strip_native_decoder_prefix(key) is not None or _native_statistics_target(key) is not None
+        }
+    return convert_ltx25_native_diffusion_decoder_state_dict(native_state_dict)
+
+
+class LTX2VideoDiffusionDecoder3d(DiffusersLTX2VideoDiffusionDecoder3d):
+    """Diffusers decoder core with the short-clip NATTEN context fix."""
+
+    def forward_stage_4(
+        self,
+        hidden_states: torch.Tensor,
+        drop_leading_frame: bool = True,
+        crop_trailing_ghost: bool = True,
+    ) -> torch.Tensor:
+        blocks = self.det_stages[-1]
+        block_mask = blocks[0].attn.build_block_mask(hidden_states)
+        for block in blocks:
+            hidden_states = block(hidden_states, block_mask)
+        hidden_states = self.upsamples[-1](hidden_states, drop_leading_frame=drop_leading_frame)
+
+        num_pad = self.trailing_pad_latent_frames
+        if crop_trailing_ghost and num_pad > 0:
+            content_frames = max(hidden_states.shape[1] - num_pad * self.temporal_compression_ratio, 1)
+            keep_frames = min(hidden_states.shape[1], max(content_frames, self.stage5_kernel[0]))
+            hidden_states = hidden_states[:, :keep_frames]
+        return hidden_states
+
+
+class LTX2VideoDiffusionDecoderModel(DiffusersLTX2VideoDiffusionDecoderModel):
+    """Diffusers 0.40 decoder plus canonical Native loading and short-clip support."""
+
+    def __init__(
+        self,
+        out_channels: int = 3,
+        latent_channels: int = 128,
+        patch_size: int = 4,
+        scaling_factor: float = 1.0,
+        decoder_head_dim: int = 64,
+        decoder_stage_channels: tuple[int, ...] = (2048, 1024, 512, 512, 256),
+        decoder_stage_depths: tuple[int, ...] = (4, 6, 4, 2, 8),
+        decoder_stage_kernels: tuple[tuple[int, int, int], ...] = ((3, 7, 7), (3, 7, 7), (3, 5, 5), (3, 5, 5)),
+        decoder_upsample_strides: tuple[tuple[int, int, int], ...] = ((1, 2, 2), (2, 1, 1), (2, 2, 2), (2, 2, 2)),
+        decoder_upsample_channel_reductions: tuple[int, ...] = (2, 2, 1, 2),
+        decoder_stage5_kernel: tuple[int, int, int] = (11, 11, 11),
+        decoder_t_emb_dim: int = 384,
+        decoder_timestep_scale_multiplier: float = 1000.0,
+        decoder_model_output_type: str = "x0",
+        decoder_num_inference_steps: int = 1,
+        spatial_compression_ratio: int = 32,
+        temporal_compression_ratio: int = 8,
+    ) -> None:
+        super().__init__(
+            out_channels=out_channels,
+            latent_channels=latent_channels,
+            patch_size=patch_size,
+            scaling_factor=scaling_factor,
+            decoder_head_dim=decoder_head_dim,
+            decoder_stage_channels=decoder_stage_channels,
+            decoder_stage_depths=decoder_stage_depths,
+            decoder_stage_kernels=decoder_stage_kernels,
+            decoder_upsample_strides=decoder_upsample_strides,
+            decoder_upsample_channel_reductions=decoder_upsample_channel_reductions,
+            decoder_stage5_kernel=decoder_stage5_kernel,
+            decoder_t_emb_dim=decoder_t_emb_dim,
+            decoder_timestep_scale_multiplier=decoder_timestep_scale_multiplier,
+            decoder_model_output_type=decoder_model_output_type,
+            decoder_num_inference_steps=decoder_num_inference_steps,
+            spatial_compression_ratio=spatial_compression_ratio,
+            temporal_compression_ratio=temporal_compression_ratio,
+        )
+        # Diffusers constructs its private decoder core internally. Switching
+        # to this behavior-only subclass preserves parameters and state-dict keys.
+        self.decoder.__class__ = LTX2VideoDiffusionDecoder3d
+        self.decoder.stage5_kernel = tuple(decoder_stage5_kernel)
+
+    @classmethod
+    def from_ltx25_native_checkpoint(
+        cls,
+        checkpoint_path: str,
+        config: Mapping[str, Any],
+        dtype: torch.dtype,
+    ) -> "LTX2VideoDiffusionDecoderModel":
+        """Construct this decoder class directly from the canonical Native checkpoint."""
+        with torch.device("meta"):
+            model = cls.from_config(dict(config))
+        state_dict = load_ltx25_native_diffusion_decoder_state_dict(checkpoint_path)
+        try:
+            model.load_state_dict(state_dict, strict=True, assign=True)
+        except RuntimeError as exc:
+            raise ValueError(f"Invalid LTX-2.5 Native DiffVAE checkpoint {checkpoint_path!r}.") from exc
+        model.to(device="cpu", dtype=dtype)
+        return model
+
+    def decode(
+        self,
+        z: torch.Tensor,
+        generator: torch.Generator | None = None,
+        num_inference_steps: int | None = None,
+        return_dict: bool = True,
+    ) -> DecoderOutput | tuple[torch.Tensor]:
+        decoded = super().decode(
+            z,
+            generator=generator,
+            num_inference_steps=num_inference_steps,
+            return_dict=False,
+        )[0]
+
+        # Short clips retain replicated trailing context through stage 5. Do
+        # not expose those context pixels to the caller.
+        target_num_frames = (z.shape[2] - 1) * self.temporal_compression_ratio + 1
+        target_height = z.shape[3] * self.spatial_compression_ratio
+        target_width = z.shape[4] * self.spatial_compression_ratio
+        decoded = decoded[:, :, :target_num_frames, :target_height, :target_width]
+
+        if not return_dict:
+            return (decoded,)
+        return DecoderOutput(sample=decoded)

vllm_omni/diffusion/models/ltx2/ltx2_diffusion_decoder_distributed.py

@@ -0,0 +1,429 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# SPDX-FileCopyrightText: Copyright 2026 Lightricks and The HuggingFace Team. All rights reserved.
+#
+# Distributed tiling is copied and modified from Diffusers' serial LTX-2.5
+# tiling at commit d035dcd7cc7c88e0a154609b62887d50bba9fdc2 (Apache-2.0).
+
+"""LTX-2.5-specific distributed execution for the diffusion VAE decoder."""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+import torch.distributed as dist
+from diffusers.utils import logging
+from diffusers.utils.torch_utils import randn_tensor
+
+from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
+    DistributedOperator,
+    DistributedVaeMixin,
+    GridSpec,
+    TileTask,
+)
+
+from .ltx2_diffusion_decoder import LTX2VideoDiffusionDecoderModel
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+def _tile_intervals(length: int, tile_size: int, stride: int, min_size: int) -> list[tuple[int, int]]:
+    """Build overlapping intervals, merging a too-small trailing remnant."""
+    if length <= tile_size:
+        return [(0, length)]
+    starts = list(range(0, length, stride))
+    while len(starts) > 1 and length - starts[-1] < min_size:
+        starts.pop()
+    return [(start, min(start + tile_size, length)) for start in starts[:-1]] + [(starts[-1], length)]
+
+
+@dataclass(frozen=True)
+class LTX2VideoDiffusionTilePlan:
+    """Geometry shared by all ranks for one distributed DiffVAE decode."""
+
+    temporal_tiles: tuple[tuple[int, int], ...]
+    height_tiles: tuple[tuple[int, int], ...]
+    width_tiles: tuple[tuple[int, int], ...]
+    num_frames: int
+    height: int
+    width: int
+    scale_t: int
+    scale_h: int
+    scale_w: int
+    stride_t: int
+    stride_h: int
+    stride_w: int
+    blend_frames: int
+    blend_height: int
+    blend_width: int
+    single_step_x0: bool
+
+
+@dataclass
+class LTX2VideoDiffusionTileTask(TileTask):
+    """One stage-4 + stage-5 tile and its rank-independent noise."""
+
+    drop_leading_frame: bool = False
+    crop_trailing_ghost: bool = False
+    noise_generator: torch.Generator | list[torch.Generator] | None = None
+    x_t: torch.Tensor | None = None
+
+
+class DistributedLTX2VideoDiffusionDecoderModel(LTX2VideoDiffusionDecoderModel, DistributedVaeMixin):
+    """LTX-2.5 diffusion decoder with model-specific tile parallelism.
+
+    Stages 1-3 run over the complete low-resolution feature volume on every
+    participating rank. Stage 4 and the diffusion stage run as independent
+    overlapping tile tasks, and rank 0 performs the reference blend/merge.
+    """
+
+    @classmethod
+    def from_pretrained(cls, *args: Any, **kwargs: Any):
+        model = super().from_pretrained(*args, **kwargs)
+        model.init_distributed()
+        return model
+
+    def set_parallel_size(self, parallel_size: int, mode: str = "tile") -> None:
+        if mode != "tile":
+            raise ValueError(f"LTX-2.5 DiffVAE only supports vae_parallel_mode='tile', got {mode!r}.")
+        super().set_parallel_size(parallel_size, mode=mode)
+
+    def _build_tiled_decode_plan(
+        self,
+        features: torch.Tensor,
+        num_inference_steps: int,
+    ) -> LTX2VideoDiffusionTilePlan:
+        """Build the rank-independent tile schedule around Diffusers' decoder."""
+        decoder = self.decoder
+        patch_size = decoder.patch_size
+        upsample_stride = decoder.upsamples[-1].stride
+        scale_t, scale_h, scale_w = (
+            upsample_stride[0],
+            upsample_stride[1] * patch_size,
+            upsample_stride[2] * patch_size,
+        )
+        tile_t = self.tile_sample_min_num_frames // scale_t
+        tile_h = self.tile_sample_min_height // scale_h
+        tile_w = self.tile_sample_min_width // scale_w
+        stride_t = self.tile_sample_stride_num_frames // scale_t
+        stride_h = self.tile_sample_stride_height // scale_h
+        stride_w = self.tile_sample_stride_width // scale_w
+        min_sizes = [
+            max(kernel_4, -(-kernel_5 // stride))
+            for kernel_4, kernel_5, stride in zip(
+                self.config.decoder_stage_kernels[-1],
+                self.config.decoder_stage5_kernel,
+                upsample_stride,
+            )
+        ]
+
+        ghost_frames = decoder.trailing_pad_latent_frames * math.prod(
+            upsample.stride[0] for upsample in decoder.upsamples[:-1]
+        )
+        num_frames = features.shape[1] - ghost_frames
+        height, width = features.shape[2:4]
+        return LTX2VideoDiffusionTilePlan(
+            temporal_tiles=tuple(_tile_intervals(num_frames, tile_t, stride_t, min_sizes[0])),
+            height_tiles=tuple(_tile_intervals(height, tile_h, stride_h, min_sizes[1])),
+            width_tiles=tuple(_tile_intervals(width, tile_w, stride_w, min_sizes[2])),
+            num_frames=num_frames,
+            height=height,
+            width=width,
+            scale_t=scale_t,
+            scale_h=scale_h,
+            scale_w=scale_w,
+            stride_t=stride_t,
+            stride_h=stride_h,
+            stride_w=stride_w,
+            blend_frames=(tile_t - stride_t) * scale_t,
+            blend_height=(tile_h - stride_h) * scale_h,
+            blend_width=(tile_w - stride_w) * scale_w,
+            single_step_x0=num_inference_steps == 1 and decoder.model_output_type == "x0",
+        )
+
+    @staticmethod
+    def _iter_tiled_decode_coords(plan: LTX2VideoDiffusionTilePlan):
+        for t_idx in range(len(plan.temporal_tiles)):
+            for h_idx in range(len(plan.height_tiles)):
+                for w_idx in range(len(plan.width_tiles)):
+                    yield t_idx, h_idx, w_idx
+
+    @staticmethod
+    def _get_tiled_feature_slice(
+        features: torch.Tensor,
+        coord: tuple[int, int, int],
+        plan: LTX2VideoDiffusionTilePlan,
+    ) -> tuple[torch.Tensor, bool, bool]:
+        t_idx, h_idx, w_idx = coord
+        t0, t1 = plan.temporal_tiles[t_idx]
+        h0, h1 = plan.height_tiles[h_idx]
+        w0, w1 = plan.width_tiles[w_idx]
+        is_origin = t0 == 0
+        is_trailing = t1 == plan.num_frames
+        feature_t1 = features.shape[1] if is_trailing else t1
+        return features[:, t0:feature_t1, h0:h1, w0:w1], is_origin, is_trailing
+
+    def _tiled_pixel_shape_from_features(
+        self,
+        feature_tile: torch.Tensor,
+        *,
+        drop_leading_frame: bool,
+        crop_trailing_ghost: bool,
+    ) -> tuple[int, int, int, int, int]:
+        """Predict one output tile's shape without materializing stage 4."""
+        decoder = self.decoder
+        stride_t, stride_h, stride_w = decoder.upsamples[-1].stride
+        num_frames = feature_tile.shape[1] * stride_t
+        if stride_t == 2 and drop_leading_frame:
+            num_frames -= 1
+        if crop_trailing_ghost and decoder.trailing_pad_latent_frames > 0:
+            content_frames = max(
+                num_frames - decoder.trailing_pad_latent_frames * decoder.temporal_compression_ratio,
+                1,
+            )
+            num_frames = min(num_frames, max(content_frames, decoder.stage5_kernel[0]))
+        return (
+            feature_tile.shape[0],
+            decoder.out_channels,
+            num_frames,
+            feature_tile.shape[2] * stride_h * decoder.patch_size,
+            feature_tile.shape[3] * stride_w * decoder.patch_size,
+        )
+
+    @staticmethod
+    def _slice_tiled_noise(
+        x_t_full: torch.Tensor,
+        coord: tuple[int, int, int],
+        tile_pixel_shape: tuple[int, int, int, int, int],
+        plan: LTX2VideoDiffusionTilePlan,
+    ) -> torch.Tensor:
+        t_idx, h_idx, w_idx = coord
+        t0, _ = plan.temporal_tiles[t_idx]
+        h0, _ = plan.height_tiles[h_idx]
+        w0, _ = plan.width_tiles[w_idx]
+        pixel_t0 = t0 * plan.scale_t - (1 if t0 != 0 and plan.scale_t == 2 else 0)
+        return x_t_full[
+            :,
+            :,
+            pixel_t0 : pixel_t0 + tile_pixel_shape[2],
+            h0 * plan.scale_h : h0 * plan.scale_h + tile_pixel_shape[3],
+            w0 * plan.scale_w : w0 * plan.scale_w + tile_pixel_shape[4],
+        ]
+
+    def _merge_tiled_decode(
+        self,
+        tiles: dict[tuple[int, int, int], torch.Tensor],
+        plan: LTX2VideoDiffusionTilePlan,
+    ) -> torch.Tensor:
+        """Blend distributed RGB tiles in Diffusers' serial order."""
+        frame_groups = []
+        for t_idx in range(len(plan.temporal_tiles)):
+            rows = [
+                [tiles[(t_idx, h_idx, w_idx)] for w_idx in range(len(plan.width_tiles))]
+                for h_idx in range(len(plan.height_tiles))
+            ]
+            result_rows = []
+            for h_idx, row in enumerate(rows):
+                result_row = []
+                for w_idx, tile in enumerate(row):
+                    if h_idx > 0:
+                        tile = self.blend_v(rows[h_idx - 1][w_idx], tile, plan.blend_height)
+                    if w_idx > 0:
+                        tile = self.blend_h(row[w_idx - 1], tile, plan.blend_width)
+                    keep_height = plan.stride_h * plan.scale_h if h_idx < len(rows) - 1 else tile.shape[3]
+                    keep_width = plan.stride_w * plan.scale_w if w_idx < len(row) - 1 else tile.shape[4]
+                    result_row.append(tile[:, :, :, :keep_height, :keep_width])
+                result_rows.append(torch.cat(result_row, dim=4))
+            frame_groups.append(torch.cat(result_rows, dim=3))
+
+        result = []
+        for t_idx, group in enumerate(frame_groups):
+            if t_idx > 0:
+                group = self.blend_t(frame_groups[t_idx - 1], group, plan.blend_frames)
+            if t_idx < len(frame_groups) - 1:
+                keep_frames = plan.stride_t * plan.scale_t - (1 if t_idx == 0 and plan.scale_t == 2 else 0)
+                group = group[:, :, :keep_frames]
+            result.append(group)
+        return torch.cat(result, dim=2)
+
+    def _default_generator(self, device: torch.device) -> torch.Generator:
+        if device.type == "cpu":
+            return torch.default_generator
+        device_module = getattr(torch, device.type, None)
+        default_generators = getattr(device_module, "default_generators", None)
+        if default_generators is None:
+            raise ValueError(
+                f"Distributed LTX-2.5 diffusion decode on {device.type!r} requires an explicit torch.Generator."
+            )
+        device_index = device.index
+        if device_index is None:
+            current_device = getattr(device_module, "current_device", None)
+            device_index = current_device() if current_device is not None else 0
+        return default_generators[device_index]
+
+    def _sync_generators(
+        self,
+        generator: torch.Generator | list[torch.Generator] | None,
+        device: torch.device,
+    ) -> torch.Generator | list[torch.Generator]:
+        """Make every rank start tile-noise generation from rank 0's state."""
+        generators = self._default_generator(device) if generator is None else generator
+        generator_list = generators if isinstance(generators, list) else [generators]
+        for item in generator_list:
+            state_on_device = item.get_state().to(device=device)
+            dist.broadcast(state_on_device, src=0, group=self.distributed_executor.group)
+            item.set_state(state_on_device.cpu())
+        return generators
+
+    @staticmethod
+    def _clone_generator(generator: torch.Generator) -> torch.Generator:
+        cloned = torch.Generator(device=generator.device)
+        cloned.set_state(generator.get_state())
+        return cloned
+
+    def _clone_generators(
+        self,
+        generators: torch.Generator | list[torch.Generator],
+    ) -> torch.Generator | list[torch.Generator]:
+        if isinstance(generators, list):
+            return [self._clone_generator(item) for item in generators]
+        return self._clone_generator(generators)
+
+    def _distributed_tile_split(
+        self,
+        z: torch.Tensor,
+        generator: torch.Generator | list[torch.Generator] | None,
+        num_inference_steps: int,
+    ) -> tuple[list[LTX2VideoDiffusionTileTask], GridSpec]:
+        decoder = self.decoder
+        features = decoder.forward_stages_1_to_3(z)
+        plan = self._build_tiled_decode_plan(features, num_inference_steps)
+        generators = self._sync_generators(generator, z.device)
+
+        x_t_full = None
+        if not plan.single_step_x0:
+            pixel_frames = plan.num_frames * plan.scale_t - (1 if plan.scale_t == 2 else 0)
+            x_t_full = randn_tensor(
+                (
+                    z.shape[0],
+                    decoder.out_channels,
+                    pixel_frames,
+                    plan.height * plan.scale_h,
+                    plan.width * plan.scale_w,
+                ),
+                generator=generators,
+                device=z.device,
+                dtype=z.dtype,
+            )
+
+        tasks = []
+        for tile_id, coord in enumerate(self._iter_tiled_decode_coords(plan)):
+            feature_tile, is_origin, is_trailing = self._get_tiled_feature_slice(features, coord, plan)
+            tile_pixel_shape = self._tiled_pixel_shape_from_features(
+                feature_tile,
+                drop_leading_frame=is_origin,
+                crop_trailing_ghost=is_trailing,
+            )
+            tile_generator = None
+            x_t = None
+            if plan.single_step_x0:
+                # Save the canonical serial state for this tile, then advance
+                # the request generator exactly as the single-rank tiled path
+                # does. Task placement therefore cannot change tile noise.
+                tile_generator = self._clone_generators(generators)
+                discarded_noise = randn_tensor(
+                    tile_pixel_shape,
+                    generator=generators,
+                    device=z.device,
+                    dtype=z.dtype,
+                )
+                del discarded_noise
+            else:
+                x_t = self._slice_tiled_noise(x_t_full, coord, tile_pixel_shape, plan)
+
+            tasks.append(
+                LTX2VideoDiffusionTileTask(
+                    tile_id=tile_id,
+                    grid_coord=coord,
+                    tensor=feature_tile,
+                    workload=math.prod(tile_pixel_shape),
+                    drop_leading_frame=is_origin,
+                    crop_trailing_ghost=is_trailing,
+                    noise_generator=tile_generator,
+                    x_t=x_t,
+                )
+            )
+
+        return tasks, GridSpec(
+            split_dims=(2, 3, 4),
+            grid_shape=(len(plan.temporal_tiles), len(plan.height_tiles), len(plan.width_tiles)),
+            tile_spec={"plan": plan},
+            output_dtype=z.dtype,
+        )
+
+    def _distributed_tile_exec(
+        self,
+        task: LTX2VideoDiffusionTileTask,
+        num_inference_steps: int,
+    ) -> torch.Tensor:
+        context = self.decoder.forward_stage_4(
+            task.tensor,
+            drop_leading_frame=task.drop_leading_frame,
+            crop_trailing_ghost=task.crop_trailing_ghost,
+        )
+        tile_pixel_shape = (
+            context.shape[0],
+            self.decoder.out_channels,
+            context.shape[1],
+            context.shape[2] * self.decoder.patch_size,
+            context.shape[3] * self.decoder.patch_size,
+        )
+        if task.x_t is None:
+            x_t = randn_tensor(
+                tile_pixel_shape,
+                generator=task.noise_generator,
+                device=context.device,
+                dtype=context.dtype,
+            )
+        else:
+            x_t = task.x_t
+        return self.decoder.denoise(context, x_t, num_inference_steps)
+
+    def _distributed_tile_merge(
+        self,
+        tiles: dict[tuple[int, int, int], torch.Tensor],
+        grid_spec: GridSpec,
+    ) -> torch.Tensor:
+        plan = grid_spec.tile_spec["plan"]
+        if not isinstance(plan, LTX2VideoDiffusionTilePlan):
+            raise TypeError(f"Expected an LTX2VideoDiffusionTilePlan, got {type(plan)!r}.")
+        return self._merge_tiled_decode(tiles, plan)
+
+    def tiled_decode(
+        self,
+        z: torch.Tensor,
+        generator: torch.Generator | list[torch.Generator] | None = None,
+        num_inference_steps: int | None = None,
+    ) -> torch.Tensor:
+        if not self.is_distributed_enabled():
+            return super().tiled_decode(z, generator=generator, num_inference_steps=num_inference_steps)
+
+        logger.debug("LTX-2.5 diffusion decoder running with distributed stage-4/stage-5 tiles")
+        num_inference_steps = num_inference_steps or self.decoder.default_num_inference_steps
+        result = self.distributed_executor.execute(
+            z,
+            DistributedOperator(
+                split=lambda tensor: self._distributed_tile_split(tensor, generator, num_inference_steps),
+                exec=lambda task: self._distributed_tile_exec(task, num_inference_steps),
+                merge=self._distributed_tile_merge,
+            ),
+            broadcast_result=False,
+        )
+        if result.numel() == 0:
+            # The base decode method crops five dimensions before the LTX
+            # runtime discards non-output-rank results.
+            return torch.empty((0, self.decoder.out_channels, 0, 0, 0), device=z.device, dtype=z.dtype)
+        return result

vllm_omni/diffusion/models/ltx2/ltx2_latents.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """Shared latent layout and normalization primitives for LTX pipelines."""
 
@@ -69,6 +69,13 @@ def unpack_latents(
     return latents.permute(0, 4, 1, 5, 2, 6, 3, 7).flatten(6, 7).flatten(4, 5).flatten(2, 3)
 
 
+def resolve_video_latent_statistics(pipeline: Any) -> tuple[torch.Tensor, torch.Tensor, float]:
+    """Return statistics owned by the active video decoder."""
+    decoder = pipeline.diffusion_decoder if getattr(pipeline, "use_diffusion_decoder", False) else pipeline.vae
+    scaling_factor = getattr(getattr(decoder, "config", None), "scaling_factor", 1.0)
+    return decoder.latents_mean, decoder.latents_std, scaling_factor
+
+
 def normalize_latents(
     latents: torch.Tensor,
     latents_mean: torch.Tensor,
@@ -202,11 +209,12 @@ def prepare_video_latents(
 ) -> torch.Tensor:
     if latents is not None:
         if latents.ndim == 5:
+            latents_mean, latents_std, scaling_factor = resolve_video_latent_statistics(pipeline)
             latents = normalize_latents(
                 latents,
-                pipeline.vae.latents_mean,
-                pipeline.vae.latents_std,
-                pipeline.vae.config.scaling_factor,
+                latents_mean,
+                latents_std,
+                scaling_factor,
             )
             latents = pack_latents(
                 latents,

vllm_omni/diffusion/models/ltx2/ltx2_phase_adapter.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """Construction and execution of a fixed LTX refinement-phase adapter.
 
@@ -420,6 +420,8 @@ def build_ltx_phase_adapter(pipeline: Any) -> LTXPhaseAdapterRuntime | None:
         pipeline.od_config.model,
         profile.artifact_repo_id,
         profile.distilled_lora_filename,
+        model_revision=getattr(pipeline.od_config, "revision", None),
+        artifact_revision=profile.artifact_revision,
     )
 
     if layer_fused and pipeline.od_config.dtype is not torch.bfloat16:

vllm_omni/diffusion/models/ltx2/ltx2_request.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """Request normalization and validation shared by LTX pipeline variants."""
 
@@ -86,6 +86,7 @@ def validate_pipeline_request(
     vae_spatial_compression_ratio: int,
     vae_temporal_compression_ratio: int,
     pipeline_name: str,
+    min_video_latent_spatial_size: tuple[int, int] | None = None,
     request_sigmas: list[float] | None = None,
     request_phase_sigmas: tuple[list[float] | None, ...] | None = None,
 ) -> None:
@@ -97,6 +98,18 @@ def validate_pipeline_request(
             f"{request_inputs.width}x{request_inputs.height}."
         )
 
+    if min_video_latent_spatial_size is not None:
+        min_latent_height, min_latent_width = min_video_latent_spatial_size
+        latent_height = request_inputs.height // vae_spatial_compression_ratio
+        latent_width = request_inputs.width // vae_spatial_compression_ratio
+        if latent_height < min_latent_height or latent_width < min_latent_width:
+            raise ValueError(
+                f"{pipeline_name} DiffVAE requires video latent spatial dimensions of at least "
+                f"{min_latent_width}x{min_latent_height} for neighborhood attention, got "
+                f"{latent_width}x{latent_height} from request resolution "
+                f"{request_inputs.width}x{request_inputs.height}."
+            )
+
     if request_inputs.num_frames < 1 or (request_inputs.num_frames - 1) % vae_temporal_compression_ratio != 0:
         raise ValueError(
             f"{pipeline_name} `num_frames` must be {vae_temporal_compression_ratio} * k + 1, "

vllm_omni/diffusion/models/ltx2/ltx2_runtime.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """Shared recipe-driven runtime for LTX pipeline variants."""
 
@@ -29,6 +29,7 @@ from vllm_omni.platforms import current_omni_platform
 from . import ltx2_latents as latent_ops
 from .ltx2_components import (
     LTXComponentProfile,
+    _ltx2_use_diffusion_decoder,
     detect_ltx_model_version,
     initialize_pipeline_components,
     resolve_ltx_component_profile,
@@ -168,6 +169,7 @@ class LTXRuntime(
                 "Use the default ulysses_mode='strict' for LTX sequence parallelism."
             )
         self.model_version = detect_ltx_model_version(od_config.model, revision=getattr(od_config, "revision", None))
+        self.use_diffusion_decoder = _ltx2_use_diffusion_decoder(od_config, self.model_version)
         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 not self.pipeline_recipe.supports_cache_dit:
@@ -178,6 +180,8 @@ class LTXRuntime(
         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)
+        if self.use_diffusion_decoder:
+            self._vae_modules.append("diffusion_decoder")
         self._resident_modules = list(self.component_profile.resident_modules)
         if self.model_version in ("2.3", "2.5"):
             self.preserve_sp_padded_audio_duration = True
@@ -249,12 +253,17 @@ class LTXRuntime(
             raise ValueError(f"{self.__class__.__name__} does not support `image` input.")
         request_sigmas = self._resolve_request_sigmas(req, sigmas)
         request_phase_sigmas = self._resolve_request_phase_sigmas(req, stage_1_sigmas, stage_2_sigmas)
+        min_video_latent_spatial_size = None
+        if getattr(self, "use_diffusion_decoder", False):
+            first_stage_kernel = self.diffusion_decoder.config.decoder_stage_kernels[0]
+            min_video_latent_spatial_size = (int(first_stage_kernel[-2]), int(first_stage_kernel[-1]))
         validate_pipeline_request(
             request_inputs,
             pipeline_recipe=self.pipeline_recipe,
             vae_spatial_compression_ratio=self.vae_spatial_compression_ratio,
             vae_temporal_compression_ratio=self.vae_temporal_compression_ratio,
             pipeline_name=self.__class__.__name__,
+            min_video_latent_spatial_size=min_video_latent_spatial_size,
             request_sigmas=request_sigmas,
             request_phase_sigmas=request_phase_sigmas,
         )
@@ -542,7 +551,10 @@ class LTXRuntime(
             return self._make_output((latents, audio_latents))
 
         latents = latents.to(connector_prompt_embeds.dtype)
-        if not self.vae.config.timestep_conditioning:
+        use_diffusion_decoder = getattr(self, "use_diffusion_decoder", False)
+        if use_diffusion_decoder:
+            timestep_decode = None
+        elif not self.vae.config.timestep_conditioning:
             timestep_decode = None
         else:
             noise = randn_tensor(latents.shape, generator=generator, device=device, dtype=latents.dtype)
@@ -559,14 +571,22 @@ class LTXRuntime(
         dist_initialized = torch.distributed.is_initialized()
         is_output_rank = not dist_initialized or torch.distributed.get_rank() == 0
         vae_decode_needs_all_ranks = False
-        is_distributed_vae_enabled = getattr(self.vae, "is_distributed_enabled", None)
+        video_decoder = self.diffusion_decoder if use_diffusion_decoder else self.vae
+        is_distributed_vae_enabled = getattr(video_decoder, "is_distributed_enabled", None)
         if self.distributed_video_decode and dist_initialized and callable(is_distributed_vae_enabled):
             # Distributed tiled decode is collective, so every rank must enter it.
             vae_decode_needs_all_ranks = bool(is_distributed_vae_enabled())
 
         should_decode_video = not self.distributed_video_decode or is_output_rank or vae_decode_needs_all_ranks
         if should_decode_video:
-            video = self.vae.decode(latents.to(self.vae.dtype), timestep_decode, return_dict=False)[0]
+            if use_diffusion_decoder:
+                video = self.diffusion_decoder.decode(
+                    latents.to(self.diffusion_decoder.dtype),
+                    generator=generator,
+                    return_dict=False,
+                )[0]
+            else:
+                video = self.vae.decode(latents.to(self.vae.dtype), timestep_decode, return_dict=False)[0]
         else:
             video = torch.empty(0, device=latents.device, dtype=latents.dtype)
 
@@ -800,11 +820,12 @@ class LTXRuntime(
             self.transformer_spatial_patch_size,
             self.transformer_temporal_patch_size,
         )
+        latents_mean, latents_std, scaling_factor = latent_ops.resolve_video_latent_statistics(self)
         latents = latent_ops.denormalize_latents(
             latents,
-            self.vae.latents_mean,
-            self.vae.latents_std,
-            self.vae.config.scaling_factor,
+            latents_mean,
+            latents_std,
+            scaling_factor,
         )
 
         audio_latents = latent_ops.unpad_audio_latents(audio_latents, forward_ctx.original_audio_num_frames)

vllm_omni/engine/async_omni_engine.py

@@ -1043,6 +1043,17 @@ class AsyncOmniEngine:
                 fastvideo_vsa_topk=kwargs.get("fastvideo_vsa_topk"),
             )
 
+        extras = dict(kwargs.get("extras") or {})
+        for key, default in (
+            ("auxiliary_text_encoder", None),
+            ("default_llama_model_id", "meta-llama/Meta-Llama-3.1-8B-Instruct"),
+        ):
+            top_level_value = kwargs.get(key)
+            if top_level_value is not None:
+                extras[key] = top_level_value
+            else:
+                extras.setdefault(key, default)
+
         stage_engine_args = {
             "max_num_seqs": kwargs.get("max_num_seqs") or 1,
             "parallel_config": parallel_config,
@@ -1103,10 +1114,7 @@ class AsyncOmniEngine:
             "enable_diffusion_pipeline_profiler": kwargs.get("enable_diffusion_pipeline_profiler", False),
             "streaming_output": kwargs.get("diffusion_streaming_output", False),
             "enable_ar_profiler": kwargs.get("enable_ar_profiler", False),
-            "extras": {
-                "auxiliary_text_encoder": kwargs.get("auxiliary_text_encoder", None),
-                "default_llama_model_id": kwargs.get("default_llama_model_id", "meta-llama/Meta-Llama-3.1-8B-Instruct"),
-            },
+            "extras": extras,
             **(
                 {
                     "profiler_config": asdict(kwargs["profiler_config"])
@@ -1196,11 +1204,29 @@ class AsyncOmniEngine:
         # Parse --stage-overrides JSON string if provided
         stage_overrides = parse_stage_overrides(stage_overrides_json)
 
+        # Unregistered diffusion checkpoints use the single-stage fallback
+        # below instead of StageConfigFactory, so fold stage-0 model extras
+        # into the fallback's input as well.  Registered pipelines still get
+        # the complete override mapping through load_and_resolve_stage_configs.
+        default_stage_kwargs = kwargs
+        stage_zero_overrides = (stage_overrides or {}).get("0", {})
+        if "extras" in stage_zero_overrides:
+            override_extras = stage_zero_overrides["extras"]
+            if not isinstance(override_extras, Mapping):
+                raise TypeError("stage 0 extras must be a mapping")
+            default_stage_kwargs = {
+                **kwargs,
+                "extras": {
+                    **(kwargs.get("extras") or {}),
+                    **override_extras,
+                },
+            }
+
         config_path, stage_configs, strategy_lb_policy = load_and_resolve_stage_configs(
             model,
             kwargs,
             trust_remote_code=trust_remote_code,
-            default_stage_cfg_factory=lambda: self._create_default_diffusion_stage_cfg(kwargs),
+            default_stage_cfg_factory=lambda: self._create_default_diffusion_stage_cfg(default_stage_kwargs),
             deploy_config_path=deploy_config_path,
             stage_overrides=stage_overrides,
             strategy_config_path=strategy_config_path,