main
py 376 lines 14.4 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Multimodal preprocessing for the DeepSeek-V4.1 vision variant.
4
5 The image transform and image-span construction are ported from the
6 official repository's ``image_processor.py`` so that token counts bit-match
7 the reference. Each ``<|deepseek_image|>`` placeholder in the prompt expands
8 to ``[IMAGE_START] + ([IMAGE] * n_llm_w + [IMAGE_NEW_LINE]) * n_llm_h +
9 [IMAGE_END]``; every span position carries ``image_token_id`` (129264) in
10 ``input_ids`` and the roles ride along in a per-image ``types`` tensor
11 (the reference's out-of-band ``token_types``). IMAGE slots receive aligner
12 rows in reading order; the delimiters take the learned ``image_start`` /
13 ``image_newline`` / ``image_end`` vectors.
14
15 The token stream matches the reference exactly: no compressor-alignment
16 pad is inserted (the reference pools image tokens across compressor-group
17 boundaries freely).
18 """
19
20 import math
21 from collections.abc import Mapping, Sequence
22 from typing import Any
23
24 import numpy as np
25 import torch
26 from PIL import Image, ImageOps
27 from transformers import BatchFeature
28
29 from vllm.config.multimodal import MultiModalDummyOptions
30 from vllm.inputs import MultiModalDataDict
31 from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems
32 from vllm.multimodal.parse import ImageSize, MultiModalDataItems
33 from vllm.multimodal.processing import (
34 BaseDummyInputsBuilder,
35 BaseMultiModalProcessor,
36 BaseProcessingInfo,
37 PromptReplacement,
38 PromptUpdate,
39 PromptUpdateDetails,
40 )
41 from vllm.transformers_utils.configs.deepseek_v41 import DeepseekV41Config
42
43 IMAGE_START, IMAGE, IMAGE_NEW_LINE, IMAGE_END = range(4)
44
45 IMAGE_PLACEHOLDER = "<|deepseek_image|>"
46
47 # The checkpoint tokenizer's id for IMAGE_PLACEHOLDER (the config's
48 # ``image_token_id``). Every image-span position carries this id, exactly
49 # like the reference; the MoE router detects image tokens as the range
50 # [IMAGE_SENTINEL_BASE_ID, IMAGE_SENTINEL_BASE_ID + 5).
51 IMAGE_SENTINEL_BASE_ID = 129264
52
53
54 def image_sentinel_mask(token_ids: torch.Tensor) -> torch.Tensor:
55 """Boolean mask for image-span positions."""
56 return token_ids == IMAGE_SENTINEL_BASE_ID
57
58
59 def validate_image_sentinel_ids(tokenizer) -> None:
60 """Check the image token id against the tokenizer."""
61 image_id = tokenizer.convert_tokens_to_ids(IMAGE_PLACEHOLDER)
62 if image_id != IMAGE_SENTINEL_BASE_ID:
63 raise ValueError(
64 f"Image placeholder {IMAGE_PLACEHOLDER!r} has id {image_id}, "
65 f"expected {IMAGE_SENTINEL_BASE_ID} (the config's "
66 "image_token_id); the DeepSeek-V4.1 vision path keys image "
67 "routing and engram masking off this id."
68 )
69
70
71 def llm_grid(best_height, best_width, patch_size, downsample_ratio):
72 """Token grid the aligner produces from a patch grid of this pixel size."""
73 return (
74 math.ceil((best_height // patch_size) / downsample_ratio),
75 math.ceil((best_width // patch_size) / downsample_ratio),
76 )
77
78
79 def num_image_tokens(n_llm_h: int, n_llm_w: int) -> int:
80 return n_llm_h * (n_llm_w + 1) + 2
81
82
83 def solve_resize_ratio(height, width, patch_size, downsample_ratio, max_n_token):
84 """Largest aspect-preserving pixel size whose token grid still fits in
85 max_n_token. Returns (best_height, best_width)."""
86 r = height / width
87 max_w_float = math.sqrt((max_n_token - 2) / r + 0.25) - 0.5
88 max_h_float = max_w_float * r
89 cell = patch_size * downsample_ratio
90 if max_w_float < 1.0: # very tall: collapse to a single column
91 return (max_n_token - 2) // 2 * cell, cell
92 if max_h_float < 1.0: # very wide: collapse to a single row
93 return cell, (max_n_token - 3) * cell
94 beta = min(
95 math.floor(max_w_float) * cell / width,
96 math.floor(max_h_float) * cell / height,
97 )
98 return (
99 math.floor(height * beta / patch_size) * patch_size,
100 math.floor(width * beta / patch_size) * patch_size,
101 )
102
103
104 def safe_resize(
105 height, width, best_height, best_width, patch_size, downsample_ratio, max_n_token
106 ):
107 """Shrink the pixel size until the image costs at most max_n_token LLM
108 tokens."""
109 n_llm_h, n_llm_w = llm_grid(best_height, best_width, patch_size, downsample_ratio)
110 if num_image_tokens(n_llm_h, n_llm_w) > max_n_token:
111 best_height, best_width = solve_resize_ratio(
112 height, width, patch_size, downsample_ratio, max_n_token
113 )
114 n_llm_h, n_llm_w = llm_grid(
115 best_height, best_width, patch_size, downsample_ratio
116 )
117 assert num_image_tokens(n_llm_h, n_llm_w) <= max_n_token
118 return n_llm_h, n_llm_w, best_height, best_width
119
120
121 def load_image(
122 image: Image.Image,
123 *,
124 patch_size: int,
125 downsample_ratio: int,
126 max_n_token: int,
127 min_pixels: int,
128 max_wh_ratio: float | None,
129 ):
130 """Transform one PIL image into ViT patches.
131
132 Same math as the reference ``load_image``, except the image is already
133 decoded (vLLM supplies PIL images instead of a record dict).
134 """
135 p = patch_size
136 image = image.convert("RGB")
137 width, height = image.size
138 if max_wh_ratio is not None and width > height * max_wh_ratio:
139 width = height * max_wh_ratio
140 if 0 < width * height < min_pixels:
141 ratio = (min_pixels / (width * height)) ** 0.5
142 width = int(width * ratio)
143 height = int(height * ratio)
144 best_width = math.ceil(width / p) * p
145 best_height = math.ceil(height / p) * p
146 n_llm_h, n_llm_w, best_height, best_width = safe_resize(
147 height, width, best_height, best_width, p, downsample_ratio, max_n_token
148 )
149 n_vit_h, n_vit_w = best_height // p, best_width // p
150 if max_wh_ratio is not None and image.width >= max_wh_ratio * image.height:
151 image = image.resize((best_width, best_height))
152 else:
153 image = ImageOps.pad(image, (best_width, best_height), color=(127, 127, 127))
154 x = torch.from_numpy(np.asarray(image, dtype=np.float32)).permute(2, 0, 1) / 255
155 x = ((x - 0.5) / 0.5).to(torch.bfloat16)
156 patches = (
157 x.reshape(3, n_vit_h, p, n_vit_w, p)
158 .permute(1, 3, 0, 2, 4)
159 .reshape(n_vit_h * n_vit_w, 3, p, p)
160 )
161 return patches, n_vit_h, n_vit_w, n_llm_h, n_llm_w
162
163
164 def image_token_types(n_llm_h: int, n_llm_w: int) -> torch.Tensor:
165 """Reading-order span layout: one IMAGE_NEW_LINE per row."""
166 types = [IMAGE_START]
167 types += ([IMAGE] * n_llm_w + [IMAGE_NEW_LINE]) * n_llm_h
168 types.append(IMAGE_END)
169 return torch.tensor(types, dtype=torch.int64)
170
171
172 class DeepseekV4VLImageProcessor:
173 """Per-image transform (the PIL-input equivalent of the reference
174 ``load_image``)."""
175
176 def __init__(self, config: DeepseekV41Config) -> None:
177 super().__init__()
178 self.patch_size = config.vision_patch_size
179 self.downsample_ratio = config.vision_downsample_ratio
180 self.max_n_token = config.vision_max_n_token
181 self.min_pixels = config.vision_min_pixels
182 self.max_wh_ratio = config.vision_max_wh_ratio
183
184 def __call__(self, image: Image.Image):
185 return load_image(
186 image,
187 patch_size=self.patch_size,
188 downsample_ratio=self.downsample_ratio,
189 max_n_token=self.max_n_token,
190 min_pixels=self.min_pixels,
191 max_wh_ratio=self.max_wh_ratio,
192 )
193
194
195 class DeepseekV4VLProcessor:
196 """Minimal stand-in for the HF processor of DeepSeek-V4.1 vision models.
197
198 The official repository ships image preprocessing as plain functions in
199 ``image_processor.py`` (no ``auto_map`` processor), so this class wraps
200 their ports directly and the model loads without ``--trust-remote-code``.
201
202 ``__call__`` returns a ``BatchFeature`` with one entry per image
203 (flattened across images):
204
205 - ``patches``: ``(sum(n_vit_h * n_vit_w), 3, p, p)`` bf16 ViT patches.
206 - ``vit_grid``: ``(num_images, 2)`` int64 ``[n_vit_h, n_vit_w]``.
207 - ``llm_grid``: ``(num_images, 2)`` int64 ``[n_llm_h, n_llm_w]``.
208 - ``types``: concatenated per-image pad-free span roles
209 (IMAGE_START/IMAGE/IMAGE_NEW_LINE/IMAGE_END); every span position
210 carries ``image_token_id`` in the prompt's token ids.
211 """
212
213 def __init__(self, config: DeepseekV41Config) -> None:
214 super().__init__()
215 self.config = config
216 self.image_processor = DeepseekV4VLImageProcessor(config)
217
218 def __call__(
219 self,
220 text: str | None = None,
221 images: Sequence[Image.Image] | None = None,
222 return_tensors: str | None = None,
223 **kwargs: Any,
224 ) -> BatchFeature:
225 patches_list = []
226 vit_grid = []
227 llm_grid_list = []
228 types_list = []
229 for image in images or []:
230 patches, n_vit_h, n_vit_w, n_llm_h, n_llm_w = self.image_processor(image)
231 patches_list.append(patches)
232 vit_grid.append((n_vit_h, n_vit_w))
233 llm_grid_list.append((n_llm_h, n_llm_w))
234 types_list.append(image_token_types(n_llm_h, n_llm_w))
235
236 if not patches_list:
237 return BatchFeature({})
238
239 return BatchFeature(
240 {
241 "patches": torch.cat(patches_list),
242 "vit_grid": torch.tensor(vit_grid, dtype=torch.int64),
243 "llm_grid": torch.tensor(llm_grid_list, dtype=torch.int64),
244 "types": torch.cat(types_list),
245 }
246 )
247
248
249 class DeepseekV4VLProcessingInfo(BaseProcessingInfo):
250 def get_hf_config(self) -> DeepseekV41Config:
251 return self.ctx.get_hf_config(DeepseekV41Config)
252
253 def get_hf_processor(self, **kwargs: object) -> DeepseekV4VLProcessor:
254 if kwargs:
255 raise ValueError(f"Unexpected processor kwargs: {sorted(kwargs)}")
256 return DeepseekV4VLProcessor(self.get_hf_config())
257
258 def get_supported_mm_limits(self) -> Mapping[str, int | None]:
259 return {"image": None}
260
261 def get_mm_max_tokens_per_item(
262 self,
263 seq_len: int,
264 mm_counts: Mapping[str, int],
265 ) -> Mapping[str, int]:
266 return {"image": self.get_hf_config().vision_max_n_token}
267
268 def get_image_placeholder_token_id(self) -> int:
269 token_id = self.get_tokenizer().convert_tokens_to_ids(IMAGE_PLACEHOLDER)
270 if token_id is None:
271 raise ValueError(f"Token not found in tokenizer: {IMAGE_PLACEHOLDER}")
272 return token_id
273
274 def get_image_size_with_most_features(self) -> ImageSize:
275 hf_config = self.get_hf_config()
276 patch_size = hf_config.vision_patch_size
277 downsample_ratio = hf_config.vision_downsample_ratio
278 budget = hf_config.vision_max_n_token
279 max_wh_ratio = hf_config.vision_max_wh_ratio
280 best_h = best_w = 0
281 # Each LLM row costs an extra newline token. Maximize the ViT grid,
282 # including the patch rounding after the optional aspect-ratio clamp.
283 for n_llm_h in range(1, (budget - 2) // 2 + 1):
284 n_vit_h = n_llm_h * downsample_ratio
285 n_vit_w = ((budget - 2) // n_llm_h - 1) * downsample_ratio
286 if max_wh_ratio is not None:
287 n_vit_w = min(n_vit_w, math.ceil(n_vit_h * max_wh_ratio))
288 if n_vit_h * n_vit_w > best_h * best_w:
289 best_h, best_w = n_vit_h, n_vit_w
290 if best_h == 0:
291 raise ValueError(
292 "Image token budget and aspect ratio allow no image patches"
293 )
294 height, width = best_h * patch_size, best_w * patch_size
295 # Avoid the min-pixels upscaler perturbing the selected aspect ratio.
296 scale = 1
297 while height * width * scale * scale < hf_config.vision_min_pixels:
298 scale *= 2
299 return ImageSize(width=width * scale, height=height * scale)
300
301
302 class DeepseekV4VLDummyInputsBuilder(
303 BaseDummyInputsBuilder[DeepseekV4VLProcessingInfo]
304 ):
305 def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
306 return IMAGE_PLACEHOLDER * mm_counts.get("image", 0)
307
308 def get_dummy_mm_data(
309 self,
310 seq_len: int,
311 mm_counts: Mapping[str, int],
312 mm_options: MultiModalDummyOptions,
313 ) -> MultiModalDataDict:
314 size = self.info.get_image_size_with_most_features()
315 return {
316 "image": self._get_dummy_images(
317 width=size.width,
318 height=size.height,
319 num_images=mm_counts.get("image", 0),
320 overrides=mm_options.get("image"),
321 ),
322 }
323
324
325 class DeepseekV4VLMultiModalProcessor(
326 BaseMultiModalProcessor[DeepseekV4VLProcessingInfo]
327 ):
328 def _get_mm_fields_config(
329 self,
330 hf_inputs: BatchFeature,
331 hf_processor_mm_kwargs: Mapping[str, object],
332 ) -> Mapping[str, MultiModalFieldConfig]:
333 vit_grid = hf_inputs.get("vit_grid")
334 llm_grid = hf_inputs.get("llm_grid")
335
336 if vit_grid is None or llm_grid is None:
337 empty = torch.empty(0, dtype=torch.long)
338 patch_sizes = types_sizes = empty
339 else:
340 patch_sizes = vit_grid.prod(-1)
341 n_llm_h, n_llm_w = llm_grid[:, 0], llm_grid[:, 1]
342 types_sizes = n_llm_h * (n_llm_w + 1) + 2
343
344 return {
345 "patches": MultiModalFieldConfig.flat_from_sizes("image", patch_sizes),
346 "vit_grid": MultiModalFieldConfig.batched("image", keep_on_cpu=True),
347 "llm_grid": MultiModalFieldConfig.batched("image", keep_on_cpu=True),
348 "types": MultiModalFieldConfig.flat_from_sizes(
349 "image", types_sizes, keep_on_cpu=True
350 ),
351 }
352
353 def _get_prompt_updates(
354 self,
355 mm_items: MultiModalDataItems,
356 hf_processor_mm_kwargs: Mapping[str, object],
357 out_mm_kwargs: MultiModalKwargsItems,
358 ) -> Sequence[PromptUpdate]:
359 image_token_id = self.info.get_image_placeholder_token_id()
360 validate_image_sentinel_ids(self.info.get_tokenizer())
361
362 def get_image_replacement(item_idx: int) -> PromptUpdateDetails:
363 types: torch.Tensor = out_mm_kwargs["image"][item_idx]["types"].data
364 # Every span position carries image_token_id; the roles live in
365 # ``types``. All of them are embed positions (delimiters get the
366 # learned vectors from embed_multimodal, not the embed table).
367 full = [image_token_id] * types.numel()
368 return PromptUpdateDetails.select_token_id(full, image_token_id)
369
370 return [
371 PromptReplacement(
372 modality="image",
373 target=[image_token_id],
374 replacement=get_image_replacement,
375 ),
376 ]