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

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

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))