From a3561ef8e49d3545c4078df43444beb4c98ae124 Mon Sep 17 00:00:00 2001 From: anmolgupt <14880251+anmolgupt@users.noreply.github.com> Date: Sat, 22 Aug 2026 21:34:38 -0700 Subject: [PATCH] [Bugfix][Model] Mistral3: fix image placeholder grid for processor size overrides (#52874) Signed-off-by: Anmol Gupta <14880251+anmolgupt@users.noreply.github.com> --- .../multimodal/processing/test_mistral3.py | 143 ++++++++++++++++++ vllm/model_executor/models/lightonocr.py | 15 +- vllm/model_executor/models/mistral3.py | 59 +++++++- 3 files changed, 210 insertions(+), 7 deletions(-) create mode 100644 tests/models/multimodal/processing/test_mistral3.py diff --git a/tests/models/multimodal/processing/test_mistral3.py b/tests/models/multimodal/processing/test_mistral3.py new file mode 100644 index 0000000000..fcb783f22f --- /dev/null +++ b/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 + ) diff --git a/vllm/model_executor/models/lightonocr.py b/vllm/model_executor/models/lightonocr.py index ccb33177fb..7f108ba556 100644 --- a/vllm/model_executor/models/lightonocr.py +++ b/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): diff --git a/vllm/model_executor/models/mistral3.py b/vllm/model_executor/models/mistral3.py index 11f214c65b..23c6323619 100644 --- a/vllm/model_executor/models/mistral3.py +++ b/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)