[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

vllm_omni/diffusion/models/ltx2/ltx2_components.py

@@ -550,6 +550,13 @@ def get_ltx2_post_process_func(od_config: Any):
         if not (isinstance(output, tuple) and len(output) == 2):
             return output
         video, audio = output
+        if isinstance(video, torch.Tensor) and video.dtype == torch.uint8:
+            if video.ndim != 5 or video.shape[-1] != 3 or not video.is_contiguous():
+                raise ValueError(
+                    "LTX uint8 video output must be contiguous BTHWC RGB, "
+                    f"got shape={tuple(video.shape)} stride={video.stride()}"
+                )
+            video = video.detach().cpu().numpy()
         if isinstance(audio, torch.Tensor):
             audio = audio.detach().cpu()
         result: dict[str, Any] = {"video": video, "audio": audio}

vllm_omni/diffusion/models/ltx2/ltx2_conditioning.py

@@ -18,6 +18,8 @@ from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img impo
 from diffusers.utils.torch_utils import randn_tensor
 from diffusers.video_processor import VideoProcessor
 
+from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import wrap_methods_by_paths
+
 from . import ltx2_latents as latent_ops
 from .ltx2_guidance import euler_step_from_velocity
 
@@ -392,6 +394,8 @@ class LTXI2VConditioningMixin:
             vae_scale_factor=self.vae_spatial_compression_ratio,
             resample="bilinear",
         )
+        if getattr(self, "enable_diffusion_pipeline_profiler", False):
+            wrap_methods_by_paths(self, ["video_processor.postprocess_video"])
 
     @staticmethod
     def _resolve_single_prompt_image(raw_image: Any) -> Any:

vllm_omni/diffusion/models/ltx2/ltx2_runtime.py

@@ -84,6 +84,45 @@ def _run_ltx_vocoder(vocoder: nn.Module, generated_mel: torch.Tensor) -> torch.T
         return vocoder(generated_mel.float()).to(input_dtype)
 
 
+def _prepare_ltx2_video_output(
+    video: torch.Tensor,
+    *,
+    do_normalize: bool,
+) -> torch.Tensor:
+    """Consume decoded BCTHW video and return contiguous uint8 BTHWC frames.
+
+    Keep denormalization, scaling, and rounding in the decoder output dtype,
+    then convert directly to uint8.  ``do_normalize=False`` means the decoder
+    output is expected to already be in ``[0, 1]``; the clamp remains a
+    defensive bound before byte conversion. Performing the reduction before
+    worker IPC cuts the video D2H payload from FP32 to uint8 without
+    materializing a full-size FP32 CUDA tensor.
+    """
+    if video.ndim != 5:
+        raise ValueError(f"Expected decoded BCTHW video, got shape {tuple(video.shape)}")
+    if video.shape[1] != 3:
+        raise ValueError(f"Expected RGB BCTHW video, got shape {tuple(video.shape)}")
+    if not video.is_floating_point():
+        raise ValueError(f"Expected floating-point decoded video, got {video.dtype}")
+
+    video = video.detach()
+    if do_normalize:
+        # Match VaeImageProcessor.denormalize in the source dtype. These
+        # in-place operations have the same dtype-rounding boundaries while
+        # avoiding another full-size decoded-video allocation.
+        video.mul_(0.5).add_(0.5).clamp_(0.0, 1.0)
+    else:
+        video.clamp_(0.0, 1.0)
+
+    # Source-dtype quantization can differ from the legacy FP32 presentation
+    # path by at most one uint8 level, while avoiding a full-size FP32 tensor.
+    video.mul_(255.0).round_()
+    return video.permute(0, 2, 3, 4, 1).to(
+        dtype=torch.uint8,
+        memory_format=torch.contiguous_format,
+    )
+
+
 def _expand_per_prompt_decode_value(
     value: float | list[float],
     *,
@@ -191,7 +230,16 @@ class LTXRuntime(
         initialize_pipeline_components(self, od_config)
         self._phase_adapter: LTXPhaseAdapterRuntime | None = build_ltx_phase_adapter(self)
         self.setup_diffusion_pipeline_profiler(
-            enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler
+            profiler_targets=[
+                "vae.encode",
+                "vae.decode",
+                "text_encoder.forward",
+                "video_processor.postprocess_video",
+                "_prepare_video_output_for_transport",
+                "audio_vae.decode",
+                "vocoder.forward",
+            ],
+            enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler,
         )
 
     def _forward_request(
@@ -534,6 +582,12 @@ class LTXRuntime(
             )
         return DiffusionOutput(output=output)
 
+    def _prepare_video_output_for_transport(self, video: torch.Tensor) -> torch.Tensor:
+        return _prepare_ltx2_video_output(
+            video,
+            do_normalize=bool(self.video_processor.config.do_normalize),
+        )
+
     def _decode_output(
         self,
         *,
@@ -599,7 +653,10 @@ class LTXRuntime(
             )
 
         if video.numel() > 0:
-            video = self.video_processor.postprocess_video(video, output_type=output_type)
+            if output_type == "np":
+                video = self._prepare_video_output_for_transport(video)
+            else:
+                video = self.video_processor.postprocess_video(video, output_type=output_type)
         generated_mel = self.audio_vae.decode(audio_latents.to(self.audio_vae.dtype), return_dict=False)[0]
         audio = _run_ltx_vocoder(self.vocoder, generated_mel)
         return self._make_output((video, audio))