Delete vit_model.py

vit_model.py

@@ -1,1083 +0,0 @@
-import json
-import types
-import math
-import torch
-from torch import Tensor, nn
-import torch.nn.functional as F
-from typing import List, Tuple, Optional, Union
-from contextlib import contextmanager
-from transformers.modeling_attn_mask_utils import (
-    _prepare_4d_causal_attention_mask_for_sdpa,  
-    _prepare_4d_causal_attention_mask_for_sdpa, 
-    _prepare_4d_causal_attention_mask,
-)
-from transformers.models.clip.configuration_clip import CLIPVisionConfig
-from transformers.modeling_outputs import BaseModelOutputWithPooling
-from .modeling_hunyuan import HunYuanDecoderLayer, HunYuanRMSNorm
-from .configuration_hunyuan import HunYuanConfig
-
-
-def NaVitForward(input_ids, encoder_input, vit, image_tensors, images_pos, vit_input_resolution, im_start_id, im_end_id, image_token_id, anyres_vit_two_views, dtype):
-    # input_ids: (B, L)
-    # encoder_input: (L, B, E)
-    # image_tensors [[Tensor],...,[Tensor]]
-    # image_pos [[Tensor],...,[Tensor]]
-    # tokenizer = get_tokenizer()
-    b = len(input_ids)
-    img_embs = None
-    all_nums = sum([len(tensors) for tensors in image_tensors]) if image_tensors else 0
-    if all_nums != 0:
-        img_embs, img_batch_pos = vit(image_tensors)
-    else:
-        # when no input image, initialize a fake tensor
-        pad_nums = 1
-        image_tensors = [[torch.rand(3, vit_input_resolution, vit_input_resolution, dtype=dtype, device=torch.cuda.current_device()) for _ in range(pad_nums)]]
-        img_embs, img_batch_pos = vit(image_tensors)
-
-    encoder_input = encoder_input.clone()
-    if all_nums > 0:
-        assert len(images_pos) == len(img_batch_pos), \
-                (len(images_pos), len(img_batch_pos))
-        start_token_id = im_start_id
-        end_token_id = im_end_id
-        placeholder_id = image_token_id
-        for idx in range(len(images_pos)):
-            assert len(images_pos[idx]) == len(img_batch_pos[idx]), \
-                (len(images_pos[idx]), len(img_batch_pos[idx]))
-            for p_img_pos_in_batch, p_batch_img_pos in zip(img_batch_pos[idx], images_pos[idx]):
-                # the positions to be filled [s_start, s_end)
-                s_idx, s_start, s_end = p_img_pos_in_batch
-                current_embs = img_embs[s_idx, s_start:s_end]
-                im_s, im_e = p_batch_img_pos
-                assert len(current_embs) == im_e - im_s, \
-                        (img_embs.shape, (s_start, s_end, s_idx), current_embs.shape, (im_s, im_e, idx))
-                if not anyres_vit_two_views:
-                    assert input_ids[idx, im_s - 1] == start_token_id, \
-                            input_ids[idx, im_s - 1]
-                    assert input_ids[idx, im_e] == end_token_id, \
-                            input_ids[idx, im_e]
-                assert (input_ids[idx, im_s:im_e] == placeholder_id).all(), \
-                        f'The tokens to be filled are not the placeholder_id {placeholder_id}: {(input_ids[idx, im_s:im_e] == placeholder_id).sum()} vs {im_e - im_s}'
-                encoder_input[idx, im_s:im_e] = current_embs
-    else:
-        # when no input image, to mask vit value
-        vit_mask = torch.zeros([1, img_embs.shape[0]], device=torch.cuda.current_device())
-        current_embs = img_embs[0, :]
-        encoder_input[0, 1:img_embs.shape[0] + 1] = encoder_input[0, 1:img_embs.shape[0] + 1] * (1 - vit_mask) + current_embs * vit_mask
-    return encoder_input, input_ids
-
-
-def VitForward(input_ids, encoder_input, vit, vit_linear_encoder, image_tensors, images_pos, vit_input_resolution, vit_mapping_type, vit_patch, vit_token):
-    vit_patch_mlp = (vit_patch > 1 and vit_mapping_type == 'mlp') or vit_patch == 0
-
-    b = len(input_ids)
-    if images_pos is None:
-        images_pos = torch.ones([len(input_ids), 1, 3])
-        images_pos[:, :, 1] = images_pos[:, :, 1]*(vit_token + 1)
-        images_pos = images_pos.long()
-
-    real_image_nums = []
-    image_tensors = image_tensors.view(b, -1, 3, vit_input_resolution, vit_input_resolution)
-    real_images = []
-
-    all_nums = 0
-    img_index = []
-    for s in range(len(images_pos)):
-        real_image_num = 0
-        for (im_s, im_e,index) in images_pos[s]:
-            if im_s == -1:
-                break
-            real_image_num += 1
-            all_nums += 1
-            img_index.append(index)
-
-        real_image_nums.append(real_image_num)
-        real_images.append(image_tensors[s][:real_image_num])
-
-    if vit_patch == 1:
-        img_index = None
-
-    if all_nums == 0:
-        # when no input image, initialize a fake tensor
-        img_input = torch.rand(b, 3, vit_input_resolution, vit_input_resolution).cuda().type(image_tensors.dtype)
-        img_embs = vit(img_input)
-        img_embs = vit_linear_encoder(img_embs)
-    else:
-        img_input = torch.cat(real_images)
-        img_embs = vit(img_input, img_index = img_index)
-        img_embs = vit_linear_encoder(img_embs)
-
-    encoder_input = encoder_input.clone()
-    start = 0
-    if all_nums > 0:
-        for s, real_image_len in enumerate(real_image_nums):
-            current_embs = img_embs[start:start + real_image_len, :] #[30, 256, 4096]
-            for ss in range(current_embs.shape[0]):
-                im_s, im_e, index = images_pos[s, ss]
-                # 子图特征更少
-                if index > 0 and vit_patch_mlp:
-                    encoder_input[s, im_s:im_e,] = current_embs[ss, :(im_e-im_s)]
-                else:
-                    encoder_input[s, im_s:im_e] = current_embs[ss, :]
-            start = start + real_image_len
-    else:
-        # when no input image, to mask vit value
-        for s in range(b):
-            vit_mask = torch.zeros([vit_token, 1]).cuda()
-            current_embs = img_embs[:, start:start + 1]
-            encoder_input[1:vit_token + 1, s] = encoder_input[1:vit_token + 1, s] * (1 - vit_mask) + current_embs[:, 0, :] * vit_mask
-            start = start + 1
-    return encoder_input, input_ids
-
-
-def group_images_by_max_seq_len(
-    images: List[List[Tensor]], patch_size: int, 
-    max_seq_len: int, adaptor_patch_size: int, 
-    add_cls_token: bool = False) -> List[List[Tensor]]:
-
-    groups = []
-    group = []
-    pos_groups = []
-    seq_len = 0
-    num_images = 0
-    for image_list in images:
-        pos_group = []
-        for image in image_list:
-            num_images += 1
-            assert isinstance(image, Tensor)
-
-            image_dims = image.shape[-2:]
-            ph, pw = map(lambda t: t // patch_size, image_dims)
-
-            image_seq_len = (ph * pw)
-            new_image_seq_len = image_seq_len
-            grouped_len = seq_len + image_seq_len
-            if add_cls_token:
-                new_image_seq_len += 1
-                grouped_len += num_images
-                
-            assert new_image_seq_len <= max_seq_len, f'image with dimensions {image_dims} exceeds maximum sequence length'
-            
-            if grouped_len > max_seq_len:
-                groups.append(group)
-                group = []
-                seq_len = 0
-                num_images = 1
-
-            group.append(image)
-            start = seq_len // (adaptor_patch_size * adaptor_patch_size)
-            end = start + image_seq_len//(adaptor_patch_size * adaptor_patch_size)
-            batch_idx = len(groups)
-            pos_group.append([batch_idx, start, end])
-            seq_len += image_seq_len
-        pos_groups.append(pos_group)
-
-    if len(group) > 0:
-        groups.append(group)
-
-    return groups, pos_groups
-
-
-class AnyResCLIPVisionEmbeddings(nn.Module):
-    def __init__(self, config: CLIPVisionConfig):
-        super().__init__()
-
-        self.config = config
-        # self.sparse_attn_mask = args.sparse_attn_mask
-        # self.use_flash_attn = args.use_flash_attn
-        self.embed_dim = config.hidden_size
-        self.image_size = config.max_image_size
-        self.patch_size = config.patch_size
-        self.max_seq_len = config.max_vit_seq_len
-        self.adaptor_patch_size = config.adaptor_patch_size
-        self.anyres_vit_two_views = config.anyres_vit_two_views
-        self.vit_add_patchemb_bias = config.vit_add_patchemb_bias
-        self.vit_remove_prenorm = config.vit_remove_prenorm
-
-        self.patch_embedding = nn.Conv2d(
-            in_channels=config.num_channels,
-            out_channels=self.embed_dim,
-            kernel_size=self.patch_size,
-            stride=self.patch_size,
-            bias=self.vit_add_patchemb_bias,
-        )
-
-        self.num_patches = (self.image_size // self.patch_size) ** 2
-        self.skip_cls_token = True
-
-        # add interpolate_pos_encoding
-        if self.anyres_vit_two_views:
-            self.num_positions = self.num_patches
-            self.position_embedding = nn.Parameter(torch.randn(1, self.num_positions, self.embed_dim) * 0.02)
-        else:
-            self.num_positions = self.num_patches + 1
-            self.register_buffer("position_ids", torch.arange(self.num_positions).expand((1, -1)))
-            # self.position_ids = torch.arange(self.num_positions).expand((1, -1))
-            self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim)
-
-        if not self.vit_remove_prenorm:
-            self.pre_layernorm = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
-
-    def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor:
-        """
-        This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher
-        resolution images.
-
-        Source:
-        https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174
-        """
-        num_patches = embeddings.shape[1]
-        position_embeddings = self.position_embedding(self.position_ids)
-        patch_pos_embed = position_embeddings[:, 1:]
-        num_positions = position_embeddings.shape[1] - 1
-        if num_patches == num_positions and height == width:
-             return patch_pos_embed
-        # class_pos_embed = position_embeddings[:, 0]
-        dim = embeddings.shape[-1]
-        h0 = height // self.patch_size
-        w0 = width // self.patch_size
-        # we add a small number to avoid floating point error in the interpolation
-        # see discussion at https://github.com/facebookresearch/dino/issues/8
-        h0, w0 = h0 + 0.1, w0 + 0.1
-        patch_pos_embed = patch_pos_embed.reshape(1, int(math.sqrt(num_positions)), int(math.sqrt(num_positions)), dim)
-        patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2)
-        raw_type = patch_pos_embed.dtype
-        patch_pos_embed = nn.functional.interpolate(
-            patch_pos_embed.to(torch.float32, non_blocking=True),
-            scale_factor=(h0 / math.sqrt(num_positions), w0 / math.sqrt(num_positions)),
-            mode="bilinear",
-            align_corners=False,
-        )
-        patch_pos_embed = patch_pos_embed.to(raw_type, non_blocking=True)
-        assert int(h0) == patch_pos_embed.shape[-2] and int(w0) == patch_pos_embed.shape[-1]
-        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
-        return patch_pos_embed
-
-    def rescale_positional_embedding(self, out_size):
-        h, w = out_size
-        pos_embed_shape = int((self.position_embedding.shape[1]) ** 0.5)
-        if (h, w) == (pos_embed_shape, pos_embed_shape):
-            return self.position_embedding
-        rescaled_positional_embedding = \
-            self.position_embedding.new_zeros(1, h*w, self.position_embedding.shape[2])
-        pe_2d = self.position_embedding[0].T.contiguous().view(1, -1, pos_embed_shape, pos_embed_shape)
-        pe_2d = F.interpolate(pe_2d, out_size, mode='bilinear', align_corners=False).view(-1, h*w)
-        rescaled_positional_embedding[0] = pe_2d.T.contiguous()
-        return rescaled_positional_embedding
-
-    def forward_single(self, pixel_values: torch.FloatTensor) -> torch.Tensor:
-        if pixel_values.ndim == 3:
-            pixel_values = pixel_values[None]
-        batch_size, num_channels, height, width = pixel_values.shape
-
-        if self.anyres_vit_two_views:
-            # padding
-            pad_h = (self.patch_size - height % self.patch_size) % self.patch_size
-            pad_w = (self.patch_size - width % self.patch_size) % self.patch_size
-            pixel_values = F.pad(pixel_values, (0, pad_w, 0, pad_h))
-
-        patch_embeds = self.patch_embedding(pixel_values)  # shape = [*, width, grid, grid]
-        b, c, h, w = patch_embeds.shape
-
-        # (b, hw, c)
-        patch_embeds = patch_embeds.flatten(2).transpose(1, 2)
-        if self.anyres_vit_two_views:
-            embeddings = patch_embeds + self.rescale_positional_embedding(out_size=(h, w))
-        else:
-            embeddings = patch_embeds + self.interpolate_pos_encoding(patch_embeds, height, width)
-        if not self.vit_remove_prenorm:
-            embeddings = self.pre_layernorm(embeddings)
-        return embeddings, (h, w)
-
-    def forward(self, images: List[List[Tensor]]):
-        '''
-        Input:
-            images: List[List[Tensor]]
-
-        Return:
-            embeddings: Tensor (B, L, E)
-            attn_mask: Tensor (B, L, 2)
-            pos_groups: List[List[(batch_idx, start, end)]]
-        '''
-        batched_images, pos_groups = group_images_by_max_seq_len(
-            images, self.patch_size, self.max_seq_len, self.adaptor_patch_size, add_cls_token=not self.skip_cls_token)
-        max_seq_len = self.max_seq_len
-
-        # batched_images is a list of a list
-        B = len(batched_images)
-        L = max_seq_len
-        E = self.embed_dim
-
-        embeddings = torch.zeros(B, L, E, dtype=self.config.torch_dtype, requires_grad=True).cuda(non_blocking=True)
-        attn_mask = embeddings.new_full((B, 1, L, L), False, dtype=torch.bool)  # True presents compute
-        assert len(images) == len(pos_groups), (len(images), len(pos_groups))
-
-        batch_images = []
-        batch_pos = []
-        for images_i, pos_group in zip(images, pos_groups):
-            assert len(images_i) == len(pos_group), (len(images_i), len(pos_group))
-            for image, pos in zip(images_i, pos_group):    
-                batch_idx, start, end = pos
-                a2 = self.adaptor_patch_size ** 2
-                # recover the real number of the input image tokens
-                start *= a2
-                end *= a2
-                emb, _ = self.forward_single(image)
-                assert emb.ndim == 3, '(B, L, E)'
-                embeddings[batch_idx, start:end] = emb
-                attn_mask[batch_idx, :, start:end, start:end] = True
-        return embeddings, attn_mask, pos_groups
-
-
-class CLIPVisionEmbeddings(nn.Module):
-    def __init__(self, config: CLIPVisionConfig, add_pre_layernorm=False, skip_cls_token=True, vit_patch=1):
-        super().__init__()
-        self.config = config
-        self.embed_dim = config.hidden_size
-        self.image_size = config.image_size
-        self.image_size = config.vit_input_resolution
-        self.patch_size = config.patch_size
-
-        self.class_embedding = nn.Parameter(torch.randn(self.embed_dim))
-
-        self.patch_embedding = nn.Conv2d(
-            in_channels=config.num_channels,
-            out_channels=self.embed_dim,
-            kernel_size=self.patch_size,
-            stride=self.patch_size,
-            bias=False,
-        )
-
-        self.num_patches = (self.image_size // self.patch_size) ** 2
-
-        self.skip_cls_token = skip_cls_token
-
-        self.num_positions = self.num_patches + 1
-
-        self.register_buffer("position_ids", torch.arange(self.num_positions).expand((1, -1)))
-        if vit_patch > 1:
-            self.position_embedding = nn.Embedding(self.num_patches * (vit_patch ** 2 + 1) + 1, self.embed_dim)
-        # 0 支持最大16张图，目前写死了，如需其他的需要额外定义参数
-        elif vit_patch == 0:
-            self.position_embedding = nn.Embedding(self.num_patches * (16 ** 2 + 1) + 1, self.embed_dim)
-        else:
-            self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim)
-
-        if add_pre_layernorm:
-            self.pre_layernorm = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
-        else:
-            self.pre_layernorm = None
-
-    def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor:
-        """
-        This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher
-        resolution images.
-
-        Source:
-        https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174
-        """
-        num_patches = embeddings.shape[1] - 1
-        position_embeddings = self.position_embedding(self.position_ids)
-        num_positions = position_embeddings.shape[1] - 1
-        if num_patches == num_positions and height == width:
-             return position_embeddings
-        class_pos_embed = position_embeddings[:, 0]
-        patch_pos_embed = position_embeddings[:, 1:]
-        dim = embeddings.shape[-1]
-        h0 = height // self.config.patch_size
-        w0 = width // self.config.patch_size
-        # we add a small number to avoid floating point error in the interpolation
-        # see discussion at https://github.com/facebookresearch/dino/issues/8
-        h0, w0 = h0 + 0.1, w0 + 0.1
-        patch_pos_embed = patch_pos_embed.reshape(1, int(math.sqrt(num_positions)), int(math.sqrt(num_positions)), dim)
-        patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2)
-        raw_type = patch_pos_embed.dtype
-        patch_pos_embed = nn.functional.interpolate(
-            patch_pos_embed.float(),
-            scale_factor=(h0 / math.sqrt(num_positions), w0 / math.sqrt(num_positions)),
-            mode="bicubic",
-            align_corners=False,
-        )
-        # print(patch_pos_embed.shape)
-        patch_pos_embed = patch_pos_embed.to(raw_type)
-        assert int(h0) == patch_pos_embed.shape[-2] and int(w0) == patch_pos_embed.shape[-1]
-        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
-        return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1)
-
-
-    def forward(self, pixel_values: torch.FloatTensor, interpolate_pos_encoding: bool = False, img_index=None) -> torch.Tensor:
-        batch_size, num_channels, height, width = pixel_values.shape
-        patch_embeds = self.patch_embedding(pixel_values)  # shape = [*, width, grid, grid]
-        patch_embeds = patch_embeds.flatten(2).transpose(1, 2)
-        if self.skip_cls_token:
-            embeddings = patch_embeds
-            if img_index is None:
-                position_ids = self.position_ids[:,1:]
-                embeddings = embeddings + self.position_embedding(position_ids)
-            else:
-                position_ids = (torch.tensor(img_index).cuda() * (self.num_positions - 1)).unsqueeze(1).repeat(1, self.num_positions - 1) \
-                    + self.position_ids.expand(batch_size, -1)[:, 1:]
-                embeddings = embeddings + self.position_embedding(position_ids)
-        else:
-            class_embeds = self.class_embedding.expand(batch_size, 1, -1)
-            embeddings = torch.cat([class_embeds, patch_embeds], dim=1)
-            if interpolate_pos_encoding:
-                embeddings = embeddings + self.interpolate_pos_encoding(embeddings, height, width)
-            else:
-                if img_index is None:
-                    embeddings = embeddings + self.position_embedding(self.position_ids)
-                else:
-                    position_ids = self.position_ids.expand(batch_size,-1)[:,0].unsqueeze(1)
-                    new_position = (torch.tensor(img_index).cuda() * (self.num_positions -1)).unsqueeze(1).repeat(1,self.num_positions-1) +  self.position_ids.expand(batch_size,-1)[:,1:]
-                    position_ids = torch.cat([position_ids,new_position],dim=1)
-                    embeddings = embeddings + self.position_embedding(position_ids)
-        if self.pre_layernorm is not None:
-            embeddings = self.pre_layernorm(embeddings)
-        return embeddings
-
-
-class NaVitTransformer(nn.Module):
-    def __init__(self, config: HunYuanConfig, vit_config: CLIPVisionConfig):
-        super().__init__()
-        self.config = config
-        self.vit_config = vit_config
-        with self.prepare_args(config, vit_config):
-            self._use_sdpa = config._attn_implementation == "sdpa"
-            self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
-            self.layers = nn.ModuleList(
-                [HunYuanDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
-            )
-
-    @contextmanager
-    def prepare_args(self, config, vit_config):
-        hidden_act = config.hidden_act
-        hidden_size = config.hidden_size
-        ffn_hidden_size = config.intermediate_size
-        num_attention_heads = config.num_attention_heads
-        num_key_value_heads = config.num_key_value_heads
-        attention_head_dim = config.attention_head_dim
-        use_qk_norm = config.use_qk_norm
-        use_rotary_pos_emb =  config.use_rotary_pos_emb
-        num_hidden_layers = config.num_hidden_layers   
-        rms_norm_eps = config.rms_norm_eps
-        attention_dropout = config.attention_dropout
-        # hidden_dropout = config.hidden_dropout
-        norm_type = config.norm_type
-        attention_bias = config.attention_bias
-        mlp_bias = config.mlp_bias
-        use_mla = config.use_mla
-        num_experts = config.num_experts
-        _attn_implementation = config._attn_implementation
-        
-        config.hidden_act = vit_config.hidden_act
-        config.hidden_size = vit_config.hidden_size
-        config.intermediate_size = vit_config.intermediate_size
-        config.num_attention_heads = vit_config.num_attention_heads
-        config.num_key_value_heads = None
-        config.attention_head_dim = vit_config.hidden_size // vit_config.num_attention_heads
-        config.use_qk_norm = False
-        config.use_rotary_pos_emb = False
-        config.num_hidden_layers = vit_config.num_hidden_layers
-        config.rms_norm_eps = vit_config.layer_norm_eps
-        config.attention_dropout = vit_config.attention_dropout
-        # config.hidden_dropout = vit_config.hidden_dropout
-        config.norm_type = config.vit_norm_type
-        config.attention_bias = True
-        config.mlp_bias = True
-        config.use_mla = False
-        config.num_experts = 1
-        config._attn_implementation = "eager"
-        
-        yield
-        config.hidden_act = hidden_act
-        config.hidden_size = hidden_size
-        config.intermediate_size = ffn_hidden_size
-        config.num_attention_heads = num_attention_heads
-        config.num_key_value_heads = num_key_value_heads
-        config.attention_head_dim = attention_head_dim
-        config.use_qk_norm = use_qk_norm
-        config.use_rotary_pos_emb = use_rotary_pos_emb
-        config.num_hidden_layers = num_hidden_layers
-        config.rms_norm_eps = rms_norm_eps
-        config.attention_dropout = attention_dropout
-        # config.hidden_dropout = hidden_dropout
-        config.attention_bias = attention_bias
-        config.mlp_bias = mlp_bias
-        config.norm_type = norm_type
-        config.use_mla = use_mla
-        config.num_experts = num_experts
-        config._attn_implementation = _attn_implementation
-
-    def forward(
-        self,
-        pixel_values: Optional[torch.FloatTensor] = None,
-    ) -> Union[Tuple, BaseModelOutputWithPooling]:
-        
-        hidden_states, attention_mask, img_pos = self.embeddings(pixel_values)  
-        attention_mask = attention_mask.int()
-        batch_size, seq_length, _ = hidden_states.shape
-        past_key_values_length = 0  
-
-        if self._use_flash_attention_2:
-            # 2d mask is passed through the layers
-            attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None
-        elif self._use_sdpa:
-            # output_attentions=True can not be supported when using SDPA, and we fall back on
-            # the manual implementation that requires a 4D causal mask in all cases.
-            attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(
-                attention_mask,
-                (batch_size, seq_length),
-                hidden_states,
-                past_key_values_length,
-            )
-        else:
-            attention_mask = _prepare_4d_causal_attention_mask(
-                attention_mask,
-                (batch_size, seq_length),
-                hidden_states,
-                past_key_values_length,
-            )
-
-        for layer_idx, decoder_layer in enumerate(self.layers):
-            layer_outputs = decoder_layer(
-                hidden_states,
-                attention_mask=attention_mask
-            )
-            hidden_states = layer_outputs[0]
-
-        return hidden_states, img_pos
-
-
-class AnyResVitTransformer(NaVitTransformer):
-    def __init__(self, config: HunYuanConfig, vit_config: CLIPVisionConfig, anyres_vit_max_image_size):
-        super().__init__(config, vit_config)
-        old_anyres_vit_max_image_size = vit_config.max_image_size
-        anyres_vit_max_image_size = anyres_vit_max_image_size or old_anyres_vit_max_image_size
-        vit_config.max_image_size = anyres_vit_max_image_size
-        vit_config.torch_dtype = config.torch_dtype
-        vit_config.anyres_vit_two_views = config.anyres_vit_two_views
-        vit_config.vit_remove_prenorm = config.vit_remove_prenorm
-        vit_config.vit_add_patchemb_bias = config.vit_add_patchemb_bias
-        self.embeddings = AnyResCLIPVisionEmbeddings(vit_config)
-        vit_config.max_image_size = old_anyres_vit_max_image_size
-
-    def fix_embeddings_fn(self, pixel_values):
-        # (B, L, E)
-        embeddings, hw = self.embeddings.forward_single(pixel_values)
-        embeddings = self.embeddings.pre_layernorm(embeddings)
-        return embeddings
-
-
-class CLIPVisionTransformer(nn.Module):
-    def __init__(self, config: HunYuanConfig, vit_config: CLIPVisionConfig):
-        super().__init__()
-        embed_dim = vit_config.hidden_size
-
-        self.pre_layrnorm = nn.LayerNorm(embed_dim, eps=vit_config.layer_norm_eps)
-        self.embeddings = CLIPVisionEmbeddings(vit_config, skip_cls_token=config.skip_cls_token, vit_patch=config.vit_patch)
-
-        with self.prepare_args(config, vit_config):
-            self.layers = nn.ModuleList(
-                [HunYuanDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
-            )
-
-    @contextmanager
-    def prepare_args(self, config, vit_config):
-        hidden_act = config.hidden_act
-        hidden_size = config.hidden_size
-        ffn_hidden_size = config.intermediate_size
-        num_attention_heads = config.num_attention_heads
-        num_key_value_heads = config.num_key_value_heads
-        attention_head_dim = config.attention_head_dim
-        use_qk_norm = config.use_qk_norm
-        use_rotary_pos_emb =  config.use_rotary_pos_emb
-        num_hidden_layers = config.num_hidden_layers   
-        rms_norm_eps = config.rms_norm_eps
-        attention_dropout = config.attention_dropout
-        # hidden_dropout = config.hidden_dropout
-        norm_type = config.norm_type
-        attention_bias = config.attention_bias
-        mlp_bias = config.mlp_bias
-        use_mla = config.use_mla
-        num_experts = config.num_experts
-        _attn_implementation = config._attn_implementation
-
-        config.hidden_act = vit_config.hidden_act
-        config.hidden_size = vit_config.hidden_size
-        config.intermediate_size = vit_config.intermediate_size
-        config.num_attention_heads = vit_config.num_attention_heads
-        config.num_key_value_heads = None
-        config.attention_head_dim = vit_config.hidden_size // vit_config.num_attention_heads
-        config.use_qk_norm = False
-        config.use_rotary_pos_emb = False
-        config.num_hidden_layers = vit_config.num_hidden_layers
-        config.rms_norm_eps = vit_config.layer_norm_eps
-        config.attention_dropout = vit_config.attention_dropout
-        # config.hidden_dropout = 0.0
-        config.norm_type = "fused"
-        config.attention_bias = True
-        config.mlp_bias = True
-        config.use_mla = False
-        config.num_experts = 1
-        config._attn_implementation = "eager"
-
-        yield
-
-        config.hidden_act = hidden_act
-        config.hidden_size = hidden_size
-        config.intermediate_size = ffn_hidden_size
-        config.num_attention_heads = num_attention_heads
-        config.num_key_value_heads = num_key_value_heads
-        config.attention_head_dim = attention_head_dim
-        config.use_qk_norm = use_qk_norm
-        config.use_rotary_pos_emb = use_rotary_pos_emb
-        config.num_hidden_layers = num_hidden_layers
-        config.rms_norm_eps = rms_norm_eps
-        config.attention_dropout = attention_dropout
-        # config.hidden_dropout = hidden_dropout
-        config.norm_type = norm_type
-        config.attention_bias = attention_bias
-        config.mlp_bias = mlp_bias
-        config.use_mla = use_mla
-        config.num_experts = num_experts
-        config._attn_implementation = _attn_implementation
-
-    def forward(
-        self,
-        pixel_values: Optional[torch.FloatTensor] = None,
-        interpolate_pos_encoding: Optional[bool] = None,
-        img_index=None
-    ) -> Union[Tuple, BaseModelOutputWithPooling]:
-        r"""
-        Returns:
-
-        """
-        hidden_states = self.embeddings(pixel_values, interpolate_pos_encoding=interpolate_pos_encoding, img_index=img_index)
-        hidden_states = self.pre_layrnorm(hidden_states)
-        batch = hidden_states.shape[0]
-        seq_len = hidden_states.shape[1]
-        device = hidden_states.device
-        attention_mask = torch.ones(batch, 1, seq_len, seq_len, dtype=torch.float32, device=device)
-
-        for layer_idx, decoder_layer in enumerate(self.layers):
-            layer_outputs = decoder_layer(
-                hidden_states,
-                attention_mask=attention_mask
-            )
-            hidden_states = layer_outputs[0]
-
-        return hidden_states
-
-
-class Vit(torch.nn.Module):
-    def __init__(self, config, resampler_token=64, pool_rate=2):
-        super().__init__()
-        self.config = config
-        self.vit_mapping_type = config.vit_mapping_type
-        self.anyres_vit_max_image_size = config.anyres_vit_max_image_size
-        self.skip_cls_token = config.skip_cls_token
-        self.pool_rate = pool_rate
-        self.vit_type = self.config.vit_type
-        self.anyres_vit_two_views = self.config.anyres_vit_two_views
-        if self.vit_type in ['Vit-g', 'Vit-bigG', 'NaVit', 'EvaVit', 'AnyResVit']:
-            self.img_init(resampler_token, config.vit_input_resolution, config.vit_mapping_type, pool_rate)
-        else:
-            raise NotImplementedError(f"unsupported vit type: {self.vit_type}")
-    
-    def img_init(self, resampler_token=64, vit_input_resolution=224, vit_mapping_type='resampler', pool_rate=2):
-        if self.vit_type == 'AnyResVit':
-            vit_config = json.load(open(f"{self.config.vit_path}/config.json"))
-            self.vit_config = types.SimpleNamespace(**vit_config["vision_config"])
-            self.vit_config.image_size = vit_input_resolution
-            self.vit = AnyResVitTransformer(self.config, self.vit_config, self.anyres_vit_max_image_size)
-        elif self.vit_type == 'Vit-g':
-            vit_config = json.load(open(f"{self.config.vit_path}/config.json"))
-            self.vit_config = types.SimpleNamespace(**{**vit_config["vision_config_dict"],**vit_config["vision_config"]})
-            self.vit_config.vit_input_resolution = vit_input_resolution
-            self.vit = CLIPVisionTransformer(self.config, self.vit_config)  
-        else:
-           assert False, "other vit_types are not supported"
-
-        if self.vit_mapping_type == 'simple_conv_mlp':
-            self.perceive = SimpleConvMlp(self.vit_config.hidden_size, self.config.hidden_size, self.config.anyres_pooling_size, \
-                self.config.vit_used_rms_norm, self.config.rms_norm_eps, poolmlp=False, twoview=True)
-        elif self.vit_mapping_type == 'oryx_mlp':
-            self.perceive = OryxMLPv2(self.vit_config.hidden_size, self.config.hidden_size, twoview=True, use_pe=False)
-        elif self.vit_mapping_type == 'mlp':
-            self.mlp_depth = 2
-            # one mlp layer already in gpt_model.py
-            mlp_hidden_size = self.vit_config.hidden_size
-            if self.vit_type in ['NaVit', 'EvaVit']:
-                mlp_hidden_size *= self.vit_config.adaptor_patch_size **2
-            if self.mlp_depth > 1:
-                mlp_modules = [torch.nn.Linear(mlp_hidden_size, self.config.hidden_size), torch.nn.GELU()]
-                if self.vit_type in ['NaVit', 'EvaVit']:
-                    for _ in range(1, self.mlp_depth):
-                        mlp_modules.append(torch.nn.Linear(self.config.hidden_size, self.config.hidden_size))
-                        mlp_modules.append(torch.nn.GELU())
-                self.perceive = torch.nn.Sequential(*mlp_modules)
-        else:
-           assert False, "other vit_mapping_types are not supported"
-
-        self.vit_patch_mlp = (self.config.vit_patch > 1 and self.vit_mapping_type == 'mlp') or self.config.vit_patch == 0     
-        for name, param in self.named_parameters():
-            setattr(param, "is_vit_param", True)
-
-    def forward(self, images, img_index=None):
-        if self.vit_type in ['AnyResVit']:
-            dtype = self.config.torch_dtype
-            device = torch.cuda.current_device()
-
-            images_size = []
-            for i in range(len(images)):
-                images_size.append([])
-                for j in range(len(images[i])): 
-                    images_size[i].append((images[i][j].size()[1] // self.vit_config.patch_size, images[i][j].size()[2] // self.vit_config.patch_size))
-        
-            images_feats, img_batch_pos = self.vit(pixel_values=images)
-            a2 = self.vit_config.adaptor_patch_size ** 2
-
-            if self.anyres_vit_two_views:
-                step = 2
-            else:
-                step = 1
-            perceive_fn = lambda x, img_size, is_video: self.perceive(x, img_size, is_video=is_video)     
-            images_list = []
-            images_fix_i = 0
-            num_img_batch_pos = len(img_batch_pos)
-            for i in range(num_img_batch_pos): # batch_id
-                for j in range(0, len(img_batch_pos[i]), step): 
-                    if self.anyres_vit_two_views:
-                        lower_idx, lower_begin, lower_end = img_batch_pos[i][j]
-                        lower_begin = lower_begin * a2
-                        lower_end = lower_end * a2
-                        higher_idx, higher_begin, higher_end = img_batch_pos[i][j + 1]
-                        higher_begin = higher_begin * a2
-                        higher_end = higher_end * a2
-                        lower_res_feat = images_feats[lower_idx, lower_begin:lower_end].unsqueeze(0)
-                        higher_res_feat = images_feats[higher_idx, higher_begin:higher_end].unsqueeze(0)
-                        lower_images_size = images_size[i][j]
-                        higher_images_size = images_size[i][j + 1]
-                        images_list.append(self.perceive(lower_res_feat, lower_images_size, higher_res_feat, higher_images_size))
-                    else:
-                        idx, begin, end = img_batch_pos[i][j]
-                        begin = begin * a2
-                        end = end * a2
-                        is_video = hasattr(images[i][j],'_is_video') and images[i][j]._is_video 
-                        images_list.append(perceive_fn(images_feats[idx, begin:end].unsqueeze(0), images_size[i][j], is_video=is_video))
-
-            images = torch.cat(images_list, dim=1)
-
-            new_batch_pos = []
-            k = 0; cur_len = 0
-            for i in range(len(images_size)):
-                new_batch_pos.append([])
-                for j in range(0, len(images_size[i]), step):
-                    new_pos = [0, cur_len, cur_len + images_list[k].size(1)]
-                    cur_len += images_list[k].size(1)
-                    k += 1
-                    new_batch_pos[i].append(new_pos)
-            return images, new_batch_pos
-        elif self.vit_type == 'Vit-g':
-            images = self.vit(pixel_values=images, interpolate_pos_encoding=False, img_index=img_index)
-        else: 
-            assert False, "other vit_types are not supported"
-        
-        if self.vit_mapping_type == 'mlp':
-            if self.vit_type in ['Vit-g'] and not self.skip_cls_token:
-                images = images[:,1:,:]
-            b, v, d = images.shape
-            s = int(math.sqrt(v))
-            images = images.reshape(b, s, s, d)
-
-
-            if self.vit_patch_mlp and img_index is not None:
-                L_tensor = torch.tensor(img_index)
-                device = images.device
-                # 获取子图位置
-                nonzero_indices = torch.nonzero(L_tensor).squeeze().to(device)
-                # 获取主图位置
-                zero_indices = torch.nonzero(L_tensor == 0).squeeze().to(device)
-
-
-                images_nonzero = torch.index_select(images,0, nonzero_indices).to(device)
-                images_zero = torch.index_select(images, 0, zero_indices).to(device)
-
-                # 子图额外多pool一次
-                pool_rate = self.pool_rate * 2
-                images_nonzero = images_nonzero.reshape(-1, s // pool_rate, pool_rate, s // pool_rate, pool_rate, d)
-                images_nonzero = images_nonzero.permute(0, 1, 3, 5, 2, 4).reshape(-1, (s // pool_rate) * (s // pool_rate), d,
-                                                                      pool_rate*pool_rate).mean(-1)
-
-                # 为了组batch折衷方案
-                images_nonzero = F.pad(images_nonzero, (0, 0, 0, (s // self.pool_rate) * (s // self.pool_rate)- (s // pool_rate) * (s // pool_rate)))
-                images_zero = images_zero.reshape(-1, s // self.pool_rate, self.pool_rate, s // self.pool_rate, self.pool_rate, d)
-                images_zero = images_zero.permute(0, 1, 3, 5, 2, 4).reshape(-1, (s // self.pool_rate) * (s // self.pool_rate), d,
-                                                                  self.pool_rate*self.pool_rate).mean(-1)
-                # 组batch
-                images = torch.zeros(b, (s // self.pool_rate) * (s // self.pool_rate), d).to(device).to(images.dtype)
-                images.index_copy_(0, nonzero_indices, images_nonzero)
-                images.index_copy_(0, zero_indices, images_zero)
-
-                if self.mlp_depth >= 2:
-                    images = self.perceive(images)
-            else:
-                if s % self.pool_rate == 0:
-                    images = images.reshape(b, s//self.pool_rate, self.pool_rate, s//self.pool_rate, self.pool_rate, d)
-                    images = images.permute(0, 1, 3, 5, 2, 4).reshape(b, (s//self.pool_rate) * (s//self.pool_rate), d, -1).mean(-1)
-                    if self.mlp_depth >= 2:
-                        images = self.perceive(images)
-                else:
-                    raise ValueError
-        return images
-
-
-class SimpleConvMlp(nn.Module):
-    def __init__(self, in_channels, out_channels, anyres_pooling_size, vit_used_rms_norm, rms_norm_eps, twoview=False, poolmlp=True, cat_extra_token=True):
-        super().__init__()
-
-        embed_std = 1 / math.sqrt(out_channels)
-        if poolmlp:
-            # if args.learnable_mlp_pooling_size is not None:
-            #     in_channels *= args.learnable_mlp_pooling_size ** 2
-            self.proj = nn.Sequential(
-                nn.Linear(in_channels, out_channels),
-                nn.GELU()
-            )
-            self.vit_linear_encoder = nn.Linear(out_channels, out_channels)
-            self.image_newline = nn.Parameter(
-                torch.randn(out_channels) * embed_std
-            )
-        else:
-            self.proj = nn.Sequential(
-                nn.Conv2d(in_channels, in_channels * 2, kernel_size=anyres_pooling_size, stride=anyres_pooling_size),
-                nn.GELU(),
-                nn.Conv2d(in_channels * 2, in_channels * 4, kernel_size=1),
-            )
-            self.mlp = nn.Linear(in_channels * 4, out_channels)
-            self.image_newline = nn.Parameter(
-                torch.randn(in_channels * 4) * embed_std
-            )
-        self.poolmlp = poolmlp
-
-        self.image_begin = nn.Parameter(
-            torch.randn(out_channels) * embed_std
-        )
-        self.image_end = nn.Parameter(
-            torch.randn(out_channels) * embed_std
-        )
-        
-        if twoview:
-            self.image_sep = nn.Parameter(
-                torch.randn(out_channels) * embed_std
-            )
-
-        self.cat_extra_token = cat_extra_token
-        self.use_rms_norm = vit_used_rms_norm
-        if self.use_rms_norm:
-            self.before_rms = HunYuanRMSNorm(in_channels, eps=rms_norm_eps)
-            self.after_rms = HunYuanRMSNorm(out_channels, eps=rms_norm_eps)
-
-    def forward(self, x, size=(16,16), x2=None, size2=(16, 16), is_video=False):
-        return self.single_forward(x=x, size=size, x2=x2, size2=size2, is_video=is_video)
-
-    def single_forward(self, x, size=(16,16), x2=None, size2=(16, 16), is_video=False): 
-        remove_vit_special_tokens = False
-        learnable_mlp_pooling_size = None         
-        if self.use_rms_norm:
-            x = self.before_rms(x)            
-        h, w = size
-        dtype = x.dtype
-        x = x.permute(0, 2, 1).reshape(x.shape[0], -1, h, w)
-        if self.poolmlp:
-            if learnable_mlp_pooling_size is None:
-                x = F.avg_pool2d(x, anyres_pooling_size)
-                x = self.proj(x.permute(0, 2, 3, 1)) # b, h, w, c
-            else:
-                x = x.permute(0, 2, 3, 1)   # b, h, w, c
-                x = x.reshape(x.shape[0], h // learnable_mlp_pooling_size, learnable_mlp_pooling_size, 
-                                    w // learnable_mlp_pooling_size, learnable_mlp_pooling_size, -1)
-                x = x.permute(0, 1, 3, 2, 4, 5).reshape(x.shape[0], h // learnable_mlp_pooling_size, w // learnable_mlp_pooling_size, -1)
-                x = self.proj(x)
-            x = self.vit_linear_encoder(x)
-            b, h, w, c = x.shape
-            if not remove_vit_special_tokens:
-                x = torch.cat([
-                    x,
-                    self.image_newline.reshape(1, 1, 1, c).expand(b, h, 1, c).to(dtype, non_blocking=True)
-                ], dim=2)
-            x = x.reshape(b, -1, c)
-        else:
-            x = self.proj(x) #b,c,h,w
-            if is_video:
-                video_avgpool_size = 2
-                stride = 2
-                x = F.avg_pool2d(x, kernel_size = video_avgpool_size, stride = stride)
-            b, c, h, w = x.shape
-            if not remove_vit_special_tokens:
-                x = torch.cat([
-                    x,
-                    self.image_newline.reshape(1, c, 1, 1).expand(b, c, h, 1).to(dtype, non_blocking=True)
-                ], dim=-1)
-            x = x.reshape(b, c, -1).permute(0, 2, 1)
-            x = self.mlp(x)
-
-
-        if x2 is not None:
-            h2, w2 = size2
-            x2 = x2.permute(0, 2, 1).reshape(x2.shape[0], -1, h2, w2)
-            if self.poolmlp:
-                x2 = F.avg_pool2d(x2, 2)
-                x2 = self.proj(x2.permute(0, 2, 3, 1)) # b, h, w, c
-                x2 = self.vit_linear_encoder(x2)
-                b2, h2, w2, c2 = x2.shape
-                if not remove_vit_special_tokens:
-                    x2 = torch.cat([
-                        x2,
-                        self.image_newline.reshape(1, 1, 1, c2).expand(b2, h2, 1, c2).to(dtype, non_blocking=True)
-                    ], dim=2)
-                x2 = x2.reshape(b2, -1, c2)
-            else:
-                x2 = self.proj(x2)
-                b2, c2, h2, w2 = x2.shape
-                if not remove_vit_special_tokens:
-                    x2 = torch.cat([
-                        x2,
-                        self.image_newline.reshape(1, c2, 1, 1).expand(b2, c2, h2, 1).to(dtype, non_blocking=True)
-                    ], dim=-1)
-                x2 = x2.reshape(b2, c2, -1).permute(0, 2, 1) #b,n,c
-                x2 = self.mlp(x2)
-
-            sep = self.image_sep.reshape(1, 1, -1).expand(b2, 1, x2.shape[-1]).to(dtype, non_blocking=True)
-
-            x = torch.cat([x, sep, x2], dim=1)
-        
-        if self.cat_extra_token:
-            begin = self.image_begin.reshape(1, 1, -1).expand(b, 1, x.shape[-1]).to(dtype, non_blocking=True)
-            end = self.image_end.reshape(1, 1, -1).expand(b, 1, x.shape[-1]).to(dtype, non_blocking=True)
-            x = torch.cat([begin, x, end], dim=1)
-
-        if self.use_rms_norm:
-            return self.after_rms(x)
-        else:
-            return x
-
-
-class NormalizedDwPooler(nn.Module):
-    def __init__(self, dim):
-        super().__init__()
-        self.dim = dim
-        self.predictor = nn.Sequential(
-            nn.Linear(dim*2, dim),
-            nn.GELU(),
-            nn.Linear(dim, dim),
-        )
-    
-    def forward(self, x, forward_type='2x'):
-        B, H, W, C = x.shape
-
-        if forward_type == '2x':
-            new_x = x.reshape(B, H//2, 2, W//2, 2, C).permute(0, 1, 3, 2, 4, 5).reshape(B, H//2, W//2, 4, C)
-            pooled_x = new_x.mean(-2, keepdim=True).expand(-1, -1, -1, 4, -1)
-            fused_x = torch.cat([new_x, pooled_x], dim=-1)
-        elif forward_type == '1x':
-            new_x = x.reshape(B, H, W, 1, C)
-            fused_x = torch.cat([new_x, new_x], dim=-1)
-        elif forward_type == '4x':
-            new_x = x.reshape(B, H//4, 4, W//4, 4, C).permute(0, 1, 3, 2, 4, 5).reshape(B, H//4, W//4, 16, C)
-            pooled_x = new_x.mean(-2, keepdim=True).expand(-1, -1, -1, 16, -1)
-            fused_x = torch.cat([new_x, pooled_x], dim=-1)
-        
-        score = self.predictor(fused_x)
-        normalized_score = F.softmax(score, dim=-2)
-        new_x = (new_x * normalized_score).sum(dim=-2)
-        return new_x
-
-
-class OryxMLPv2(nn.Module):
-    def __init__(self, in_channels, out_channels, twoview=False, use_pe=False):
-        super().__init__()
-        
-        self.proj1 = nn.Linear(in_channels, out_channels)
-        self.proj2 = nn.Linear(out_channels, out_channels)
-        self.act = nn.GELU()
-        self.pooler = NormalizedDwPooler(out_channels)
-        embed_std = 1 / math.sqrt(out_channels)
-
-        self.use_pe = use_pe
-        if not use_pe:
-            self.image_newline = nn.Parameter(
-                torch.randn(out_channels) * embed_std
-            )
-        self.image_begin = nn.Parameter(
-            torch.randn(out_channels) * embed_std
-        )
-        self.image_end = nn.Parameter(
-            torch.randn(out_channels) * embed_std
-        )
-        
-        if twoview:
-            self.image_sep = nn.Parameter(
-                torch.randn(out_channels) * embed_std
-            )
-
-    def forward(self, x, size=(16,16), x2=None, size2=(16, 16), is_video=False):
-        h, w = size
-        dtype = x.dtype
-        x = x.reshape(x.shape[0], h, w, -1)
-        # x = self.pooler(x, forward_type=REGIONAL_POOL)
-        # x = self.proj(x) #b,h,w, c
-        x = self.proj1(x)
-        x = self.pooler(x, forward_type='2x')
-        x = self.act(x)
-        x = self.proj2(x)
-
-
-        b, h, w, c = x.shape
-        if not self.use_pe:
-            x = torch.cat([
-                x,
-                self.image_newline.reshape(1, 1, 1, c).expand(b, h, 1, c).to(dtype)
-            ], dim=2)
-        else:
-            pe_h = torch.arange(h, dtype=torch.long, device=x.device).reshape(1, h, 1, 1).expand(b, h, w, 1).reshape(b, h*w, 1)
-            pe_w = torch.arange(w, dtype=torch.long, device=x.device).reshape(1, 1, w, 1).expand(b, h, w, 1).reshape(b, h*w, 1)
-            pe = torch.cat([pe_h, pe_w], dim=-1)
-
-        x = x.reshape(b, -1, c)
-
-        if x2 is not None:
-            h2, w2 = size2
-            x2 = x2.reshape(x2.shape[0], h2, w2, -1)
-            # x2 = self.pooler(x2, forward_type=REGIONAL_POOL)
-            ## x2 = self.proj(x2) #b,h,w, c
-            x2 = self.proj1(x2)
-            x2 = self.pooler(x2, forward_type='2x')
-            x2 = self.act(x2)
-            x2 = self.proj2(x2)
-
-            b2, h2, w2, c2 = x2.shape
-            if not self.use_pe:
-                x2 = torch.cat([
-                    x2,
-                    self.image_newline.reshape(1, 1, 1, c).expand(b, h2, 1, c).to(dtype)
-                ], dim=2)
-            x2 = x2.reshape(b, -1, c)
-            sep = self.image_sep.reshape(1, 1, -1).expand(b, 1, c2).to(dtype)
-            x = torch.cat([x, sep, x2], dim=1)
-        
-        begin = self.image_begin.reshape(1, 1, -1).expand(b, 1, c).to(dtype)
-        end = self.image_end.reshape(1, 1, -1).expand(b, 1, c).to(dtype)
-        x = torch.cat([begin, x, end], dim=1)
-        # print(x.shape, x2.shape, h, w, h2, w2)
-        # print("vit rank = " + str(torch.distributed.get_rank()) +" x = " + str(x))
-        if self.use_pe:
-            zero_pad = torch.zeros(b, 1, 2, device=x.device, dtype=torch.long)
-            pe = torch.cat([zero_pad, pe, zero_pad], dim=1)
-            assert pe.shape[1] == x.shape[1]
-            return x, pe
-        else:
-            nseq = x.shape[1]
-            fake_pe = torch.zeros(b, nseq, 2, device=x.device, dtype=torch.long)
-            return x #, fake_pe
-