mlx-community/FastVLM-0.5B-bf16

1import re2import torch3from transformers import ProcessorMixin, BatchFeature, CLIPImageProcessorFast4from transformers.image_processing_utils import BaseImageProcessor5from transformers.image_utils import ImageInput6from typing import Any, Dict, List, Optional, Union7from PIL import Image8 9from .llava_qwen import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN10 11# Adapted from transformers.models.llava_next.image_processing_llava_next.expand_to_square12def expand_to_square(image: torch.Tensor, background_color=0) -> torch.Tensor:13    """14    Expands an image to a square by adding a background color.15    """16    c, height, width = image.shape17    if width == height:18        return image19    elif width > height:20        result = torch.ones((c, width, width), dtype=image.dtype) * background_color21        result[:, (width - height) // 2 : (width - height) // 2 + height, :] = image22        return result23    else:24        result = torch.ones((c, height, height), dtype=image.dtype) * background_color25        result[:, :, (height - width) // 2 : (height - width) // 2 + width] = image26        return result27 28 29class FastVLMImageProcessor(CLIPImageProcessorFast):30    def _preprocess(self, images, **kwargs):31        image_sizes = [image.shape[-2:][::-1] for image in images]32        images = [expand_to_square(image) for image in images]33        images = super()._preprocess(images, **kwargs)34        pixel_values = torch.stack(images.pixel_values, dim=0)35        return BatchFeature(data={"pixel_values": pixel_values, "image_sizes": image_sizes})36 37class FastVLMProcessor(ProcessorMixin):38    attributes = ["tokenizer", "image_processor"]39    image_processor_class = "AutoImageProcessor"40    tokenizer_class = "AutoTokenizer"41 42    def __init__(43        self,44        tokenizer,45        image_processor,46        chat_template=None,47        **kwargs48    ):49        super().__init__(tokenizer, image_processor, chat_template=chat_template, **kwargs)50 51    def __call__(52        self,53        images: ImageInput = None,54        text: Optional[Union[str, List[str]]] = None,55        return_tensors: Optional[str] = "pt",56        **kwargs,57    ) -> BatchFeature:58        if isinstance(text, str):59            text = [text]60        elif not isinstance(text, list) and not isinstance(text[0], str):61            raise TypeError("Invalid input text. Please provide a string, or a list of strings")62 63        image_inputs = {}64        if images is not None:65            image_inputs = self.image_processor(images=images)66            67            image_token = torch.tensor([[IMAGE_TOKEN_INDEX]], dtype=torch.int64)68            input_ids = torch.tensor([], dtype=torch.int64)69            attention_mask = torch.tensor([], dtype=torch.int64)70            for prompt in text:71                image_indexes = [m.start() for m in re.finditer(DEFAULT_IMAGE_TOKEN, prompt)]72                if len(image_indexes) > 1:73                    raise ValueError(74                        f"Expected up to 1 image tokens per prompt, got {len(image_indexes)} instead."75                    )76 77                # DEFAULT_IMAGE_TOKEN is -200, not in the vocab (so we can't tokenize the full string)78                pre, _, post = prompt.partition(DEFAULT_IMAGE_TOKEN)79                pre_ids  = self.tokenizer(pre, return_tensors="pt", add_special_tokens=False).input_ids80                post_ids = self.tokenizer(post, return_tensors="pt", add_special_tokens=False).input_ids81 82                sample_ids = torch.cat([pre_ids, image_token, post_ids], dim=1).to(dtype=torch.int64)83                sample_mask = torch.ones_like(sample_ids)84 85                input_ids = torch.cat([input_ids, sample_ids], dim=0)86                attention_mask = torch.cat([attention_mask, sample_mask], dim=0)87 88        return BatchFeature(data={"input_ids": input_ids, "attention_mask": attention_mask, **image_inputs}, tensor_type=return_tensors)89