[Bugfix][MiniCPM-o] Take the duplex vision tile from the image-processor config (#8462)

Signed-off-by: Tianyao Wu <rayroy31@gmail.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>

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

@@ -20,18 +20,24 @@ from __future__ import annotations
 
 import base64
 import io
+import json
 from dataclasses import dataclass
+from pathlib import Path
+from types import SimpleNamespace
 
 import pytest
 from PIL import Image
 
 from vllm_omni.engine.duplex.config import DuplexSessionConfig
+from vllm_omni.engine.duplex.contracts import DuplexFence
+from vllm_omni.model_executor.models.minicpmo_4_5.duplex import plugin
 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,
+    build_duplex_data_plane_prompt,
     duplex_scheduler_token_budget,
 )
 from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import MiniCPMVImageProcessor
@@ -156,7 +162,14 @@ def test_a_frame_whose_header_will_not_parse_keeps_the_sliced_reservation(first_
 
 
 def test_no_camera_track_reserves_nothing() -> None:
-    for payload in ({}, {"video_frames": []}, {"video_frames": "not a list"}, {"video_frames": [None, ""]}, None):
+    payloads: tuple[object, ...] = (
+        {},
+        {"video_frames": []},
+        {"video_frames": "not a list"},
+        {"video_frames": [None, ""]},
+        None,
+    )
+    for payload in payloads:
         assert _duplex_vision_tokens(payload, tile_pixels=TILE_448) == 0
 
 
@@ -170,39 +183,130 @@ def test_the_audio_budget_is_untouched_by_the_camera_track() -> None:
     assert duplex_scheduler_token_budget(with_frames, tile_pixels=TILE_448) == 12 + 2 * TOKENS_PER_BLOCK
 
 
+# ---- the first append decodes its frame once ----
+
+
+def test_a_first_append_decodes_the_base_frame_once(monkeypatch: pytest.MonkeyPatch) -> None:
+    """The first-units branch reuses the vision count instead of decoding again."""
+    calls: list[str] = []
+    real = plugin._duplex_base_frame_blocks
+
+    def counting(frame: str, tile_pixels: int | None) -> int:
+        calls.append(frame)
+        return real(frame, tile_pixels)
+
+    monkeypatch.setattr(plugin, "_duplex_base_frame_blocks", counting)
+    two_seconds = base64.b64encode(b"\x00" * 4 * 32000).decode()
+    frame = _frame((960, 540))
+
+    prompt = build_duplex_data_plane_prompt(
+        request_id="tile-request",
+        fence=DuplexFence("sid", turn_id=1),
+        session_config={},
+        runtime_config={"duplex_first_append_context_tokens": 5, "duplex_vision_tile_pixels": TILE_448},
+        seq=1,
+        turn_seq=1,
+        payload={"audio": two_seconds, "format": "pcm_f32le", "video_frames": [frame, frame]},
+        final=False,
+    )
+
+    assert len(calls) == 1
+    # 5 context + one unit of 12 - 1, plus the sliced base frame (3 blocks) and the composite (1 block).
+    assert len(prompt["prompt_token_ids"]) == 5 + 12 - 1 + 4 * 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:
+    model: object = None
     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"],
+_ABSENT = object()
+
+
+def _checkpoint(tmp_path: Path, processor: object = _ABSENT, config_side: int | None = None) -> _ModelConfig:
+    """A local checkpoint: ``processor`` is ``preprocessor_config.json``'s ``scale_resolution``.
+
+    ``config_side`` is the tile ``config.json`` carries, in the released
+    checkpoint's layout (``slice_config`` dict plus ``image_size``). It goes on
+    ``hf_config`` as vLLM's ``ModelConfig`` would load it, so a test can make
+    the two sources disagree.
+    """
+    if processor is not _ABSENT:
+        (tmp_path / "preprocessor_config.json").write_text(
+            json.dumps({"image_processor_type": "MiniCPMVImageProcessor", "scale_resolution": processor})
         )
-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
+    hf_config = None
+    if config_side is not None:
+        hf_config = SimpleNamespace(
+            slice_config={"max_slice_nums": 1, "scale_resolution": config_side}, image_size=config_side
+        )
+    return _ModelConfig(model=str(tmp_path), hf_config=hf_config)
+
+
+def test_the_tile_is_read_from_the_released_checkpoint_layout(tmp_path: Path) -> None:
+    assert _model_vision_tile_pixels(_checkpoint(tmp_path, processor=448, config_side=448)) == TILE_448
+
+
+def test_the_tile_follows_the_processor_not_config_json(tmp_path: Path) -> None:
+    """Stage0 slices with the processor, so ``config.json`` alone cannot move the tile."""
+    assert _model_vision_tile_pixels(_checkpoint(tmp_path, processor=448, config_side=560)) == TILE_448
+    assert _model_vision_tile_pixels(_checkpoint(tmp_path, processor=336, config_side=448)) == 336 * 336
+
+
+def test_a_nested_processor_config_wins_like_it_does_for_the_processor(tmp_path: Path) -> None:
+    """transformers 5 saves the image processor nested in ``processor_config.json`` and reads that first."""
+    model_config = _checkpoint(tmp_path, processor=448, config_side=448)
+    nested = tmp_path / "processor_config.json"
+
+    nested.write_text(json.dumps({"processor_class": "MiniCPMOProcessor"}))
+    assert _model_vision_tile_pixels(model_config) == TILE_448
+
+    nested.write_text(json.dumps({"image_processor": {"scale_resolution": 336}}))
+    assert _model_vision_tile_pixels(model_config) == 336 * 336
+
+    (tmp_path / "preprocessor_config.json").unlink()
+    assert _model_vision_tile_pixels(model_config) == 336 * 336
+
+
+def test_a_config_json_override_does_not_shrink_the_reservation(tmp_path: Path) -> None:
+    """A 500x500 base frame with a 560 tile in ``config.json`` and 448 in the processor.
+
+    Reading ``config.json`` reserved two blocks (132 tokens) for a frame the
+    processor slices into three plus the composite (264 tokens).
+    """
+    runtime_config: dict[str, object] = {}
+    _apply_default_scheduler_policy(
+        runtime_config,
+        config=DuplexSessionConfig(),
+        tokenizer=None,
+        model_config=_checkpoint(tmp_path, processor=448, config_side=560),
+    )
+    frame = _frame((500, 500))
+    expected_blocks = _blocks_the_model_makes((500, 500), BASE_FRAME_MAX_SLICES, 448) + _blocks_the_model_makes(
+        (500, 500), 1, 448
+    )
+
+    reserved = _duplex_vision_tokens(
+        {"video_frames": [frame, frame]}, tile_pixels=_duplex_vision_tile_pixels(runtime_config)
+    )
+
+    assert reserved == expected_blocks * TOKENS_PER_BLOCK == 264
+
+
+@pytest.mark.parametrize("value", [0, -1, "448", 448.0, True, None], ids=repr)
+def test_a_nonsense_processor_tile_is_no_tile(tmp_path: Path, value: object) -> None:
+    assert _model_vision_tile_pixels(_checkpoint(tmp_path, processor=value, config_side=448)) is None
+
+
+def test_without_a_processor_config_there_is_no_tile(tmp_path: Path) -> None:
+    """``config.json`` alone is not what Stage0 slices with, so it is not used."""
+    assert _model_vision_tile_pixels(_checkpoint(tmp_path, config_side=448)) is None
+    (tmp_path / "preprocessor_config.json").write_text("not json")
+    assert _model_vision_tile_pixels(_ModelConfig(model=str(tmp_path))) is None
 
 
 def test_no_model_config_means_no_tile() -> None:
@@ -210,7 +314,7 @@ def test_no_model_config_means_no_tile() -> None:
     assert _model_vision_tile_pixels(_ModelConfig()) is None
 
 
-def test_the_tile_reaches_the_budget_through_the_runtime_config() -> None:
+def test_the_tile_reaches_the_budget_through_the_runtime_config(tmp_path: Path) -> None:
     """``_apply_default_scheduler_policy`` writes it; ``build_duplex_data_plane_prompt`` reads it back."""
     runtime_config: dict[str, object] = {}
 
@@ -218,7 +322,7 @@ def test_the_tile_reaches_the_budget_through_the_runtime_config() -> None:
         runtime_config,
         config=DuplexSessionConfig(),
         tokenizer=None,
-        model_config=_ModelConfig(hf_config=_HFConfig(slice_config={"scale_resolution": 448})),
+        model_config=_checkpoint(tmp_path, processor=448),
     )
 
     assert runtime_config["duplex_vision_tile_pixels"] == TILE_448

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

