| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Multimodal preprocessing for the DeepSeek-V4 vision variants |
| 4 | (DeepSeek-V4-Flash-Vision-Exp). |
| 5 | |
| 6 | The image transform and sentinel-block construction are ported from the |
| 7 | official repository's ``image_processor.py`` so that token counts bit-match |
| 8 | the reference. Each ``<｜deepseek_image｜>`` placeholder in the prompt expands |
| 9 | to a variable-length block of sentinel tokens; only positions with |
| 10 | ``type == IMAGE`` receive vision embeddings, the other sentinels are looked |
| 11 | up from learned embedding vectors in the model. |
| 12 | |
| 13 | Unlike the reference (which uses out-of-vocab ids ``vocab_size + type``), |
| 14 | the sentinel block borrows five consecutive reserved tokenizer tokens |
| 15 | (``<|place_holder_mm_span_0431|>`` .. ``_0435|>``): they are special tokens |
| 16 | the tokenizer never emits from plain text, so the ids stay in-vocabulary and |
| 17 | work with stock token validation, logprobs and detokenization. The ids are |
| 18 | pure markers — every sentinel position's embedding is overwritten with |
| 19 | vision/sentinel vectors before the decoder layers see it, exactly like the |
| 20 | reference's out-of-vocab scheme. |
| 21 | """ |
| 22 | |
| 23 | import math |
| 24 | from collections.abc import Mapping, Sequence |
| 25 | from typing import Any |
| 26 | |
| 27 | import numpy as np |
| 28 | import torch |
| 29 | from PIL import Image, ImageOps |
| 30 | from transformers import BatchFeature |
| 31 | from typing_extensions import assert_never |
| 32 | |
| 33 | from vllm.config.multimodal import MultiModalDummyOptions |
| 34 | from vllm.inputs import MultiModalDataDict |
| 35 | from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems |
| 36 | from vllm.multimodal.parse import ImageSize, MultiModalDataItems |
| 37 | from vllm.multimodal.processing import ( |
| 38 | BaseDummyInputsBuilder, |
| 39 | BaseMultiModalProcessor, |
| 40 | BaseProcessingInfo, |
| 41 | PromptReplacement, |
| 42 | PromptUpdate, |
| 43 | PromptUpdateDetails, |
| 44 | ) |
| 45 | from vllm.multimodal.processing.processor import ( |
| 46 | MultiModalPromptUpdates, |
| 47 | MultiModalPromptUpdatesApplyResult, |
| 48 | PlaceholderFeaturesInfo, |
| 49 | UpdateMode, |
| 50 | _plan_prompt_updates, |
| 51 | ) |
| 52 | from vllm.transformers_utils.configs.deepseek_v4 import DeepseekV4Config |
| 53 | |
| 54 | IMAGE_START, IMAGE_PAD, IMAGE, IMAGE_NEW_LINE, IMAGE_END = range(5) |
| 55 | COMPRESS_PAD_TO = 4 |
| 56 | |
| 57 | IMAGE_PLACEHOLDER = "<｜deepseek_image｜>" |
| 58 | |
| 59 | # Sentinel roles borrow five consecutive ``<|place_holder_mm_span_XXXX|>`` |
| 60 | # tokens (reserved special tokens, never emitted from plain text). The order |
| 61 | # must match IMAGE_START..IMAGE_END above so that |
| 62 | # ``id == IMAGE_SENTINEL_BASE_ID + type``. |
| 63 | IMAGE_SENTINEL_BASE_ID = 129257 |
| 64 | IMAGE_SENTINEL_TOKEN_NAMES = ( |
| 65 | "<|place_holder_mm_span_0431|>", # IMAGE_START |
| 66 | "<|place_holder_mm_span_0432|>", # IMAGE_PAD |
| 67 | "<|place_holder_mm_span_0433|>", # IMAGE |
| 68 | "<|place_holder_mm_span_0434|>", # IMAGE_NEW_LINE |
| 69 | "<|place_holder_mm_span_0435|>", # IMAGE_END |
| 70 | ) |
| 71 | |
| 72 | |
| 73 | def image_sentinel_mask(token_ids: torch.Tensor) -> torch.Tensor: |
| 74 | """Boolean mask for image-block sentinel positions (in-vocab ids).""" |
| 75 | return (token_ids >= IMAGE_SENTINEL_BASE_ID) & ( |
| 76 | token_ids < IMAGE_SENTINEL_BASE_ID + len(IMAGE_SENTINEL_TOKEN_NAMES) |
| 77 | ) |
| 78 | |
| 79 | |
| 80 | def validate_image_sentinel_ids(tokenizer) -> None: |
| 81 | """Check the borrowed sentinel ids against the tokenizer.""" |
| 82 | for i, name in enumerate(IMAGE_SENTINEL_TOKEN_NAMES): |
| 83 | token_id = tokenizer.convert_tokens_to_ids(name) |
| 84 | if token_id != IMAGE_SENTINEL_BASE_ID + i: |
| 85 | raise ValueError( |
| 86 | f"Image sentinel token {name!r} has id {token_id}, expected " |
| 87 | f"{IMAGE_SENTINEL_BASE_ID + i}; the DeepSeek-V4 vision " |
| 88 | "sentinel block requires these consecutive reserved ids." |
| 89 | ) |
| 90 | |
| 91 | |
| 92 | def grid_tokens(best_height, best_width, patch_size, downsample_ratio): |
| 93 | """Number of LLM tokens the aligner grid occupies (N-layout, including |
| 94 | row/align padding).""" |
| 95 | n_llm_h = math.ceil((best_height // patch_size) / downsample_ratio) |
| 96 | n_llm_w = math.ceil((best_width // patch_size) / downsample_ratio) |
| 97 | num_tokens = n_llm_h * (n_llm_w + 1) + 2 |
| 98 | if n_llm_h % 2 == 1: |
| 99 | num_tokens += n_llm_w + 1 |
| 100 | num_tokens += (n_llm_h + 1) // 2 * (n_llm_w + 1) % 2 * 2 |
| 101 | return n_llm_h, n_llm_w, num_tokens |
| 102 | |
| 103 | |
| 104 | def solve_resize_ratio(height, width, patch_size, downsample_ratio, max_n_token): |
| 105 | r = height / width |
| 106 | max_w_float = math.sqrt((max_n_token - 2) / r + 0.25) - 0.5 |
| 107 | max_h_float = max_w_float * r |
| 108 | if max_w_float < 1.0: |
| 109 | max_w = 1 |
| 110 | max_h = (max_n_token - 2) // (max_w + 1) |
| 111 | if max_h % 2 == 1: |
| 112 | max_h -= 1 |
| 113 | best_width = max_w * patch_size * downsample_ratio |
| 114 | best_height = max_h * patch_size * downsample_ratio |
| 115 | elif max_h_float < 2.0: |
| 116 | max_h = 2 |
| 117 | max_w = ((max_n_token - 2) // max_h) - 1 |
| 118 | assert max_w > 1 |
| 119 | best_width = max_w * patch_size * downsample_ratio |
| 120 | best_height = max_h * patch_size * downsample_ratio |
| 121 | else: |
| 122 | max_w = math.floor(max_w_float) |
| 123 | max_h = math.floor(max_h_float) |
| 124 | if max_h % 2 == 1: |
| 125 | max_h -= 1 |
| 126 | beta = min( |
| 127 | max_w * patch_size * downsample_ratio / width, |
| 128 | max_h * patch_size * downsample_ratio / height, |
| 129 | ) |
| 130 | best_width = math.floor(width * beta / patch_size) * patch_size |
| 131 | best_height = math.floor(height * beta / patch_size) * patch_size |
| 132 | n_llm_h, n_llm_w, num_tokens = grid_tokens( |
| 133 | best_height, best_width, patch_size, downsample_ratio |
| 134 | ) |
| 135 | return n_llm_h, n_llm_w, best_height, best_width, num_tokens |
| 136 | |
| 137 | |
| 138 | def safe_resize( |
| 139 | height, width, best_height, best_width, patch_size, downsample_ratio, max_n_token |
| 140 | ): |
| 141 | max_n_token -= COMPRESS_PAD_TO - 1 |
| 142 | n_llm_h, n_llm_w, num_tokens = grid_tokens( |
| 143 | best_height, best_width, patch_size, downsample_ratio |
| 144 | ) |
| 145 | budget = max_n_token |
| 146 | while num_tokens > max_n_token: |
| 147 | n_llm_h, n_llm_w, best_height, best_width, num_tokens = solve_resize_ratio( |
| 148 | height, width, patch_size, downsample_ratio, budget |
| 149 | ) |
| 150 | budget -= 1 |
| 151 | return n_llm_h, n_llm_w, best_height, best_width |
| 152 | |
| 153 | |
| 154 | def load_image( |
| 155 | image: Image.Image, |
| 156 | *, |
| 157 | patch_size: int, |
| 158 | downsample_ratio: int, |
| 159 | max_n_token: int, |
| 160 | min_pixels: int, |
| 161 | max_wh_ratio: float | None, |
| 162 | ): |
| 163 | """Transform one PIL image into ViT patches. |
| 164 | |
| 165 | Same math as the reference ``load_image``, except the image is already |
| 166 | decoded (vLLM supplies PIL images instead of a record dict). |
| 167 | """ |
| 168 | p = patch_size |
| 169 | image = image.convert("RGB") |
| 170 | width, height = image.size |
| 171 | if max_wh_ratio is not None and width > height * max_wh_ratio: |
| 172 | width = height * max_wh_ratio |
| 173 | if 0 < width * height < min_pixels: |
| 174 | ratio = (min_pixels / (width * height)) ** 0.5 |
| 175 | width = int(width * ratio) |
| 176 | height = int(height * ratio) |
| 177 | best_width = math.ceil(width / p) * p |
| 178 | best_height = math.ceil(height / p) * p |
| 179 | n_llm_h, n_llm_w, best_height, best_width = safe_resize( |
| 180 | height, width, best_height, best_width, p, downsample_ratio, max_n_token |
| 181 | ) |
| 182 | n_vit_h, n_vit_w = best_height // p, best_width // p |
| 183 | if max_wh_ratio is not None and image.width >= max_wh_ratio * image.height: |
| 184 | image = image.resize((best_width, best_height)) |
| 185 | else: |
| 186 | image = ImageOps.pad(image, (best_width, best_height), color=(127, 127, 127)) |
| 187 | x = torch.from_numpy(np.asarray(image, dtype=np.float32)).permute(2, 0, 1) / 255 |
| 188 | x = ((x - 0.5) / 0.5).to(torch.bfloat16) |
| 189 | patches = ( |
| 190 | x.reshape(3, n_vit_h, p, n_vit_w, p) |
| 191 | .permute(1, 3, 0, 2, 4) |
| 192 | .reshape(n_vit_h * n_vit_w, 3, p, p) |
| 193 | ) |
| 194 | return patches, n_vit_h, n_vit_w, n_llm_h, n_llm_w |
| 195 | |
| 196 | |
| 197 | def build_image_block(n_llm_h: int, n_llm_w: int, start_pos: int): |
| 198 | """Builds the N-layout token types (final order) and the aligner-row order |
| 199 | for IMAGE slots.""" |
| 200 | compress_pad = COMPRESS_PAD_TO - 1 - start_pos % COMPRESS_PAD_TO |
| 201 | pad_h = n_llm_h % 2 |
| 202 | rows = n_llm_h + pad_h |
| 203 | row_len = n_llm_w + 1 |
| 204 | pad_last = rows // 2 * row_len % 2 * 2 |
| 205 | types = torch.tensor( |
| 206 | ([IMAGE] * n_llm_w + [IMAGE_NEW_LINE]) * n_llm_h |
| 207 | + [IMAGE_PAD] * (row_len * pad_h), |
| 208 | dtype=torch.int64, |
| 209 | ) |
| 210 | order = torch.arange(rows * row_len).view(rows // 2, 2, row_len) |
| 211 | order = order.transpose(1, 2).reshape(-1) |
| 212 | image_idx = torch.full((rows * row_len,), -1, dtype=torch.int64) |
| 213 | image_idx.view(rows, row_len)[:n_llm_h, :n_llm_w] = torch.arange( |
| 214 | n_llm_h * n_llm_w |
| 215 | ).view(n_llm_h, n_llm_w) |
| 216 | perm = image_idx[order] |
| 217 | perm = perm[perm >= 0] |
| 218 | types = torch.cat( |
| 219 | [ |
| 220 | torch.full((compress_pad,), IMAGE_PAD, dtype=torch.int64), |
| 221 | torch.tensor([IMAGE_START]), |
| 222 | types[order], |
| 223 | torch.full((pad_last,), IMAGE_PAD, dtype=torch.int64), |
| 224 | torch.tensor([IMAGE_END]), |
| 225 | ] |
| 226 | ) |
| 227 | return types, perm |
| 228 | |
| 229 | |
| 230 | def build_image_block_pad_free(n_llm_h: int, n_llm_w: int): |
| 231 | """``build_image_block`` without the position-dependent compressor pad. |
| 232 | |
| 233 | ``start_pos`` is chosen so that ``compress_pad == 0``; the pad is instead |
| 234 | prepended when the block is spliced into the final prompt, where its |
| 235 | position is known. |
| 236 | """ |
| 237 | return build_image_block(n_llm_h, n_llm_w, COMPRESS_PAD_TO - 1) |
| 238 | |
| 239 | |
| 240 | class DeepseekV4VLImageProcessor: |
| 241 | """Per-image transform (the PIL-input equivalent of the reference |
| 242 | ``load_image``).""" |
| 243 | |
| 244 | def __init__(self, config: DeepseekV4Config) -> None: |
| 245 | super().__init__() |
| 246 | self.patch_size = config.vision_patch_size |
| 247 | self.downsample_ratio = config.vision_downsample_ratio |
| 248 | self.max_n_token = config.vision_max_n_token |
| 249 | self.min_pixels = config.vision_min_pixels |
| 250 | self.max_wh_ratio = config.vision_max_wh_ratio |
| 251 | |
| 252 | def __call__(self, image: Image.Image): |
| 253 | return load_image( |
| 254 | image, |
| 255 | patch_size=self.patch_size, |
| 256 | downsample_ratio=self.downsample_ratio, |
| 257 | max_n_token=self.max_n_token, |
| 258 | min_pixels=self.min_pixels, |
| 259 | max_wh_ratio=self.max_wh_ratio, |
| 260 | ) |
| 261 | |
| 262 | |
| 263 | class DeepseekV4VLProcessor: |
| 264 | """Minimal stand-in for the HF processor of DeepSeek-V4 vision models. |
| 265 | |
| 266 | The official repository ships image preprocessing as plain functions in |
| 267 | ``image_processor.py`` (no ``auto_map`` processor), so this class wraps |
| 268 | their ports directly and the model loads without ``--trust-remote-code``. |
| 269 | |
| 270 | ``__call__`` returns a ``BatchFeature`` with one entry per image |
| 271 | (flattened across images): |
| 272 | |
| 273 | - ``patches``: ``(sum(n_vit_h * n_vit_w), 3, p, p)`` bf16 ViT patches. |
| 274 | - ``vit_grid``: ``(num_images, 2)`` int64 ``[n_vit_h, n_vit_w]``. |
| 275 | - ``llm_grid``: ``(num_images, 2)`` int64 ``[n_llm_h, n_llm_w]``. |
| 276 | - ``perm``: concatenated per-image ``(n_llm_h * n_llm_w,)`` int64 index |
| 277 | selecting aligner outputs into the final N-layout order. |
| 278 | - ``types``: concatenated per-image pad-free sentinel block types; |
| 279 | ``block ids = IMAGE_SENTINEL_BASE_ID + types``. |
| 280 | """ |
| 281 | |
| 282 | def __init__(self, config: DeepseekV4Config) -> None: |
| 283 | super().__init__() |
| 284 | self.config = config |
| 285 | self.image_processor = DeepseekV4VLImageProcessor(config) |
| 286 | |
| 287 | def __call__( |
| 288 | self, |
| 289 | text: str | None = None, |
| 290 | images: Sequence[Image.Image] | None = None, |
| 291 | return_tensors: str | None = None, |
| 292 | **kwargs: Any, |
| 293 | ) -> BatchFeature: |
| 294 | patches_list = [] |
| 295 | vit_grid = [] |
| 296 | llm_grid = [] |
| 297 | perm_list = [] |
| 298 | types_list = [] |
| 299 | for image in images or []: |
| 300 | patches, n_vit_h, n_vit_w, n_llm_h, n_llm_w = self.image_processor(image) |
| 301 | types, perm = build_image_block_pad_free(n_llm_h, n_llm_w) |
| 302 | patches_list.append(patches) |
| 303 | vit_grid.append((n_vit_h, n_vit_w)) |
| 304 | llm_grid.append((n_llm_h, n_llm_w)) |
| 305 | perm_list.append(perm) |
| 306 | types_list.append(types) |
| 307 | |
| 308 | if not patches_list: |
| 309 | return BatchFeature({}) |
| 310 | |
| 311 | return BatchFeature( |
| 312 | { |
| 313 | "patches": torch.cat(patches_list), |
| 314 | "vit_grid": torch.tensor(vit_grid, dtype=torch.int64), |
| 315 | "llm_grid": torch.tensor(llm_grid, dtype=torch.int64), |
| 316 | "perm": torch.cat(perm_list), |
| 317 | "types": torch.cat(types_list), |
| 318 | } |
| 319 | ) |
| 320 | |
| 321 | |
| 322 | class DeepseekV4VLProcessingInfo(BaseProcessingInfo): |
| 323 | def get_hf_config(self) -> DeepseekV4Config: |
| 324 | return self.ctx.get_hf_config(DeepseekV4Config) |
| 325 | |
| 326 | def get_hf_processor(self, **kwargs: object) -> DeepseekV4VLProcessor: |
| 327 | if kwargs: |
| 328 | raise ValueError(f"Unexpected processor kwargs: {sorted(kwargs)}") |
| 329 | return DeepseekV4VLProcessor(self.get_hf_config()) |
| 330 | |
| 331 | def get_supported_mm_limits(self) -> Mapping[str, int | None]: |
| 332 | return {"image": None} |
| 333 | |
| 334 | def get_mm_max_tokens_per_item( |
| 335 | self, |
| 336 | seq_len: int, |
| 337 | mm_counts: Mapping[str, int], |
| 338 | ) -> Mapping[str, int]: |
| 339 | # ``safe_resize`` reserves COMPRESS_PAD_TO - 1 tokens of |
| 340 | # vision_max_n_token for the compressor-alignment pad, so the full |
| 341 | # sentinel block is bounded by vision_max_n_token; the margin is kept |
| 342 | # in case that reservation changes. |
| 343 | return { |
| 344 | "image": self.get_hf_config().vision_max_n_token + COMPRESS_PAD_TO - 1, |
| 345 | } |
| 346 | |
| 347 | def get_image_placeholder_token_id(self) -> int: |
| 348 | token_id = self.get_tokenizer().convert_tokens_to_ids(IMAGE_PLACEHOLDER) |
| 349 | if token_id is None: |
| 350 | raise ValueError(f"Token not found in tokenizer: {IMAGE_PLACEHOLDER}") |
| 351 | return token_id |
| 352 | |
| 353 | def get_image_size_with_most_features(self) -> ImageSize: |
| 354 | hf_config = self.get_hf_config() |
| 355 | patch_size = hf_config.vision_patch_size |
| 356 | downsample_ratio = hf_config.vision_downsample_ratio |
| 357 | max_wh_ratio = hf_config.vision_max_wh_ratio |
| 358 | # Every aligner row costs an extra IMAGE_NEW_LINE token, so height is |
| 359 | # more expensive than width and a square is *not* optimal: search the |
| 360 | # LLM grids that fit the budget for the one with the most ViT patches |
| 361 | # (ties broken by the longer sentinel block). |
| 362 | budget = hf_config.vision_max_n_token - (COMPRESS_PAD_TO - 1) |
| 363 | step = patch_size * downsample_ratio |
| 364 | best_grid = (1, 1) |
| 365 | best_key = (0, 0) |
| 366 | for n_llm_h in range(1, budget): |
| 367 | if ( |
| 368 | grid_tokens(n_llm_h * step, step, patch_size, downsample_ratio)[2] |
| 369 | > budget |
| 370 | ): |
| 371 | break |
| 372 | for n_llm_w in range(1, budget): |
| 373 | if max_wh_ratio is not None and n_llm_w > n_llm_h * max_wh_ratio: |
| 374 | break |
| 375 | num_tokens = grid_tokens( |
| 376 | n_llm_h * step, n_llm_w * step, patch_size, downsample_ratio |
| 377 | )[2] |
| 378 | if num_tokens > budget: |
| 379 | break |
| 380 | key = (n_llm_h * n_llm_w, num_tokens) |
| 381 | if key > best_key: |
| 382 | best_key, best_grid = key, (n_llm_h, n_llm_w) |
| 383 | n_llm_h, n_llm_w = best_grid |
| 384 | return ImageSize(width=n_llm_w * step, height=n_llm_h * step) |
| 385 | |
| 386 | |
| 387 | class DeepseekV4VLDummyInputsBuilder( |
| 388 | BaseDummyInputsBuilder[DeepseekV4VLProcessingInfo] |
| 389 | ): |
| 390 | def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: |
| 391 | return IMAGE_PLACEHOLDER * mm_counts.get("image", 0) |
| 392 | |
| 393 | def get_dummy_mm_data( |
| 394 | self, |
| 395 | seq_len: int, |
| 396 | mm_counts: Mapping[str, int], |
| 397 | mm_options: MultiModalDummyOptions, |
| 398 | ) -> MultiModalDataDict: |
| 399 | size = self.info.get_image_size_with_most_features() |
| 400 | return { |
| 401 | "image": self._get_dummy_images( |
| 402 | width=size.width, |
| 403 | height=size.height, |
| 404 | num_images=mm_counts.get("image", 0), |
| 405 | overrides=mm_options.get("image"), |
| 406 | ), |
| 407 | } |
| 408 | |
| 409 | |
| 410 | class DeepseekV4VLMultiModalProcessor( |
| 411 | BaseMultiModalProcessor[DeepseekV4VLProcessingInfo] |
| 412 | ): |
| 413 | def _get_mm_fields_config( |
| 414 | self, |
| 415 | hf_inputs: BatchFeature, |
| 416 | hf_processor_mm_kwargs: Mapping[str, object], |
| 417 | ) -> Mapping[str, MultiModalFieldConfig]: |
| 418 | vit_grid = hf_inputs.get("vit_grid") |
| 419 | llm_grid = hf_inputs.get("llm_grid") |
| 420 | |
| 421 | if vit_grid is None or llm_grid is None: |
| 422 | empty = torch.empty(0, dtype=torch.long) |
| 423 | patch_sizes = perm_sizes = types_sizes = empty |
| 424 | else: |
| 425 | patch_sizes = vit_grid.prod(-1) |
| 426 | perm_sizes = llm_grid.prod(-1) |
| 427 | n_llm_h, n_llm_w = llm_grid[:, 0], llm_grid[:, 1] |
| 428 | # Pad-free block length; same formula as ``grid_tokens`` given |
| 429 | # the LLM grid. |
| 430 | types_sizes = ( |
| 431 | n_llm_h * (n_llm_w + 1) |
| 432 | + 2 |
| 433 | + (n_llm_h % 2) * (n_llm_w + 1) |
| 434 | + (n_llm_h + 1) // 2 * (n_llm_w + 1) % 2 * 2 |
| 435 | ) |
| 436 | |
| 437 | return { |
| 438 | "patches": MultiModalFieldConfig.flat_from_sizes("image", patch_sizes), |
| 439 | "vit_grid": MultiModalFieldConfig.batched("image", keep_on_cpu=True), |
| 440 | "llm_grid": MultiModalFieldConfig.batched("image", keep_on_cpu=True), |
| 441 | "perm": MultiModalFieldConfig.flat_from_sizes("image", perm_sizes), |
| 442 | "types": MultiModalFieldConfig.flat_from_sizes( |
| 443 | "image", types_sizes, keep_on_cpu=True |
| 444 | ), |
| 445 | } |
| 446 | |
| 447 | def _get_prompt_updates( |
| 448 | self, |
| 449 | mm_items: MultiModalDataItems, |
| 450 | hf_processor_mm_kwargs: Mapping[str, object], |
| 451 | out_mm_kwargs: MultiModalKwargsItems, |
| 452 | ) -> Sequence[PromptUpdate]: |
| 453 | image_token_id = self.info.get_image_placeholder_token_id() |
| 454 | validate_image_sentinel_ids(self.info.get_tokenizer()) |
| 455 | image_embed_id = IMAGE_SENTINEL_BASE_ID + IMAGE |
| 456 | |
| 457 | def get_image_replacement(item_idx: int) -> PromptUpdateDetails: |
| 458 | types: torch.Tensor = out_mm_kwargs["image"][item_idx]["types"].data |
| 459 | full = (IMAGE_SENTINEL_BASE_ID + types).tolist() |
| 460 | return PromptUpdateDetails.select_token_id(full, image_embed_id) |
| 461 | |
| 462 | return [ |
| 463 | PromptReplacement( |
| 464 | modality="image", |
| 465 | target=[image_token_id], |
| 466 | replacement=get_image_replacement, |
| 467 | ), |
| 468 | ] |
| 469 | |
| 470 | def _apply_token_matches_with_placeholders( |
| 471 | self, |
| 472 | token_ids: list[int], |
| 473 | mm_prompt_updates: MultiModalPromptUpdates, |
| 474 | ) -> tuple[ |
| 475 | list[int], |
| 476 | MultiModalPromptUpdatesApplyResult, |
| 477 | Mapping[str, list[PlaceholderFeaturesInfo]], |
| 478 | ]: |
| 479 | # Same as the base implementation, except that each image block gets |
| 480 | # its compressor-alignment pad (``3 - start_pos % 4`` IMAGE_PAD |
| 481 | # sentinels) prepended while splicing: the pad depends on the block's |
| 482 | # final position in the prompt, which is unknown when the (cacheable) |
| 483 | # prompt updates are built. |
| 484 | matched_updates, result = _plan_prompt_updates(token_ids, mm_prompt_updates) |
| 485 | placeholders: dict[str, list[PlaceholderFeaturesInfo]] = { |
| 486 | modality: [] for modality in mm_prompt_updates |
| 487 | } |
| 488 | |
| 489 | pad_id = IMAGE_SENTINEL_BASE_ID + IMAGE_PAD |
| 490 | |
| 491 | new_token_ids = list[int]() |
| 492 | prev_end_idx = 0 |
| 493 | for matched_update in matched_updates: |
| 494 | update = matched_update.update |
| 495 | match = matched_update.match |
| 496 | |
| 497 | if update.mode == UpdateMode.INSERT: |
| 498 | end_idx_to_insert = match.end_idx |
| 499 | elif update.mode == UpdateMode.REPLACE: |
| 500 | end_idx_to_insert = match.start_idx |
| 501 | else: |
| 502 | assert_never(update.mode) |
| 503 | |
| 504 | new_token_ids.extend(token_ids[prev_end_idx:end_idx_to_insert]) |
| 505 | start_idx = len(new_token_ids) |
| 506 | |
| 507 | tokens = list(update.content.full) |
| 508 | if tokens and update.modality == "image": |
| 509 | compress_pad = COMPRESS_PAD_TO - 1 - start_idx % COMPRESS_PAD_TO |
| 510 | tokens = [pad_id] * compress_pad + tokens |
| 511 | |
| 512 | if tokens: |
| 513 | content_is_embed = update.content.is_embed |
| 514 | is_embed = ( |
| 515 | content_is_embed(tokens) if content_is_embed is not None else None |
| 516 | ) |
| 517 | placeholders[update.modality].append( |
| 518 | PlaceholderFeaturesInfo( |
| 519 | modality=update.modality, |
| 520 | item_idx=update.item_idx, |
| 521 | start_idx=start_idx, |
| 522 | tokens=tokens, |
| 523 | is_embed=is_embed, |
| 524 | ) |
| 525 | ) |
| 526 | new_token_ids.extend(tokens) |
| 527 | |
| 528 | prev_end_idx = match.end_idx |
| 529 | |
| 530 | new_token_ids.extend(token_ids[prev_end_idx:]) |
| 531 | return new_token_ids, result, placeholders |