""" HyperCLOVAX Omni Image Processor (Fast) Implements dynamic resolution image processing: - Smart resize: adjusts image to fit within min_pixels and max_pixels - Vision token calculation: token reduction using merge_size - Discrete image processing: separate processing for discrete vision tokens Based on BaseImageProcessorFast with torchvision resize. """ import math import os from typing import List, Optional, Tuple, Union import torch from torchvision.transforms.v2 import functional as F from transformers.image_processing_utils import BatchFeature from transformers.image_processing_utils_fast import ( BaseImageProcessorFast, DefaultFastImageProcessorKwargs, group_images_by_shape, reorder_images, ) from transformers.image_utils import ( OPENAI_CLIP_MEAN, OPENAI_CLIP_STD, ImageInput, PILImageResampling, SizeDict, ) from transformers.processing_utils import ImagesKwargs, Unpack from transformers.utils import TensorType, logging logger = logging.get_logger(__name__) def smart_resize( height: int, width: int, factor: int = 28, min_pixels: int = 56 * 56, max_pixels: int = 14 * 14 * 4 * 1280, ) -> Tuple[int, int]: """Smart resize for dynamic resolution. Adjusts image dimensions to satisfy: 1. Both dimensions are divisible by factor. 2. Total pixel count is between min_pixels and max_pixels. Args: height: Original image height. width: Original image width. factor: Rounding unit (default: 28 = patch_size * merge_size). min_pixels: Minimum pixel count (default: 3136). max_pixels: Maximum pixel count (default: 1003520). Returns: Tuple of (new_height, new_width). """ if max(height, width) / min(height, width) > 200: raise ValueError( f"absolute aspect ratio must be smaller than 200, got {max(height, width) / min(height, width)}" ) h_bar = round(height / factor) * factor w_bar = round(width / factor) * factor if h_bar * w_bar > max_pixels: beta = math.sqrt((height * width) / max_pixels) h_bar = max(factor, math.floor(height / beta / factor) * factor) w_bar = max(factor, math.floor(width / beta / factor) * factor) elif h_bar * w_bar < min_pixels: beta = math.sqrt(min_pixels / (height * width)) h_bar = math.ceil(height * beta / factor) * factor w_bar = math.ceil(width * beta / factor) * factor return h_bar, w_bar def calculate_num_patches( height: int, width: int, patch_size: int = 14, merge_size: int = 2, ) -> int: """Calculate the number of vision tokens. Args: height: Image height (should be divisible by patch_size * merge_size). width: Image width (should be divisible by patch_size * merge_size). patch_size: ViT patch size (default: 14). merge_size: Merge size for token reduction (default: 2). Returns: Number of vision tokens. """ grid_h = height // patch_size grid_w = width // patch_size num_patches = (grid_h // merge_size) * (grid_w // merge_size) return num_patches class HyperCLOVAXOmniFastImageProcessorKwargs(DefaultFastImageProcessorKwargs, total=False): min_pixels: Optional[int] max_pixels: Optional[int] patch_size: Optional[int] temporal_patch_size: Optional[int] merge_size: Optional[int] # Token parameters image_token: Optional[str] image_start_token: Optional[str] image_end_token: Optional[str] # Discrete image parameters discrete_image_size: Optional[int] discrete_token_size: Optional[int] discrete_image_ratios: Optional[List] discrete_image_token: Optional[str] discrete_image_start_token: Optional[str] discrete_image_end_token: Optional[str] use_discrete_image_token: Optional[bool] vision_eol_token: Optional[str] vision_eof_token: Optional[str] class HyperCLOVAXOmniImageProcessor(BaseImageProcessorFast): """Fast image processor for HyperCLOVAX Omni. Uses torchvision-based resize for dynamic resolution processing: 1. Smart resize: adjusts image size to be within min_pixels and max_pixels. 2. Vision token calculation: uses merge_size for token reduction. 3. Discrete image processing: separate processing for discrete vision tokens. """ # Class-level defaults resample = PILImageResampling.BICUBIC image_mean = OPENAI_CLIP_MEAN image_std = OPENAI_CLIP_STD do_resize = True do_rescale = True do_normalize = True do_convert_rgb = True size = {"shortest_edge": 3136, "longest_edge": 2073600} default_to_square = False min_pixels = 3136 max_pixels = 2073600 patch_size = 14 temporal_patch_size = 2 merge_size = 2 model_input_names = ["pixel_values"] valid_kwargs = HyperCLOVAXOmniFastImageProcessorKwargs def __init__(self, **kwargs): # Handle size <-> min_pixels/max_pixels size = kwargs.pop("size", None) min_pixels = kwargs.pop("min_pixels", None) max_pixels = kwargs.pop("max_pixels", None) size = {**self.size} if size is None else size if min_pixels is not None: size["shortest_edge"] = min_pixels size.pop("min_pixels", None) if max_pixels is not None: size["longest_edge"] = max_pixels size.pop("max_pixels", None) if "shortest_edge" not in size or "longest_edge" not in size: raise ValueError("size must contain 'shortest_edge' and 'longest_edge' keys.") # Set defaults for discrete parameters (not in valid_kwargs) kwargs.setdefault("image_token", "<|IMAGE_PAD|>") kwargs.setdefault("image_start_token", "<|image_start|>") kwargs.setdefault("image_end_token", "<|image_end|>") kwargs.setdefault("discrete_image_size", 384) kwargs.setdefault("discrete_token_size", 27) if "discrete_image_ratios" not in kwargs or kwargs["discrete_image_ratios"] is None: kwargs["discrete_image_ratios"] = [] kwargs.setdefault("discrete_image_token", "<|DISCRETE_IMAGE_PAD|>") kwargs.setdefault("discrete_image_start_token", "<|discrete_image_start|>") kwargs.setdefault("discrete_image_end_token", "<|discrete_image_end|>") kwargs.setdefault("use_discrete_image_token", False) kwargs.setdefault("vision_eol_token", "<|vision_eol|>") kwargs.setdefault("vision_eof_token", "<|vision_eof|>") super().__init__(size=size, min_pixels=min_pixels, max_pixels=max_pixels, **kwargs) # Ensure min_pixels/max_pixels are always set from size if self.min_pixels is None: self.min_pixels = self.size["shortest_edge"] if self.max_pixels is None: self.max_pixels = self.size["longest_edge"] # Build ratio -> token mapping from discrete_image_ratios ratios = self.discrete_image_ratios if self.discrete_image_ratios is not None else [] self.discrete_image_ratio_tokens = { f"{r[0]}:{r[1]}": f"<|vision_ratio_{r[0]}:{r[1]}|>" for r in ratios } def _further_process_kwargs( self, size: Optional[SizeDict] = None, min_pixels: Optional[int] = None, max_pixels: Optional[int] = None, **kwargs, ) -> dict: """Synchronize size <-> min_pixels/max_pixels.""" if min_pixels is not None and max_pixels is not None: size = {"shortest_edge": min_pixels, "longest_edge": max_pixels} elif size is not None: if "shortest_edge" not in size or "longest_edge" not in size: raise ValueError("size must contain 'shortest_edge' and 'longest_edge' keys.") min_pixels = size["shortest_edge"] max_pixels = size["longest_edge"] else: size = {**self.size} return super()._further_process_kwargs(size=size, min_pixels=min_pixels, max_pixels=max_pixels, **kwargs) def _find_best_ratio_token( self, original_size: List[int], discrete_image_ratios: Optional[List[List[int]]] = None, ) -> Tuple[int, int]: """Find the best ratio token based on the original image aspect ratio. Args: original_size: Original [height, width] of the image. discrete_image_ratios: List of [h, w] ratio pairs. Defaults to self.discrete_image_ratios. Returns: Best matching (h_ratio, w_ratio) tuple. """ discrete_image_ratios = discrete_image_ratios if discrete_image_ratios is not None else self.discrete_image_ratios if not discrete_image_ratios: return (1, 1) h, w = original_size if h == 0 or w == 0: return (1, 1) ratios = [i / j for i, j in discrete_image_ratios] diffs = [abs(w / h - r) for r in ratios] best_size_idx = diffs.index(min(diffs)) return discrete_image_ratios[best_size_idx] def _preprocess_continuous_image( self, images: list, do_resize: bool, size: SizeDict, interpolation: Optional[F.InterpolationMode], do_rescale: bool, rescale_factor: float, do_normalize: bool, image_mean: Optional[Union[float, list]], image_std: Optional[Union[float, list]], patch_size: int, temporal_patch_size: int, merge_size: int, disable_grouping: Optional[bool], ) -> dict: """Preprocess images for continuous vision features. Performs smart resize -> rescale+normalize -> patchify. Args: images: List of image tensors to preprocess. do_resize: Whether to perform resizing. size: SizeDict containing min_pixels/max_pixels. interpolation: Interpolation method. do_rescale: Whether to perform rescaling. rescale_factor: Rescale factor. do_normalize: Whether to perform normalization. image_mean: Normalization mean. image_std: Normalization standard deviation. patch_size: ViT patch size. temporal_patch_size: Temporal patch size. merge_size: Token reduction merge size. disable_grouping: Whether to disable image grouping. Returns: Dictionary with: - "pixel_values": Tensor of shape (N, num_patches, patch_dim). - "image_grid_thw": Tensor of shape (N, 3). - "num_image_tokens": Tensor of shape (N,) with per-image token counts. """ # 1. Group & smart resize grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping) resized_images_grouped = {} for shape, stacked_images in grouped_images.items(): height, width = stacked_images.shape[-2:] if do_resize: resized_height, resized_width = smart_resize( height, width, factor=patch_size * merge_size, min_pixels=size["shortest_edge"], max_pixels=size["longest_edge"], ) stacked_images = self.resize( image=stacked_images, size=SizeDict(height=resized_height, width=resized_width), interpolation=interpolation, ) resized_images_grouped[shape] = stacked_images resized_images = reorder_images(resized_images_grouped, grouped_images_index) # 2. Group again -> fused rescale+normalize -> patchify grouped_images, grouped_images_index = group_images_by_shape(resized_images, disable_grouping=disable_grouping) processed_images_grouped = {} processed_grids = {} for shape, stacked_images in grouped_images.items(): resized_height, resized_width = stacked_images.shape[-2:] patches = self.rescale_and_normalize( stacked_images, do_rescale, rescale_factor, do_normalize, image_mean, image_std ) # Add temporal dimension for images (ndim == 4 means no temporal dim) if patches.ndim == 4: patches = patches.unsqueeze(1) # Pad temporal dimension to be divisible by temporal_patch_size if patches.shape[1] % temporal_patch_size != 0: repeats = patches[:, -1:].repeat(1, temporal_patch_size - 1, 1, 1, 1) patches = torch.cat([patches, repeats], dim=1) batch_size, grid_t, channel = patches.shape[:3] grid_t = grid_t // temporal_patch_size grid_h, grid_w = resized_height // patch_size, resized_width // patch_size # Patchify: reshape -> permute -> flatten patches = patches.view( batch_size, grid_t, temporal_patch_size, channel, grid_h // merge_size, merge_size, patch_size, grid_w // merge_size, merge_size, patch_size, ) patches = patches.permute(0, 1, 4, 7, 5, 8, 3, 2, 6, 9) flatten_patches = patches.reshape( batch_size, grid_t * grid_h * grid_w, channel * temporal_patch_size * patch_size * patch_size, ) processed_images_grouped[shape] = flatten_patches processed_grids[shape] = [[grid_t, grid_h, grid_w]] * batch_size processed_images = reorder_images(processed_images_grouped, grouped_images_index) processed_grids = reorder_images(processed_grids, grouped_images_index) pixel_values = torch.cat(processed_images, dim=0) image_grid_thw = torch.tensor(processed_grids) num_image_tokens = image_grid_thw.prod(dim=1) // (merge_size ** 2) return { "pixel_values": pixel_values, "image_grid_thw": image_grid_thw, "num_image_tokens": num_image_tokens, } def _preprocess_discrete_image( self, images: list, original_sizes: List[Tuple[int, int]], interpolation: Optional[F.InterpolationMode], ) -> dict: """Preprocess images for discrete vision tokens. Resizes each image to a fixed size (discrete_image_size) and finds the closest aspect ratio token. Args: images: List of image tensors to preprocess. original_sizes: List of (height, width) tuples for each image. interpolation: Interpolation method. Returns: Dictionary with: - "discrete_pixel_values": Tensor of shape (N, C, discrete_image_size, discrete_image_size). - "discrete_image_ratios": Tensor of shape (N, 2). - "num_discrete_image_tokens": Tensor of shape (N,) with per-image discrete token counts. """ discrete_pixel_values_list = [] discrete_image_ratios_list = [] for i, img in enumerate(images): orig_h, orig_w = original_sizes[i] best_ratio = self._find_best_ratio_token([orig_h, orig_w]) # Resize to fixed discrete_image_size x discrete_image_size (torchvision) discrete_img = F.resize( img.unsqueeze(0), [self.discrete_image_size, self.discrete_image_size], interpolation=interpolation, antialias=True, ) discrete_img = discrete_img.squeeze(0) # Match torchvision to_tensor: float32 / 255.0 (no normalize) discrete_img = discrete_img.to(torch.float32) / 255.0 discrete_pixel_values_list.append(discrete_img) discrete_image_ratios_list.append(best_ratio) n = len(images) discrete_token_size = self.discrete_token_size # ratio_token(1) + discrete_token_size rows * (discrete_token_size tokens + vision_eol(1)) + vision_eof(1) num_discrete_per_image = 1 + discrete_token_size * (discrete_token_size + 1) + 1 return { "discrete_pixel_values": torch.stack(discrete_pixel_values_list), "discrete_image_ratios": torch.tensor(discrete_image_ratios_list), "num_discrete_image_tokens": torch.full((n,), num_discrete_per_image, dtype=torch.long), } def _preprocess( self, images: list, do_resize: bool, size: SizeDict, interpolation: Optional[F.InterpolationMode], do_rescale: bool, rescale_factor: float, do_normalize: bool, image_mean: Optional[Union[float, list]], image_std: Optional[Union[float, list]], patch_size: int, temporal_patch_size: int, merge_size: int, disable_grouping: Optional[bool], return_tensors: Optional[Union[str, TensorType]], **kwargs, ) -> BatchFeature: """Main preprocessing entry point called by BaseImageProcessorFast. Performs continuous image processing and optionally discrete image processing. Args: images: List of image tensors. do_resize: Whether to perform resizing. size: SizeDict containing min_pixels/max_pixels. interpolation: Interpolation method. do_rescale: Whether to perform rescaling. rescale_factor: Rescale factor. do_normalize: Whether to perform normalization. image_mean: Normalization mean. image_std: Normalization standard deviation. patch_size: ViT patch size. temporal_patch_size: Temporal patch size. merge_size: Token reduction merge size. disable_grouping: Whether to disable image grouping. return_tensors: Desired tensor type for outputs. Returns: BatchFeature containing pixel_values, image_grid_thw, and optionally discrete processing results. """ # Record original sizes for discrete processing before any transforms if self.use_discrete_image_token: original_sizes = [(img.shape[-2], img.shape[-1]) for img in images] # Continuous processing continuous_result = self._preprocess_continuous_image( images, do_resize=do_resize, size=size, interpolation=interpolation, do_rescale=do_rescale, rescale_factor=rescale_factor, do_normalize=do_normalize, image_mean=image_mean, image_std=image_std, patch_size=patch_size, temporal_patch_size=temporal_patch_size, merge_size=merge_size, disable_grouping=disable_grouping, ) data = continuous_result # Discrete processing if self.use_discrete_image_token: discrete_result = self._preprocess_discrete_image( images, original_sizes=original_sizes, interpolation=interpolation, ) data.update(discrete_result) return BatchFeature(data=data, tensor_type=return_tensors) def get_num_image_tokens( self, image_width: Optional[int] = None, image_height: Optional[int] = None, pixel_values: Optional[torch.Tensor] = None, include_boundary_tokens: bool = False, min_pixels: Optional[int] = None, max_pixels: Optional[int] = None, patch_size: Optional[int] = None, merge_size: Optional[int] = None, return_tuple: Optional[bool] = None, ) -> int: """Compute the number of image tokens for the given input. Args: image_width: Image width (used when pixel_values is None). image_height: Image height (used when pixel_values is None). pixel_values: Pre-computed pixel values tensor. include_boundary_tokens: Whether to include start/end boundary tokens. min_pixels: Minimum pixel count. Defaults to self.min_pixels. max_pixels: Maximum pixel count. Defaults to self.max_pixels. patch_size: ViT patch size. Defaults to self.patch_size. merge_size: Token reduction merge size. Defaults to self.merge_size. return_tuple: If True, return (continuous, discrete) tuple. Otherwise return the sum. Returns: Token count as int, or (continuous, discrete) tuple if return_tuple is True. """ patch_size = patch_size if patch_size is not None else self.patch_size merge_size = merge_size if merge_size is not None else self.merge_size min_pixels = min_pixels if min_pixels is not None else self.min_pixels max_pixels = max_pixels if max_pixels is not None else self.max_pixels num_continuous_tokens, num_discrete_tokens = 0, 0 if pixel_values is None: factor = patch_size * merge_size resized_height, resized_width = smart_resize( image_height, image_width, factor=factor, min_pixels=min_pixels, max_pixels=max_pixels ) num_patches = calculate_num_patches(resized_height, resized_width, patch_size, merge_size) num_continuous_tokens = num_patches // (merge_size ** 2) elif len(pixel_values.shape) == 2: num_continuous_tokens = pixel_values.shape[0] // (merge_size ** 2) else: num_continuous_tokens = sum( pv.shape[0] // (merge_size ** 2) for pv in pixel_values ) if include_boundary_tokens: num_continuous_tokens += 2 if self.use_discrete_image_token: discrete_token_size = self.discrete_token_size num_discrete_tokens = discrete_token_size ** 2 if include_boundary_tokens: num_discrete_tokens += 2 if return_tuple: return (num_continuous_tokens, num_discrete_tokens) else: return num_continuous_tokens + num_discrete_tokens def save_pretrained( self, save_directory: Union[str, os.PathLike], *args, **kwargs, ) -> None: """Save the processor to a directory. Registers for auto class before saving. Args: save_directory: Directory path to save the processor. """ self.register_for_auto_class() super().save_pretrained(save_directory, *args, **kwargs)