from __future__ import annotations import math import torch from PIL import Image from torchvision.transforms.v2 import functional as tvF from transformers.image_utils import PILImageResampling from models.config import VLMConfig """ Explanation: Assume patch_size = 16, pooling_kernel_size = 3, and max_teacher_patches = 1,500 Suppose we have an input image with (height = 500, width = 600) total_px = height * width = 500 * 600 = 300,000 target_px = max_teacher_patches * (patch_size**2) = 1,500 * (16**2) = 384,000 factor = math.sqrt(target_px / total_px) = sqrt(384,000 / 300,000) = 1.1313 - factor by which each side will be multiplied ideal_height = factor * height = 500 * 1.1313 = 565.65 ideal_width = factor * width = 600 * 1.1313 = 678.78 side_mult = pooling_kernel_size * patch_size = 3 * 16 = 48 - each side must be a multiple of 48 target_height = int(math.floor(ideal_height / side_mult)) * side_mult = int(math.floor(565.65 / 48)) * 48 = 528 target_width = int(math.floor(ideal_width / side_mult)) * side_mult = int(math.floor(678.78 / 48)) * 48 = 672\ max_side_length = (max_teacher_patches // pooling_kernel_size**2) * side_mult = (1,500 // 3**2) * 48 = 166 * 48 = 7,968 [ (max_teacher_patches // pooling_kernel_size**2) counts the maximum number of model patches (48x48) and * side_mult turns it into a number of pixels ] """ def get_aspect_ratio_preserving_size( height: int, width: int, teacher_patch_size: int, max_teacher_patches: int, pooling_kernel_size: int, ) -> tuple[int, int]: """ Purpose: Determine target dimensions for image that is resized to preserve aspect ratio so it fits within the patch budget. Target dimensions are the largest that: 1) Produce at most `max_teacher_patches` patches when patchified with `patch_size` 2) Have height and width divisible by `pooling_kernel_size * patch_size` (i.e. size of model patch) Parameters: * height (int) : height of the image to resize * width (int) : width of the image to resize * teacher_patch_size (int) : length, in pixels, of one side of a patch that teacher uses * max_teacher_patches (int) : maximum number of patches (of size teacher_patch_size) that a resized image may contain * pooling_kernel_size (int) : length of one side of a pooling kernel that pools multiple teacher patches Returns: A tuple containing (target_height, target_width) - new dimensions the image should have after resizing """ total_px = height * width target_px = max_teacher_patches * (teacher_patch_size**2) factor = math.sqrt(target_px / total_px) ideal_height = factor * height ideal_width = factor * width side_mult = pooling_kernel_size * teacher_patch_size # Round down to nearest multiple of side_mult target_height = int(math.floor(ideal_height / side_mult)) * side_mult target_width = int(math.floor(ideal_width / side_mult)) * side_mult # Handle edge cases where one or both dimensions round to 0 if target_height == 0 and target_width == 0: raise ValueError( "Attempting to resize to a 0 x 0 image. Resized height should be divisble by " f"`pooling_kernel_size * patch_size`={pooling_kernel_size * teacher_patch_size}." ) max_side_length = (max_teacher_patches // pooling_kernel_size**2) * side_mult if target_height == 0: target_height = side_mult target_width = min( int(math.floor(width / height)) * side_mult, max_side_length, ) elif target_width == 0: target_width = side_mult target_height = min( int(math.floor(height / width)) * side_mult, max_side_length, ) if target_height * target_width > target_px: raise ValueError( f"Resizing [{height}x{width}] to [{target_height}x{target_width}] " f"but this exceeds {max_teacher_patches} patches with patch_size {teacher_patch_size}" ) return target_height, target_width def convert_image_to_patches(image: torch.Tensor, patch_size: int) -> torch.Tensor: """ Purpose: Convert 3D tensor image of shape (num_channels, image_height, image_width) into 2D tensor of patches of shape (num_patches_height * num_patches_width, patch_size * patch_size * num_channels). Parameters: * image (torch.Tensor) : tensor representing an image. Has dimensions (num_channels, image_height, image_width) * patch_size (int) : length, in pixels, of one side of a patch Returns: 2D tensor of patches of shape (num_patches, flat_patch_size) = (num_patches_height * num_patches_width, patch_size * patch_size * num_channels) """ num_channels, image_height, image_width = image.shape num_patches_height = image_height // patch_size num_patches_width = image_width // patch_size patched_image = image.reshape(num_channels, num_patches_height, patch_size, num_patches_width, patch_size) patched_image = patched_image.permute(1, 3, 2, 4, 0) patched_image = patched_image.reshape(num_patches_height * num_patches_width, -1) return patched_image """ Explanation: flat_patches is a sequence of flattened patches. It's shape is (num_patches, flat_patch_size). positions is a sequence of [x, y] coordinate pairs. It's shape is (num_patches, 2). Assume we have flat_patches = [ [a, b, c, d], [e, f, g, h], [i, j, k, l] ] and positions = [ [0, 0], [0, 1], [1, 0] ]. Assume target_length = 8 current_length = flat_patches.shape[0] = 3 padding_length = target_length - current_length = 8 - 3 = 5 padding = [0, 0] * (image.ndim - 1) + [0, padding_length] = [0, 0] * (2 - 1) + [0, 5] = = [0, 0] + [0, 5] = [0, 0, 0, 5] pos_padding = (0, 0, 0, padding_length) = (0, 0, 0, 5) image = torch.nn.functional.pad(image, padding, mode="constant", value=0) = [ [a, b, c, d], [e, f, g, h], [i, j, k, l], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0] ] positions = torch.nn.functional.pad(positions, pos_padding, mode="constant", value=-1) = [ [0, 0], [0, 1], [1, 0], [-1, -1], [-1, -1], [-1, -1], [-1, -1], [-1, -1]] """ def pad_along_first_dim( flat_patches: torch.Tensor, positions: torch.Tensor, target_length: int ) -> tuple[torch.Tensor, torch.Tensor]: """ Purpose: Given a sequence of patch embeddings (flat_patches) and a sequence of flat patches' original 2D positions (positions), add padding to both so that they have specified length (target_length). Parameters: * flat_patches (torch.Tensor) : a 2D tensor of shape (num_patches, flat_patch_size) that contains a sequence of flattened patches * positions (torch.Tensor) : a 2D tensor of shape (num_patchs, flat_patch_size) s.t. positions[i] is the original 2D position of flattened patch flat_patches[i] Returns: A tuple (flat_patches, positions), where flat_patches has shape (target_length, flat_patch_size) and positions has shape (target_length, 2) """ current_length = flat_patches.shape[0] padding_length = target_length - current_length # the number of padding tokens to append if padding_length > 0: # padding and pos_padding are both tuples with 4 elements because the pad parameter of torch.nn.functional.pad() # is a tuple/list of even length giving padding amounts in pairs (left, right), ordered from the last dimension backward padding = [0, 0] * (flat_patches.ndim - 1) + [0, padding_length] pos_padding = (0, 0, 0, padding_length) # Specifying value=0 tells torch.nn.functional.pad() to pad flat_patches with vectors [0, ..., 0] # of size flat_patch_size flat_patches = torch.nn.functional.pad(flat_patches, padding, mode="constant", value=0) # Specifying value=0 tells torch.nn.functional.pad() to pad positions with vectors [-1, -1] positions = torch.nn.functional.pad(positions, pos_padding, mode="constant", value=-1) return flat_patches, positions def patches_merge( patches: torch.Tensor, positions_xy: torch.Tensor, length: int, ) -> tuple[torch.Tensor, torch.Tensor]: """ Purpose: Merge k×k groups of small patches into larger patches. Given `L` input patches of dimension `D = patch_size² × 3`, merge groups of `k×k` spatially adjacent patches into `length` output patches of dimension `(k × patch_size)² × 3`. The spatial grouping is determined by integer-dividing the XY positions by `k`. Parameters: small_patches: (*, L, D) — input patches. positions_xy: (*, L, 2) — integer XY positions for each patch (-1 for padding). length: target number of output patches. Must satisfy L = length × k². Returns: merged_patches: (*, length, k²×D) — merged patch features. merged_positions: (*, length, 2) — new XY positions for merged patches. """ patch_size = math.isqrt(patches.shape[-1] // 3) if patches.shape[-1] != patch_size * patch_size * 3: raise ValueError(f"Patch dimension {patches.shape[-1]} is not a valid `patch_size * patch_size * 3`") k = math.isqrt(patches.shape[-2] // length) if k * k * length != patches.shape[-2]: raise ValueError(f"Cannot merge {patches.shape} to {length}") # Compute target ordering for reordering patches into kernel-grouped order. # This ensures patches within each k×k kernel are contiguous. max_x = positions_xy[..., 0].max(dim=-1, keepdim=True)[0] + 1 kernel_idxs = torch.div(positions_xy, k, rounding_mode="floor") num_patches_from_top_left = k * k * kernel_idxs[..., 0] + k * max_x * kernel_idxs[..., 1] position_within_kernel = torch.remainder(positions_xy, k) num_patches_from_top_left_of_kernel = position_within_kernel[..., 0] + position_within_kernel[..., 1] * k target_ordering = num_patches_from_top_left_of_kernel + num_patches_from_top_left # Reorder patches by computing the inverse permutation via argsort, # then gathering patches into kernel-grouped order. perm = target_ordering.long().argsort(dim=-1) # inverse permutation # Expand perm indices to match patch feature dimension for gathering perm_expanded = perm.unsqueeze(-1).expand_as(patches) kernel_ordered_patches = patches.gather(-2, perm_expanded) batch_shape = patches.shape[:-2] # Reshape: (*, length*k*k, patch_size*patch_size*3) → (*, length, (k*patch_size)*(k*patch_size)*3) kernel_ordered_patches = kernel_ordered_patches.reshape(*batch_shape, length, k * k, patch_size, patch_size, 3) # Rearrange (l, a*b, p, q, c) → (l, a*p, b*q, c) kernel_ordered_patches = kernel_ordered_patches.reshape(*batch_shape, length, k, k, patch_size, patch_size, 3) kernel_ordered_patches = kernel_ordered_patches.permute( *range(len(batch_shape)), -6, -5, -3, -4, -2, -1 ) # (..., l, k, p, k, q, c) merged_patches = kernel_ordered_patches.reshape(*batch_shape, length, k * patch_size * k * patch_size * 3) # Compute new positions for merged patches perm_pos = perm.unsqueeze(-1).expand_as(positions_xy) kernel_ordered_positions = positions_xy.float().gather(-2, perm_pos.long()) # Handle padding: preserve -1 positions padding = (positions_xy == -1).all(dim=-1, keepdim=True) # (..., L, 1) kernel_ordered_positions = kernel_ordered_positions * (~padding).float() + positions_xy.float() * padding.float() # Reshape positions and take min within each kernel to get the merged position kernel_ordered_positions = kernel_ordered_positions.reshape(*batch_shape, length, k * k, 2) new_positions = torch.div(kernel_ordered_positions, k, rounding_mode="floor") # For each merged patch, take the minimum position across the kernel new_positions = new_positions.min(dim=-2)[0].to(torch.long) return merged_patches, new_positions class ImageProcessor: resample: PILImageResampling = PILImageResampling.BICUBIC do_resize: bool = True do_rescale: bool = True rescale_factor: float = 1 / 255 # uint8 [0, 255] -> float [0, 1] do_normalize: bool = False def __init__(self, cfg: VLMConfig) -> None: self.max_soft_tokens = cfg.max_soft_tokens self.pooling_kernel_size = cfg.pooling_kernel_size self.teacher_patch_size = cfg.teacher_patch_size self.max_teacher_patches = cfg.max_teacher_patches # maximum number of teacher patches a resized image may consist of def aspect_ratio_preserving_resize( self, image: torch.Tensor, ) -> torch.Tensor: """ Purpose: Resize image to preserve aspect ratio so it fits within the patch budget. Target dimensions are the largest that: 1) Produce at most `max_teacher_patches` patches when patchified with `patch_size` 2) Have height and width divisible by `pooling_kernel_size * patch_size` (i.e. size of model patch) Parameters: * image (torch.Tensor) : tensor representing an image. Has shape (C, H, W) Returns: Resized image - a tensor of shape (C, target_height, target_width) """ height, width = image.shape[-2], image.shape[-1] target_height, target_width = get_aspect_ratio_preserving_size( height=height, width=width, teacher_patch_size=self.teacher_patch_size, max_teacher_patches=self.max_teacher_patches, pooling_kernel_size=self.pooling_kernel_size, ) if target_height == height and target_width == width: return image return tvF.resize( image, size=[target_height, target_width], interpolation=self.resample, antialias=True, ) def preprocess( self, image: torch.Tensor, ) -> tuple: """ Purpose: Convert the image into a tensor of flattened model-size patches. Also returns a sequence of flattened patches' (x, y) positions and the number of model-sized patches. Parameters: * self * image (torch.Tensor) : a tensor of shape (C, H, W) Returns: A tuple with three elements: * merged_patches (torch.Tensor) - flattened model-sized patches. Has shape (num_image_tokens, model_patch_size²*3) * merged_positions (torch.Tensor) - flattened patches' 2D positions. Has shape (num_image_tokens, 2) * num_image_tokens (int) - the number of soft tokens in the image """ # Step 1: Aspect-ratio-preserving resize if self.do_resize: image = self.aspect_ratio_preserving_resize(image=image) # Step 2: Rescale pixel values (typically to [0, 1]) and optionally identity normalize image = image.to(torch.float32) * self.rescale_factor # Is this sufficient? # Step 3: Patchify into teacher-size patches (16px) # (num_channels, height, width) → (num_teacher_patches, patch_size²*3) patch_height = image.shape[-2] // self.teacher_patch_size patch_width = image.shape[-1] // self.teacher_patch_size teacher_patches = convert_image_to_patches(image, self.teacher_patch_size) # Step 4: Compute teacher-level position IDs device = image.device patch_grid = torch.meshgrid( torch.arange(patch_width, device=device), torch.arange(patch_height, device=device), indexing="xy", ) teacher_positions = torch.stack(patch_grid, dim=-1).reshape(teacher_patches.shape[0], 2) # Step 5: Merge k×k teacher patches into model patches via patches_merge # (num_teacher_patches, 768) → (num_model_patches, 6912) num_model_patches = teacher_patches.shape[0] // (self.pooling_kernel_size**2) merged_patches, merged_positions = patches_merge( teacher_patches.unsqueeze(0), # add batch dim for patches_merge teacher_positions.unsqueeze(0), num_model_patches, ) merged_patches = merged_patches.squeeze(0) # remove batch dim merged_positions = merged_positions.squeeze(0) num_image_tokens = merged_patches.shape[0] return merged_patches, merged_positions, num_image_tokens def __call__(self, image: Image.Image,) -> tuple: """ Purpose: Convert the image into a tensor of flattened model-size patches. Also returns a sequence of flattened patches' (x, y) positions and the number of model-sized patches. Parameters: * self * image (torch.Tensor) : a tensor of shape (C, H, W) Returns: A tuple with three elements: * merged_patches (torch.Tensor) - flattened model-sized patches. Has shape (num_image_tokens, model_patch_size²*3) * merged_positions (torch.Tensor) - flattened patches' 2D positions. Has shape (num_image_tokens, 2) * num_image_tokens (int) - the number of soft tokens in the image """ # Convert each image in the list into a tensor image_as_tensor = tvF.pil_to_tensor(image) # uint8, shape (B, C, H, W), values [0, 255] return self.preprocess(image_as_tensor)