From 022d7a7695ffc3ce7329c38656e388cc11aecb8d Mon Sep 17 00:00:00 2001 From: R0CKSTAR Date: Fri, 14 Aug 2026 15:49:26 +0800 Subject: [PATCH] [Perf][Diffusion] Avoid redundant MiniMax-H3 reference video scans (#6064) Signed-off-by: Xiaodong Ye --- .../minimax_h3/test_minimax_h3_contract.py | 210 +++++++++++++++++- .../models/minimax_h3/reference_video.py | 133 ++++++++--- 2 files changed, 300 insertions(+), 43 deletions(-) diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py index 60ed7090f..f872e381f 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import json +import sys from multiprocessing.reduction import ForkingPickler from types import SimpleNamespace from unittest.mock import Mock @@ -1428,34 +1429,219 @@ def test_r7_r8_ref2va_video_segment_matrix(monkeypatch, tmp_path, case, start_ti assert transcode_calls[0][1]["target_frame_count"] == 209 -def test_ref2va_qwen_sampling_decodes_prepared_video_once(monkeypatch): +def test_ref2va_qwen_sampling_uses_one_selective_decode(monkeypatch): from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module - decoded = np.arange(25 * 2 * 2 * 3, dtype=np.uint8).reshape(25, 2, 2, 3) - load_calls = [] + decoded = np.arange(3 * 2 * 2 * 3, dtype=np.uint8).reshape(3, 2, 2, 3) + decode_calls = [] monkeypatch.setattr( reference_video_module, "_probe_video", - lambda _path: {"frame_count": len(decoded)}, + lambda _path: { + "frame_count": 25, + "width": 2, + "height": 2, + }, ) monkeypatch.setattr( reference_video_module, - "load_video_frames", - lambda path: load_calls.append(path) or decoded, + "_decode_video_frames_ffmpeg", + lambda path, **kwargs: decode_calls.append((path, kwargs)) or decoded, ) sampled = reference_video_module.sample_reference_video_frames("prepared.mp4") - assert load_calls == ["prepared.mp4"] - for frame, expected in zip( - sampled["frames"], - (decoded[0], decoded[12], decoded[24]), - strict=True, - ): + assert decode_calls == [ + ( + "prepared.mp4", + { + "frame_count": 25, + "indices": [0, 12, 24], + "width": 2, + "height": 2, + }, + ) + ] + for frame, expected in zip(sampled["frames"], decoded, strict=True): np.testing.assert_array_equal(frame, expected) assert sampled["block_timestamps"] == [0.25, 1.0] +@pytest.mark.parametrize( + ("indices", "frame_count"), + [([0, 12], 25), (None, 2)], +) +def test_ref2va_ffmpeg_decode_reads_exact_rgb_frames(monkeypatch, indices, frame_count): + from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module + + expected = np.arange(2 * 2 * 3 * 3, dtype=np.uint8).reshape(2, 2, 3, 3) + observed = [] + + def fake_run(command, **kwargs): + observed.append((command, kwargs)) + return SimpleNamespace(stdout=expected.tobytes()) + + monkeypatch.setattr(reference_video_module.subprocess, "run", fake_run) + + actual = reference_video_module._decode_video_frames_ffmpeg( + "prepared.mp4", + frame_count=frame_count, + indices=indices, + width=3, + height=2, + ) + + np.testing.assert_array_equal(actual, expected) + command, kwargs = observed[0] + assert command[0] == "ffmpeg" + assert command[command.index("-frames:v") + 1] == "2" + if indices is None: + assert "-vf" not in command + else: + assert command[command.index("-vf") + 1] == "select=eq(n\\,0)+eq(n\\,12)" + assert command[-1] == "pipe:1" + assert kwargs == {"check": True, "capture_output": True} + + +def test_ref2va_load_video_frames_uses_ffmpeg_without_decord(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module + + expected = np.zeros((2, 2, 3, 3), dtype=np.uint8) + decode_calls = [] + monkeypatch.setitem(sys.modules, "decord", None) + monkeypatch.setattr( + reference_video_module, + "_probe_video", + lambda _path: { + "frame_count": 2, + "width": 3, + "height": 2, + }, + ) + monkeypatch.setattr( + reference_video_module, + "_decode_video_frames_ffmpeg", + lambda path, **kwargs: decode_calls.append((path, kwargs)) or expected, + ) + + actual = reference_video_module.load_video_frames("prepared.mp4") + + assert actual is expected + assert decode_calls == [ + ( + "prepared.mp4", + { + "frame_count": 2, + "width": 3, + "height": 2, + }, + ) + ] + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"frame_count": 0, "width": 3, "height": 2}, "video has no frames"), + ({"frame_count": 1, "width": 0, "height": 2}, "invalid dimensions"), + ( + {"frame_count": 1, "indices": [], "width": 3, "height": 2}, + "invalid frame indices", + ), + ], +) +def test_ref2va_ffmpeg_decode_rejects_invalid_metadata(kwargs, message): + from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module + from vllm_omni.errors import OmniClientError + + with pytest.raises(OmniClientError, match=message): + reference_video_module._decode_video_frames_ffmpeg("prepared.mp4", **kwargs) + + +def test_ref2va_probe_uses_container_frame_count_without_scan(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module + + calls = [] + probe = { + "streams": [ + { + "codec_type": "video", + "width": 4, + "height": 2, + "r_frame_rate": "24/1", + "nb_frames": "3", + "duration": "0.125", + "sample_aspect_ratio": "1:1", + } + ], + "format": {"duration": "0.125", "size": "123"}, + } + monkeypatch.setattr( + reference_video_module.subprocess, + "run", + lambda command, **kwargs: calls.append((command, kwargs)) or SimpleNamespace(stdout=json.dumps(probe)), + ) + + metadata = reference_video_module._probe_video("prepared.mp4") + + assert metadata["frame_count"] == 3 + assert len(calls) == 1 + assert "-count_frames" not in calls[0][0] + + +def test_ref2va_probe_counts_frames_only_when_metadata_is_missing(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module + + calls = [] + stream = { + "codec_type": "video", + "width": 4, + "height": 2, + "r_frame_rate": "24/1", + "duration": "0.125", + "sample_aspect_ratio": "1:1", + } + first_probe = { + "streams": [stream], + "format": {"duration": "0.125", "size": "123"}, + } + second_probe = { + "streams": [{**stream, "nb_read_frames": "3"}], + "format": {"duration": "0.125", "size": "123"}, + } + monkeypatch.setattr( + reference_video_module.subprocess, + "run", + lambda command, **kwargs: calls.append((command, kwargs)) + or SimpleNamespace(stdout=json.dumps(first_probe if len(calls) == 1 else second_probe)), + ) + + metadata = reference_video_module._probe_video("prepared.mp4") + + assert metadata["frame_count"] == 3 + assert len(calls) == 2 + assert "-count_frames" in calls[1][0] + + +def test_ref2va_ffmpeg_decode_rejects_incomplete_output(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module + from vllm_omni.errors import OmniClientError + + monkeypatch.setattr( + reference_video_module.subprocess, + "run", + lambda *_args, **_kwargs: SimpleNamespace(stdout=b"short"), + ) + + with pytest.raises(OmniClientError, match="decoded 5 bytes.*expected 18"): + reference_video_module._decode_video_frames_ffmpeg( + "prepared.mp4", + frame_count=1, + width=3, + height=2, + ) + + def test_ref2va_two_video_recipe_tolerates_container_rounding(monkeypatch, tmp_path): from vllm_omni.diffusion.models.minimax_h3 import reference_video as reference_video_module diff --git a/vllm_omni/diffusion/models/minimax_h3/reference_video.py b/vllm_omni/diffusion/models/minimax_h3/reference_video.py index 0ecd69f31..18c8de5b4 100644 --- a/vllm_omni/diffusion/models/minimax_h3/reference_video.py +++ b/vllm_omni/diffusion/models/minimax_h3/reference_video.py @@ -44,31 +44,33 @@ def _nearest_multiple(value: float, multiple: int) -> int: def _probe_video(path: str) -> dict[str, Any]: - result = subprocess.run( - [ - "ffprobe", - "-v", - "error", - "-count_frames", - "-show_entries", - ( - "stream=codec_type,codec_name,width,height,r_frame_rate,nb_read_frames,nb_frames,duration," - "sample_aspect_ratio:format=format_name,size,duration" - ), - "-of", - "json", - path, - ], - check=True, - capture_output=True, - text=True, + show_entries = ( + "stream=codec_type,codec_name,width,height,r_frame_rate,nb_read_frames,nb_frames,duration," + "sample_aspect_ratio:format=format_name,size,duration" ) + command = ["ffprobe", "-v", "error", "-show_entries", show_entries, "-of", "json", path] + result = subprocess.run(command, check=True, capture_output=True, text=True) probe = json.loads(result.stdout) streams = probe.get("streams") or [] video_streams = [stream for stream in streams if stream.get("codec_type") == "video"] if not video_streams: raise OmniClientError(f"media has no video stream: {path}") stream = video_streams[0] + raw_count = stream.get("nb_read_frames") or stream.get("nb_frames") + if raw_count in (None, "", "N/A"): + result = subprocess.run( + ["ffprobe", "-v", "error", "-count_frames", *command[3:]], + check=True, + capture_output=True, + text=True, + ) + probe = json.loads(result.stdout) + streams = probe.get("streams") or [] + video_streams = [stream for stream in streams if stream.get("codec_type") == "video"] + if not video_streams: + raise OmniClientError(f"media has no video stream: {path}") + stream = video_streams[0] + raw_count = stream.get("nb_read_frames") or stream.get("nb_frames") try: numerator, denominator = str(stream["r_frame_rate"]).split("/", 1) fps = float(numerator) / float(denominator) @@ -76,7 +78,6 @@ def _probe_video(path: str) -> dict[str, Any]: raise OmniClientError(f"cannot determine video FPS: {path}") from exc if not math.isfinite(fps) or fps <= 0: raise OmniClientError(f"cannot determine video FPS: {path}") - raw_count = stream.get("nb_read_frames") or stream.get("nb_frames") if raw_count in (None, "", "N/A"): raise OmniClientError(f"cannot determine video frame count: {path}") raw_duration = stream.get("duration") @@ -404,13 +405,13 @@ def load_video_frames(path: str) -> np.ndarray: try: import decord except ImportError: - import av - - with av.open(path) as container: - frames = [frame.to_ndarray(format="rgb24") for frame in container.decode(video=0)] - if not frames: - raise OmniClientError(f"video has no frames: {path}") - return np.stack(frames) + meta = _probe_video(path) + return _decode_video_frames_ffmpeg( + path, + frame_count=int(meta["frame_count"]), + width=int(meta["width"]), + height=int(meta["height"]), + ) reader = decord.VideoReader(path) if len(reader) <= 0: @@ -419,6 +420,70 @@ def load_video_frames(path: str) -> np.ndarray: return frames.asnumpy() if hasattr(frames, "asnumpy") else np.asarray(frames) +def _decode_video_frames_ffmpeg( + path: str, + *, + frame_count: int, + width: int, + height: int, + indices: list[int] | None = None, +) -> np.ndarray: + frame_count = int(frame_count) + width = int(width) + height = int(height) + if frame_count <= 0: + raise OmniClientError(f"video has no frames: {path}") + if width <= 0 or height <= 0: + raise OmniClientError(f"video has invalid dimensions {width}x{height}: {path}") + if indices is not None and (not indices or any(index < 0 or index >= frame_count for index in indices)): + raise OmniClientError(f"invalid frame indices for {path}: {indices}") + + output_frame_count = frame_count if indices is None else len(indices) + command = [ + "ffmpeg", + "-loglevel", + "error", + "-threads", + "0", + "-i", + path, + "-map", + "0:v:0", + "-an", + ] + if indices is not None: + select = "+".join(f"eq(n\\,{index})" for index in indices) + command.extend(["-vf", f"select={select}"]) + command.extend( + [ + "-vsync", + "0", + "-frames:v", + str(output_frame_count), + "-f", + "rawvideo", + "-pix_fmt", + "rgb24", + "pipe:1", + ] + ) + result = subprocess.run( + command, + check=True, + capture_output=True, + ) + frame_size = width * height * 3 + expected_size = output_frame_count * frame_size + if len(result.stdout) != expected_size: + raise OmniClientError(f"decoded {len(result.stdout)} bytes from {path}, expected {expected_size}") + return np.frombuffer(result.stdout, dtype=np.uint8).reshape( + output_frame_count, + height, + width, + 3, + ) + + def sample_reference_video_frames(prepared_path: str) -> dict[str, Any]: meta = _probe_video(prepared_path) ratio = MINIMAX_H3_FPS / MINIMAX_H3_QWEN_VIDEO_SAMPLE_FPS @@ -434,11 +499,17 @@ def sample_reference_video_frames(prepared_path: str) -> dict[str, Any]: if not indices: raise OmniClientError(f"no frames sampled from {prepared_path}") - # The preparation step emits a lossless RGB stream. Decode it once and - # sample from that array so Qwen3VL sees exactly the same pixels as the - # video VAE, without starting one full-video ffmpeg decode per sample. - decoded_frames = load_video_frames(prepared_path) - frames = [np.asarray(decoded_frames[source_index]) for source_index in indices] + # The preparation step emits a lossless RGB stream. Decode only the sampled + # frames in one ffmpeg process so Qwen3VL sees the exact prepared pixels + # without decoding the entire high-bitrate stream into host memory. + decoded_frames = _decode_video_frames_ffmpeg( + prepared_path, + frame_count=int(meta["frame_count"]), + indices=indices, + width=meta["width"], + height=meta["height"], + ) + frames = [np.asarray(frame) for frame in decoded_frames] timestamps = [index / MINIMAX_H3_QWEN_VIDEO_SAMPLE_FPS for index in range(len(indices))] timestamps += [timestamps[-1]] * ((-len(timestamps)) % MINIMAX_H3_QWEN_TEMPORAL_PATCH)