fix_import (#1)

- fix import (b330978e4027db0529fd11c72c013fba5145c917)
- fix import (6ff8b355df43fecb64be8fb9c1ad700765b38fd3)


Co-authored-by: Haiping Wu <haipingwu@users.noreply.huggingface.co>

image_embedding_phi3_v.py

@@ -1,322 +0,0 @@
-# coding=utf-8
-# Copyright 2024 Microsoft and the HuggingFace Inc. team. All rights reserved.
-#
-# 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.
-
-import torch
-from torch import nn
-from transformers import CLIPVisionConfig, CLIPVisionModel, PretrainedConfig
-from transformers.models.clip.modeling_clip import CLIPAttention
-from transformers.utils import logging
-
-try:
-    from flash_attn import flash_attn_func
-except ImportError:
-    pass
-
-logger = logging.get_logger(__name__)
-
-
-MAX_INPUT_ID = int(1e9)
-
-CLIP_VIT_LARGE_PATCH14_336_CONFIG = CLIPVisionConfig(
-  attention_dropout=0.0,
-  dropout=0.0,
-  hidden_act="quick_gelu",
-  hidden_size=1024,
-  image_size=336,
-  initializer_factor=1.0,
-  initializer_range=0.02,
-  intermediate_size=4096,
-  layer_norm_eps=1e-05,
-  num_attention_heads=16,
-  num_channels=3,
-  num_hidden_layers=24,
-  patch_size=14,
-  projection_dim=768
-)
-
-class CLIPAttentionFA2(CLIPAttention):
-    """Add flash attention 2 to CLIPAttention. (This is only used in the vision encoder)"""
-
-    def forward(self,
-        hidden_states,
-        attention_mask=None,
-        causal_attention_mask=None,
-        output_attentions=False,
-    ):
-        """Input shape: Batch x Time x Channel"""
-
-        assert attention_mask is None, "CLIPAttentionFA2 does not support attention_mask"
-        assert causal_attention_mask is None, "CLIPAttentionFA2 does not support causal_attention_mask"
-        assert output_attentions is False, "CLIPAttentionFA2 does not support output_attentions"
-
-        bsz, tgt_len, embed_dim = hidden_states.size()
-        query_states = self.q_proj(hidden_states).reshape(bsz, tgt_len, self.num_heads, self.head_dim)
-        key_states = self.k_proj(hidden_states).reshape(bsz, tgt_len, self.num_heads, self.head_dim)
-        value_states = self.v_proj(hidden_states).reshape(bsz, tgt_len, self.num_heads, self.head_dim)
-
-        attn_output = flash_attn_func(
-            query_states,
-            key_states,
-            value_states,
-            dropout_p=self.dropout if self.training else 0.0,
-            softmax_scale=self.scale,
-            causal=False,
-        ).reshape(bsz, tgt_len, embed_dim)
-
-        attn_output = self.out_proj(attn_output)
-        return attn_output, None
-
-
-class Phi3ImageEmbedding(nn.Module):
-    """Phi3 Image embedding."""
-
-    def __init__(self, config: PretrainedConfig, wte=None, **kwargs) -> None:
-        super().__init__()
-
-        # n_embed or hidden_size
-        hidden_size = config.n_embd if hasattr(config, 'n_embd') else config.hidden_size
-        if hasattr(config, 'embd_pdrop') or hasattr(config, 'embed_pdrop'):
-            embd_drop = config.embd_pdrop if hasattr(config, 'embd_pdrop') else config.embed_pdrop
-            self.drop = nn.Dropout(embd_drop)
-        else:
-            self.drop = None
-
-        self.wte = wte
-
-        if isinstance(config.img_processor, dict) and config.img_processor.get('name', None) == 'clip_vision_model':
-            assert 'model_name' in config.img_processor, 'model_name must be provided for CLIPVisionModel'
-            assert 'image_dim_out' in config.img_processor, 'image_dim_out must be provided for CLIPVisionModel'
-            assert 'num_img_tokens' in config.img_processor, 'num_img_tokens must be provided for CLIPVisionModel'
-            assert config.img_processor['model_name'] == 'openai/clip-vit-large-patch14-336'
-            clip_config = CLIP_VIT_LARGE_PATCH14_336_CONFIG
-            self.img_processor = CLIPVisionModel(clip_config)
-            image_dim_out = config.img_processor['image_dim_out']
-            self.num_img_tokens = config.img_processor['num_img_tokens']
-
-            # FA2 in CLIP
-            if config._attn_implementation == 'flash_attention_2':
-                for layer in self.img_processor.vision_model.encoder.layers:
-                    clip_fa2 = CLIPAttentionFA2(clip_config)
-                    del layer.self_attn
-                    layer.self_attn = clip_fa2
-        else:
-            raise NotImplementedError(f'img_processor = {config.img_processor}, not implemented')
-
-        self.image_dim_out = image_dim_out
-        self.img_sizes = None
-
-        # global_gn and sub_gn for hd transform, serves as line separator
-        self.use_hd_transform = kwargs.get('use_hd_transform', False)
-        self.with_learnable_separator = kwargs.get('with_learnable_separator', False)
-        self.hd_transform_order = kwargs.get('hd_transform_order', 'glb_sub')
-        # with_hd_transform and with_learnable_separator should have same value
-        assert self.use_hd_transform == self.with_learnable_separator, 'use_hd_transform and with_learnable_separator should have same value'
-        if self.with_learnable_separator:
-            assert self.use_hd_transform, 'learnable separator is only for hd transform'
-            # 1024 * 4, merge spatial to channel dimension
-            self.glb_GN = nn.Parameter(torch.zeros([1, 1, self.image_dim_out * 4]))
-            self.sub_GN = nn.Parameter(torch.zeros([1, 1, 1, self.image_dim_out * 4]))
-            logger.info(f'learnable separator enabled for hd transform, hd_transform_order = {self.hd_transform_order}')
-
-        projection_cls = kwargs.get('projection_cls', 'linear')
-        if projection_cls == 'linear':
-            self.img_projection = nn.Linear(image_dim_out, hidden_size)
-        elif projection_cls == 'mlp' and self.use_hd_transform:
-            dim_projection = hidden_size
-            depth = 2
-            layers = [nn.Linear(image_dim_out * 4, dim_projection)]
-            for _ in range(1, depth):
-                layers.extend([nn.GELU(),
-                                nn.Linear(dim_projection, dim_projection)])
-            self.img_projection = nn.Sequential(*layers)
-        elif projection_cls == 'mlp':
-            dim_projection = hidden_size
-            depth = 2
-            layers = [nn.Linear(image_dim_out, dim_projection)]
-            for _ in range(1, depth):
-                layers.extend([nn.GELU(),
-                                nn.Linear(dim_projection, dim_projection)])
-            self.img_projection = nn.Sequential(*layers)
-        else:
-            raise NotImplementedError(f'projection_cls = {projection_cls}, not implemented')
-
-        self.vocab_size = config.vocab_size
-        self.img_features = None
-
-        if isinstance(config.img_processor, dict):
-            self.layer_idx = config.img_processor.get('layer_idx', -2)
-            self.type_feature = config.img_processor.get('type_feature', 'patch')
-        else:
-            self.layer_idx = -2
-            self.type_feature = 'patch'
-
-
-    def set_img_features(self, img_features: torch.FloatTensor) -> None:
-        self.img_features = img_features
-
-    def set_img_sizes(self, img_sizes: torch.LongTensor) -> None:
-        self.img_sizes = img_sizes
-
-    def get_img_features(self, img_embeds: torch.FloatTensor) -> torch.FloatTensor:
-        LAYER_IDX = self.layer_idx
-        TYPE_FEATURE = self.type_feature
-
-        img_processor_output = self.img_processor(img_embeds, output_hidden_states=True)
-        img_feature = img_processor_output.hidden_states[LAYER_IDX]
-
-        if TYPE_FEATURE == "patch":
-            patch_feature = img_feature[:, 1:]
-            return patch_feature
-
-        raise NotImplementedError
-
-    def forward(
-        self, input_ids: torch.LongTensor, pixel_values: torch.FloatTensor, image_sizes=None
-    ) -> torch.FloatTensor:
-        input_shape = input_ids.size()
-        input_ids = input_ids.view(-1, input_shape[-1])
-
-        # positions for image tokens
-        positions = torch.nonzero((input_ids < 0) & (input_ids > -MAX_INPUT_ID), as_tuple=True)
-        has_image = len(positions[0].tolist()) > 0
-        input_ids = input_ids.clamp_min(0).clamp_max(self.vocab_size).detach()
-        hidden_states = self.wte(input_ids)
-
-        if has_image:
-            assert self.use_hd_transform
-            num_images, num_crops, c, h, w = pixel_values.shape
-            assert c == 3 and h == w == 336
-            img_features = self.get_img_features(pixel_values.flatten(0, 1)).reshape(
-                num_images, num_crops, -1, self.image_dim_out
-            )
-            image_features_proj = self.hd_feature_transform(img_features, image_sizes)
-            hidden_states = hidden_states.index_put(
-                positions, image_features_proj, accumulate=False
-            )
-
-        if self.drop is not None:
-            hidden_states = self.drop(hidden_states)
-
-        return hidden_states
-
-    def hd_feature_transform(self, image_features, image_sizes):
-        """
-        image_features: (num_images, num_crops+1, 24*24, 1024)
-        """
-        assert (
-            self.hd_transform_order == 'sub_glb'
-        ), f'hd_transform_order `{self.hd_transform_order}` not implemented'
-        if isinstance(self.img_projection, nn.Sequential):
-            target_device = self.img_projection[0].bias.device
-            target_dtype = self.img_projection[0].bias.dtype
-        else:  # It's a single nn.Linear layer
-            target_device = self.img_projection.bias.device
-            target_dtype = self.img_projection.bias.dtype
-
-        global_image_features = image_features[:, 0]  # (num_images, 24*24, 1024)
-        # global feature can be viewed as a special HD case with num_crops 1x1
-        global_image_features_hd = self.reshape_hd_patches_2x2merge(global_image_features, 1, 1)
-        global_image_features_hd_newline = self.add_image_newline(global_image_features_hd)
-
-        all_image_embeddings = []
-        # need a for loop to process each image because of different image sizes
-        # (patch arrangement is different for each image)
-        for i, img_size in enumerate(image_sizes):
-            h, w = img_size
-            h_crop = h // 336
-            w_crop = w // 336
-            num_crops = h_crop * w_crop
-
-            # NOTE: real num_crops is padded
-            # (num_crops, 24*24, 1024)
-            sub_image_features = image_features[i, 1 : 1 + num_crops]
-            sub_image_features_hd = self.reshape_hd_patches_2x2merge(
-                sub_image_features, h_crop, w_crop
-            )
-            sub_image_features_hd_newline = self.add_image_newline(sub_image_features_hd)
-
-            # [sub features, separator, global features]
-            all_image_embeddings.extend(
-                [
-                    sub_image_features_hd_newline.squeeze(0),  # (h_crop*12*(w_crop*12+1), 4096)
-                    self.glb_GN.squeeze(0),
-                    global_image_features_hd_newline[i],
-                ]
-            )
-
-        image_features_proj = self.img_projection(
-            torch.cat(all_image_embeddings, dim=0).to(target_device).to(target_dtype)
-        )
-
-        return image_features_proj
-
-    def reshape_hd_patches_2x2merge(self, image_features, h_crop, w_crop):
-        """
-        image_features: (num_images*num_crops, 24*24, 1024)
-        output: (num_images, h_crop*12, w_crop*12, 4096), h_crop*w_crop == num_crops
-        """
-        N, L, C = image_features.shape
-        assert L == 24 * 24 and C == 1024 and N % (h_crop * w_crop) == 0
-        num_images = N // (h_crop * w_crop)
-        H = int(L**0.5)
-        image_features_hd = (
-            image_features.reshape(N, H, H, C)  # N, 24, 24, 1024
-            .reshape(N, H // 2, 2, H // 2, 2, C)  # N, 12, 2, 12, 2, 1024
-            .permute(0, 1, 3, 2, 4, 5)  # N, 12, 12, 2, 2, 1024
-            .reshape(N, -1, 4 * C)  # N, 144, 4096
-            .reshape(
-                num_images, h_crop, w_crop, H // 2, H // 2, -1
-            )  # n_img, h_crop, w_crop, 12, 12, 4096
-            .permute(0, 1, 3, 2, 4, 5)  # n_img, h_crop, 12, w_crop, 12, 4096
-            .reshape(
-                num_images, h_crop * H // 2, w_crop * H // 2, 4 * C
-            )  # n_img, h_crop*12, w_crop*12, 4096
-        )
-
-        # alternative implementation using einops
-        # from einops import rearrange
-        # image_features_nhwc = rearrange(
-        #     image_features,
-        #     'N (H W) c -> N H W c',
-        #     H=H,
-        #     W=H,
-        # )
-        # image_features_2x2merge = rearrange(
-        #     image_features_nhwc,
-        #     'N (h h_pool) (w w_pool) c -> N h w (h_pool w_pool c)',
-        #     h_pool=2,
-        #     w_pool=2,
-        # )
-        # image_features_hd = rearrange(
-        #     image_features_2x2merge,
-        #     '(n_img h_crop w_crop) h w C -> n_img (h_crop h) (w_crop w) C',
-        #     h_crop=h_crop,
-        #     w_crop=w_crop,
-        # )
-
-        return image_features_hd
-
-    def add_image_newline(self, image_features_hd):
-        """
-        image_features_hd: (num_images, h_crop*12, w_crop*12, 4096)
-        output: (num_images, (h_crop*12) * (w_crop*12+1), 4096)
-        """
-        num_images, h, w, hid_dim = image_features_hd.shape
-        # add the newline token to the HD image feature patches
-        newline_embeddings = self.sub_GN.expand(num_images, h, -1, -1)  # (n_img, h, 1, hid_dim)
-        image_features_hd_newline = torch.cat(
-            [image_features_hd, newline_embeddings], dim=2
-        ).reshape(num_images, -1, hid_dim)
-        return image_features_hd_newline

image_processing_phi3_v.py

@@ -1,274 +0,0 @@
-# coding=utf-8
-# Copyright 2024 Microsoft and the HuggingFace Inc. team. All rights reserved.
-#
-# 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.
-
-"""Image processor class for Phi3-V."""
-
-from typing import List, Optional, Union
-
-import numpy as np
-
-from transformers.image_processing_utils import BaseImageProcessor, BatchFeature
-from transformers.image_transforms import (
-    convert_to_rgb,
-)
-from transformers.image_utils import (
-    OPENAI_CLIP_MEAN,
-    OPENAI_CLIP_STD,
-    ImageInput,
-    make_list_of_images,
-    valid_images,
-)
-from transformers.utils import TensorType, is_vision_available, logging
-
-from transformers import AutoImageProcessor
-
-logger = logging.get_logger(__name__)
-
-
-if is_vision_available():
-    from PIL import Image
-
-import torch
-import torchvision
-
-def padding_336(b):
-    width, height = b.size
-    tar = int(np.ceil(height / 336) * 336)
-    top_padding = int((tar - height)/2)
-    bottom_padding = tar - height - top_padding
-    left_padding = 0
-    right_padding = 0
-    b = torchvision.transforms.functional.pad(b, [left_padding, top_padding, right_padding, bottom_padding], fill=[255,255,255])
-
-    return b
-
-def calc_padded_size(width, height, padding_unit=336):  
-    target_height = int(np.ceil(height / padding_unit) * padding_unit)  
-    top_padding = int((target_height - height) / 2)  
-    bottom_padding = target_height - height - top_padding  
-    left_padding = 0  
-    right_padding = 0  
-    padded_width = width + left_padding + right_padding  
-    padded_height = height + top_padding + bottom_padding  
-    return padded_width, padded_height  
-
-def HD_transform(img, hd_num=16):
-    width, height = img.size
-    trans = False
-    if width < height:
-        img = img.transpose(Image.TRANSPOSE)
-        trans = True
-        width, height = img.size
-    ratio = (width/ height)
-    scale = 1
-    while scale*np.ceil(scale/ratio) <= hd_num:
-        scale += 1
-    scale -= 1
-    new_w = int(scale * 336)
-    new_h = int(new_w / ratio)
-
-    img = torchvision.transforms.functional.resize(img, [new_h, new_w],)
-    img = padding_336(img)
-    width, height = img.size
-    if trans:
-        img = img.transpose(Image.TRANSPOSE)
-
-    return img
-
-def calc_hd_transform_size(width, height, hd_num=16):  
-    transposed = False  
-    if width < height:  
-        width, height = height, width  
-        transposed = True  
-  
-    ratio = width / height  
-    scale = 1  
-    while scale * np.ceil(scale / ratio) <= hd_num:  
-        scale += 1  
-    scale -= 1  
-  
-    new_width = int(scale * 336)  
-    new_height = int(new_width / ratio)  
-  
-    padded_width, padded_height = calc_padded_size(new_width, new_height)  
-      
-    if transposed:  
-        padded_width, padded_height = padded_height, padded_width  
-  
-    return padded_width, padded_height  
-
-def pad_to_max_num_crops_tensor(images, max_crops=5):
-    """
-    images: B x 3 x H x W, B<=max_crops
-    """
-    B, _, H, W = images.shape
-    if B < max_crops:
-        pad = torch.zeros(max_crops - B, 3, H, W, dtype=images.dtype, device=images.device)
-        images = torch.cat([images, pad], dim=0)
-    return images
-
-
-class Phi3VImageProcessor(BaseImageProcessor):
-    r"""
-    Constructs a Phi3 image processor. Based on [`CLIPImageProcessor`] with incorporation of additional techniques
-    for processing high resolution images as explained in the [InternLM-XComposer2-4KHD](https://arxiv.org/pdf/2404.06512)
-
-    Args:
-        image_mean (`float` or `List[float]`, *optional*, defaults to `[0.48145466, 0.4578275, 0.40821073]`):
-            Mean to use if normalizing the image. This is a float or list of floats the length of the number of
-            channels in the image. Can be overridden by the `image_mean` parameter in the `preprocess` method.
-        image_std (`float` or `List[float]`, *optional*, defaults to `[0.26862954, 0.26130258, 0.27577711]`):
-            Standard deviation to use if normalizing the image. This is a float or list of floats the length of the
-            number of channels in the image. Can be overridden by the `image_std` parameter in the `preprocess` method.
-            Can be overridden by the `image_std` parameter in the `preprocess` method.
-        do_convert_rgb (`bool`, *optional*, defaults to `True`):
-            Whether to convert the image to RGB.
-    """
-
-    model_input_names = ["pixel_values"]
-
-    def __init__(
-        self,
-        num_crops: int = 1,
-        image_mean: Optional[Union[float, List[float]]] = None,
-        image_std: Optional[Union[float, List[float]]] = None,
-        do_convert_rgb: bool = True,
-        **kwargs,
-    ) -> None:
-        super().__init__(**kwargs)
-        self.num_crops = num_crops
-        self.image_mean = image_mean if image_mean is not None else OPENAI_CLIP_MEAN
-        self.image_std = image_std if image_std is not None else OPENAI_CLIP_STD
-        self.do_convert_rgb = do_convert_rgb
-    
-    def calc_num_image_tokens(
-            self, 
-            images: ImageInput 
-    ):
-        """ Calculate the number of image tokens for each image.
-        Args:
-            images (`ImageInput`):
-                Image to preprocess. Expects a single or batch of images with pixel values ranging from 0 to 255. If
-                passing in images with pixel values between 0 and 1, set `do_rescale=False`.
-        """
-        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 = [image.convert('RGB') for image in images]
-        # (H, W, C)
-        elems = [HD_transform(im, hd_num = self.num_crops) for im in images] 
-        shapes = [[im.size[1], im.size[0]] for im in elems]
-        num_img_tokens = [int((h//336*w//336+1)*144 + 1 + (h//336+1)*12) for h, w in shapes]
-        return num_img_tokens
-
-    def calc_num_image_tokens_from_image_size(self, width, height):
-        """
-        Calculate the number of image tokens for a given image size.
-        Args:
-            width (`int`): Width of the image.
-            height (`int`): Height of the image.
-        """
-        new_width, new_height = calc_hd_transform_size(width, height, hd_num=self.num_crops)  
-        num_img_tokens = int((new_height // 336 * new_width // 336 + 1) * 144 + 1 + (new_height // 336 + 1) * 12)  
-        return num_img_tokens
-
-    def preprocess(
-        self,
-        images: ImageInput,
-        image_mean: Optional[Union[float, List[float]]] = None,
-        image_std: Optional[Union[float, List[float]]] = None,
-        do_convert_rgb: bool = None,
-        return_tensors: Optional[Union[str, TensorType]] = None,
-    ):
-        """
-        Args:
-            images (`ImageInput`):
-                Image to preprocess. Expects a single or batch of images with pixel values ranging from 0 to 255. If
-                passing in images with pixel values between 0 and 1, set `do_rescale=False`.
-            image_mean (`float` or `List[float]`, *optional*, defaults to `self.image_mean`):
-                Image mean to use for normalization. Only has an effect if `do_normalize` is set to `True`.
-            image_std (`float` or `List[float]`, *optional*, defaults to `self.image_std`):
-                Image standard deviation to use for normalization. Only has an effect if `do_normalize` is set to
-                `True`.
-            do_convert_rgb (`bool`, *optional*, defaults to `self.do_convert_rgb`):
-                Whether to convert the image to RGB.
-            return_tensors (`str` or `TensorType`, *optional*):
-                The type of tensors to return. Can be one of:
-                - Unset: Return a list of `np.ndarray`.
-                - `TensorType.TENSORFLOW` or `'tf'`: Return a batch of type `tf.Tensor`.
-                - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`.
-                - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`.
-                - `TensorType.JAX` or `'jax'`: Return a batch of type `jax.numpy.ndarray`.
-        """
-        image_mean = image_mean if image_mean is not None else self.image_mean
-        image_std = image_std if image_std is not None else self.image_std
-        do_convert_rgb = do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb
-
-        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."
-            )
-
-        if do_convert_rgb:
-            images = [convert_to_rgb(image) for image in images]
-
-        image_sizes = []
-        img_processor = torchvision.transforms.Compose([
-            torchvision.transforms.ToTensor(),
-            torchvision.transforms.Normalize(image_mean, image_std)
-        ])
-
-        # PIL images
-        # HD_transform pad images to size of multiiply of 336, 336
-        # convert to RGB first
-        images = [image.convert('RGB') for image in images]
-        elems = [HD_transform(im, hd_num = self.num_crops) for im in images] 
-        # tensor transform and normalize
-        hd_images = [img_processor(im) for im in elems]
-        # create global image 
-        global_image = [torch.nn.functional.interpolate(im.unsqueeze(0).float(), size=(336, 336), mode='bicubic',).to(im.dtype) for im in hd_images]
-
-        # [(3, h, w)], where h, w is multiple of 336
-        shapes = [[im.size(1), im.size(2)] for im in hd_images]
-        num_img_tokens = [int(((h//336)*(w//336)+1)*144 + 1 + (h//336+1)*12) for h, w in shapes]
-        # reshape to channel dimension -> (num_images, num_crops, 3, 336, 336)
-        # (1, 3, h//336, 336, w//336, 336) -> (1, h//336, w//336, 3, 336, 336) -> (h//336*w//336, 3, 336, 336)
-        hd_images_reshape = [im.reshape(1, 3, h//336, 336, w//336, 336).permute(0,2,4,1,3,5).reshape(-1, 3, 336, 336).contiguous() for im, (h, w) in zip(hd_images, shapes)]
-        # concat global image and local image
-        hd_images_reshape = [torch.cat([_global_image] + [_im], dim=0) for _global_image, _im in zip(global_image, hd_images_reshape)]
-
-        # pad to max_num_crops
-        image_transformed = [pad_to_max_num_crops_tensor(im, self.num_crops+1) for im in hd_images_reshape]
-        image_transformed = torch.stack(image_transformed, dim=0)
-        image_sizes = [torch.LongTensor(_shapes) for _shapes in shapes]
-        padded_images = image_transformed
-        image_sizes = shapes
-
-        data = {"pixel_values": padded_images, 
-                "image_sizes": image_sizes,
-                "num_img_tokens": num_img_tokens
-                }
-
-        return BatchFeature(data=data, tensor_type=return_tensors)
-
-AutoImageProcessor.register("Phi3VImageProcessor", Phi3VImageProcessor)

modeling_phi3_v.py

@@ -45,8 +45,6 @@ from transformers.utils import (
     replace_return_docstrings,
 )
 from .configuration_phi3_v import Phi3VConfig
-from .image_embedding_phi3_v import Phi3ImageEmbedding
-
 
 try:
     from flash_attn import flash_attn_func, flash_attn_varlen_func
@@ -56,6 +54,310 @@ try:
 except ImportError:
     pass
 
+import torch
+from torch import nn
+from transformers import CLIPVisionConfig, CLIPVisionModel, PretrainedConfig
+from transformers.models.clip.modeling_clip import CLIPAttention
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+
+MAX_INPUT_ID = int(1e9)
+
+CLIP_VIT_LARGE_PATCH14_336_CONFIG = CLIPVisionConfig(
+  attention_dropout=0.0,
+  dropout=0.0,
+  hidden_act="quick_gelu",
+  hidden_size=1024,
+  image_size=336,
+  initializer_factor=1.0,
+  initializer_range=0.02,
+  intermediate_size=4096,
+  layer_norm_eps=1e-05,
+  num_attention_heads=16,
+  num_channels=3,
+  num_hidden_layers=24,
+  patch_size=14,
+  projection_dim=768
+)
+
+class CLIPAttentionFA2(CLIPAttention):
+    """Add flash attention 2 to CLIPAttention. (This is only used in the vision encoder)"""
+
+    def forward(self,
+        hidden_states,
+        attention_mask=None,
+        causal_attention_mask=None,
+        output_attentions=False,
+    ):
+        """Input shape: Batch x Time x Channel"""
+
+        assert attention_mask is None, "CLIPAttentionFA2 does not support attention_mask"
+        assert causal_attention_mask is None, "CLIPAttentionFA2 does not support causal_attention_mask"
+        assert output_attentions is False, "CLIPAttentionFA2 does not support output_attentions"
+
+        bsz, tgt_len, embed_dim = hidden_states.size()
+        query_states = self.q_proj(hidden_states).reshape(bsz, tgt_len, self.num_heads, self.head_dim)
+        key_states = self.k_proj(hidden_states).reshape(bsz, tgt_len, self.num_heads, self.head_dim)
+        value_states = self.v_proj(hidden_states).reshape(bsz, tgt_len, self.num_heads, self.head_dim)
+
+        attn_output = flash_attn_func(
+            query_states,
+            key_states,
+            value_states,
+            dropout_p=self.dropout if self.training else 0.0,
+            softmax_scale=self.scale,
+            causal=False,
+        ).reshape(bsz, tgt_len, embed_dim)
+
+        attn_output = self.out_proj(attn_output)
+        return attn_output, None
+
+
+class Phi3ImageEmbedding(nn.Module):
+    """Phi3 Image embedding."""
+
+    def __init__(self, config: PretrainedConfig, wte=None, **kwargs) -> None:
+        super().__init__()
+
+        # n_embed or hidden_size
+        hidden_size = config.n_embd if hasattr(config, 'n_embd') else config.hidden_size
+        if hasattr(config, 'embd_pdrop') or hasattr(config, 'embed_pdrop'):
+            embd_drop = config.embd_pdrop if hasattr(config, 'embd_pdrop') else config.embed_pdrop
+            self.drop = nn.Dropout(embd_drop)
+        else:
+            self.drop = None
+
+        self.wte = wte
+
+        if isinstance(config.img_processor, dict) and config.img_processor.get('name', None) == 'clip_vision_model':
+            assert 'model_name' in config.img_processor, 'model_name must be provided for CLIPVisionModel'
+            assert 'image_dim_out' in config.img_processor, 'image_dim_out must be provided for CLIPVisionModel'
+            assert 'num_img_tokens' in config.img_processor, 'num_img_tokens must be provided for CLIPVisionModel'
+            assert config.img_processor['model_name'] == 'openai/clip-vit-large-patch14-336'
+            clip_config = CLIP_VIT_LARGE_PATCH14_336_CONFIG
+            self.img_processor = CLIPVisionModel(clip_config)
+            image_dim_out = config.img_processor['image_dim_out']
+            self.num_img_tokens = config.img_processor['num_img_tokens']
+
+            # FA2 in CLIP
+            if config._attn_implementation == 'flash_attention_2':
+                for layer in self.img_processor.vision_model.encoder.layers:
+                    clip_fa2 = CLIPAttentionFA2(clip_config)
+                    del layer.self_attn
+                    layer.self_attn = clip_fa2
+        else:
+            raise NotImplementedError(f'img_processor = {config.img_processor}, not implemented')
+
+        self.image_dim_out = image_dim_out
+        self.img_sizes = None
+
+        # global_gn and sub_gn for hd transform, serves as line separator
+        self.use_hd_transform = kwargs.get('use_hd_transform', False)
+        self.with_learnable_separator = kwargs.get('with_learnable_separator', False)
+        self.hd_transform_order = kwargs.get('hd_transform_order', 'glb_sub')
+        # with_hd_transform and with_learnable_separator should have same value
+        assert self.use_hd_transform == self.with_learnable_separator, 'use_hd_transform and with_learnable_separator should have same value'
+        if self.with_learnable_separator:
+            assert self.use_hd_transform, 'learnable separator is only for hd transform'
+            # 1024 * 4, merge spatial to channel dimension
+            self.glb_GN = nn.Parameter(torch.zeros([1, 1, self.image_dim_out * 4]))
+            self.sub_GN = nn.Parameter(torch.zeros([1, 1, 1, self.image_dim_out * 4]))
+            logger.info(f'learnable separator enabled for hd transform, hd_transform_order = {self.hd_transform_order}')
+
+        projection_cls = kwargs.get('projection_cls', 'linear')
+        if projection_cls == 'linear':
+            self.img_projection = nn.Linear(image_dim_out, hidden_size)
+        elif projection_cls == 'mlp' and self.use_hd_transform:
+            dim_projection = hidden_size
+            depth = 2
+            layers = [nn.Linear(image_dim_out * 4, dim_projection)]
+            for _ in range(1, depth):
+                layers.extend([nn.GELU(),
+                                nn.Linear(dim_projection, dim_projection)])
+            self.img_projection = nn.Sequential(*layers)
+        elif projection_cls == 'mlp':
+            dim_projection = hidden_size
+            depth = 2
+            layers = [nn.Linear(image_dim_out, dim_projection)]
+            for _ in range(1, depth):
+                layers.extend([nn.GELU(),
+                                nn.Linear(dim_projection, dim_projection)])
+            self.img_projection = nn.Sequential(*layers)
+        else:
+            raise NotImplementedError(f'projection_cls = {projection_cls}, not implemented')
+
+        self.vocab_size = config.vocab_size
+        self.img_features = None
+
+        if isinstance(config.img_processor, dict):
+            self.layer_idx = config.img_processor.get('layer_idx', -2)
+            self.type_feature = config.img_processor.get('type_feature', 'patch')
+        else:
+            self.layer_idx = -2
+            self.type_feature = 'patch'
+
+
+    def set_img_features(self, img_features: torch.FloatTensor) -> None:
+        self.img_features = img_features
+
+    def set_img_sizes(self, img_sizes: torch.LongTensor) -> None:
+        self.img_sizes = img_sizes
+
+    def get_img_features(self, img_embeds: torch.FloatTensor) -> torch.FloatTensor:
+        LAYER_IDX = self.layer_idx
+        TYPE_FEATURE = self.type_feature
+
+        img_processor_output = self.img_processor(img_embeds, output_hidden_states=True)
+        img_feature = img_processor_output.hidden_states[LAYER_IDX]
+
+        if TYPE_FEATURE == "patch":
+            patch_feature = img_feature[:, 1:]
+            return patch_feature
+
+        raise NotImplementedError
+
+    def forward(
+        self, input_ids: torch.LongTensor, pixel_values: torch.FloatTensor, image_sizes=None
+    ) -> torch.FloatTensor:
+        input_shape = input_ids.size()
+        input_ids = input_ids.view(-1, input_shape[-1])
+
+        # positions for image tokens
+        positions = torch.nonzero((input_ids < 0) & (input_ids > -MAX_INPUT_ID), as_tuple=True)
+        has_image = len(positions[0].tolist()) > 0
+        input_ids = input_ids.clamp_min(0).clamp_max(self.vocab_size).detach()
+        hidden_states = self.wte(input_ids)
+
+        if has_image:
+            assert self.use_hd_transform
+            num_images, num_crops, c, h, w = pixel_values.shape
+            assert c == 3 and h == w == 336
+            img_features = self.get_img_features(pixel_values.flatten(0, 1)).reshape(
+                num_images, num_crops, -1, self.image_dim_out
+            )
+            image_features_proj = self.hd_feature_transform(img_features, image_sizes)
+            hidden_states = hidden_states.index_put(
+                positions, image_features_proj, accumulate=False
+            )
+
+        if self.drop is not None:
+            hidden_states = self.drop(hidden_states)
+
+        return hidden_states
+
+    def hd_feature_transform(self, image_features, image_sizes):
+        """
+        image_features: (num_images, num_crops+1, 24*24, 1024)
+        """
+        assert (
+            self.hd_transform_order == 'sub_glb'
+        ), f'hd_transform_order `{self.hd_transform_order}` not implemented'
+        if isinstance(self.img_projection, nn.Sequential):
+            target_device = self.img_projection[0].bias.device
+            target_dtype = self.img_projection[0].bias.dtype
+        else:  # It's a single nn.Linear layer
+            target_device = self.img_projection.bias.device
+            target_dtype = self.img_projection.bias.dtype
+
+        global_image_features = image_features[:, 0]  # (num_images, 24*24, 1024)
+        # global feature can be viewed as a special HD case with num_crops 1x1
+        global_image_features_hd = self.reshape_hd_patches_2x2merge(global_image_features, 1, 1)
+        global_image_features_hd_newline = self.add_image_newline(global_image_features_hd)
+
+        all_image_embeddings = []
+        # need a for loop to process each image because of different image sizes
+        # (patch arrangement is different for each image)
+        for i, img_size in enumerate(image_sizes):
+            h, w = img_size
+            h_crop = h // 336
+            w_crop = w // 336
+            num_crops = h_crop * w_crop
+
+            # NOTE: real num_crops is padded
+            # (num_crops, 24*24, 1024)
+            sub_image_features = image_features[i, 1 : 1 + num_crops]
+            sub_image_features_hd = self.reshape_hd_patches_2x2merge(
+                sub_image_features, h_crop, w_crop
+            )
+            sub_image_features_hd_newline = self.add_image_newline(sub_image_features_hd)
+
+            # [sub features, separator, global features]
+            all_image_embeddings.extend(
+                [
+                    sub_image_features_hd_newline.squeeze(0),  # (h_crop*12*(w_crop*12+1), 4096)
+                    self.glb_GN.squeeze(0),
+                    global_image_features_hd_newline[i],
+                ]
+            )
+
+        image_features_proj = self.img_projection(
+            torch.cat(all_image_embeddings, dim=0).to(target_device).to(target_dtype)
+        )
+
+        return image_features_proj
+
+    def reshape_hd_patches_2x2merge(self, image_features, h_crop, w_crop):
+        """
+        image_features: (num_images*num_crops, 24*24, 1024)
+        output: (num_images, h_crop*12, w_crop*12, 4096), h_crop*w_crop == num_crops
+        """
+        N, L, C = image_features.shape
+        assert L == 24 * 24 and C == 1024 and N % (h_crop * w_crop) == 0
+        num_images = N // (h_crop * w_crop)
+        H = int(L**0.5)
+        image_features_hd = (
+            image_features.reshape(N, H, H, C)  # N, 24, 24, 1024
+            .reshape(N, H // 2, 2, H // 2, 2, C)  # N, 12, 2, 12, 2, 1024
+            .permute(0, 1, 3, 2, 4, 5)  # N, 12, 12, 2, 2, 1024
+            .reshape(N, -1, 4 * C)  # N, 144, 4096
+            .reshape(
+                num_images, h_crop, w_crop, H // 2, H // 2, -1
+            )  # n_img, h_crop, w_crop, 12, 12, 4096
+            .permute(0, 1, 3, 2, 4, 5)  # n_img, h_crop, 12, w_crop, 12, 4096
+            .reshape(
+                num_images, h_crop * H // 2, w_crop * H // 2, 4 * C
+            )  # n_img, h_crop*12, w_crop*12, 4096
+        )
+
+        # alternative implementation using einops
+        # from einops import rearrange
+        # image_features_nhwc = rearrange(
+        #     image_features,
+        #     'N (H W) c -> N H W c',
+        #     H=H,
+        #     W=H,
+        # )
+        # image_features_2x2merge = rearrange(
+        #     image_features_nhwc,
+        #     'N (h h_pool) (w w_pool) c -> N h w (h_pool w_pool c)',
+        #     h_pool=2,
+        #     w_pool=2,
+        # )
+        # image_features_hd = rearrange(
+        #     image_features_2x2merge,
+        #     '(n_img h_crop w_crop) h w C -> n_img (h_crop h) (w_crop w) C',
+        #     h_crop=h_crop,
+        #     w_crop=w_crop,
+        # )
+
+        return image_features_hd
+
+    def add_image_newline(self, image_features_hd):
+        """
+        image_features_hd: (num_images, h_crop*12, w_crop*12, 4096)
+        output: (num_images, (h_crop*12) * (w_crop*12+1), 4096)
+        """
+        num_images, h, w, hid_dim = image_features_hd.shape
+        # add the newline token to the HD image feature patches
+        newline_embeddings = self.sub_GN.expand(num_images, h, -1, -1)  # (n_img, h, 1, hid_dim)
+        image_features_hd_newline = torch.cat(
+            [image_features_hd, newline_embeddings], dim=2
+        ).reshape(num_images, -1, hid_dim)
+        return image_features_hd_newline
+
+
 logger = logging.get_logger(__name__)
 
 _CHECKPOINT_FOR_DOC = "microsoft/Phi-3-vision-128k-instruct"

processing_phi3_v.py

@@ -27,7 +27,268 @@ from transformers.image_utils import ImageInput
 from transformers.processing_utils import ProcessorMixin
 from transformers.tokenization_utils_base import PaddingStrategy, TextInput, TruncationStrategy
 from transformers.utils import TensorType
-from .image_processing_phi3_v import Phi3VImageProcessor 
+
+
+"""Image processor class for Phi3-V."""
+
+from typing import List, Optional, Union
+
+import numpy as np
+
+from transformers.image_processing_utils import BaseImageProcessor, BatchFeature
+from transformers.image_transforms import (
+    convert_to_rgb,
+)
+from transformers.image_utils import (
+    OPENAI_CLIP_MEAN,
+    OPENAI_CLIP_STD,
+    ImageInput,
+    make_list_of_images,
+    valid_images,
+)
+from transformers.utils import TensorType, is_vision_available, logging
+
+from transformers import AutoImageProcessor
+
+logger = logging.get_logger(__name__)
+
+
+if is_vision_available():
+    from PIL import Image
+
+import torch
+import torchvision
+
+def padding_336(b):
+    width, height = b.size
+    tar = int(np.ceil(height / 336) * 336)
+    top_padding = int((tar - height)/2)
+    bottom_padding = tar - height - top_padding
+    left_padding = 0
+    right_padding = 0
+    b = torchvision.transforms.functional.pad(b, [left_padding, top_padding, right_padding, bottom_padding], fill=[255,255,255])
+
+    return b
+
+def calc_padded_size(width, height, padding_unit=336):  
+    target_height = int(np.ceil(height / padding_unit) * padding_unit)  
+    top_padding = int((target_height - height) / 2)  
+    bottom_padding = target_height - height - top_padding  
+    left_padding = 0  
+    right_padding = 0  
+    padded_width = width + left_padding + right_padding  
+    padded_height = height + top_padding + bottom_padding  
+    return padded_width, padded_height  
+
+def HD_transform(img, hd_num=16):
+    width, height = img.size
+    trans = False
+    if width < height:
+        img = img.transpose(Image.TRANSPOSE)
+        trans = True
+        width, height = img.size
+    ratio = (width/ height)
+    scale = 1
+    while scale*np.ceil(scale/ratio) <= hd_num:
+        scale += 1
+    scale -= 1
+    new_w = int(scale * 336)
+    new_h = int(new_w / ratio)
+
+    img = torchvision.transforms.functional.resize(img, [new_h, new_w],)
+    img = padding_336(img)
+    width, height = img.size
+    if trans:
+        img = img.transpose(Image.TRANSPOSE)
+
+    return img
+
+def calc_hd_transform_size(width, height, hd_num=16):  
+    transposed = False  
+    if width < height:  
+        width, height = height, width  
+        transposed = True  
+  
+    ratio = width / height  
+    scale = 1  
+    while scale * np.ceil(scale / ratio) <= hd_num:  
+        scale += 1  
+    scale -= 1  
+  
+    new_width = int(scale * 336)  
+    new_height = int(new_width / ratio)  
+  
+    padded_width, padded_height = calc_padded_size(new_width, new_height)  
+      
+    if transposed:  
+        padded_width, padded_height = padded_height, padded_width  
+  
+    return padded_width, padded_height  
+
+def pad_to_max_num_crops_tensor(images, max_crops=5):
+    """
+    images: B x 3 x H x W, B<=max_crops
+    """
+    B, _, H, W = images.shape
+    if B < max_crops:
+        pad = torch.zeros(max_crops - B, 3, H, W, dtype=images.dtype, device=images.device)
+        images = torch.cat([images, pad], dim=0)
+    return images
+
+
+class Phi3VImageProcessor(BaseImageProcessor):
+    r"""
+    Constructs a Phi3 image processor. Based on [`CLIPImageProcessor`] with incorporation of additional techniques
+    for processing high resolution images as explained in the [InternLM-XComposer2-4KHD](https://arxiv.org/pdf/2404.06512)
+
+    Args:
+        image_mean (`float` or `List[float]`, *optional*, defaults to `[0.48145466, 0.4578275, 0.40821073]`):
+            Mean to use if normalizing the image. This is a float or list of floats the length of the number of
+            channels in the image. Can be overridden by the `image_mean` parameter in the `preprocess` method.
+        image_std (`float` or `List[float]`, *optional*, defaults to `[0.26862954, 0.26130258, 0.27577711]`):
+            Standard deviation to use if normalizing the image. This is a float or list of floats the length of the
+            number of channels in the image. Can be overridden by the `image_std` parameter in the `preprocess` method.
+            Can be overridden by the `image_std` parameter in the `preprocess` method.
+        do_convert_rgb (`bool`, *optional*, defaults to `True`):
+            Whether to convert the image to RGB.
+    """
+
+    model_input_names = ["pixel_values"]
+
+    def __init__(
+        self,
+        num_crops: int = 1,
+        image_mean: Optional[Union[float, List[float]]] = None,
+        image_std: Optional[Union[float, List[float]]] = None,
+        do_convert_rgb: bool = True,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.num_crops = num_crops
+        self.image_mean = image_mean if image_mean is not None else OPENAI_CLIP_MEAN
+        self.image_std = image_std if image_std is not None else OPENAI_CLIP_STD
+        self.do_convert_rgb = do_convert_rgb
+    
+    def calc_num_image_tokens(
+            self, 
+            images: ImageInput 
+    ):
+        """ Calculate the number of image tokens for each image.
+        Args:
+            images (`ImageInput`):
+                Image to preprocess. Expects a single or batch of images with pixel values ranging from 0 to 255. If
+                passing in images with pixel values between 0 and 1, set `do_rescale=False`.
+        """
+        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 = [image.convert('RGB') for image in images]
+        # (H, W, C)
+        elems = [HD_transform(im, hd_num = self.num_crops) for im in images] 
+        shapes = [[im.size[1], im.size[0]] for im in elems]
+        num_img_tokens = [int((h//336*w//336+1)*144 + 1 + (h//336+1)*12) for h, w in shapes]
+        return num_img_tokens
+
+    def calc_num_image_tokens_from_image_size(self, width, height):
+        """
+        Calculate the number of image tokens for a given image size.
+        Args:
+            width (`int`): Width of the image.
+            height (`int`): Height of the image.
+        """
+        new_width, new_height = calc_hd_transform_size(width, height, hd_num=self.num_crops)  
+        num_img_tokens = int((new_height // 336 * new_width // 336 + 1) * 144 + 1 + (new_height // 336 + 1) * 12)  
+        return num_img_tokens
+
+    def preprocess(
+        self,
+        images: ImageInput,
+        image_mean: Optional[Union[float, List[float]]] = None,
+        image_std: Optional[Union[float, List[float]]] = None,
+        do_convert_rgb: bool = None,
+        return_tensors: Optional[Union[str, TensorType]] = None,
+    ):
+        """
+        Args:
+            images (`ImageInput`):
+                Image to preprocess. Expects a single or batch of images with pixel values ranging from 0 to 255. If
+                passing in images with pixel values between 0 and 1, set `do_rescale=False`.
+            image_mean (`float` or `List[float]`, *optional*, defaults to `self.image_mean`):
+                Image mean to use for normalization. Only has an effect if `do_normalize` is set to `True`.
+            image_std (`float` or `List[float]`, *optional*, defaults to `self.image_std`):
+                Image standard deviation to use for normalization. Only has an effect if `do_normalize` is set to
+                `True`.
+            do_convert_rgb (`bool`, *optional*, defaults to `self.do_convert_rgb`):
+                Whether to convert the image to RGB.
+            return_tensors (`str` or `TensorType`, *optional*):
+                The type of tensors to return. Can be one of:
+                - Unset: Return a list of `np.ndarray`.
+                - `TensorType.TENSORFLOW` or `'tf'`: Return a batch of type `tf.Tensor`.
+                - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`.
+                - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`.
+                - `TensorType.JAX` or `'jax'`: Return a batch of type `jax.numpy.ndarray`.
+        """
+        image_mean = image_mean if image_mean is not None else self.image_mean
+        image_std = image_std if image_std is not None else self.image_std
+        do_convert_rgb = do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb
+
+        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."
+            )
+
+        if do_convert_rgb:
+            images = [convert_to_rgb(image) for image in images]
+
+        image_sizes = []
+        img_processor = torchvision.transforms.Compose([
+            torchvision.transforms.ToTensor(),
+            torchvision.transforms.Normalize(image_mean, image_std)
+        ])
+
+        # PIL images
+        # HD_transform pad images to size of multiiply of 336, 336
+        # convert to RGB first
+        images = [image.convert('RGB') for image in images]
+        elems = [HD_transform(im, hd_num = self.num_crops) for im in images] 
+        # tensor transform and normalize
+        hd_images = [img_processor(im) for im in elems]
+        # create global image 
+        global_image = [torch.nn.functional.interpolate(im.unsqueeze(0).float(), size=(336, 336), mode='bicubic',).to(im.dtype) for im in hd_images]
+
+        # [(3, h, w)], where h, w is multiple of 336
+        shapes = [[im.size(1), im.size(2)] for im in hd_images]
+        num_img_tokens = [int(((h//336)*(w//336)+1)*144 + 1 + (h//336+1)*12) for h, w in shapes]
+        # reshape to channel dimension -> (num_images, num_crops, 3, 336, 336)
+        # (1, 3, h//336, 336, w//336, 336) -> (1, h//336, w//336, 3, 336, 336) -> (h//336*w//336, 3, 336, 336)
+        hd_images_reshape = [im.reshape(1, 3, h//336, 336, w//336, 336).permute(0,2,4,1,3,5).reshape(-1, 3, 336, 336).contiguous() for im, (h, w) in zip(hd_images, shapes)]
+        # concat global image and local image
+        hd_images_reshape = [torch.cat([_global_image] + [_im], dim=0) for _global_image, _im in zip(global_image, hd_images_reshape)]
+
+        # pad to max_num_crops
+        image_transformed = [pad_to_max_num_crops_tensor(im, self.num_crops+1) for im in hd_images_reshape]
+        image_transformed = torch.stack(image_transformed, dim=0)
+        image_sizes = [torch.LongTensor(_shapes) for _shapes in shapes]
+        padded_images = image_transformed
+        image_sizes = shapes
+
+        data = {"pixel_values": padded_images, 
+                "image_sizes": image_sizes,
+                "num_img_tokens": num_img_tokens
+                }
+
+        return BatchFeature(data=data, tensor_type=return_tensors)
+
+AutoImageProcessor.register("Phi3VImageProcessor", Phi3VImageProcessor)
+
 transformers.Phi3VImageProcessor = Phi3VImageProcessor 
 
 class Phi3VProcessor(ProcessorMixin):