Update Notice.txt

Notice.txt

@@ -1,1083 +1,160 @@
-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
+Usage and Legal Notices:
 
+Tencent is pleased to support the open source community by making Tencent Hunyuan A13B available.
+
+Copyright (C) Tencent.  All rights reserved. The below software and/or models in this distribution may have been modified by Tencent ("Tencent Modifications"). All Tencent Modifications are Copyright (C) Tencent.
+
+Tencent Hunyuan A13B is licensed under TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT, which can be found in this repository called "LICENSE", except for the third-party components listed below. Tencent Hunyuan A13B does not impose any additional limitations beyond what is outlined in the respective licenses of these third-party components. Users must comply with all terms and conditions of original licenses of these third-party components and must ensure that the usage of the third party components adheres to all relevant laws and regulations. 
+
+For avoidance of doubts, Tencent Hunyuan A13B refers to the inference code, training code, parameters and the weights of Tencent Hunyuan A13B only, which are made publicly available by Tencent in accordance with the TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT.
+
+
+Other dependencies and licenses:
+
+
+Open Source Software Licensed under the Apache License Version 2.0:
+The below software in this distribution may have been modified by Tencent ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2025 Tencent.
+--------------------------------------------------------------------
+1. pytorch
+Copyright 2016-2017 TorchAPI
+Copyright 2016-2017 Contributors
+
+2. VLLM
+Copyright (c) vllm original author and authors
+Please note this software has been modified by Tencent in this distribution.
+
+3. transformers
+Copyright 2018- The Hugging Face team. All rights reserved.
+
+4. accelerate
+Copyright (c) accelerate original author and authors
+
+
+Terms of the Apache License Version 2.0:
+--------------------------------------------------------------------
+Apache License 
+
+Version 2.0, January 2004
+
+http://www.apache.org/licenses/ 
+
+TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+1. Definitions.
+
+"License" shall mean the terms and conditions for use, reproduction, and distribution as defined by Sections 1 through 9 of this document.
+
+"Licensor" shall mean the copyright owner or entity authorized by the copyright owner that is granting the License.
+
+"Legal Entity" shall mean the union of the acting entity and all other entities that control, are controlled by, or are under common control with that entity. For the purposes of this definition, "control" means (i) the power, direct or indirect, to cause the direction or management of such entity, whether by contract or otherwise, or (ii) ownership of fifty percent (50%) or more of the outstanding shares, or (iii) beneficial ownership of such entity.
+
+"You" (or "Your") shall mean an individual or Legal Entity exercising permissions granted by this License.
+
+"Source" form shall mean the preferred form for making modifications, including but not limited to software source code, documentation source, and configuration files.
+
+"Object" form shall mean any form resulting from mechanical transformation or translation of a Source form, including but not limited to compiled object code, generated documentation, and conversions to other media types.
+
+"Work" shall mean the work of authorship, whether in Source or Object form, made available under the License, as indicated by a copyright notice that is included in or attached to the work (an example is provided in the Appendix below).
+
+"Derivative Works" shall mean any work, whether in Source or Object form, that is based on (or derived from) the Work and for which the editorial revisions, annotations, elaborations, or other modifications represent, as a whole, an original work of authorship. For the purposes of this License, Derivative Works shall not include works that remain separable from, or merely link (or bind by name) to the interfaces of, the Work and Derivative Works thereof.
+
+"Contribution" shall mean any work of authorship, including the original version of the Work and any modifications or additions to that Work or Derivative Works thereof, that is intentionally submitted to Licensor for inclusion in the Work by the copyright owner or by an individual or Legal Entity authorized to submit on behalf of the copyright owner. For the purposes of this definition, "submitted" means any form of electronic, verbal, or written communication sent to the Licensor or its representatives, including but not limited to communication on electronic mailing lists, source code control systems, and issue tracking systems that are managed by, or on behalf of, the Licensor for the purpose of discussing and improving the Work, but excluding communication that is conspicuously marked or otherwise designated in writing by the copyright owner as "Not a Contribution."
+
+"Contributor" shall mean Licensor and any individual or Legal Entity on behalf of whom a Contribution has been received by Licensor and subsequently incorporated within the Work.
+
+2. Grant of Copyright License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable copyright license to reproduce, prepare Derivative Works of, publicly display, publicly perform, sublicense, and distribute the Work and such Derivative Works in Source or Object form.
+
+3. Grant of Patent License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable (except as stated in this section) patent license to make, have made, use, offer to sell, sell, import, and otherwise transfer the Work, where such license applies only to those patent claims licensable by such Contributor that are necessarily infringed by their Contribution(s) alone or by combination of their Contribution(s) with the Work to which such Contribution(s) was submitted. If You institute patent litigation against any entity (including a cross-claim or counterclaim in a lawsuit) alleging that the Work or a Contribution incorporated within the Work constitutes direct or contributory patent infringement, then any patent licenses granted to You under this License for that Work shall terminate as of the date such litigation is filed.
+
+4. Redistribution. You may reproduce and distribute copies of the Work or Derivative Works thereof in any medium, with or without modifications, and in Source or Object form, provided that You meet the following conditions:
+
+You must give any other recipients of the Work or Derivative Works a copy of this License; and 
+
+You must cause any modified files to carry prominent notices stating that You changed the files; and 
+
+You must retain, in the Source form of any Derivative Works that You distribute, all copyright, patent, trademark, and attribution notices from the Source form of the Work, excluding those notices that do not pertain to any part of the Derivative Works; and 
+
+If the Work includes a "NOTICE" text file as part of its distribution, then any Derivative Works that You distribute must include a readable copy of the attribution notices contained within such NOTICE file, excluding those notices that do not pertain to any part of the Derivative Works, in at least one of the following places: within a NOTICE text file distributed as part of the Derivative Works; within the Source form or documentation, if provided along with the Derivative Works; or, within a display generated by the Derivative Works, if and wherever such third-party notices normally appear. The contents of the NOTICE file are for informational purposes only and do not modify the License. You may add Your own attribution notices within Derivative Works that You distribute, alongside or as an addendum to the NOTICE text from the Work, provided that such additional attribution notices cannot be construed as modifying the License. 
+
+You may add Your own copyright statement to Your modifications and may provide additional or different license terms and conditions for use, reproduction, or distribution of Your modifications, or for any such Derivative Works as a whole, provided Your use, reproduction, and distribution of the Work otherwise complies with the conditions stated in this License. 
+
+5. Submission of Contributions. Unless You explicitly state otherwise, any Contribution intentionally submitted for inclusion in the Work by You to the Licensor shall be under the terms and conditions of this License, without any additional terms or conditions. Notwithstanding the above, nothing herein shall supersede or modify the terms of any separate license agreement you may have executed with Licensor regarding such Contributions.
+
+6. Trademarks. This License does not grant permission to use the trade names, trademarks, service marks, or product names of the Licensor, except as required for reasonable and customary use in describing the origin of the Work and reproducing the content of the NOTICE file.
+
+7. Disclaimer of Warranty. Unless required by applicable law or agreed to in writing, Licensor provides the Work (and each Contributor provides its Contributions) on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied, including, without limitation, any warranties or conditions of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A PARTICULAR PURPOSE. You are solely responsible for determining the appropriateness of using or redistributing the Work and assume any risks associated with Your exercise of permissions under this License.
+
+8. Limitation of Liability. In no event and under no legal theory, whether in tort (including negligence), contract, or otherwise, unless required by applicable law (such as deliberate and grossly negligent acts) or agreed to in writing, shall any Contributor be liable to You for damages, including any direct, indirect, special, incidental, or consequential damages of any character arising as a result of this License or out of the use or inability to use the Work (including but not limited to damages for loss of goodwill, work stoppage, computer failure or malfunction, or any and all other commercial damages or losses), even if such Contributor has been advised of the possibility of such damages.
+
+9. Accepting Warranty or Additional Liability. While redistributing the Work or Derivative Works thereof, You may choose to offer, and charge a fee for, acceptance of support, warranty, indemnity, or other liability obligations and/or rights consistent with this License. However, in accepting such obligations, You may act only on Your own behalf and on Your sole responsibility, not on behalf of any other Contributor, and only if You agree to indemnify, defend, and hold each Contributor harmless for any liability incurred by, or claims asserted against, such Contributor by reason of your accepting any such warranty or additional liability.
+
+END OF TERMS AND CONDITIONS
+
+
+
+Open Source Software Licensed under the BSD 3-Clause License and Other Licenses of the Third-Party Components therein:
+--------------------------------------------------------------------
+1. pytorch
+Copyright (c) 2016-     Facebook, Inc            (Adam Paszke)
+Copyright (c) 2014-     Facebook, Inc            (Soumith Chintala)
+Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)
+Copyright (c) 2012-2014 Deepmind Technologies    (Koray Kavukcuoglu)
+Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)
+Copyright (c) 2011-2013 NYU                      (Clement Farabet)
+Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)
+Copyright (c) 2006      Idiap Research Institute (Samy Bengio)
+Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)
+
+
+Terms of the BSD 3-Clause:
+--------------------------------------------------------------------
+Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+For the license of other third party components, please refer to the following URL:
+https://github.com/pytorch/pytorch/blob/v2.1.1/NOTICE
+https://github.com/pytorch/pytorch/tree/v2.1.1/third_party
+
+
+Open Source Software Licensed under the BSD 3-Clause License:
+--------------------------------------------------------------------
+1. flash_attn
+Copyright (c) 2022, the respective contributors, as shown by the AUTHORS file.
+All rights reserved.
+
+
+A copy of the BSD 3-Clause is included in this file.
+
+
+
+Open Source Software Licensed under the Apache License Version 2.0 and Other Licenses of the Third-Party Components therein:
+The below software in this distribution is modified by Tencent ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2025 Tencent.
+--------------------------------------------------------------------
+1. sglang
+Copyright 2023-2024 SGLang Team
+
+
+A copy of the Apache 2.0 is included in this file.
+
+For the license of other third party components, please refer to the following URL:
+https://github.com/sgl-project/sglang/tree/v0.4.7/3rdparty/amd
+
+
+
+Open Source Software Licensed under the Apache License Version 2.0 and Other Licenses of the Third-Party Components therein:
+The below software in this distribution is modified by Tencent ("Tencent Modifications"). All Tencent Modifications are Copyright (C) 2025 Tencent.
+--------------------------------------------------------------------
+1. TensorRT-LLM
+Copyright (c) TensorRT-LLM original author and authors
+
+
+A copy of the Apache 2.0 is included in this file.
+
+For the license of other third party components, please refer to the following URL:
+https://github.com/NVIDIA/TensorRT-LLM/tree/v0.20.0/3rdparty