Upload folder using huggingface_hub

image_processing_minicpmv.py

@@ -0,0 +1,402 @@
+from typing import Optional, Union, Dict, Any
+
+import torch
+import math
+import PIL.Image
+import PIL.ImageSequence
+import numpy as np
+import PIL
+from PIL import Image
+
+from transformers.utils import TensorType, requires_backends, is_torch_dtype, is_torch_device
+from transformers.image_processing_utils import BaseImageProcessor, BatchFeature
+from transformers import AutoImageProcessor
+from transformers.image_transforms import to_channel_dimension_format
+from transformers.image_utils import (
+    ImageInput, 
+    make_list_of_images, 
+    valid_images, 
+    is_torch_tensor, 
+    to_numpy_array, 
+    infer_channel_dimension_format,
+    ChannelDimension
+)
+
+
+def recursive_converter(converter, value):
+    if isinstance(value, list):
+        new_value = []
+        for v in value:
+            new_value += [recursive_converter(converter, v)]
+        return new_value
+    else:
+        return converter(value)
+
+
+class MiniCPMVBatchFeature(BatchFeature):
+    r"""
+    Extend from BatchFeature for supporting various image size
+    """
+    def __init__(self, data: Optional[Dict[str, Any]] = None, tensor_type: Union[None, str, TensorType] = None):
+        super().__init__(data)
+        self.convert_to_tensors(tensor_type=tensor_type)
+
+    def convert_to_tensors(self, tensor_type: Optional[Union[str, TensorType]] = None):
+        if tensor_type is None:
+            return self
+        
+        is_tensor, as_tensor = self._get_is_as_tensor_fns(tensor_type)
+
+        def converter(value):
+            try:
+                if not is_tensor(value):
+                    tensor = as_tensor(value)
+                    return tensor
+            except:  # noqa E722
+                if key == "overflowing_values":
+                    raise ValueError("Unable to create tensor returning overflowing values of different lengths. ")
+                raise ValueError(
+                    "Unable to create tensor, you should probably activate padding "
+                    "with 'padding=True' to have batched tensors with the same length."
+                )
+
+
+        for key, value in self.items():
+            self[key] = recursive_converter(converter, value)
+        return self
+            
+    def to(self, *args, **kwargs) -> "MiniCPMVBatchFeature":
+        requires_backends(self, ["torch"])
+        import torch
+
+        def cast_tensor(v):
+            # check if v is a floating point
+            if torch.is_floating_point(v):
+                # cast and send to device
+                return v.to(*args, **kwargs)
+            elif device is not None:
+                return v.to(device=device)
+            else:
+                return v
+
+        new_data = {}
+        device = kwargs.get("device")
+        # Check if the args are a device or a dtype
+        if device is None and len(args) > 0:
+            # device should be always the first argument
+            arg = args[0]
+            if is_torch_dtype(arg):
+                # The first argument is a dtype
+                pass
+            elif isinstance(arg, str) or is_torch_device(arg) or isinstance(arg, int):
+                device = arg
+            else:
+                # it's something else
+                raise ValueError(f"Attempting to cast a BatchFeature to type {str(arg)}. This is not supported.")
+        # We cast only floating point tensors to avoid issues with tokenizers casting `LongTensor` to `FloatTensor`
+        for k, v in self.items():
+            new_data[k] = recursive_converter(cast_tensor, v)
+        self.data = new_data
+        return self
+
+
+class MiniCPMVImageProcessor(BaseImageProcessor):
+    model_input_names = ["pixel_values"]
+
+    def __init__(
+            self, 
+            max_slice_nums=9,
+            scale_resolution=448,
+            patch_size=14,
+            **kwargs):
+        super().__init__(**kwargs)
+        self.max_slice_nums = max_slice_nums
+        self.scale_resolution = scale_resolution
+        self.patch_size = patch_size
+        self.image_feature_size = kwargs.pop("image_feature_size", 64)
+        self.im_start_token = kwargs.pop("im_start", "<image>")
+        self.im_end_token = kwargs.pop("im_end", "</image>")
+        self.slice_start_token = kwargs.pop("slice_start", "<slice>")
+        self.slice_end_token = kwargs.pop("slice_end", "</slice>")
+        self.unk_token = kwargs.pop("unk", "<unk>")
+        self.mean = np.array(kwargs.pop("norm_mean", [0.5, 0.5, 0.5]))
+        self.std = np.array(kwargs.pop("norm_std", [0.5, 0.5, 0.5]))
+        self.version = kwargs.pop("version", 2.0)
+
+    def ensure_divide(self, length, patch_size):
+        return max(round(length / patch_size) * patch_size, patch_size)
+
+    def find_best_resize(self,
+                         original_size,
+                         scale_resolution,
+                         patch_size,
+                         allow_upscale=False):
+        width, height = original_size
+        if (width * height >
+                scale_resolution * scale_resolution) or allow_upscale:
+            r = width / height
+            height = int(scale_resolution / math.sqrt(r))
+            width = int(height * r)
+        best_width = self.ensure_divide(width, patch_size)
+        best_height = self.ensure_divide(height, patch_size)
+        return (best_width, best_height)
+
+    def get_refine_size(self,
+                        original_size,
+                        grid,
+                        scale_resolution,
+                        patch_size,
+                        allow_upscale=False):
+        width, height = original_size
+        grid_x, grid_y = grid
+
+        refine_width = self.ensure_divide(width, grid_x)
+        refine_height = self.ensure_divide(height, grid_y)
+
+        grid_width = refine_width / grid_x
+        grid_height = refine_height / grid_y
+
+        best_grid_size = self.find_best_resize((grid_width, grid_height),
+                                               scale_resolution,
+                                               patch_size,
+                                               allow_upscale=allow_upscale)
+        refine_size = (best_grid_size[0] * grid_x, best_grid_size[1] * grid_y)
+        return refine_size
+
+    def split_to_patches(self, image, grid):
+        patches = []
+        width, height = image.size
+        grid_x = int(width / grid[0])
+        grid_y = int(height / grid[1])
+        for i in range(0, height, grid_y):
+            images = []
+            for j in range(0, width, grid_x):
+                box = (j, i, j + grid_x, i + grid_y)
+                patch = image.crop(box)
+                images.append(patch)
+            patches.append(images)
+        return patches
+
+    def slice_image(
+        self, image, max_slice_nums=9, scale_resolution=448, patch_size=14, never_split=False
+    ):
+        original_size = image.size
+        original_width, original_height = original_size
+        log_ratio = math.log(original_width / original_height)
+        ratio = original_width * original_height / (scale_resolution * scale_resolution)
+        multiple = min(math.ceil(ratio), max_slice_nums)
+
+        source_image = None
+        best_grid = None
+        patches = []
+
+        if multiple <= 1 or never_split:
+            # dont need to slice, upsample
+            best_size = self.find_best_resize(
+                original_size, scale_resolution, patch_size, allow_upscale=True
+            )
+            source_image = image.resize(best_size, resample=Image.Resampling.BICUBIC)
+        else:
+            candidate_split_grids_nums = []
+            for i in [multiple - 1, multiple, multiple + 1]:
+                if i == 1 or i > max_slice_nums:
+                    continue
+                candidate_split_grids_nums.append(i)
+
+            # source image, down-sampling and ensure divided by patch_size
+            best_resize = self.find_best_resize(original_size, scale_resolution, patch_size)
+            source_image = image.copy().resize(best_resize, resample=Image.Resampling.BICUBIC)
+            candidate_grids = []
+
+            # find best grid
+            for split_grids_nums in candidate_split_grids_nums:
+                m = 1
+                while m <= split_grids_nums:
+                    if split_grids_nums % m == 0:
+                        candidate_grids.append([m, split_grids_nums // m])
+                    m += 1
+
+            best_grid = [1, 1]
+            min_error = float("inf")
+            for grid in candidate_grids:
+                error = abs(log_ratio - math.log(grid[0] / grid[1]))
+                if error < min_error:
+                    best_grid = grid
+                    min_error = error
+
+            refine_size = self.get_refine_size(
+                original_size, best_grid, scale_resolution, patch_size, allow_upscale=True
+            )
+
+            refine_image = image.resize(refine_size, resample=Image.Resampling.BICUBIC)
+            patches = self.split_to_patches(refine_image, best_grid)
+
+        return source_image, patches, best_grid
+
+    def get_grid_placeholder(self, grid):
+        if grid is None:
+            return ""
+        image_placeholder = (
+            self.im_start_token 
+            + self.unk_token * self.image_feature_size
+            + self.im_end_token
+        )
+
+        cols = grid[0]
+        rows = grid[1]
+        slices = []
+        for i in range(rows):
+            lines = []
+            for j in range(cols):
+                lines.append(image_placeholder)
+            slices.append("".join(lines))
+            
+        slice_placeholder = self.slice_start_token + "\n".join(slices) + self.slice_end_token
+        return slice_placeholder
+
+    def get_sliced_images(self, image):
+        slice_images = []
+
+        source_image, patches, sliced_grid = self.slice_image(
+            image,
+            self.max_slice_nums,  # default: 9
+            self.scale_resolution,  # default: 448
+            self.patch_size  # default: 14
+        )
+        slice_images.append(source_image)
+
+        if len(patches) > 0:
+            for i in range(len(patches)):
+                for j in range(len(patches[0])):
+                    slice_images.append(patches[i][j])
+        return slice_images
+
+    def get_sliced_grid(self, image_size):
+        original_width, original_height = image_size
+        log_ratio = math.log(original_width / original_height)
+        ratio = original_width * original_height / (self.scale_resolution * self.scale_resolution)
+        multiple = min(math.ceil(ratio), self.max_slice_nums)
+        if multiple <= 1:
+            return None
+        candidate_split_grids_nums = []
+        for i in [multiple - 1, multiple, multiple + 1]:
+            if i == 1 or i > self.max_slice_nums:
+                continue
+            candidate_split_grids_nums.append(i)
+        
+        candidate_grids = []
+        for split_grids_nums in candidate_split_grids_nums:
+            m = 1
+            while m <= split_grids_nums:
+                if split_grids_nums % m == 0:
+                    candidate_grids.append([m, split_grids_nums // m])
+                m += 1
+
+        best_grid = [1, 1]
+        min_error = float("inf")
+        for grid in candidate_grids:
+            error = abs(log_ratio - math.log(grid[0] / grid[1]))
+            if error < min_error:
+                best_grid = grid
+                min_error = error
+        
+        return best_grid
+
+    def get_slice_image_placeholder(self, image_size):
+        grid = self.get_sliced_grid(image_size=image_size)
+        return (
+            self.im_start_token 
+            + self.unk_token * self.image_feature_size 
+            + self.im_end_token
+        ) + self.get_grid_placeholder(grid=grid)
+
+    def to_pil_image(self, image, rescale=None) -> PIL.Image.Image:
+        """
+        Converts `image` to a PIL Image. Optionally rescales it and puts the channel dimension back as the last axis if
+        needed.
+
+        Args:
+            image (`PIL.Image.Image` or `numpy.ndarray` or `torch.Tensor`):
+                The image to convert to the PIL Image format.
+            rescale (`bool`, *optional*):
+                Whether or not to apply the scaling factor (to make pixel values integers between 0 and 255). Will
+                default to `True` if the image type is a floating type, `False` otherwise.
+        """
+        if isinstance(image, PIL.Image.Image):
+            return image
+        if is_torch_tensor(image):
+            image = image.numpy()
+
+        if isinstance(image, np.ndarray):
+            if rescale is None:
+                # rescale default to the array being of floating type.
+                rescale = isinstance(image.flat[0], np.floating)
+            # If the channel as been moved to first dim, we put it back at the end.
+            if image.ndim == 3 and image.shape[0] in [1, 3]:
+                image = image.transpose(1, 2, 0)
+            if rescale:
+                image = image * 255
+            image = image.astype(np.uint8)
+            return PIL.Image.fromarray(image)
+        return image
+
+    def reshape_by_patch(self, image):
+        """
+        :param image: shape [3, H, W]
+        :param patch_size:
+        :return: [3, patch_size, HW/patch_size]
+        """
+        image = torch.from_numpy(image)
+        patch_size = self.patch_size
+        patches = torch.nn.functional.unfold(
+            image,
+            (patch_size, patch_size),
+            stride=(patch_size, patch_size)
+        )
+
+        patches = patches.reshape(image.size(0), patch_size, patch_size, -1)
+        patches = patches.permute(0, 1, 3, 2).reshape(image.size(0), patch_size, -1)
+        return patches.numpy()
+
+    def preprocess(
+            self, 
+            images: ImageInput,
+            do_pad: Optional[bool] = True, # TODO: add pad for MiniCPM-Llama3-V-2_5
+            return_tensors: Optional[Union[str, TensorType]] = None
+        ) -> MiniCPMVBatchFeature:
+        images = make_list_of_images(images)
+
+        if 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."
+            )
+        
+        images = [self.to_pil_image(image).convert("RGB") for image in images]
+        input_data_format = infer_channel_dimension_format(np.array(images[0]))
+
+        new_images = []
+        image_sizes = [image.size for image in images]
+        tgt_sizes = []
+        for image in images:
+            image_patches = self.get_sliced_images(image)
+            image_patches = [to_numpy_array(image).astype(np.float32) / 255 for image in image_patches]
+            image_patches = [
+                self.normalize(image=image, mean=self.mean, std=self.std, input_data_format=input_data_format)
+                    for image in image_patches
+            ]
+            image_patches = [
+                to_channel_dimension_format(image, ChannelDimension.FIRST, input_channel_dim=input_data_format) 
+                    for image in image_patches
+            ]
+            for slice_image in image_patches:
+                new_images.append(self.reshape_by_patch(slice_image))
+                tgt_sizes.append(np.array((slice_image.shape[1] // self.patch_size, slice_image.shape[2] // self.patch_size)))
+
+        if tgt_sizes:
+            tgt_sizes = np.vstack(tgt_sizes)
+        return MiniCPMVBatchFeature(
+            data={"pixel_values": [new_images], "image_sizes": [image_sizes], "tgt_sizes": [tgt_sizes]}, tensor_type=return_tensors
+        )
+
+AutoImageProcessor.register("MiniCPMVImageProcessor", MiniCPMVImageProcessor)

modeling_minicpmv.py

@@ -0,0 +1,364 @@
+import math
+import json
+import torch
+from threading import Thread
+from copy import deepcopy
+from PIL import Image
+from torchvision import transforms
+from transformers import LlamaPreTrainedModel, LlamaForCausalLM, TextIteratorStreamer
+from transformers.models.idefics2.modeling_idefics2 import Idefics2VisionTransformer
+from transformers import AutoProcessor
+
+from .configuration_minicpm import MiniCPMVConfig
+from .resampler import Resampler
+
+IMAGENET_INCEPTION_MEAN = (0.5, 0.5, 0.5) # timm.data.IMAGENET_INCEPTION_MEAN
+IMAGENET_INCEPTION_STD = (0.5, 0.5, 0.5)  # timm.data.IMAGENET_INCEPTION_STD
+
+class MiniCPMVPreTrainedModel(LlamaPreTrainedModel):
+    config_class = MiniCPMVConfig
+
+
+class MiniCPMV(MiniCPMVPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.llm = LlamaForCausalLM(config)
+        self.vpm = self.init_vision_module()
+        self.vision_dim = self.vpm.embed_dim
+        self.embed_dim = self.llm.config.hidden_size
+        self.resampler = self.init_resampler(self.embed_dim, self.vision_dim)
+        self.transform = self.init_transform()
+
+    def init_vision_module(self):
+        # same as HuggingFaceM4/siglip-so400m-14-980-flash-attn2-navit
+        model = Idefics2VisionTransformer(self.config.vision_config)
+        if self.config.drop_vision_last_layer:
+            model.encoder.layers = model.encoder.layers[:-1]
+
+        setattr(model, 'embed_dim', model.embeddings.embed_dim)
+        setattr(model, 'patch_size', model.embeddings.patch_size)
+
+        return model
+
+    def init_resampler(self, embed_dim, vision_dim):
+        return Resampler(
+            num_queries=self.config.query_num,
+            embed_dim=embed_dim,
+            num_heads=embed_dim // 128,
+            kv_dim=vision_dim,
+            adaptive=True
+        )
+
+    def init_transform(self):
+        return transforms.Compose(
+            [
+                transforms.ToTensor(),
+                transforms.Normalize(
+                    mean=IMAGENET_INCEPTION_MEAN, std=IMAGENET_INCEPTION_STD
+                ),
+            ]
+        )
+
+    def get_input_embeddings(self):
+        return self.llm.get_input_embeddings()
+
+    def set_input_embeddings(self, value):
+        self.llm.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.llm.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.llm.lm_head = new_embeddings
+
+    def set_decoder(self, decoder):
+        self.llm = decoder
+
+    def get_decoder(self):
+        return self.llm
+
+    def get_vllm_embedding(self, data):
+        if 'vision_hidden_states' not in data:
+            dtype = self.llm.model.embed_tokens.weight.dtype
+            device = self.llm.model.embed_tokens.weight.device
+            tgt_sizes = data['tgt_sizes']
+            pixel_values_list = data['pixel_values']
+            vision_hidden_states = []
+            all_pixel_values = []
+            img_cnt = []
+            for pixel_values in pixel_values_list:
+                img_cnt.append(len(pixel_values))
+                all_pixel_values.extend([i.flatten(end_dim=1).permute(1, 0) for i in pixel_values])
+                
+            # exist image
+            if all_pixel_values:
+                tgt_sizes = torch.vstack(tgt_sizes).type(torch.int32)
+
+                if self.config.batch_vision_input:
+                    max_patches = torch.max(tgt_sizes[:, 0] * tgt_sizes[:, 1])
+
+                    all_pixel_values = torch.nn.utils.rnn.pad_sequence(all_pixel_values, batch_first=True,
+                                                                       padding_value=0.0)
+                    B, L, _ = all_pixel_values.shape
+                    all_pixel_values = all_pixel_values.permute(0, 2, 1).reshape(B, 3, -1, L)
+
+                    patch_attn_mask = torch.zeros((B, 1, max_patches), dtype=torch.bool, device=device)
+                    for i in range(B):
+                        patch_attn_mask[i, :tgt_sizes[i][0] * tgt_sizes[i][1]] = True
+
+                    vision_embedding = self.vpm(all_pixel_values.type(dtype), patch_attention_mask=patch_attn_mask).last_hidden_state
+                    vision_embedding = self.resampler(vision_embedding, tgt_sizes)
+                else:
+                    # get vision_embedding foreach
+                    vision_embedding = []
+                    for single_tgt_size, single_pixel_values in zip(tgt_sizes, all_pixel_values):
+                        single_pixel_values = single_pixel_values.unsqueeze(0)
+                        B, L, _ = single_pixel_values.shape
+                        single_pixel_values = single_pixel_values.permute(0, 2, 1).reshape(B, 3, -1, L)
+                        single_vision_embedding = self.vpm(single_pixel_values.type(dtype)).last_hidden_state
+                        single_vision_embedding = self.resampler(single_vision_embedding, single_tgt_size.unsqueeze(0))
+                        vision_embedding.append(single_vision_embedding)
+                    vision_embedding = torch.vstack(vision_embedding)
+
+                start = 0
+                for pixel_values in pixel_values_list:
+                    img_cnt = len(pixel_values)
+                    if img_cnt > 0:
+                        vision_hidden_states.append(vision_embedding[start: start + img_cnt])
+                        start += img_cnt
+                    else:
+                        vision_hidden_states.append([])
+            else: # no image
+                if self.training:
+                    dummy_image = torch.zeros(
+                        (1, 3, 224, 224),
+                        device=device, dtype=dtype
+                    )
+                    tgt_sizes = torch.Tensor([[(224 // self.config.patch_size), math.ceil(224 / self.config.patch_size)]]).type(torch.int32)
+                    dummy_feature = self.resampler(self.vpm(dummy_image).last_hidden_state, tgt_sizes)
+                else:
+                    dummy_feature = []
+                for _ in range(len(pixel_values_list)):
+                    vision_hidden_states.append(dummy_feature)
+
+        else:
+            vision_hidden_states = data['vision_hidden_states']
+
+        if hasattr(self.llm.config, 'scale_emb'):
+            vllm_embedding = self.llm.model.embed_tokens(data['input_ids']) * self.llm.config.scale_emb
+        else:
+            vllm_embedding = self.llm.model.embed_tokens(data['input_ids'])
+
+        vision_hidden_states = [i.type(vllm_embedding.dtype) if isinstance(
+            i, torch.Tensor) else i for i in vision_hidden_states]
+
+        bs = len(data['input_ids'])
+        for i in range(bs):
+            cur_vs_hs = vision_hidden_states[i]
+            if len(cur_vs_hs) > 0:
+                cur_vllm_emb = vllm_embedding[i]
+                cur_image_bound = data['image_bound'][i]
+                if len(cur_image_bound) > 0:
+                    image_indices = torch.stack(
+                        [torch.arange(r[0], r[1], dtype=torch.long) for r in cur_image_bound]
+                    ).to(vllm_embedding.device)
+
+                    cur_vllm_emb.scatter_(0, image_indices.view(-1, 1).repeat(1, cur_vllm_emb.shape[-1]),
+                                          cur_vs_hs.view(-1, cur_vs_hs.shape[-1]))
+                elif self.training:
+                    cur_vllm_emb += cur_vs_hs[0].mean() * 0
+        
+        return vllm_embedding, vision_hidden_states
+
+    def forward(self, data, **kwargs):
+        vllm_embedding, vision_hidden_states = self.get_vllm_embedding(data)
+        position_ids = data["position_ids"]
+        if position_ids.dtype != torch.int64:
+            position_ids = position_ids.long()
+
+        return self.llm(
+            input_ids=None,
+            position_ids=position_ids,
+            inputs_embeds=vllm_embedding,
+            **kwargs
+        )
+
+    def _decode_text(self, result_ids, tokenizer):
+        result_text = []
+        for result in result_ids:
+            result = result[result != 0]
+            if result[0] == tokenizer.bos_id:
+                result = result[1:]
+            if result[-1] == tokenizer.eos_id or result[-1] == tokenizer.eot_id:
+                result = result[:-1]
+            result_text.append(tokenizer.decode(result).strip())
+        return result_text
+
+    def _decode(self, inputs_embeds, tokenizer, decode_text=False, **kwargs):
+        terminators = [
+            tokenizer.eos_token_id,
+            tokenizer.convert_tokens_to_ids("<|eot_id|>")
+        ]
+        output = self.llm.generate(
+            inputs_embeds=inputs_embeds,
+            pad_token_id=0,
+            eos_token_id=terminators,
+            **kwargs
+        )
+        if decode_text:
+            return self._decode_text(output, tokenizer)
+        return output
+    
+    def _decode_stream(self, inputs_embeds, tokenizer, **kwargs):
+        terminators = [
+            tokenizer.eos_token_id,
+            tokenizer.convert_tokens_to_ids("<|eot_id|>")
+        ]
+        streamer = TextIteratorStreamer(tokenizer=tokenizer)
+        generation_kwargs = {
+            'inputs_embeds': inputs_embeds,
+            'pad_token_id': 0,
+            'eos_token_id': terminators,
+            'streamer': streamer
+        }
+        generation_kwargs.update(kwargs)
+
+        thread = Thread(target=self.llm.generate, kwargs=generation_kwargs)
+        thread.start()
+    
+        return streamer
+
+    def generate(
+        self,
+        model_inputs,
+        tokenizer=None,
+        vision_hidden_states=None,
+        stream=False,
+        **kwargs
+    ):
+        bs = len(model_inputs["input_ids"])
+        img_list = model_inputs["pixel_values"]
+        tgt_sizes = model_inputs["tgt_sizes"]
+        if img_list is None:
+            img_list = [[] for i in range(bs)]
+        assert bs == len(img_list)
+        if vision_hidden_states is None:
+            pixel_values = []
+            for i in range(bs):
+                img_inps = []
+                for img in img_list[i]:
+                    img_inps.append(img.to(self.device))
+                if img_inps:
+                    pixel_values.append(img_inps)
+                else:
+                    pixel_values.append([])
+            model_inputs["pixel_values"] = pixel_values
+            model_inputs['tgt_sizes'] = tgt_sizes
+        else:
+            model_inputs["vision_hidden_states"] = vision_hidden_states
+
+        (
+            input_embeds,
+            vision_hidden_states,
+        ) = self.get_vllm_embedding(model_inputs)            
+
+        # output_ids = self._decode(input_embeds, tokenizer, **kwargs)
+        if stream:
+            kwargs.pop("decode_text")
+            result = self._decode_stream(input_embeds, tokenizer, **kwargs)
+        else:
+            result = self._decode(input_embeds, tokenizer, **kwargs)
+
+        return result
+
+    def chat(
+        self,
+        image,
+        msgs,
+        tokenizer,
+        processor=None,
+        vision_hidden_states=None,
+        max_new_tokens=1024,
+        sampling=True,
+        max_inp_length=2048,
+        system_prompt='',
+        stream=False,
+        **kwargs
+    ):
+        if processor is None:
+            processor = AutoProcessor.from_pretrained(self.config._name_or_path, trust_remote_code=True)
+        if isinstance(msgs, str):
+            msgs = json.loads(msgs)
+        copy_msgs = deepcopy(msgs)
+
+        assert len(msgs) > 0, "msgs is empty"
+        assert sampling or not stream, "if use stream mode, make sure sampling=True"
+
+        if image is not None and isinstance(copy_msgs[0]["content"], str):
+            # copy_msgs[0]['content'] = '(<image>./</image>)\n' + copy_msgs[0]['content']
+            copy_msgs[0]["content"] = [image, copy_msgs[0]["content"]]
+
+        images = []
+        for i, msg in enumerate(copy_msgs):
+            role = msg["role"]
+            content = msg["content"]
+            assert role in ["user", "assistant"]
+            if i == 0:
+                assert role == "user", "The role of first msg should be user"
+            if isinstance(content, str):
+                content = [content]
+            cur_msgs = []
+            for c in content:
+                if isinstance(c, Image.Image):
+                    images.append(c)
+                    cur_msgs.append("(<image>./</image>)")
+                elif isinstance(c, str):
+                    cur_msgs.append(c)
+            msg["content"] = "\n".join(cur_msgs)
+
+        if system_prompt:
+            sys_msg = {'role': 'system', 'content': system_prompt}
+            copy_msgs = [sys_msg] + copy_msgs        
+
+        prompt = processor.tokenizer.apply_chat_template(copy_msgs, tokenize=False, add_generation_prompt=True)
+        inputs = processor(prompt, images, return_tensors="pt", max_length=max_inp_length).to(self.device)
+
+        if sampling:
+            generation_config = {
+                "top_p": 0.8,
+                "top_k": 100,
+                "temperature": 0.7,
+                "do_sample": True,
+                "repetition_penalty": 1.05
+            }
+        else:
+            generation_config = {
+                "num_beams": 3,
+                "repetition_penalty": 1.2,
+            }
+
+        generation_config.update(
+            (k, kwargs[k]) for k in generation_config.keys() & kwargs.keys()
+        )
+        with torch.inference_mode():
+            res = self.generate(
+                inputs,
+                tokenizer=tokenizer,
+                max_new_tokens=max_new_tokens,
+                vision_hidden_states=vision_hidden_states,
+                stream=stream,
+                decode_text=True,
+                **generation_config
+            )
+        
+        if stream:
+            def stream_gen():
+                for text in res:
+                    text = text.replace(tokenizer.eot_token, '').replace(tokenizer.eos_token, '')
+                    yield text
+            return stream_gen()
+
+        else:
+            answer = res[0]
+            return answer

processing_minicpmv.py

@@ -0,0 +1,244 @@
+# coding=utf-8
+# Copyright 2024 The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Processor class for MiniCPMV.
+"""
+
+from typing import List, Optional, Union, Dict, Any
+import torch
+import re
+
+from transformers.image_processing_utils import BatchFeature
+from transformers.image_utils import ImageInput
+from transformers.processing_utils import ProcessorMixin
+from transformers.tokenization_utils_base import PaddingStrategy, PreTokenizedInput, TextInput, TruncationStrategy
+from transformers.utils import TensorType, requires_backends, is_torch_dtype, is_torch_device
+
+from .image_processing_minicpmv import MiniCPMVBatchFeature
+
+
+class MiniCPMVProcessor(ProcessorMixin):
+    r"""
+    Constructs a MiniCPMV processor which wraps a MiniCPMV image processor and a MiniCPMV tokenizer into a single processor.
+
+    [`MiniCPMVProcessor`] offers all the functionalities of [`MiniCPMVImageProcessor`] and [`LlamaTokenizerWrapper`]. See the
+    [`~MiniCPMVProcessor.__call__`] and [`~MiniCPMVProcessor.decode`] for more information.
+
+    Args:
+        image_processor ([`MiniCPMVImageProcessor`], *optional*):
+            The image processor is a required input.
+        tokenizer ([`LlamaTokenizerWrapper`], *optional*):
+            The tokenizer is a required input.
+    """
+    attributes = ["image_processor", "tokenizer"]
+    image_processor_class = "AutoImageProcessor"
+    tokenizer_class = "AutoTokenizer"
+
+    def __init__(self, image_processor=None, tokenizer=None):
+        super().__init__(image_processor, tokenizer)
+        self.version = image_processor.version
+    
+    def __call__(
+        self,
+        text: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]],
+        images: ImageInput = None,
+        padding: Union[bool, str, PaddingStrategy] = False,
+        truncation: Union[bool, str, TruncationStrategy] = None,
+        max_length: Optional[int] = None,
+        do_pad: Optional[bool] = True,
+        return_tensors: Optional[Union[str, TensorType]] = TensorType.PYTORCH,
+    ) -> MiniCPMVBatchFeature:
+        """
+        Only support for single input for now. Batched input is coming soon.
+
+        Args:
+            text (`str`):
+                The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings
+                (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set
+                `is_split_into_words=True` (to lift the ambiguity with a batch of sequences).
+            images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `List[PIL.Image.Image]`, `List[np.ndarray]`, `List[torch.Tensor]`):
+                The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch
+                tensor. Both channels-first and channels-last formats are supported.
+            padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `False`):
+                Select a strategy to pad the returned sequences (according to the model's padding side and padding
+                index) among:
+                - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single
+                  sequence if provided).
+                - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum
+                  acceptable input length for the model if that argument is not provided.
+                - `False` or `'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of different
+                  lengths).
+            max_length (`int`, *optional*):
+                Maximum length of the returned list and optionally padding length (see above).
+            do_pad (`bool`, *optional*, defaults to self.do_pad):
+                Whether to pad the image. If `True` will pad the images in the batch to the largest image in the batch
+                and create a pixel mask. Padding will be applied to the bottom and right of the image with zeros.
+            truncation (`bool`, *optional*):
+                Activates truncation to cut input sequences longer than `max_length` to `max_length`.
+            return_tensors (`str` or [`~utils.TensorType`], *optional*):
+                If set, will return tensors of a particular framework. Acceptable values are:
+
+                - `'tf'`: Return TensorFlow `tf.constant` objects.
+                - `'pt'`: Return PyTorch `torch.Tensor` objects.
+                - `'np'`: Return NumPy `np.ndarray` objects.
+                - `'jax'`: Return JAX `jnp.ndarray` objects.
+
+        Returns:
+            [`BatchFeature`]: A [`BatchFeature`] with the following fields:
+
+            - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`.
+            - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when
+              `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not
+              `None`).
+            - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`.
+        """
+        if images is not None:
+            image_inputs = self.image_processor(images, do_pad=do_pad, return_tensors=return_tensors)
+        return self._convert_images_texts_to_inputs(image_inputs, text, max_length=max_length)
+    
+    # Copied from transformers.models.clip.processing_clip.CLIPProcessor.batch_decode with CLIP->Llama
+    def batch_decode(self, *args, **kwargs):
+        """
+        This method forwards all its arguments to LlamaTokenizerFast's [`~PreTrainedTokenizer.batch_decode`]. Please
+        refer to the docstring of this method for more information.
+        """
+        output_ids = args[0]
+        result_text = []
+        for result in output_ids:
+            result = result[result != 0]
+            if result[0] == self.tokenizer.bos_id:
+                result = result[1:]
+            if result[-1] == self.tokenizer.eos_id:
+                result = result[:-1]
+            result_text.append(self.tokenizer.decode(result, *args[1:], **kwargs).strip())
+        return result_text
+        # return self.tokenizer.batch_decode(*args, **kwargs)
+    
+    # Copied from transformers.models.clip.processing_clip.CLIPProcessor.decode with CLIP->Llama
+    def decode(self, *args, **kwargs):
+        """
+        This method forwards all its arguments to LlamaTokenizerFast's [`~PreTrainedTokenizer.decode`]. Please refer to
+        the docstring of this method for more information.
+        """
+        result = args[0]
+        result = result[result != 0]
+        if result[0] == self.tokenizer.bos_id:
+            result = result[1:]
+        if result[-1] == self.tokenizer.eos_id or (hasattr(self.tokenizer, "eot_id") and result[-1] == self.tokenizer.eot_id):
+            result = result[:-1]
+        return self.tokenizer.decode(result, *args[1:], **kwargs).strip()
+
+    def _convert(
+        self, input_str, max_inp_length: Optional[int] = None
+    ):
+        if self.version == 2.5 or self.tokenizer.add_bos_token:
+            input_ids = self.tokenizer.encode(input_str)
+        else:
+            input_ids = [self.tokenizer.bos_id] + self.tokenizer.encode(input_str)
+        if max_inp_length is not None:
+            input_ids = input_ids[:max_inp_length]
+        input_ids = torch.tensor(input_ids, dtype=torch.int32)
+
+        image_start_tokens = torch.where(input_ids == self.tokenizer.im_start_id)[0]
+        image_start_tokens += 1
+        image_end_tokens = torch.where(input_ids == self.tokenizer.im_end_id)[0]
+        valid_image_nums = max(len(image_start_tokens), len(image_end_tokens))
+        image_bounds = torch.hstack(
+            [
+                image_start_tokens[:valid_image_nums].unsqueeze(-1),
+                image_end_tokens[:valid_image_nums].unsqueeze(-1),
+            ]
+        )
+        return input_ids.unsqueeze(0), image_bounds
+
+    def _convert_images_texts_to_inputs(self, images, texts, do_pad=False, truncation=None, max_length=None, return_tensors=None):
+        if not len(images):
+            model_inputs = self.tokenizer(texts, return_tensors=return_tensors, padding=do_pad, truncation=truncation, max_length=max_length)
+            return MiniCPMVBatchFeature(data={**model_inputs})
+        
+        pattern = "(<image>./</image>)"
+        images, image_sizes, tgt_sizes = images["pixel_values"], images["image_sizes"], images["tgt_sizes"]
+
+        image_tags = re.findall(pattern, texts)
+        assert len(image_tags) == len(image_sizes[0])
+        text_chunks = texts.split(pattern)
+        final_texts = ""
+        for i in range(len(image_tags)):
+            final_texts = final_texts + text_chunks[i] + self.image_processor.get_slice_image_placeholder(image_sizes[0][i])
+        final_texts += text_chunks[-1]
+        input_ids, image_bounds = self._convert(final_texts, max_length)
+        return MiniCPMVBatchFeature(data={
+            "input_ids": input_ids,
+            "pixel_values": images,
+            "image_sizes": image_sizes,
+            "image_bound": [image_bounds],
+            "tgt_sizes": tgt_sizes
+        })
+
+    @property
+    # Copied from transformers.models.clip.processing_clip.CLIPProcessor.model_input_names
+    def model_input_names(self):
+        tokenizer_input_names = self.tokenizer.model_input_names
+        image_processor_input_names = self.image_processor.model_input_names
+        return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names))
+
+
+    def pad(self, orig_items, key, max_length=None, padding_value=0, padding_side="left"):
+        items = []
+        if isinstance(orig_items[0][key], list):
+            assert isinstance(orig_items[0][key][0], torch.Tensor)
+            for it in orig_items:
+                for tr in it[key]:
+                    items.append({key: tr})
+        else:
+            assert isinstance(orig_items[0][key], torch.Tensor)
+            items = orig_items
+
+        batch_size = len(items)
+        shape = items[0][key].shape
+        dim = len(shape)
+        assert dim <= 3
+        if max_length is None:
+            max_length = 0
+        max_length = max(max_length, max(item[key].shape[-1] for item in items))
+        min_length = min(item[key].shape[-1] for item in items)
+        dtype = items[0][key].dtype
+
+        if dim == 1:
+            return torch.cat([item[key] for item in items], dim=0)
+        elif dim == 2:
+            if max_length == min_length:
+                return torch.cat([item[key] for item in items], dim=0)
+            tensor = torch.zeros((batch_size, max_length), dtype=dtype) + padding_value
+        else:
+            tensor = (
+                torch.zeros((batch_size, max_length, shape[-1]), dtype=dtype)
+                + padding_value
+            )
+
+        for i, item in enumerate(items):
+            if dim == 2:
+                if padding_side == "left":
+                    tensor[i, -len(item[key][0]) :] = item[key][0].clone()
+                else:
+                    tensor[i, : len(item[key][0])] = item[key][0].clone()
+            elif dim == 3:
+                if padding_side == "left":
+                    tensor[i, -len(item[key][0]) :, :] = item[key][0].clone()
+                else:
+                    tensor[i, : len(item[key][0]), :] = item[key][0].clone()
+
+        return tensor
+