# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 # This file was automatically generated from src/transformers/models/molmo2/modular_molmo2.py. # Do NOT edit this file manually as any edits will be overwritten by the generation of # the file from the modular. If any change should be done, please apply the change to the # modular_molmo2.py file directly. One of our CI enforces this. # 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 # Copyright 2026 The HuggingFace Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import itertools import math import torch from torch.nn import functional as F from ...image_processing_backends import TorchvisionBackend from ...image_processing_utils import BatchFeature from ...image_transforms import group_images_by_shape, reorder_images from ...image_utils import IMAGENET_STANDARD_MEAN, IMAGENET_STANDARD_STD, ImageInput, PILImageResampling, SizeDict from ...processing_utils import ImagesKwargs, Unpack from ...utils import TensorType, auto_docstring class Molmo2ImageProcessorKwargs(ImagesKwargs, total=False): """ max_crops (`int`, *optional*, defaults to 8): Maximum number of crops to use per image. overlap_margins (`list[int]`, *optional*, defaults to `[4, 4]`): Overlap margins (in patches) for overlapping crop extraction. patch_size (`int`, *optional*, defaults to 14): The spatial patch size of the vision encoder. pooling_size (`list[int]`, *optional*, defaults to `[2, 2]`): The pooling size of the vision adapter. """ max_crops: int overlap_margins: list[int] patch_size: int pooling_size: list[int] def select_tiling(height: int, width: int, patch_size: int, max_num_crops: int) -> tuple[int, int]: """Same as `get_optimal_tiled_canvas` in cohere2_vision, with the candidates in (height, width) order and ties broken on height, as in the original Molmo2 processor.""" tilings = [ (tile_height, tile_width) for tile_height, tile_width in itertools.product(range(1, max_num_crops + 1), repeat=2) if tile_height * tile_width <= max_num_crops ] tilings.sort(key=lambda x: (x[0] * x[1], x[0])) candidate_resolutions = torch.tensor(tilings, dtype=torch.int32) * patch_size original_size = torch.tensor([height, width], dtype=torch.float32) required_scales = candidate_resolutions.to(torch.float32) / original_size required_scale = required_scales.amin(dim=-1, keepdim=True) if torch.all(required_scale < 1): return tilings[int(required_scale.argmax())] required_scale = torch.where(required_scale < 1.0, 10e9, required_scale) return tilings[int(required_scale.argmin())] def compute_crop_geometry( height: int, width: int, max_crops: int, overlap_margins: list[int], crop_size: int, patch_size: int, ) -> tuple[int, int, int, int]: """Tiling `(tiling_h, tiling_w)` and canvas pixel size `(canvas_height, canvas_width)` for one image.""" left_margin, right_margin = overlap_margins crop_patches = crop_size // patch_size window_patches = crop_patches - (left_margin + right_margin) window_size = window_patches * patch_size margin_size = (left_margin + right_margin) * patch_size tiling_h, tiling_w = select_tiling(height - margin_size, width - margin_size, window_size, max_crops) canvas_height = tiling_h * window_size + margin_size canvas_width = tiling_w * window_size + margin_size return tiling_h, tiling_w, canvas_height, canvas_width def batch_pixels_to_patches(images: torch.Tensor, patch_size: int) -> torch.Tensor: """Reshape images of [num_images, h, w, channels] -> [num_images, n_patches, pixels_per_patch]""" num_crops, height, width, channels = images.shape num_patches_height = height // patch_size num_patches_width = width // patch_size images = images.reshape(num_crops, num_patches_height, patch_size, num_patches_width, patch_size, channels) images = images.transpose(2, 3) images = images.reshape(num_crops, num_patches_height * num_patches_width, patch_size * patch_size * channels) return images def arange_for_pooling( index_grid: torch.Tensor, pool_h: int, pool_w: int, ) -> tuple[torch.Tensor, int, int]: """Group a [h, w] patch-index grid into pooling windows: `[num_rows * num_cols, pool_h * pool_w]`, padded with -1. Also returns the pooled grid's `(num_rows, num_cols)`.""" height_padding = pool_h * ((index_grid.shape[0] + pool_h - 1) // pool_h) - index_grid.shape[0] width_padding = pool_w * ((index_grid.shape[1] + pool_w - 1) // pool_w) - index_grid.shape[1] index_grid = F.pad( index_grid, (width_padding // 2, (width_padding + 1) // 2, height_padding // 2, (height_padding + 1) // 2), mode="constant", value=-1, ) num_rows, num_cols = index_grid.shape[0] // pool_h, index_grid.shape[1] // pool_w pooling_indices = ( index_grid.reshape(num_rows, pool_h, num_cols, pool_w) .permute(0, 2, 1, 3) .reshape(num_rows * num_cols, pool_h * pool_w) ) return pooling_indices, num_rows, num_cols def resize_and_normalize_image( backend: "TorchvisionBackend", image: torch.Tensor, output_size: list[int], resample: PILImageResampling, do_rescale: bool, rescale_factor: float, do_normalize: bool, image_mean: list[float], image_std: list[float], ) -> torch.Tensor: # Under `torch.compile`, torchvision's uint8 bilinear resize can produce values slightly outside 0-255, # which overflow and wrap around (-1 becomes 255, so a near-black pixel becomes white). Resizing in # float, then rounding and clamping back, avoids that and gives the exact same bytes in eager mode. is_integer_input = not image.dtype.is_floating_point resized = backend.resize( image.float() if is_integer_input else image, size=SizeDict(height=output_size[0], width=output_size[1]), resample=resample, antialias=False, ) if is_integer_input: resized = resized.round().clamp(0, 255).to(image.dtype) return backend.rescale_and_normalize(resized, do_rescale, rescale_factor, do_normalize, image_mean, image_std) def build_resized_image( backend: "TorchvisionBackend", images: torch.Tensor, crop_size: int, resample: PILImageResampling, do_rescale: bool, rescale_factor: float, do_normalize: bool, image_mean: list[float], image_std: list[float], image_patch_size: int, ) -> tuple[torch.Tensor, torch.Tensor]: # `images`: a batch of same-shape images `[N, C, H, W]`; resize the whole batch at once. resized = resize_and_normalize_image( backend, images, [crop_size, crop_size], resample, do_rescale=do_rescale, rescale_factor=rescale_factor, do_normalize=do_normalize, image_mean=image_mean, image_std=image_std, ) # [N, C, S, S] -> [N, 1, S, S, C]: one global (low-res) view per image. resized = resized.permute(0, 2, 3, 1).unsqueeze(1) # The per-patch index grid depends only on the (shared) shape, so it is built once. patches_per_side = crop_size // image_patch_size resized_index_grid = torch.arange(patches_per_side**2, dtype=torch.int32, device=images.device).reshape( patches_per_side, patches_per_side ) return resized, resized_index_grid @auto_docstring class Molmo2ImageProcessor(TorchvisionBackend): valid_kwargs = Molmo2ImageProcessorKwargs model_input_names = ["pixel_values", "image_token_pooling", "image_grids", "image_num_crops"] resample = PILImageResampling.BILINEAR image_mean = IMAGENET_STANDARD_MEAN image_std = IMAGENET_STANDARD_STD size = {"height": 378, "width": 378} do_resize = True do_rescale = True do_normalize = True do_convert_rgb = True max_crops = 8 overlap_margins = [4, 4] patch_size = 14 pooling_size = [2, 2] def __init__(self, **kwargs: Unpack[Molmo2ImageProcessorKwargs]): super().__init__(**kwargs) def _build_overlapping_crops( self, images: torch.Tensor, max_crops: int, overlap_margins: list[int], crop_size: int, resample: PILImageResampling, do_rescale: bool, rescale_factor: float, do_normalize: bool, image_mean: list[float], image_std: list[float], image_patch_size: int, ) -> tuple[torch.Tensor, torch.Tensor]: """Tile a batch of same-shape images `[N, C, H, W]` into overlapping square crops. The tiling and the per-patch index grid depend only on the (shared) shape, so they are computed once and the resize/unfold are batched over N.""" left_margin, right_margin = overlap_margins crop_patches = crop_size // image_patch_size window_patches = crop_patches - (left_margin + right_margin) window_size = window_patches * image_patch_size num_images, num_channels, original_height, original_width = images.shape tiling_h, tiling_w, canvas_height, canvas_width = compute_crop_geometry( original_height, original_width, max_crops=max_crops, overlap_margins=overlap_margins, crop_size=crop_size, patch_size=image_patch_size, ) canvas = resize_and_normalize_image( self, images, [canvas_height, canvas_width], resample, do_rescale=do_rescale, rescale_factor=rescale_factor, do_normalize=do_normalize, image_mean=image_mean, image_std=image_std, ) # [N, C, canvas_height, canvas_width] -> unfold spatial dims -> [N, C, tiling_h, tiling_w, crop, crop] crops = canvas.unfold(2, crop_size, window_size).unfold(3, crop_size, window_size) crops = ( crops.permute(0, 2, 3, 4, 5, 1) .reshape(num_images, tiling_h * tiling_w, crop_size, crop_size, num_channels) .contiguous() ) patch_index_grid = torch.arange( tiling_h * tiling_w * crop_patches * crop_patches, dtype=torch.int32, device=images.device ).reshape(tiling_h, tiling_w, crop_patches, crop_patches) if left_margin: patch_index_grid[1:, :, :left_margin, :] = -1 patch_index_grid[:, 1:, :, :left_margin] = -1 if right_margin: patch_index_grid[:-1, :, -right_margin:, :] = -1 patch_index_grid[:, :-1, :, -right_margin:] = -1 patch_index_grid = patch_index_grid.permute(0, 2, 1, 3).reshape(-1) patch_index_grid = patch_index_grid[patch_index_grid >= 0].reshape( canvas_height // image_patch_size, canvas_width // image_patch_size ) return crops, patch_index_grid def _image_batch_to_patches_and_grids( self, images: torch.Tensor, max_crops: int, overlap_margins: list[int], crop_size: int, resample: PILImageResampling, do_rescale: bool, rescale_factor: float, do_normalize: bool, image_mean: list[float], image_std: list[float], image_patch_size: int, image_pooling_w: int, image_pooling_h: int, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Process a batch of same-shape images `[N, C, H, W]`. The grid and pooling indices are shape-dependent (shared across the batch) and returned expanded to `[N, ...]` so the caller can reorder them per-image alongside the batched patch tensor.""" patches_per_crop = (crop_size // image_patch_size) ** 2 crops, patch_index_grid = self._build_overlapping_crops( images, max_crops=max_crops, overlap_margins=overlap_margins, crop_size=crop_size, resample=resample, do_rescale=do_rescale, rescale_factor=rescale_factor, do_normalize=do_normalize, image_mean=image_mean, image_std=image_std, image_patch_size=image_patch_size, ) pooling_indices, num_patch_rows, num_patch_cols = arange_for_pooling( patch_index_grid, image_pooling_h, image_pooling_w ) resized, resized_index_grid = build_resized_image( self, images, crop_size=crop_size, resample=resample, do_rescale=do_rescale, rescale_factor=rescale_factor, do_normalize=do_normalize, image_mean=image_mean, image_std=image_std, image_patch_size=image_patch_size, ) # [N, 1, S, S, C] + [N, ncrops, S, S, C] -> [N, 1 + ncrops, S, S, C] crops = torch.cat([resized, crops], dim=1) resized_pooling_indices, resized_h, resized_w = arange_for_pooling( resized_index_grid, image_pooling_h, image_pooling_w ) pooling_indices = torch.where(pooling_indices >= 0, pooling_indices + patches_per_crop, -1) pooling_indices = torch.cat([resized_pooling_indices, pooling_indices]) image_grid = torch.tensor( [[resized_h, resized_w, num_patch_rows, num_patch_cols]], dtype=torch.int64, device=images.device ) # [N, total_crops, S, S, C] -> patches [N, total_crops, n_patch, pixels_per_patch] num_images, total_crops = crops.shape[0], crops.shape[1] pixels_per_patch = image_patch_size * image_patch_size * crops.shape[-1] patches = batch_pixels_to_patches( crops.reshape(num_images * total_crops, *crops.shape[2:]), image_patch_size ).reshape(num_images, total_crops, -1, pixels_per_patch) # The grid and pooling indices are shared by all images of this shape; expand them to one row # per image so the batch can later be restored to the original image order. return image_grid.expand(num_images, -1), patches, pooling_indices.unsqueeze(0).expand(num_images, -1, -1) @auto_docstring def preprocess( self, images: ImageInput, **kwargs: Unpack[Molmo2ImageProcessorKwargs], ) -> BatchFeature: return super().preprocess(images, **kwargs) def _preprocess( self, images: list["torch.Tensor"], do_resize: bool, size: SizeDict, resample: PILImageResampling, do_rescale: bool, rescale_factor: float, do_normalize: bool, image_mean: list[float], image_std: list[float], do_convert_rgb: bool, max_crops: int, overlap_margins: list[int], patch_size: int, pooling_size: list[int], disable_grouping: bool | None, return_tensors: str | TensorType | None, **kwargs, ) -> BatchFeature: if size.height != size.width: raise ValueError(f"Molmo2 only supports a square `size`, got height={size.height}, width={size.width}.") crop_size = size.height image_pooling_h, image_pooling_w = pooling_size # Group images by shape and process each unique shape as a single batch (the tiling and all # patch/pooling indices are shape-dependent only), then restore the original order. device = images[0].device grouped_images, grouped_index = group_images_by_shape(images, disable_grouping=disable_grouping) grids_grouped: dict = {} patches_grouped: dict = {} pooled_grouped: dict = {} for shape, stacked_images in grouped_images.items(): image_grid, patches, pooled_indices = self._image_batch_to_patches_and_grids( stacked_images, max_crops=max_crops, overlap_margins=overlap_margins, crop_size=crop_size, resample=resample, do_rescale=do_rescale, rescale_factor=rescale_factor, do_normalize=do_normalize, image_mean=image_mean, image_std=image_std, image_patch_size=patch_size, image_pooling_w=image_pooling_w, image_pooling_h=image_pooling_h, ) grids_grouped[shape] = image_grid patches_grouped[shape] = patches pooled_grouped[shape] = pooled_indices grids = reorder_images(grids_grouped, grouped_index) patches = reorder_images(patches_grouped, grouped_index) pooled = reorder_images(pooled_grouped, grouped_index) all_crops: list[torch.Tensor] = [] all_pooled: list[torch.Tensor] = [] patch_offset = 0 for crops, pooled_indices in zip(patches, pooled): all_pooled.append(torch.where(pooled_indices >= 0, pooled_indices + patch_offset, pooled_indices)) all_crops.append(crops) patch_offset += crops.shape[0] * crops.shape[1] data = { "pixel_values": torch.cat(all_crops, dim=0), "image_token_pooling": torch.cat(all_pooled, dim=0), "image_grids": torch.stack(grids, dim=0), "image_num_crops": torch.tensor([crops.shape[0] for crops in patches], dtype=torch.int64, device=device), } return BatchFeature(data=data, tensor_type=return_tensors) def get_number_of_image_patches(self, height: int, width: int, images_kwargs=None) -> int: image_grid, _ = self.get_image_grid_and_crops(height, width, images_kwargs) resized_h, resized_w, high_res_rows, high_res_cols = image_grid return resized_h * resized_w + high_res_rows * high_res_cols def get_image_grid_and_crops(self, height: int, width: int, images_kwargs=None) -> tuple[list[int], int]: """Return the `[low_res_h, low_res_w, high_res_h, high_res_w]` pooled grid and the crop count for one image.""" if images_kwargs is None: images_kwargs = {} max_crops = images_kwargs.get("max_crops", self.max_crops) overlap_margins = images_kwargs.get("overlap_margins", self.overlap_margins) patch_size = images_kwargs.get("patch_size", self.patch_size) pooling_size = images_kwargs.get("pooling_size", self.pooling_size) size = images_kwargs.get("size", self.size) base_h = size["height"] if isinstance(size, dict) else size.height base_w = size["width"] if isinstance(size, dict) else size.width pooling_h, pooling_w = pooling_size tiling_h, tiling_w, canvas_height, canvas_width = compute_crop_geometry( height, width, max_crops=max_crops, overlap_margins=overlap_margins, crop_size=base_h, patch_size=patch_size, ) num_patch_rows_high = math.ceil((canvas_height // patch_size) / pooling_h) num_patch_cols_high = math.ceil((canvas_width // patch_size) / pooling_w) resized_h = math.ceil((base_h // patch_size) / pooling_h) resized_w = math.ceil((base_w // patch_size) / pooling_w) return [resized_h, resized_w, num_patch_rows_high, num_patch_cols_high], tiling_h * tiling_w + 1 __all__ = ["Molmo2ImageProcessor"]