[Bugfix][Model] Mistral3: fix image placeholder grid for processor size overrides (#52874)

Signed-off-by: Anmol Gupta <14880251+anmolgupt@users.noreply.github.com>

tests/models/multimodal/processing/test_mistral3.py

@@ -0,0 +1,143 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for Mistral3's multimodal preprocessing kwargs."""
+
+import pytest
+import torch
+from PIL import Image
+from transformers import BatchFeature
+
+from vllm.model_executor.models.lightonocr import LightOnOCRProcessingInfo
+from vllm.model_executor.models.mistral3 import Mistral3HFEncoderInfo
+from vllm.model_executor.models.pixtral import PixtralHFEncoderInfo
+from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.multimodal.inputs import MultiModalKwargsItems
+
+from ...utils import build_model_context
+
+# This repo ships both params.json (Mistral) and config.json (HF). Auto config
+# selects PixtralForConditionalGeneration; force HF to exercise Mistral3.
+_MODEL_CONFIG_KWARGS = {"config_format": "hf"}
+_MODEL_ID = "mistralai/Mistral-Small-3.1-24B-Instruct-2503"
+_LIGHTON_MODEL_ID = "lightonai/LightOnOCR-1B-1025"
+
+
+def _processed_pixel_values(
+    hf_processor,
+    images: list[Image.Image],
+    mm_processor_kwargs: dict[str, object],
+) -> list[torch.Tensor]:
+    """Resize via the HF image processor and un-pad to per-image H×W."""
+    image_processor = hf_processor.image_processor
+    hf_out = image_processor(images=images, return_tensors="pt", **mm_processor_kwargs)
+    pixel_values = hf_out["pixel_values"]
+    image_sizes = hf_out["image_sizes"]
+    return [p[:, :h, :w] for p, (h, w) in zip(pixel_values, image_sizes)]
+
+
+def _placeholder_count_from_prompt_updates(
+    processor,
+    images: list[Image.Image],
+    pixel_values: list[torch.Tensor],
+    mm_processor_kwargs: dict[str, object],
+) -> int:
+    hf_inputs = BatchFeature({"pixel_values": pixel_values})
+    fields_config = processor._get_mm_fields_config(hf_inputs, mm_processor_kwargs)
+    out_mm_kwargs = MultiModalKwargsItems.from_hf_inputs(hf_inputs, fields_config)
+    # Prompt updates use raw PIL sizes and must predict the processed grid.
+    mm_items = processor.info.parse_mm_data({"image": images})
+    updates = processor._get_prompt_updates(
+        mm_items, mm_processor_kwargs, out_mm_kwargs
+    )
+    image_token_id = processor.info.get_hf_config().image_token_index
+
+    total = 0
+    for item_idx in range(len(images)):
+        details = updates[0].resolve(item_idx).content
+        total += details.full.count(image_token_id)
+    return total
+
+
+def _expected_placeholder_tokens_per_image(
+    hf_processor,
+    pixel_values: torch.Tensor,
+) -> int:
+    """Count projected tokens from the actual HF-processed H×W."""
+    image_h, image_w = pixel_values.shape[-2:]
+    patch_size = hf_processor.image_processor.patch_size
+    if isinstance(patch_size, dict):
+        patch_h = patch_size["height"]
+        patch_w = patch_size["width"]
+    else:
+        patch_h = patch_w = int(patch_size)
+
+    spatial_merge_size = getattr(hf_processor, "spatial_merge_size", 1)
+    assert image_h % patch_h == 0
+    assert image_w % patch_w == 0
+
+    return (image_h // (patch_h * spatial_merge_size)) * (
+        image_w // (patch_w * spatial_merge_size)
+    )
+
+
+@pytest.mark.parametrize("model_id", [_MODEL_ID])
+@pytest.mark.parametrize(
+    ("mm_processor_kwargs", "image_size", "expected_toks_per_img"),
+    [
+        ({}, (448, 448), 256),
+        ({"size": {"longest_edge": 1008}}, (1540, 1540), 1296),
+        ({"size": {"longest_edge": 1288}}, (1536, 1187), 1656),
+        ({"size": {"longest_edge": 1008}}, (29, 29), 1),
+        ({"size": {"longest_edge": 1000}}, (1540, 1700), 1152),
+    ],
+)
+@pytest.mark.parametrize("num_imgs", [1, 2])
+@pytest.mark.parametrize("kwargs_on_init", [True, False])
+def test_processor_size_override(
+    model_id: str,
+    mm_processor_kwargs: dict[str, object],
+    image_size: tuple[int, int],
+    expected_toks_per_img: int,
+    num_imgs: int,
+    kwargs_on_init: bool,
+):
+    ctx = build_model_context(
+        model_id,
+        mm_processor_kwargs=mm_processor_kwargs if kwargs_on_init else None,
+        limit_mm_per_prompt={"image": num_imgs},
+        model_config_kwargs=_MODEL_CONFIG_KWARGS,
+    )
+    processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
+    hf_processor_mm_kwargs = {} if kwargs_on_init else mm_processor_kwargs
+    hf_processor = processor.info.get_hf_processor(**hf_processor_mm_kwargs)
+
+    dummy_image = Image.new("RGB", image_size, color=(127, 127, 127))
+    images = [dummy_image] * num_imgs
+    merged_mm_kwargs = processor.info.ctx.get_merged_mm_kwargs(hf_processor_mm_kwargs)
+    pixel_values = _processed_pixel_values(hf_processor, images, merged_mm_kwargs)
+
+    image_token_count = _placeholder_count_from_prompt_updates(
+        processor, images, pixel_values, hf_processor_mm_kwargs
+    )
+    expected_from_pixel_values = _expected_placeholder_tokens_per_image(
+        hf_processor, pixel_values[0]
+    )
+    assert expected_from_pixel_values == expected_toks_per_img
+    assert image_token_count == expected_from_pixel_values * num_imgs
+
+
+def test_lightonocr_keeps_vision_config_image_size():
+    ctx = build_model_context(
+        _LIGHTON_MODEL_ID,
+        mm_processor_kwargs={"size": {"longest_edge": 1008}},
+        model_config_kwargs=_MODEL_CONFIG_KWARGS,
+    )
+    processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
+
+    assert isinstance(processor.info, LightOnOCRProcessingInfo)
+    encoder_info = processor.info.get_vision_encoder_info()
+    assert isinstance(encoder_info, PixtralHFEncoderInfo)
+    assert not isinstance(encoder_info, Mistral3HFEncoderInfo)
+    assert encoder_info.get_image_size() == (
+        processor.info.get_hf_config().vision_config.image_size
+    )

vllm/model_executor/models/lightonocr.py

@@ -38,7 +38,18 @@ from vllm.multimodal.processing import (
 _I = TypeVar("_I", bound=Mistral3ProcessingInfo)
 
 
-class LightOnOCRMultiModalProcessor(BaseMultiModalProcessor[Mistral3ProcessingInfo]):
+class LightOnOCRProcessingInfo(Mistral3ProcessingInfo):
+    def get_vision_encoder_info(
+        self,
+        mm_processor_kwargs: Mapping[str, object] | None = None,
+    ) -> PixtralHFEncoderInfo:
+        # LightOnOCR prompt updates use vision_config.image_size directly;
+        # preserve that behavior rather than inheriting Mistral3's processor
+        # longest_edge override.
+        return PixtralHFEncoderInfo(self.get_hf_config())
+
+
+class LightOnOCRMultiModalProcessor(BaseMultiModalProcessor[LightOnOCRProcessingInfo]):
     def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str:
         return self.dummy_inputs.get_dummy_text(mm_counts)
 
@@ -124,7 +135,7 @@ class LightOnOCRMultiModalProcessor(BaseMultiModalProcessor[Mistral3ProcessingIn
 
 @MULTIMODAL_REGISTRY.register_processor(
     LightOnOCRMultiModalProcessor,
-    info=Mistral3ProcessingInfo,
+    info=LightOnOCRProcessingInfo,
     dummy_inputs=Mistral3DummyInputsBuilder,
 )
 class LightOnOCRForConditionalGeneration(Mistral3ForConditionalGeneration):

vllm/model_executor/models/mistral3.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import math
 from collections.abc import Iterable, Mapping, Sequence
 from typing import Annotated, Literal
 
@@ -22,7 +23,11 @@ from vllm.multimodal.inputs import (
     MultiModalFieldConfig,
     MultiModalKwargsItems,
 )
-from vllm.multimodal.parse import ImageProcessorItems, ImageSize, MultiModalDataItems
+from vllm.multimodal.parse import (
+    ImageProcessorItems,
+    ImageSize,
+    MultiModalDataItems,
+)
 from vllm.multimodal.processing import (
     BaseDummyInputsBuilder,
     BaseMultiModalProcessor,
@@ -50,7 +55,6 @@ from .utils import (
     init_vllm_registered_model,
     maybe_prefix,
 )
-from .vision import get_vision_encoder_info
 
 
 class Mistral3ImagePixelInputs(TensorSchema):
@@ -170,12 +174,57 @@ class Mistral3MultiModalProjector(nn.Module):
         return hidden_states
 
 
+class Mistral3HFEncoderInfo(PixtralHFEncoderInfo):
+    def __init__(self, hf_config: Mistral3Config, image_size: int) -> None:
+        super().__init__(hf_config)
+        self.image_size = image_size
+
+    def get_image_size(self) -> int:
+        return self.image_size
+
+    def get_patch_grid_size(
+        self,
+        *,
+        image_width: int,
+        image_height: int,
+    ) -> tuple[int, int]:
+        max_size = self.get_image_size()
+        ratio = max(image_width / max_size, image_height / max_size)
+
+        if ratio > 1:
+            image_width = math.floor(image_width / ratio)
+            image_height = math.floor(image_height / ratio)
+
+        patch_size = self.vision_config.patch_size
+        spatial_merge_size = self.hf_config.spatial_merge_size
+
+        # The HF processor rounds each dimension up to the vision patch size
+        # before the projector drops incomplete spatial-merge groups. This is
+        # not equivalent to rounding directly to the merged patch size.
+        num_width_patches = (image_width - 1) // patch_size + 1
+        num_height_patches = (image_height - 1) // patch_size + 1
+
+        ncols = num_width_patches // spatial_merge_size
+        nrows = num_height_patches // spatial_merge_size
+        return ncols, nrows
+
+
 class Mistral3ProcessingInfo(BaseProcessingInfo):
     def get_hf_config(self) -> Mistral3Config:
         return self.ctx.get_hf_config(Mistral3Config)
 
-    def get_vision_encoder_info(self):
-        return get_vision_encoder_info(self.get_hf_config())
+    def get_vision_encoder_info(
+        self,
+        mm_processor_kwargs: Mapping[str, object] | None = None,
+    ) -> Mistral3HFEncoderInfo:
+        processor = self.get_hf_processor()
+        size = processor.image_processor.size
+        merged_kwargs = self.ctx.get_merged_mm_kwargs(mm_processor_kwargs or {})
+        if override_size := merged_kwargs.get("size"):
+            size = size | override_size
+
+        image_size = size["longest_edge"]
+        return Mistral3HFEncoderInfo(self.get_hf_config(), image_size)
 
     def get_hf_processor(self, **kwargs: object):
         return self.ctx.get_hf_processor(PixtralProcessor, **kwargs)
@@ -284,7 +333,7 @@ class Mistral3MultiModalProcessor(BaseMultiModalProcessor[Mistral3ProcessingInfo
         image_end_id = vocab[processor.image_end_token]
 
         assert isinstance(hf_config.vision_config, PixtralVisionConfig)
-        encoder_info = PixtralHFEncoderInfo(hf_config)
+        encoder_info = self.info.get_vision_encoder_info(hf_processor_mm_kwargs)
 
         def get_replacement(item_idx: int):
             images = mm_items.get_items("image", ImageProcessorItems)