anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Perf][Diffusion] Avoid redundant MiniMax-H3 reference video scans (#6064)

Browse Source 
Signed-off-by: Xiaodong Ye <xiaodong.ye@mthreads.com>
This commit is contained in:
R0CKSTAR authored and GitHub committed 2026-08-14 15:49:26 +08:00
1 parent b300f66656
commit 022d7a7695
2 files changed
+300 -43
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py
+198 -12
View File
@@ -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
vllm_omni/diffusion/models/minimax_h3/reference_video.py
+102 -31
View File
@@ -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)