| 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 | ] |