mlx-community/Molmo2-8B-4bit

1"""Image processor class for Molmo2"""2from typing import Optional, Union3import numpy as np4import einops5import torch6import torchvision.transforms7 8from transformers.image_utils import (9    IMAGENET_STANDARD_MEAN,10    IMAGENET_STANDARD_STD,11    ImageInput,12    PILImageResampling,13    make_flat_list_of_images,14    valid_images,15    to_numpy_array,16)17from transformers.image_transforms import convert_to_rgb18from transformers.processing_utils import ImagesKwargs19from transformers.image_processing_utils import BaseImageProcessor, get_size_dict20from transformers.utils import logging21from transformers.feature_extraction_utils import BatchFeature22from transformers.utils import TensorType, logging23 24 25logger = logging.get_logger(__name__)26 27 28def normalize_image(29    image: np.ndarray,30    image_mean: list[float],31    image_std: list[float],32) -> np.ndarray:33    image -= np.array(image_mean, dtype=np.float32)[None, None, :]34    image /= np.array(image_std, dtype=np.float32)[None, None, :]35    return image36 37 38def resize_image(39    image: np.ndarray,40    desired_output_size: list[int],41    resample: PILImageResampling,42) -> np.ndarray:43    image = torch.permute(torch.from_numpy(image), [2, 0, 1])44    dtype = image.dtype45    if torch.is_floating_point(image):46        in_min = 0.047        in_max = 1.048        resized = torchvision.transforms.Resize(49            desired_output_size,50            resample,51            antialias=False,52        )(image)53        resized = torch.clip(resized, 0.0, 1.0).to(dtype)54    else:55        assert image.dtype == torch.uint8, "SigLIP expects float images or uint8 images, but got {}".format(image.dtype)56        in_min = 0.057        in_max = 255.058        resized = torchvision.transforms.Resize(59            desired_output_size,60            resample,61            antialias=False,62        )(image)63        resized = torch.clip(resized, 0, 255).to(dtype)64 65    resized = resized.to(torch.float32)66    resized = (resized - in_min) / (in_max - in_min)67 68    resized = torch.permute(resized, [1, 2, 0]).numpy()69 70    return resized71 72 73def select_tiling(h, w, patch_size, max_num_crops):74    """Divide in image of size [w, h] in up to max_num_patches of size patch_size"""75    original_size = np.stack([h, w])  # [1, 2]76    original_res = h * w77    tilings = []78    for i in range(1, max_num_crops + 1):79        for j in range(1, max_num_crops + 1):80            if i*j <= max_num_crops:81                tilings.append((i, j))82    # sort so argmin and argmax favour smaller tilings in the event of a tie83    tilings.sort(key=lambda x: (x[0]*x[1], x[0]))84    candidate_tilings = np.array(tilings, dtype=np.int32)  # [n_resolutions, 2]85    candidate_resolutions = candidate_tilings * patch_size  # [n_resolutions, 2]86 87    # How much we would need to scale the image to fit exactly in each tiling88    original_size = np.stack([h, w], dtype=np.float32)  # [1, 2]89 90    # The original size can be zero in rare cases if the image is smaller than the margin91    # In those cases letting the scale become infinite means the tiling is based on the92    # other side, or falls back to the smallest tiling93    with np.errstate(divide='ignore'):94        required_scale_d = candidate_resolutions.astype(np.float32) / original_size,95    required_scale = np.min(required_scale_d, axis=-1, keepdims=True)  # [n_resolutions, 1]96    if np.all(required_scale < 1):97        # We are forced to downscale, so try to minimize the amount of downscaling98        ix = np.argmax(required_scale)99    else:100        # Pick the resolution that required the least upscaling so that it most closely fits the image101        required_scale = np.where(required_scale < 1.0, 10e9, required_scale)102        ix = np.argmin(required_scale)103    return candidate_tilings[ix]104 105 106def build_resized_image(107    image: np.ndarray,108    base_image_input_size: list[int],109    resample: PILImageResampling,110    image_mean: list[float],111    image_std: list[float],112    image_patch_size: int,113) -> tuple[np.ndarray, np.ndarray]:114    resized = resize_image(115        image, base_image_input_size, resample,116    )117    resized = normalize_image(resized, image_mean, image_std)118    if len(resized.shape) == 3:119        resized = np.expand_dims(resized, 0)120    crop_patch_w = base_image_input_size[1] // image_patch_size121    crop_patch_h = base_image_input_size[0] // image_patch_size122    resize_idx = np.arange(crop_patch_w*crop_patch_h).reshape([crop_patch_h, crop_patch_w])123    return resized, resize_idx124 125 126def build_overlapping_crops(127    image: np.ndarray,128    max_crops: int,129    overlap_margins: list[int],130    base_image_input_size: list[int],131    resample: PILImageResampling,132    image_mean: list[float],133    image_std: list[float],134    image_patch_size: int,135) -> tuple[np.ndarray, np.ndarray]:136    """Decompose an image into a set of overlapping crops137 138    :return crop_arr: [n_crops, h, w, 3] The crops139    :return patch_idx: [overlap_patch_h, overlap_patch_w] For each patch in the resized image140                        the crops were extracted from, what patch in `crop_arr` it corresponds to141    """142    original_image_h, original_image_w = image.shape[:2]143    crop_size = base_image_input_size[0]144    assert base_image_input_size[0] == base_image_input_size[1]145 146    left_margin, right_margin = overlap_margins147    total_margin_pixels = image_patch_size * (right_margin + left_margin)  # pixels removed per dim148    crop_patches = base_image_input_size[0] // image_patch_size  # patches per crop dim149    crop_window_patches = crop_patches - (right_margin + left_margin)  # usable patches150    crop_window_size = crop_window_patches * image_patch_size151    crop_patch_w = base_image_input_size[1] // image_patch_size152    crop_patch_h = base_image_input_size[0] // image_patch_size153    original_image_h, original_image_w = image.shape[:2]154    crop_size = base_image_input_size[0]155 156    # Decide how to tile the image, to account for the overlap margins we compute the tiling157    # as if we had an image without the margins and were using a crop size without the margins158    tiling = select_tiling(159        original_image_h - total_margin_pixels,160        original_image_w - total_margin_pixels,161        crop_window_size,162        max_crops,163    )164 165    src = resize_image(166        image,167        [tiling[0]*crop_window_size+total_margin_pixels, tiling[1]*crop_window_size+total_margin_pixels],168        resample,169    )170    src = normalize_image(src, image_mean, image_std)171 172    # Now we have to split the image into crops, and track what patches came from173    # where in `patch_idx_arr`174    n_crops = tiling[0] * tiling[1]175    crop_arr = np.zeros([n_crops, crop_size, crop_size, 3], dtype=src.dtype)176    patch_idx_arr = np.zeros([n_crops, crop_patch_h, crop_patch_w], dtype=np.int32)177    on_crop = 0178    for i in range(tiling[0]):179        # Slide over `src` by `crop_window_size` steps, but extract crops of size `crops_size`180        # which results in overlapping crop windows181        y0 = i*crop_window_size182        for j in range(tiling[1]):183            x0 = j*crop_window_size184            crop_arr[on_crop] = src[y0:y0+crop_size, x0:x0+crop_size]185            patch_idx = np.arange(crop_patch_w*crop_patch_h).reshape(crop_patch_h, crop_patch_w)186            patch_idx += on_crop * crop_patch_h * crop_patch_w187 188            # Mask out idx that are in the overlap region189            if i != 0:190                patch_idx[:left_margin, :] = -1191            if j != 0:192                patch_idx[:, :left_margin] = -1193            if i != tiling[0]-1:194                patch_idx[-right_margin:, :] = -1195            if j != tiling[1]-1:196                patch_idx[:, -right_margin:] = -1197            patch_idx_arr[on_crop] = patch_idx198            on_crop += 1199 200    # `patch_idx_arr` is ordered crop-by-crop, here we transpose `patch_idx_arr`201    # so it is ordered left-to-right order202    patch_idx_arr = np.reshape(203        patch_idx_arr,204        [tiling[0], tiling[1], crop_patch_h, crop_patch_w]205    )206    patch_idx_arr = np.transpose(patch_idx_arr, [0, 2, 1, 3])207    patch_idx_arr = np.reshape(patch_idx_arr, [-1])208 209    # Now get the parts not in the overlap region, so it should map each patch in `src`210    # to the correct patch it should come from in `crop_arr`211    patch_idx_arr = patch_idx_arr[patch_idx_arr >= 0].reshape(212        src.shape[0]//image_patch_size,213        src.shape[1]//image_patch_size,214    )215    return crop_arr, patch_idx_arr216 217 218def batch_pixels_to_patches(array: np.ndarray, patch_size: int) -> np.ndarray:219    """Reshape images of [n_images, h, w, 3] -> [n_images, n_patches, pixels_per_patch]"""220    if len(array.shape) == 3:221        n_crops, h, w = array.shape222        h_patches = h//patch_size223        w_patches = w//patch_size224        array = np.reshape(array, [n_crops, h_patches, patch_size, w_patches, patch_size])225        array = np.transpose(array, [0, 1, 3, 2, 4])226        array = np.reshape(array, [n_crops, h_patches*w_patches, patch_size*patch_size])227        return array228    else:229        n_crops, h, w, c = array.shape230        h_patches = h//patch_size231        w_patches = w//patch_size232        array = np.reshape(array, [n_crops, h_patches, patch_size, w_patches, patch_size, c])233        array = np.transpose(array, [0, 1, 3, 2, 4, 5])234        array = np.reshape(array, [n_crops, h_patches*w_patches, patch_size*patch_size*c])235        return array236 237 238def arange_for_pooling(239    idx_arr: np.ndarray,240    pool_h: int,241    pool_w: int,242) -> np.ndarray:243    h_pad = pool_h * ((idx_arr.shape[0] + pool_h - 1) // pool_h) - idx_arr.shape[0]244    w_pad = pool_w * ((idx_arr.shape[1] + pool_w - 1) // pool_w) - idx_arr.shape[1]245    idx_arr = np.pad(idx_arr, [[h_pad//2, (h_pad+1)//2], [w_pad//2, (w_pad+1)//2]],246                     mode='constant',constant_values=-1)247    return einops.rearrange(248        idx_arr, "(h dh) (w dw) -> h w (dh dw)", dh=pool_h, dw=pool_w)249 250 251def image_to_patches_and_grids(252    image: np.ndarray,253    max_crops: int,254    overlap_margins: list[int],255    base_image_input_size: list[int],256    resample: PILImageResampling,257    image_mean: list[float],258    image_std: list[float],259    image_patch_size: int,260    image_pooling_w: int,261    image_pooling_h: int,262) -> tuple[np.ndarray, np.ndarray, np.ndarray]:263    """264    :return image_grids, the shape of each (low-res, high-res) image after pooling265    :return crops, the image crops to processes with the ViT266    :return pooled_patch_idx, for each patch_id tokens in `image_tokens`, the indices of the267                                patches in `crops` to pool for that token, masked with -1268    """269    if isinstance(base_image_input_size, int):270        base_image_input_size = (base_image_input_size, base_image_input_size)271    272    base_image_input_d = image_patch_size273    pooling_w = image_pooling_w274    pooling_h = image_pooling_h275    crop_patch_w = base_image_input_size[1] // base_image_input_d276    crop_patch_h = base_image_input_size[0] // base_image_input_d277 278    crop_arr, patch_idx_arr = build_overlapping_crops(279        image,280        max_crops,281        overlap_margins,282        base_image_input_size,283        resample,284        image_mean,285        image_std,286        image_patch_size,287    )288    pooling_idx = arange_for_pooling(patch_idx_arr, pooling_h, pooling_w)289    h, w = pooling_idx.shape[:2]290    pooling_idx = pooling_idx.reshape([-1, pooling_h*pooling_w])291    292    # Finally do the same for the global image293    resized, resize_idx = build_resized_image(294        image,295        base_image_input_size,296        resample,297        image_mean,298        image_std,299        image_patch_size,300    )301    crop_arr = np.concatenate([resized, crop_arr], 0)302 303    resize_idx = arange_for_pooling(resize_idx, pooling_h, pooling_w)304    resized_h, resized_w = resize_idx.shape[:2]305    resize_idx = resize_idx.reshape([-1, pooling_h*pooling_w])306 307    # Global image goes first, so the order of patches in previous crops gets increased308    pooling_idx = np.where(309        pooling_idx >= 0,310        pooling_idx + crop_patch_h*crop_patch_w,311        -1312    )313    pooling_idx = np.concatenate([resize_idx, pooling_idx])314    image_grid = [np.array([resized_h, resized_w, h, w])]315 316    return (317        np.stack(image_grid, 0),318        batch_pixels_to_patches(crop_arr, image_patch_size),319        pooling_idx320    )321 322 323class Molmo2ImagesKwargs(ImagesKwargs, total=False):324    max_crops: Optional[int]325    overlap_margins: Optional[list[int]]326    patch_size: Optional[int]327    pooling_size: Optional[list[int]]328 329 330class Molmo2ImageProcessor(BaseImageProcessor):331    r"""332    Constructs a Molmo2 image processor that preprocesses images for the model.333 334    Args:335        size (`dict[str, int]` *optional*, defaults to `{"height": 378, "width": 378}`):336            Size of the image after resizing.337        resample (`PILImageResampling`, *optional*, defaults to `Resampling.BILINEAR`):338            Resampling filter to use when resizing the image.339        image_mean (`float` or `list[float]`, *optional*, defaults to `[0.5, 0.5, 0.5]`):340            Mean to use if normalizing the image. This is a float or list of floats for each channel in the image.341        image_std (`float` or `list[float]`, *optional*, defaults to `[0.5, 0.5, 0.5]`):342            Standard deviation to use if normalizing the image. This is a float or list of floats for each channel in the image.343        do_convert_rgb (`bool`, *optional*, defaults to `True`):344            Whether to convert the image to RGB.345        max_crops (`int`, *optional*, defaults to `8`):346            Maximum number of crops to use per image.347        overlap_margins (`list[int]`, *optional*, defaults to `[4, 4]`):348            Overlap margins to use.349        patch_size (`int`, *optional*, defaults to 14):350            The spatial patch size of the vision encoder.351        pooling_size (`list[int]`, *optional*, defaults to `[2, 2]`):352            The pooling size of the vision adapter.353    """354 355    model_input_names = ["pixel_values", "image_token_pooling", "image_grids", "image_num_crops"]356 357    def __init__(358        self,359        size: Optional[dict[str, int]] = None,360        resample: PILImageResampling = PILImageResampling.BILINEAR,361        image_mean: Optional[Union[float, list[float]]] = None,362        image_std: Optional[Union[float, list[float]]] = None,363        do_convert_rgb: bool = True,364        max_crops: int = 8,365        overlap_margins: list[int] = [4, 4],366        patch_size: int = 14,367        pooling_size: list[int] = [2, 2],368        **kwargs,369    ) -> None:370        super().__init__(**kwargs)371        size = size if size is not None else {"height": 378, "width": 378}372        size = get_size_dict(size, default_to_square=True)373        self.size = size374 375        self.resample = resample376        self.image_mean = image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN377        self.image_std = image_std if image_std is not None else IMAGENET_STANDARD_STD378        self.do_convert_rgb = do_convert_rgb379 380        self.max_crops = max_crops381        self.overlap_margins = overlap_margins382        self.patch_size = patch_size383        self.pooling_size = pooling_size384    385    def preprocess(386        self,387        images: ImageInput,388        size: Optional[dict[str, int]] = None,389        resample: Optional[PILImageResampling] = None,390        image_mean: Optional[Union[float, list[float]]] = None,391        image_std: Optional[Union[float, list[float]]] = None,392        do_convert_rgb: Optional[bool] = None,393        max_crops: Optional[int] = None,394        overlap_margins: Optional[list[int]] = None,395        patch_size: Optional[int] = None,396        pooling_size: Optional[list[int]] = None,397        return_tensors: Optional[Union[str, TensorType]] = None,398        **kwargs,399    ) -> BatchFeature:400        """401        Args:402            images (`ImageInput`):403                Image to preprocess.404            size (`dict[str, int]`, *optional*, defaults to `self.size`):405                Size of the image after resizing.406            resample (`PILImageResampling`, *optional*, defaults to `self.resample`):407                Resampling filter to use when resizing the image. This can be one of the enum `PILImageResampling`. Only408                has an effect if `do_resize` is set to `True`.409            image_mean (`float` or `list[float]`, *optional*, defaults to `self.image_mean`):410                Image mean to use for normalization. Only has an effect if `do_normalize` is set to `True`.411            image_std (`float` or `list[float]`, *optional*, defaults to `self.image_std`):412                Image standard deviation to use for normalization. Only has an effect if `do_normalize` is set to413                `True`.414            do_convert_rgb (`bool`, *optional*, defaults to `self.do_convert_rgb`):415                Whether to convert the image to RGB.416            max_crops (`int`, *optional*, defaults to `self.max_crops`):417                Maximum number of crops to use per image.418            overlap_margins (`list[int]`, *optional*, defaults to `self.overlap_margins`):419                Overlap margins to use.420            patch_size (`int`, *optional*, defaults to `self.patch_size`):421                The spatial patch size of the vision encoder.422            pooling_size (`list[int]`, *optional*, defaults to `self.pooling_size`):423                The pooling size of the vision adapter.424            return_tensors (`str` or `TensorType`, *optional*):425                The type of tensors to return. Can be one of:426                - Unset: Return a list of `np.ndarray`.427                - `TensorType.TENSORFLOW` or `'tf'`: Return a batch of type `tf.Tensor`.428                - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`.429                - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`.430                - `TensorType.JAX` or `'jax'`: Return a batch of type `jax.numpy.ndarray`.431 432        Returns:433            A `BatchFeature` containing the following keys:434                - `pixel_values`: The preprocessed images.435                - `image_token_pooling`: The indices of the patches in `crops` to pool for each token in `image_tokens`.436                - `image_grids`: The image grids.437                - `image_num_crops`: The number of crops for each image.438        """439        if size is not None:440            if "height" not in size or "width" not in size:441                raise ValueError("size must contain 'height' and 'width' keys.")442        else:443            size = {**self.size}444        445        base_image_input_size = [size["height"], size["width"]]446        447        resample = resample or self.resample448        image_mean = image_mean or self.image_mean449        image_std = image_std or self.image_std450        do_convert_rgb = do_convert_rgb or self.do_convert_rgb451 452        max_crops = max_crops or self.max_crops453        overlap_margins = overlap_margins or self.overlap_margins454        patch_size = patch_size or self.patch_size455        pooling_size = pooling_size or self.pooling_size456 457        image_pooling_h, image_pooling_w = pooling_size458 459        if images is not None:460            images = self.fetch_images(images)461            images = make_flat_list_of_images(images)462        463        if images is not None and not valid_images(images):464            raise ValueError(465                "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "466                "torch.Tensor, tf.Tensor or jax.ndarray."467            )468 469        if do_convert_rgb:470            images = [convert_to_rgb(image) for image in images]471 472        # All transformations expect numpy arrays.473        images = [to_numpy_array(image) for image in images]474 475        data = {}476        if images is not None:477            batch_grids = []478            batch_crops = []479            batch_pooled_patches_idx = []480            batch_num_crops = []481 482            for image in images:483                image_grid, crops, pooled_idx = image_to_patches_and_grids(484                    image,485                    max_crops,486                    overlap_margins,487                    base_image_input_size,488                    resample,489                    image_mean,490                    image_std,491                    patch_size,492                    image_pooling_w,493                    image_pooling_h,494                )495                batch_grids.append(image_grid)496                batch_crops.append(crops)497                batch_pooled_patches_idx.append(pooled_idx)498                batch_num_crops.append(crops.shape[0])499            500            pixel_values = np.concatenate(batch_crops, 0)501            image_token_pooling = np.concatenate(batch_pooled_patches_idx, 0)502            image_grids = np.concatenate(batch_grids, 0)503            image_num_crops = np.array(batch_num_crops)504 505            data.update(506                pixel_values=pixel_values,507                image_token_pooling=image_token_pooling,508                image_grids=image_grids,509                image_num_crops=image_num_crops,510            )511 512        return BatchFeature(data, tensor_type=return_tensors)513 514 515Molmo2ImageProcessor.register_for_auto_class()516