[Bugfix][MiniCPM-o] Size the duplex HD slice reservation from the frame (#7654)

Signed-off-by: Tianyao Wu <rayroy31@gmail.com>
Co-authored-by: Claude <noreply@anthropic.com>

tests/model_executor/models/minicpmo_4_5/duplex/test_scheduler_token_budget.py

@@ -0,0 +1,245 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""What a duplex append reserves in the scheduler, and why that number.
+
+The reservation is not advisory: ``build_duplex_data_plane_prompt`` turns it
+into ``[scheduler_token_id] * budget``, and a unit whose embeddings outnumber
+its slots has its tail dropped with a warning rather than failing
+(``MiniCPMO45OmniModel.get_input_embeddings``). Over-reserving only wastes
+slots, so every uncertainty here has to resolve upwards.
+
+The camera side depends on HD slicing, which depends on the frame size
+*relative to the checkpoint's normalization tile*. The tile is configuration,
+not a constant, so the cases below drive the arithmetic at more than one tile
+size and check it against ``MiniCPMVImageProcessor.get_sliced_grid`` -- the
+same grid search the checkpoint's own processor runs.
+"""
+
+from __future__ import annotations
+
+import base64
+import io
+from dataclasses import dataclass
+
+import pytest
+from PIL import Image
+
+from vllm_omni.engine.duplex.config import DuplexSessionConfig
+from vllm_omni.model_executor.models.minicpmo_4_5.duplex.plugin import (
+    PRIVATE_RUNTIME_CONFIG_KEYS,
+    _apply_default_scheduler_policy,
+    _duplex_vision_tile_pixels,
+    _duplex_vision_tokens,
+    _model_vision_tile_pixels,
+    duplex_scheduler_token_budget,
+)
+from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import MiniCPMVImageProcessor
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+TOKENS_PER_BLOCK = 66
+#: ``_official_max_slice_nums`` gives the first frame of a stacked pair 2 and
+#: every other frame 1.
+BASE_FRAME_MAX_SLICES = 2
+#: What the released MiniCPM-o 4.5 checkpoint configures.
+TILE_448 = 448 * 448
+
+SIZES = [(224, 224), (320, 180), (448, 448), (448, 449), (640, 480), (960, 540), (1280, 720), (540, 960)]
+
+
+def _frame(size: tuple[int, int], fmt: str = "JPEG") -> str:
+    buffer = io.BytesIO()
+    Image.new("RGB", size).save(buffer, format=fmt)
+    return base64.b64encode(buffer.getvalue()).decode()
+
+
+def _blocks_the_model_makes(size: tuple[int, int], max_slice_nums: int, scale_resolution: int) -> int:
+    """Source tile plus the HD patches, straight from the model's grid search."""
+    processor = MiniCPMVImageProcessor(scale_resolution=scale_resolution)
+    grid = processor.get_sliced_grid(image_size=size, max_slice_nums=max_slice_nums)
+    return 1 + grid[0] * grid[1] if grid else 1
+
+
+@pytest.mark.parametrize("scale_resolution", [336, 448, 560])
+@pytest.mark.parametrize("size", SIZES, ids=[f"{w}x{h}" for w, h in SIZES])
+@pytest.mark.parametrize("fmt", ["JPEG", "PNG"])
+def test_a_stacked_pair_reserves_what_the_model_slices_it_into(
+    size: tuple[int, int], fmt: str, scale_resolution: int
+) -> None:
+    """The tile comes from the checkpoint, so the arithmetic has to follow it."""
+    frame = _frame(size, fmt)
+    expected_blocks = _blocks_the_model_makes(size, BASE_FRAME_MAX_SLICES, scale_resolution) + _blocks_the_model_makes(
+        size, 1, scale_resolution
+    )
+
+    reserved = _duplex_vision_tokens({"video_frames": [frame, frame]}, tile_pixels=scale_resolution**2)
+
+    assert reserved == expected_blocks * TOKENS_PER_BLOCK
+
+
+def test_a_small_frame_no_longer_reserves_three_hd_blocks() -> None:
+    """A 448x448 frame is one tile, so a stacked pair is two blocks, not four.
+
+    The size-independent count reserved 264 tokens for this pair. 132 of them
+    were for patches the model never produces.
+    """
+    frame = _frame((448, 448))
+
+    assert _duplex_vision_tokens({"video_frames": [frame, frame]}, tile_pixels=TILE_448) == 2 * TOKENS_PER_BLOCK
+
+
+def test_one_tile_is_the_whole_slicing_decision() -> None:
+    """``ceil(w * h / tile)`` capped at 2: a single pixel over the tile slices."""
+
+    def pair(size: tuple[int, int]) -> int:
+        frame = _frame(size)
+        return _duplex_vision_tokens({"video_frames": [frame, frame]}, tile_pixels=TILE_448)
+
+    assert pair((448, 448)) == 2 * TOKENS_PER_BLOCK
+    assert pair((448, 449)) == 4 * TOKENS_PER_BLOCK
+
+
+def test_only_the_first_frame_of_a_pair_is_hd_sliced() -> None:
+    """``_official_max_slice_nums`` is ``[2, 1]``: the composite is never sliced.
+
+    Sizing the wrong half of the pair would pass every same-size case, so the
+    two orders have to disagree.
+    """
+    small, large = _frame((448, 448)), _frame((1280, 720))
+
+    assert _duplex_vision_tokens({"video_frames": [small, large]}, tile_pixels=TILE_448) == 2 * TOKENS_PER_BLOCK
+    assert _duplex_vision_tokens({"video_frames": [large, small]}, tile_pixels=TILE_448) == 4 * TOKENS_PER_BLOCK
+
+
+def test_a_lone_frame_is_never_hd_sliced() -> None:
+    """``_official_max_slice_nums(1)`` is ``[1]``: one frame is one tile at any size."""
+    assert _duplex_vision_tokens({"video_frames": [_frame((1280, 720))]}, tile_pixels=TILE_448) == TOKENS_PER_BLOCK
+
+
+def test_extra_frames_past_the_pair_each_cost_one_block() -> None:
+    """The wire caps a frame list at two; the in-process API does not."""
+    small, large = _frame((448, 448)), _frame((1280, 720))
+
+    assert _duplex_vision_tokens({"video_frames": [small] * 3}, tile_pixels=TILE_448) == 3 * TOKENS_PER_BLOCK
+    assert _duplex_vision_tokens({"video_frames": [large] * 3}, tile_pixels=TILE_448) == 5 * TOKENS_PER_BLOCK
+
+
+def test_an_unknown_tile_size_keeps_the_sliced_reservation() -> None:
+    """No tile, no shrinking. A checkpoint whose config cannot be read keeps the old number."""
+    frame = _frame((224, 224))
+
+    assert _duplex_vision_tokens({"video_frames": [frame, frame]}) == 4 * TOKENS_PER_BLOCK
+    assert _duplex_vision_tokens({"video_frames": [frame, frame]}, tile_pixels=None) == 4 * TOKENS_PER_BLOCK
+
+
+@pytest.mark.parametrize(
+    "first_frame",
+    [
+        pytest.param("data:image/jpeg;base64," + _frame((224, 224)), id="data url"),
+        pytest.param("!!! not base64 !!!", id="corrupt base64"),
+        pytest.param(base64.b64encode(b"\x00" * 64).decode(), id="base64 of something else"),
+        pytest.param(base64.b64encode(base64.b64decode(_frame((224, 224)))[:8]).decode(), id="truncated jpeg"),
+    ],
+)
+def test_a_frame_whose_header_will_not_parse_keeps_the_sliced_reservation(first_frame: str) -> None:
+    """The Realtime wire screens most of these; the in-process API does not.
+
+    ``validate_realtime_video_frames`` rejects bad base64 and anything whose
+    magic bytes are not JPEG or PNG, so on a websocket session only a truncated
+    but well-headed frame gets this far. ``DuplexOmni.append_audio`` submits
+    without that screen, so the reservation still has to hold on its own.
+    """
+    good = _frame((224, 224))
+
+    assert _duplex_vision_tokens({"video_frames": [first_frame, good]}, tile_pixels=TILE_448) == 4 * TOKENS_PER_BLOCK
+
+
+def test_no_camera_track_reserves_nothing() -> None:
+    for payload in ({}, {"video_frames": []}, {"video_frames": "not a list"}, {"video_frames": [None, ""]}, None):
+        assert _duplex_vision_tokens(payload, tile_pixels=TILE_448) == 0
+
+
+def test_the_audio_budget_is_untouched_by_the_camera_track() -> None:
+    """One second of pcm_f32le is 12 slots; the frames are added on top."""
+    one_second = base64.b64encode(b"\x00" * 4 * 16000).decode()
+    audio_only = {"audio": one_second, "format": "pcm_f32le"}
+    with_frames = {**audio_only, "video_frames": [_frame((448, 448)), _frame((448, 448))]}
+
+    assert duplex_scheduler_token_budget(audio_only, tile_pixels=TILE_448) == 12
+    assert duplex_scheduler_token_budget(with_frames, tile_pixels=TILE_448) == 12 + 2 * TOKENS_PER_BLOCK
+
+
+# ---- where the tile comes from ----
+
+
+@dataclass(frozen=True)
+class _SliceConfig:
+    scale_resolution: int
+
+
+@dataclass(frozen=True)
+class _HFConfig:
+    slice_config: object = None
+    image_size: object = None
+
+
+@dataclass(frozen=True)
+class _ModelConfig:
+    hf_config: object = None
+
+
+@pytest.mark.parametrize(
+    ("hf_config", "expected"),
+    [
+        (_HFConfig(slice_config={"max_slice_nums": 1, "scale_resolution": 448}, image_size=448), TILE_448),
+        (_HFConfig(slice_config=_SliceConfig(scale_resolution=336), image_size=448), 336 * 336),
+        (_HFConfig(image_size=560), 560 * 560),
+        (_HFConfig(), None),
+        (_HFConfig(slice_config={"scale_resolution": 0}, image_size=0), None),
+    ],
+    ids=["released-checkpoint", "attribute", "image-size-fallback", "nothing-to-read", "nonsense"],
+)
+def test_the_tile_is_read_from_the_checkpoint(hf_config: object, expected: int | None) -> None:
+    """First row is the released MiniCPM-o 4.5 ``config.json``, where ``slice_config`` is a dict."""
+    assert _model_vision_tile_pixels(_ModelConfig(hf_config=hf_config)) == expected
+
+
+def test_no_model_config_means_no_tile() -> None:
+    assert _model_vision_tile_pixels(None) is None
+    assert _model_vision_tile_pixels(_ModelConfig()) is None
+
+
+def test_the_tile_reaches_the_budget_through_the_runtime_config() -> None:
+    """``_apply_default_scheduler_policy`` writes it; ``build_duplex_data_plane_prompt`` reads it back."""
+    runtime_config: dict[str, object] = {}
+
+    _apply_default_scheduler_policy(
+        runtime_config,
+        config=DuplexSessionConfig(),
+        tokenizer=None,
+        model_config=_ModelConfig(hf_config=_HFConfig(slice_config={"scale_resolution": 448})),
+    )
+
+    assert runtime_config["duplex_vision_tile_pixels"] == TILE_448
+    assert _duplex_vision_tile_pixels(runtime_config) == TILE_448
+
+
+def test_a_checkpoint_that_cannot_be_read_leaves_the_key_out() -> None:
+    """No key, no tile, and no tile is the sliced reservation."""
+    runtime_config: dict[str, object] = {}
+
+    _apply_default_scheduler_policy(runtime_config, config=DuplexSessionConfig(), tokenizer=None)
+
+    assert "duplex_vision_tile_pixels" not in runtime_config
+    assert _duplex_vision_tile_pixels(runtime_config) is None
+
+
+def test_a_junk_tile_in_the_runtime_config_is_ignored() -> None:
+    for value in (0, -1, "448", 448.0, None):
+        assert _duplex_vision_tile_pixels({"duplex_vision_tile_pixels": value}) is None
+
+
+def test_the_tile_is_server_owned() -> None:
+    """A client that could set the tile could shrink the reservation under the worker."""
+    assert "duplex_vision_tile_pixels" in PRIVATE_RUNTIME_CONFIG_KEYS

vllm_omni/model_executor/models/minicpmo_4_5/duplex/plugin.py

@@ -63,6 +63,15 @@ _DUPLEX_VISION_TOKENS_PER_FRAME = 66
 # Official stacked pair uses max_slice_nums=[2, 1]: the current frame is HD
 # sliced (1 source + 2 patches on 960x540) and the composite is not.
 _DUPLEX_HD_SLICES_PER_BASE_FRAME = 3
+# How many blocks that slicing costs depends on the frame, and at
+# ``max_slice_nums=2`` the processor's grid search has only two outcomes.
+# ``MiniCPMVImageProcessor.get_sliced_grid`` takes
+# ``multiple = min(ceil(w * h / scale_resolution**2), max_slice_nums)`` and
+# returns no grid at all for ``multiple <= 1``; at 2 the only candidate split
+# is 2, so a frame is either unsliced or a 2-cell grid. One normalization tile
+# is therefore the whole decision -- but ``scale_resolution`` is the
+# checkpoint's, not a constant, so it is read from the model rather than
+# assumed.
 
 PRIVATE_RUNTIME_CONFIG_KEYS = frozenset(
     {
@@ -70,6 +79,7 @@ PRIVATE_RUNTIME_CONFIG_KEYS = frozenset(
         "duplex_stage_max_tokens",
         "duplex_stage0_max_tokens",
         "duplex_scheduler_token_id",
+        "duplex_vision_tile_pixels",
         "duplex_first_append_context_tokens",
         "ref_audio_data",
         "ref_audio_format",
@@ -86,29 +96,68 @@ class MiniCPMO45ClientRuntimeConfigError(DuplexRuntimeConfigError):
 # ---- engine policy helpers: scheduler token budget ----
 
 
-def _duplex_frame_count(payload: object) -> int:
+def _duplex_frames(payload: object) -> list[str]:
     if not isinstance(payload, dict):
-        return 0
+        return []
     frames = payload.get("video_frames")
     if not isinstance(frames, list):
-        return 0
-    return sum(1 for frame in frames if isinstance(frame, str) and frame)
+        return []
+    return [frame for frame in frames if isinstance(frame, str) and frame]
 
 
-def _duplex_vision_tokens(payload: object) -> int:
+def _duplex_base_frame_blocks(frame: str, tile_pixels: int | None) -> int:
+    """Blocks the HD-sliced frame of a stacked pair costs, read from the frame.
+
+    Wire format is the one Stage0 decodes: bare base64 JPEG/PNG, and only the
+    header is parsed -- the pixels are the worker's job.
+
+    Every uncertainty resolves to the sliced count, because the two directions
+    are not symmetric. Over-reserving wastes scheduler slots. Under-reserving
+    hands the worker fewer prompt slots than it has embeddings, and
+    ``MiniCPMO45OmniModel`` then drops the tail of the unit with a warning
+    rather than failing, which is silent audio loss.
+    """
+    if tile_pixels is None:
+        return _DUPLEX_HD_SLICES_PER_BASE_FRAME
+
+    from io import BytesIO
+
+    from PIL import Image
+
+    try:
+        raw = b64decode(frame, validate=True)
+        with Image.open(BytesIO(raw)) as image:
+            width, height = image.size
+    except (BinasciiError, ValueError, OSError, Image.DecompressionBombError):
+        return _DUPLEX_HD_SLICES_PER_BASE_FRAME
+    if width * height <= tile_pixels:
+        return 1
+    return _DUPLEX_HD_SLICES_PER_BASE_FRAME
+
+
+def _duplex_vision_tile_pixels(runtime_config: object) -> int | None:
+    """Area of the tile this model normalizes a frame to, or ``None`` if unknown."""
+    if not isinstance(runtime_config, dict):
+        return None
+    value = runtime_config.get("duplex_vision_tile_pixels")
+    return value if isinstance(value, int) and value > 0 else None
+
+
+def _duplex_vision_tokens(payload: object, *, tile_pixels: int | None = None) -> int:
     """Scheduler slots for this append's camera track.
 
     Audio is never stacked: a unit still carries one second of soundtrack.
     ``stack_frames`` only adds a second *image*. Official HD on that pair is
-    ``[2, 1]``, so the base frame reserves three 66-token blocks and every
-    extra frame reserves one.
+    ``[2, 1]``, so only the first frame is sliced and every other frame
+    reserves one 66-token block.
     """
-    count = _duplex_frame_count(payload)
-    if count <= 0:
+    frames = _duplex_frames(payload)
+    if not frames:
         return 0
-    if count >= 2:
-        return (_DUPLEX_HD_SLICES_PER_BASE_FRAME + (count - 1)) * _DUPLEX_VISION_TOKENS_PER_FRAME
-    return count * _DUPLEX_VISION_TOKENS_PER_FRAME
+    if len(frames) == 1:
+        return _DUPLEX_VISION_TOKENS_PER_FRAME
+    blocks = _duplex_base_frame_blocks(frames[0], tile_pixels) + (len(frames) - 1)
+    return blocks * _DUPLEX_VISION_TOKENS_PER_FRAME
 
 
 def _duplex_pcm_sample_count(payload: object) -> int | None:
@@ -136,8 +185,8 @@ def duplex_first_append_unit_count(payload: object) -> int | None:
     return max(1, sample_count // _DUPLEX_CHUNK_SAMPLES - 1)
 
 
-def duplex_scheduler_token_budget(payload: object, *, default: int = 64) -> int:
-    vision_tokens = _duplex_vision_tokens(payload)
+def duplex_scheduler_token_budget(payload: object, *, default: int = 64, tile_pixels: int | None = None) -> int:
+    vision_tokens = _duplex_vision_tokens(payload, tile_pixels=tile_pixels)
     sample_count = _duplex_pcm_sample_count(payload)
     if sample_count is None:
         return max(1, int(default)) + vision_tokens
@@ -185,13 +234,15 @@ def build_duplex_data_plane_prompt(
     payload: object,
     final: bool,
 ) -> dict[str, object]:
-    token_budget = duplex_scheduler_token_budget(payload)
+    tile_pixels = _duplex_vision_tile_pixels(runtime_config)
+    token_budget = duplex_scheduler_token_budget(payload, tile_pixels=tile_pixels)
     if seq <= 1:
         context_reserve = duplex_first_append_context_reserve(runtime_config)
         token_budget += context_reserve
         first_units = duplex_first_append_unit_count(payload)
         if first_units is not None:
-            token_budget = context_reserve + first_units * 12 - 1 + _duplex_vision_tokens(payload)
+            vision_tokens = _duplex_vision_tokens(payload, tile_pixels=tile_pixels)
+            token_budget = context_reserve + first_units * 12 - 1 + vision_tokens
     if seq > 1 and duplex_payload_is_exact_chunks(payload):
         token_budget += 1
     if final and duplex_payload_is_exact_chunks(payload):
@@ -466,11 +517,34 @@ def _apply_first_append_context_tokens(
     runtime_config["duplex_first_append_context_tokens"] = len(prefix_ids) + ref_tokens + len(suffix_ids)
 
 
+def _model_vision_tile_pixels(model_config: ModelConfig | None) -> int | None:
+    """Area of one normalization tile, from the checkpoint that will do the slicing.
+
+    ``MiniCPMVImageProcessor`` is built with ``scale_resolution=config.image_size``
+    and Stage0 loads the checkpoint's own processor, so this is per-checkpoint
+    configuration. ``None`` when it cannot be read, which keeps the reservation
+    at the sliced count.
+    """
+    hf_config = getattr(model_config, "hf_config", None)
+    if hf_config is None:
+        return None
+    slice_config = getattr(hf_config, "slice_config", None)
+    side = getattr(slice_config, "scale_resolution", None)
+    if not isinstance(side, int):
+        side = slice_config.get("scale_resolution") if isinstance(slice_config, dict) else None
+    if not isinstance(side, int):
+        side = getattr(hf_config, "image_size", None)
+    if not isinstance(side, int) or side <= 0:
+        return None
+    return side * side
+
+
 def _apply_default_scheduler_policy(
     runtime_config: dict[str, object],
     *,
     config: DuplexSessionConfig,
     tokenizer: PreTrainedTokenizerBase | None,
+    model_config: ModelConfig | None = None,
 ) -> None:
     stage0_max_tokens = config.max_tokens if isinstance(config.max_tokens, int) and config.max_tokens > 0 else 20
     runtime_config["duplex_stage_max_tokens"] = {"0": stage0_max_tokens, "1": 8192}
@@ -491,6 +565,9 @@ def _apply_default_scheduler_policy(
     scheduler_token_id = _scheduler_token_id(tokenizer)
     if scheduler_token_id is not None:
         runtime_config["duplex_scheduler_token_id"] = scheduler_token_id
+    tile_pixels = _model_vision_tile_pixels(model_config)
+    if tile_pixels is not None:
+        runtime_config["duplex_vision_tile_pixels"] = tile_pixels
 
 
 class MiniCPMO45DuplexPlugin(DuplexModelPlugin):
@@ -653,7 +730,7 @@ class MiniCPMO45DuplexPlugin(DuplexModelPlugin):
         if isinstance(initial_user_text, str) and initial_user_text:
             runtime_config["initial_user_text"] = initial_user_text
         tokenizer = await self._tokenizer_for(model_config)
-        _apply_default_scheduler_policy(runtime_config, config=config, tokenizer=tokenizer)
+        _apply_default_scheduler_policy(runtime_config, config=config, tokenizer=tokenizer, model_config=model_config)
 
         ref_audio = config.ref_audio
         extra_ref_audio = extra_body.get("ref_audio")