add processor & image processor

configuration.json

@@ -0,0 +1 @@
+{"framework":"Pytorch","task":"multimodal-dialogue"}

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

@@ -9,6 +9,7 @@ from PIL import Image
 from torchvision import transforms
 from transformers import LlamaTokenizer, LlamaPreTrainedModel, LlamaForCausalLM, AutoModel, PreTrainedTokenizerFast, TextIteratorStreamer
 from transformers.models.idefics2.modeling_idefics2 import Idefics2VisionTransformer
+from transformers import AutoProcessor
 
 from .configuration_minicpm import MiniCPMVConfig
 from .resampler import Resampler
@@ -42,13 +43,13 @@ class MiniCPMV(MiniCPMVPreTrainedModel):
 
         return model
 
-    def init_resampler(self, embed_dim, vision_dim,):
+    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,
+            adaptive=True
         )
 
     def init_transform(self):
@@ -60,17 +61,29 @@ class MiniCPMV(MiniCPMVPreTrainedModel):
                 ),
             ]
         )
-    
+
     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
+            dtype = self.vpm.embeddings.position_embedding.weight.dtype
+            device = self.vpm.embeddings.position_embedding.weight.device
             tgt_sizes = data['tgt_sizes']
             pixel_values_list = data['pixel_values']
             vision_hidden_states = []
@@ -78,9 +91,7 @@ class MiniCPMV(MiniCPMVPreTrainedModel):
             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
+                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)
 
@@ -107,7 +118,6 @@ class MiniCPMV(MiniCPMVPreTrainedModel):
                         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)
 
@@ -148,18 +158,19 @@ class MiniCPMV(MiniCPMVPreTrainedModel):
             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]
+                cur_image_bound = data['image_bounds'][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
         
+        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"]
@@ -173,47 +184,18 @@ class MiniCPMV(MiniCPMVPreTrainedModel):
             **kwargs
         )
 
-    def _convert_to_tensors(
-        self, tokenizer, input_ids, max_inp_length: Optional[int] = None
-    ):
-        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 == tokenizer.im_start_id)[0]
-        # 跳过 im_start
-        image_start_tokens += 1
-        image_end_tokens = torch.where(input_ids == tokenizer.im_end_id)[0]
-        valid_image_nums = max(len(image_start_tokens), len(image_end_tokens))
-        image_bound = torch.hstack(
-            [
-                image_start_tokens[:valid_image_nums].unsqueeze(-1),
-                image_end_tokens[:valid_image_nums].unsqueeze(-1),
-            ]
-        )
-
-        model_input = {}
-        model_input["input_ids"] = input_ids.unsqueeze(0).to(self.device)
-        model_input["image_bound"] = image_bound
-
-        return model_input
-
-    def _process_list(
-        self, tokenizer, input_id_list, max_inp_length: Optional[int] = None
-    ):
-        pad_keys = ["input_ids"]
-        input_tensors = []
-        for input_ids in input_id_list:
-            input_tensors.append(
-                self._convert_to_tensors(tokenizer, input_ids, max_inp_length)
-            )
-        padded = {}
-        for key in pad_keys:
-            padded[key] = pad(input_tensors, key, padding_side="left").to(self.device)
-        padded["image_bound"] = [i["image_bound"] for i in input_tensors]
-        return padded
+    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, **kwargs):
+    def _decode(self, inputs_embeds, tokenizer, decode_text=False, **kwargs):
         terminators = [
             tokenizer.eos_token_id,
             tokenizer.convert_tokens_to_ids("<|eot_id|>")
@@ -224,7 +206,9 @@ class MiniCPMV(MiniCPMVPreTrainedModel):
             eos_token_id=terminators,
             **kwargs
         )
-        return self._decode_text(output, tokenizer)
+        if decode_text:
+            return self._decode_text(output, tokenizer)
+        return output
     
     def _decode_stream(self, inputs_embeds, tokenizer, **kwargs):
         terminators = [
@@ -245,93 +229,20 @@ class MiniCPMV(MiniCPMVPreTrainedModel):
     
         return streamer
 
-    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 slice_image(self, image):
-        return slice_image(
-            image,
-            self.config.slice_config.max_slice_nums,
-            self.config.slice_config.scale_resolution,
-            self.config.slice_config.patch_size,
-        )
-
-    def get_slice_image_placeholder(self, image, tokenizer):
-        image_placeholder = (
-            tokenizer.im_start
-            + tokenizer.unk_token * self.config.query_num
-            + tokenizer.im_end
-        )
-
-        slice_images = []
-
-        source_image, patches, best_grid = slice_image(
-            image,
-            self.config.slice_config.max_slice_nums,
-            self.config.slice_config.scale_resolution,
-            self.config.slice_config.patch_size,
-        )
-
-        slice_images.append(source_image)
-        final_placeholder = image_placeholder
-
-        if len(patches) > 0:
-            for i in range(len(patches)):
-                for j in range(len(patches[0])):
-                    slice_images.append(patches[i][j])
-
-            final_placeholder += get_grid_placeholder(
-                tokenizer, best_grid, self.config.query_num
-            )
-
-        return slice_images, final_placeholder
-
-    def reshape_by_patch(self, image_tensor):
-        """
-        :param image_tensor: shape [3, H, W]
-        :param patch_size:
-        :return: [3, patch_size, HW/patch_size]
-        """
-        patch_size = self.config.patch_size
-        patches = torch.nn.functional.unfold(
-            image_tensor,
-            (patch_size, patch_size),
-            stride=(patch_size, patch_size)
-        )
-
-        patches = patches.reshape(image_tensor.size(0), patch_size, patch_size, -1)
-        patches = patches.permute(0, 1, 3, 2).reshape(image_tensor.size(0), patch_size, -1)
-        return patches
-
     def generate(
         self,
-        input_id_list=None,
-        img_list=None,
-        tgt_sizes=None,
+        model_inputs,
         tokenizer=None,
-        max_inp_length: Optional[int] = None,
         vision_hidden_states=None,
-        return_vision_hidden_states=False,
         stream=False,
         **kwargs
     ):
-
-        assert input_id_list is not None
-        bs = len(input_id_list)
-        if img_list == None:
+        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)
-
-        model_inputs = self._process_list(tokenizer, input_id_list, max_inp_length)
-
         if vision_hidden_states is None:
             pixel_values = []
             for i in range(bs):
