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

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

tests/diffusion/models/ltx2/test_ltx2_conditioning.py

@@ -0,0 +1,219 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Unit tests for LTX image-to-video input and conditioning behavior."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_ltx23_request_pipe(cls):
+    pipe = object.__new__(cls)
+    torch.nn.Module.__init__(pipe)
+    pipe.device = torch.device("cpu")
+    pipe.tokenizer_max_length = 99
+    return pipe
+
+
+class TestLTXImageToVideoForwardStages:
+    def test_forward_resolves_request_image_and_delegates_to_shared_forward_impl(self):
+        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
+
+        pipe = _make_ltx23_request_pipe(LTX23ImageToVideoPipeline)
+        image = torch.zeros(3, 8, 8)
+        req = DiffusionRequestBatch(
+            [
+                OmniDiffusionRequest(
+                    prompt={
+                        "prompt": "make the image move",
+                        "negative_prompt": "jitter",
+                        "multi_modal_data": {"image": image},
+                    },
+                    sampling_params=OmniDiffusionSamplingParams(
+                        height=384,
+                        width=512,
+                        num_frames=25,
+                        num_inference_steps=2,
+                    ),
+                    request_id="ltx23-i2v-forward-stage-delegation",
+                )
+            ]
+        )
+        seen = {}
+
+        def fake_forward_impl(req_arg, request_inputs, **kwargs):
+            seen["req"] = req_arg
+            seen["request_inputs"] = request_inputs
+            seen["kwargs"] = kwargs
+            return ["i2v-delegated"]
+
+        object.__setattr__(pipe, "_forward_impl", fake_forward_impl)
+
+        output = pipe.forward(req, noise_scale=0.5)
+
+        assert output == ["i2v-delegated"]
+        assert seen["req"] is req
+        assert seen["request_inputs"].prompt == ["make the image move"]
+        assert seen["request_inputs"].negative_prompt == ["jitter"]
+        assert seen["kwargs"]["image"] is image
+        assert seen["kwargs"]["noise_scale"] == 0.5
+
+    def test_denoise_timestep_kwargs_masks_video_only(self):
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline
+
+        pipe = object.__new__(LTX23ImageToVideoPipeline)
+        ts = torch.tensor([2.0, 4.0])
+        denoise_ctx = SimpleNamespace(
+            conditioning_mask=torch.tensor([[1.0, 0.0]]),
+            conditioning_mask_for_model=torch.tensor([[1.0, 0.0], [0.0, 1.0]]),
+        )
+
+        kwargs = pipe._denoise_timestep_kwargs(
+            ts,
+            SimpleNamespace(cfg_parallel_ready=False),
+            denoise_ctx,
+        )
+
+        torch.testing.assert_close(kwargs["timestep"], torch.tensor([[0.0, 2.0], [4.0, 0.0]]))
+        torch.testing.assert_close(kwargs["audio_timestep"], ts)
+        torch.testing.assert_close(kwargs["sigma"], ts)
+
+
+class TestLTXImageToVideoConditioning:
+    def test_ltx23_i2v_supports_image_input(self):
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline
+
+        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 import LTX23ImageToVideoPipeline
+
+        image = object()
+
+        assert LTX23ImageToVideoPipeline._resolve_single_prompt_image([image]) is image
+        with pytest.raises(ValueError, match="exactly one image per prompt"):
+            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 import LTX23ImageToVideoPipeline
+
+        image = torch.zeros(1, 3, 4, 4)
+        additional = {
+            "preprocessed_image": None,
+            "pixel_values": image,
+            "image": torch.ones_like(image),
+        }
+
+        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.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)
+        pipe.vae_spatial_compression_ratio = 1
+        pipe.vae_temporal_compression_ratio = 1
+        pipe.transformer_spatial_patch_size = 1
+        pipe.transformer_temporal_patch_size = 1
+
+        def fake_randn_tensor(shape, generator=None, device=None, dtype=None):
+            raise AssertionError("packed I2V latents should not be noised")
+
+        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]]])
+
+        out, conditioning_mask = pipe.prepare_latents(
+            image=None,
+            batch_size=1,
+            num_channels_latents=2,
+            height=1,
+            width=1,
+            num_frames=3,
+            noise_scale=1.0,
+            dtype=torch.float32,
+            device=torch.device("cpu"),
+            latents=latents,
+        )
+
+        torch.testing.assert_close(conditioning_mask, torch.tensor([[1.0, 0.0, 0.0]]))
+        torch.testing.assert_close(out, latents)
+
+    def test_ltx23_i2v_5d_latents_noise_preserves_conditioning_frame(self, monkeypatch):
+        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)
+        pipe.vae_spatial_compression_ratio = 1
+        pipe.vae_temporal_compression_ratio = 1
+        pipe.transformer_spatial_patch_size = 1
+        pipe.transformer_temporal_patch_size = 1
+        pipe.vae = SimpleNamespace(
+            latents_mean=torch.zeros(2),
+            latents_std=torch.ones(2),
+            config=SimpleNamespace(scaling_factor=1.0),
+        )
+
+        def fake_randn_tensor(shape, generator=None, device=None, dtype=None):
+            return torch.ones(shape, device=device, dtype=dtype)
+
+        monkeypatch.setattr(ltx2_conditioning, "randn_tensor", fake_randn_tensor)
+        monkeypatch.setattr(ltx2_latents, "randn_tensor", fake_randn_tensor)
+
+        latents = torch.tensor([[[[[10.0]], [[20.0]], [[30.0]]], [[[11.0]], [[21.0]], [[31.0]]]]])
+
+        out, conditioning_mask = pipe.prepare_latents(
+            image=None,
+            batch_size=1,
+            num_channels_latents=2,
+            height=1,
+            width=1,
+            num_frames=3,
+            noise_scale=1.0,
+            dtype=torch.float32,
+            device=torch.device("cpu"),
+            latents=latents,
+        )
+
+        torch.testing.assert_close(conditioning_mask, torch.tensor([[1.0, 0.0, 0.0]]))
+        torch.testing.assert_close(out, torch.tensor([[[10.0, 11.0], [1.0, 1.0], [1.0, 1.0]]]))
+
+    def test_ltx23_i2v_video_step_preserves_conditioning_frame(self):
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline
+
+        pipe = object.__new__(LTX23ImageToVideoPipeline)
+        torch.nn.Module.__init__(pipe)
+        pipe.transformer_spatial_patch_size = 1
+        pipe.transformer_temporal_patch_size = 1
+
+        class FakeScheduler:
+            def step(self, noise_pred, t, latents, return_dict=False):
+                return (latents + noise_pred + t,)
+
+        pipe.scheduler = FakeScheduler()
+        latents = torch.tensor([[[1.0], [2.0], [3.0]]])
+        noise_pred = torch.full_like(latents, 10.0)
+
+        out = pipe._step_video_latents_i2v(
+            noise_pred,
+            latents,
+            torch.tensor(0.5),
+            latent_num_frames=3,
+            latent_height=1,
+            latent_width=1,
+        )
+
+        torch.testing.assert_close(out[:, :1], latents[:, :1])
+        torch.testing.assert_close(out[:, 1:], latents[:, 1:] + noise_pred[:, 1:] + 0.5)