| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import math |
| 4 | from typing import Any, TypedDict |
| 5 | |
| 6 | import numpy as np |
| 7 | import torch |
| 8 | import torch.nn.functional as F |
| 9 | from PIL import Image |
| 10 | from transformers import ( |
| 11 | BatchFeature, |
| 12 | ImageProcessingMixin, |
| 13 | ProcessorMixin, |
| 14 | TensorType, |
| 15 | ) |
| 16 | from transformers.processing_utils import ProcessingKwargs |
| 17 | from typing_extensions import Unpack |
| 18 | |
| 19 | from vllm.tokenizers.hf import HfTokenizer |
| 20 | |
| 21 | MAX_PIXELS = 60_000_000 # 60-megapixel ceiling ≈ 8200 × 7300 px |
| 22 | |
| 23 | # Vision preprocessing constants |
| 24 | VISION_MEAN = (0.5, 0.5, 0.5) |
| 25 | VISION_STD = (0.5, 0.5, 0.5) |
| 26 | VISION_SCALE = 1 / 255 |
| 27 | |
| 28 | |
| 29 | def _make_writeable(arr: np.ndarray) -> np.ndarray: |
| 30 | """Return *arr* itself if it is already writeable, otherwise try to flip the |
| 31 | write flag in-place and finally fall back to `arr.copy()`. |
| 32 | This guarantees the buffer handed to `torch.from_numpy()` is always |
| 33 | writeable, silencing the PyTorch warning about undefined behaviour. |
| 34 | """ |
| 35 | if arr.flags.writeable: |
| 36 | return arr |
| 37 | |
| 38 | # First, try the cheap path — in-place flag toggle (works for mmap'd arrays |
| 39 | # and some shared memory buffers): |
| 40 | try: |
| 41 | arr.setflags(write=True) |
| 42 | return arr # success: no data copy |
| 43 | except ValueError: |
| 44 | # Buffer is inherently read-only (e.g. backed by PyAV / PIL): make copy |
| 45 | return arr.copy() |
| 46 | |
| 47 | |
| 48 | def extract_image_pil(image: Image.Image) -> torch.Tensor: |
| 49 | if image.width * image.height > MAX_PIXELS: |
| 50 | raise ValueError( |
| 51 | f"Image (w={image.width}, h={image.height}) > MAX=`{MAX_PIXELS}`" |
| 52 | ) |
| 53 | img = image if image.mode == "RGB" else image.convert("RGB") |
| 54 | arr = np.asarray(img) |
| 55 | arr = _make_writeable(arr) |
| 56 | return torch.from_numpy(arr) |
| 57 | |
| 58 | |
| 59 | def get_image_size_for_max_num_patches( |
| 60 | image_height: int, |
| 61 | image_width: int, |
| 62 | patch_size: int, |
| 63 | max_num_patches: int, |
| 64 | min_num_patches: int | None = None, |
| 65 | eps: float = 1e-5, |
| 66 | pixel_shuffle_scale: int = 1, |
| 67 | ) -> tuple[int, int]: |
| 68 | r"""Compute a target resolution whose patch grid satisfies patching parametrization. |
| 69 | |
| 70 | Args: |
| 71 | image_height (`int`): |
| 72 | Height in pixels of the source image prior to any resizing. |
| 73 | image_width (`int`): |
| 74 | Width in pixels of the source image prior to any resizing. |
| 75 | patch_size (`int`): |
| 76 | Size of the square patch used by the vision encoder. |
| 77 | max_num_patches (`int`): |
| 78 | Upper bound on `(height / patch_size) * (width / patch_size)` after |
| 79 | resizing. |
| 80 | min_num_patches (`int`, *optional*): |
| 81 | Lower bound on the number of patches. When provided the image will |
| 82 | be scaled up if necessary. |
| 83 | eps (`float`, *optional*, defaults to 1e-5): |
| 84 | Convergence tolerance for the internal binary search to determine |
| 85 | the target dimensions. |
| 86 | pixel_shuffle_scale (`int`, *optional*, defaults to 1): |
| 87 | Additional stride multiplier applied when pixel shuffle later |
| 88 | reduces spatial resolution. |
| 89 | |
| 90 | Returns: |
| 91 | `tuple[int, int]`: Height and width (in pixels) that are multiples of |
| 92 | `patch_size * pixel_shuffle_scale` and respect both the maximum and |
| 93 | optional minimum patch-count constraints. |
| 94 | |
| 95 | """ |
| 96 | |
| 97 | def get_scaled_image_size(scale, original_size, patch_size, pixel_shuffle_scale): |
| 98 | scaled_size = scale * original_size |
| 99 | divisor = patch_size * pixel_shuffle_scale |
| 100 | scaled_size = math.ceil(scaled_size / divisor) * divisor |
| 101 | scaled_size = max(divisor, scaled_size) |
| 102 | return int(scaled_size) |
| 103 | |
| 104 | # Ensure divisibility |
| 105 | divisor = patch_size * pixel_shuffle_scale |
| 106 | adjusted_height = math.ceil(image_height / divisor) * divisor |
| 107 | adjusted_height = max(divisor, adjusted_height) |
| 108 | adjusted_width = math.ceil(image_width / divisor) * divisor |
| 109 | adjusted_width = max(divisor, adjusted_width) |
| 110 | |
| 111 | num_patches = (adjusted_height / patch_size) * (adjusted_width / patch_size) |
| 112 | |
| 113 | if min_num_patches is not None and num_patches < min_num_patches: |
| 114 | # Scale up |
| 115 | scale_min, scale_max = 1.0, 100.0 |
| 116 | while (scale_max - scale_min) >= eps: |
| 117 | scale = (scale_min + scale_max) / 2 |
| 118 | target_height = get_scaled_image_size( |
| 119 | scale, image_height, patch_size, pixel_shuffle_scale |
| 120 | ) |
| 121 | target_width = get_scaled_image_size( |
| 122 | scale, image_width, patch_size, pixel_shuffle_scale |
| 123 | ) |
| 124 | num_patches = (target_height / patch_size) * (target_width / patch_size) |
| 125 | if num_patches >= min_num_patches: |
| 126 | scale_max = scale |
| 127 | else: |
| 128 | scale_min = scale |
| 129 | scale = scale_max |
| 130 | target_height = get_scaled_image_size( |
| 131 | scale, image_height, patch_size, pixel_shuffle_scale |
| 132 | ) |
| 133 | target_width = get_scaled_image_size( |
| 134 | scale, image_width, patch_size, pixel_shuffle_scale |
| 135 | ) |
| 136 | return target_height, target_width |
| 137 | elif num_patches <= max_num_patches: |
| 138 | return adjusted_height, adjusted_width |
| 139 | else: |
| 140 | # Scale down |
| 141 | scale_min, scale_max = eps / 10, 1.0 |
| 142 | while (scale_max - scale_min) >= eps: |
| 143 | scale = (scale_min + scale_max) / 2 |
| 144 | target_height = get_scaled_image_size( |
| 145 | scale, image_height, patch_size, pixel_shuffle_scale |
| 146 | ) |
| 147 | target_width = get_scaled_image_size( |
| 148 | scale, image_width, patch_size, pixel_shuffle_scale |
| 149 | ) |
| 150 | num_patches = (target_height / patch_size) * (target_width / patch_size) |
| 151 | if num_patches <= max_num_patches: |
| 152 | scale_min = scale |
| 153 | else: |
| 154 | scale_max = scale |
| 155 | scale = scale_min |
| 156 | target_height = get_scaled_image_size( |
| 157 | scale, image_height, patch_size, pixel_shuffle_scale |
| 158 | ) |
| 159 | target_width = get_scaled_image_size( |
| 160 | scale, image_width, patch_size, pixel_shuffle_scale |
| 161 | ) |
| 162 | return target_height, target_width |
| 163 | |
| 164 | |
| 165 | _MEAN_TENSOR = torch.tensor(VISION_MEAN, dtype=torch.float32).view(1, 1, 1, -1) |
| 166 | _STD_TENSOR = torch.tensor(VISION_STD, dtype=torch.float32).view(1, 1, 1, -1) |
| 167 | |
| 168 | |
| 169 | def prepare_image_tensor( |
| 170 | image: torch.Tensor, |
| 171 | scale: float = VISION_SCALE, |
| 172 | ) -> torch.Tensor: |
| 173 | r"""Standardize RGB images prior to patch extraction via rescaling and whitening. |
| 174 | |
| 175 | Args: |
| 176 | image (`torch.Tensor`): |
| 177 | Tensor with shape `(..., height, width, 3)` containing RGB values. |
| 178 | The tensor is converted to floating point if needed. |
| 179 | scale (`float`, *optional*, defaults to `VISION_SCALE`): |
| 180 | Scalar multiplier applied before normalization. |
| 181 | |
| 182 | Returns: |
| 183 | `torch.Tensor`: Normalized tensor with the same shape as the input and |
| 184 | dtype `torch.float32`. |
| 185 | |
| 186 | """ |
| 187 | if not torch.is_floating_point(image): |
| 188 | image = image.float() |
| 189 | rescaled = image * scale |
| 190 | |
| 191 | # Use precomputed tensors and move to the correct device if needed |
| 192 | mean_tensor = _MEAN_TENSOR.to(image.device) |
| 193 | std_tensor = _STD_TENSOR.to(image.device) |
| 194 | |
| 195 | normalized = (rescaled - mean_tensor) / std_tensor |
| 196 | return normalized |
| 197 | |
| 198 | |
| 199 | def patchify_vision(image: torch.Tensor, patch_size: int) -> torch.Tensor: |
| 200 | r"""Convert normalized images into flattened ViT-style patches. |
| 201 | |
| 202 | Args: |
| 203 | image (`torch.Tensor`): |
| 204 | Tensor of shape `(num_images, height, width, channels)`. |
| 205 | patch_size (`int`): |
| 206 | Edge length of the square patches |
| 207 | |
| 208 | Returns: |
| 209 | `torch.Tensor`: |
| 210 | Patch tensor where each position stores the flattened pixels |
| 211 | belonging to that patch. |
| 212 | |
| 213 | Raises: |
| 214 | ValueError: If `height` or `width` is not divisible by `patch_size`. |
| 215 | |
| 216 | """ |
| 217 | num_images, height, width, channels = image.shape |
| 218 | if height % patch_size or width % patch_size: |
| 219 | raise ValueError( |
| 220 | "Dimensions of images " |
| 221 | f"{image.shape} are not divisible by patch_size={patch_size}." |
| 222 | ) |
| 223 | patches = image.reshape( |
| 224 | num_images, |
| 225 | height // patch_size, |
| 226 | patch_size, |
| 227 | width // patch_size, |
| 228 | patch_size, |
| 229 | channels, |
| 230 | ) |
| 231 | patches = patches.permute(0, 1, 3, 2, 4, 5) |
| 232 | patches = patches.reshape( |
| 233 | num_images, |
| 234 | height // patch_size, |
| 235 | width // patch_size, |
| 236 | channels * patch_size * patch_size, |
| 237 | ) |
| 238 | return patches |
| 239 | |
| 240 | |
| 241 | def process_vision_for_patches( |
| 242 | images: torch.Tensor, |
| 243 | patch_size: int, |
| 244 | max_num_patches: int, |
| 245 | min_num_patches: int | None = None, |
| 246 | pixel_shuffle_scale: int = 1, |
| 247 | ) -> tuple[torch.Tensor, list[int]]: |
| 248 | r"""Resize, normalize, and patchify RGB images for the vision encoder. |
| 249 | |
| 250 | Args: |
| 251 | images (`torch.Tensor`): |
| 252 | Either `(height, width, channels)` for a single image or |
| 253 | `(num_images, height, width, channels)` for a batch. Channels are |
| 254 | expected to be RGB. |
| 255 | patch_size (`int`): |
| 256 | Edge length of square patches; implicitly controls resize grid granularity. |
| 257 | max_num_patches (`int`): |
| 258 | Maximum number of patches allowed after resizing. |
| 259 | min_num_patches (`int`, *optional*): |
| 260 | Minimum number of patches. If provided, the routine upsamples images |
| 261 | as needed to satisfy the lower bound. |
| 262 | pixel_shuffle_scale (`int`, *optional*, defaults to 1): |
| 263 | Pixel shuffle scale factor; influences the target grid that the |
| 264 | function produces. |
| 265 | |
| 266 | Returns: |
| 267 | `tuple[torch.Tensor, list[int]]`: A pair `(patches, dims_virtual)` |
| 268 | where `patches` has shape `(num_images, target_h / patch_size, target_w |
| 269 | / patch_size, channels * patch_size**2)` and `dims_virtual` encodes |
| 270 | effective `(images, height, width)` dimensions after optional pixel |
| 271 | shuffling. |
| 272 | |
| 273 | """ |
| 274 | # Add batch dim if single image |
| 275 | if images.dim() == 3: |
| 276 | images = images.unsqueeze(0) |
| 277 | |
| 278 | # Permute to channel first for resize |
| 279 | images = images.permute(0, 3, 1, 2) |
| 280 | |
| 281 | # Get target dimensions |
| 282 | _, _, orig_height, orig_width = images.shape |
| 283 | target_height, target_width = get_image_size_for_max_num_patches( |
| 284 | orig_height, |
| 285 | orig_width, |
| 286 | patch_size, |
| 287 | max_num_patches, |
| 288 | min_num_patches=min_num_patches, |
| 289 | pixel_shuffle_scale=pixel_shuffle_scale, |
| 290 | ) |
| 291 | |
| 292 | # Resize |
| 293 | images = F.interpolate( |
| 294 | images, |
| 295 | size=(target_height, target_width), |
| 296 | mode="bilinear", |
| 297 | align_corners=False, |
| 298 | ) |
| 299 | |
| 300 | # Back to channel last |
| 301 | images = images.permute(0, 2, 3, 1) |
| 302 | |
| 303 | # Normalize |
| 304 | images = prepare_image_tensor(images) |
| 305 | |
| 306 | # Patchify |
| 307 | patches = patchify_vision(images, patch_size=patch_size) |
| 308 | |
| 309 | # Calculate dimensions for the patches |
| 310 | n_images, h_patches, w_patches, _ = patches.shape |
| 311 | dims_virtual = ( |
| 312 | [1, h_patches, w_patches] |
| 313 | if pixel_shuffle_scale == 1 |
| 314 | else [1, h_patches // pixel_shuffle_scale, w_patches // pixel_shuffle_scale] |
| 315 | ) |
| 316 | |
| 317 | return patches, dims_virtual |
| 318 | |
| 319 | |
| 320 | class IsaacImagesKwargs(TypedDict, total=False): |
| 321 | patch_size: int |
| 322 | max_num_patches: int |
| 323 | min_num_patches: int |
| 324 | pixel_shuffle_scale: int |
| 325 | |
| 326 | |
| 327 | class IsaacProcessorKwargs(ProcessingKwargs, total=False): # type: ignore[call-arg] |
| 328 | images_kwargs: IsaacImagesKwargs |
| 329 | _defaults = { |
| 330 | "text_kwargs": {"padding": False}, |
| 331 | "images_kwargs": {}, |
| 332 | } |
| 333 | |
| 334 | |
| 335 | class IsaacImageProcessor(ImageProcessingMixin): |
| 336 | model_input_names = ["pixel_values", "image_grid_thw"] |
| 337 | |
| 338 | def __init__( |
| 339 | self, |
| 340 | patch_size: int = 16, |
| 341 | vision_max_num_patches: int = 6144, |
| 342 | vision_min_num_patches: int = 256, |
| 343 | pixel_shuffle_scale: int = 2, |
| 344 | ) -> None: |
| 345 | self.patch_size = patch_size |
| 346 | self.vision_max_num_patches = vision_max_num_patches |
| 347 | self.vision_min_num_patches = vision_min_num_patches |
| 348 | self.pixel_shuffle_scale = pixel_shuffle_scale |
| 349 | |
| 350 | def __call__( |
| 351 | self, |
| 352 | images: Image.Image | list[Image.Image], |
| 353 | return_tensors: str | TensorType | None = None, |
| 354 | **kwargs: Unpack[IsaacImagesKwargs], |
| 355 | ) -> BatchFeature: |
| 356 | """Preprocess images into format compatible with vLLM input processing.""" |
| 357 | if not isinstance(images, list): |
| 358 | images = [images] |
| 359 | |
| 360 | all_pixel_values: list[torch.Tensor] = [] |
| 361 | all_image_grids: list[torch.Tensor] = [] |
| 362 | |
| 363 | for image in images: |
| 364 | image_tensor = extract_image_pil(image) |
| 365 | |
| 366 | patches, dims_virtual = process_vision_for_patches( |
| 367 | image_tensor, |
| 368 | patch_size=kwargs.get("patch_size", self.patch_size), |
| 369 | max_num_patches=kwargs.get( |
| 370 | "max_num_patches", self.vision_max_num_patches |
| 371 | ), |
| 372 | min_num_patches=kwargs.get( |
| 373 | "min_num_patches", self.vision_min_num_patches |
| 374 | ), |
| 375 | pixel_shuffle_scale=kwargs.get( |
| 376 | "pixel_shuffle_scale", self.pixel_shuffle_scale |
| 377 | ), |
| 378 | ) |
| 379 | |
| 380 | # Isaac packs a dummy temporal dim for images |
| 381 | patches = patches.unsqueeze(1) # [N, T=1, Hp, Wp, D] |
| 382 | |
| 383 | hp, wp, dim = patches.shape[-3], patches.shape[-2], patches.shape[-1] |
| 384 | current_num_patches = hp * wp |
| 385 | pixel_values = patches.reshape(current_num_patches, dim) # [N_tokens, D] |
| 386 | |
| 387 | # Use real patch dimensions for image_grid_thw, not virtual dimensions |
| 388 | # This ensures the vision model receives correct grid info for pixel shuffle |
| 389 | dims_real = [1, hp, wp] # Real patch dimensions |
| 390 | image_grid_thw = torch.tensor(dims_real).unsqueeze(0) |
| 391 | |
| 392 | all_pixel_values.append(pixel_values) |
| 393 | all_image_grids.append(image_grid_thw) |
| 394 | |
| 395 | if all_pixel_values: |
| 396 | final_pixel_values = torch.cat(all_pixel_values, dim=0) |
| 397 | final_image_grids = torch.cat(all_image_grids, dim=0) |
| 398 | else: |
| 399 | final_pixel_values = torch.empty(0, 0) |
| 400 | final_image_grids = torch.empty(0, 3) |
| 401 | |
| 402 | return BatchFeature( |
| 403 | data={ |
| 404 | "pixel_values": final_pixel_values, |
| 405 | "image_grid_thw": final_image_grids, |
| 406 | }, |
| 407 | tensor_type=return_tensors, |
| 408 | ) |
| 409 | |
| 410 | |
| 411 | class IsaacProcessor(ProcessorMixin): |
| 412 | attributes = ["image_processor", "tokenizer"] |
| 413 | |
| 414 | def __init__( |
| 415 | self, |
| 416 | image_processor: IsaacImageProcessor, |
| 417 | tokenizer: HfTokenizer, |
| 418 | image_token: str = "<image>", |
| 419 | ): |
| 420 | self.image_processor = image_processor |
| 421 | self.tokenizer = tokenizer |
| 422 | |
| 423 | self.image_token = image_token |
| 424 | |
| 425 | def __call__( |
| 426 | self, |
| 427 | text: str | list[str] | None = None, |
| 428 | images: Image.Image | list[Image.Image] | None = None, |
| 429 | return_tensors: str | TensorType | None = None, |
| 430 | **kwargs: Unpack[IsaacProcessorKwargs], # type: ignore[misc] |
| 431 | ) -> BatchFeature: |
| 432 | output_kwargs = self._merge_kwargs( |
| 433 | IsaacProcessorKwargs, |
| 434 | tokenizer_init_kwargs=self.tokenizer.init_kwargs, |
| 435 | **kwargs, |
| 436 | ) |
| 437 | |
| 438 | if images is not None: |
| 439 | image_inputs = self.image_processor( |
| 440 | images, **output_kwargs["images_kwargs"] |
| 441 | ) |
| 442 | image_grid_thw = image_inputs["image_grid_thw"] |
| 443 | else: |
| 444 | image_inputs = {} |
| 445 | image_grid_thw = [] |
| 446 | |
| 447 | if text is not None: |
| 448 | if not isinstance(text, list): |
| 449 | text = [text] |
| 450 | |
| 451 | if image_inputs: |
| 452 | text = text.copy() # below lines change text in-place |
| 453 | merge_length = self.image_processor.pixel_shuffle_scale**2 |
| 454 | index = 0 |
| 455 | for i in range(len(text)): |
| 456 | while self.image_token in text[i]: |
| 457 | num_image_tokens = image_grid_thw[index].prod() // merge_length |
| 458 | text[i] = text[i].replace( |
| 459 | self.image_token, "<|placeholder|>" * num_image_tokens, 1 |
| 460 | ) |
| 461 | index += 1 |
| 462 | text[i] = text[i].replace("<|placeholder|>", "<|image_pad|>") |
| 463 | |
| 464 | text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"]) |
| 465 | else: |
| 466 | text_inputs = {} |
| 467 | |
| 468 | return BatchFeature( |
| 469 | data={**text_inputs, **image_inputs}, |
| 470 | tensor_type=return_tensors, |
| 471 | ) |
| 472 | |
| 473 | def apply_chat_template( |
| 474 | self, |
| 475 | messages: list[dict[str, Any]], |
| 476 | tokenize: bool = False, |
| 477 | add_generation_prompt: bool = False, |
| 478 | **kwargs, |
| 479 | ) -> Any: |
| 480 | # Convert mixed content messages to simple text format |
| 481 | processed_messages = [] |
| 482 | |
| 483 | for message in messages: |
| 484 | if "content" in message and isinstance(message["content"], list): |
| 485 | # Handle mixed content (text + image) |
| 486 | text_parts = [] |
| 487 | for content_item in message["content"]: |
| 488 | if content_item.get("type") == "text": |
| 489 | text_parts.append(content_item.get("text", "")) |
| 490 | elif content_item.get("type") == "image": |
| 491 | # Replace image with vision token |
| 492 | text_parts.append(self.image_token) |
| 493 | |
| 494 | processed_message = { |
| 495 | "role": message.get("role", "user"), |
| 496 | "content": "".join(text_parts), |
| 497 | } |
| 498 | processed_messages.append(processed_message) |
| 499 | else: |
| 500 | # Regular text message |
| 501 | processed_messages.append(message) |
| 502 | |
| 503 | kwargs["return_dict"] = False |
| 504 | return self.tokenizer.apply_chat_template( |
| 505 | processed_messages, |
| 506 | tokenize=tokenize, |
| 507 | add_generation_prompt=add_generation_prompt, |
| 508 | **kwargs, |
| 509 | ) |