[Bugfix][MiniCPM-o] Attach a unit's whole frame group and size the vision slot budget like Stage0 (#7271)

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

examples/online_serving/minicpmo/README.md

@@ -183,10 +183,20 @@ python examples/online_serving/minicpmo/realtime_duplex_demo.py \
     --output-dir /tmp/minicpmo_realtime_duplex_video_demo
 ```
 
-Detail inside a composite is capped by `scale_resolution=448` at
-`max_slice_nums=1`: official suggests HD slicing (`max_slice_nums=[2, 1]`) for
-stacked frames, which the duplex adapter does not implement yet. Reading small
-text or digits out of a wide scene is limited by that, not by frame timing.
+Stage 0 processes a stacked unit the way official HD mode does
+(`max_slice_nums=[2, 1]`): the base frame keeps its own 66-token block plus the
+patches `get_sliced_grid` cuts for its size at `scale_resolution=448`, and the
+composite adds one block. The wire still rejects a caller-supplied
+`max_slice_nums`, so slicing is Stage 0's decision, not the client's.
+
+That detail costs context. One 960x540 frame is 66 vision tokens per unit and a
+stacked pair is 264, on top of the unit's audio, so the Stage 0 prompt grows by
+277 tokens per second stacked against 79 with a single frame. Against the
+40960-token context this model ships with, a listen-heavy call reaches the limit
+after about 2.5 minutes stacked and about 8.5 minutes unstacked (measured on a
+264 s 960x540 clip). Duplex sessions cannot yet evict old units, so keep
+multi-minute calls on the default `--stack-frames 1`; a stacked session that
+outgrows the context ends with a context-length error.
 
 ## Open the experimental browser client
 

tests/engine/duplex/test_duplex_runtime.py

@@ -23,6 +23,7 @@ from vllm_omni.engine.duplex.runtime import (
 )
 from vllm_omni.model_executor.models.minicpmo_4_5.duplex.runtime import (
     MiniCPMO45DuplexRuntimeExtension,
+    _duplex_hd_slice_count,
     build_duplex_data_plane_prompt,
     duplex_scheduler_token_budget,
 )
@@ -481,3 +482,82 @@ def test_placeholder_budget_is_planned_inside_omni_engine_boundary():
     assert len(prompt["prompt_token_ids"]) == 16
     assert prompt["model_intermediate_buffer"]["duplex"]["fence"] == fence
     assert prompt["model_intermediate_buffer"]["duplex"]["scheduler_token_budget"] == 16
+
+
+def _jpeg_b64(width: int, height: int) -> str:
+    import base64
+    from io import BytesIO
+
+    from PIL import Image
+
+    buffer = BytesIO()
+    Image.new("RGB", (width, height), (32, 64, 96)).save(buffer, format="JPEG")
+    return base64.b64encode(buffer.getvalue()).decode("ascii")
+
+
+def _pcm_unit_payload(frames: list[str]) -> dict[str, object]:
+    import base64
+
+    return {
+        "audio": base64.b64encode(b"\x00" * (16000 * 4)).decode("ascii"),
+        "format": "pcm_f32le",
+        "sample_rate_hz": 16000,
+        "video_frames": frames,
+    }
+
+
+def test_duplex_scheduler_token_budget_one_frame_is_one_block():
+    audio_only = duplex_scheduler_token_budget(_pcm_unit_payload([]))
+    one_frame = duplex_scheduler_token_budget(_pcm_unit_payload([_jpeg_b64(960, 540)]))
+
+    assert one_frame - audio_only == 66
+
+
+@pytest.mark.parametrize(
+    ("base_size", "expected_blocks"),
+    [
+        ((960, 540), 4),  # official camera frame: source + 2 HD patches + composite
+        ((540, 960), 4),  # portrait phone capture
+        ((640, 360), 4),
+        ((448, 448), 2),  # fits one tile: the processor does not slice it
+        ((320, 240), 2),
+    ],
+)
+def test_duplex_scheduler_token_budget_matches_stage0_hd_slicing(base_size, expected_blocks):
+    """A stacked pair is processed with max_slice_nums=[2, 1]; the budget must
+    equal the blocks Stage0 builds for that base-frame size."""
+    audio_only = duplex_scheduler_token_budget(_pcm_unit_payload([]))
+    pair = duplex_scheduler_token_budget(_pcm_unit_payload([_jpeg_b64(*base_size), _jpeg_b64(448, 448)]))
+
+    assert pair - audio_only == expected_blocks * 66
+
+
+@pytest.mark.parametrize(
+    ("size", "expected_patches"),
+    [
+        # multiple = min(ceil(w * h / 448**2), max_slice_nums); <= 1 is never sliced.
+        ((448, 448), 0),  # exactly one tile
+        ((320, 240), 0),  # smaller than one tile
+        ((447, 449), 0),  # just under one tile
+        ((449, 449), 2),  # just over: the processor cuts a 2-tile grid
+        ((640, 360), 2),
+        ((960, 540), 2),  # official camera frame
+        ((540, 960), 2),  # same frame rotated
+        ((1920, 1080), 2),  # capped by max_slice_nums=2, not by area
+    ],
+)
+def test_duplex_hd_slice_count_matches_official_grid(size, expected_patches):
+    """Lock the port of ``MiniCPMVImageProcessor.get_sliced_grid``.
+
+    Stage0 slices a stacked unit with ``max_slice_nums=[2, 1]``, so the patch
+    count for the base frame decides the scheduler budget. A frame that fits one
+    ``scale_resolution=448`` tile is not sliced at all.
+    """
+    assert _duplex_hd_slice_count(size, 2) == expected_patches
+
+
+def test_duplex_scheduler_token_budget_unreadable_base_frame_keeps_hd_fallback():
+    audio_only = duplex_scheduler_token_budget(_pcm_unit_payload([]))
+    pair = duplex_scheduler_token_budget(_pcm_unit_payload(["not-an-image", _jpeg_b64(448, 448)]))
+
+    assert pair - audio_only == 4 * 66

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

@@ -177,3 +177,108 @@ def test_pcm_commit_reservation_rollback_restores_residual_audio():
     )
     assert retried.payload is not None
     assert base64.b64decode(retried.payload["audio"]) == (base64.b64decode(original["audio"]) + b"\x00" * (8_000 * 4))
+
+
+def _frame_payload(samples: int, frames: list[str]) -> dict[str, object]:
+    payload = pcm_payload(samples)
+    payload["video_frames"] = frames
+    return payload
+
+
+def test_append_attaches_every_frame_of_the_unit_closing_append():
+    """A base frame and its stacked composite arrive on one append and must
+    enter the same model unit (official ``frame_list``), not one per unit."""
+    buffer = MiniCPMO45PcmAppendBuffer()
+
+    emitted = buffer.append(_frame_payload(16_000, ["base-0", "stack-0"]), chunk_period_ms=1_000)
+
+    assert emitted is not None
+    assert emitted["video_frames"] == ["base-0", "stack-0"]
+    assert buffer._frame_queue == []
+
+
+def test_stacked_frames_do_not_accumulate_across_units():
+    buffer = MiniCPMO45PcmAppendBuffer()
+    attached: list[list[str]] = []
+
+    for unit in range(64):
+        emitted = buffer.append(
+            _frame_payload(16_000, [f"base-{unit}", f"stack-{unit}"]),
+            chunk_period_ms=1_000,
+        )
+        assert emitted is not None
+        attached.append(list(emitted["video_frames"]))
+
+    assert attached[0] == ["base-0", "stack-0"]
+    assert attached[-1] == ["base-63", "stack-63"]
+    assert buffer._frame_queue == []
+
+
+def test_frames_from_partial_appends_ride_the_unit_they_close():
+    """200 ms client chunks: the frames arrive with the 5th chunk of each
+    second and attach to the unit that chunk completes, one group per unit."""
+    buffer = MiniCPMO45PcmAppendBuffer()
+
+    emitted_units: list[dict[str, object]] = []
+    for unit in range(3):
+        for chunk in range(5):
+            frames = [f"base-{unit}", f"stack-{unit}"] if chunk == 4 else []
+            emitted = buffer.append(_frame_payload(3_200, frames), chunk_period_ms=1_000)
+            if emitted is not None:
+                emitted_units.append(emitted)
+
+    assert [unit["video_frames"] for unit in emitted_units] == [
+        ["base-0", "stack-0"],
+        ["base-1", "stack-1"],
+        ["base-2", "stack-2"],
+    ]
+    assert buffer._frame_queue == []
+
+
+def test_single_frame_per_unit_still_attaches_one_frame_per_unit():
+    buffer = MiniCPMO45PcmAppendBuffer()
+
+    first = buffer.append(_frame_payload(16_000, ["base-0"]), chunk_period_ms=1_000)
+    second = buffer.append(_frame_payload(16_000, ["base-1"]), chunk_period_ms=1_000)
+
+    assert first is not None and first["video_frames"] == ["base-0"]
+    assert second is not None and second["video_frames"] == ["base-1"]
+
+
+def test_frame_groups_stay_queued_when_audio_outruns_units():
+    """Two appends land before one unit closes: the later group waits for the
+    next unit instead of merging into the first."""
+    buffer = MiniCPMO45PcmAppendBuffer()
+
+    assert buffer.append(_frame_payload(8_000, ["base-0", "stack-0"]), chunk_period_ms=1_000) is None
+    first = buffer.append(_frame_payload(16_000, ["base-1", "stack-1"]), chunk_period_ms=1_000)
+    second = buffer.append(pcm_payload(8_000), chunk_period_ms=1_000)
+
+    assert first is not None and first["video_frames"] == ["base-0", "stack-0"]
+    assert second is not None and second["video_frames"] == ["base-1", "stack-1"]
+    assert buffer._frame_queue == []
+
+
+def test_rollback_restores_frame_groups_in_wire_order():
+    buffer = MiniCPMO45PcmAppendBuffer()
+
+    first = buffer.prepare_append(
+        _frame_payload(16_000, ["base-0", "stack-0"]),
+        operation_id="append-0",
+        chunk_period_ms=1_000,
+    )
+    second = buffer.prepare_append(
+        _frame_payload(16_000, ["base-1", "stack-1"]),
+        operation_id="append-1",
+        chunk_period_ms=1_000,
+    )
+    assert first is not None and second is not None
+    assert first.payload is not None and first.payload["video_frames"] == ["base-0", "stack-0"]
+
+    first.rollback()
+
+    assert buffer._frame_queue == [["base-0", "stack-0"], ["base-1", "stack-1"]]
+    retried = buffer.flush(chunk_period_ms=1_000)
+    assert retried is not None
+    assert retried["video_frames"] == ["base-0", "stack-0"]
+    assert buffer._frame_queue == [["base-1", "stack-1"]]

vllm_omni/entrypoints/duplex/realtime_input.py

@@ -1068,14 +1068,14 @@ class RealtimeInputTranslator(RealtimeStateOwner):
         Wire contract matches the official MiniCPM-o duplex loop: one base
         base64 JPEG per ~1 s audio chunk, optionally followed by that unit's
         stacked composite tiling the sub-frames captured inside it (at most 2
-        images either way). HD slicing (max_slice_nums > 1) is not implemented
-        by the duplex adapter yet and is rejected explicitly rather than
-        silently ignored; note official suggests it for composites
-        (``max_slice_nums=[2, 1]``), so stacked detail is capped at
-        scale_resolution until that lands.
+        images either way). A caller-supplied ``max_slice_nums`` is rejected
+        rather than silently ignored: slicing is Stage 0's decision here, and
+        Stage 0 already applies the official HD suggestion for a stacked unit
+        (``max_slice_nums=[2, 1]``). The wire simply does not let the client
+        choose it.
         """
         if max_slice_nums not in (None, 1):
-            return "max_slice_nums > 1 (HD slicing) is not implemented by the duplex Realtime adapter"
+            return "max_slice_nums is not selectable on the wire; Stage0 slices stacked units itself"
         if not isinstance(video_frames, list):
             return "video_frames must be a list of base64-encoded images"
         frames = [frame for frame in video_frames if frame is not None]

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

@@ -116,9 +116,12 @@ class MiniCPMO45PcmAppendBuffer:
         self._turn_had_speech = False
         self._reservation_seq = 0
         self._reservations: list[MiniCPMO45PcmAppendReservation] = []
-        # Omni duplex: queued camera frames (base64 JPEG), consumed FIFO at
-        # one frame per emitted model unit alongside the unit's audio.
-        self._frame_queue: list[str] = []
+        # Omni duplex: queued camera frames (base64 JPEG), grouped per client
+        # append and consumed FIFO at one group per emitted model unit. A
+        # group is the base frame plus its optional stacked composite (the
+        # wire contract allows at most 2 images per append); both must ride
+        # the same unit, mirroring official ``frame_list``.
+        self._frame_queue: list[list[str]] = []
 
     def clear(self) -> None:
         for reservation in self._reservations:
@@ -251,7 +254,9 @@ class MiniCPMO45PcmAppendBuffer:
         )
         frames_in = payload.get("video_frames")
         if isinstance(frames_in, list):
-            self._frame_queue.extend(frame for frame in frames_in if isinstance(frame, str) and frame)
+            frame_group = [frame for frame in frames_in if isinstance(frame, str) and frame]
+            if frame_group:
+                self._frame_queue.append(frame_group)
         self._turn_had_speech = self._turn_had_speech or bool(payload.get("is_speech", False))
         if not allow_emit:
             return None
@@ -284,16 +289,17 @@ class MiniCPMO45PcmAppendBuffer:
         out.pop("video_frames", None)
         out["audio"] = base64.b64encode(emit_raw).decode("ascii")
         out["sample_rate_hz"] = sample_rate_hz
-        # Omni duplex: attach at most one queued camera frame per emitted
-        # model unit (official cadence: one frame per 1 s chunk). The engine
-        # budgets 66 scheduler slots per attached frame from this payload.
-        # Official omni cadence is one frame per ~1 s chunk, and the first
-        # append consumes extra samples (1035 ms first window), so per-unit
-        # attachment could outrun the units Stage0 actually builds. Attach at
-        # most ONE frame per emitted payload; the rest stay queued.
+        # Omni duplex: attach at most one queued frame *group* per emitted
+        # model unit (official cadence: one ``frame_list`` per 1 s chunk). The
+        # engine budgets scheduler slots for every frame in this payload.
+        # The first append consumes extra samples (1035 ms first window), so
+        # per-unit attachment could outrun the units Stage0 actually builds.
+        # Attach ONE group per emitted payload; later groups stay queued. A
+        # group is what one client append carried (base frame + optional
+        # stacked composite), so the composite never lags its base frame.
         attached_frames: list[str] = []
         if emit_samples + pad_samples >= min_samples and self._frame_queue:
-            attached_frames = [self._frame_queue.pop(0)]
+            attached_frames = list(self._frame_queue.pop(0))
             out["video_frames"] = attached_frames
         out["force_listen"] = any(span.force_listen for span in reserved_spans)
         out["is_speech"] = any(span.is_speech for span in reserved_spans)
@@ -400,9 +406,9 @@ class MiniCPMO45PcmAppendBuffer:
         restored = b"".join(item._raw for item in rolled_back)
         self._buffer[:0] = restored
         self._prepend_spans([span for item in rolled_back for span in item._spans])
-        restored_frames = [frame for item in rolled_back for frame in item._video_frames]
-        if restored_frames:
-            self._frame_queue[:0] = restored_frames
+        restored_groups = [list(item._video_frames) for item in rolled_back if item._video_frames]
+        if restored_groups:
+            self._frame_queue[:0] = restored_groups
         self._sample_rate_hz = self._sample_rate_hz or reservation._sample_rate_hz
         self._turn_had_speech = self._turn_had_speech or any(item._turn_had_speech for item in rolled_back)
         for item in rolled_back:

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

@@ -26,10 +26,12 @@ class MiniCPMO45DuplexPolicy:
     SAMPLES_PER_AUDIO_TOKEN = 1600
     # Vision framing contract (omni duplex). Official streaming_prefill feeds
     # each frame as <image> + 64 resampler embeddings + </image> inside the
-    # unit, ahead of the unit's audio embeddings (max_slice_nums=1 in
-    # streaming, so exactly one 64-token block per frame). A unit may carry
-    # its base frame plus an optional stacked composite of the sub-frames
-    # captured inside that unit (2 blocks).
+    # unit, ahead of the unit's audio embeddings. A lone frame is one block.
+    # A unit carrying its base frame plus the stacked composite of the
+    # sub-frames captured inside it is processed with the official HD
+    # suggestion (max_slice_nums=[2, 1]): the base frame adds the patches the
+    # processor cuts for its size (two for a 960x540 capture, none for a frame
+    # that fits one 448x448 tile) and the composite adds one block.
     VISION_EMBEDS_PER_FRAME = 64
     VISION_TOKENS_PER_FRAME = VISION_EMBEDS_PER_FRAME + 2  # <image> + embeds + </image>
     DEFAULT_MAX_NEW_SPEAK_TOKENS_PER_CHUNK = 20

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

@@ -3,6 +3,7 @@
 
 from __future__ import annotations
 
+import math
 from base64 import b64decode
 from binascii import Error as BinasciiError
 from typing import Any, cast
@@ -23,17 +24,73 @@ _DUPLEX_SAMPLES_PER_AUDIO_TOKEN = 1600
 # matching MiniCPMO45DuplexPolicy.VISION_TOKENS_PER_FRAME.
 _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.
+# sliced (1 source + up to 2 patches, e.g. on 960x540) and the composite is
+# not. The patch count depends on the frame size (see
+# ``_duplex_hd_slice_count``); the constant is the fallback when the frame
+# header cannot be read.
 _DUPLEX_HD_SLICES_PER_BASE_FRAME = 3
+_DUPLEX_HD_MAX_SLICE_NUMS = 2
+_DUPLEX_SCALE_RESOLUTION = 448
+
+
+def _duplex_frames(payload: object) -> list[str]:
+    if not isinstance(payload, dict):
+        return []
+    frames = payload.get("video_frames")
+    if not isinstance(frames, list):
+        return []
+    return [frame for frame in frames if isinstance(frame, str) and frame]
 
 
 def _duplex_frame_count(payload: object) -> int:
-    if not isinstance(payload, dict):
+    return len(_duplex_frames(payload))
+
+
+def _duplex_frame_size(frame_b64: str) -> tuple[int, int] | None:
+    """Pixel size of a base64 JPEG/PNG frame from its header, or ``None``."""
+    from io import BytesIO
+
+    try:
+        from PIL import Image
+
+        with Image.open(BytesIO(b64decode(frame_b64, validate=True))) as image:
+            width, height = image.size
+    except Exception:  # noqa: BLE001 - Stage0 rejects bad frames with a reason
+        return None
+    if width <= 0 or height <= 0:
+        return None
+    return int(width), int(height)
+
+
+def _duplex_hd_slice_count(image_size: tuple[int, int], max_slice_nums: int) -> int:
+    """Number of HD patches ``MiniCPMVImageProcessor.get_sliced_grid`` adds.
+
+    Port of the official grid selection (``scale_resolution=448``): an image
+    whose area fits in one 448x448 tile is not sliced at all, otherwise the
+    grid closest to the image aspect ratio among ``multiple-1 .. multiple+1``
+    splits is used. Stage0 runs the same processor, so the scheduler budget
+    and the worker-built embeddings agree.
+    """
+    width, height = image_size
+    ratio = width * height / (_DUPLEX_SCALE_RESOLUTION * _DUPLEX_SCALE_RESOLUTION)
+    multiple = min(math.ceil(ratio), max_slice_nums)
+    if multiple <= 1:
         return 0
-    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)
+    log_ratio = math.log(width / height)
+    best_grid = (1, 1)
+    min_error = float("inf")
+    for split_grids_nums in (multiple - 1, multiple, multiple + 1):
+        if split_grids_nums == 1 or split_grids_nums > max_slice_nums:
+            continue
+        for m in range(1, split_grids_nums + 1):
+            if split_grids_nums % m:
+                continue
+            grid = (m, split_grids_nums // m)
+            error = abs(log_ratio - math.log(grid[0] / grid[1]))
+            if error < min_error:
+                best_grid = grid
+                min_error = error
+    return best_grid[0] * best_grid[1]
 
 
 def _duplex_vision_tokens(payload: object) -> int:
@@ -41,15 +98,25 @@ def _duplex_vision_tokens(payload: object) -> int:
 
     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.
+    ``max_slice_nums=[2, 1]``: the base frame keeps its source block plus the
+    HD patches the processor cuts for its size (two for a 960x540 camera
+    frame, none for anything that fits in one 448x448 tile) and every extra
+    frame is one block. A single frame is never sliced. The count must match
+    the embeddings Stage0 builds exactly: surplus slots become pad
+    embeddings inside the KV.
     """
-    count = _duplex_frame_count(payload)
+    frames = _duplex_frames(payload)
+    count = len(frames)
     if count <= 0:
         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 count == 1:
+        return _DUPLEX_VISION_TOKENS_PER_FRAME
+    base_size = _duplex_frame_size(frames[0])
+    if base_size is None:
+        base_blocks = _DUPLEX_HD_SLICES_PER_BASE_FRAME
+    else:
+        base_blocks = 1 + _duplex_hd_slice_count(base_size, _DUPLEX_HD_MAX_SLICE_NUMS)
+    return (base_blocks + (count - 1)) * _DUPLEX_VISION_TOKENS_PER_FRAME
 
 
 def _duplex_pcm_sample_count(payload: object) -> int | None: