Update LICENSE.txt

LICENSE.txt

@@ -1,1083 +1,77 @@
-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
-
+TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT
+Tencent Hunyuan A13B Release Date: June 27, 2025
+THIS LICENSE AGREEMENT DOES NOT APPLY IN THE EUROPEAN UNION, UNITED KINGDOM AND SOUTH KOREA AND IS EXPRESSLY LIMITED TO THE TERRITORY, AS DEFINED BELOW.
+By clicking to agree or by using, reproducing, modifying, distributing, performing or displaying any portion or element of the Tencent Hunyuan Works, including via any Hosted Service, You will be deemed to have recognized and accepted the content of this Agreement, which is effective immediately.
+1.	DEFINITIONS.
+a.	“Acceptable Use Policy” shall mean the policy made available by Tencent as set forth in the Exhibit A.
+b.	“Agreement” shall mean the terms and conditions for use, reproduction, distribution, modification, performance and displaying of Tencent Hunyuan Works or any portion or element thereof set forth herein.
+c.	“Documentation” shall mean the specifications, manuals and documentation for Tencent Hunyuan made publicly available by Tencent.
+d.	“Hosted Service” shall mean a hosted service offered via an application programming interface (API), web access, or any other electronic or remote means.
+e.	“Licensee,” “You” or “Your” shall mean a natural person or legal entity exercising the rights granted by this Agreement and/or using the Tencent Hunyuan Works for any purpose and in any field of use.
+f.	“Materials” shall mean, collectively, Tencent’s proprietary Tencent Hunyuan and Documentation (and any portion thereof) as made available by Tencent under this Agreement.
+g.	“Model Derivatives” shall mean all: (i) modifications to Tencent Hunyuan or any Model Derivative of Tencent Hunyuan; (ii) works based on Tencent Hunyuan or any Model Derivative of Tencent Hunyuan; or (iii) any other machine learning model which is created by transfer of patterns of the weights, parameters, operations, or Output of Tencent Hunyuan or any Model Derivative of Tencent Hunyuan, to that model in order to cause that model to perform similarly to Tencent Hunyuan or a Model Derivative of Tencent Hunyuan, including distillation methods, methods that use intermediate data representations, or methods based on the generation of synthetic data Outputs by Tencent Hunyuan or a Model Derivative of Tencent Hunyuan for training that model. For clarity, Outputs by themselves are not deemed Model Derivatives.
+h.	“Output” shall mean the information and/or content output of Tencent Hunyuan or a Model Derivative that results from operating or otherwise using Tencent Hunyuan or a Model Derivative, including via a Hosted Service.
+i.	“Tencent,” “We” or “Us” shall mean the applicable entity or entities in the Tencent corporate family that own(s) intellectual property or other rights embodied in or utilized by the Materials.
+j.	“Tencent Hunyuan” shall mean the large language models, text/image/video/audio/3D generation models, and multimodal large language models and their software and algorithms, including trained model weights, parameters (including optimizer states), machine-learning model code, inference-enabling code, training-enabling code, fine-tuning enabling code and other elements of the foregoing made publicly available by Us, including, without limitation to, Tencent Hunyuan A13B released at [https://github.com/Tencent-Hunyuan/Hunyuan-A13B].
+k.	“Tencent Hunyuan Works” shall mean: (i) the Materials; (ii) Model Derivatives; and (iii) all derivative works thereof.
+l.	“Territory” shall mean the worldwide territory, excluding the territory of the European Union, United Kingdom and South Korea. 
+m.	“Third Party” or “Third Parties” shall mean individuals or legal entities that are not under common control with Us or You.
+n.	“including” shall mean including but not limited to.
+2.	GRANT OF RIGHTS.
+We grant You, for the Territory only, a non-exclusive, non-transferable and royalty-free limited license under Tencent’s intellectual property or other rights owned by Us embodied in or utilized by the Materials to use, reproduce, distribute, create derivative works of (including Model Derivatives), and make modifications to the Materials, only in accordance with the terms of this Agreement and the Acceptable Use Policy, and You must not violate (or encourage or permit anyone else to violate) any term of this Agreement or the Acceptable Use Policy.
+3.	DISTRIBUTION.
+You may, subject to Your compliance with this Agreement, distribute or make available to Third Parties the Tencent Hunyuan Works, exclusively in the Territory, provided that You meet all of the following conditions:
+a.	You must provide all such Third Party recipients of the Tencent Hunyuan Works or products or services using them a copy of this Agreement;
+b.	You must cause any modified files to carry prominent notices stating that You changed the files;
+c.	You are encouraged to: (i) publish at least one technology introduction blogpost or one public statement expressing Your experience of using the Tencent Hunyuan Works; and (ii) mark the products or services developed by using the Tencent Hunyuan Works to indicate that the product/service is “Powered by Tencent Hunyuan”; and
+d.	All distributions to Third Parties (other than through a Hosted Service) must be accompanied by a “Notice” text file that contains the following notice: “Tencent Hunyuan is licensed under the Tencent Hunyuan Community License Agreement, Copyright © 2025 Tencent. All Rights Reserved. The trademark rights of “Tencent Hunyuan” are owned by Tencent or its affiliate.”
+You may add Your own copyright statement to Your modifications and, except as set forth in this Section and in Section 5, may provide additional or different license terms and conditions for use, reproduction, or distribution of Your modifications, or for any such Model Derivatives as a whole, provided Your use, reproduction, modification, distribution, performance and display of the work otherwise complies with the terms and conditions of this Agreement (including as regards the Territory). If You receive Tencent Hunyuan Works from a Licensee as part of an integrated end user product, then this Section 3 of this Agreement will not apply to You.
+4.	ADDITIONAL COMMERCIAL TERMS.
+If, on the Tencent Hunyuan version release date, the monthly active users of all products or services made available by or for Licensee is greater than 100 million monthly active users in the preceding calendar month, You must request a license from Tencent, which Tencent may grant to You in its sole discretion, and You are not authorized to exercise any of the rights under this Agreement unless or until Tencent otherwise expressly grants You such rights.
+5.	RULES OF USE.
+a.	Your use of the Tencent Hunyuan Works must comply with applicable laws and regulations (including trade compliance laws and regulations) and adhere to the Acceptable Use Policy for the Tencent Hunyuan Works, which is hereby incorporated by reference into this Agreement. You must include the use restrictions referenced in these Sections 5(a) and 5(b) as an enforceable provision in any agreement (e.g., license agreement, terms of use, etc.) governing the use and/or distribution of Tencent Hunyuan Works and You must provide notice to subsequent users to whom You distribute that Tencent Hunyuan Works are subject to the use restrictions in these Sections 5(a) and 5(b).
+b.	You must not use the Tencent Hunyuan Works or any Output or results of the Tencent Hunyuan Works to improve any other AI model (other than Tencent Hunyuan or Model Derivatives thereof).
+c.	You must not use, reproduce, modify, distribute, or display the Tencent Hunyuan Works, Output or results of the Tencent Hunyuan Works outside the Territory. Any such use outside the Territory is unlicensed and unauthorized under this Agreement.
+6.	INTELLECTUAL PROPERTY.
+a.	Subject to Tencent’s ownership of Tencent Hunyuan Works made by or for Tencent and intellectual property rights therein, conditioned upon Your compliance with the terms and conditions of this Agreement, as between You and Tencent, You will be the owner of any derivative works and modifications of the Materials and any Model Derivatives that are made by or for You.
+b.	No trademark licenses are granted under this Agreement, and in connection with the Tencent Hunyuan Works, Licensee may not use any name or mark owned by or associated with Tencent or any of its affiliates, except as required for reasonable and customary use in describing and distributing the Tencent Hunyuan Works. Tencent hereby grants You a license to use “Tencent Hunyuan” (the “Mark”) in the Territory solely as required to comply with the provisions of Section 3(c), provided that You comply with any applicable laws related to trademark protection. All goodwill arising out of Your use of the Mark will inure to the benefit of Tencent.
+c.	If You commence a lawsuit or other proceedings (including a cross-claim or counterclaim in a lawsuit) against Us or any person or entity alleging that the Materials or any Output, or any portion of any of the foregoing, infringe any intellectual property or other right owned or licensable by You, then all licenses granted to You under this Agreement shall terminate as of the date such lawsuit or other proceeding is filed. You will defend, indemnify and hold harmless Us from and against any claim by any Third Party arising out of or related to Your or the Third Party’s use or distribution of the Tencent Hunyuan Works.
+d.	Tencent claims no rights in Outputs You generate. You and Your users are solely responsible for Outputs and their subsequent uses.
+7.	DISCLAIMERS OF WARRANTY AND LIMITATIONS OF LIABILITY.
+a.	We are not obligated to support, update, provide training for, or develop any further version of the Tencent Hunyuan Works or to grant any license thereto.
+b.	UNLESS AND ONLY TO THE EXTENT REQUIRED BY APPLICABLE LAW, THE TENCENT HUNYUAN WORKS AND ANY OUTPUT AND RESULTS THEREFROM ARE PROVIDED “AS IS” WITHOUT ANY EXPRESS OR IMPLIED WARRANTIES OF ANY KIND INCLUDING ANY WARRANTIES OF TITLE, MERCHANTABILITY, NONINFRINGEMENT, COURSE OF DEALING, USAGE OF TRADE, OR FITNESS FOR A PARTICULAR PURPOSE. YOU ARE SOLELY RESPONSIBLE FOR DETERMINING THE APPROPRIATENESS OF USING, REPRODUCING, MODIFYING, PERFORMING, DISPLAYING OR DISTRIBUTING ANY OF THE TENCENT HUNYUAN WORKS OR OUTPUTS AND ASSUME ANY AND ALL RISKS ASSOCIATED WITH YOUR OR A THIRD PARTY’S USE OR DISTRIBUTION OF ANY OF THE TENCENT HUNYUAN WORKS OR OUTPUTS AND YOUR EXERCISE OF RIGHTS AND PERMISSIONS UNDER THIS AGREEMENT.
+c.	TO THE FULLEST EXTENT PERMITTED BY APPLICABLE LAW, IN NO EVENT SHALL TENCENT OR ITS AFFILIATES BE LIABLE UNDER ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, TORT, NEGLIGENCE, PRODUCTS LIABILITY, OR OTHERWISE, FOR ANY DAMAGES, INCLUDING ANY DIRECT, INDIRECT, SPECIAL, INCIDENTAL, EXEMPLARY, CONSEQUENTIAL OR PUNITIVE DAMAGES, OR LOST PROFITS OF ANY KIND ARISING FROM THIS AGREEMENT OR RELATED TO ANY OF THE TENCENT HUNYUAN WORKS OR OUTPUTS, EVEN IF TENCENT OR ITS AFFILIATES HAVE BEEN ADVISED OF THE POSSIBILITY OF ANY OF THE FOREGOING.
+8.	SURVIVAL AND TERMINATION.
+a.	The term of this Agreement shall commence upon Your acceptance of this Agreement or access to the Materials and will continue in full force and effect until terminated in accordance with the terms and conditions herein.
+b.	We may terminate this Agreement if You breach any of the terms or conditions of this Agreement. Upon termination of this Agreement, You must promptly delete and cease use of the Tencent Hunyuan Works. Sections 6(a), 6(c), 7 and 9 shall survive the termination of this Agreement.
+9.	GOVERNING LAW AND JURISDICTION.
+a.	This Agreement and any dispute arising out of or relating to it will be governed by the laws of the Hong Kong Special Administrative Region of the People’s Republic of China, without regard to conflict of law principles, and the UN Convention on Contracts for the International Sale of Goods does not apply to this Agreement.
+b.	Exclusive jurisdiction and venue for any dispute arising out of or relating to this Agreement will be a court of competent jurisdiction in the Hong Kong Special Administrative Region of the People’s Republic of China, and Tencent and Licensee consent to the exclusive jurisdiction of such court with respect to any such dispute.
+ 
+EXHIBIT A
+ACCEPTABLE USE POLICY
+
+Tencent reserves the right to update this Acceptable Use Policy from time to time.
+Last modified: November 5, 2024
+
+Tencent endeavors to promote safe and fair use of its tools and features, including Tencent Hunyuan. You agree not to use Tencent Hunyuan or Model Derivatives:
+1.	Outside the Territory;
+2.	In any way that violates any applicable national, federal, state, local, international or any other law or regulation;
+3.	To harm Yourself or others;
+4.	To repurpose or distribute output from Tencent Hunyuan or any Model Derivatives to harm Yourself or others; 
+5.	To override or circumvent the safety guardrails and safeguards We have put in place;
+6.	For the purpose of exploiting, harming or attempting to exploit or harm minors in any way;
+7.	To generate or disseminate verifiably false information and/or content with the purpose of harming others or influencing elections;
+8.	To generate or facilitate false online engagement, including fake reviews and other means of fake online engagement;
+9.	To intentionally defame, disparage or otherwise harass others;
+10.	To generate and/or disseminate malware (including ransomware) or any other content to be used for the purpose of harming electronic systems;
+11.	To generate or disseminate personal identifiable information with the purpose of harming others;
+12.	To generate or disseminate information (including images, code, posts, articles), and place the information in any public context (including –through the use of bot generated tweets), without expressly and conspicuously identifying that the information and/or content is machine generated;
+13.	To impersonate another individual without consent, authorization, or legal right;
+14.	To make high-stakes automated decisions in domains that affect an individual’s safety, rights or wellbeing (e.g., law enforcement, migration, medicine/health, management of critical infrastructure, safety components of products, essential services, credit, employment, housing, education, social scoring, or insurance);
+15.	In a manner that violates or disrespects the social ethics and moral standards of other countries or regions;
+16.	To perform, facilitate, threaten, incite, plan, promote or encourage violent extremism or terrorism;
+17.	For any use intended to discriminate against or harm individuals or groups based on protected characteristics or categories, online or offline social behavior or known or predicted personal or personality characteristics;
+18.	To intentionally exploit any of the vulnerabilities of a specific group of persons based on their age, social, physical or mental characteristics, in order to materially distort the behavior of a person pertaining to that group in a manner that causes or is likely to cause that person or another person physical or psychological harm;
+19.	For military purposes;
+20.	To engage in the unauthorized or unlicensed practice of any profession including, but not limited to, financial, legal, medical/health, or other professional practices.