| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import torch |
| 4 | from mistral_common.protocol.instruct.chunk import ImageChunk |
| 5 | from mistral_common.tokens.tokenizers.multimodal import ImageEncoder |
| 6 | from PIL import Image |
| 7 | from transformers import BatchFeature, ImageProcessingMixin, ProcessorMixin, TensorType |
| 8 | from transformers.image_utils import ImageInput |
| 9 | |
| 10 | from vllm.tokenizers.mistral import MistralTokenizer |
| 11 | |
| 12 | |
| 13 | class MistralCommonImageProcessor(ImageProcessingMixin): |
| 14 | """Provide a HF-compatible interface for |
| 15 | `mistral_common.tokens.tokenizers.multimodal.ImageEncoder`. |
| 16 | """ |
| 17 | |
| 18 | def __init__(self, mm_encoder: ImageEncoder) -> None: |
| 19 | self.mm_encoder = mm_encoder |
| 20 | |
| 21 | def __call__( |
| 22 | self, |
| 23 | images: ImageInput, |
| 24 | return_tensors: str | TensorType | None = None, |
| 25 | **kwargs, |
| 26 | ) -> BatchFeature: |
| 27 | images_lst = [images] if not isinstance(images, list) else images |
| 28 | |
| 29 | images_processed = list[torch.Tensor]() |
| 30 | |
| 31 | for image in images_lst: |
| 32 | image_inputs = self.mm_encoder(ImageChunk(image=image)) |
| 33 | image_processed = torch.tensor(image_inputs.image) |
| 34 | |
| 35 | images_processed.append(image_processed) |
| 36 | |
| 37 | return BatchFeature({"images": images_processed}, tensor_type=return_tensors) |
| 38 | |
| 39 | def get_number_of_image_patches( |
| 40 | self, |
| 41 | height: int, |
| 42 | width: int, |
| 43 | ) -> tuple[int, int, int]: |
| 44 | image = Image.new("RGB", (width, height)) |
| 45 | ncols, nrows = self.mm_encoder._image_to_num_tokens(image) |
| 46 | return ncols * nrows, nrows, ncols |
| 47 | |
| 48 | # Copied from Transformers (Apache-2.0): |
| 49 | # https://github.com/huggingface/transformers/blob/d20946079fd422335fbae3eeb98b7cd88334612f/src/transformers/image_processing_base.py#L473 |
| 50 | def fetch_images(self, image_url_or_urls): |
| 51 | from transformers.image_utils import is_valid_image, load_image |
| 52 | |
| 53 | if isinstance(image_url_or_urls, (list, tuple)): |
| 54 | return [self.fetch_images(x) for x in image_url_or_urls] |
| 55 | if isinstance(image_url_or_urls, str): |
| 56 | return load_image(image_url_or_urls) |
| 57 | if is_valid_image(image_url_or_urls): |
| 58 | return image_url_or_urls |
| 59 | raise TypeError( |
| 60 | "only a single or a list of entries is supported but got " |
| 61 | f"type={type(image_url_or_urls)}" |
| 62 | ) |
| 63 | |
| 64 | |
| 65 | class MistralCommonPixtralProcessor(ProcessorMixin): |
| 66 | attributes = ["image_processor", "tokenizer"] |
| 67 | |
| 68 | def __init__( |
| 69 | self, |
| 70 | tokenizer: MistralTokenizer, |
| 71 | image_processor: MistralCommonImageProcessor, |
| 72 | ) -> None: |
| 73 | self.tokenizer = tokenizer.transformers_tokenizer |
| 74 | self.image_processor = image_processor |
| 75 | |
| 76 | image_special_ids = self.image_processor.mm_encoder.special_ids |
| 77 | self.image_break_id = image_special_ids.img_break |
| 78 | self.image_token_id = image_special_ids.img |
| 79 | self.image_end_id = image_special_ids.img_end |