[Model][Frontend] MiniMax-H3: Parallelize MP4 response frame conversion (#6499)

Signed-off-by: MosCloud <bwzhang1991@163.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/entrypoints/openai_api/test_api_server_guards.py

@@ -398,6 +398,12 @@ async def test_api_server_assembly_replaces_upstream_routes_and_mounts_omni_rout
     async def fake_omni_init_app_state(engine_client, state, args):
         state.engine_client = engine_client
         state.initialized_by_omni = True
+        state.openai_serving_video = SimpleNamespace(
+            shutdown=lambda: captured.__setitem__("video_shutdown", True),
+        )
+        state.openai_serving_speech = SimpleNamespace(
+            shutdown=lambda: captured.__setitem__("speech_shutdown", True),
+        )
 
     async def fake_storage_start():
         captured["storage_started"] = True
@@ -433,6 +439,8 @@ async def test_api_server_assembly_replaces_upstream_routes_and_mounts_omni_rout
     assert captured["supported_tasks"] == ("generate",)
     assert captured["storage_started"] is True
     assert captured["restrictions"] == {}
+    assert captured["video_shutdown"] is True
+    assert captured["speech_shutdown"] is True
     assert sock.closed is True
     assert served_app.state.initialized_by_omni is True
 

tests/entrypoints/openai_api/test_video_api_utils.py

@@ -3,6 +3,7 @@
 """Unit tests for OpenAI-compatible video API encoding helpers."""
 
 import base64
+import threading
 from io import BytesIO
 
 import av
@@ -249,6 +250,95 @@ def test_channel_first_video_tensor_uses_direct_planar_mux(monkeypatch):
     assert options["fps"] == 12.0
 
 
+def test_planar_converter_serial_and_parallel_outputs_match():
+    frames = [np.full((2, 3, 3), fill_value=value / 10, dtype=np.float32) for value in (1, 4, 7)]
+
+    serial_converter = video_api_utils._PlanarFrameConverter(max_workers=1)
+    parallel_converter = video_api_utils._PlanarFrameConverter(max_workers=2)
+    try:
+        serial = [
+            frame.to_ndarray(format="rgb24") for frame in serial_converter.iter_frames(frames, np.dtype(np.float32))
+        ]
+        parallel = [
+            frame.to_ndarray(format="rgb24") for frame in parallel_converter.iter_frames(frames, np.dtype(np.float32))
+        ]
+    finally:
+        serial_converter.shutdown()
+        parallel_converter.shutdown()
+
+    np.testing.assert_array_equal(np.stack(serial), np.stack(parallel))
+
+
+def test_planar_converter_preserves_fifo_order_and_bounds_pending_futures(monkeypatch):
+    converter = video_api_utils._PlanarFrameConverter(max_workers=2)
+    started = threading.Event()
+    release = threading.Event()
+    four_submitted = threading.Event()
+    outstanding = 0
+    peak_outstanding = 0
+    lock = threading.Lock()
+    real_submit = converter._executor.submit
+
+    def fake_build(frame, _common_dtype):
+        started.set()
+        release.wait(timeout=2)
+        return int(frame[0, 0, 0])
+
+    def submit(*args, **kwargs):
+        nonlocal outstanding, peak_outstanding
+        with lock:
+            outstanding += 1
+            peak_outstanding = max(peak_outstanding, outstanding)
+            if peak_outstanding >= 4:
+                four_submitted.set()
+        future = real_submit(*args, **kwargs)
+
+        def done(_future):
+            nonlocal outstanding
+            with lock:
+                outstanding -= 1
+
+        future.add_done_callback(done)
+        return future
+
+    monkeypatch.setattr(converter, "_build_frame", fake_build)
+    monkeypatch.setattr(converter._executor, "submit", submit)
+    frames = [np.full((1, 1, 3), value, dtype=np.uint8) for value in range(8)]
+    result: list[int] = []
+    iterator = converter.iter_frames(frames, np.dtype(np.uint8))
+    consumer = threading.Thread(target=lambda: result.extend(iterator))
+    consumer.start()
+    try:
+        assert started.wait(timeout=2)
+        assert four_submitted.wait(timeout=2)
+        release.set()
+        consumer.join(timeout=2)
+        assert not consumer.is_alive()
+    finally:
+        release.set()
+        iterator.close()
+        converter.shutdown()
+
+    assert result == list(range(8))
+    assert peak_outstanding <= 4
+
+
+def test_planar_converter_reuses_executor_and_shutdowns_after_cancel():
+    converter = video_api_utils._PlanarFrameConverter(max_workers=2)
+    executor = converter._executor
+    frames = [np.zeros((2, 2, 3), dtype=np.uint8) for _ in range(4)]
+    try:
+        list(converter.iter_frames(frames, np.dtype(np.uint8)))
+        iterator = converter.iter_frames(frames, np.dtype(np.uint8))
+        next(iterator)
+        iterator.close()
+        assert converter._executor is executor
+    finally:
+        converter.shutdown()
+
+    assert executor._shutdown
+
+
 @pytest.mark.parametrize(
     ("video", "expected_path", "expected_reason", "expected_mux"),
     [
@@ -347,6 +437,35 @@ def test_video_only_does_not_log_default_audio_sample_rate(monkeypatch):
     assert log_messages == []
 
 
+def test_planar_mux_accepts_none_rate_and_defaults_audio_to_44100():
+    def make_frame():
+        return next(
+            video_api_utils._iter_planar_video_frames(
+                [np.zeros((2, 2, 3), dtype=np.uint8)],
+                np.dtype(np.uint8),
+            )
+        )
+
+    video_only = media_utils.mux_av_video_audio_bytes(
+        [make_frame()],
+        width=2,
+        height=2,
+        audio_sample_rate=None,
+    )
+    with av.open(BytesIO(video_only), mode="r", format="mp4") as container:
+        assert not container.streams.audio
+
+    with_audio = media_utils.mux_av_video_audio_bytes(
+        [make_frame()],
+        width=2,
+        height=2,
+        audio_waveform=np.zeros((1, 1024), dtype=np.float32),
+        audio_sample_rate=None,
+    )
+    with av.open(BytesIO(with_audio), mode="r", format="mp4") as container:
+        assert container.streams.audio[0].rate == 44100
+
+
 def test_interleaved_video_uses_legacy_fallback_automatically(monkeypatch):
     calls = []
     path_logs = []
@@ -383,6 +502,7 @@ def test_interleaved_video_uses_legacy_fallback_automatically(monkeypatch):
     assert "reason=non_contiguous_rgb_planes" in path_logs[0]
     assert "audio_present=False" in path_logs[0]
     assert "effective_audio_sample_rate=None" in path_logs[0]
+    assert "effective_frame_conversion_workers=0" in path_logs[0]
     assert len(coerce_calls) == 1
 
 

tests/entrypoints/openai_api/test_video_server.py

@@ -95,12 +95,13 @@ class FakeAsyncOmni:
         yield MockVideoResult(videos)
 
 
-def test_raw_and_base64_encoders_receive_no_policy_config(mocker: MockerFixture):
+def test_raw_and_base64_encoders_receive_persistent_converter(mocker: MockerFixture):
     engine = FakeAsyncOmni()
     handler = OmniOpenAIServingVideo.for_diffusion(
         engine,
         model_name="test-model",
     )
+    assert handler._video_frame_converter.max_workers == 8
     raw_encoder = mocker.patch(
         "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
         return_value=b"encoded-video",
@@ -117,8 +118,9 @@ def test_raw_and_base64_encoders_receive_no_policy_config(mocker: MockerFixture)
 
     asyncio.run(_generate_both_response_types())
 
-    assert "encoding_config" not in raw_encoder.call_args.kwargs
-    assert "encoding_config" not in base64_encoder.call_args.kwargs
+    assert raw_encoder.call_args.kwargs["frame_converter"] is handler._video_frame_converter
+    assert base64_encoder.call_args.kwargs["frame_converter"] is handler._video_frame_converter
+    handler.shutdown()
 
 
 def test_resolve_diffusion_od_config_falls_back_to_attribute():
@@ -1231,8 +1233,9 @@ def test_audio_sample_rate_comes_from_model_config(test_client, mocker: MockerFi
         audio=None,
         audio_sample_rate=None,
         video_codec_options=None,
+        frame_converter=None,
     ):
-        del video, fps, audio, video_codec_options
+        del video, fps, audio, video_codec_options, frame_converter
         audio_sample_rates.append(audio_sample_rate)
         return b"fake-video"
 

vllm_omni/diffusion/utils/media_utils.py

@@ -210,7 +210,7 @@ def mux_av_video_audio_bytes(
     audio_waveform: np.ndarray | None = None,
     *,
     fps: float = 25.0,
-    audio_sample_rate: int = 44100,
+    audio_sample_rate: int | None = None,
     video_codec: str = "h264",
     audio_codec: str = "aac",
     crf: str = "18",
@@ -236,6 +236,7 @@ def mux_av_video_audio_bytes(
         samples: np.ndarray | None = None
         layout: str | None = None
         if audio_waveform is not None:
+            effective_audio_sample_rate = 44100 if audio_sample_rate is None else audio_sample_rate
             samples = audio_waveform.astype(np.float32)
             if samples.ndim == 1:
                 samples = samples.reshape(1, -1)
@@ -243,7 +244,7 @@ def mux_av_video_audio_bytes(
                 samples = np.ascontiguousarray(samples.T)
             num_channels = samples.shape[0]
             layout = "stereo" if num_channels >= 2 else "mono"
-            a_stream = cast(av.AudioStream, container.add_stream(audio_codec, rate=audio_sample_rate))
+            a_stream = cast(av.AudioStream, container.add_stream(audio_codec, rate=effective_audio_sample_rate))
             a_stream.layout = layout
 
         for frame in video_frames:
@@ -256,12 +257,12 @@ def mux_av_video_audio_bytes(
             if samples is None or layout is None:
                 raise ValueError("Audio samples were not prepared for muxing.")
             audio_frame = av.AudioFrame.from_ndarray(samples, format="fltp", layout=layout)
-            audio_frame.sample_rate = audio_sample_rate
+            audio_frame.sample_rate = effective_audio_sample_rate
             # AAC has a one-frame encoder delay. Mark the input waveform as
             # starting at t=0 so the MP4 muxer writes the corresponding negative
             # priming timestamp instead of exposing the delay as leading silence.
             audio_frame.pts = 0
-            audio_frame.time_base = Fraction(1, audio_sample_rate)
+            audio_frame.time_base = Fraction(1, effective_audio_sample_rate)
             for packet in a_stream.encode(audio_frame):
                 container.mux(packet)
             for packet in a_stream.encode():

vllm_omni/entrypoints/openai/api_server.py

@@ -615,6 +615,9 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None,
             await shutdown_task
         finally:
             state = getattr(app, "state", None)
+            serving_video = getattr(state, "openai_serving_video", None) if state is not None else None
+            if serving_video is not None:
+                serving_video.shutdown()
             serving_speech = getattr(state, "openai_serving_speech", None) if state is not None else None
             if serving_speech is not None:
                 serving_speech.shutdown()

vllm_omni/entrypoints/openai/serving_video.py

@@ -32,6 +32,7 @@ from vllm_omni.entrypoints.openai.stage_params import (
 from vllm_omni.entrypoints.openai.utils import is_video_generation_pipeline, parse_lora_request
 from vllm_omni.entrypoints.openai.video_api_utils import (
     _encode_video_bytes,
+    _PlanarFrameConverter,
     encode_video_base64,
 )
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt
@@ -44,6 +45,8 @@ from vllm_omni.outputs.output_metadata import (
 
 logger = init_logger(__name__)
 
+_VIDEO_RESPONSE_FRAME_CONVERSION_WORKERS = 8
+
 if TYPE_CHECKING:
     from vllm_omni.diffusion.data import OmniDiffusionConfig
 
@@ -96,6 +99,11 @@ class OmniOpenAIServingVideo:
         self._engine_client = engine_client
         self._model_name = model_name
         self._stage_configs = stage_configs
+        self._video_frame_converter = _PlanarFrameConverter(max_workers=_VIDEO_RESPONSE_FRAME_CONVERSION_WORKERS)
+        logger.info(
+            "Video response frame conversion pool configured: workers=%d",
+            self._video_frame_converter.max_workers,
+        )
 
     def _resolve_diffusion_od_config(self) -> OmniDiffusionConfig | SimpleNamespace | None:
         get_od_config = getattr(self._engine_client, "get_diffusion_od_config", None)
@@ -170,6 +178,9 @@ class OmniOpenAIServingVideo:
             stage_configs=stage_configs,
         )
 
+    def shutdown(self) -> None:
+        self._video_frame_converter.shutdown()
+
     async def _run_and_extract(
         self,
         request: VideoGenerationRequest,
@@ -361,6 +372,7 @@ class OmniOpenAIServingVideo:
                         video,
                         fps=artifacts.output_fps,
                         video_codec_options=video_codec_options,
+                        frame_converter=self._video_frame_converter,
                     )
                     if artifacts.audios[idx] is None
                     else encode_video_base64(
@@ -369,6 +381,7 @@ class OmniOpenAIServingVideo:
                         audio=artifacts.audios[idx],
                         audio_sample_rate=artifacts.audio_sample_rate,
                         video_codec_options=video_codec_options,
+                        frame_converter=self._video_frame_converter,
                     )
                 ),
                 action=artifacts.actions[idx],
@@ -425,6 +438,7 @@ class OmniOpenAIServingVideo:
             fps=artifacts.output_fps,
             **({"audio": audio, "audio_sample_rate": artifacts.audio_sample_rate} if audio is not None else {}),
             video_codec_options=video_codec_options,
+            frame_converter=self._video_frame_converter,
         )
         _t_encode_ms = (time.perf_counter() - _t_encode_start) * 1000
         logger.info("Video response encoding (MP4 bytes): %.2f ms", _t_encode_ms)

vllm_omni/entrypoints/openai/video_api_utils.py

@@ -10,9 +10,12 @@ import base64
 import binascii
 import os
 import tempfile
-from collections.abc import Iterator
+import threading
+from collections import deque
+from collections.abc import Generator
+from concurrent.futures import Future, ThreadPoolExecutor
 from io import BytesIO
-from typing import TYPE_CHECKING, Any, Literal, cast
+from typing import TYPE_CHECKING, Any, Literal
 
 import httpx
 import numpy as np
@@ -576,6 +579,112 @@ def _direct_planar_fallback_reason(
     return None
 
 
+class _PlanarFrameConverter:
+    """Convert response frames with one reusable, bounded executor."""
+
+    def __init__(self, max_workers: int) -> None:
+        self._max_workers = max_workers
+        if self._max_workers < 1:
+            raise ValueError("max_workers must be positive")
+        self._executor = (
+            ThreadPoolExecutor(max_workers=self._max_workers, thread_name_prefix="video-planar")
+            if self._max_workers > 1
+            else None
+        )
+        self._scratch_local = threading.local()
+        self._shutdown = False
+
+    @property
+    def max_workers(self) -> int:
+        return self._max_workers
+
+    def _channel_scratch(self, height: int, width: int, common_dtype: np.dtype) -> np.ndarray:
+        scratch_dtype = np.dtype(common_dtype)
+        if not np.issubdtype(scratch_dtype, np.floating):
+            scratch_dtype = np.dtype(np.float64)
+        scratch = getattr(self._scratch_local, "channel", None)
+        if scratch is None or scratch.shape != (height, width) or scratch.dtype != scratch_dtype:
+            scratch = np.empty((height, width), dtype=scratch_dtype)
+            self._scratch_local.channel = scratch
+        return scratch
+
+    def _build_frame(self, frame: np.ndarray, common_dtype: np.dtype) -> av.VideoFrame:
+        """Build one quantized GBR PyAV frame."""
+        import av
+
+        height, width = frame.shape[:2]
+        scratch = None if frame.dtype == np.uint8 else self._channel_scratch(height, width, common_dtype)
+        av_frame = av.VideoFrame(width, height, format="gbrp")
+        for plane, channel in zip(av_frame.planes, (1, 2, 0)):
+            if plane.height < height or plane.line_size < width:
+                raise ValueError("PyAV video plane is smaller than the requested frame dimensions.")
+            plane_view = np.frombuffer(
+                memoryview(plane),
+                dtype=np.uint8,
+                count=plane.height * plane.line_size,
+            ).reshape(plane.height, plane.line_size)
+            plane_view.fill(0)
+            if frame.dtype == np.uint8:
+                plane_view[:height, :width] = frame[..., channel]
+            else:
+                np.copyto(scratch, frame[..., channel], casting="unsafe")
+                np.clip(scratch, 0.0, 1.0, out=scratch)
+                scratch *= 255.0
+                np.rint(scratch, out=scratch)
+                plane_view[:height, :width] = scratch
+        return av_frame
+
+    def iter_frames(
+        self,
+        frames: list[np.ndarray],
+        common_dtype: np.dtype,
+    ) -> Generator[av.VideoFrame, None, None]:
+        """Yield converted frames in order with at most two batches in flight."""
+        effective_workers = min(len(frames), self._max_workers)
+        if effective_workers <= 1:
+            for frame in frames:
+                yield self._build_frame(frame, common_dtype)
+            return
+
+        executor = self._executor
+        assert executor is not None
+        pending: deque[Future[av.VideoFrame]] = deque()
+        next_index = 0
+        max_pending = 2 * effective_workers
+        try:
+            while next_index < len(frames) and len(pending) < max_pending:
+                pending.append(executor.submit(self._build_frame, frames[next_index], common_dtype))
+                next_index += 1
+
+            while pending:
+                converted_frame = pending.popleft().result()
+                if next_index < len(frames):
+                    pending.append(executor.submit(self._build_frame, frames[next_index], common_dtype))
+                    next_index += 1
+                yield converted_frame
+        finally:
+            for future in pending:
+                future.cancel()
+
+    def shutdown(self) -> None:
+        if not self._shutdown:
+            self._shutdown = True
+            if self._executor is not None:
+                self._executor.shutdown(wait=True, cancel_futures=True)
+
+
+def _iter_planar_video_frames(
+    frames: list[np.ndarray],
+    common_dtype: np.dtype,
+) -> Generator[av.VideoFrame, None, None]:
+    """Yield serial planar frames for direct utility callers."""
+    converter = _PlanarFrameConverter(max_workers=1)
+    try:
+        yield from converter.iter_frames(frames, common_dtype)
+    finally:
+        converter.shutdown()
+
+
 def _log_video_encoding_path(
     *,
     selected_path: str,
@@ -585,12 +694,13 @@ def _log_video_encoding_path(
     fps: int,
     audio: AudioInput | None,
     audio_sample_rate: int | None,
+    effective_frame_conversion_workers: int,
     reason: str | None = None,
 ) -> None:
     reason_field = "" if reason is None else f" reason={reason}"
     logger.info(
         "Video response encoding route selected: selected_path=%s%s frames=%d frame_shape=%s dtype=%s fps=%s "
-        "audio_present=%s effective_audio_sample_rate=%s",
+        "audio_present=%s effective_audio_sample_rate=%s effective_frame_conversion_workers=%s",
         selected_path,
         reason_field,
         len(frames),
@@ -599,6 +709,7 @@ def _log_video_encoding_path(
         fps,
         audio is not None,
         audio_sample_rate,
+        effective_frame_conversion_workers,
     )
 
 
@@ -611,40 +722,6 @@ def _resolve_audio_sample_rate(audio: AudioInput | None, audio_sample_rate: int
     return audio_sample_rate or DEFAULT_AUDIO_SAMPLE_RATE
 
 
-def _iter_planar_video_frames(
-    frames: list[np.ndarray],
-    common_dtype: np.dtype,
-) -> Iterator[av.VideoFrame]:
-    """Yield planar PyAV frames while retaining only one channel scratch buffer."""
-    import av
-
-    height, width = frames[0].shape[:2]
-    scratch_dtype = np.float64 if np.issubdtype(common_dtype, np.bool_) else common_dtype
-    scratch = None if common_dtype == np.uint8 else np.empty((height, width), dtype=scratch_dtype)
-
-    for frame in frames:
-        av_frame = av.VideoFrame(width, height, format="gbrp")
-        for plane, channel in zip(av_frame.planes, (1, 2, 0)):
-            if plane.height < height or plane.line_size < width:
-                raise ValueError("PyAV video plane is smaller than the requested frame dimensions.")
-            plane_view = np.frombuffer(
-                plane,
-                dtype=np.uint8,
-                count=plane.height * plane.line_size,
-            ).reshape(plane.height, plane.line_size)
-            plane_view.fill(0)
-            if frame.dtype == np.uint8:
-                plane_view[:height, :width] = frame[..., channel]
-            else:
-                scratch_buffer = cast(np.ndarray, scratch)
-                np.copyto(scratch_buffer, frame[..., channel], casting="unsafe")
-                np.clip(scratch_buffer, 0.0, 1.0, out=scratch_buffer)
-                scratch_buffer *= 255.0
-                np.rint(scratch_buffer, out=scratch_buffer)
-                plane_view[:height, :width] = scratch_buffer
-        yield av_frame
-
-
 def _encode_prepared_video_bytes_legacy(
     frames: list[np.ndarray],
     frame_shape: tuple[int, ...],
@@ -693,8 +770,9 @@ def _encode_video_bytes(
     audio: Any | None = None,
     audio_sample_rate: int | None = None,
     video_codec_options: dict[str, str] | None = None,
+    frame_converter: _PlanarFrameConverter | None = None,
 ) -> bytes:
-    """Encode a video payload through the automatic capability dispatcher."""
+    """Encode a video payload through the direct planar or legacy path."""
     from vllm_omni.diffusion.utils.media_utils import mux_av_video_audio_bytes
 
     # Prepare once so validation is shared by both paths and malformed common
@@ -712,6 +790,7 @@ def _encode_video_bytes(
             fps=fps,
             audio=audio,
             audio_sample_rate=effective_audio_sample_rate,
+            effective_frame_conversion_workers=0,
         )
         return _encode_prepared_video_bytes_legacy(
             frames,
@@ -723,6 +802,8 @@ def _encode_video_bytes(
             video_codec_options=video_codec_options,
         )
 
+    owns_frame_converter = frame_converter is None
+    converter = frame_converter or _PlanarFrameConverter(max_workers=1)
     _log_video_encoding_path(
         selected_path="direct_planar",
         frames=frames,
@@ -731,17 +812,24 @@ def _encode_video_bytes(
         fps=fps,
         audio=audio,
         audio_sample_rate=effective_audio_sample_rate,
+        effective_frame_conversion_workers=min(len(frames), converter.max_workers),
     )
     audio_np = _coerce_audio_to_numpy(audio) if audio is not None else None
-    return mux_av_video_audio_bytes(
-        _iter_planar_video_frames(frames, common_dtype),
-        width=frame_shape[1],
-        height=frame_shape[0],
-        audio_waveform=audio_np,
-        fps=float(fps),
-        audio_sample_rate=effective_audio_sample_rate,
-        video_codec_options=video_codec_options,
-    )
+    video_frames = converter.iter_frames(frames, common_dtype)
+    try:
+        return mux_av_video_audio_bytes(
+            video_frames,
+            width=frame_shape[1],
+            height=frame_shape[0],
+            audio_waveform=audio_np,
+            fps=float(fps),
+            audio_sample_rate=effective_audio_sample_rate,
+            video_codec_options=video_codec_options,
+        )
+    finally:
+        video_frames.close()
+        if owns_frame_converter:
+            converter.shutdown()
 
 
 class FragmentedMP4VideoEncoder:
@@ -799,6 +887,7 @@ def encode_video_base64(
     audio: Any | None = None,
     audio_sample_rate: int | None = None,
     video_codec_options: dict[str, str] | None = None,
+    frame_converter: _PlanarFrameConverter | None = None,
 ) -> str:
     """Encode a video (frames/array/tensor) to base64 MP4."""
     video_bytes = _encode_video_bytes(
@@ -807,5 +896,6 @@ def encode_video_base64(
         audio=audio,
         audio_sample_rate=audio_sample_rate,
         video_codec_options=video_codec_options,
+        frame_converter=frame_converter,
     )
     return base64.b64encode(video_bytes).decode("utf-8")