@@ -347,19 +258,17 @@ class MiniCPMV(MiniCPMVPreTrainedModel):
         else:
             model_inputs["vision_hidden_states"] = vision_hidden_states
 
-        with torch.inference_mode():
-            (
-                model_inputs["inputs_embeds"],
-                vision_hidden_states,
-            ) = self.get_vllm_embedding(model_inputs)
+        (
+            input_embeds,
+            vision_hidden_states,
+        ) = self.get_vllm_embedding(model_inputs)            
 
-            if stream:
-                result = self._decode_stream(model_inputs["inputs_embeds"], tokenizer, **kwargs)
-            else:
-                result = self._decode(model_inputs["inputs_embeds"], tokenizer, **kwargs)
-
-        if return_vision_hidden_states:
-            return result, vision_hidden_states
+        # 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
 
@@ -368,6 +277,7 @@ class MiniCPMV(MiniCPMVPreTrainedModel):
         image,
         msgs,
         tokenizer,
+        processor=None,
         vision_hidden_states=None,
         max_new_tokens=1024,
         sampling=True,
@@ -376,61 +286,22 @@ class MiniCPMV(MiniCPMVPreTrainedModel):
         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(copy_msgs) > 0, 'msgs is empty'
+        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, copy_msgs[0]['content']]
-
-        images = []
-        tgt_sizes = []
-        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):
-                    image = c
-                    if self.config.slice_mode:
-                        slice_images, image_placeholder = self.get_slice_image_placeholder(
-                            image, tokenizer
-                        )
-                        cur_msgs.append(image_placeholder)
-                        for slice_image in slice_images:
-                            slice_image = self.transform(slice_image)
-                            H, W = slice_image.shape[1:]
-                            images.append(self.reshape_by_patch(slice_image))
-                            tgt_sizes.append(torch.Tensor([H // self.config.patch_size, W // self.config.patch_size]).type(torch.int32))
-                    else:
-                        images.append(self.transform(image))
-                        cur_msgs.append(
-                            tokenizer.im_start
-                            + tokenizer.unk_token * self.config.query_num
-                            + tokenizer.im_end
-                        )
-                elif isinstance(c, str):
-                    cur_msgs.append(c)
-
-
-            msg['content'] = '\n'.join(cur_msgs)
-        if tgt_sizes:
-            tgt_sizes = torch.vstack(tgt_sizes)
-
+        if image is not None and isinstance(msgs[0]['content'], str):
+            msgs[0]['content'] = '(<image>./</image>)\n' + msgs[0]['content']
         if system_prompt:
             sys_msg = {'role': 'system', 'content': system_prompt}
-            copy_msgs = [sys_msg] + copy_msgs
+            copy_msgs = [sys_msg] + copy_msgs        
 
-        input_ids = tokenizer.apply_chat_template(copy_msgs, tokenize=True, add_generation_prompt=False)
+        prompt = processor.tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
+        inputs = processor(prompt, [image], return_tensors="pt", max_length=max_inp_length).to(self.device)
 
         if sampling:
             generation_config = {
@@ -449,21 +320,17 @@ class MiniCPMV(MiniCPMVPreTrainedModel):
         generation_config.update(
             (k, kwargs[k]) for k in generation_config.keys() & kwargs.keys()
         )
-
         with torch.inference_mode():
-            res, vision_hidden_states = self.generate(
-                input_id_list=[input_ids],
-                max_inp_length=max_inp_length,
-                img_list=[images],
-                tgt_sizes=[tgt_sizes],
+            res = self.generate(
+                inputs,
                 tokenizer=tokenizer,
                 max_new_tokens=max_new_tokens,
                 vision_hidden_states=vision_hidden_states,
-                return_vision_hidden_states=True,
                 stream=stream,
+                decode_text=True,
                 **generation_config
             )
-
+        
         if stream:
             def stream_gen():
                 for text in res:
@@ -474,229 +341,3 @@ class MiniCPMV(MiniCPMVPreTrainedModel):
         else:
             answer = res[0]
             return answer
-
-
-class PreTrainedTokenizerFastWrapper(PreTrainedTokenizerFast):
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-        self.eot_token = "<|eot_id|>"
-        self.im_start = "<image>"
-        self.im_end = "</image>"
-        self.ref_start = "<ref>"
-        self.ref_end = "</ref>"
-        self.box_start = "<box>"
-        self.box_end = "</box>"
-        self.quad_start = "<quad>"
-        self.quad_end = "</quad>"
-        self.slice_start = "<slice>"
-        self.slice_end = "</slice>"
-
-    @property
-    def eos_id(self):
-        return self.eos_token_id
-
-    @property
-    def bos_id(self):
-        return self.bos_token_id
-
-    @property
-    def unk_id(self):
-        return self.unk_token_id
-
-    @property
-    def eot_id(self):
-        return self.convert_tokens_to_ids(self.eot_token)
-
-    @property
-    def im_start_id(self):
-        return self.convert_tokens_to_ids(self.im_start)
-
-    @property
-    def im_end_id(self):
-        return self.convert_tokens_to_ids(self.im_end)
-
-    @staticmethod
-    def escape(text: str) -> str:
-        return text
-
-    @staticmethod
-    def unescape(text: str) -> str:
-        return text
-
-
-def pad(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
-
-
-def slice_image(
-    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 = find_best_resize(
-            original_size, scale_resolution, patch_size, allow_upscale=True
-        )
-        source_image = image.resize(best_size, 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 = find_best_resize(original_size, scale_resolution, patch_size)
-        source_image = image.copy().resize(best_resize, 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 = get_refine_size(
-            original_size, best_grid, scale_resolution, patch_size, allow_upscale=True
-        )
-
-        refine_image = image.resize(refine_size, Image.Resampling.BICUBIC)
-        patches = split_to_patches(refine_image, best_grid)
-
-    return source_image, patches, best_grid
-
-
-def ensure_divide(length, patch_size):
-    return max(round(length / patch_size) * patch_size, patch_size)
-
-
-def find_best_resize(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 = ensure_divide(width, patch_size)
-    best_height = ensure_divide(height, patch_size)
-    return (best_width, best_height)
-
-
-def get_refine_size(
-    original_size, grid, scale_resolution, patch_size, allow_upscale=False
-):
-    width, height = original_size
-    grid_x, grid_y = grid
-
-    refine_width = ensure_divide(width, grid_x)
-    refine_height = ensure_divide(height, grid_y)
-
-    grid_width = refine_width / grid_x
-    grid_height = refine_height / grid_y
-
-    best_grid_size = 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(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 get_grid_placeholder(tokenizer, grid, query_num):
-    image_placeholder = (
-        tokenizer.im_start + tokenizer.unk_token * query_num + tokenizer.im_end
-    )
-
-    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 = tokenizer.slice_start + "\n".join(slices) + tokenizer.slice_end
-    return slice_placeholder

preprocessor_config.json

@@ -0,0 +1,20 @@
+{
+    "image_processor_type": "MiniCPMVImageProcessor",
+    "auto_map": {
+        "AutoProcessor": "processing_minicpmv.MiniCPMVProcessor",
+        "AutoImageProcessor": "image_processing_minicpmv.MiniCPMVImageProcessor"
+      },
+    "processor_class": "MiniCPMVProcessor",
+    "max_slice_nums": 9,
+    "scale_resolution": 448,
+    "patch_size": 14,
+    "image_feature_size": 96,
+    "im_start": "<image>",
+    "im_end": "</image>",
+    "slice_start": "<slice>",
+    "slice_end": "</slice>",
+    "unk": "<unk>",
+    "norm_mean": [0.5, 0.5, 0.5],
+    "norm_std": [0.5, 0.5, 0.5],
+    "version": 2.5
+}

processing_minicpmv.py

@@ -0,0 +1,247 @@
+# 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:
+        """
+        Main method to prepare for the model one or several sequences(s) and image(s). This method forwards the `text`
+        and `kwargs` arguments to LlamaTokenizerFast's [`~LlamaTokenizerFast.__call__`] if `text` is not `None` to encode
+        the text. To prepare the image(s), this method forwards the `images` and `kwrags` arguments to
+        LlavaNextImageProcessor's [`~LlavaNextImageProcessor.__call__`] if `images` is not `None`. Please refer to the doctsring
+        of the above two methods for more information.
+
+        Args:
+            text (`str`, `List[str]`, `List[List[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)
+        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[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_bounds": [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

resampler.py

@@ -1,17 +1,9 @@
 from functools import partial
 import numpy as np
-import warnings
-from typing import Optional, Tuple
+
 import torch
 from torch import nn
-from torch import Tensor
-import torch.nn.functional as F 
-from torch.nn.functional import *
-from torch.nn.modules.activation import *
 from torch.nn.init import trunc_normal_
-from torch.nn.init import constant_, xavier_normal_, xavier_uniform_
-from transformers import PreTrainedModel
-from transformers.integrations import is_deepspeed_zero3_enabled
 
 def get_2d_sincos_pos_embed(embed_dim, image_size):
     """
@@ -63,7 +55,7 @@ def get_1d_sincos_pos_embed_from_grid_new(embed_dim, pos):
     emb = np.concatenate([emb_sin, emb_cos], axis=-1)  # (H, W, D)
     return emb
 
-    
+
 class Resampler(nn.Module):
     """
     A 2D perceiver-resampler network with one cross attention layers by
@@ -90,13 +82,14 @@ class Resampler(nn.Module):
         self.max_size = max_size
 
         self.query = nn.Parameter(torch.zeros(self.num_queries, embed_dim))
+        trunc_normal_(self.query, std=.02)
 
         if kv_dim is not None and kv_dim != embed_dim:
             self.kv_proj = nn.Linear(kv_dim, embed_dim, bias=False)
         else:
             self.kv_proj = nn.Identity()
 
-        self.attn = MultiheadAttention(embed_dim, num_heads)
+        self.attn = nn.MultiheadAttention(embed_dim, num_heads)
         self.ln_q = norm_layer(embed_dim)
         self.ln_kv = norm_layer(embed_dim)
 
@@ -104,10 +97,9 @@ class Resampler(nn.Module):
         self.proj = nn.Parameter((embed_dim ** -0.5) * torch.randn(embed_dim, embed_dim))
 
         self._set_2d_pos_cache(self.max_size)
+        self.apply(self._init_weights)
 
     def _set_2d_pos_cache(self, max_size, device='cpu'):
-        if is_deepspeed_zero3_enabled():
-            device='cuda'
         pos_embed = torch.from_numpy(get_2d_sincos_pos_embed(self.embed_dim, max_size)).float().to(device)
         self.register_buffer("pos_embed", pos_embed, persistent=False)
 
@@ -168,645 +160,4 @@ class Resampler(nn.Module):
         return x
 
     def _repeat(self, query, N: int):
-        return query.unsqueeze(1).repeat(1, N, 1)
-
-
-class MultiheadAttention(nn.MultiheadAttention):
-    def __init__(self, embed_dim, num_heads, dropout=0., bias=True, add_bias_kv=False, 
-                 add_zero_attn=False, kdim=None, vdim=None, batch_first=False, device=None, dtype=None):
-        super().__init__(embed_dim, num_heads, dropout, bias, add_bias_kv, add_zero_attn, kdim, vdim, batch_first, device, dtype)
-
-        # rewrite out_proj layer，with nn.Linear
-        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias, device=device, dtype=dtype)
-
-    def forward(
-                self,
-                query: Tensor,
-                key: Tensor,
-                value: Tensor,
-                key_padding_mask: Optional[Tensor] = None,
-                need_weights: bool = True,
-                attn_mask: Optional[Tensor] = None,
-                average_attn_weights: bool = True,
-                is_causal : bool = False) -> Tuple[Tensor, Optional[Tensor]]:
-        why_not_fast_path = ''
-        if ((attn_mask is not None and torch.is_floating_point(attn_mask))
-           or (key_padding_mask is not None) and torch.is_floating_point(key_padding_mask)):
-            why_not_fast_path = "floating-point masks are not supported for fast path."
-
-        is_batched = query.dim() == 3
-
-        key_padding_mask = F._canonical_mask(
-            mask=key_padding_mask,
-            mask_name="key_padding_mask",
-            other_type=F._none_or_dtype(attn_mask),
-            other_name="attn_mask",
-            target_type=query.dtype
-        )
-
-        attn_mask = F._canonical_mask(
-            mask=attn_mask,
-            mask_name="attn_mask",
-            other_type=None,
-            other_name="",
-            target_type=query.dtype,
-            check_other=False,
-        )
-
-
-        if not is_batched:
-            why_not_fast_path = f"input not batched; expected query.dim() of 3 but got {query.dim()}"
-        elif query is not key or key is not value:
-            # When lifting this restriction, don't forget to either
-            # enforce that the dtypes all match or test cases where
-            # they don't!
-            why_not_fast_path = "non-self attention was used (query, key, and value are not the same Tensor)"
-        elif self.in_proj_bias is not None and query.dtype != self.in_proj_bias.dtype:
-            why_not_fast_path = f"dtypes of query ({query.dtype}) and self.in_proj_bias ({self.in_proj_bias.dtype}) don't match"
-        elif self.in_proj_weight is None:
-            why_not_fast_path = "in_proj_weight was None"
-        elif query.dtype != self.in_proj_weight.dtype:
-            # this case will fail anyway, but at least they'll get a useful error message.
-            why_not_fast_path = f"dtypes of query ({query.dtype}) and self.in_proj_weight ({self.in_proj_weight.dtype}) don't match"
-        elif self.training:
-            why_not_fast_path = "training is enabled"
-        elif (self.num_heads % 2) != 0:
-            why_not_fast_path = "self.num_heads is not even"
-        elif not self.batch_first:
-            why_not_fast_path = "batch_first was not True"
-        elif self.bias_k is not None:
-            why_not_fast_path = "self.bias_k was not None"
-        elif self.bias_v is not None:
-            why_not_fast_path = "self.bias_v was not None"
-        elif self.add_zero_attn:
-            why_not_fast_path = "add_zero_attn was enabled"
-        elif not self._qkv_same_embed_dim:
-            why_not_fast_path = "_qkv_same_embed_dim was not True"
-        elif query.is_nested and (key_padding_mask is not None or attn_mask is not None):
-            why_not_fast_path = "supplying both src_key_padding_mask and src_mask at the same time \
-                                 is not supported with NestedTensor input"
-        elif torch.is_autocast_enabled():
-            why_not_fast_path = "autocast is enabled"
-
-        if not why_not_fast_path:
-            tensor_args = (
-                query,
-                key,
-                value,
-                self.in_proj_weight,
-                self.in_proj_bias,
-                self.out_proj.weight,
-                self.out_proj.bias,
-            )
-            # We have to use list comprehensions below because TorchScript does not support
-            # generator expressions.
-            if torch.overrides.has_torch_function(tensor_args):
-                why_not_fast_path = "some Tensor argument has_torch_function"
-            elif _is_make_fx_tracing():
-                why_not_fast_path = "we are running make_fx tracing"
-            elif not all(_check_arg_device(x) for x in tensor_args):
-                why_not_fast_path = ("some Tensor argument's device is neither one of "
-                                     f"cpu, cuda or {torch.utils.backend_registration._privateuse1_backend_name}")
-            elif torch.is_grad_enabled() and any(_arg_requires_grad(x) for x in tensor_args):
-                why_not_fast_path = ("grad is enabled and at least one of query or the "
-                                     "input/output projection weights or biases requires_grad")
-            if not why_not_fast_path:
-                merged_mask, mask_type = self.merge_masks(attn_mask, key_padding_mask, query)
-
-                if self.in_proj_bias is not None and self.in_proj_weight is not None:
-                    return torch._native_multi_head_attention(
-                        query,
-                        key,
-                        value,
-                        self.embed_dim,
-                        self.num_heads,
-                        self.in_proj_weight,
-                        self.in_proj_bias,
-                        self.out_proj.weight,
-                        self.out_proj.bias,
-                        merged_mask,
-                        need_weights,
-                        average_attn_weights,
-                        mask_type)
-
-        any_nested = query.is_nested or key.is_nested or value.is_nested
-        assert not any_nested, ("MultiheadAttention does not support NestedTensor outside of its fast path. " +
-                                f"The fast path was not hit because {why_not_fast_path}")
-
-        if self.batch_first and is_batched:
-            # make sure that the transpose op does not affect the "is" property
-            if key is value:
-                if query is key:
-                    query = key = value = query.transpose(1, 0)
-                else:
-                    query, key = (x.transpose(1, 0) for x in (query, key))
-                    value = key
-            else:
-                query, key, value = (x.transpose(1, 0) for x in (query, key, value))
-        
-        if not self._qkv_same_embed_dim:
-            attn_output, attn_output_weights = self.multi_head_attention_forward(
-                query, key, value, self.embed_dim, self.num_heads,
-                self.in_proj_weight, self.in_proj_bias,
-                self.bias_k, self.bias_v, self.add_zero_attn,
-                self.dropout, self.out_proj.weight, self.out_proj.bias,
-                training=self.training,
-                key_padding_mask=key_padding_mask, need_weights=need_weights,
-                attn_mask=attn_mask,
-                use_separate_proj_weight=True,
-                q_proj_weight=self.q_proj_weight, k_proj_weight=self.k_proj_weight,
-                v_proj_weight=self.v_proj_weight,
-                average_attn_weights=average_attn_weights,
-                is_causal=is_causal)
-        else:
-            attn_output, attn_output_weights = self.multi_head_attention_forward(
-                query, key, value, self.embed_dim, self.num_heads,
-                self.in_proj_weight, self.in_proj_bias,
-                self.bias_k, self.bias_v, self.add_zero_attn,
-                self.dropout, self.out_proj.weight, self.out_proj.bias,
-                training=self.training,
-                key_padding_mask=key_padding_mask,
-                need_weights=need_weights,
-                attn_mask=attn_mask,
-                average_attn_weights=average_attn_weights,
-                is_causal=is_causal)
-        if self.batch_first and is_batched:
-            return attn_output.transpose(1, 0), attn_output_weights
-        else:
-            return attn_output, attn_output_weights
-            
-    def multi_head_attention_forward(
-        self,
-        query: Tensor,
-        key: Tensor,
-        value: Tensor,
-        embed_dim_to_check: int,
-        num_heads: int,
-        in_proj_weight: Optional[Tensor],
-        in_proj_bias: Optional[Tensor],
-        bias_k: Optional[Tensor],
-        bias_v: Optional[Tensor],
-        add_zero_attn: bool,
-        dropout_p: float,
-        out_proj_weight: Tensor,
-        out_proj_bias: Optional[Tensor],
-        training: bool = True,
-        key_padding_mask: Optional[Tensor] = None,
-        need_weights: bool = True,
-        attn_mask: Optional[Tensor] = None,
-        use_separate_proj_weight: bool = False,
-        q_proj_weight: Optional[Tensor] = None,
-        k_proj_weight: Optional[Tensor] = None,
-        v_proj_weight: Optional[Tensor] = None,
-        static_k: Optional[Tensor] = None,
-        static_v: Optional[Tensor] = None,
-        average_attn_weights: bool = True,
-        is_causal: bool = False,
-    ) -> Tuple[Tensor, Optional[Tensor]]:
-        tens_ops = (query, key, value, in_proj_weight, in_proj_bias, bias_k, bias_v, out_proj_weight, out_proj_bias)
-        if has_torch_function(tens_ops):
-            return handle_torch_function(
-                multi_head_attention_forward,
-                tens_ops,
-                query,
-                key,
-                value,
-                embed_dim_to_check,
-                num_heads,
-                in_proj_weight,
-                in_proj_bias,
-                bias_k,
-                bias_v,
-                add_zero_attn,
-                dropout_p,
-                out_proj_weight,
-                out_proj_bias,
-                training=training,
-                key_padding_mask=key_padding_mask,
-                need_weights=need_weights,
-                attn_mask=attn_mask,
-                is_causal=is_causal,
-                use_separate_proj_weight=use_separate_proj_weight,
-                q_proj_weight=q_proj_weight,
-                k_proj_weight=k_proj_weight,
-                v_proj_weight=v_proj_weight,
-                static_k=static_k,
-                static_v=static_v,
-                average_attn_weights=average_attn_weights,
-            )
-    
-        is_batched = _mha_shape_check(query, key, value, key_padding_mask, attn_mask, num_heads)
-    
-        # For unbatched input, we unsqueeze at the expected batch-dim to pretend that the input
-        # is batched, run the computation and before returning squeeze the
-        # batch dimension so that the output doesn't carry this temporary batch dimension.
-        if not is_batched:
-            # unsqueeze if the input is unbatched
-            query = query.unsqueeze(1)
-            key = key.unsqueeze(1)
-            value = value.unsqueeze(1)
-            if key_padding_mask is not None:
-                key_padding_mask = key_padding_mask.unsqueeze(0)
-    
-        # set up shape vars
-        tgt_len, bsz, embed_dim = query.shape
-        src_len, _, _ = key.shape
-    
-        key_padding_mask = _canonical_mask(
-            mask=key_padding_mask,
-            mask_name="key_padding_mask",
-            other_type=_none_or_dtype(attn_mask),
-            other_name="attn_mask",
-            target_type=query.dtype
-        )
-    
-        if is_causal and attn_mask is None:
-            raise RuntimeError(
-                "Need attn_mask if specifying the is_causal hint. "
-                "You may use the Transformer module method "
-                "`generate_square_subsequent_mask` to create this mask."
-            )
-    
-        if is_causal and key_padding_mask is None and not need_weights:
-            # when we have a kpm or need weights, we need attn_mask
-            # Otherwise, we use the is_causal hint go as is_causal
-            # indicator to SDPA.
-            attn_mask = None
-        else:
-            attn_mask = _canonical_mask(
-                mask=attn_mask,
-                mask_name="attn_mask",
-                other_type=None,
-                other_name="",
-                target_type=query.dtype,
-                check_other=False,
-            )
-    
-            if key_padding_mask is not None:
-                # We have the attn_mask, and use that to merge kpm into it.
-                # Turn off use of is_causal hint, as the merged mask is no
-                # longer causal.
-                is_causal = False
-    
-        assert embed_dim == embed_dim_to_check, \
-            f"was expecting embedding dimension of {embed_dim_to_check}, but got {embed_dim}"
-        if isinstance(embed_dim, torch.Tensor):
-            # embed_dim can be a tensor when JIT tracing
-            head_dim = embed_dim.div(num_heads, rounding_mode='trunc')
-        else:
-            head_dim = embed_dim // num_heads
-        assert head_dim * num_heads == embed_dim, f"embed_dim {embed_dim} not divisible by num_heads {num_heads}"
-        if use_separate_proj_weight:
-            # allow MHA to have different embedding dimensions when separate projection weights are used
-            assert key.shape[:2] == value.shape[:2], \
-                f"key's sequence and batch dims {key.shape[:2]} do not match value's {value.shape[:2]}"
-        else:
-            assert key.shape == value.shape, f"key shape {key.shape} does not match value shape {value.shape}"
-    
-        #
-        # compute in-projection
-        #
-        if not use_separate_proj_weight:
-            assert in_proj_weight is not None, "use_separate_proj_weight is False but in_proj_weight is None"
-            q, k, v = _in_projection_packed(query, key, value, in_proj_weight, in_proj_bias)
-        else:
-            assert q_proj_weight is not None, "use_separate_proj_weight is True but q_proj_weight is None"
-            assert k_proj_weight is not None, "use_separate_proj_weight is True but k_proj_weight is None"
-            assert v_proj_weight is not None, "use_separate_proj_weight is True but v_proj_weight is None"
-            if in_proj_bias is None:
-                b_q = b_k = b_v = None
-            else:
-                b_q, b_k, b_v = in_proj_bias.chunk(3)
-            q, k, v = _in_projection(query, key, value, q_proj_weight, k_proj_weight, v_proj_weight, b_q, b_k, b_v)
-    
-        # prep attention mask
-    
-        if attn_mask is not None:
-            # ensure attn_mask's dim is 3
-            if attn_mask.dim() == 2:
-                correct_2d_size = (tgt_len, src_len)
-                if attn_mask.shape != correct_2d_size:
-                    raise RuntimeError(f"The shape of the 2D attn_mask is {attn_mask.shape}, but should be {correct_2d_size}.")
-                attn_mask = attn_mask.unsqueeze(0)
-            elif attn_mask.dim() == 3:
-                correct_3d_size = (bsz * num_heads, tgt_len, src_len)
-                if attn_mask.shape != correct_3d_size:
-                    raise RuntimeError(f"The shape of the 3D attn_mask is {attn_mask.shape}, but should be {correct_3d_size}.")
-            else:
-                raise RuntimeError(f"attn_mask's dimension {attn_mask.dim()} is not supported")
-    
-        # add bias along batch dimension (currently second)
-        if bias_k is not None and bias_v is not None:
-            assert static_k is None, "bias cannot be added to static key."
-            assert static_v is None, "bias cannot be added to static value."
-            k = torch.cat([k, bias_k.repeat(1, bsz, 1)])
-            v = torch.cat([v, bias_v.repeat(1, bsz, 1)])
-            if attn_mask is not None:
-                attn_mask = pad(attn_mask, (0, 1))
-            if key_padding_mask is not None:
-                key_padding_mask = pad(key_padding_mask, (0, 1))
-        else:
-            assert bias_k is None
-            assert bias_v is None
-    
-        #
-        # reshape q, k, v for multihead attention and make em batch first
-        #
-        q = q.view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1)
-        if static_k is None:
-            k = k.view(k.shape[0], bsz * num_heads, head_dim).transpose(0, 1)
-        else:
-            # TODO finish disentangling control flow so we don't do in-projections when statics are passed
-            assert static_k.size(0) == bsz * num_heads, \
-                f"expecting static_k.size(0) of {bsz * num_heads}, but got {static_k.size(0)}"
-            assert static_k.size(2) == head_dim, \
-                f"expecting static_k.size(2) of {head_dim}, but got {static_k.size(2)}"
-            k = static_k
-        if static_v is None:
-            v = v.view(v.shape[0], bsz * num_heads, head_dim).transpose(0, 1)
-        else:
-            # TODO finish disentangling control flow so we don't do in-projections when statics are passed
-            assert static_v.size(0) == bsz * num_heads, \
-                f"expecting static_v.size(0) of {bsz * num_heads}, but got {static_v.size(0)}"
-            assert static_v.size(2) == head_dim, \
-                f"expecting static_v.size(2) of {head_dim}, but got {static_v.size(2)}"
-            v = static_v
-    
-        # add zero attention along batch dimension (now first)
-        if add_zero_attn:
-            zero_attn_shape = (bsz * num_heads, 1, head_dim)
-            k = torch.cat([k, torch.zeros(zero_attn_shape, dtype=k.dtype, device=k.device)], dim=1)
-            v = torch.cat([v, torch.zeros(zero_attn_shape, dtype=v.dtype, device=v.device)], dim=1)
-            if attn_mask is not None:
-                attn_mask = pad(attn_mask, (0, 1))
-            if key_padding_mask is not None:
-                key_padding_mask = pad(key_padding_mask, (0, 1))
-    
-        # update source sequence length after adjustments
-        src_len = k.size(1)
-    
-        # merge key padding and attention masks
-        if key_padding_mask is not None:
-            assert key_padding_mask.shape == (bsz, src_len), \
-                f"expecting key_padding_mask shape of {(bsz, src_len)}, but got {key_padding_mask.shape}"
-            key_padding_mask = key_padding_mask.view(bsz, 1, 1, src_len).   \
-                expand(-1, num_heads, -1, -1).reshape(bsz * num_heads, 1, src_len)
-            if attn_mask is None:
-                attn_mask = key_padding_mask
-            else:
-                attn_mask = attn_mask + key_padding_mask
-    
-        # adjust dropout probability
-        if not training:
-            dropout_p = 0.0
-    
-        #
-        # (deep breath) calculate attention and out projection
-        #
-    
-        if need_weights:
-            B, Nt, E = q.shape
-            q_scaled = q / math.sqrt(E)
-    
-            assert not (is_causal and attn_mask is None), "FIXME: is_causal not implemented for need_weights"
-    
-            if attn_mask is not None:
-                attn_output_weights = torch.baddbmm(attn_mask, q_scaled, k.transpose(-2, -1))
-            else:
-                attn_output_weights = torch.bmm(q_scaled, k.transpose(-2, -1))
-            attn_output_weights = softmax(attn_output_weights, dim=-1)
-            if dropout_p > 0.0:
-                attn_output_weights = dropout(attn_output_weights, p=dropout_p)
-    
-            attn_output = torch.bmm(attn_output_weights, v)
-    
-            attn_output = attn_output.transpose(0, 1).contiguous().view(tgt_len * bsz, embed_dim)
-            attn_output = self.out_proj(attn_output)
-            attn_output = attn_output.view(tgt_len, bsz, attn_output.size(1))
-    
-            # optionally average attention weights over heads
-            attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)
-            if average_attn_weights:
-                attn_output_weights = attn_output_weights.mean(dim=1)
-    
-            if not is_batched:
-                # squeeze the output if input was unbatched
-                attn_output = attn_output.squeeze(1)
-                attn_output_weights = attn_output_weights.squeeze(0)
-            return attn_output, attn_output_weights
-        else:
-            # attn_mask can be either (L,S) or (N*num_heads, L, S)
-            # if attn_mask's shape is (1, L, S) we need to unsqueeze to (1, 1, L, S)
-            # in order to match the input for SDPA of (N, num_heads, L, S)
-            if attn_mask is not None:
-                if attn_mask.size(0) == 1 and attn_mask.dim() == 3:
-                    attn_mask = attn_mask.unsqueeze(0)
-                else:
-                    attn_mask = attn_mask.view(bsz, num_heads, -1, src_len)
-    
-            q = q.view(bsz, num_heads, tgt_len, head_dim)
-            k = k.view(bsz, num_heads, src_len, head_dim)
-            v = v.view(bsz, num_heads, src_len, head_dim)
-    
-            attn_output = F.scaled_dot_product_attention(q, k, v, attn_mask, dropout_p, is_causal)
-            attn_output = attn_output.permute(2, 0, 1, 3).contiguous().view(bsz * tgt_len, embed_dim)
-    
-            attn_output = self.out_proj(attn_output)
-            attn_output = attn_output.view(tgt_len, bsz, attn_output.size(1))
-            if not is_batched:
-                # squeeze the output if input was unbatched
-                attn_output = attn_output.squeeze(1)
-            return attn_output, None
-
-
-def _mha_shape_check(query: Tensor, key: Tensor, value: Tensor,
-                     key_padding_mask: Optional[Tensor], attn_mask: Optional[Tensor], num_heads: int):
-    # Verifies the expected shape for `query, `key`, `value`, `key_padding_mask` and `attn_mask`
-    # and returns if the input is batched or not.
-    # Raises an error if `query` is not 2-D (unbatched) or 3-D (batched) tensor.
-
-    # Shape check.
-    if query.dim() == 3:
-        # Batched Inputs
-        is_batched = True
-        assert key.dim() == 3 and value.dim() == 3, \
-            ("For batched (3-D) `query`, expected `key` and `value` to be 3-D"
-             f" but found {key.dim()}-D and {value.dim()}-D tensors respectively")
-        if key_padding_mask is not None:
-            assert key_padding_mask.dim() == 2, \
-                ("For batched (3-D) `query`, expected `key_padding_mask` to be `None` or 2-D"
-                 f" but found {key_padding_mask.dim()}-D tensor instead")
-        if attn_mask is not None:
-            assert attn_mask.dim() in (2, 3), \
-                ("For batched (3-D) `query`, expected `attn_mask` to be `None`, 2-D or 3-D"
-                 f" but found {attn_mask.dim()}-D tensor instead")
-    elif query.dim() == 2:
-        # Unbatched Inputs
-        is_batched = False
-        assert key.dim() == 2 and value.dim() == 2, \
-            ("For unbatched (2-D) `query`, expected `key` and `value` to be 2-D"
-             f" but found {key.dim()}-D and {value.dim()}-D tensors respectively")
-
-        if key_padding_mask is not None:
-            assert key_padding_mask.dim() == 1, \
-                ("For unbatched (2-D) `query`, expected `key_padding_mask` to be `None` or 1-D"
-                 f" but found {key_padding_mask.dim()}-D tensor instead")
-
-        if attn_mask is not None:
-            assert attn_mask.dim() in (2, 3), \
-                ("For unbatched (2-D) `query`, expected `attn_mask` to be `None`, 2-D or 3-D"
-                 f" but found {attn_mask.dim()}-D tensor instead")
-            if attn_mask.dim() == 3:
-                expected_shape = (num_heads, query.shape[0], key.shape[0])
-                assert attn_mask.shape == expected_shape, \
-                    (f"Expected `attn_mask` shape to be {expected_shape} but got {attn_mask.shape}")
-    else:
-        raise AssertionError(
-            f"query should be unbatched 2D or batched 3D tensor but received {query.dim()}-D query tensor")
-
-    return is_batched
-
-
-def _canonical_mask(
-        mask: Optional[Tensor],
-        mask_name: str,
-        other_type: Optional[DType],
-        other_name: str,
-        target_type: DType,
-        check_other: bool = True,
-) -> Optional[Tensor]:
-
-    if mask is not None:
-        _mask_dtype = mask.dtype
-        _mask_is_float = torch.is_floating_point(mask)
-        if _mask_dtype != torch.bool and not _mask_is_float:
-            raise AssertionError(
-                f"only bool and floating types of {mask_name} are supported")
-        if check_other and other_type is not None:
-            if _mask_dtype != other_type:
-                warnings.warn(
-                    f"Support for mismatched {mask_name} and {other_name} "
-                    "is deprecated. Use same type for both instead."
-                )
-        if not _mask_is_float:
-            mask = (
-                torch.zeros_like(mask, dtype=target_type)
-                .masked_fill_(mask, float("-inf"))
-            )
-    return mask
-
-
-def _none_or_dtype(input: Optional[Tensor]) -> Optional[DType]:
-    if input is None:
-        return None
-    elif isinstance(input, torch.Tensor):
-        return input.dtype
-    raise RuntimeError("input to _none_or_dtype() must be None or torch.Tensor")
-
-def _in_projection_packed(
-    q: Tensor,
-    k: Tensor,
-    v: Tensor,
-    w: Tensor,
-    b: Optional[Tensor] = None,
-) -> List[Tensor]:
-    r"""
-    Performs the in-projection step of the attention operation, using packed weights.
-    Output is a triple containing projection tensors for query, key and value.
-    Args:
-        q, k, v: query, key and value tensors to be projected. For self-attention,
-            these are typically the same tensor; for encoder-decoder attention,
-            k and v are typically the same tensor. (We take advantage of these
-            identities for performance if they are present.) Regardless, q, k and v
-            must share a common embedding dimension; otherwise their shapes may vary.
-        w: projection weights for q, k and v, packed into a single tensor. Weights
-            are packed along dimension 0, in q, k, v order.
-        b: optional projection biases for q, k and v, packed into a single tensor
-            in q, k, v order.
-    Shape:
-        Inputs:
-        - q: :math:`(..., E)` where E is the embedding dimension
-        - k: :math:`(..., E)` where E is the embedding dimension
-        - v: :math:`(..., E)` where E is the embedding dimension
-        - w: :math:`(E * 3, E)` where E is the embedding dimension
-        - b: :math:`E * 3` where E is the embedding dimension
-        Output:
-        - in output list :math:`[q', k', v']`, each output tensor will have the
-            same shape as the corresponding input tensor.
-    """
-    E = q.size(-1)
-    if k is v:
-        if q is k:
-            # self-attention
-            proj = linear(q, w, b)
-            # reshape to 3, E and not E, 3 is deliberate for better memory coalescing and keeping same order as chunk()
-            proj = proj.unflatten(-1, (3, E)).unsqueeze(0).transpose(0, -2).squeeze(-2).contiguous()
-            return proj[0], proj[1], proj[2]
-        else:
-            # encoder-decoder attention
-            w_q, w_kv = w.split([E, E * 2])
-            if b is None:
-                b_q = b_kv = None
-            else:
-                b_q, b_kv = b.split([E, E * 2])
-            q_proj = linear(q, w_q, b_q)
-            kv_proj = linear(k, w_kv, b_kv)
-            # reshape to 2, E and not E, 2 is deliberate for better memory coalescing and keeping same order as chunk()
-            kv_proj = kv_proj.unflatten(-1, (2, E)).unsqueeze(0).transpose(0, -2).squeeze(-2).contiguous()
-            return (q_proj, kv_proj[0], kv_proj[1])
-    else:
-        w_q, w_k, w_v = w.chunk(3)
-        if b is None:
-            b_q = b_k = b_v = None
-        else:
-            b_q, b_k, b_v = b.chunk(3)
-        return linear(q, w_q, b_q), linear(k, w_k, b_k), linear(v, w_v, b_v)
-
-
-def _in_projection(
-    q: Tensor,
-    k: Tensor,
-    v: Tensor,
-    w_q: Tensor,
-    w_k: Tensor,
-    w_v: Tensor,
-    b_q: Optional[Tensor] = None,
-    b_k: Optional[Tensor] = None,
-    b_v: Optional[Tensor] = None,
-) -> Tuple[Tensor, Tensor, Tensor]:
-    r"""
-    Performs the in-projection step of the attention operation. This is simply
-    a triple of linear projections, with shape constraints on the weights which
-    ensure embedding dimension uniformity in the projected outputs.
-    Output is a triple containing projection tensors for query, key and value.
-    Args:
-        q, k, v: query, key and value tensors to be projected.
-        w_q, w_k, w_v: weights for q, k and v, respectively.
-        b_q, b_k, b_v: optional biases for q, k and v, respectively.
-    Shape:
-        Inputs:
-        - q: :math:`(Qdims..., Eq)` where Eq is the query embedding dimension and Qdims are any
-            number of leading dimensions.
-        - k: :math:`(Kdims..., Ek)` where Ek is the key embedding dimension and Kdims are any
-            number of leading dimensions.
-        - v: :math:`(Vdims..., Ev)` where Ev is the value embedding dimension and Vdims are any
-            number of leading dimensions.
-        - w_q: :math:`(Eq, Eq)`
-        - w_k: :math:`(Eq, Ek)`
-        - w_v: :math:`(Eq, Ev)`
-        - b_q: :math:`(Eq)`
-        - b_k: :math:`(Eq)`
-        - b_v: :math:`(Eq)`
-        Output: in output triple :math:`(q', k', v')`,
-         - q': :math:`[Qdims..., Eq]`
-         - k': :math:`[Kdims..., Eq]`
-         - v': :math:`[Vdims..., Eq]`
-    """
-    Eq, Ek, Ev = q.size(-1), k.size(-1), v.size(-1)
-    assert w_q.shape == (Eq, Eq), f"expecting query weights shape of {(Eq, Eq)}, but got {w_q.shape}"
-    assert w_k.shape == (Eq, Ek), f"expecting key weights shape of {(Eq, Ek)}, but got {w_k.shape}"
-    assert w_v.shape == (Eq, Ev), f"expecting value weights shape of {(Eq, Ev)}, but got {w_v.shape}"
-    assert b_q is None or b_q.shape == (Eq,), f"expecting query bias shape of {(Eq,)}, but got {b_q.shape}"
-    assert b_k is None or b_k.shape == (Eq,), f"expecting key bias shape of {(Eq,)}, but got {b_k.shape}"
-    assert b_v is None or b_v.shape == (Eq,), f"expecting value bias shape of {(Eq,)}, but got {b_v.shape}"
-    return linear(q, w_q, b_q), linear(k, w_k, b_k), linear(v, w_v, b_v)
+        return query.unsqueeze(1).repeat(1, N, 1)

tokenization_minicpmv_fast.py

@@ -0,0 +1,51 @@
+import json
+
+from transformers.tokenization_utils_fast import PreTrainedTokenizerFast
+
+
+class MiniCPMVTokenizerFast(PreTrainedTokenizerFast):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.eot_token = "<|eot_id|>"
+        self.im_start = "<image>"
+        self.im_end = "</image>"
+        self.ref_start = "<ref>"
+        self.ref_end = "</ref>"
+        self.box_start = "<box>"
+        self.box_end = "</box>"
+        self.quad_start = "<quad>"
+        self.quad_end = "</quad>"
+        self.slice_start = "<slice>"
+        self.slice_end = "</slice>"
+
+    @property
+    def eos_id(self):
+        return self.eos_token_id
+
+    @property
+    def bos_id(self):
+        return self.bos_token_id
+
+    @property
+    def unk_id(self):
+        return self.unk_token_id
+
+    @property
+    def eot_id(self):
+        return self.convert_tokens_to_ids(self.eot_token)
+
+    @property
+    def im_start_id(self):
+        return self.convert_tokens_to_ids(self.im_start)
+
+    @property
+    def im_end_id(self):
+        return self.convert_tokens_to_ids(self.im_end)
+
+    @staticmethod
+    def escape(text: str) -> str:
+        return text
+
+    @staticmethod
+    def unescape(text: str) -> str:
+        return text

tokenizer_config.json

@@ -2051,7 +2051,7 @@
   },
   "auto_map": {
     "AutoTokenizer": [
-      "modeling_minicpmv.PreTrainedTokenizerFastWrapper",
+      "tokenization_minicpmv_fast.MiniCPMVTokenizerFast",
       null
     ]
   },
@@ -2063,10 +2063,10 @@
     "input_ids",
     "attention_mask"
   ],
-  "model_max_length": 1000000000000000019884624838656,
+  "model_max_length": 2048,
   "pad_token": "!",
   "padding_side": "right",
-  "tokenizer_class": "PreTrainedTokenizerFastWrapper",
+  "tokenizer_class": "MiniCPMVTokenizerFast",
   "truncation_side": "right",
   "unk_token": "<unk>"
 }