[Perf][Model] Optimize LTX video output transport (#7000)

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

tests/diffusion/models/ltx2/test_ltx2_output_cuda.py

@@ -0,0 +1,136 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""CUDA coverage for LTX decoded-video output preparation."""
+
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+from diffusers.video_processor import VideoProcessor
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.diffusion.models.ltx2.ltx2_runtime import _prepare_ltx2_video_output
+from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@pytest.mark.parametrize("do_normalize", [True, False])
+def test_ltx_uint8_output_runs_on_cuda_with_bounded_rounding_drift(dtype, do_normalize):
+    processor = VideoProcessor(vae_scale_factor=8, do_normalize=do_normalize)
+    source = torch.linspace(
+        -1.2 if do_normalize else 0.0,
+        1.2 if do_normalize else 1.0,
+        2 * 3 * 4 * 5 * 7,
+        device="cuda",
+        dtype=torch.float32,
+    ).reshape(2, 3, 4, 5, 7)
+    source = source.to(dtype)
+
+    legacy = processor.postprocess_video(source.clone(), output_type="np")
+    if do_normalize:
+        expected = np.clip(legacy, 0.0, 1.0)
+    else:
+        # With normalization disabled, VideoProcessor's contract is already-[0, 1] input.
+        np.testing.assert_array_less(legacy, 1.0 + np.finfo(np.float32).eps)
+        np.testing.assert_array_less(-np.finfo(np.float32).eps, legacy)
+        expected = legacy
+    expected = np.rint(expected * 255.0).astype(np.uint8)
+
+    prepared = _prepare_ltx2_video_output(source.clone(), do_normalize=do_normalize)
+
+    assert prepared.is_cuda
+    assert prepared.shape == (2, 4, 5, 7, 3)
+    assert prepared.dtype == torch.uint8
+    assert prepared.is_contiguous()
+    delta = np.abs(prepared.cpu().numpy().astype(np.int16) - expected.astype(np.int16))
+    assert delta.max() <= 1
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_ltx_uint8_output_defensively_clamps_unnormalized_values_on_cuda():
+    source = (
+        torch.tensor([-0.2, 0.0, 0.5, 1.0, 1.2], device="cuda", dtype=torch.bfloat16)
+        .reshape(1, 1, 1, 1, 5)
+        .expand(1, 3, 1, 1, 5)
+        .clone()
+    )
+
+    prepared = _prepare_ltx2_video_output(source, do_normalize=False)
+
+    expected = torch.tensor([0, 0, 128, 255, 255], device="cuda", dtype=torch.uint8).reshape(1, 1, 1, 5)
+    assert torch.equal(prepared[..., 0], expected)
+    assert torch.equal(prepared[..., 1], expected)
+    assert torch.equal(prepared[..., 2], expected)
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("use_diffusion_decoder", [False, True], ids=["conv-vae", "diff-vae"])
+def test_ltx_bf16_decoder_output_uses_cuda_transport_path(use_diffusion_decoder):
+    decoded_video = torch.linspace(
+        -1.2,
+        1.2,
+        1 * 3 * 2 * 3 * 4,
+        device="cuda",
+        dtype=torch.float32,
+    ).reshape(1, 3, 2, 3, 4)
+    decoded_video = decoded_video.to(torch.bfloat16)
+    processor = VideoProcessor(vae_scale_factor=8)
+    legacy = processor.postprocess_video(decoded_video.clone(), output_type="np")
+    expected = np.rint(np.clip(legacy, 0.0, 1.0) * 255.0).astype(np.uint8)
+    decode_calls = []
+
+    class VideoDecoder:
+        dtype = torch.bfloat16
+        config = SimpleNamespace(timestep_conditioning=False)
+
+        def __init__(self, name):
+            self.name = name
+
+        def decode(self, *_args, **_kwargs):
+            decode_calls.append(self.name)
+            return (decoded_video.clone(),)
+
+    class AudioVae:
+        dtype = torch.bfloat16
+
+        def decode(self, audio_latents, *, return_dict):
+            assert return_dict is False
+            return (audio_latents,)
+
+    pipe = object.__new__(LTX2Pipeline)
+    torch.nn.Module.__init__(pipe)
+    pipe.reports_stage_durations = False
+    pipe.distributed_video_decode = False
+    pipe.use_diffusion_decoder = use_diffusion_decoder
+    pipe.vae = VideoDecoder("conv-vae")
+    pipe.diffusion_decoder = VideoDecoder("diff-vae")
+    pipe.audio_vae = AudioVae()
+    pipe.vocoder = torch.nn.Identity()
+    pipe.video_processor = processor
+
+    output = pipe._decode_output(
+        latents=torch.ones(1, 1, device="cuda", dtype=torch.bfloat16),
+        audio_latents=torch.ones(1, 1, device="cuda", dtype=torch.bfloat16),
+        output_type="np",
+        connector_prompt_embeds=torch.ones(1, 1, device="cuda", dtype=torch.bfloat16),
+        generator=None,
+        device=torch.device("cuda"),
+        decode_timestep=0.0,
+        decode_noise_scale=None,
+        prompt_batch_size=1,
+    )
+
+    prepared = output.output[0]
+    selected_decoder = "diff-vae" if use_diffusion_decoder else "conv-vae"
+    assert decode_calls == [selected_decoder]
+    assert prepared.is_cuda
+    assert prepared.shape == (1, 2, 3, 4, 3)
+    assert prepared.dtype == torch.uint8
+    assert prepared.is_contiguous()
+    delta = np.abs(prepared.cpu().numpy().astype(np.int16) - expected.astype(np.int16))
+    assert delta.max() <= 1

tests/diffusion/models/ltx2/test_ltx2_pipeline.py

@@ -10,8 +10,10 @@ from dataclasses import replace
 from types import SimpleNamespace
 from typing import Any
 
+import numpy as np
 import pytest
 import torch
+from diffusers.video_processor import VideoProcessor
 
 from vllm_omni.diffusion.data import DiffusionOutput
 from vllm_omni.diffusion.models.ltx2 import (
@@ -65,7 +67,10 @@ from vllm_omni.diffusion.models.ltx2.ltx2_request import (
     validate_ltx_checkpoint,
     validate_pipeline_request,
 )
-from vllm_omni.diffusion.models.ltx2.ltx2_runtime import LTXRuntime
+from vllm_omni.diffusion.models.ltx2.ltx2_runtime import (
+    LTXRuntime,
+    _prepare_ltx2_video_output,
+)
 from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import (
     LTX2DistilledOneStagePipeline,
     LTX2I2VDMD2Pipeline,
@@ -123,6 +128,103 @@ def _resolve_request_inputs_for_test(
     )
 
 
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32])
+@pytest.mark.parametrize("do_normalize", [True, False])
+def test_ltx_device_uint8_output_stays_within_one_level_of_legacy_rounding(dtype, do_normalize):
+    processor = VideoProcessor(vae_scale_factor=8, do_normalize=do_normalize)
+    source = torch.linspace(
+        -1.2 if do_normalize else 0.0,
+        1.2 if do_normalize else 1.0,
+        2 * 3 * 4 * 5 * 7,
+        dtype=torch.float32,
+    ).reshape(2, 3, 4, 5, 7)
+    source = source.to(dtype)
+
+    legacy = processor.postprocess_video(source.clone(), output_type="np")
+    if do_normalize:
+        expected = np.clip(legacy, 0.0, 1.0)
+    else:
+        # With normalization disabled, VideoProcessor's contract is already-[0, 1] input.
+        np.testing.assert_array_less(legacy, 1.0 + np.finfo(np.float32).eps)
+        np.testing.assert_array_less(-np.finfo(np.float32).eps, legacy)
+        expected = legacy
+    expected = np.rint(expected * 255.0).astype(np.uint8)
+
+    prepared = _prepare_ltx2_video_output(
+        source.clone(),
+        do_normalize=do_normalize,
+    )
+
+    assert prepared.shape == (2, 4, 5, 7, 3)
+    assert prepared.dtype == torch.uint8
+    assert prepared.is_contiguous()
+    delta = np.abs(prepared.numpy().astype(np.int16) - expected.astype(np.int16))
+    assert delta.max() <= (0 if dtype == torch.float32 else 1)
+
+
+@pytest.mark.parametrize(
+    ("video", "error"),
+    [
+        (torch.zeros(1, 3, 4, 5), "BCTHW"),
+        (torch.zeros(1, 4, 2, 4, 5), "RGB"),
+        (torch.zeros(1, 3, 2, 4, 5, dtype=torch.uint8), "floating-point"),
+    ],
+)
+def test_ltx_device_uint8_output_rejects_invalid_contract(video, error):
+    with pytest.raises(ValueError, match=error):
+        _prepare_ltx2_video_output(video, do_normalize=True)
+
+
+def test_ltx_profiler_covers_video_output_transport(tmp_path, monkeypatch):
+    from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
+
+    (tmp_path / "model_index.json").write_text(json.dumps({"vocoder": ["ltx2", "LTX2Vocoder"]}))
+
+    def stub_components(pipe, od_config):
+        pipe.od_config = od_config
+        pipe.vae_spatial_compression_ratio = 32
+
+    profiler_setup: dict[str, Any] = {}
+
+    def capture_profiler_setup(_pipe, **kwargs):
+        profiler_setup.update(kwargs)
+        _pipe.enable_diffusion_pipeline_profiler = False
+
+    monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)
+    monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", capture_profiler_setup)
+
+    LTX2Pipeline(od_config=SimpleNamespace(model=str(tmp_path), enable_diffusion_pipeline_profiler=False))
+
+    assert "video_processor.postprocess_video" in profiler_setup["profiler_targets"]
+    assert "_prepare_video_output_for_transport" in profiler_setup["profiler_targets"]
+
+
+def test_ltx_i2v_rewraps_replaced_video_processor_for_profiler(tmp_path, monkeypatch):
+    from vllm_omni.diffusion.models.ltx2 import ltx2_conditioning, ltx2_runtime
+
+    (tmp_path / "model_index.json").write_text(json.dumps({"vocoder": ["ltx2", "LTX2Vocoder"]}))
+
+    def stub_components(pipe, od_config):
+        pipe.od_config = od_config
+        pipe.vae_spatial_compression_ratio = 32
+
+    def enable_profiler(pipe, **_kwargs):
+        pipe.enable_diffusion_pipeline_profiler = True
+
+    wrapped: list[tuple[Any, list[str]]] = []
+    monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)
+    monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", enable_profiler)
+    monkeypatch.setattr(
+        ltx2_conditioning,
+        "wrap_methods_by_paths",
+        lambda pipe, paths: wrapped.append((pipe.video_processor, paths)),
+    )
+
+    pipe = LTX2Pipeline(od_config=SimpleNamespace(model=str(tmp_path), enable_diffusion_pipeline_profiler=True))
+
+    assert wrapped == [(pipe.video_processor, ["video_processor.postprocess_video"])]
+
+
 def test_ltx_public_entries_share_runtime_and_keep_recipe_boundaries():
     from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import (
         DistributedAutoencoderKLLTX2Video,
@@ -2140,6 +2242,12 @@ class TestPostProcessFunction:
             assert "audio" in result
             assert result["audio_sample_rate"] == 48000
 
+            prepared_video = torch.zeros(1, 4, 64, 64, 3, dtype=torch.uint8)
+            prepared_result = func((prepared_video, audio))
+            assert isinstance(prepared_result["video"], np.ndarray)
+            assert prepared_result["video"].shape == (1, 4, 64, 64, 3)
+            assert prepared_result["video"].dtype == np.uint8
+
     def test_post_process_without_vocoder_config(self):
         """Post-process func should work without vocoder config (no audio_sample_rate key)."""
         import torch