fount-vlm-nell-02 0.3.11__py3-none-any.whl

mlx_vlm/models/molmo2/processing.py

@@ -0,0 +1,773 @@
+"""Image processor for Molmo2 - MLX-native implementation without torch dependency."""
+
+import logging
+from typing import List, Optional, Tuple
+
+import numpy as np
+from PIL import Image
+from transformers import AutoProcessor
+from transformers.feature_extraction_utils import BatchFeature
+from transformers.image_processing_utils import BaseImageProcessor, get_size_dict
+from transformers.image_transforms import convert_to_rgb
+from transformers.image_utils import (
+    ImageInput,
+    PILImageResampling,
+    make_flat_list_of_images,
+    to_numpy_array,
+    valid_images,
+)
+from transformers.processing_utils import ProcessorMixin
+
+logger = logging.getLogger(__name__)
+
+# Special tokens
+IMAGE_PATCH_TOKEN = "<im_patch>"
+IMAGE_LOW_RES_TOKEN = "<im_low>"
+IM_START_TOKEN = "<im_start>"
+LOW_RES_IMAGE_START_TOKEN = "<low_res_im_start>"
+FRAME_START_TOKEN = "<frame_start>"
+IM_END_TOKEN = "<im_end>"
+FRAME_END_TOKEN = "<frame_end>"
+IM_COL_TOKEN = "<im_col>"
+IMAGE_PROMPT = "<|image|>"
+
+
+def normalize_image(
+    image: np.ndarray,
+    image_mean: List[float],
+    image_std: List[float],
+) -> np.ndarray:
+    """Normalize image with mean and std."""
+    image = image.astype(np.float32)
+    image -= np.array(image_mean, dtype=np.float32)[None, None, :]
+    image /= np.array(image_std, dtype=np.float32)[None, None, :]
+    return image
+
+
+def resize_image_pil(
+    image: np.ndarray,
+    desired_output_size: List[int],
+    resample: PILImageResampling,
+) -> np.ndarray:
+    """Resize image using PIL instead of torch."""
+    # Convert numpy to PIL
+    if image.max() <= 1.0:
+        image_uint8 = (image * 255).astype(np.uint8)
+    else:
+        image_uint8 = image.astype(np.uint8)
+
+    pil_image = Image.fromarray(image_uint8)
+
+    # Map PILImageResampling to PIL resampling
+    resample_map = {
+        PILImageResampling.NEAREST: Image.Resampling.NEAREST,
+        PILImageResampling.BILINEAR: Image.Resampling.BILINEAR,
+        PILImageResampling.BICUBIC: Image.Resampling.BICUBIC,
+        PILImageResampling.LANCZOS: Image.Resampling.LANCZOS,
+        PILImageResampling.BOX: Image.Resampling.BOX,
+        PILImageResampling.HAMMING: Image.Resampling.HAMMING,
+    }
+    pil_resample = resample_map.get(resample, Image.Resampling.BILINEAR)
+
+    # Resize (PIL uses width, height order)
+    resized = pil_image.resize(
+        (desired_output_size[1], desired_output_size[0]), resample=pil_resample
+    )
+
+    # Convert back to numpy and normalize to [0, 1]
+    resized_np = np.array(resized, dtype=np.float32) / 255.0
+    return resized_np
+
+
+def select_tiling(
+    h: int, w: int, patch_size: int, max_num_crops: int
+) -> Tuple[int, int]:
+    """Divide an image of size [w, h] into up to max_num_crops of size patch_size."""
+    tilings = []
+    for i in range(1, max_num_crops + 1):
+        for j in range(1, max_num_crops + 1):
+            if i * j <= max_num_crops:
+                tilings.append((i, j))
+    tilings.sort(key=lambda x: (x[0] * x[1], x[0]))
+    candidate_tilings = np.array(tilings, dtype=np.int32)
+    candidate_resolutions = candidate_tilings * patch_size
+
+    original_size = np.array([h, w], dtype=np.float32)
+
+    with np.errstate(divide="ignore"):
+        required_scale = np.min(
+            candidate_resolutions.astype(np.float32) / original_size, axis=-1
+        )
+
+    if np.all(required_scale < 1):
+        ix = np.argmax(required_scale)
+    else:
+        required_scale = np.where(required_scale < 1.0, 1e9, required_scale)
+        ix = np.argmin(required_scale)
+
+    return tuple(candidate_tilings[ix])
+
+
+def build_resized_image(
+    image: np.ndarray,
+    base_image_input_size: List[int],
+    resample: PILImageResampling,
+    image_mean: List[float],
+    image_std: List[float],
+    image_patch_size: int,
+) -> Tuple[np.ndarray, np.ndarray]:
+    """Build a single resized image crop."""
+    resized = resize_image_pil(image, base_image_input_size, resample)
+    resized = normalize_image(resized, image_mean, image_std)
+    if len(resized.shape) == 3:
+        resized = np.expand_dims(resized, 0)
+    crop_patch_w = base_image_input_size[1] // image_patch_size
+    crop_patch_h = base_image_input_size[0] // image_patch_size
+    resize_idx = np.arange(crop_patch_w * crop_patch_h).reshape(
+        [crop_patch_h, crop_patch_w]
+    )
+    return resized, resize_idx
+
+
+def build_overlapping_crops(
+    image: np.ndarray,
+    max_crops: int,
+    overlap_margins: List[int],
+    base_image_input_size: List[int],
+    resample: PILImageResampling,
+    image_mean: List[float],
+    image_std: List[float],
+    image_patch_size: int,
+) -> Tuple[np.ndarray, np.ndarray]:
+    """Decompose an image into overlapping crops."""
+    crop_size = base_image_input_size[0]
+    assert base_image_input_size[0] == base_image_input_size[1]
+
+    left_margin, right_margin = overlap_margins
+    total_margin_pixels = image_patch_size * (right_margin + left_margin)
+    crop_patches = base_image_input_size[0] // image_patch_size
+    crop_window_patches = crop_patches - (right_margin + left_margin)
+    crop_window_size = crop_window_patches * image_patch_size
+    crop_patch_w = base_image_input_size[1] // image_patch_size
+    crop_patch_h = base_image_input_size[0] // image_patch_size
+    original_image_h, original_image_w = image.shape[:2]
+
+    tiling = select_tiling(
+        original_image_h - total_margin_pixels,
+        original_image_w - total_margin_pixels,
+        crop_window_size,
+        max_crops,
+    )
+
+    src = resize_image_pil(
+        image,
+        [
+            tiling[0] * crop_window_size + total_margin_pixels,
+            tiling[1] * crop_window_size + total_margin_pixels,
+        ],
+        resample,
+    )
+    src = normalize_image(src, image_mean, image_std)
+
+    n_crops = tiling[0] * tiling[1]
+    crop_arr = np.zeros([n_crops, crop_size, crop_size, 3], dtype=src.dtype)
+    patch_idx_arr = np.zeros([n_crops, crop_patch_h, crop_patch_w], dtype=np.int32)
+
+    on_crop = 0
+    for i in range(tiling[0]):
+        y0 = i * crop_window_size
+        for j in range(tiling[1]):
+            x0 = j * crop_window_size
+            crop_arr[on_crop] = src[y0 : y0 + crop_size, x0 : x0 + crop_size]
+            patch_idx = np.arange(crop_patch_w * crop_patch_h).reshape(
+                crop_patch_h, crop_patch_w
+            )
+            patch_idx += on_crop * crop_patch_h * crop_patch_w
+
+            if i != 0:
+                patch_idx[:left_margin, :] = -1
+            if j != 0:
+                patch_idx[:, :left_margin] = -1
+            if i != tiling[0] - 1:
+                patch_idx[-right_margin:, :] = -1
+            if j != tiling[1] - 1:
+                patch_idx[:, -right_margin:] = -1
+            patch_idx_arr[on_crop] = patch_idx
+            on_crop += 1
+
+    # Transpose the patch_idx_arr to get the full index array
+    patch_idx_full = np.zeros(
+        [
+            tiling[0] * crop_window_patches + left_margin + right_margin,
+            tiling[1] * crop_window_patches + left_margin + right_margin,
+        ],
+        dtype=np.int32,
+    )
+    for i in range(tiling[0]):
+        for j in range(tiling[1]):
+            crop_idx = i * tiling[1] + j
+            y_start = i * crop_window_patches
+            x_start = j * crop_window_patches
+            patch_idx_full[
+                y_start : y_start + crop_patch_h, x_start : x_start + crop_patch_w
+            ] = np.where(
+                patch_idx_arr[crop_idx] >= 0,
+                patch_idx_arr[crop_idx],
+                patch_idx_full[
+                    y_start : y_start + crop_patch_h, x_start : x_start + crop_patch_w
+                ],
+            )
+
+    return crop_arr, patch_idx_full, tiling
+
+
+def batch_pixels_to_patches(crops: np.ndarray, patch_size: int) -> np.ndarray:
+    """Convert image crops to patches."""
+    n_crops, h, w, c = crops.shape
+    n_patches_h = h // patch_size
+    n_patches_w = w // patch_size
+    n_patches = n_patches_h * n_patches_w
+    patch_dim = patch_size * patch_size * c
+
+    # Reshape to patches
+    crops = crops.reshape(n_crops, n_patches_h, patch_size, n_patches_w, patch_size, c)
+    crops = crops.transpose(0, 1, 3, 2, 4, 5)
+    crops = crops.reshape(n_crops, n_patches, patch_dim)
+    return crops
+
+
+def arange_for_pooling(
+    idx_arr: np.ndarray,
+    pool_h: int,
+    pool_w: int,
+) -> np.ndarray:
+    """Build pooling indices using centered padding (matches HuggingFace implementation)."""
+    h, w = idx_arr.shape
+    # Calculate padding to make dimensions divisible by pool size (centered padding)
+    h_pad = pool_h * ((h + pool_h - 1) // pool_h) - h
+    w_pad = pool_w * ((w + pool_w - 1) // pool_w) - w
+
+    # Apply centered padding
+    idx_arr = np.pad(
+        idx_arr,
+        [[h_pad // 2, (h_pad + 1) // 2], [w_pad // 2, (w_pad + 1) // 2]],
+        mode="constant",
+        constant_values=-1,
+    )
+
+    # Rearrange into pooling windows: (h dh) (w dw) -> h w (dh dw)
+    padded_h, padded_w = idx_arr.shape
+    out_h = padded_h // pool_h
+    out_w = padded_w // pool_w
+
+    # Reshape to separate pooling dimensions
+    idx_arr = idx_arr.reshape(out_h, pool_h, out_w, pool_w)
+    # Transpose to get (out_h, out_w, pool_h, pool_w)
+    idx_arr = idx_arr.transpose(0, 2, 1, 3)
+    # Reshape to (out_h, out_w, pool_h * pool_w)
+    idx_arr = idx_arr.reshape(out_h, out_w, pool_h * pool_w)
+
+    return idx_arr
+
+
+def image_to_patches_and_grids(
+    image: np.ndarray,
+    max_crops: int,
+    overlap_margins: List[int],
+    base_image_input_size: List[int],
+    resample: PILImageResampling,
+    image_mean: List[float],
+    image_std: List[float],
+    image_patch_size: int,
+    image_pooling_w: int,
+    image_pooling_h: int,
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Tuple[int, int]]:
+    """Convert image to patches with pooling information.
+
+    Returns crops and pooling indices in the order expected by the model:
+    - Crops: [low_res_crop, high_res_crops...] (low-res first)
+    - Pooling indices: [low_res_pooling, high_res_pooling] (low-res first)
+    - Image grid: [lo_h, lo_w, hi_h, hi_w] (low-res dimensions first)
+    """
+    crop_patch_w = base_image_input_size[1] // image_patch_size
+    crop_patch_h = base_image_input_size[0] // image_patch_size
+
+    # Build overlapping crops for high-res
+    crop_arr, patch_idx_arr, tiling = build_overlapping_crops(
+        image,
+        max_crops,
+        overlap_margins,
+        base_image_input_size,
+        resample,
+        image_mean,
+        image_std,
+        image_patch_size,
+    )
+
+    # Build pooling indices for high-res using centered padding (matches HF)
+    pooling_idx = arange_for_pooling(patch_idx_arr, image_pooling_h, image_pooling_w)
+    h, w = pooling_idx.shape[:2]
+    pooling_idx = pooling_idx.reshape(-1, image_pooling_h * image_pooling_w)
+
+    # Build resized image for low-res
+    resize_arr, resize_idx = build_resized_image(
+        image,
+        base_image_input_size,
+        resample,
+        image_mean,
+        image_std,
+        image_patch_size,
+    )
+
+    # Combine crops: LOW-RES FIRST (matches HuggingFace)
+    all_crops = np.concatenate([resize_arr, crop_arr], axis=0)
+
+    # Build pooling indices for low-res
+    resize_pooling_idx = arange_for_pooling(
+        resize_idx, image_pooling_h, image_pooling_w
+    )
+    resized_h, resized_w = resize_pooling_idx.shape[:2]
+    resize_pooling_idx = resize_pooling_idx.reshape(
+        -1, image_pooling_h * image_pooling_w
+    )
+
+    # LOW-RES crop is first, so offset HIGH-RES indices by the number of low-res patches
+    # (matches HuggingFace: "Global image goes first, so the order of patches in previous crops gets increased")
+    pooling_idx = np.where(
+        pooling_idx >= 0,
+        pooling_idx
+        + crop_patch_h * crop_patch_w,  # Offset by one crop (the low-res crop)
+        -1,
+    )
+
+    # Concatenate pooling indices: LOW-RES FIRST (matches HuggingFace)
+    pooling_idx = np.concatenate([resize_pooling_idx, pooling_idx], axis=0)
+
+    # Image grid format: [resized_h, resized_w, h, w] = [lo_h, lo_w, hi_h, hi_w]
+    # (matches HuggingFace order)
+    image_grid = np.array([[resized_h, resized_w, h, w]], dtype=np.int32)
+
+    return (
+        image_grid,
+        batch_pixels_to_patches(all_crops, image_patch_size),
+        pooling_idx,
+        (h, w),  # Return high-res pooled dims for token generation
+    )
+
+
+class Molmo2ImageProcessor(BaseImageProcessor):
+    """
+    MLX-native image processor for Molmo2 that doesn't require torch.
+    """
+
+    model_input_names = [
+        "pixel_values",
+        "image_token_pooling",
+        "image_grids",
+        "image_num_crops",
+    ]
+
+    def __init__(
+        self,
+        size: Optional[dict] = None,
+        resample: PILImageResampling = PILImageResampling.BILINEAR,
+        image_mean: Optional[List[float]] = None,
+        image_std: Optional[List[float]] = None,
+        do_convert_rgb: bool = True,
+        max_crops: int = 8,
+        overlap_margins: List[int] = None,
+        patch_size: int = 14,
+        pooling_size: List[int] = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        size = size if size is not None else {"height": 378, "width": 378}
+        size = get_size_dict(size, default_to_square=True)
+        self.size = size
+
+        self.resample = resample
+        self.image_mean = image_mean if image_mean is not None else [0.5, 0.5, 0.5]
+        self.image_std = image_std if image_std is not None else [0.5, 0.5, 0.5]
+        self.do_convert_rgb = do_convert_rgb
+
+        self.max_crops = max_crops
+        self.overlap_margins = (
+            overlap_margins if overlap_margins is not None else [4, 4]
+        )
+        self.patch_size = patch_size
+        self.pooling_size = pooling_size if pooling_size is not None else [2, 2]
+
+    def preprocess(
+        self,
+        images: ImageInput,
+        size: Optional[dict] = None,
+        resample: Optional[PILImageResampling] = None,
+        image_mean: Optional[List[float]] = None,
+        image_std: Optional[List[float]] = None,
+        do_convert_rgb: Optional[bool] = None,
+        max_crops: Optional[int] = None,
+        overlap_margins: Optional[List[int]] = None,
+        patch_size: Optional[int] = None,
+        pooling_size: Optional[List[int]] = None,
+        return_tensors: Optional[str] = None,
+        **kwargs,
+    ) -> BatchFeature:
+        """Preprocess images for Molmo2."""
+        if size is not None:
+            if "height" not in size or "width" not in size:
+                raise ValueError("size must contain 'height' and 'width' keys.")
+        else:
+            size = {**self.size}
+
+        base_image_input_size = [size["height"], size["width"]]
+
+        resample = resample or self.resample
+        image_mean = image_mean or self.image_mean
+        image_std = image_std or self.image_std
+        do_convert_rgb = (
+            do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb
+        )
+
+        max_crops = max_crops or self.max_crops
+        overlap_margins = overlap_margins or self.overlap_margins
+        patch_size = patch_size or self.patch_size
+        pooling_size = pooling_size or self.pooling_size
+
+        image_pooling_h, image_pooling_w = pooling_size
+
+        if images is not None:
+            images = make_flat_list_of_images(images)
+
+        if images is not None and not valid_images(images):
+            raise ValueError(
+                "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
+                "torch.Tensor, tf.Tensor or jax.ndarray."
+            )
+
+        if do_convert_rgb:
+            images = [convert_to_rgb(image) for image in images]
+
+        images = [to_numpy_array(image) for image in images]
+
+        data = {}
+        if images is not None:
+            batch_grids = []
+            batch_crops = []
+            batch_pooled_patches_idx = []
+            batch_num_crops = []
+            batch_hi_pooled_dims = []
+
+            for image in images:
+                image_grid, crops, pooled_idx, hi_pooled_dims = (
+                    image_to_patches_and_grids(
+                        image,
+                        max_crops,
+                        overlap_margins,
+                        base_image_input_size,
+                        resample,
+                        image_mean,
+                        image_std,
+                        patch_size,
+                        image_pooling_w,
+                        image_pooling_h,
+                    )
+                )
+                batch_grids.append(image_grid)
+                batch_crops.append(crops)
+                batch_pooled_patches_idx.append(pooled_idx)
+                batch_num_crops.append(crops.shape[0])
+                batch_hi_pooled_dims.append(hi_pooled_dims)
+
+            pixel_values = np.concatenate(batch_crops, 0)
+            image_token_pooling = np.concatenate(batch_pooled_patches_idx, 0)
+            image_grids = np.concatenate(batch_grids, 0)
+            image_num_crops = np.array(batch_num_crops)
+
+            data.update(
+                pixel_values=pixel_values,
+                image_token_pooling=image_token_pooling,
+                image_grids=image_grids,
+                image_num_crops=image_num_crops,
+                _hi_pooled_dims=batch_hi_pooled_dims,  # Internal use for token generation
+            )
+
+        return BatchFeature(data, tensor_type=return_tensors)
+
+
+class Molmo2Processor(ProcessorMixin):
+    """
+    Processor for Molmo2 that combines image processor and tokenizer.
+    """
+
+    attributes = ["image_processor", "tokenizer"]
+    image_processor_class = "Molmo2ImageProcessor"
+    tokenizer_class = "AutoTokenizer"
+
+    def __init__(
+        self,
+        image_processor=None,
+        tokenizer=None,
+        image_use_col_tokens: bool = True,
+        use_single_crop_col_tokens: Optional[bool] = None,
+        use_single_crop_start_token: bool = True,
+        **kwargs,
+    ):
+        if image_processor is None:
+            image_processor = Molmo2ImageProcessor()
+        super().__init__(image_processor, tokenizer, **kwargs)
+        self.image_use_col_tokens = image_use_col_tokens
+        self.use_single_crop_col_tokens = use_single_crop_col_tokens
+        self.use_single_crop_start_token = use_single_crop_start_token
+        self.image_placeholder_token = IMAGE_PROMPT
+
+    def get_image_tokens(self, image_grid: np.ndarray) -> str:
+        """Generate image token string from image grid.
+
+        Args:
+            image_grid: Array of [resized_h, resized_w, height, width]
+                        = [lo_pooled_h, lo_pooled_w, hi_pooled_h, hi_pooled_w]
+                        (matches HuggingFace format)
+
+        Returns:
+            String of image tokens to insert into prompt.
+        """
+        # Unpack in HuggingFace order: [lo_h, lo_w, hi_h, hi_w]
+        resized_h, resized_w, height, width = image_grid
+
+        # Build high-res tokens first (will be appended after low-res)
+        per_row = [IMAGE_PATCH_TOKEN] * width
+        if self.image_use_col_tokens:
+            per_row = per_row + [IM_COL_TOKEN]
+        hi_res_tokens = [IM_START_TOKEN] + per_row * height + [IM_END_TOKEN]
+
+        # Build low-res tokens
+        per_row = [IMAGE_PATCH_TOKEN] * resized_w
+        use_single_crop_col_tokens = (
+            self.image_use_col_tokens
+            if self.use_single_crop_col_tokens is None
+            else self.use_single_crop_col_tokens
+        )
+        image_start_token = (
+            LOW_RES_IMAGE_START_TOKEN
+            if self.use_single_crop_start_token
+            else IM_START_TOKEN
+        )
+        if use_single_crop_col_tokens:
+            per_row = per_row + [IM_COL_TOKEN]
+        lo_res_tokens = [image_start_token] + per_row * resized_h + [IM_END_TOKEN]
+
+        # Low-res comes first, then high-res (matches HuggingFace)
+        all_tokens = lo_res_tokens + hi_res_tokens
+        return "".join(all_tokens)
+
+    def __call__(
+        self,
+        text=None,
+        images=None,
+        padding=False,
+        truncation=None,
+        max_length=None,
+        return_tensors=None,
+        **kwargs,
+    ):
+        """Process text and images for the model."""
+        encoding = {}
+        image_grids = None
+
+        if images is not None:
+            image_inputs = self.image_processor(images, return_tensors=None)
+            image_grids = image_inputs.get("image_grids")
+            # Remove internal keys before adding to encoding
+            hi_pooled_dims = image_inputs.pop("_hi_pooled_dims", None)
+            encoding.update(image_inputs)
+
+        if text is not None:
+            # Expand image placeholders with actual image tokens
+            if image_grids is not None:
+                if isinstance(text, str):
+                    text = [text]
+                    was_string = True
+                else:
+                    text = list(text)
+                    was_string = False
+
+                image_idx = 0
+                for i in range(len(text)):
+                    num_images = text[i].count(self.image_placeholder_token)
+                    for _ in range(num_images):
+                        if image_idx < len(image_grids):
+                            image_tokens = self.get_image_tokens(image_grids[image_idx])
+                            text[i] = text[i].replace(
+                                self.image_placeholder_token, image_tokens, 1
+                            )
+                            image_idx += 1
+
+                if was_string:
+                    text = text[0]
+
+            text_inputs = self.tokenizer(
+                text,
+                padding=padding,
+                truncation=truncation,
+                max_length=max_length,
+                return_tensors=return_tensors,
+                **kwargs,
+            )
+            encoding.update(text_inputs)
+
+        # Convert to requested tensor type
+        if return_tensors is not None:
+            encoding = BatchFeature(encoding, tensor_type=return_tensors)
+
+        return encoding
+
+    def apply_chat_template(
+        self,
+        conversation,
+        chat_template=None,
+        add_generation_prompt=False,
+        tokenize=False,
+        **kwargs,
+    ):
+        """Apply chat template to conversation."""
+        if chat_template is None:
+            chat_template = getattr(self, "chat_template", None)
+        if chat_template is None:
+            chat_template = getattr(self.tokenizer, "chat_template", None)
+        if chat_template is None:
+            # Default Molmo2 chat template
+            chat_template = (
+                "{% for message in messages %}"
+                "{% if message['role'] == 'user' %}"
+                "User: {{ message['content'] }}\n"
+                "{% elif message['role'] == 'assistant' %}"
+                "Assistant: {{ message['content'] }}\n"
+                "{% endif %}"
+                "{% endfor %}"
+                "{% if add_generation_prompt %}Assistant: {% endif %}"
+            )
+
+        from jinja2 import Environment
+
+        # Use Environment with loopcontrols extension to support {% continue %} and {% break %}
+        env = Environment(extensions=["jinja2.ext.loopcontrols"])
+        template = env.from_string(chat_template)
+        rendered = template.render(
+            messages=conversation,
+            add_generation_prompt=add_generation_prompt,
+            **kwargs,
+        )
+
+        if tokenize:
+            return self.tokenizer.encode(rendered)
+        return rendered
+
+    def batch_decode(self, *args, **kwargs):
+        return self.tokenizer.batch_decode(*args, **kwargs)
+
+    def decode(self, *args, **kwargs):
+        return self.tokenizer.decode(*args, **kwargs)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+        """Load processor from pretrained model."""
+        import json
+        from pathlib import Path
+
+        from huggingface_hub import hf_hub_download
+        from transformers import AutoTokenizer
+
+        kwargs.pop("trust_remote_code", None)
+
+        model_path = Path(pretrained_model_name_or_path)
+        is_local = model_path.exists() and model_path.is_dir()
+
+        # Load tokenizer
+        tokenizer = AutoTokenizer.from_pretrained(
+            str(model_path) if is_local else pretrained_model_name_or_path,
+            trust_remote_code=True,
+            local_files_only=is_local,
+        )
+
+        # Load image processor config
+        image_processor_config = {}
+        try:
+            if is_local:
+                config_path = model_path / "preprocessor_config.json"
+            else:
+                config_path = Path(
+                    hf_hub_download(
+                        pretrained_model_name_or_path, "preprocessor_config.json"
+                    )
+                )
+            if config_path.exists():
+                with open(config_path, "r", encoding="utf-8") as f:
+                    config = json.load(f)
+                for key in [
+                    "size",
+                    "max_crops",
+                    "overlap_margins",
+                    "patch_size",
+                    "pooling_size",
+                    "image_mean",
+                    "image_std",
+                    "do_convert_rgb",
+                ]:
+                    if key in config:
+                        image_processor_config[key] = config[key]
+        except Exception:
+            pass
+
+        image_processor = Molmo2ImageProcessor(**image_processor_config)
+
+        # Load chat template
+        chat_template = getattr(tokenizer, "chat_template", None)
+
+        return cls(
+            image_processor=image_processor,
+            tokenizer=tokenizer,
+        )
+
+
+# Patch AutoProcessor for Molmo2 models
+import json
+from pathlib import Path
+
+_original_auto_processor_from_pretrained_molmo2 = AutoProcessor.from_pretrained
+
+
+@classmethod
+def _patched_auto_processor_from_pretrained_molmo2(
+    cls, pretrained_model_name_or_path, **kwargs
+):
+    """Patched from_pretrained that returns Molmo2Processor for molmo2 models."""
+    from huggingface_hub import hf_hub_download
+
+    model_path = Path(pretrained_model_name_or_path)
+    is_local = model_path.exists() and model_path.is_dir()
+
+    # Check if this is a molmo2 model
+    is_molmo2 = False
+    try:
+        if is_local:
+            config_path = model_path / "config.json"
+        else:
+            config_path = Path(
+                hf_hub_download(pretrained_model_name_or_path, "config.json")
+            )
+        with open(config_path, "r", encoding="utf-8") as f:
+            config = json.load(f)
+        model_type = config.get("model_type", "").lower()
+        is_molmo2 = model_type == "molmo2"
+    except Exception:
+        pass
+
+    if is_molmo2:
+        return Molmo2Processor.from_pretrained(pretrained_model_name_or_path, **kwargs)
+
+    return _original_auto_processor_from_pretrained_molmo2.__func__(
+        cls, pretrained_model_name_or_path, **kwargs
+    )
+
+
+AutoProcessor.from_pretrained = _patched_auto_processor_from_pretrained_molmo2