vit_model.py

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