# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from itertools import product from math import ceil import numpy as np import torch from PIL import Image from torchvision import transforms from torchvision.transforms.functional import InterpolationMode from transformers import ( BatchFeature, ImageProcessingMixin, ProcessorMixin, TensorType, ) from vllm.tokenizers import TokenizerLike MAX_IMAGE_SIZE: int = 3024 ImageWithPatches = tuple[Image.Image, list[Image.Image], list[bool]] class Step3VisionProcessor: def __init__(self, size, interpolation_mode="bicubic", patch_size=None): mean = [0.48145466, 0.4578275, 0.40821073] std = [0.26862954, 0.26130258, 0.27577711] patch_size = patch_size if patch_size is not None else size self.transform = transforms.Compose( [ transforms.ToTensor(), transforms.Normalize(mean, std), transforms.Resize( (size, size), interpolation=InterpolationMode.BICUBIC if interpolation_mode == "bicubic" else InterpolationMode.BILINEAR, antialias=True, ), ] ) self.patch_transform = ( transforms.Compose( [ transforms.ToTensor(), transforms.Normalize(mean, std), transforms.Resize( (patch_size, patch_size), interpolation=InterpolationMode.BICUBIC if interpolation_mode == "bicubic" else InterpolationMode.BILINEAR, antialias=True, ), ] ) if patch_size is not None else None ) def __call__(self, image, is_patch=False): if is_patch: assert self.patch_transform is not None return {"pixel_values": self.patch_transform(image).unsqueeze(0)} return {"pixel_values": self.transform(image).unsqueeze(0)} class ImagePatcher: def __init__(self, enable_patch: bool = True) -> None: self.enable_patch = enable_patch def determine_window_size(self, long: int, short: int) -> int: if long < 728: return short if long / short > 1.5 else 0 return min(short, 504) if long / short > 4 else 504 def slide_window( self, width: int, height: int, sizes: list[tuple[int, int]], steps: list[tuple[int, int]], img_rate_thr: float = 0.6, ) -> tuple[list[tuple[int, int, int, int]], tuple[int, int]]: assert 1 >= img_rate_thr >= 0, "The `in_rate_thr` should lie in 0~1" windows = [] # Sliding windows. for size, step in zip(sizes, steps): size_w, size_h = size step_w, step_h = step x_num = 1 if width <= size_w else ceil((width - size_w) / step_w + 1) x_start = [step_w * i for i in range(x_num)] if len(x_start) > 1 and x_start[-1] + size_w > width: x_start[-1] = width - size_w y_num = 1 if height <= size_h else ceil((height - size_h) / step_h + 1) y_start = [step_h * i for i in range(y_num)] if len(y_start) > 1 and y_start[-1] + size_h > height: y_start[-1] = height - size_h start = np.array(list(product(y_start, x_start)), dtype=int) start[:, [0, 1]] = start[:, [1, 0]] windows.append(np.concatenate([start, start + size], axis=1)) windows = np.concatenate(windows, axis=0) return [ (int(box[0]), int(box[1]), int(box[2] - box[0]), int(box[3] - box[1])) for box in windows ], (x_num, y_num) def square_pad(self, img: Image.Image) -> Image.Image: w, h = img.size if w == h: return img size = max(w, h) padded = Image.new(img.mode, (size, size), 0) padded.paste(img, (0, 0)) return padded def get_image_size_for_padding( self, img_width: int, img_height: int ) -> tuple[int, int]: ratio = img_width / img_height if min(img_height, img_width) < 32 and (ratio > 4 or ratio < 1 / 4): new_size = max(img_height, img_width) return new_size, new_size return img_width, img_height def get_image_size_for_preprocess( self, img_width: int, img_height: int ) -> tuple[int, int]: if max(img_height, img_width) > MAX_IMAGE_SIZE: scale_factor = MAX_IMAGE_SIZE / max(img_height, img_width) img_width = int(img_width * scale_factor) img_height = int(img_height * scale_factor) return img_width, img_height def get_image_size_for_crop( self, img_width: int, img_height: int, window_size: int ): w_ratio = img_width / window_size h_ratio = img_height / window_size if w_ratio < 1: width_new = img_width else: decimal_w = w_ratio - img_width // window_size w_ratio = int(w_ratio) + 1 if decimal_w > 0.2 else int(w_ratio) width_new = window_size * w_ratio if h_ratio < 1: height_new = img_height else: decimal_h = h_ratio - img_height // window_size h_ratio = int(h_ratio) + 1 if decimal_h > 0.2 else int(h_ratio) height_new = window_size * h_ratio return int(width_new), int(height_new) def patch_crop(self, img: Image.Image, i: int, j: int, th: int, tw: int): target = img.crop((j, i, j + tw, i + th)) return target def get_num_patches(self, img_width: int, img_height: int) -> tuple[int, int]: img_width, img_height = self.get_image_size_for_padding(img_width, img_height) img_width, img_height = self.get_image_size_for_preprocess( img_width, img_height ) window_size = self.determine_window_size( max(img_height, img_width), min(img_height, img_width) ) if window_size == 0 or not self.enable_patch: return 0, 0 else: img_width, img_height = self.get_image_size_for_crop( img_width, img_height, window_size ) center_list, (x_num, y_num) = self.slide_window( img_width, img_height, [(window_size, window_size)], [(window_size, window_size)], ) full_rows = (len(center_list) - 1) // x_num + 1 if len(center_list) > 0 and len(center_list) % x_num == 0: full_rows -= 1 return len(center_list), full_rows def __call__( self, img: Image.Image ) -> tuple[Image.Image, list[Image.Image], list[bool]]: img_width, img_height = img.size new_img_width, new_img_height = self.get_image_size_for_padding( img_width, img_height ) if new_img_width != img_width or new_img_height != img_height: img = self.square_pad(img) img_width, img_height = img.size new_img_width, new_img_height = self.get_image_size_for_preprocess( img_width, img_height ) img = img.resize((new_img_width, new_img_height), Image.Resampling.BILINEAR) window_size = self.determine_window_size( max(new_img_height, new_img_width), min(new_img_height, new_img_width) ) if window_size == 0 or not self.enable_patch: return img, [], [] else: new_img_width, new_img_height = self.get_image_size_for_crop( new_img_width, new_img_height, window_size ) if (new_img_width, new_img_height) != (img_width, img_height): img_for_crop = img.resize( (new_img_width, new_img_height), Image.Resampling.BILINEAR ) else: img_for_crop = img patches = [] newlines = [] center_list, (x_num, y_num) = self.slide_window( new_img_width, new_img_height, [(window_size, window_size)], [(window_size, window_size)], ) for patch_id, center_lf_point in enumerate(center_list): x, y, patch_w, patch_h = center_lf_point big_patch = self.patch_crop(img_for_crop, y, x, patch_h, patch_w) patches.append(big_patch) if (patch_id + 1) % x_num == 0: newlines.append(patch_id) if newlines and newlines[-1] == len(patches) - 1: newlines.pop() return ( img, patches, [i in newlines for i in range(len(patches))], ) class Step3VLImageProcessor(ImageProcessingMixin): def __init__( self, image_size: int = 728, patch_size: int = 504, num_image_feature_size: int = 169, num_patch_feature_size: int = 81, enable_patch: bool = True, ) -> None: self.image_size = image_size self.patch_size = patch_size self.num_image_feature_size = num_image_feature_size self.num_patch_feature_size = num_patch_feature_size self.image_preprocessor = Step3VisionProcessor( image_size, "bilinear", patch_size ) self.patcher = ImagePatcher(enable_patch=enable_patch) def get_num_image_tokens(self, img_width: int, img_height: int) -> int: num_patches, num_newlines = self.patcher.get_num_patches(img_width, img_height) return ( num_patches * (self.num_patch_feature_size + 2) + self.num_image_feature_size + 2 + num_newlines ) def _split_images(self, images: list[Image.Image]) -> list[ImageWithPatches]: result = [] for img in images: result.append(self.patcher(img)) return result def _convert_images_to_pixel_values( self, images: list[Image.Image], is_patch: bool = False, ) -> list[torch.Tensor]: return [ self.image_preprocessor(img, is_patch=is_patch)["pixel_values"] for img in images ] def __call__( self, images: Image.Image | list[Image.Image], return_tensors: str | TensorType | None = None, ) -> BatchFeature: if not isinstance(images, list): images = [images] split_images_data = self._split_images(images) pixel_values_lst = [] patch_pixel_values_lst = [] patch_newline_mask_lst = [] num_patches = [] for raw_img, img_patches, patch_newline_mask in split_images_data: pixel_values_lst.extend(self._convert_images_to_pixel_values([raw_img])) num_patches.append(len(img_patches)) patch_pixel_values_lst.extend( self._convert_images_to_pixel_values(img_patches, is_patch=True) ) patch_newline_mask_lst.extend(patch_newline_mask) pixel_values = torch.cat(pixel_values_lst) patch_size = self.patch_size image_inputs = { "pixel_values": pixel_values, "num_patches": num_patches, "patch_pixel_values": ( torch.cat(patch_pixel_values_lst) if patch_pixel_values_lst else pixel_values.new_empty((0, 3, patch_size, patch_size)) ), "patch_newline_mask": torch.tensor( patch_newline_mask_lst, dtype=torch.bool ), } return BatchFeature(image_inputs, tensor_type=return_tensors) class Step3VLProcessor(ProcessorMixin): attributes = ["image_processor", "tokenizer"] def __init__( self, image_processor: Step3VLImageProcessor, tokenizer: TokenizerLike, ) -> None: self.image_processor = image_processor self.tokenizer = tokenizer self.image_start_token = image_start_token = "" self.image_end_token = image_end_token = "" self.patch_start_token = patch_start_token = "" self.patch_end_token = patch_end_token = "" self.patch_newline_token = patch_newline_token = "" self.image_start_token_id = tokenizer.convert_tokens_to_ids(image_start_token) self.image_end_token_id = tokenizer.convert_tokens_to_ids(image_end_token) self.patch_start_token_id = tokenizer.convert_tokens_to_ids(patch_start_token) self.patch_end_token_id = tokenizer.convert_tokens_to_ids(patch_end_token) self.patch_newline_token_id = tokenizer.convert_tokens_to_ids( patch_newline_token ) self.image_token = image_token = "" self.image_feature_tokens = image_token * image_processor.num_image_feature_size self.patch_feature_tokens = image_token * image_processor.num_patch_feature_size self.image_token_id = image_token_id = tokenizer.convert_tokens_to_ids( image_token ) self.image_feature_token_ids = [ image_token_id ] * image_processor.num_image_feature_size self.patch_feature_token_ids = [ image_token_id ] * image_processor.num_patch_feature_size def _get_patch_repl_text( self, num_patches: int, patch_newline_mask: list[bool], ) -> str: assert len(patch_newline_mask) == num_patches parts = [] for i in range(num_patches): parts.extend( [ self.patch_start_token, self.patch_feature_tokens, self.patch_end_token, ] ) if patch_newline_mask[i]: parts.append(self.patch_newline_token) return "".join(parts) def _get_patch_repl_ids( self, num_patches: int, patch_newline_mask: list[bool], ) -> list[int]: assert len(patch_newline_mask) == num_patches parts = [] for i in range(num_patches): parts.extend( [ self.patch_start_token_id, *self.patch_feature_token_ids, self.patch_end_token_id, ] ) if patch_newline_mask[i]: parts.append(self.patch_newline_token_id) return parts def _get_image_repl_text( self, num_images: int, ) -> str: parts = [ self.image_start_token, self.image_feature_tokens, self.image_end_token, ] * num_images return "".join(parts) def _get_image_repl_ids( self, num_images: int, ) -> list[int]: part = [ self.image_start_token_id, *self.image_feature_token_ids, self.image_end_token_id, ] return part * num_images def get_image_repl_feature_text( self, num_images: int, num_patches: int, patch_new_line_idx: list[bool], ) -> str: patch_repl = self._get_patch_repl_text(num_patches, patch_new_line_idx) image_repl = self._get_image_repl_text(num_images) return patch_repl + image_repl def get_image_repl_feature_ids( self, num_images: int, num_patches: int, patch_new_line_idx: list[bool], ) -> list[int]: patch_repl = self._get_patch_repl_ids(num_patches, patch_new_line_idx) image_repl = self._get_image_repl_ids(num_images) return patch_repl + image_repl def replace_placeholder(self, text: str, placeholder: str, repls: list[str]) -> str: parts = text.split(placeholder) if len(parts) - 1 != len(repls): raise ValueError( "The number of placeholders does not match the number of replacements." ) result = [parts[0]] for i, repl in enumerate(repls): result.append(repl) result.append(parts[i + 1]) return "".join(result) def __call__( self, text: str | list[str] | None = None, images: Image.Image | list[Image.Image] | None = None, return_tensors: str | TensorType | None = None, ) -> BatchFeature: if images is not None: image_inputs = self.image_processor( images=images, return_tensors=return_tensors, ) num_patches = image_inputs["num_patches"] patch_newline_mask = image_inputs["patch_newline_mask"] else: image_inputs = {} num_patches = [] patch_newline_mask = [] if text is not None: if not isinstance(text, list): text = [text] if image_inputs: image_token = self.image_token image_repl_str_lst = [] start = 0 for n_patches in num_patches: image_repl_str = self.get_image_repl_feature_text( 1, n_patches, patch_newline_mask[start : start + n_patches] ) image_repl_str_lst.append(image_repl_str) start += n_patches text = [ self.replace_placeholder(t, image_token, image_repl_str_lst) for t in text ] text_inputs = self.tokenizer(text) else: text_inputs = {} return BatchFeature( data={**text_inputs, **image_inputs}, tensor_type=return_tensors, )