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

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

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