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