Update mllm_encoder/mllm_encoder.py

mllm_encoder/mllm_encoder.py

@@ -2328,9 +2328,9 @@ class MLLMEncoder(ModelMixin, ConfigMixin):
         self.video_queries = nn.Parameter(
             torch.randn(num_video_queries, hidden_size) * 0.02
         )
-        self.ref_queries = nn.Parameter(
-            torch.randn(num_ref_queries, hidden_size) * 0.02
-        )
+        # self.ref_queries = nn.Parameter(
+        #     torch.randn(num_ref_queries, hidden_size) * 0.02
+        # )
 
         # Connector MLP: MLLM hidden → DiT dim
         self.connector = nn.Sequential(
@@ -2342,13 +2342,13 @@ class MLLMEncoder(ModelMixin, ConfigMixin):
         nn.init.zeros_(self.connector[2].bias)
 
         # Ref connector MLP (separate from main connector)
-        self.ref_connector = nn.Sequential(
-            nn.Linear(hidden_size, dit_dim),
-            nn.GELU(approximate="tanh"),
-            nn.Linear(dit_dim, dit_dim),
-        )
-        nn.init.zeros_(self.ref_connector[2].weight)
-        nn.init.zeros_(self.ref_connector[2].bias)
+        # self.ref_connector = nn.Sequential(
+        #     nn.Linear(hidden_size, dit_dim),
+        #     nn.GELU(approximate="tanh"),
+        #     nn.Linear(dit_dim, dit_dim),
+        # )
+        # nn.init.zeros_(self.ref_connector[2].weight)
+        # nn.init.zeros_(self.ref_connector[2].bias)
 
         # Qwen VL model and processor (loaded lazily)
         self.qwen_model = None
@@ -2705,31 +2705,31 @@ class MLLMEncoder(ModelMixin, ConfigMixin):
         learnable_query_features = self.connector(learnable_query_features)
 
         # Extract ref image features if in ref mode
-        if ref_image:
-            vision_start_id = self.processor.tokenizer.convert_tokens_to_ids(
-                "<|vision_start|>"
-            )
-            vision_end_id = self.processor.tokenizer.convert_tokens_to_ids(
-                "<|vision_end|>"
-            )
-            input_ids = inputs.input_ids[0]
-            vision_start_indices = (input_ids == vision_start_id).nonzero(
-                as_tuple=True
-            )[-1]
-            if len(vision_start_indices) > 0:
-                last_vision_start = vision_start_indices[-1]
-                remaining_ids = input_ids[last_vision_start:]
-                end_relative_idx = (remaining_ids == vision_end_id).nonzero(
-                    as_tuple=True
-                )[-1]
-                if len(end_relative_idx) > 0:
-                    last_vision_end = last_vision_start + end_relative_idx[0]
-                    ref_image_features = hidden_states[
-                        :, last_vision_start + 1 : last_vision_end, :
-                    ]
-                    ref_image_features = self.ref_connector(ref_image_features)
-                    learnable_query_features = torch.cat(
-                        [ref_image_features, learnable_query_features], dim=1
-                    )
+        # if ref_image:
+        #     vision_start_id = self.processor.tokenizer.convert_tokens_to_ids(
+        #         "<|vision_start|>"
+        #     )
+        #     vision_end_id = self.processor.tokenizer.convert_tokens_to_ids(
+        #         "<|vision_end|>"
+        #     )
+        #     input_ids = inputs.input_ids[0]
+        #     vision_start_indices = (input_ids == vision_start_id).nonzero(
+        #         as_tuple=True
+        #     )[-1]
+        #     if len(vision_start_indices) > 0:
+        #         last_vision_start = vision_start_indices[-1]
+        #         remaining_ids = input_ids[last_vision_start:]
+        #         end_relative_idx = (remaining_ids == vision_end_id).nonzero(
+        #             as_tuple=True
+        #         )[-1]
+        #         if len(end_relative_idx) > 0:
+        #             last_vision_end = last_vision_start + end_relative_idx[0]
+        #             ref_image_features = hidden_states[
+        #                 :, last_vision_start + 1 : last_vision_end, :
+        #             ]
+        #             ref_image_features = self.ref_connector(ref_image_features)
+        #             learnable_query_features = torch.cat(
+        #                 [ref_image_features, learnable_query_features], dim=1
+        #             )
 
         return learnable_query_features