[Bugfix][Model] Mistral3: fix image placeholder grid for processor size overrides (#52874)
Signed-off-by: Anmol Gupta <14880251+anmolgupt@users.noreply.github.com>
anmolgupt committed
Aug 22, 2026 at 21:34 UTC
a3561ef8e49d3545c4078df43444beb4c98ae124
3 files changed
+210
-7
tests/models/multimodal/processing/test_mistral3.py
new
+143
@@ -0,0 +1,143 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for Mistral3's multimodal preprocessing kwargs."""
4
+
5
+import pytest
6
+import torch
7
+from PIL import Image
8
+from transformers import BatchFeature
9
+
10
+from vllm.model_executor.models.lightonocr import LightOnOCRProcessingInfo
11
+from vllm.model_executor.models.mistral3 import Mistral3HFEncoderInfo
12
+from vllm.model_executor.models.pixtral import PixtralHFEncoderInfo
13
+from vllm.multimodal import MULTIMODAL_REGISTRY
14
+from vllm.multimodal.inputs import MultiModalKwargsItems
15
+
16
+from ...utils import build_model_context
17
+
18
+# This repo ships both params.json (Mistral) and config.json (HF). Auto config
19
+# selects PixtralForConditionalGeneration; force HF to exercise Mistral3.
20
+_MODEL_CONFIG_KWARGS = {"config_format": "hf"}
21
+_MODEL_ID = "mistralai/Mistral-Small-3.1-24B-Instruct-2503"
22
+_LIGHTON_MODEL_ID = "lightonai/LightOnOCR-1B-1025"
23
+
24
+
25
+def _processed_pixel_values(
26
+ hf_processor,
27
+ images: list[Image.Image],
28
+ mm_processor_kwargs: dict[str, object],
29
+) -> list[torch.Tensor]:
30
+ """Resize via the HF image processor and un-pad to per-image H×W."""
31
+ image_processor = hf_processor.image_processor
32
+ hf_out = image_processor(images=images, return_tensors="pt", **mm_processor_kwargs)
33
+ pixel_values = hf_out["pixel_values"]
34
+ image_sizes = hf_out["image_sizes"]
35
+ return [p[:, :h, :w] for p, (h, w) in zip(pixel_values, image_sizes)]
36
+
37
+
38
+def _placeholder_count_from_prompt_updates(
39
+ processor,
40
+ images: list[Image.Image],
41
+ pixel_values: list[torch.Tensor],
42
+ mm_processor_kwargs: dict[str, object],
43
+) -> int:
44
+ hf_inputs = BatchFeature({"pixel_values": pixel_values})
45
+ fields_config = processor._get_mm_fields_config(hf_inputs, mm_processor_kwargs)
46
+ out_mm_kwargs = MultiModalKwargsItems.from_hf_inputs(hf_inputs, fields_config)
47
+ # Prompt updates use raw PIL sizes and must predict the processed grid.
48
+ mm_items = processor.info.parse_mm_data({"image": images})
49
+ updates = processor._get_prompt_updates(
50
+ mm_items, mm_processor_kwargs, out_mm_kwargs
51
+ )
52
+ image_token_id = processor.info.get_hf_config().image_token_index
53
+
54
+ total = 0
55
+ for item_idx in range(len(images)):
56
+ details = updates[0].resolve(item_idx).content
57
+ total += details.full.count(image_token_id)
58
+ return total
59
+
60
+
61
+def _expected_placeholder_tokens_per_image(
62
+ hf_processor,
63
+ pixel_values: torch.Tensor,
64
+) -> int:
65
+ """Count projected tokens from the actual HF-processed H×W."""
66
+ image_h, image_w = pixel_values.shape[-2:]
67
+ patch_size = hf_processor.image_processor.patch_size
68
+ if isinstance(patch_size, dict):
69
+ patch_h = patch_size["height"]
70
+ patch_w = patch_size["width"]
71
+ else:
72
+ patch_h = patch_w = int(patch_size)
73
+
74
+ spatial_merge_size = getattr(hf_processor, "spatial_merge_size", 1)
75
+ assert image_h % patch_h == 0
76
+ assert image_w % patch_w == 0
77
+
78
+ return (image_h // (patch_h * spatial_merge_size)) * (
79
+ image_w // (patch_w * spatial_merge_size)
80
+ )
81
+
82
+
83
+@pytest.mark.parametrize("model_id", [_MODEL_ID])
84
+@pytest.mark.parametrize(
85
+ ("mm_processor_kwargs", "image_size", "expected_toks_per_img"),
86
+ [
87
+ ({}, (448, 448), 256),
88
+ ({"size": {"longest_edge": 1008}}, (1540, 1540), 1296),
89
+ ({"size": {"longest_edge": 1288}}, (1536, 1187), 1656),
90
+ ({"size": {"longest_edge": 1008}}, (29, 29), 1),
91
+ ({"size": {"longest_edge": 1000}}, (1540, 1700), 1152),
92
+ ],
93
+)
94
+@pytest.mark.parametrize("num_imgs", [1, 2])
95
+@pytest.mark.parametrize("kwargs_on_init", [True, False])
96
+def test_processor_size_override(
97
+ model_id: str,
98
+ mm_processor_kwargs: dict[str, object],
99
+ image_size: tuple[int, int],
100
+ expected_toks_per_img: int,
101
+ num_imgs: int,
102
+ kwargs_on_init: bool,
103
+):
104
+ ctx = build_model_context(
105
+ model_id,
106
+ mm_processor_kwargs=mm_processor_kwargs if kwargs_on_init else None,
107
+ limit_mm_per_prompt={"image": num_imgs},
108
+ model_config_kwargs=_MODEL_CONFIG_KWARGS,
109
+ )
110
+ processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
111
+ hf_processor_mm_kwargs = {} if kwargs_on_init else mm_processor_kwargs
112
+ hf_processor = processor.info.get_hf_processor(**hf_processor_mm_kwargs)
113
+
114
+ dummy_image = Image.new("RGB", image_size, color=(127, 127, 127))
115
+ images = [dummy_image] * num_imgs
116
+ merged_mm_kwargs = processor.info.ctx.get_merged_mm_kwargs(hf_processor_mm_kwargs)
117
+ pixel_values = _processed_pixel_values(hf_processor, images, merged_mm_kwargs)
118
+
119
+ image_token_count = _placeholder_count_from_prompt_updates(
120
+ processor, images, pixel_values, hf_processor_mm_kwargs
121
+ )
122
+ expected_from_pixel_values = _expected_placeholder_tokens_per_image(
123
+ hf_processor, pixel_values[0]
124
+ )
125
+ assert expected_from_pixel_values == expected_toks_per_img
126
+ assert image_token_count == expected_from_pixel_values * num_imgs
127
+
128
+
129
+def test_lightonocr_keeps_vision_config_image_size():
130
+ ctx = build_model_context(
131
+ _LIGHTON_MODEL_ID,
132
+ mm_processor_kwargs={"size": {"longest_edge": 1008}},
133
+ model_config_kwargs=_MODEL_CONFIG_KWARGS,
134
+ )
135
+ processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
136
+
137
+ assert isinstance(processor.info, LightOnOCRProcessingInfo)
138
+ encoder_info = processor.info.get_vision_encoder_info()
139
+ assert isinstance(encoder_info, PixtralHFEncoderInfo)
140
+ assert not isinstance(encoder_info, Mistral3HFEncoderInfo)
141
+ assert encoder_info.get_image_size() == (
142
+ processor.info.get_hf_config().vision_config.image_size
143
+ )
vllm/model_executor/models/lightonocr.py
+13
-2
@@ -38,7 +38,18 @@ from vllm.multimodal.processing import (
38
_I = TypeVar("_I", bound=Mistral3ProcessingInfo)
39
40
41
-class LightOnOCRMultiModalProcessor(BaseMultiModalProcessor[Mistral3ProcessingInfo]):
41
+class LightOnOCRProcessingInfo(Mistral3ProcessingInfo):
42
+ def get_vision_encoder_info(
43
+ self,
44
+ mm_processor_kwargs: Mapping[str, object] | None = None,
45
+ ) -> PixtralHFEncoderInfo:
46
+ # LightOnOCR prompt updates use vision_config.image_size directly;
47
+ # preserve that behavior rather than inheriting Mistral3's processor
48
+ # longest_edge override.
49
+ return PixtralHFEncoderInfo(self.get_hf_config())
50
+
51
+
52
+class LightOnOCRMultiModalProcessor(BaseMultiModalProcessor[LightOnOCRProcessingInfo]):
53
def _get_hf_processor_text(self, mm_counts: Mapping[str, int]) -> str:
54
return self.dummy_inputs.get_dummy_text(mm_counts)
55
@@ -124,7 +135,7 @@ class LightOnOCRMultiModalProcessor(BaseMultiModalProcessor[Mistral3ProcessingIn
135
136
@MULTIMODAL_REGISTRY.register_processor(
137
LightOnOCRMultiModalProcessor,
127
- info=Mistral3ProcessingInfo,
138
+ info=LightOnOCRProcessingInfo,
139
dummy_inputs=Mistral3DummyInputsBuilder,
140
)
141
class LightOnOCRForConditionalGeneration(Mistral3ForConditionalGeneration):
vllm/model_executor/models/mistral3.py
+54
-5
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import math
5
from collections.abc import Iterable, Mapping, Sequence
6
from typing import Annotated, Literal
7
@@ -22,7 +23,11 @@ from vllm.multimodal.inputs import (
23
MultiModalFieldConfig,
24
MultiModalKwargsItems,
25
)
25
-from vllm.multimodal.parse import ImageProcessorItems, ImageSize, MultiModalDataItems
26
+from vllm.multimodal.parse import (
27
+ ImageProcessorItems,
28
+ ImageSize,
29
+ MultiModalDataItems,
30
+)
31
from vllm.multimodal.processing import (
32
BaseDummyInputsBuilder,
33
BaseMultiModalProcessor,
@@ -50,7 +55,6 @@ from .utils import (
55
init_vllm_registered_model,
56
maybe_prefix,
57
)
53
-from .vision import get_vision_encoder_info
58
59
60
class Mistral3ImagePixelInputs(TensorSchema):
@@ -170,12 +174,57 @@ class Mistral3MultiModalProjector(nn.Module):
174
return hidden_states
175
176
177
+class Mistral3HFEncoderInfo(PixtralHFEncoderInfo):
178
+ def __init__(self, hf_config: Mistral3Config, image_size: int) -> None:
179
+ super().__init__(hf_config)
180
+ self.image_size = image_size
181
+
182
+ def get_image_size(self) -> int:
183
+ return self.image_size
184
+
185
+ def get_patch_grid_size(
186
+ self,
187
+ *,
188
+ image_width: int,
189
+ image_height: int,
190
+ ) -> tuple[int, int]:
191
+ max_size = self.get_image_size()
192
+ ratio = max(image_width / max_size, image_height / max_size)
193
+
194
+ if ratio > 1:
195
+ image_width = math.floor(image_width / ratio)
196
+ image_height = math.floor(image_height / ratio)
197
+
198
+ patch_size = self.vision_config.patch_size
199
+ spatial_merge_size = self.hf_config.spatial_merge_size
200
+
201
+ # The HF processor rounds each dimension up to the vision patch size
202
+ # before the projector drops incomplete spatial-merge groups. This is
203
+ # not equivalent to rounding directly to the merged patch size.
204
+ num_width_patches = (image_width - 1) // patch_size + 1
205
+ num_height_patches = (image_height - 1) // patch_size + 1
206
+
207
+ ncols = num_width_patches // spatial_merge_size
208
+ nrows = num_height_patches // spatial_merge_size
209
+ return ncols, nrows
210
+
211
+
212
class Mistral3ProcessingInfo(BaseProcessingInfo):
213
def get_hf_config(self) -> Mistral3Config:
214
return self.ctx.get_hf_config(Mistral3Config)
215
177
- def get_vision_encoder_info(self):
178
- return get_vision_encoder_info(self.get_hf_config())
216
+ def get_vision_encoder_info(
217
+ self,
218
+ mm_processor_kwargs: Mapping[str, object] | None = None,
219
+ ) -> Mistral3HFEncoderInfo:
220
+ processor = self.get_hf_processor()
221
+ size = processor.image_processor.size
222
+ merged_kwargs = self.ctx.get_merged_mm_kwargs(mm_processor_kwargs or {})
223
+ if override_size := merged_kwargs.get("size"):
224
+ size = size | override_size
225
+
226
+ image_size = size["longest_edge"]
227
+ return Mistral3HFEncoderInfo(self.get_hf_config(), image_size)
228
229
def get_hf_processor(self, **kwargs: object):
230
return self.ctx.get_hf_processor(PixtralProcessor, **kwargs)
@@ -284,7 +333,7 @@ class Mistral3MultiModalProcessor(BaseMultiModalProcessor[Mistral3ProcessingInfo
333
image_end_id = vocab[processor.image_end_token]
334
335
assert isinstance(hf_config.vision_config, PixtralVisionConfig)
287
- encoder_info = PixtralHFEncoderInfo(hf_config)
336
+ encoder_info = self.info.get_vision_encoder_info(hf_processor_mm_kwargs)
337
338
def get_replacement(item_idx: int):
339
images = mm_items.get_items("image", ImageProcessorItems)