Upload processor

modeling_yangjian.py

@@ -262,14 +262,14 @@ class YangJianCompareVisualEncoder(nn.Module):
         super().__init__()
         self.config = config
         self.hidden_size = config.hidden_size
-        self.token_size = 100  * (config.spatial_merge_size**2) if "compare_token_size" not in config else config.compare_token_size  * (config.spatial_merge_size**2)
-        
+        # self.token_size = 100  * (config.spatial_merge_size**2) if "compare_token_size" not in config else config.compare_token_size  * (config.spatial_merge_size**2)
+        self.token_size = 100 if "compare_token_size" not in config else config.compare_token_size
         # Encoder 部分：双向图像特征交互
         # 第一个cross attention: previous attend to current
         self.encoder_cross_attn1 = OptimizedCrossAttention(config, is_cross_attention=True)
         # 第二个cross attention: current attend to previous
         self.encoder_cross_attn2 = OptimizedCrossAttention(config, is_cross_attention=True)
-        
+
         self.encoder_norm1 = Qwen2RMSNorm(self.hidden_size, eps=1e-6)
         self.encoder_norm2 = Qwen2RMSNorm(self.hidden_size, eps=1e-6)
         self.encoder_norm3 = Qwen2RMSNorm(self.hidden_size, eps=1e-6)
@@ -290,6 +290,8 @@ class YangJianCompareVisualEncoder(nn.Module):
         self.decoder_norm2 = Qwen2RMSNorm(self.hidden_size, eps=1e-6)
         self.decoder_mlp = Qwen2_5_VLMLP(config)
 
+        self.compare_projector = nn.Linear(config.hidden_size, config.out_hidden_size)
+
     def _ensure_device_dtype_consistency(self, target_tensor):
         """
         确保所有模块组件都在目标张量的设备上并使用相同的数据类型
@@ -391,8 +393,19 @@ class YangJianCompareVisualEncoder(nn.Module):
             torch.ones(batch_size, self.token_size, dtype=torch.bool, device=device),  # query掩码
             attention_masks  # encoded特征的掩码
         )
+
+        # 记录每个batch的token数量
+        batch_size = compare_visual_embeds.size(0)
+        token_size = compare_visual_embeds.size(1)
+        # 将所有batch的数据拼接在一起
+        # [batch_size * token_size, hidden_size]
+        flattened_embeds = compare_visual_embeds.view(-1, compare_visual_embeds.size(-1))
+        merged = self.compare_projector(flattened_embeds)  # [batch_size * token_size, merged_hidden_size]
+        merged_token_size = token_size
+        # [batch_size, merged_token_size, merged_hidden_size]
+        compare_visual_embeds = merged.view(batch_size, merged_token_size, -1)
         
-        return compare_visual_embeds  # [batch_size, token_size, hidden_size]
+        return compare_visual_embeds  # [batch_size, token_size, out_hidden_size]
     
     def _encoder_forward(self, current_features, previous_features, current_mask=None, previous_mask=None):
         """
@@ -447,8 +460,9 @@ class YangJianCompareVisualEncoder(nn.Module):
         residual = current_features
         mlp_input2 = self.encoder_norm4(current_features)
         mlp_output2 = self.encoder_mlp2(mlp_input2)
-        current_features = residual + mlp_output2
-        
+        # current_features = residual + mlp_output2
+        # 修改为减法
+        current_features = residual - mlp_output2
         return current_features
     
     def _decoder_forward(self, queries, encoded_features, query_mask=None, encoded_mask=None):
@@ -548,18 +562,7 @@ class YangJianVisionTransformerPretrainedModel(Qwen2_5_VisionTransformerPretrain
         splited_hidden_states_before_merger = torch.split(hidden_states, split_sizes)
         # [total_images, token_size, hidden_size]
         compare_visual_embeds = self.compare_visual_encoder(splited_hidden_states_before_merger)
-        # 记录每个batch的token数量
-        batch_size = compare_visual_embeds.size(0)
-        token_size = compare_visual_embeds.size(1)
-        # 将所有batch的数据拼接在一起
-        # [batch_size * token_size, hidden_size]
-        flattened_embeds = compare_visual_embeds.view(-1, compare_visual_embeds.size(-1))
-        # 一次性进行merger操作
-        # 假设merger会将token数量变为原来的1/4
-        merged = self.merger(flattened_embeds)  # [(batch_size * token_size)/4, merged_hidden_size]
-        merged_token_size = token_size // self.spatial_merge_size**2  
-        # [batch_size, merged_token_size, merged_hidden_size]
-        compare_visual_embeds = merged.view(batch_size, merged_token_size, -1)
+
         
         hidden_states = self.merger(hidden_states)
         reverse_indices = torch.argsort(window_index)
@@ -853,8 +856,10 @@ class YangJianVLModel(Qwen2_5_VLModel):
                     if ed_image < ed_video:
                         # 如果当前是图片,则需要插入 compare_token_size 个图像对比的token的position
                         compare_t_index = t_index[-1].repeat(self.compare_token_size)
-                        compare_h_index = torch.arange(self.compare_token_size)
-                        compare_w_index = torch.arange(self.compare_token_size)
+                        # compare_h_index = torch.arange(self.compare_token_size)
+                        # compare_w_index = torch.arange(self.compare_token_size)
+                        compare_h_index = compare_t_index
+                        compare_w_index = compare_t_index
                         llm_pos_ids_list.append(torch.stack([compare_t_index, compare_h_index, compare_w_index]) + text_len + st_idx)
                         st = st + self.compare_token_size