lemuralabs/Step-3.7-Flash-OptiQ-3.7bpw-mlx
394
1from transformers import BaseImageProcessor, ImageProcessingMixin2from transformers.processing_utils import ImagesKwargs, MultiModalData, ProcessingKwargs, ProcessorMixin, Unpack, VideosKwargs3import math4from typing import Iterable, Optional, Tuple, List, TypedDict, Literal, Union, overload5 6from PIL import Image7import torch8import numpy as np9import torchvision10from torch import nn11from torch.nn import functional as F, LayerNorm12from torchvision.transforms.functional import InterpolationMode13from transformers.activations import ACT2FN14from torchvision import transforms15from torchvision.transforms.functional import InterpolationMode16from transformers.feature_extraction_utils import BatchFeature, TensorType17from transformers.image_utils import ImageInput18from transformers.processing_utils import ProcessingKwargs, ProcessorMixin, Unpack19from transformers.tokenization_utils_tokenizers import TokenizersBackend20from math import ceil21from itertools import product22 23 24 25MAX_IMAGE_SIZE: int = 302426 27class Step3VLImagePixelInputs(TypedDict):28 type: Literal["pixel_values"]29 pixel_values: torch.Tensor30 patch_pixel_values: Optional[torch.Tensor]31 num_patches: list[int]32 33 34class Step3VLImageEmbeddingInputs(TypedDict):35 type: Literal["image_embeds"]36 image_embeds: torch.Tensor37 38 39ImageWithPatches = tuple[Image.Image, list[Image.Image], list[int] | None]40 41 42class GPUToTensor(torch.nn.Module):43 44 def forward(self, raw_image: Union[np.ndarray,45 Image.Image]) -> torch.Tensor:46 if isinstance(raw_image, Image.Image):47 return transforms.ToTensor()(raw_image)48 if raw_image.ndim == 2:49 raw_image = raw_image[:, :, None].repeat(3, -1)50 if torch.cuda.is_available():51 device = torch.device("cuda")52 else:53 device = torch.device("cpu")54 image_tensor = torch.from_numpy(raw_image).to(device)55 image_tensor = torch.permute(image_tensor, (2, 0, 1)).contiguous()56 if image_tensor.dtype == torch.uint8:57 image_tensor = image_tensor.to(torch.float32).div(255)58 return image_tensor59 60class Step3VisionProcessor(BaseImageProcessor):61 62 def __init__(self, size, interpolation_mode="bicubic", patch_size=None):63 mean = [0.48145466, 0.4578275, 0.40821073]64 std = [0.26862954, 0.26130258, 0.27577711]65 patch_size = patch_size if patch_size is not None else size66 67 self.transform = transforms.Compose([68 GPUToTensor(),69 transforms.Normalize(mean, std),70 transforms.Resize(71 (size, size),72 interpolation=InterpolationMode.BICUBIC if interpolation_mode73 == "bicubic" else InterpolationMode.BILINEAR,74 antialias=True),75 ])76 77 self.patch_transform = transforms.Compose([78 GPUToTensor(),79 transforms.Normalize(mean, std),80 transforms.Resize(81 (patch_size, patch_size),82 interpolation=InterpolationMode.BICUBIC if interpolation_mode83 == "bicubic" else InterpolationMode.BILINEAR,84 antialias=True),85 ]) if patch_size is not None else None86 87 def __call__(self, image, is_patch=False):88 if is_patch:89 return {"pixel_values": self.patch_transform(image).unsqueeze(0)}90 else:91 return {"pixel_values": self.transform(image).unsqueeze(0)}92 93class ImagePatcher:94 def determine_window_size(self, long: int, short: int) -> int:95 if long <= 728:96 return short if long / short > 1.5 else 097 return min(short, 504) if long / short > 4 else 50498 def slide_window(99 self,100 width: int,101 height: int,102 sizes: list[tuple[int, int]],103 steps: list[tuple[int, int]],104 img_rate_thr: float = 0.6,105 ) -> tuple[list[tuple[int, int, int, int]], tuple[int, int]]:106 assert 1 >= img_rate_thr >= 0, "The `in_rate_thr` should lie in 0~1"107 windows = []108 # Sliding windows.109 for size, step in zip(sizes, steps):110 size_w, size_h = size111 step_w, step_h = step112 113 x_num = 1 if width <= size_w else ceil((width - size_w) / step_w +114 1)115 x_start = [step_w * i for i in range(x_num)]116 if len(x_start) > 1 and x_start[-1] + size_w > width:117 x_start[-1] = width - size_w118 119 y_num = 1 if height <= size_h else ceil((height - size_h) /120 step_h + 1)121 y_start = [step_h * i for i in range(y_num)]122 if len(y_start) > 1 and y_start[-1] + size_h > height:123 y_start[-1] = height - size_h124 125 start = np.array(list(product(y_start, x_start)), dtype=int)126 start[:, [0, 1]] = start[:, [1, 0]]127 windows.append(np.concatenate([start, start + size], axis=1))128 windows = np.concatenate(windows, axis=0)129 130 return [(int(box[0]), int(box[1]), int(box[2] - box[0]),131 int(box[3] - box[1])) for box in windows], (x_num, y_num)132 133 def square_pad(self, img: Image.Image) -> Image.Image:134 w, h = img.size135 if w == h:136 return img137 size = max(w, h)138 padded = Image.new(img.mode, (size, size), 0)139 padded.paste(img, (0, 0))140 return padded141 142 def get_image_size_for_padding(self, img_width: int,143 img_height: int) -> tuple[int, int]:144 ratio = img_width / img_height145 if min(img_height, img_width) < 32 and (ratio > 4 or ratio < 1 / 4):146 new_size = max(img_height, img_width)147 return new_size, new_size148 return img_width, img_height149 150 def get_image_size_for_preprocess(self, img_width: int,151 img_height: int) -> tuple[int, int]:152 153 if max(img_height, img_width) > MAX_IMAGE_SIZE:154 scale_factor = MAX_IMAGE_SIZE / max(img_height, img_width)155 img_width = int(img_width * scale_factor)156 img_height = int(img_height * scale_factor)157 return img_width, img_height158 159 def get_image_size_for_crop(self, img_width: int, img_height: int,160 window_size: int):161 w_ratio = img_width / window_size162 h_ratio = img_height / window_size163 164 if w_ratio < 1:165 width_new = img_width166 else:167 decimal_w = w_ratio - img_width // window_size168 w_ratio = int(w_ratio) + 1 if decimal_w > 0.2 else int(w_ratio)169 width_new = window_size * w_ratio170 if h_ratio < 1:171 height_new = img_height172 else:173 decimal_h = h_ratio - img_height // window_size174 h_ratio = int(h_ratio) + 1 if decimal_h > 0.2 else int(h_ratio)175 height_new = window_size * h_ratio176 return int(width_new), int(height_new)177 178 def patch_crop(self, img: Image.Image, i: int, j: int, th: int, tw: int):179 target = img.crop((j, i, j + tw, i + th))180 return target181 182 def get_num_patches(self, img_width: int,183 img_height: int) -> tuple[int, int]:184 img_width, img_height = self.get_image_size_for_padding(185 img_width, img_height)186 img_width, img_height = self.get_image_size_for_preprocess(187 img_width, img_height)188 window_size = self.determine_window_size(max(img_height, img_width),189 min(img_height, img_width))190 if window_size == 0:191 return 0, 0192 else:193 img_width, img_height = self.get_image_size_for_crop(194 img_width, img_height, window_size)195 center_list, (x_num, y_num) = self.slide_window(196 img_width, img_height, [(window_size, window_size)],197 [(window_size, window_size)])198 full_rows = (len(center_list) - 1) // x_num + 1199 if len(center_list) > 0 and len(center_list) % x_num == 0:200 full_rows -= 1201 return len(center_list), full_rows202 203 def __call__(204 self, img: Image.Image205 ) -> tuple[Image.Image, list[Image.Image], list[bool] | None]:206 img_width, img_height = img.size207 new_img_width, new_img_height = self.get_image_size_for_padding(208 img_width, img_height)209 if new_img_width != img_width or new_img_height != img_height:210 img = self.square_pad(img)211 img_width, img_height = img.size212 213 new_img_width, new_img_height = self.get_image_size_for_preprocess(214 img_width, img_height)215 img = img.resize((new_img_width, new_img_height),216 Image.Resampling.BILINEAR)217 window_size = self.determine_window_size(218 max(new_img_height, new_img_width),219 min(new_img_height, new_img_width))220 # return img, [], None221 if window_size == 0:222 return img, [], None223 else:224 new_img_width, new_img_height = self.get_image_size_for_crop(225 new_img_width, new_img_height, window_size)226 if (new_img_width, new_img_height) != (img_width, img_height):227 img_for_crop = img.resize((new_img_width, new_img_height),228 Image.Resampling.BILINEAR)229 else:230 img_for_crop = img231 232 patches = []233 newlines = []234 center_list, (x_num, y_num) = self.slide_window(235 new_img_width, new_img_height, [(window_size, window_size)],236 [(window_size, window_size)])237 for patch_id, center_lf_point in enumerate(center_list):238 x, y, patch_w, patch_h = center_lf_point239 big_patch = self.patch_crop(img_for_crop, y, x, patch_h,240 patch_w)241 patches.append(big_patch)242 if (patch_id + 1) % x_num == 0:243 newlines.append(patch_id)244 245 if newlines and newlines[-1] == len(patches) - 1:246 newlines.pop()247 248 return img, patches, [i in newlines for i in range(len(patches))] if len(patches) > 0 else None249 250 251 252 253class Step3VLProcessor(ProcessorMixin):254 # Align ProcessorMixin with our custom components.255 # We only have an image processor (not a feature extractor) plus a tokenizer.256 attributes = ["tokenizer"]257 tokenizer_class = "AutoTokenizer"258 259 @classmethod260 def _load_tokenizer_from_pretrained(261 cls, sub_processor_type, pretrained_model_name_or_path, subfolder="", **kwargs262 ):263 return TokenizersBackend.from_pretrained(264 pretrained_model_name_or_path,265 subfolder=subfolder,266 **kwargs,267 )268 269 def __init__(270 self,271 tokenizer=None,272 chat_template=None,273 **kwargs274 ) -> None:275 self.image_size = 728276 self.patch_size = 504277 278 self.image_preprocessor = Step3VisionProcessor(self.image_size,279 "bilinear",280 self.patch_size)281 282 self.num_image_feature_size = 169283 self.num_patch_feature_size = 81284 self.image_token = "<im_patch>"285 self.image_feature_placeholder = (self.image_token *286 self.num_image_feature_size)287 self.patch_feature_placeholder = (self.image_token *288 self.num_patch_feature_size)289 super().__init__(tokenizer=tokenizer, chat_template=chat_template, **kwargs)290 self.patcher = ImagePatcher()291 292 @property293 def image_token_id(self) -> int:294 return self.tokenizer.get_vocab()[self.image_token]295 296 def get_num_image_tokens(self, img_width: int, img_height: int) -> int:297 num_patches, num_newlines = self.patcher.get_num_patches(298 img_width, img_height)299 300 return num_patches * (301 self.num_patch_feature_size +302 2) + self.num_image_feature_size + 2 + num_newlines303 304 def _split_images(self,305 images: list[Image.Image]) -> list[ImageWithPatches]:306 result = []307 for img in images:308 result.append(self.patcher(img))309 return result310 311 def _convert_images_to_pixel_values(312 self,313 images: list[Image.Image],314 is_patch: bool = False,315 ) -> list[torch.Tensor]:316 return [317 self.image_preprocessor(img, is_patch=is_patch)["pixel_values"]318 for img in images319 ]320 321 def _get_patch_repl(322 self,323 num_patches: int,324 patch_newline_mask: list[bool] | None,325 ) -> tuple[str, list[int]]:326 text = ""327 token_ids = []328 for i in range(num_patches):329 assert len(patch_newline_mask) == num_patches330 text += f"<patch_start>{self.patch_feature_placeholder}<patch_end>"331 token_ids.extend(332 [self.tokenizer.convert_tokens_to_ids("<patch_start>")] +333 [self.image_token_id] * self.num_patch_feature_size +334 [self.tokenizer.convert_tokens_to_ids("<patch_end>")])335 if patch_newline_mask and patch_newline_mask[i]:336 text += "<patch_newline>"337 token_ids.append(338 self.tokenizer.convert_tokens_to_ids("<patch_newline>"))339 return text, token_ids340 341 def _get_image_repl(342 self,343 num_images: int,344 ) -> tuple[str, list[int]]:345 text = f"<im_start>{self.image_feature_placeholder}<im_end>"346 token_ids = [347 self.tokenizer.convert_tokens_to_ids("<im_start>")348 ] + [self.image_token_id] * self.num_image_feature_size + [349 self.tokenizer.convert_tokens_to_ids("<im_end>")350 ]351 return text * num_images, token_ids * num_images352 353 def _get_image_repl_features(354 self,355 num_images: int,356 num_patches: int,357 patch_new_line_idx: Optional[list[bool]],358 ) -> tuple[str, list[int]]:359 if num_patches > 0:360 patch_repl, patch_repl_ids = self._get_patch_repl(361 num_patches, patch_new_line_idx)362 else:363 patch_repl = ""364 patch_repl_ids = []365 image_repl, image_repl_ids = self._get_image_repl(num_images)366 return patch_repl + image_repl, patch_repl_ids + image_repl_ids367 368 def replace_placeholder(self, text: str, placeholder: str,369 repls: list[str]) -> str:370 parts = text.split(placeholder)371 372 if len(parts) - 1 != len(repls):373 raise ValueError(374 "The number of placeholders does not match the number of replacements." # noqa: E501375 )376 377 result = [parts[0]]378 for i, repl in enumerate(repls):379 result.append(repl)380 result.append(parts[i + 1])381 382 return "".join(result)383 384 def __call__(385 self,386 text: Optional[Union[str, list[str]]] = None,387 images: ImageInput | None = None,388 return_tensors: Optional[Union[str, TensorType]] = None,389 **kwargs,390 ) -> BatchFeature:391 392 if images is not None:393 images = self.image_preprocessor.fetch_images(images)394 if text is None:395 text = []396 if not isinstance(text, list):397 text = [text]398 if images is None:399 images = []400 elif not isinstance(images, list):401 images = [images]402 elif isinstance(images[0], list):403 images = images[0]404 405 if len(images) == 0:406 image_inputs = {}407 text_inputs = self.tokenizer(text)408 else:409 splitted_images_data = self._split_images(images)410 pixel_values_lst = []411 patch_pixel_values_lst = []412 patch_newline_mask_lst = []413 image_repl_str_lst = []414 image_repl_ids_lst = []415 num_patches = []416 for raw_img, img_patches, patch_newline_mask in splitted_images_data: # noqa: E501417 pixel_values_lst.extend(418 self._convert_images_to_pixel_values([raw_img]))419 420 if len(img_patches) > 0:421 patch_pixel_values_lst.extend(422 self._convert_images_to_pixel_values(img_patches,423 is_patch=True))424 num_patches.append(len(img_patches))425 426 image_repl_str, image_repl_ids = self._get_image_repl_features(427 1, len(img_patches), patch_newline_mask)428 image_repl_str_lst.append(image_repl_str)429 image_repl_ids_lst.extend(image_repl_ids)430 431 if patch_newline_mask is not None:432 patch_newline_mask_lst.extend(patch_newline_mask)433 434 image_inputs = {435 "pixel_values": torch.cat(pixel_values_lst),436 "num_patches": num_patches,437 }438 if patch_pixel_values_lst:439 image_inputs["patch_pixel_values"] = torch.cat(440 patch_pixel_values_lst)441 if patch_newline_mask_lst:442 image_inputs["patch_newline_mask"] = torch.tensor(443 patch_newline_mask_lst, dtype=torch.bool)444 445 text = [446 self.replace_placeholder(t, self.image_token,447 image_repl_str_lst) for t in text448 ]449 text_inputs = self.tokenizer(text)450 451 return BatchFeature(452 {453 **text_inputs,454 **image_inputs,455 },456 tensor_type=return_tensors,457 )458 459 # Copied from transformers.models.clip.processing_clip.CLIPProcessor.batch_decode with CLIP->Gemma460 def batch_decode(self, *args, **kwargs):461 """462 This method forwards all its arguments to GemmaTokenizerFast's [`~PreTrainedTokenizer.batch_decode`]. Please463 refer to the docstring of this method for more information.464 """465 return self.tokenizer.batch_decode(*args, **kwargs)466 467 # Copied from transformers.models.clip.processing_clip.CLIPProcessor.decode with CLIP->Gemma468 def decode(self, *args, **kwargs):469 """470 This method forwards all its arguments to GemmaTokenizerFast's [`~PreTrainedTokenizer.decode`]. Please refer to471 the docstring of this method for more information.472 """473 return self.tokenizer.decode(*args, **kwargs)474 475__all__ = ["Step3VLProcessor"]476