[Diffusion][Performance] Optimize MiniMax-H3 video output transfer (#6824)

Signed-off-by: lishunyang12 <lishunyang12@163.com>

tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py

@@ -57,7 +57,8 @@ def test_h3_prepares_resolved_cache_state_immediately_before_denoise():
         return torch.zeros(1), torch.zeros(1)
 
     pipeline.diffuse = diffuse
-    pipeline.decode = Mock(return_value=(torch.zeros(1), torch.zeros(1)))
+    decoded_audio = torch.zeros(1)
+    pipeline.decode = Mock(return_value=(torch.zeros(1, 3, 1, 1, 1), decoded_audio))
     sampling = OmniDiffusionSamplingParams(
         quality="high",
         width=1344,
@@ -86,7 +87,9 @@ def test_h3_prepares_resolved_cache_state_immediately_before_denoise():
         num_inference_steps=50,
         extra_args={"task": "t2va", "aspect_ratio": "16:9"},
     )
-    assert output.output == pipeline.decode.return_value
+    assert output.output[0].shape == (1, 1, 1, 1, 3)
+    assert output.output[0].dtype == torch.uint8
+    assert output.output[1] is decoded_audio
 
 
 def test_pipeline_import_registry_and_component_discovery():
@@ -505,6 +508,26 @@ def test_joint_postprocess_is_multiprocessing_picklable():
     assert result["fps"] == 24
 
 
+def test_joint_video_output_is_quantized_before_transfer():
+    from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
+        _minimax_h3_post_process,
+        _prepare_minimax_h3_video_output,
+    )
+
+    video = torch.linspace(-0.1, 1.1, 2 * 3 * 2 * 4 * 5, dtype=torch.float64).reshape(2, 3, 2, 4, 5)
+    expected = torch.round(video.float().clamp(0, 1) * 255).to(torch.uint8).permute(0, 2, 3, 4, 1)
+    prepared = _prepare_minimax_h3_video_output(video)
+
+    assert prepared.dtype == torch.uint8
+    assert prepared.is_contiguous()
+    assert prepared.shape == (2, 2, 4, 5, 3)
+    torch.testing.assert_close(prepared, expected, rtol=0, atol=0)
+
+    result = _minimax_h3_post_process((prepared, torch.zeros(1, 2, 6)))
+    assert result["video"][0].dtype == np.uint8
+    np.testing.assert_array_equal(result["video"][0], expected[0].numpy())
+
+
 def test_cfg_parallel_is_rejected_for_distilled_checkpoint():
     from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
 
@@ -627,7 +650,7 @@ def _distilled_pipeline(diffuse_calls, base_schedule_by_partition):
         return torch.zeros(1), torch.zeros(1)
 
     pipeline.diffuse = diffuse
-    pipeline.decode = Mock(return_value=(torch.zeros(1), torch.zeros(1)))
+    pipeline.decode = Mock(return_value=(torch.zeros(1, 3, 1, 1, 1), torch.zeros(1)))
     return pipeline
 
 

tests/entrypoints/openai_api/test_video_api_utils.py

@@ -223,6 +223,19 @@ def test_channel_first_video_tensor_is_converted_to_channel_last():
     np.testing.assert_array_equal(frames, video.permute(1, 2, 3, 0).numpy())
 
 
+def test_contiguous_uint8_video_reaches_mux_without_conversion(monkeypatch):
+    video = np.arange(2 * 4 * 5 * 3, dtype=np.uint8).reshape(2, 4, 5, 3)
+
+    def fail_prepare(*args, **kwargs):
+        raise AssertionError("ready-to-encode uint8 video must not be normalized")
+
+    monkeypatch.setattr(video_api_utils, "_prepare_video_frames", fail_prepare)
+
+    frames = video_api_utils._coerce_video_to_uint8_frames(video)
+
+    assert frames is video
+
+
 def test_channel_first_video_tensor_uses_direct_planar_mux(monkeypatch):
     calls = []
 

vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py

@@ -361,7 +361,10 @@ def _minimax_h3_post_process(output, output_type: str = "np"):
     if output_type == "latent":
         return output
     if output_type == "np":
-        video = video.detach().float().cpu().permute(0, 2, 3, 4, 1).clamp(0, 1).numpy()
+        if video.dtype == torch.uint8 and video.ndim == 5 and video.shape[-1] == 3:
+            video = video.detach().cpu().numpy()
+        else:
+            video = video.detach().float().cpu().permute(0, 2, 3, 4, 1).clamp(0, 1).numpy()
         audio = audio.detach().float().cpu().numpy()
         video = [sample for sample in video]
     return {
@@ -372,6 +375,16 @@ def _minimax_h3_post_process(output, output_type: str = "np"):
     }
 
 
+def _prepare_minimax_h3_video_output(video: torch.Tensor) -> torch.Tensor:
+    """Quantize decoded frames in place before worker-to-engine transfer."""
+    video = video.detach().float()
+    video.clamp_(0, 1).mul_(255).round_()
+    return video.permute(0, 2, 3, 4, 1).to(
+        dtype=torch.uint8,
+        memory_format=torch.contiguous_format,
+    )
+
+
 def _register_dlo_component_cache(cache: BoundedAllocatorCache, *components: Any) -> None:
     for component in components:
         if component is not None:
@@ -2406,10 +2419,10 @@ class MiniMaxH3Pipeline(
                 height=context["height"],
                 width=context["width"],
             )
-            videos.append(video)
+            videos.append(_prepare_minimax_h3_video_output(video))
             audios.append(audio)
-        video = torch.cat(videos, dim=0)
-        audio = torch.cat(audios, dim=0)
+        video = videos[0] if len(videos) == 1 else torch.cat(videos, dim=0)
+        audio = audios[0] if len(audios) == 1 else torch.cat(audios, dim=0)
         return DiffusionOutput(
             output=(video, audio),
             post_process_func=get_minimax_h3_post_process_func(self.od_config),
@@ -2711,6 +2724,7 @@ class MiniMaxH3Pipeline(
             height=shape["height"],
             width=shape["width"],
         )
+        video = _prepare_minimax_h3_video_output(video)
         return DiffusionOutput(
             output=(video, audio),
             post_process_func=get_minimax_h3_post_process_func(self.od_config),

vllm_omni/entrypoints/openai/video_api_utils.py

@@ -553,6 +553,14 @@ def _coerce_prepared_video_to_uint8_frames(
 
 def _coerce_video_to_uint8_frames(video: Any) -> np.ndarray:
     """Convert a video payload into contiguous uint8 frames shaped (F, H, W, 3)."""
+    if (
+        isinstance(video, np.ndarray)
+        and video.dtype == np.uint8
+        and video.ndim == 4
+        and video.shape[-1] == 3
+        and video.flags.c_contiguous
+    ):
+        return video
     frames, frame_shape, common_dtype = _prepare_video_frames(video)
     return _coerce_prepared_video_to_uint8_frames(frames, frame_shape, common_dtype)