@@ -186,16 +186,21 @@ 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, tile_pixels: int | None = None) -> int:
-    vision_tokens = _duplex_vision_tokens(payload, tile_pixels=tile_pixels)
+def _duplex_audio_token_budget(payload: object, *, default: int = 64) -> int:
     sample_count = _duplex_pcm_sample_count(payload)
     if sample_count is None:
-        return max(1, int(default)) + vision_tokens
+        return max(1, int(default))
     sample_count = max(1, sample_count)
     if sample_count % _DUPLEX_CHUNK_SAMPLES == 0:
         units = sample_count // _DUPLEX_CHUNK_SAMPLES
-        return units * (2 + _DUPLEX_CHUNK_SAMPLES // _DUPLEX_SAMPLES_PER_AUDIO_TOKEN) + vision_tokens
-    return max(16, min(768, sample_count // _DUPLEX_SAMPLES_PER_AUDIO_TOKEN + 8)) + vision_tokens
+        return units * (2 + _DUPLEX_CHUNK_SAMPLES // _DUPLEX_SAMPLES_PER_AUDIO_TOKEN)
+    return max(16, min(768, sample_count // _DUPLEX_SAMPLES_PER_AUDIO_TOKEN + 8))
+
+
+def duplex_scheduler_token_budget(payload: object, *, default: int = 64, tile_pixels: int | None = None) -> int:
+    return _duplex_audio_token_budget(payload, default=default) + _duplex_vision_tokens(
+        payload, tile_pixels=tile_pixels
+    )
 
 
 def duplex_first_append_context_reserve(runtime_config: object) -> int:
@@ -235,14 +240,14 @@ def build_duplex_data_plane_prompt(
     payload: object,
     final: bool,
 ) -> dict[str, object]:
-    tile_pixels = _duplex_vision_tile_pixels(runtime_config)
-    token_budget = duplex_scheduler_token_budget(payload, tile_pixels=tile_pixels)
+    # Computed once: with a known tile this decodes the stacked pair's first frame.
+    vision_tokens = _duplex_vision_tokens(payload, tile_pixels=_duplex_vision_tile_pixels(runtime_config))
+    token_budget = _duplex_audio_token_budget(payload) + vision_tokens
     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:
-            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
@@ -479,23 +484,31 @@ def _apply_first_append_context_tokens(
 
 
 def _model_vision_tile_pixels(model_config: ModelConfig | None) -> int | None:
-    """Area of one normalization tile, from the checkpoint that will do the slicing.
+    """Area of one normalization tile, as the processor that does the slicing sees it.
 
-    ``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.
+    Stage0 slices with the checkpoint's own processor
+    (``MiniCPMO45Stage0DuplexRuntime._load_processor_from_path``), whose
+    ``MiniCPMVImageProcessor`` takes ``scale_resolution`` from the checkpoint's
+    image-processor config: a nested ``image_processor`` entry in
+    ``processor_config.json`` if there is one, else ``preprocessor_config.json``.
+    ``get_image_processor_dict`` is that lookup. ``hf_config`` has its own
+    ``slice_config.scale_resolution`` and ``image_size``; they agree on the
+    released checkpoint, but reading them here would under-reserve as soon as
+    one side changes alone (``--hf-overrides`` reaches only ``hf_config``).
+    Local files only, like ``_load_tokenizer``. ``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:
+    model_path = getattr(model_config, "model", None)
+    if not isinstance(model_path, str) or not model_path:
         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:
+    try:
+        from transformers.image_processing_base import ImageProcessingMixin
+
+        processor_dict, _ = ImageProcessingMixin.get_image_processor_dict(model_path, local_files_only=True)
+        side = processor_dict.get("scale_resolution")
+    except Exception:
+        return None
+    if not isinstance(side, int) or isinstance(side, bool) or side <= 0:
         return None
     return side * side