[Refactor] Unify the LTX-2 and LTX-2.3 pipeline runtime (#5147)

Signed-off-by: mglyn <1203789601@qq.com>

tests/diffusion/models/dmd2/test_dmd2_request_sanitization.py

@@ -6,12 +6,17 @@ import pytest
 import torch
 
 from vllm_omni.diffusion.models.flux.pipeline_flux import FluxDMD2Pipeline, FluxPipeline
-from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline, LTX2T2VDMD2Pipeline
-from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_image2video import LTX2I2VDMD2Pipeline, LTX2ImageToVideoPipeline
+from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import (
+    LTX2I2VDMD2Pipeline,
+    LTX2ImageToVideoPipeline,
+    LTX2Pipeline,
+    LTX2T2VDMD2Pipeline,
+)
 from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image import QwenImageDMD2Pipeline, QwenImagePipeline
 from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import Wan22Pipeline, WanT2VDMD2Pipeline
 from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v import Wan22I2VPipeline, WanI2VDMD2Pipeline
 from vllm_omni.diffusion.request import OmniDiffusionRequest, OmniDiffusionSamplingParams
+from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -152,6 +157,28 @@ def test_string_prompt_not_mutated(pipeline):
     assert req.prompt == "a cat dancing"
 
 
+def test_request_batch_sanitizes_each_request(pipeline):
+    requests = [
+        _make_request(
+            prompt={"prompt": f"prompt-{index}", "negative_prompt": "blurry"},
+            num_inference_steps=40 + index,
+            guidance_scale=5.0 + index,
+            guidance_scale_provided=True,
+        )
+        for index in range(2)
+    ]
+
+    pipeline._sanitize_dmd2_request(DiffusionRequestBatch(requests))
+
+    for request in requests:
+        assert request.sampling_params.num_inference_steps == pipeline.dmd2_config.num_inference_steps
+        assert request.sampling_params.guidance_scale == pipeline.dmd2_config.guidance_scale
+        assert request.sampling_params.guidance_scale_provided is False
+        assert request.sampling_params.do_classifier_free_guidance is False
+        assert request.sampling_params.is_cfg_negative is False
+        assert "negative_prompt" not in request.prompt
+
+
 # ---------------------------------------------------------------------------
 # Clean request — nothing changes
 # ---------------------------------------------------------------------------

tests/diffusion/models/dmd2/test_dmd2_scheduler.py

@@ -6,8 +6,12 @@ import pytest
 import torch
 
 from vllm_omni.diffusion.models.flux.pipeline_flux import FluxDMD2Pipeline, FluxPipeline
-from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline, LTX2T2VDMD2Pipeline
-from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_image2video import LTX2I2VDMD2Pipeline, LTX2ImageToVideoPipeline
+from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import (
+    LTX2I2VDMD2Pipeline,
+    LTX2ImageToVideoPipeline,
+    LTX2Pipeline,
+    LTX2T2VDMD2Pipeline,
+)
 from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image import QwenImageDMD2Pipeline, QwenImagePipeline
 from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import Wan22Pipeline, WanT2VDMD2Pipeline
 from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v import Wan22I2VPipeline, WanI2VDMD2Pipeline

tests/diffusion/models/ltx2/test_ltx2_cfg_parallel_adaptation.py

@@ -1,58 +0,0 @@
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-def _make_pipeline(sequence_parallel_size: int = 1) -> LTX2Pipeline:
-    pipeline = object.__new__(LTX2Pipeline)
-    torch.nn.Module.__init__(pipeline)
-    pipeline.audio_vae_temporal_compression_ratio = 4
-    pipeline.audio_vae_mel_compression_ratio = 4
-    pipeline.od_config = SimpleNamespace(parallel_config=SimpleNamespace(sequence_parallel_size=sequence_parallel_size))
-    # Mock audio_vae with identity normalization (mean=0, std=1) so
-    # _normalize_audio_latents is a no-op and test values are preserved.
-    pipeline.audio_vae = SimpleNamespace(
-        latents_mean=torch.tensor(0.0),
-        latents_std=torch.tensor(1.0),
-    )
-    return pipeline
-
-
-def test_prepare_audio_latents_pads_packed_sequence_dim_for_provided_latents():
-    pipeline = _make_pipeline(sequence_parallel_size=4)
-    latents = torch.arange(40, dtype=torch.float32).view(1, 10, 4)
-
-    padded, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents(
-        batch_size=1,
-        num_channels_latents=2,
-        num_mel_bins=8,
-        audio_latent_length=10,
-        dtype=torch.float32,
-        device=torch.device("cpu"),
-        latents=latents,
-    )
-
-    assert original_num_frames == 10
-    assert padded_num_frames == 12
-    assert padded.shape == (1, 12, 4)
-    torch.testing.assert_close(padded[:, :10], latents)
-    torch.testing.assert_close(padded[:, 10:], torch.zeros(1, 2, 4))
-
-
-def test_unpad_audio_latents_restores_original_frames_before_unpack():
-    pipeline = _make_pipeline()
-    original = torch.arange(40, dtype=torch.float32).view(1, 10, 4)
-    padded = torch.cat([original, torch.full((1, 2, 4), 999.0)], dim=1)
-
-    unpadded = pipeline._unpad_audio_latents(padded, 10)
-    unpacked = pipeline._unpack_audio_latents(unpadded, latent_length=10, num_mel_bins=2)
-    expected = pipeline._unpack_audio_latents(original, latent_length=10, num_mel_bins=2)
-
-    assert unpacked.shape == (1, 2, 10, 2)
-    assert not (unpacked == 999.0).any()
-    torch.testing.assert_close(unpacked, expected)

tests/diffusion/models/ltx2/test_ltx2_conditioning.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""Unit tests for the LTX-2.3 image-to-video pipeline."""
+"""Unit tests for LTX image-to-video input and conditioning behavior."""
 
 from types import SimpleNamespace
 
@@ -19,9 +19,9 @@ def _make_ltx23_request_pipe(cls):
     return pipe
 
 
-class TestLTX23ImageToVideoForwardStages:
+class TestLTXImageToVideoForwardStages:
     def test_forward_resolves_request_image_and_delegates_to_shared_forward_impl(self):
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline
         from vllm_omni.diffusion.request import OmniDiffusionRequest
         from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
         from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -66,7 +66,7 @@ class TestLTX23ImageToVideoForwardStages:
         assert seen["kwargs"]["noise_scale"] == 0.5
 
     def test_denoise_timestep_kwargs_masks_video_only(self):
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline
 
         pipe = object.__new__(LTX23ImageToVideoPipeline)
         ts = torch.tensor([2.0, 4.0])
@@ -86,16 +86,14 @@ class TestLTX23ImageToVideoForwardStages:
         torch.testing.assert_close(kwargs["sigma"], ts)
 
 
-class TestLTX23ImageToVideoPipeline:
-    def test_ltx23_i2v_pipeline_reuses_ltx23_semantics(self):
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
+class TestLTXImageToVideoConditioning:
+    def test_ltx23_i2v_supports_image_input(self):
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline
 
-        assert issubclass(LTX23ImageToVideoPipeline, LTX23Pipeline)
         assert LTX23ImageToVideoPipeline.support_image_input is True
 
     def test_ltx23_i2v_rejects_multi_image_prompt_list(self):
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline
 
         image = object()
 
@@ -104,7 +102,7 @@ class TestLTX23ImageToVideoPipeline:
             LTX23ImageToVideoPipeline._resolve_single_prompt_image([object(), object()])
 
     def test_ltx23_i2v_additional_image_resolution_is_tensor_safe(self):
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline
 
         image = torch.zeros(1, 3, 4, 4)
         additional = {
@@ -116,8 +114,9 @@ class TestLTX23ImageToVideoPipeline:
         assert LTX23ImageToVideoPipeline._resolve_additional_image(additional) is image
 
     def test_ltx23_i2v_packed_latents_are_not_noised(self, monkeypatch):
-        import vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video as ltx23_i2v
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
+        import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as ltx2_conditioning
+        import vllm_omni.diffusion.models.ltx2.ltx2_latents as ltx2_latents
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline
 
         pipe = object.__new__(LTX23ImageToVideoPipeline)
         torch.nn.Module.__init__(pipe)
@@ -129,7 +128,8 @@ class TestLTX23ImageToVideoPipeline:
         def fake_randn_tensor(shape, generator=None, device=None, dtype=None):
             raise AssertionError("packed I2V latents should not be noised")
 
-        monkeypatch.setattr(ltx23_i2v, "randn_tensor", fake_randn_tensor)
+        monkeypatch.setattr(ltx2_conditioning, "randn_tensor", fake_randn_tensor)
+        monkeypatch.setattr(ltx2_latents, "randn_tensor", fake_randn_tensor)
 
         latents = torch.tensor([[[10.0, 11.0], [20.0, 21.0], [30.0, 31.0]]])
 
@@ -150,9 +150,9 @@ class TestLTX23ImageToVideoPipeline:
         torch.testing.assert_close(out, latents)
 
     def test_ltx23_i2v_5d_latents_noise_preserves_conditioning_frame(self, monkeypatch):
-        import vllm_omni.diffusion.models.ltx2.pipeline_ltx2 as ltx2
-        import vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video as ltx23_i2v
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
+        import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as ltx2_conditioning
+        import vllm_omni.diffusion.models.ltx2.ltx2_latents as ltx2_latents
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline
 
         pipe = object.__new__(LTX23ImageToVideoPipeline)
         torch.nn.Module.__init__(pipe)
@@ -169,8 +169,8 @@ class TestLTX23ImageToVideoPipeline:
         def fake_randn_tensor(shape, generator=None, device=None, dtype=None):
             return torch.ones(shape, device=device, dtype=dtype)
 
-        monkeypatch.setattr(ltx23_i2v, "randn_tensor", fake_randn_tensor)
-        monkeypatch.setattr(ltx2, "randn_tensor", fake_randn_tensor)
+        monkeypatch.setattr(ltx2_conditioning, "randn_tensor", fake_randn_tensor)
+        monkeypatch.setattr(ltx2_latents, "randn_tensor", fake_randn_tensor)
 
         latents = torch.tensor([[[[[10.0]], [[20.0]], [[30.0]]], [[[11.0]], [[21.0]], [[31.0]]]]])
 
@@ -191,7 +191,7 @@ class TestLTX23ImageToVideoPipeline:
         torch.testing.assert_close(out, torch.tensor([[[10.0, 11.0], [1.0, 1.0], [1.0, 1.0]]]))
 
     def test_ltx23_i2v_video_step_preserves_conditioning_frame(self):
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline
 
         pipe = object.__new__(LTX23ImageToVideoPipeline)
         torch.nn.Module.__init__(pipe)
@@ -217,28 +217,3 @@ class TestLTX23ImageToVideoPipeline:
 
         torch.testing.assert_close(out[:, :1], latents[:, :1])
         torch.testing.assert_close(out[:, 1:], latents[:, 1:] + noise_pred[:, 1:] + 0.5)
-
-
-class TestLTX23ImageToVideoModule:
-    """Test that pipeline_ltx2_3_image2video.py exposes I2V entry points."""
-
-    def test_i2v_classes_importable(self):
-        """I2V classes must be importable from the implementation module."""
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
-
-        assert LTX23ImageToVideoPipeline is not None
-
-    def test_post_process_func_importable(self):
-        """get_ltx2_post_process_func must be importable for registry lookup."""
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import get_ltx2_post_process_func
-
-        assert callable(get_ltx2_post_process_func)
-
-    def test_i2v_class_matches_package_export(self):
-        """Package export must point at the I2V implementation module."""
-        from vllm_omni.diffusion.models.ltx2 import LTX23ImageToVideoPipeline as PackageExported
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import (
-            LTX23ImageToVideoPipeline as ModuleExported,
-        )
-
-        assert PackageExported is ModuleExported

tests/diffusion/models/ltx2/test_ltx2_hsdp.py

@@ -1,25 +0,0 @@
-import pytest
-import torch.nn as nn
-
-from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-def test_ltx2_exposes_hsdp_shard_conditions_for_transformer_blocks():
-    model = object.__new__(LTX2VideoTransformer3DModel)
-    nn.Module.__init__(model)
-    model.transformer_blocks = nn.ModuleList([nn.Linear(4, 4) for _ in range(2)])
-    model.norm_out = nn.LayerNorm(4)
-
-    conditions = getattr(model, "_hsdp_shard_conditions", None)
-
-    assert conditions is not None
-    assert len(conditions) == 1
-
-    matched = []
-    for name, module in model.named_modules():
-        if any(cond(name, module) for cond in conditions):
-            matched.append(name)
-
-    assert matched == ["transformer_blocks.0", "transformer_blocks.1"]

tests/diffusion/models/ltx2/test_ltx2_latents.py

@@ -0,0 +1,121 @@
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.ltx2 import ltx2_latents as latent_ops
+from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline, LTX23Pipeline
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_pipeline(pipeline_cls, sequence_parallel_size: int = 1):
+    pipeline = object.__new__(pipeline_cls)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.audio_vae_temporal_compression_ratio = 4
+    pipeline.audio_vae_mel_compression_ratio = 4
+    pipeline.od_config = SimpleNamespace(parallel_config=SimpleNamespace(sequence_parallel_size=sequence_parallel_size))
+    # Mock audio_vae with identity normalization (mean=0, std=1).
+    pipeline.audio_vae = SimpleNamespace(
+        latents_mean=torch.tensor(0.0),
+        latents_std=torch.tensor(1.0),
+    )
+    return pipeline
+
+
+def test_prepare_audio_latents_pads_generated_dummy_length_for_sp():
+    pipeline = _make_pipeline(LTX23Pipeline, sequence_parallel_size=2)
+
+    latents, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents(
+        batch_size=1,
+        num_channels_latents=8,
+        num_mel_bins=64,
+        audio_latent_length=1,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
+    )
+
+    assert original_num_frames == 1
+    assert padded_num_frames == 2
+    assert latents.shape == (1, 2, 128)
+
+
+def test_prepare_audio_latents_pads_packed_sequence_dim_for_provided_latents():
+    pipeline = _make_pipeline(LTX2Pipeline, sequence_parallel_size=4)
+    latents = torch.arange(40, dtype=torch.float32).view(1, 10, 4)
+
+    padded, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents(
+        batch_size=1,
+        num_channels_latents=2,
+        num_mel_bins=8,
+        audio_latent_length=10,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
+        latents=latents,
+    )
+
+    assert original_num_frames == 10
+    assert padded_num_frames == 12
+    assert padded.shape == (1, 12, 4)
+    torch.testing.assert_close(padded[:, :10], latents)
+    torch.testing.assert_close(padded[:, 10:], torch.zeros(1, 2, 4))
+
+
+def test_unpad_audio_latents_restores_original_frames_before_unpack():
+    original = torch.arange(40, dtype=torch.float32).view(1, 10, 4)
+    padded = torch.cat([original, torch.full((1, 2, 4), 999.0)], dim=1)
+
+    unpadded = latent_ops.unpad_audio_latents(padded, 10)
+    unpacked = latent_ops.unpack_audio_latents(unpadded, latent_length=10, num_mel_bins=2)
+    expected = latent_ops.unpack_audio_latents(original, latent_length=10, num_mel_bins=2)
+
+    assert unpacked.shape == (1, 2, 10, 2)
+    assert not (unpacked == 999.0).any()
+    torch.testing.assert_close(unpacked, expected)
+
+
+def test_prepare_audio_latents_accepts_already_padded_4d_latents_for_sp():
+    pipeline = _make_pipeline(LTX23Pipeline, sequence_parallel_size=4)
+    latents = torch.arange(96, dtype=torch.float32).view(1, 2, 12, 4)
+
+    audio_latent_length = pipeline._resolve_audio_latent_length(10, latents)
+    padded, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents(
+        batch_size=1,
+        num_channels_latents=2,
+        num_mel_bins=16,
+        audio_latent_length=audio_latent_length,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
+        latents=latents,
+    )
+
+    assert audio_latent_length == 10
+    assert original_num_frames == 10
+    assert padded_num_frames == 12
+    assert padded.shape == (1, 12, 8)
+    torch.testing.assert_close(padded, latent_ops.pack_audio_latents(latents))
+
+
+def test_resolve_audio_latent_length_preserves_legacy_4d_shape_inference():
+    pipeline = _make_pipeline(LTX23Pipeline, sequence_parallel_size=4)
+    latents = torch.zeros(1, 2, 13, 4)
+
+    audio_latent_length = pipeline._resolve_audio_latent_length(10, latents)
+
+    assert audio_latent_length == 13
+
+
+def test_prepare_audio_latents_rejects_incompatible_provided_length():
+    pipeline = _make_pipeline(LTX23Pipeline, sequence_parallel_size=4)
+    latents = torch.zeros(1, 11, 4)
+
+    with pytest.raises(ValueError, match="incompatible audio frame count"):
+        pipeline.prepare_audio_latents(
+            batch_size=1,
+            num_channels_latents=2,
+            num_mel_bins=8,
+            audio_latent_length=10,
+            dtype=torch.float32,
+            device=torch.device("cpu"),
+            latents=latents,
+        )

tests/diffusion/models/ltx2/test_ltx2_parallel.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""Unit tests for LTX-2.3 forward-parallel behavior."""
+"""Unit tests for LTX guidance and forward-parallel behavior."""
 
 from types import SimpleNamespace
 
@@ -11,29 +11,11 @@ import torch
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
-class TestCFGParallelIntegration:
-    def test_ltx23_pipeline_has_cfg_parallel_mixin(self):
-        """LTX23Pipeline must use the shared CFG parallel implementation."""
-        from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
-
-        assert issubclass(LTX23Pipeline, CFGParallelMixin)
-
-    def test_ltx2_transformer_has_dit_cache_config(self):
-        """Ensure LTX2 has a Cache DiT adapter config and that it uses separate CFG."""
-        from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
-        from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel
-
-        adapter_config = getattr(LTX2VideoTransformer3DModel, "_cache_dit_adapter_config")
-        assert isinstance(adapter_config, CacheDiTAdapterConfig)
-        assert adapter_config.has_separate_cfg
-
-
 class TestCFGParallelHelpers:
     """Test LTX-2.3 CFG helper math without loading model weights."""
 
     def test_combine_cfg_noise_matches_x0_space_formula(self):
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23Pipeline
 
         pipe = object.__new__(LTX23Pipeline)
         video_sample = torch.tensor([[[1.0, -2.0]]])
@@ -69,7 +51,8 @@ class TestCFGParallelHelpers:
         assert torch.allclose(audio_combined, expected_audio)
 
     def test_two_rank_cfg_parallel_smoke_uses_rank_local_branch_and_x0_formula(self, monkeypatch):
-        from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2_3 as ltx23
+        from vllm_omni.diffusion.models.ltx2 import ltx2_guidance
+        from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2 as ltx23
 
         pipe = object.__new__(ltx23.LTX23Pipeline)
         video_sample = torch.tensor([[[1.0, -2.0]]])
@@ -89,17 +72,17 @@ class TestCFGParallelHelpers:
                     return [video_pos, video_neg]
                 return [audio_pos, audio_neg]
 
-        monkeypatch.setattr(ltx23, "get_classifier_free_guidance_world_size", lambda: 2)
-        monkeypatch.setattr(ltx23, "get_cfg_group", lambda: FakeCfgGroup())
+        monkeypatch.setattr(ltx2_guidance, "get_classifier_free_guidance_world_size", lambda: 2)
+        monkeypatch.setattr(ltx2_guidance, "get_cfg_group", lambda: FakeCfgGroup())
 
-        expected_video = ltx23.LTX23Pipeline._combine_x0_space_cfg(
+        expected_video = ltx2_guidance.combine_velocity_via_x0(
             video_sample,
             video_pos,
             video_neg,
             video_sigma,
             scale,
         )
-        expected_audio = ltx23.LTX23Pipeline._combine_x0_space_cfg(
+        expected_audio = ltx2_guidance.combine_velocity_via_x0(
             audio_sample,
             audio_pos,
             audio_neg,
@@ -109,7 +92,7 @@ class TestCFGParallelHelpers:
 
         for rank, expected_branch in ((0, "positive"), (1, "negative")):
             calls = []
-            monkeypatch.setattr(ltx23, "get_classifier_free_guidance_rank", lambda rank=rank: rank)
+            monkeypatch.setattr(ltx2_guidance, "get_classifier_free_guidance_rank", lambda rank=rank: rank)
 
             def fake_predict_noise(**kwargs):
                 calls.append(kwargs["branch"])
@@ -140,17 +123,20 @@ class TestCFGParallelHelpers:
 class TestCFGParallelForwardPath:
     """Test the LTX-2.3 CFG-parallel denoising path without loading model weights."""
 
-    def test_forward_collates_request_prompt_embeds_and_mask_aliases(self, monkeypatch):
-        from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2_3 as ltx23
+    @pytest.mark.parametrize("pipeline_cls_name", ["LTX2Pipeline", "LTX23Pipeline"])
+    def test_forward_collates_request_prompt_embeds_and_mask_aliases(self, monkeypatch, pipeline_cls_name):
+        from vllm_omni.diffusion.models.ltx2 import ltx2_pipeline_runtime
+        from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2 as ltx23
         from vllm_omni.diffusion.request import OmniDiffusionRequest
         from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
         from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
-        pipe = object.__new__(ltx23.LTX23Pipeline)
+        pipeline_cls = getattr(ltx23, pipeline_cls_name)
+        pipe = object.__new__(pipeline_cls)
         torch.nn.Module.__init__(pipe)
         pipe.device = torch.device("cpu")
         pipe.tokenizer_max_length = 4
-        monkeypatch.setattr(ltx23, "get_classifier_free_guidance_world_size", lambda: 1)
+        monkeypatch.setattr(ltx2_pipeline_runtime, "get_classifier_free_guidance_world_size", lambda: 1)
 
         class StopAtEncodePromptError(Exception):
             pass
@@ -247,7 +233,8 @@ class TestCFGParallelForwardPath:
         audio_sampling_rate,
         expected_frame_rate,
     ):
-        from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2_3 as ltx23
+        from vllm_omni.diffusion.models.ltx2 import ltx2_denoise, ltx2_guidance, ltx2_pipeline_runtime
+        from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2 as ltx23
         from vllm_omni.diffusion.request import OmniDiffusionRequest
         from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
         from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -291,15 +278,16 @@ class TestCFGParallelForwardPath:
                     return [audio_pos, audio_neg]
                 raise AssertionError(f"Unexpected gathered tensor: {tensor}")
 
-        monkeypatch.setattr(ltx23, "get_classifier_free_guidance_world_size", lambda: 2)
-        monkeypatch.setattr(ltx23, "get_classifier_free_guidance_rank", lambda: cfg_rank)
-        monkeypatch.setattr(ltx23, "get_cfg_group", lambda: FakeCfgGroup())
+        monkeypatch.setattr(ltx2_pipeline_runtime, "get_classifier_free_guidance_world_size", lambda: 2)
+        monkeypatch.setattr(ltx2_guidance, "get_classifier_free_guidance_world_size", lambda: 2)
+        monkeypatch.setattr(ltx2_guidance, "get_classifier_free_guidance_rank", lambda: cfg_rank)
+        monkeypatch.setattr(ltx2_guidance, "get_cfg_group", lambda: FakeCfgGroup())
 
         def fake_retrieve_timesteps(scheduler, num_inference_steps, device, timesteps, sigmas=None, mu=None):
             scheduler.sigmas = torch.tensor([0.25, 0.25], device=device)
             return torch.tensor([1.0, 0.5], device=device), 2
 
-        monkeypatch.setattr(ltx23, "retrieve_timesteps", fake_retrieve_timesteps)
+        monkeypatch.setattr(ltx2_denoise, "retrieve_timesteps", fake_retrieve_timesteps)
 
         class FakeScheduler:
             def __init__(self, name="video", calls=None):
@@ -402,14 +390,14 @@ class TestCFGParallelForwardPath:
 
         output = pipe.forward(DiffusionRequestBatch(requests=[req]))[0]
 
-        expected_video_noise = ltx23.LTX23Pipeline._combine_x0_space_cfg(
+        expected_video_noise = ltx2_guidance.combine_velocity_via_x0(
             video_latents,
             video_pos,
             video_neg,
             pipe.scheduler.sigmas[0],
             4.0,
         )
-        expected_audio_noise = ltx23.LTX23Pipeline._combine_x0_space_cfg(
+        expected_audio_noise = ltx2_guidance.combine_velocity_via_x0(
             audio_latents,
             audio_pos,
             audio_neg,

tests/diffusion/models/ltx2/test_ltx2_pipeline.py

@@ -1,16 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""Unit tests for LTX-2.3 pipeline integration.
-
-These tests verify:
-- Pipeline is properly registered in the diffusion registry
-- Post-process function is registered
-- Cache-DiT enablers are registered
-- Pipeline does NOT inherit from LTX2Pipeline
-- Vocoder sample rate detection logic
-- Re-export module works correctly
-"""
+"""Unit tests for the shared LTX pipeline runtime and public contracts."""
 
 import json
 import os
@@ -20,25 +11,37 @@ from types import SimpleNamespace
 import pytest
 import torch
 
+from vllm_omni.diffusion.data import DiffusionOutput
+from vllm_omni.diffusion.models.ltx2.ltx2_components import (
+    LTX2_COMPONENT_PROFILE,
+    LTX23_COMPONENT_PROFILE,
+)
+from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import LTXI2VConditioningMixin
+from vllm_omni.diffusion.models.ltx2.ltx2_denoise import LTXDenoiseExecutor, LTXPhaseResult
+from vllm_omni.diffusion.models.ltx2.ltx2_guidance import (
+    LTX_LEGACY_VELOCITY_GUIDANCE,
+    LTX_OFFICIAL_X0_GUIDANCE,
+)
+from vllm_omni.diffusion.models.ltx2.ltx2_latents import LTXAVState
+from vllm_omni.diffusion.models.ltx2.ltx2_pipeline_runtime import LTXPipelineRuntime
+from vllm_omni.diffusion.models.ltx2.ltx2_recipes import LTX2_ONE_STAGE_RECIPE, LTX23_ONE_STAGE_RECIPE
+from vllm_omni.diffusion.models.ltx2.ltx2_request import LTXRequestInputs
+from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import (
+    LTX2ImageToVideoPipeline,
+    LTX2Pipeline,
+    LTX23ImageToVideoPipeline,
+    LTX23Pipeline,
+    LTXOneStagePipeline,
+)
+from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import (
+    LTX2ImageToVideoTwoStagesPipeline,
+    LTX2TwoStagesPipeline,
+)
+
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
-def _make_ltx23_pipeline(sequence_parallel_size: int = 1):
-    from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
-
-    pipeline = object.__new__(LTX23Pipeline)
-    torch.nn.Module.__init__(pipeline)
-    pipeline.audio_vae_temporal_compression_ratio = 4
-    pipeline.audio_vae_mel_compression_ratio = 4
-    pipeline.od_config = SimpleNamespace(parallel_config=SimpleNamespace(sequence_parallel_size=sequence_parallel_size))
-    pipeline.audio_vae = SimpleNamespace(
-        latents_mean=torch.tensor(0.0),
-        latents_std=torch.tensor(1.0),
-    )
-    return pipeline
-
-
-def _make_ltx23_request_pipe(cls):
+def _make_ltx_request_pipe(cls):
     pipe = object.__new__(cls)
     torch.nn.Module.__init__(pipe)
     pipe.device = torch.device("cpu")
@@ -73,10 +76,263 @@ def _resolve_request_inputs_for_test(pipe, req):
     )
 
 
-class TestLTX23RequestParsing:
-    def test_t2v_and_i2v_share_request_input_resolution(self):
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline
+def test_ltx_versions_share_runtime_without_cross_version_inheritance():
+    from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import (
+        DistributedAutoencoderKLLTX2Video,
+    )
+
+    assert issubclass(LTX2Pipeline, LTXPipelineRuntime)
+    assert issubclass(LTX23Pipeline, LTXPipelineRuntime)
+    assert issubclass(LTX2TwoStagesPipeline, LTXPipelineRuntime)
+    assert not issubclass(LTX23Pipeline, LTX2Pipeline)
+    assert LTX2Pipeline.component_profile is LTX2_COMPONENT_PROFILE
+    assert LTX23Pipeline.component_profile is LTX23_COMPONENT_PROFILE
+    assert LTX2Pipeline.one_stage_recipe is LTX2_ONE_STAGE_RECIPE
+    assert LTX23Pipeline.one_stage_recipe is LTX23_ONE_STAGE_RECIPE
+    assert LTX2_COMPONENT_PROFILE.video_vae_cls is DistributedAutoencoderKLLTX2Video
+    assert LTX23_COMPONENT_PROFILE.video_vae_cls is DistributedAutoencoderKLLTX2Video
+
+
+def test_ltx_one_stage_variants_expose_shared_request_batch_and_distributed_decode_contracts():
+    one_stage_variants = (
+        LTX2Pipeline,
+        LTX2ImageToVideoPipeline,
+        LTX23Pipeline,
+        LTX23ImageToVideoPipeline,
+    )
+    assert all(pipeline_cls.supports_request_batch for pipeline_cls in one_stage_variants)
+    assert all(pipeline_cls.distributed_video_decode for pipeline_cls in one_stage_variants)
+    assert not LTX2TwoStagesPipeline.supports_request_batch
+    assert not LTX2ImageToVideoTwoStagesPipeline.supports_request_batch
+
+
+@pytest.mark.parametrize("pipeline_cls", [LTX2Pipeline, LTX23Pipeline])
+def test_ltx_request_batch_decode_splits_video_and_audio_per_request(pipeline_cls):
+    from vllm_omni.diffusion.request import OmniDiffusionRequest
+    from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+    requests = [
+        OmniDiffusionRequest(
+            prompt=f"prompt-{index}",
+            sampling_params=OmniDiffusionSamplingParams(),
+            request_id=f"request-{index}",
+        )
+        for index in range(2)
+    ]
+    request_batch = DiffusionRequestBatch(requests)
+    video = torch.tensor([[1.0], [2.0]])
+    audio = torch.tensor([[3.0], [4.0]])
+    pipe = object.__new__(pipeline_cls)
+    torch.nn.Module.__init__(pipe)
+    object.__setattr__(pipe, "_decode_output", lambda **_kwargs: DiffusionOutput(output=(video, audio)))
+    forward_ctx = SimpleNamespace(
+        req=request_batch,
+        request_inputs=SimpleNamespace(
+            output_type="np",
+            generator=None,
+            decode_timestep=0.0,
+            decode_noise_scale=None,
+        ),
+        prompt_context=SimpleNamespace(connector_prompt_embeds=torch.empty(2, 1)),
+        device=torch.device("cpu"),
+        batch_size=2,
+        num_videos_per_prompt=1,
+    )
+
+    outputs = pipe.decode_phase(
+        LTXPhaseResult(
+            forward_context=forward_ctx,
+            video=video,
+            audio=audio,
+        )
+    )
+
+    assert len(outputs) == 2
+    torch.testing.assert_close(outputs[0].output[0], video[:1])
+    torch.testing.assert_close(outputs[0].output[1], audio[:1])
+    torch.testing.assert_close(outputs[1].output[0], video[1:])
+    torch.testing.assert_close(outputs[1].output[1], audio[1:])
+
+
+def test_ltx_one_stage_variants_share_forward_template():
+    assert LTX2Pipeline._forward_impl is LTXOneStagePipeline._forward_impl
+    assert LTX2ImageToVideoPipeline._forward_impl is LTXOneStagePipeline._forward_impl
+    assert LTX23Pipeline._forward_impl is LTXOneStagePipeline._forward_impl
+    assert LTX23ImageToVideoPipeline._forward_impl is LTXOneStagePipeline._forward_impl
+    assert LTX2ImageToVideoPipeline.forward is LTX23ImageToVideoPipeline.forward
+
+
+def test_ltx_versions_share_request_prompt_and_step_templates():
+    shared_methods = (
+        "_get_gemma_prompt_embeds",
+        "encode_prompt",
+        "check_inputs",
+        "_resolve_request_inputs",
+        "_prepare_prompt_context",
+        "_denoise_step",
+    )
+    for method_name in shared_methods:
+        base_method = getattr(LTXPipelineRuntime, method_name)
+        assert getattr(LTX2Pipeline, method_name) is base_method
+        assert getattr(LTX23Pipeline, method_name) is base_method
+
+
+def test_ltx_versions_select_guidance_without_overriding_control_flow():
+    assert LTX2Pipeline.guidance_strategy is LTX_LEGACY_VELOCITY_GUIDANCE
+    assert LTX23Pipeline.guidance_strategy is LTX_OFFICIAL_X0_GUIDANCE
+    assert LTX2Pipeline._predict_noise_for_step is LTXPipelineRuntime._predict_noise_for_step
+    assert LTX23Pipeline._predict_noise_for_step is LTXPipelineRuntime._predict_noise_for_step
+    assert LTX2Pipeline.combine_cfg_noise is LTXPipelineRuntime.combine_cfg_noise
+    assert LTX23Pipeline.combine_cfg_noise is LTXPipelineRuntime.combine_cfg_noise
+
+
+def test_ltx2_two_stage_variants_share_stage_orchestration():
+    assert issubclass(LTX2ImageToVideoTwoStagesPipeline, LTX2TwoStagesPipeline)
+    assert issubclass(LTX2ImageToVideoTwoStagesPipeline, LTXI2VConditioningMixin)
+    assert LTX2ImageToVideoTwoStagesPipeline._run_two_stage is LTX2TwoStagesPipeline._run_two_stage
+
+
+def test_ltx_variants_share_denoise_loop_and_i2v_conditioning():
+    assert "_denoise_loop" not in LTX2Pipeline.__dict__
+    assert "_denoise_loop" not in LTX23Pipeline.__dict__
+    assert issubclass(LTX2ImageToVideoPipeline, LTXI2VConditioningMixin)
+    assert issubclass(LTX23ImageToVideoPipeline, LTXI2VConditioningMixin)
+    assert LTX2ImageToVideoPipeline.prepare_latents is LTXI2VConditioningMixin.prepare_latents
+    assert LTX23ImageToVideoPipeline.prepare_latents is LTXI2VConditioningMixin.prepare_latents
+    assert LTX2Pipeline.prepare_latents is LTXPipelineRuntime.prepare_latents
+    assert LTX23Pipeline.prepare_latents is LTXPipelineRuntime.prepare_latents
+    assert LTX2Pipeline.prepare_audio_latents is LTXPipelineRuntime.prepare_audio_latents
+    assert LTX23Pipeline.prepare_audio_latents is LTXPipelineRuntime.prepare_audio_latents
+    assert LTX2Pipeline.decode_phase is LTXPipelineRuntime.decode_phase
+    assert LTX23Pipeline.decode_phase is LTXPipelineRuntime.decode_phase
+
+
+def test_denoise_executor_owns_progress_and_interrupt():
+    updates = []
+
+    class Progress:
+        def __enter__(self):
+            return self
+
+        def __exit__(self, exc_type, exc, traceback):
+            return False
+
+        def update(self):
+            updates.append(True)
+
+    pipeline = SimpleNamespace(
+        interrupt=False,
+        progress_bar=lambda total: Progress(),
+    )
+    seen: list[tuple[int, float]] = []
+    timesteps = torch.tensor([3.0, 2.0, 1.0])
+    initial_state = LTXAVState(video=torch.tensor(0.0), audio=torch.tensor(10.0))
+
+    def step(index, timestep, state):
+        seen.append((index, timestep.item()))
+        pipeline.interrupt = True
+        return LTXAVState(video=state.video + 1, audio=state.audio + 1)
+
+    state = LTXDenoiseExecutor.run(pipeline, initial_state, timesteps, step)
+
+    assert seen == [(0, 3.0)]
+    assert updates == [True]
+    torch.testing.assert_close(state.video, torch.tensor(1.0))
+    torch.testing.assert_close(state.audio, torch.tensor(11.0))
+
+
+def test_ltx2_two_stage_reuses_prompt_context_between_phases():
+    request_inputs = LTXRequestInputs(
+        prompt="prompt",
+        negative_prompt="negative",
+        height=32,
+        width=32,
+        num_frames=1,
+        frame_rate=24.0,
+        num_inference_steps=4,
+        guidance_scale=4.0,
+        guidance_rescale=0.0,
+        num_videos_per_prompt=1,
+        generator=None,
+        latents=None,
+        audio_latents=None,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+        prompt_attention_mask=None,
+        negative_prompt_attention_mask=None,
+        decode_timestep=0.25,
+        decode_noise_scale=0.5,
+        output_type="np",
+        max_sequence_length=16,
+    )
+    prompt_context = object()
+    phase_calls = []
+
+    def resolve_request_inputs(req, **kwargs):
+        return request_inputs
+
+    def run_phase(req, inputs, *, prompt_context=None, **kwargs):
+        phase_calls.append((inputs, prompt_context))
+        if len(phase_calls) == 1:
+            assert prompt_context is None
+            assert inputs.decode_timestep == 0.25
+            assert inputs.decode_noise_scale == 0.5
+            context = prompt_context_sentinel
+            video = torch.tensor([1.0])
+            audio = torch.tensor([2.0])
+        else:
+            assert prompt_context is prompt_context_sentinel
+            torch.testing.assert_close(inputs.latents, torch.tensor([11.0]))
+            torch.testing.assert_close(inputs.audio_latents, torch.tensor([2.0]))
+            assert inputs.guidance_scale == 1.0
+            assert inputs.num_inference_steps == 3
+            assert inputs.decode_timestep == 0.0
+            assert inputs.decode_noise_scale is None
+            context = prompt_context
+            video = torch.tensor([3.0])
+            audio = torch.tensor([4.0])
+        return LTXPhaseResult(
+            forward_context=SimpleNamespace(prompt_context=context),
+            video=video,
+            audio=audio,
+        )
+
+    def decode_phase(phase):
+        return DiffusionOutput(output=(phase.video, phase.audio))
+
+    class FakeUpsampler(torch.nn.Module):
+        def forward(self, *, latents, output_type, return_dict):
+            assert output_type == "latent"
+            assert not return_dict
+            return (latents + 10,)
+
+    prompt_context_sentinel = prompt_context
+    pipeline = object.__new__(LTX2TwoStagesPipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.distilled = True
+    pipeline.upsample_pipe = FakeUpsampler()
+    object.__setattr__(pipeline, "_resolve_request_inputs", resolve_request_inputs)
+    object.__setattr__(pipeline, "run_phase", run_phase)
+    object.__setattr__(pipeline, "decode_phase", decode_phase)
+
+    output = pipeline.forward(SimpleNamespace())
+
+    assert len(phase_calls) == 2
+    assert phase_calls[1][1] is prompt_context_sentinel
+    torch.testing.assert_close(output.output[0], torch.tensor([3.0]))
+    torch.testing.assert_close(output.output[1], torch.tensor([4.0]))
+
+
+class TestLTXRequestParsing:
+    @pytest.mark.parametrize(
+        ("t2v_cls", "i2v_cls"),
+        [
+            (LTX2Pipeline, LTX2ImageToVideoPipeline),
+            (LTX23Pipeline, LTX23ImageToVideoPipeline),
+        ],
+    )
+    def test_t2v_and_i2v_share_request_input_resolution(self, t2v_cls, i2v_cls):
         from vllm_omni.diffusion.request import OmniDiffusionRequest
         from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
         from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -124,11 +380,11 @@ class TestLTX23RequestParsing:
         )
 
         resolved_t2v = _resolve_request_inputs_for_test(
-            _make_ltx23_request_pipe(LTX23Pipeline),
+            _make_ltx_request_pipe(t2v_cls),
             req,
         )
         resolved_i2v = _resolve_request_inputs_for_test(
-            _make_ltx23_request_pipe(LTX23ImageToVideoPipeline),
+            _make_ltx_request_pipe(i2v_cls),
             req,
         )
 
@@ -169,12 +425,12 @@ class TestLTX23RequestParsing:
         )
 
     def test_request_input_resolution_rejects_mixed_precomputed_prompt_fields(self):
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23Pipeline
         from vllm_omni.diffusion.request import OmniDiffusionRequest
         from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
         from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
-        pipe = _make_ltx23_request_pipe(LTX23Pipeline)
+        pipe = _make_ltx_request_pipe(LTX23Pipeline)
         req = DiffusionRequestBatch(
             [
                 OmniDiffusionRequest(
@@ -209,9 +465,9 @@ class TestLTX23RequestParsing:
 
 class TestLTX23ForwardStages:
     def test_encode_prompt_repeats_precomputed_embeds_for_num_outputs(self):
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23Pipeline
 
-        pipe = _make_ltx23_request_pipe(LTX23Pipeline)
+        pipe = _make_ltx_request_pipe(LTX23Pipeline)
         prompt_embeds = torch.tensor([[[1.0], [2.0]], [[3.0], [4.0]]])
         negative_prompt_embeds = torch.tensor([[[-1.0], [-2.0]], [[-3.0], [-4.0]]])
         prompt_attention_mask = torch.tensor([[True, False], [False, True]])
@@ -251,12 +507,12 @@ class TestLTX23ForwardStages:
         )
 
     def test_t2v_forward_delegates_to_shared_forward_impl(self):
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23Pipeline
         from vllm_omni.diffusion.request import OmniDiffusionRequest
         from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
         from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
-        pipe = _make_ltx23_request_pipe(LTX23Pipeline)
+        pipe = _make_ltx_request_pipe(LTX23Pipeline)
         req = DiffusionRequestBatch(
             [
                 OmniDiffusionRequest(
@@ -311,36 +567,10 @@ class TestLTX23ForwardStages:
         }
 
 
-class TestPipelineIndependence:
-    """Verify LTX23Pipeline is fully independent from LTX2Pipeline."""
-
-    def test_ltx23_pipeline_does_not_inherit_from_ltx2(self):
-        """LTX23Pipeline must NOT inherit from LTX2Pipeline."""
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
-
-        assert not issubclass(LTX23Pipeline, LTX2Pipeline), (
-            "LTX23Pipeline should be fully independent and not inherit from LTX2Pipeline"
-        )
-
-    def test_ltx23_pipeline_is_nn_module(self):
-        """LTX23Pipeline must be an nn.Module."""
-        import torch.nn as nn
-
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
-
-        assert issubclass(LTX23Pipeline, nn.Module)
-
-    def test_ltx23_pipeline_has_progress_bar(self):
-        """LTX23Pipeline must mix in ProgressBarMixin."""
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
-        from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
-
-        assert issubclass(LTX23Pipeline, ProgressBarMixin)
-
+class TestPipelineComponents:
     def test_ltx23_pipeline_declares_offload_components(self):
         """LTX23Pipeline must expose LTX-2.3-specific modules to offload discovery."""
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23Pipeline
         from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
 
         pipe = object.__new__(LTX23Pipeline)
@@ -363,17 +593,10 @@ class TestPipelineIndependence:
         assert modules.resident_names == ["vocoder"]
         assert len(modules.vaes) == 2
 
-    def test_ltx23_pipeline_has_diffusion_pipeline_profiler_mixin(self):
-        """LTX23Pipeline must support lightweight stage timing."""
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
-        from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
-
-        assert issubclass(LTX23Pipeline, DiffusionPipelineProfilerMixin)
-
 
 class TestLTX23DecodeConditioning:
     def test_decode_conditioning_expands_per_prompt_values_to_effective_batch(self):
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import _expand_per_prompt_decode_value
+        from vllm_omni.diffusion.models.ltx2.ltx2_pipeline_runtime import _expand_per_prompt_decode_value
 
         assert _expand_per_prompt_decode_value(
             [0.1, 0.2],
@@ -395,7 +618,7 @@ class TestLTX23DecodeConditioning:
         ) == [0.1, 0.2, 0.3, 0.4]
 
     def test_decode_conditioning_rejects_ambiguous_lengths(self):
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import _expand_per_prompt_decode_value
+        from vllm_omni.diffusion.models.ltx2.ltx2_pipeline_runtime import _expand_per_prompt_decode_value
 
         with pytest.raises(ValueError, match="decode_timestep"):
             _expand_per_prompt_decode_value(
@@ -409,28 +632,15 @@ class TestLTX23DecodeConditioning:
 class TestRegistryIntegration:
     """Verify all LTX-2.3 pipeline variants are registered."""
 
-    def test_pipeline_models_registered(self):
-        """LTX-2.3 pipeline variants must be in _DIFFUSION_MODELS."""
-        from vllm_omni.diffusion.registry import _DIFFUSION_MODELS
-
-        expected = [
-            "LTX23Pipeline",
-            "LTX23ImageToVideoPipeline",
-        ]
-        for name in expected:
-            assert name in _DIFFUSION_MODELS, f"{name} not found in _DIFFUSION_MODELS"
-
     def test_pipeline_module_paths(self):
         """Registry entries must point to the correct modules."""
         from vllm_omni.diffusion.registry import _DIFFUSION_MODELS
 
-        # T2V -> pipeline_ltx2_3
-        assert _DIFFUSION_MODELS["LTX23Pipeline"] == ("ltx2", "pipeline_ltx2_3", "LTX23Pipeline")
+        assert _DIFFUSION_MODELS["LTX23Pipeline"] == ("ltx2", "pipeline_ltx2", "LTX23Pipeline")
 
-        # I2V -> pipeline_ltx2_3_image2video
         assert _DIFFUSION_MODELS["LTX23ImageToVideoPipeline"] == (
             "ltx2",
-            "pipeline_ltx2_3_image2video",
+            "pipeline_ltx2",
             "LTX23ImageToVideoPipeline",
         )
 
@@ -439,6 +649,12 @@ class TestRegistryIntegration:
         from vllm_omni.diffusion.registry import _DIFFUSION_POST_PROCESS_FUNCS
 
         expected = [
+            "LTX2Pipeline",
+            "LTX2ImageToVideoPipeline",
+            "LTX2TwoStagesPipeline",
+            "LTX2ImageToVideoTwoStagesPipeline",
+            "LTX2T2VDMD2Pipeline",
+            "LTX2I2VDMD2Pipeline",
             "LTX23Pipeline",
             "LTX23ImageToVideoPipeline",
         ]
@@ -446,6 +662,35 @@ class TestRegistryIntegration:
             assert name in _DIFFUSION_POST_PROCESS_FUNCS, f"{name} not in _DIFFUSION_POST_PROCESS_FUNCS"
             assert _DIFFUSION_POST_PROCESS_FUNCS[name] == "get_ltx2_post_process_func"
 
+    @pytest.mark.parametrize(
+        "model_class_name",
+        [
+            "LTX2Pipeline",
+            "LTX2ImageToVideoPipeline",
+            "LTX2TwoStagesPipeline",
+            "LTX2ImageToVideoTwoStagesPipeline",
+            "LTX2T2VDMD2Pipeline",
+            "LTX2I2VDMD2Pipeline",
+            "LTX23Pipeline",
+            "LTX23ImageToVideoPipeline",
+        ],
+    )
+    def test_post_process_func_resolves_from_every_entry_module(self, model_class_name):
+        from vllm_omni.diffusion.registry import get_diffusion_post_process_func
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            vocoder_dir = os.path.join(tmpdir, "vocoder")
+            os.makedirs(vocoder_dir)
+            with open(os.path.join(vocoder_dir, "config.json"), "w") as config_file:
+                json.dump({"output_sampling_rate": 48000}, config_file)
+
+            post_process = get_diffusion_post_process_func(
+                SimpleNamespace(model_class_name=model_class_name, model=tmpdir)
+            )
+
+        result = post_process((torch.zeros(1), torch.zeros(1)))
+        assert result["audio_sample_rate"] == 48000
+
     def test_cache_dit_for_ltx2_does_not_have_custom_enablers_registered(self):
         """Pipeline variants are *not* registered in CUSTOM_DIT_ENABLERS."""
         from vllm_omni.diffusion.cache.cache_dit_backend import CUSTOM_DIT_ENABLERS
@@ -465,7 +710,7 @@ class TestVocoderSampleRateDetection:
 
     def test_detects_48khz_from_config(self):
         """Should detect output_sampling_rate=48000 from vocoder/config.json."""
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import _detect_vocoder_output_sample_rate
+        from vllm_omni.diffusion.models.ltx2.ltx2_components import _detect_vocoder_output_sample_rate
 
         with tempfile.TemporaryDirectory() as tmpdir:
             vocoder_dir = os.path.join(tmpdir, "vocoder")
@@ -478,7 +723,7 @@ class TestVocoderSampleRateDetection:
 
     def test_returns_none_for_no_output_sr(self):
         """Should return None if vocoder config has no output_sampling_rate."""
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import _detect_vocoder_output_sample_rate
+        from vllm_omni.diffusion.models.ltx2.ltx2_components import _detect_vocoder_output_sample_rate
 
         with tempfile.TemporaryDirectory() as tmpdir:
             vocoder_dir = os.path.join(tmpdir, "vocoder")
@@ -491,7 +736,7 @@ class TestVocoderSampleRateDetection:
 
     def test_returns_none_for_missing_directory(self):
         """Should return None if vocoder directory doesn't exist."""
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import _detect_vocoder_output_sample_rate
+        from vllm_omni.diffusion.models.ltx2.ltx2_components import _detect_vocoder_output_sample_rate
 
         result = _detect_vocoder_output_sample_rate("/nonexistent/path")
         assert result is None
@@ -504,7 +749,7 @@ class TestPostProcessFunction:
         """Post-process func should include audio_sample_rate when detected."""
         import torch
 
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import get_ltx2_post_process_func
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import get_ltx2_post_process_func
 
         with tempfile.TemporaryDirectory() as tmpdir:
             vocoder_dir = os.path.join(tmpdir, "vocoder")
@@ -531,7 +776,7 @@ class TestPostProcessFunction:
         """Post-process func should work without vocoder config (no audio_sample_rate key)."""
         import torch
 
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import get_ltx2_post_process_func
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import get_ltx2_post_process_func
 
         class MockConfig:
             model = "/nonexistent/path"
@@ -546,101 +791,3 @@ class TestPostProcessFunction:
         assert "video" in result
         assert "audio" in result
         assert "audio_sample_rate" not in result
-
-
-class TestInitExports:
-    """Test that __init__.py exports all LTX-2.3 classes."""
-
-    def test_all_ltx23_classes_exported(self):
-        """All LTX-2.3 pipeline classes must be in the ltx2 package __all__."""
-        from vllm_omni.diffusion.models import ltx2
-
-        expected_classes = [
-            "LTX23Pipeline",
-            "LTX23ImageToVideoPipeline",
-        ]
-        for name in expected_classes:
-            assert hasattr(ltx2, name), f"{name} not exported from ltx2 package"
-            assert name in ltx2.__all__, f"{name} not in ltx2.__all__"
-
-
-class TestAudioLatentSPPadding:
-    def test_prepare_audio_latents_pads_generated_dummy_length_for_sp(self):
-        pipeline = _make_ltx23_pipeline(sequence_parallel_size=2)
-
-        latents, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents(
-            batch_size=1,
-            num_channels_latents=8,
-            num_mel_bins=64,
-            audio_latent_length=1,
-            dtype=torch.float32,
-            device=torch.device("cpu"),
-        )
-
-        assert original_num_frames == 1
-        assert padded_num_frames == 2
-        assert latents.shape == (1, 2, 128)
-
-    def test_prepare_audio_latents_pads_provided_packed_sequence_dim_for_sp(self):
-        pipeline = _make_ltx23_pipeline(sequence_parallel_size=4)
-        latents = torch.arange(40, dtype=torch.float32).view(1, 10, 4)
-
-        padded, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents(
-            batch_size=1,
-            num_channels_latents=2,
-            num_mel_bins=8,
-            audio_latent_length=10,
-            dtype=torch.float32,
-            device=torch.device("cpu"),
-            latents=latents,
-        )
-
-        assert original_num_frames == 10
-        assert padded_num_frames == 12
-        assert padded.shape == (1, 12, 4)
-        torch.testing.assert_close(padded[:, :10], latents)
-        torch.testing.assert_close(padded[:, 10:], torch.zeros(1, 2, 4))
-
-    def test_prepare_audio_latents_accepts_already_padded_4d_latents_for_sp(self):
-        pipeline = _make_ltx23_pipeline(sequence_parallel_size=4)
-        latents = torch.arange(96, dtype=torch.float32).view(1, 2, 12, 4)
-
-        audio_latent_length = pipeline._resolve_audio_latent_length(10, latents)
-        padded, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents(
-            batch_size=1,
-            num_channels_latents=2,
-            num_mel_bins=16,
-            audio_latent_length=audio_latent_length,
-            dtype=torch.float32,
-            device=torch.device("cpu"),
-            latents=latents,
-        )
-
-        assert audio_latent_length == 10
-        assert original_num_frames == 10
-        assert padded_num_frames == 12
-        assert padded.shape == (1, 12, 8)
-        torch.testing.assert_close(padded, pipeline._pack_audio_latents(latents))
-
-    def test_resolve_audio_latent_length_preserves_legacy_4d_shape_inference(self):
-        pipeline = _make_ltx23_pipeline(sequence_parallel_size=4)
-        latents = torch.zeros(1, 2, 13, 4)
-
-        audio_latent_length = pipeline._resolve_audio_latent_length(10, latents)
-
-        assert audio_latent_length == 13
-
-    def test_prepare_audio_latents_rejects_incompatible_provided_length(self):
-        pipeline = _make_ltx23_pipeline(sequence_parallel_size=4)
-        latents = torch.zeros(1, 11, 4)
-
-        with pytest.raises(ValueError, match="incompatible audio frame count"):
-            pipeline.prepare_audio_latents(
-                batch_size=1,
-                num_channels_latents=2,
-                num_mel_bins=8,
-                audio_latent_length=10,
-                dtype=torch.float32,
-                device=torch.device("cpu"),
-                latents=latents,
-            )

tests/diffusion/models/ltx2/test_ltx2_transformer.py

@@ -1,9 +1,13 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import pytest
 from torch import nn
 
-from vllm_omni.diffusion.models.ltx2.ltx2_transformer import _make_rms_norm
+from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig
+from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel, _make_rms_norm
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
 def test_ltx_rms_norm_no_affine_identity_weight_is_non_persistent_buffer():
@@ -20,3 +24,29 @@ def test_ltx_rms_norm_affine_weight_remains_parameter():
     assert isinstance(dict(norm.named_parameters())["weight"], nn.Parameter)
     assert "weight" not in dict(norm.named_buffers())
     assert "weight" in norm.state_dict()
+
+
+def test_ltx_transformer_has_separate_cfg_cache_dit_config():
+    adapter_config = getattr(LTX2VideoTransformer3DModel, "_cache_dit_adapter_config")
+
+    assert isinstance(adapter_config, CacheDiTAdapterConfig)
+    assert adapter_config.has_separate_cfg
+
+
+def test_ltx_transformer_exposes_hsdp_shard_conditions_for_blocks():
+    model = object.__new__(LTX2VideoTransformer3DModel)
+    nn.Module.__init__(model)
+    model.transformer_blocks = nn.ModuleList([nn.Linear(4, 4) for _ in range(2)])
+    model.norm_out = nn.LayerNorm(4)
+
+    conditions = getattr(model, "_hsdp_shard_conditions", None)
+
+    assert conditions is not None
+    assert len(conditions) == 1
+
+    matched = []
+    for name, module in model.named_modules():
+        if any(condition(name, module) for condition in conditions):
+            matched.append(name)
+
+    assert matched == ["transformer_blocks.0", "transformer_blocks.1"]

tests/diffusion/models/ltx2/test_ltx2_vae.py

@@ -11,24 +11,56 @@ import torch
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
-class TestLTX23OutputRank:
-    def test_single_process_rank_is_output_rank(self, monkeypatch):
-        import vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 as ltx23
+class TestLTXOutputRank:
+    @pytest.mark.parametrize(
+        ("distributed_vae_state", "expected_decode_calls"),
+        [(False, 0), (True, 1), (RuntimeError("unavailable"), 0)],
+    )
+    def test_non_output_rank_only_enters_collective_vae_decode(
+        self,
+        monkeypatch,
+        distributed_vae_state,
+        expected_decode_calls,
+    ):
+        import vllm_omni.diffusion.models.ltx2.pipeline_ltx2 as ltx
 
-        monkeypatch.setattr(ltx23.torch.distributed, "is_initialized", lambda: False)
+        class FakeVae:
+            dtype = torch.float32
+            config = SimpleNamespace(timestep_conditioning=False)
 
-        assert ltx23._is_output_rank() is True
-        assert ltx23._should_decode_video_on_rank(SimpleNamespace(is_distributed_enabled=lambda: False)) is True
+            def __init__(self):
+                self.decode_calls = 0
 
-    def test_non_output_rank_skips_decode_unless_vae_decode_is_distributed(self, monkeypatch):
-        import vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 as ltx23
+            def is_distributed_enabled(self):
+                if isinstance(distributed_vae_state, Exception):
+                    raise distributed_vae_state
+                return distributed_vae_state
 
-        monkeypatch.setattr(ltx23.torch.distributed, "is_initialized", lambda: True)
-        monkeypatch.setattr(ltx23.torch.distributed, "get_rank", lambda: 1)
+            def decode(self, *_args, **_kwargs):
+                self.decode_calls += 1
+                return (torch.ones(1, 1),)
 
-        assert ltx23._is_output_rank() is False
-        assert ltx23._should_decode_video_on_rank(SimpleNamespace(is_distributed_enabled=lambda: False)) is False
-        assert ltx23._should_decode_video_on_rank(SimpleNamespace(is_distributed_enabled=lambda: True)) is True
+        monkeypatch.setattr(ltx.torch.distributed, "is_initialized", lambda: True)
+        monkeypatch.setattr(ltx.torch.distributed, "get_rank", lambda: 1)
+        pipe = object.__new__(ltx.LTX2Pipeline)
+        torch.nn.Module.__init__(pipe)
+        pipe.vae = FakeVae()
+
+        output = pipe._decode_output(
+            latents=torch.ones(1, 1),
+            audio_latents=torch.ones(1, 1),
+            output_type="np",
+            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.vae.decode_calls == expected_decode_calls
+        assert output.output[0].numel() == 0
+        assert output.output[1].numel() == 0
 
 
 class TestLTX23VaeDistributedDecode: