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

mlx_vlm/models/paddleocr_vl/paddleocr_vl.py

@@ -0,0 +1,207 @@
+from typing import Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from ..base import InputEmbeddingsFeatures, install_auto_processor_patch
+from .config import ModelConfig
+from .language import LanguageModel
+from .processing_paddleocr_vl import PaddleOCRVLProcessor
+from .vision import VisionModel
+
+install_auto_processor_patch("paddleocr_vl", PaddleOCRVLProcessor)
+
+
+class Model(nn.Module):
+    def __init__(self, config: ModelConfig):
+        super().__init__()
+        self.config = config
+        self.visual = VisionModel(config.vision_config)
+        self.language_model = LanguageModel(config.text_config, config)
+
+    def get_input_embeddings(
+        self,
+        input_ids: Optional[mx.array] = None,
+        pixel_values: Optional[mx.array] = None,
+        **kwargs,
+    ):
+        image_grid_thw = kwargs.pop("image_grid_thw", None)
+        video_grid_thw = kwargs.pop("video_grid_thw", None)
+        mask = kwargs.pop("mask", None)
+        grid_thw = image_grid_thw if image_grid_thw is not None else video_grid_thw
+
+        if pixel_values is None:
+            # Reset position state for text-only generation
+            self.language_model._position_ids = None
+            self.language_model._rope_deltas = None
+            return InputEmbeddingsFeatures(
+                inputs_embeds=self.language_model.model.embed_tokens(input_ids)
+            )
+
+        dtype = self.visual.embeddings.patch_embedding.weight.dtype
+        pixel_values = mx.array(pixel_values, dtype=dtype)
+
+        # Get the input embeddings from the language model
+        inputs_embeds = self.language_model.model.embed_tokens(input_ids)
+
+        # Get the ouptut hidden states from the vision model
+        hidden_states = self.visual(pixel_values, grid_thw, output_hidden_states=False)
+
+        # Insert special image tokens in the input_ids
+        final_inputs_embeds = self.merge_input_ids_with_image_features(
+            self.config.image_token_id,
+            hidden_states,
+            inputs_embeds,
+            input_ids,
+        )
+
+        # Pre-calculate position_ids for chunked prefill
+        if image_grid_thw is not None or video_grid_thw is not None:
+            position_ids, rope_deltas = self.language_model.get_rope_index(
+                input_ids, image_grid_thw, video_grid_thw, mask
+            )
+            self.language_model._position_ids = position_ids
+            self.language_model._rope_deltas = rope_deltas
+
+        return InputEmbeddingsFeatures(inputs_embeds=final_inputs_embeds)
+
+    @staticmethod
+    def merge_input_ids_with_image_features(
+        image_token_id,
+        image_features,
+        inputs_embeds,
+        input_ids,
+    ):
+        """Merge image features into input embeddings at image token positions.
+
+        Args:
+            image_features: Vision features from the vision tower [num_features, hidden_dim]
+            inputs_embeds: Input embeddings [batch_size, seq_len, hidden_dim]
+            input_ids: Input token IDs [batch_size, seq_len]
+
+        Returns:
+            Updated input embeddings with image features inserted
+        """
+
+        # Positions of <image> tokens in input_ids
+        image_positions = input_ids == image_token_id
+
+        # Get dimensions
+        batch_size, seq_len = input_ids.shape
+
+        # Process each batch item
+        batch_outputs = []
+        feature_start_idx = 0
+
+        for batch_idx in range(batch_size):
+            # Get mask for this batch
+            image_mask = image_positions[batch_idx]
+            num_positions = mx.sum(image_mask).item()
+
+            if num_positions > 0:
+                # Extract features for this batch
+                batch_features = image_features[
+                    feature_start_idx : feature_start_idx + num_positions
+                ]
+
+                # Validate we have the right number of features
+                if batch_features.shape[0] != num_positions:
+                    raise ValueError(
+                        f"Number of image token positions ({num_positions}) does not match "
+                        f"number of image features ({batch_features.shape[0]}) for batch {batch_idx}"
+                    )
+
+                # Create indices for gathering
+                cumsum = mx.cumsum(image_mask.astype(mx.int32))
+                feature_indices = mx.where(image_mask, cumsum - 1, 0)
+
+                # Gather features
+                gathered_features = batch_features[feature_indices]
+
+                # Combine with original embeddings
+                image_mask_expanded = mx.expand_dims(image_mask, axis=-1)
+                batch_output = mx.where(
+                    image_mask_expanded, gathered_features, inputs_embeds[batch_idx]
+                )
+
+                feature_start_idx += num_positions
+            else:
+                # No image tokens in this batch item
+                batch_output = inputs_embeds[batch_idx]
+
+            batch_outputs.append(batch_output)
+
+        # Stack all batch outputs
+        return mx.stack(batch_outputs, axis=0)
+
+    @property
+    def layers(self):
+        return self.language_model.model.layers
+
+    def __call__(
+        self,
+        input_ids: mx.array,
+        pixel_values: Optional[mx.array] = None,
+        mask: Optional[mx.array] = None,
+        cache=None,
+        **kwargs,
+    ):
+
+        input_embeddings_features = self.get_input_embeddings(
+            input_ids, pixel_values, **kwargs
+        )
+        kwargs = {
+            "pixel_values": pixel_values,
+            **kwargs,
+        }
+        logits = self.language_model(
+            input_ids,
+            input_embeddings_features.inputs_embeds,
+            mask=mask,
+            cache=cache,
+            **kwargs,
+        )
+        return logits
+
+    def sanitize(self, weights):
+        _keys_to_ignore_on_load_unexpected = [
+            "packing_position_embedding",
+            "vision_model.head",
+        ]
+
+        def transform_key(key):
+            if "visual.vision_model" in key:
+                if "embeddings" in key or "post_layernorm" in key:
+                    key = key.replace("visual.vision_model", "visual")
+                elif "encoder" in key:
+                    key = key.replace("visual.vision_model.encoder", "visual")
+            elif "mlp_AR" in key:
+                key = key.replace("mlp_AR", "visual.projector")
+            elif "model" in key:
+                key = key.replace("model", "language_model.model")
+            elif "lm_head" in key:
+                key = key.replace("lm_head", "language_model.lm_head")
+
+            return key
+
+        new_weights = {}
+        for k, v in weights.items():
+            if (
+                "packing_position_embedding" in k
+                or "vision_model.head" in k
+                or ("visual" in k and "k_proj" in k)
+                or ("visual" in k and "v_proj" in k)
+            ):
+                continue
+            elif "visual" in k and "q_proj" in k:
+                new_key = transform_key(k)
+                k_proj = weights.get(k.replace("q_proj", "k_proj"), None)
+                v_proj = weights.get(k.replace("q_proj", "v_proj"), None)
+                if k_proj is not None and v_proj is not None:
+                    merged_tensor = mx.concatenate([v, k_proj, v_proj], axis=0)
+                    merged_key = new_key.replace("q_proj", "qkv")
+                    new_weights[merged_key] = merged_tensor
+            else:
+                new_weights[transform_key(k)] = v
+
+        return new_weights

mlx_vlm/models/paddleocr_vl/processing_paddleocr_vl.py

@@ -0,0 +1,425 @@
+import json
+import math
+from pathlib import Path
+from typing import List, Optional, Union
+
+import numpy as np
+from transformers import AutoTokenizer
+from transformers.feature_extraction_utils import BatchFeature
+from transformers.image_processing_utils import BaseImageProcessor
+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
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+
+def smart_resize(
+    height: int,
+    width: int,
+    factor: int,
+    min_pixels: int,
+    max_pixels: int,
+):
+    if height < factor:
+        width = round((width * factor) / height)
+        height = factor
+
+    if width < factor:
+        height = round((height * factor) / width)
+        width = factor
+
+    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 = math.floor(height / beta / factor) * factor
+        w_bar = 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
+
+
+class ImageProcessor(BaseImageProcessor):
+    """
+    MLX-native image processor for PaddleOCRVL that doesn't require torch.
+    """
+
+    model_input_names = ["pixel_values"]
+
+    def __init__(
+        self,
+        do_resize: bool = True,
+        size: dict[str, int] | None = None,
+        resample: PILImageResampling = PILImageResampling.BICUBIC,
+        do_rescale: bool = True,
+        rescale_factor: int | float = 1 / 255,
+        do_normalize: bool = True,
+        image_mean: float | list[float] | None = None,
+        image_std: float | list[float] | None = None,
+        do_convert_rgb: bool = True,
+        min_pixels: int = 147384,
+        max_pixels: int = 2822400,
+        patch_size: int = 14,
+        temporal_patch_size: int = 1,
+        merge_size: int = 2,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        if 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."
+                )
+        else:
+            size = {"shortest_edge": 147384, "longest_edge": 2822400}
+        # backward compatibility: override size with min_pixels and max_pixels if they are provided
+        if min_pixels is not None:
+            size["shortest_edge"] = min_pixels
+        if max_pixels is not None:
+            size["longest_edge"] = max_pixels
+        self.min_pixels = size["shortest_edge"]
+        self.max_pixels = size["longest_edge"]
+        self.size = size
+        self.do_resize = do_resize
+        self.resample = resample
+        self.do_rescale = do_rescale
+        self.rescale_factor = rescale_factor
+        self.do_normalize = do_normalize
+        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.patch_size = patch_size
+        self.temporal_patch_size = temporal_patch_size
+        self.merge_size = merge_size
+        self.do_convert_rgb = do_convert_rgb
+
+    def preprocess(
+        self,
+        images: ImageInput,
+        do_resize: Optional[bool] = None,
+        size: Optional[dict[str, int]] = None,
+        min_pixels: Optional[int] = None,
+        max_pixels: Optional[int] = None,
+        resample: Optional[PILImageResampling] = None,
+        do_rescale: Optional[bool] = None,
+        rescale_factor: Optional[float] = None,
+        do_normalize: Optional[bool] = None,
+        image_mean: Optional[Union[float, list[float]]] = None,
+        image_std: Optional[Union[float, list[float]]] = None,
+        patch_size: Optional[int] = None,
+        temporal_patch_size: Optional[int] = None,
+        merge_size: Optional[int] = None,
+        do_convert_rgb: Optional[bool] = None,
+        return_tensors: Optional[str] = None,
+        **kwargs,
+    ) -> BatchFeature:
+        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
+
+        if 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."
+                )
+        elif min_pixels is not None and max_pixels is not None:
+            # backward compatibility: override size with min_pixels and max_pixels if they are provided
+            size = {"shortest_edge": min_pixels, "longest_edge": max_pixels}
+
+        do_resize = do_resize if do_resize is not None else self.do_resize
+        resample = resample if resample is not None else self.resample
+        do_rescale = do_rescale if do_rescale is not None else self.do_rescale
+        rescale_factor = (
+            rescale_factor if rescale_factor is not None else self.rescale_factor
+        )
+        do_normalize = do_normalize if do_normalize is not None else self.do_normalize
+        image_mean = image_mean if image_mean is not None else self.image_mean
+        image_std = image_std if image_std is not None else self.image_std
+        patch_size = patch_size if patch_size is not None else self.patch_size
+        temporal_patch_size = (
+            temporal_patch_size
+            if temporal_patch_size is not None
+            else self.temporal_patch_size
+        )
+        merge_size = merge_size if merge_size is not None else self.merge_size
+        do_convert_rgb = (
+            do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb
+        )
+
+        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]
+
+        data = {}
+        pixel_values, vision_grid_thws = [], []
+        if images is not None:
+            processed_images = []
+            for image in images:
+                width, height = image.size
+                resized_height, resized_width = smart_resize(
+                    height,
+                    width,
+                    factor=patch_size * merge_size,
+                    min_pixels=min_pixels,
+                    max_pixels=max_pixels,
+                )
+                image = image.resize((resized_width, resized_height), resample)
+                img_array = to_numpy_array(image)
+
+                if do_rescale:
+                    img_array = img_array / 255.0
+
+                if do_normalize:
+                    mean = np.array(self.image_mean).reshape(1, 1, 3)
+                    std = np.array(self.image_std).reshape(1, 1, 3)
+                    img_array = (img_array - mean) / std
+
+                processed_images.append(img_array)
+
+            patches = np.array(processed_images)
+
+            if patches.shape[1] > 3:
+                patches = patches.transpose(0, 3, 1, 2)
+            if patches.shape[0] == 1:
+                patches = np.tile(patches, (temporal_patch_size, 1, 1, 1))
+
+            channel = patches.shape[1]
+            grid_t = patches.shape[0] // temporal_patch_size
+            grid_h, grid_w = (
+                resized_height // patch_size,
+                resized_width // patch_size,
+            )
+            patches = patches.reshape(
+                grid_t,
+                temporal_patch_size,
+                channel,
+                grid_h,
+                patch_size,
+                grid_w,
+                patch_size,
+            )
+            patches = patches.transpose(0, 3, 5, 2, 1, 4, 6)
+            if temporal_patch_size != 1:
+                raise ValueError(
+                    f"temporal_patch_size must be 1!, but got {temporal_patch_size}!"
+                )
+            flatten_patches = patches.reshape(
+                grid_t * grid_h * grid_w, channel, patch_size, patch_size
+            )
+            image_grid_thw = (grid_t, grid_h, grid_w)
+            pixel_values.extend(flatten_patches)
+            vision_grid_thws.append(image_grid_thw)
+
+        pixel_values = np.array([pixel_values])
+        vision_grid_thws = np.array(vision_grid_thws)
+        data.update({"pixel_values": pixel_values, "image_grid_thw": vision_grid_thws})
+
+        return BatchFeature(data, tensor_type=return_tensors)
+
+
+class PaddleOCRVLProcessor(ProcessorMixin):
+    attributes = ["image_processor", "tokenizer"]
+    valid_kwargs = ["chat_template"]
+    image_processor_class = "AutoImageProcessor"
+    tokenizer_class = "AutoTokenizer"
+
+    def __init__(
+        self,
+        image_processor=None,
+        tokenizer=None,
+        chat_template=None,
+        **kwargs,
+    ):
+
+        if image_processor is None:
+            image_processor = ImageProcessor(**kwargs)
+
+        self.tokenizer = tokenizer
+        self.image_token = (
+            "<|IMAGE_PLACEHOLDER|>"
+            if not hasattr(tokenizer, "image_token")
+            else tokenizer.image_token
+        )
+        self.image_processor = image_processor
+
+        super().__init__(image_processor, tokenizer, chat_template=chat_template)
+
+    def __call__(
+        self,
+        images=None,
+        text: Union[str, List[str]] = None,
+        **kwargs,
+    ) -> BatchFeature:
+        """Process images and text for the model.
+
+        Args:
+            images: Single image or list of images
+            text: Single text or list of texts
+            videos: Video inputs (not currently supported)
+            **kwargs: Additional arguments passed to tokenizer
+
+        Returns:
+            BatchFeature with:
+                - input_ids: Token IDs with image placeholders replaced
+                - attention_mask: Attention mask
+                - pixel_values: Processed image patches
+                - image_grid_thw: Grid dimensions for each image
+                - position_ids: 4D position IDs for xdrope
+        """
+        image_inputs = {}
+
+        if images is not None:
+            image_inputs = self.image_processor(images=images)
+            image_grid_thw = image_inputs["image_grid_thw"]
+
+        if text is None:
+            text = [""]
+        elif not isinstance(text, list):
+            text = [text]
+
+        text = [t for t in text]  # Copy to avoid modifying original
+
+        if images is not None:
+            index = 0
+            for i in range(len(text)):
+                while self.image_token in text[i]:
+                    text[i] = text[i].replace(
+                        self.image_token,
+                        "<|placeholder|>"
+                        * (
+                            image_grid_thw[index].prod()
+                            // self.image_processor.merge_size
+                            // self.image_processor.merge_size
+                        ),
+                        1,
+                    )
+                    index += 1
+                text[i] = text[i].replace("<|placeholder|>", self.image_token)
+
+        # Pop return_tensors to handle it ourselves at the end
+        return_tensors = kwargs.pop("return_tensors", None)
+
+        # Tokenize text
+        text_inputs = self.tokenizer(text, add_special_tokens=False, **kwargs)
+
+        # Get input_ids and convert to numpy array for processing
+        input_ids = text_inputs["input_ids"]
+        if hasattr(input_ids, "tolist"):
+            # Handle mlx arrays or torch tensors
+            input_ids = np.array(input_ids.tolist())
+        elif isinstance(input_ids, list):
+            input_ids = np.array(input_ids)
+
+        return BatchFeature(
+            data={**text_inputs, **image_inputs},
+            tensor_type=return_tensors,
+        )
+
+    def batch_decode(self, *args, **kwargs):
+        """Decode token IDs to text."""
+        return self.tokenizer.batch_decode(*args, **kwargs)
+
+    def decode(self, *args, **kwargs):
+        """Decode token IDs to text."""
+        return self.tokenizer.decode(*args, **kwargs)
+
+    def apply_chat_template(self, *args, **kwargs):
+        """Apply chat template using the tokenizer."""
+        return self.tokenizer.apply_chat_template(*args, **kwargs)
+
+    @property
+    def model_input_names(self):
+        """Return combined input names from tokenizer and image processor."""
+        tokenizer_input_names = (
+            self.tokenizer.model_input_names if self.tokenizer else []
+        )
+        image_processor_input_names = self.image_processor.model_input_names
+        return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names))
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+        """Load processor from pretrained model path."""
+        import warnings
+
+        from huggingface_hub import hf_hub_download
+
+        trust_remote_code = kwargs.pop("trust_remote_code", True)
+
+        model_path = Path(pretrained_model_name_or_path)
+        is_local = model_path.exists() and model_path.is_dir()
+
+        # Suppress warning about mrope_section in rope_parameters
+        with warnings.catch_warnings():
+            warnings.filterwarnings(
+                "ignore", message="Unrecognized keys in `rope_parameters`"
+            )
+            tokenizer = AutoTokenizer.from_pretrained(
+                str(model_path) if is_local else pretrained_model_name_or_path,
+                trust_remote_code=trust_remote_code,
+                local_files_only=is_local,
+                **kwargs,
+            )
+
+        # Load image processor config from preprocessor_config.json
+        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:
+                    preprocessor_config = json.load(f)
+                # Extract relevant image processor parameters
+                relevant_keys = [
+                    "min_pixels",
+                    "max_pixels",
+                    "patch_size",
+                    "temporal_patch_size",
+                    "merge_size",
+                    "image_mean",
+                    "image_std",
+                    "do_resize",
+                    "do_rescale",
+                    "do_normalize",
+                    "do_convert_rgb",
+                ]
+                for key in relevant_keys:
+                    if key in preprocessor_config:
+                        image_processor_config[key] = preprocessor_config[key]
+
+        except Exception:
+            pass
+
+        image_processor = ImageProcessor(**image_processor_config)
+        return cls(image_processor=image_processor, tokenizer=tokenizer, **kwargs)
+
+
+__all__ = [
+    "PaddleOCRVLProcessor",
+    "ImageProcessor",
+]