| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | # ruff: noqa: E501 |
| 5 | # coding=utf-8 |
| 6 | # adapted from https://github.com/AIDC-AI/Ovis/blob/35ab51a1a1e3542fa6db260a1084cefbc8f164bb/ovis/vllm/processing_ovis.py |
| 7 | # Copyright 2025 The Qwen Team and The HuggingFace Inc. team. All rights reserved. |
| 8 | # |
| 9 | # This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX |
| 10 | # and OPT implementations in this library. It has been modified from its |
| 11 | # original forms to accommodate minor architectural differences compared |
| 12 | # to GPT-NeoX and OPT used by the Meta AI team that trained the model. |
| 13 | # |
| 14 | # Licensed under the Apache License, Version 2.0 (the "License"); |
| 15 | # you may not use this file except in compliance with the License. |
| 16 | # You may obtain a copy of the License at |
| 17 | # |
| 18 | # http://www.apache.org/licenses/LICENSE-2.0 |
| 19 | # |
| 20 | # Unless required by applicable law or agreed to in writing, software |
| 21 | # distributed under the License is distributed on an "AS IS" BASIS, |
| 22 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 23 | # See the License for the specific language governing permissions and |
| 24 | # limitations under the License. |
| 25 | from functools import cached_property |
| 26 | |
| 27 | import PIL |
| 28 | import torch |
| 29 | from transformers import BatchFeature |
| 30 | from transformers.image_utils import ImageInput |
| 31 | from transformers.processing_utils import ProcessingKwargs, ProcessorMixin, Unpack |
| 32 | from transformers.tokenization_utils_base import PreTokenizedInput, TextInput |
| 33 | |
| 34 | from vllm.multimodal.image import convert_image_mode |
| 35 | |
| 36 | __all__ = ["OvisProcessor"] |
| 37 | IGNORE_ID = -100 |
| 38 | |
| 39 | |
| 40 | class OvisProcessorKwargs(ProcessingKwargs, total=False): # type: ignore[call-arg] |
| 41 | _defaults = { |
| 42 | "text_kwargs": { |
| 43 | "padding": False, |
| 44 | }, |
| 45 | "images_kwargs": { |
| 46 | "do_convert_rgb": True, |
| 47 | "return_tensors": "pt", |
| 48 | }, |
| 49 | } |
| 50 | |
| 51 | |
| 52 | class OvisProcessor(ProcessorMixin): |
| 53 | r"""Constructs an Ovis processor which wraps an Ovis image processor and a Qwen2 tokenizer into a single processor. |
| 54 | [`OvisProcessor`] offers all the functionalities of [`Qwen2VLImageProcessor`] and [`Qwen2TokenizerFast`]. See the |
| 55 | [`~OvisProcessor.__call__`] and [`~OvisProcessor.decode`] for more information. |
| 56 | |
| 57 | Args: |
| 58 | image_processor ([`Qwen2VLImageProcessor`], *optional*): |
| 59 | The image processor is a required input. |
| 60 | tokenizer ([`Qwen2TokenizerFast`], *optional*): |
| 61 | The tokenizer is a required input. |
| 62 | chat_template (`str`, *optional*): A Jinja template which will be used to convert lists of messages |
| 63 | in a chat into a tokenizable string. |
| 64 | |
| 65 | """ |
| 66 | |
| 67 | attributes = ["image_processor", "tokenizer"] |
| 68 | valid_kwargs = ["chat_template", "image_pad_token", "image_segment_len"] |
| 69 | |
| 70 | image_processor_class = "AutoImageProcessor" |
| 71 | tokenizer_class = "AutoTokenizer" |
| 72 | |
| 73 | def __init__( |
| 74 | self, |
| 75 | image_processor=None, |
| 76 | tokenizer=None, |
| 77 | chat_template=None, |
| 78 | image_pad_token=None, |
| 79 | image_segment_len=255, |
| 80 | **kwargs, |
| 81 | ): |
| 82 | self.image_token = "<image>" |
| 83 | self.image_pad_token = image_pad_token |
| 84 | self.image_segment_len = image_segment_len |
| 85 | super().__init__(image_processor, tokenizer, chat_template=chat_template) |
| 86 | |
| 87 | @cached_property |
| 88 | def extra_special_tokens(self): |
| 89 | image_pad_token_id = self.tokenizer.get_vocab()[self.image_pad_token] |
| 90 | extra_special_tokens = { |
| 91 | "image_token": -200, |
| 92 | "image_atom": -300, |
| 93 | "image_start": -301, |
| 94 | "image_prefix": -302, |
| 95 | "image_col_sep": -303, |
| 96 | "image_row_sep": -304, |
| 97 | "image_end": -305, |
| 98 | "image_pad": image_pad_token_id, |
| 99 | } |
| 100 | return extra_special_tokens |
| 101 | |
| 102 | def __call__( |
| 103 | self, |
| 104 | images: ImageInput = None, |
| 105 | text: TextInput |
| 106 | | PreTokenizedInput |
| 107 | | list[TextInput] |
| 108 | | list[PreTokenizedInput] = None, |
| 109 | **kwargs: Unpack[OvisProcessorKwargs], |
| 110 | ) -> BatchFeature: |
| 111 | """Main method to prepare for the model one or several sequences(s) and image(s). This method forwards the `text` |
| 112 | and `kwargs` arguments to Qwen2TokenizerFast's [`~Qwen2TokenizerFast.__call__`] if `text` is not `None` to encode |
| 113 | the text. To prepare the vision inputs, this method forwards the `vision_infos` and `kwargs` arguments to |
| 114 | Qwen2VLImageProcessor's [`~Qwen2VLImageProcessor.__call__`] if `vision_infos` is not `None`. |
| 115 | |
| 116 | Args: |
| 117 | images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `list[PIL.Image.Image]`, `list[np.ndarray]`, `list[torch.Tensor]`): |
| 118 | The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch |
| 119 | tensor. Both channels-first and channels-last formats are supported. |
| 120 | text (`str`, `list[str]`, `list[list[str]]`): |
| 121 | The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings |
| 122 | (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set |
| 123 | `is_split_into_words=True` (to lift the ambiguity with a batch of sequences). |
| 124 | videos (`np.ndarray`, `torch.Tensor`, `list[np.ndarray]`, `list[torch.Tensor]`): |
| 125 | The image or batch of videos to be prepared. Each video can be a 4D NumPy array or PyTorch |
| 126 | tensor, or a nested list of 3D frames. Both channels-first and channels-last formats are supported. |
| 127 | return_tensors (`str` or [`~utils.TensorType`], *optional*): |
| 128 | If set, will return tensors of a particular framework. Acceptable values are: |
| 129 | - `'tf'`: Return TensorFlow `tf.constant` objects. |
| 130 | - `'pt'`: Return PyTorch `torch.Tensor` objects. |
| 131 | - `'np'`: Return NumPy `np.ndarray` objects. |
| 132 | - `'jax'`: Return JAX `jnp.ndarray` objects. |
| 133 | |
| 134 | Returns: |
| 135 | [`BatchFeature`]: A [`BatchFeature`] with the following fields: |
| 136 | - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`. |
| 137 | - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when |
| 138 | `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not |
| 139 | `None`). |
| 140 | - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`. |
| 141 | - **pixel_values_videos** -- Pixel values of videos to be fed to a model. Returned when `videos` is not `None`. |
| 142 | - **image_grid_thw** -- List of image 3D grid in LLM. Returned when `images` is not `None`. |
| 143 | - **video_grid_thw** -- List of video 3D grid in LLM. Returned when `videos` is not `None`. |
| 144 | - **second_per_grid_ts** -- List of video seconds per time grid. Returned when `videos` is not `None`. |
| 145 | |
| 146 | """ |
| 147 | max_partition = kwargs.pop("max_partition", 9) |
| 148 | covering_threshold = kwargs.pop("covering_threshold", 0.9) |
| 149 | |
| 150 | output_kwargs = self._merge_kwargs( |
| 151 | OvisProcessorKwargs, |
| 152 | tokenizer_init_kwargs=self.tokenizer.init_kwargs, |
| 153 | **kwargs, |
| 154 | ) |
| 155 | |
| 156 | # Process all images first |
| 157 | image_features = {} |
| 158 | if images is not None: |
| 159 | processed_images = [] |
| 160 | image_placeholders_list = [] |
| 161 | grids = [] |
| 162 | |
| 163 | # Process each image |
| 164 | for image in images if isinstance(images, list) else [images]: |
| 165 | pixel_values, image_placeholders, grid = self.preprocess_image( |
| 166 | image=image, |
| 167 | max_partition=max_partition, |
| 168 | covering_threshold=covering_threshold, |
| 169 | **output_kwargs["images_kwargs"], |
| 170 | ) |
| 171 | processed_images.append(pixel_values) |
| 172 | image_placeholders_list.append(image_placeholders) |
| 173 | grids.append(grid) |
| 174 | |
| 175 | # assign all processed images |
| 176 | if processed_images: |
| 177 | image_features["image_placeholders"] = image_placeholders_list |
| 178 | |
| 179 | # Process text input |
| 180 | if text is not None: |
| 181 | if not isinstance(text, list): |
| 182 | text = [text] |
| 183 | |
| 184 | tokenized_batched_text = self._tokenize_with_image_symbol(text) |
| 185 | image_token_id = self.get_token_value("image_token") |
| 186 | replaced_ids_list = [] |
| 187 | idx = 0 |
| 188 | for ids_tensor in tokenized_batched_text: |
| 189 | if ( |
| 190 | image_token_id in ids_tensor |
| 191 | and "image_placeholders" in image_features |
| 192 | ): |
| 193 | if idx < len(image_features["image_placeholders"]): |
| 194 | # Converts in list for ease of use |
| 195 | ids_list = ids_tensor.tolist() |
| 196 | |
| 197 | new_ids = [] |
| 198 | |
| 199 | # replace placeholders |
| 200 | for i, token_id in enumerate(ids_list): |
| 201 | if token_id == image_token_id: |
| 202 | placeholder_ids = image_features["image_placeholders"][ |
| 203 | idx |
| 204 | ] |
| 205 | new_ids.extend(placeholder_ids) |
| 206 | idx += 1 |
| 207 | else: |
| 208 | new_ids.append(token_id) |
| 209 | |
| 210 | # Converts back to tensors |
| 211 | ids_tensor = torch.tensor(new_ids, dtype=torch.long) |
| 212 | else: |
| 213 | raise RuntimeError( |
| 214 | "Mismatch between the images you provided and the number of placeholder present in the text" |
| 215 | ) |
| 216 | |
| 217 | replaced_ids_list.append(ids_tensor) |
| 218 | |
| 219 | if replaced_ids_list: |
| 220 | replaced_and_tokenized_ids = torch.stack(replaced_ids_list) |
| 221 | else: |
| 222 | replaced_and_tokenized_ids = torch.tensor([], dtype=torch.long) |
| 223 | |
| 224 | # Create the output with text features |
| 225 | output = BatchFeature( |
| 226 | data={ |
| 227 | "input_ids": replaced_and_tokenized_ids, |
| 228 | } |
| 229 | ) |
| 230 | |
| 231 | # Add image features if present |
| 232 | if image_features: |
| 233 | output["pixel_values"] = processed_images |
| 234 | output["grids"] = grids |
| 235 | |
| 236 | return output |
| 237 | |
| 238 | # If only images were provided |
| 239 | return BatchFeature(data=image_features) |
| 240 | |
| 241 | def _tokenize_with_image_symbol(self, text_list: list[str]) -> torch.LongTensor: |
| 242 | batch_token_ids = [] |
| 243 | for text in text_list: |
| 244 | text_chunks = [ |
| 245 | self.tokenizer(chunk, add_special_tokens=False).input_ids |
| 246 | for chunk in text.split(self.image_token) |
| 247 | ] |
| 248 | token_ids = [] |
| 249 | num_chuck = len(text_chunks) |
| 250 | for i, chunk in enumerate(text_chunks): |
| 251 | token_ids.extend(chunk) |
| 252 | if i < num_chuck - 1: |
| 253 | token_ids.append(self.get_token_value("image_token")) |
| 254 | batch_token_ids.append(token_ids) |
| 255 | return torch.tensor(batch_token_ids, dtype=torch.long) |
| 256 | |
| 257 | def get_image_size(self): |
| 258 | size = self.image_processor.size |
| 259 | if "shortest_edge" in size: |
| 260 | width = height = size["shortest_edge"] |
| 261 | elif "height" in size and "width" in size: |
| 262 | width = size["width"] |
| 263 | height = size["height"] |
| 264 | else: |
| 265 | raise ValueError("Can't parse image size from image_processor config.") |
| 266 | return height, width |
| 267 | |
| 268 | def get_token_value(self, tok): |
| 269 | return self.extra_special_tokens[tok] |
| 270 | |
| 271 | def construct_image_indicators(self, grid): |
| 272 | image_placeholders = [ |
| 273 | self.get_token_value("image_start"), |
| 274 | self.get_token_value("image_atom"), |
| 275 | self.get_token_value("image_prefix"), |
| 276 | ] |
| 277 | if grid[0] * grid[1] > 1: |
| 278 | for r in range(grid[0]): |
| 279 | for c in range(grid[1]): |
| 280 | image_placeholders.append(self.get_token_value("image_atom")) |
| 281 | if c < grid[1] - 1: |
| 282 | image_placeholders.append(self.get_token_value("image_col_sep")) |
| 283 | if r < grid[0] - 1: |
| 284 | image_placeholders.append(self.get_token_value("image_row_sep")) |
| 285 | image_placeholders.append(self.get_token_value("image_end")) |
| 286 | return image_placeholders |
| 287 | |
| 288 | def construct_image_placeholders(self, grid): |
| 289 | image_placeholders = self.construct_image_indicators(grid) |
| 290 | |
| 291 | image_atom_token_id = self.get_token_value("image_atom") |
| 292 | # Extract the padding token ID from tokenizer |
| 293 | image_padding_token_id = self.get_token_value("image_pad") |
| 294 | |
| 295 | # Create a new list with padding tokens inserted |
| 296 | padded_placeholder_tokens = [] |
| 297 | for token in image_placeholders: |
| 298 | padded_placeholder_tokens.append(image_padding_token_id) |
| 299 | if token == image_atom_token_id: |
| 300 | padded_placeholder_tokens.extend( |
| 301 | [image_padding_token_id] * self.image_segment_len |
| 302 | ) |
| 303 | return padded_placeholder_tokens |
| 304 | |
| 305 | def preprocess_image( |
| 306 | self, |
| 307 | image: PIL.Image.Image, |
| 308 | max_partition, |
| 309 | covering_threshold, |
| 310 | do_convert_rgb, |
| 311 | return_tensors, |
| 312 | ): |
| 313 | def _preprocess(img: PIL.Image.Image, side): |
| 314 | # first resize and preprocess |
| 315 | w, h = img.size |
| 316 | if w == h: |
| 317 | new_width = new_height = side |
| 318 | elif w > h: |
| 319 | new_width = side |
| 320 | new_height = int(h / w * new_width) |
| 321 | else: |
| 322 | new_height = side |
| 323 | new_width = int(w / h * new_height) |
| 324 | new_size = dict(height=new_height, width=new_width) |
| 325 | pixel_values = self.image_processor.preprocess( |
| 326 | img, size=new_size, return_tensors=return_tensors |
| 327 | )["pixel_values"] |
| 328 | |
| 329 | # then pad to square |
| 330 | square_values = torch.zeros( |
| 331 | [1, 3, side, side], dtype=pixel_values.dtype, device=pixel_values.device |
| 332 | ) |
| 333 | new_height, new_width = pixel_values.shape[2:] |
| 334 | if new_height == new_width: |
| 335 | square_values[:, :, :, :] = pixel_values |
| 336 | elif new_height > new_width: |
| 337 | from_index = (side - new_width) // 2 |
| 338 | square_values[:, :, :, from_index : from_index + new_width] = ( |
| 339 | pixel_values |
| 340 | ) |
| 341 | else: |
| 342 | from_index = (side - new_height) // 2 |
| 343 | square_values[:, :, from_index : from_index + new_height, :] = ( |
| 344 | pixel_values |
| 345 | ) |
| 346 | |
| 347 | return square_values |
| 348 | |
| 349 | def _partition(img, grid) -> list[tuple[int, int, int, int]]: |
| 350 | w, h = img.size |
| 351 | row_height = h // grid[0] |
| 352 | col_width = w // grid[1] |
| 353 | |
| 354 | partition = [] |
| 355 | for row in range(grid[0]): |
| 356 | for col in range(grid[1]): |
| 357 | left = col * col_width |
| 358 | upper = row * row_height |
| 359 | right = w if col == grid[1] - 1 else (col + 1) * col_width |
| 360 | lower = h if row == grid[0] - 1 else (row + 1) * row_height |
| 361 | partition.append((left, upper, right, lower)) |
| 362 | |
| 363 | return partition |
| 364 | |
| 365 | def _covering_area(left, upper, right, lower, side): |
| 366 | w = right - left |
| 367 | h = lower - upper |
| 368 | w, h = max(w, h), min(w, h) |
| 369 | if w > side: |
| 370 | h = h / w * side |
| 371 | w = side |
| 372 | return w * h |
| 373 | |
| 374 | def _get_best_grid(img, side): |
| 375 | img_area = img.size[0] * img.size[1] |
| 376 | |
| 377 | candidate_grids = [] |
| 378 | for i in range(1, max_partition + 1): |
| 379 | for j in range(1, max_partition + 1): |
| 380 | if i * j <= max_partition: |
| 381 | candidate_grids.append((i, j)) |
| 382 | |
| 383 | all_grids = [] |
| 384 | good_grids = [] |
| 385 | for grid in candidate_grids: |
| 386 | partition = _partition(img, grid) |
| 387 | covering_ratio = ( |
| 388 | sum([_covering_area(*p, side) for p in partition]) / img_area |
| 389 | ) |
| 390 | assert covering_ratio <= 1.0 |
| 391 | all_grids.append((grid, covering_ratio)) |
| 392 | if covering_ratio > covering_threshold: |
| 393 | good_grids.append((grid, covering_ratio)) |
| 394 | |
| 395 | if len(good_grids) > 0: |
| 396 | # pick the good partition with minimum #sub_images and break the tie using covering_ratio |
| 397 | return sorted(good_grids, key=lambda x: (x[0][0] * x[0][1], -x[1]))[0][ |
| 398 | 0 |
| 399 | ] |
| 400 | else: |
| 401 | # pick the partition with maximum covering_ratio and break the tie using #sub_images |
| 402 | return sorted(all_grids, key=lambda x: (-x[1], x[0][0] * x[0][1]))[0][0] |
| 403 | |
| 404 | if do_convert_rgb: |
| 405 | image = convert_image_mode(image, "RGB") |
| 406 | |
| 407 | sides = self.get_image_size() |
| 408 | if sides[0] != sides[1]: |
| 409 | raise ValueError("get_image_size() returns non-square size") |
| 410 | side = sides[0] |
| 411 | grid = _get_best_grid(image, side) |
| 412 | partition = _partition(image, grid) |
| 413 | crops = [image.crop(p) for p in partition] |
| 414 | if len(crops) > 1: |
| 415 | crops.insert(0, image) |
| 416 | pixel_values = torch.cat([_preprocess(crop, side) for crop in crops], dim=0) |
| 417 | image_placeholders = self.construct_image_placeholders(grid) |
| 418 | return torch.tensor(pixel_values), image_placeholders, torch.tensor(grid) |
| 419 | |
| 420 | def batch_decode(self, *args, **kwargs): |
| 421 | """This method forwards all its arguments to Qwen2TokenizerFast's [`~PythonBackend.batch_decode`]. Please |
| 422 | refer to the docstring of this method for more information. |
| 423 | """ |
| 424 | return self.tokenizer.batch_decode(*args, **kwargs) |
| 425 | |
| 426 | def decode(self, *args, **kwargs): |
| 427 | """This method forwards all its arguments to Qwen2TokenizerFast's [`~PythonBackend.decode`]. Please refer to |
| 428 | the docstring of this method for more information. |
| 429 | """ |
| 430 | return self.tokenizer.decode(*args, **kwargs) |
| 431 | |
| 432 | def post_process_image_text_to_text(self, generated_outputs): |
| 433 | """Post-process the output of the model to decode the text. |
| 434 | |
| 435 | Args: |
| 436 | generated_outputs (`torch.Tensor` or `np.ndarray`): |
| 437 | The output of the model `generate` function. The output is expected to be a tensor of shape `(batch_size, sequence_length)` |
| 438 | or `(sequence_length,)`. |
| 439 | |
| 440 | Returns: |
| 441 | `list[str]`: The decoded text. |
| 442 | |
| 443 | """ |
| 444 | return self.tokenizer.batch_decode( |
| 445 | generated_outputs, |
| 446 | skip_special_tokens=True, |
| 447 | clean_up_tokenization_spaces=False, |
| 448 | ) |
| 449 | |
| 450 | @property |
| 451 | def model_input_names(self): |
| 452 | tokenizer_input_names = self.tokenizer.model_input_names |
| 453 | image_processor_input_names = self.image_processor.model_input_names |
| 454 | names_from_processor = list( |
| 455 | dict.fromkeys(tokenizer_input_names + image_processor_input_names) |
| 456 | ) |
| 457 | return names_from_processor + ["second_per_grid_ts"] |