fix: sync all compat fixes from AD-Copilot (imports, get_rope_index kwargs, remove sdpa hardcode)

modeling_ad_copilot.py

@@ -0,0 +1,758 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np  
+from typing import Any, Callable, Optional, Union
+
+from transformers import Qwen2_5_VLForConditionalGeneration, AutoModelForImageTextToText
+from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import (
+    Qwen2_5_VisionTransformerPretrainedModel,
+    Qwen2_5_VLModel,
+    Qwen2_5_VLMLP,
+    ALL_ATTENTION_FUNCTIONS,
+)
+try:
+    from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import Qwen2RMSNorm
+except ImportError:
+    from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm
+try:
+    from transformers.image_utils import ImageInput
+except ImportError:
+    ImageInput = Any
+try:
+    from transformers.tokenization_utils import TextInput, PreTokenizedInput
+except ImportError:
+    TextInput = str
+    PreTokenizedInput = list
+try:
+    from transformers.video_utils import VideoInput
+except ImportError:
+    VideoInput = Any
+try:
+    from transformers.feature_extraction_utils import BatchFeature
+except ImportError:
+    from transformers import BatchFeature
+
+from transformers import Qwen2_5_VLProcessor, Qwen2_5_VLConfig
+try:
+    from transformers.models.qwen2_5_vl.processing_qwen2_5_vl import Qwen2_5_VLProcessorKwargs
+except ImportError:
+    Qwen2_5_VLProcessorKwargs = dict
+
+class ADCopilotConfig(Qwen2_5_VLConfig):
+    model_type = "ad_copilot"
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.vision_config.compare_token_size = 100
+        self.architectures = ["ADCopilotVLForConditionalGeneration"]
+        self.sequence_compare = True
+        
+class ADCopilotProcessor(Qwen2_5_VLProcessor):
+    config_class = ADCopilotConfig
+    def __init__(self, image_processor=None, tokenizer=None, video_processor=None, chat_template=None, **kwargs):
+        super().__init__(image_processor, tokenizer, video_processor, chat_template, **kwargs)
+        self.compare_token_size = 100 if "compare_token_size" not in kwargs else kwargs["compare_token_size"]
+
+    def __call__(
+        self,
+        images: ImageInput = None,
+        text: Union[TextInput, PreTokenizedInput, list[TextInput], list[PreTokenizedInput]] = None,
+        videos: VideoInput = None,
+        **kwargs,
+    ) -> BatchFeature:
+        """
+        Main method to prepare for the model one or several sequences(s) and image(s). This method forwards the `text`
+        and `kwargs` arguments to Qwen2TokenizerFast's [`~Qwen2TokenizerFast.__call__`] if `text` is not `None` to encode
+        the text. To prepare the vision inputs, this method forwards the `vision_infos` and `kwrags` arguments to
+        Qwen2VLImageProcessor's [`~Qwen2VLImageProcessor.__call__`] if `vision_infos` is not `None`.
+
+        Args:
+            images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `list[PIL.Image.Image]`, `list[np.ndarray]`, `list[torch.Tensor]`):
+                The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch
+                tensor. Both channels-first and channels-last formats are supported.
+            text (`str`, `list[str]`, `list[list[str]]`):
+                The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings
+                (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set
+                `is_split_into_words=True` (to lift the ambiguity with a batch of sequences).
+            videos (`np.ndarray`, `torch.Tensor`, `list[np.ndarray]`, `list[torch.Tensor]`):
+                The image or batch of videos to be prepared. Each video can be a 4D NumPy array or PyTorch
+                tensor, or a nested list of 3D frames. Both channels-first and channels-last formats are supported.
+            return_tensors (`str` or [`~utils.TensorType`], *optional*):
+                If set, will return tensors of a particular framework. Acceptable values are:
+                - `'tf'`: Return TensorFlow `tf.constant` objects.
+                - `'pt'`: Return PyTorch `torch.Tensor` objects.
+                - `'np'`: Return NumPy `np.ndarray` objects.
+                - `'jax'`: Return JAX `jnp.ndarray` objects.
+
+        Returns:
+            [`BatchFeature`]: A [`BatchFeature`] with the following fields:
+
+            - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`.
+            - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when
+              `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not
+              `None`).
+            - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`.
+            - **pixel_values_videos** -- Pixel values of videos to be fed to a model. Returned when `videos` is not `None`.
+            - **image_grid_thw** -- List of image 3D grid in LLM. Returned when `images` is not `None`.
+            - **video_grid_thw** -- List of video 3D grid in LLM. Returned when `videos` is not `None`.
+            - **second_per_grid_ts** -- List of video seconds per time grid. Returned when `videos` is not `None`.
+        """
+        output_kwargs = self._merge_kwargs(
+            Qwen2_5_VLProcessorKwargs,
+            tokenizer_init_kwargs=self.tokenizer.init_kwargs,
+            **kwargs,
+        )
+
+        image_inputs = videos_inputs = {}
+        if images is not None:
+            image_inputs = self.image_processor(images=images, **output_kwargs["images_kwargs"])
+            image_grid_thw = image_inputs["image_grid_thw"]
+
+        if videos is not None:
+            fps = output_kwargs["videos_kwargs"].get("fps", 2.0)
+            videos_inputs = self.video_processor(videos=videos, **output_kwargs["videos_kwargs"])
+            video_grid_thw = videos_inputs["video_grid_thw"]
+
+            if isinstance(fps, (int, float)):
+                second_per_grid_ts = [self.video_processor.temporal_patch_size / fps] * len(video_grid_thw)
+            elif hasattr(fps, "__len__") and len(fps) == len(video_grid_thw):
+                second_per_grid_ts = [self.video_processor.temporal_patch_size / tmp for tmp in fps]
+            else:
+                raise ValueError(
+                    f"The length of fps ({len(fps) if hasattr(fps, '__len__') else fps}) must be equal to the length of video_grid_thw ({len(video_grid_thw)}) or fps should be a single number."
+                )
+            videos_inputs.update({"second_per_grid_ts": second_per_grid_ts})
+
+        if not isinstance(text, list):
+            text = [text]
+
+        text = text.copy()  # below lines change text in-place
+        if images is not None:
+            merge_length = self.image_processor.merge_size**2
+            index = 0
+            for i in range(len(text)):
+                while self.image_token in text[i]:
+                    num_image_tokens = image_grid_thw[index].prod() // merge_length
+                    # text[i] = text[i].replace(self.image_token, "<|placeholder|>" * (num_image_tokens), 1)
+                    text[i] = text[i].replace(self.image_token, "<|placeholder|>" * (num_image_tokens + self.compare_token_size), 1)
+                    index += 1
+                text[i] = text[i].replace("<|placeholder|>", self.image_token)
+
+        if videos is not None:
+            merge_length = self.video_processor.merge_size**2
+            index = 0
+            for i in range(len(text)):
+                while self.video_token in text[i]:
+                    num_video_tokens = video_grid_thw[index].prod() // merge_length
+                    text[i] = text[i].replace(self.video_token, "<|placeholder|>" * num_video_tokens, 1)
+                    index += 1
+                text[i] = text[i].replace("<|placeholder|>", self.video_token)
+
+        return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None)
+        return_mm_token_type_ids = output_kwargs["text_kwargs"].pop("return_mm_token_type_ids", None)
+        text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"])
+        self._check_special_mm_tokens(text, text_inputs, modalities=["image", "video"])
+
+        if return_mm_token_type_ids:
+            array_ids = np.array(text_inputs["input_ids"])
+            mm_token_type_ids = np.zeros_like(text_inputs["input_ids"])
+            mm_token_type_ids[array_ids == self.image_token_id] = 1
+            text_inputs["mm_token_type_ids"] = mm_token_type_ids.tolist()
+
+        return BatchFeature(data={**text_inputs, **image_inputs, **videos_inputs}, tensor_type=return_tensors)
+
+
+class OptimizedCrossAttention(nn.Module):
+    """
+    仿照 Qwen2_5_VLVisionAttention 结构的优化 Cross Attention
+    """
+    def __init__(self, config, is_cross_attention=True):
+        super().__init__()
+        self.config = config
+        self.dim = config.hidden_size
+        self.num_heads = config.num_heads
+        self.head_dim = self.dim // self.num_heads
+        self.scaling = self.head_dim**-0.5
+        self.attention_dropout = 0.0
+        self.is_causal = False  # cross attention 不需要因果掩码
+        self.is_cross_attention = is_cross_attention
+        
+        if is_cross_attention:
+            # Cross attention: Q 来自一个序列，K、V 来自另一个序列
+            self.q_proj = nn.Linear(self.dim, self.dim, bias=True)
+            self.kv = nn.Linear(self.dim, self.dim * 2, bias=True)  # 融合 K、V
+        else:
+            # Self attention: Q、K、V 来自同一个序列
+            self.qkv = nn.Linear(self.dim, self.dim * 3, bias=True)  # 融合 Q、K、V
+        
+        self.proj = nn.Linear(self.dim, self.dim, bias=True)
+        
+    def forward(
+        self,
+        query_states: torch.Tensor,
+        key_value_states: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        cu_seqlens: Optional[torch.Tensor] = None,   # 只FA2用
+        kv_cu_seqlens: Optional[torch.Tensor] = None,# 只FA2用
+        **kwargs,
+    ) -> torch.Tensor:
+        # 允许 query_states [B,T,d] 或 [T,d]，自动扩展 batch 维
+        orig_2d = False
+        if query_states.dim() == 2:
+            query_states = query_states.unsqueeze(0)
+            orig_2d = True
+
+        batch_size, seq_len_q, _ = query_states.shape
+
+        # Q/K/V投影
+        if self.is_cross_attention and key_value_states is not None:
+            if key_value_states.dim() == 2:
+                key_value_states = key_value_states.unsqueeze(0)
+            q = self.q_proj(query_states)
+            kv = self.kv(key_value_states)
+            seq_len_kv = kv.shape[1]
+            k, v = kv.reshape(batch_size, seq_len_kv, 2, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4).unbind(0)
+            q = q.reshape(batch_size, seq_len_q, self.num_heads, self.head_dim).transpose(1, 2)
+        else:
+            if key_value_states is None:
+                key_value_states = query_states
+            qkv = self.qkv(query_states)
+            q, k, v = qkv.reshape(batch_size, seq_len_q, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4).unbind(0)
+
+        # 选用哪个 attention kernel
+        attn_impl = getattr(self.config, '_attn_implementation', 'sdpa')
+        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS[attn_impl]
+
+        # ========= 支持 FA2 ==========
+        if attn_impl == "flash_attention_2":
+            # Qwen2_5 之所以能支持 FA2，是因为准备了 flatten+cu_seqlens
+            # 这里假设 query_states/key_value_states 按 batch 维是变长的
+
+            # 检查 cu_seqlens，有就用，否则尝试自动生成
+            if cu_seqlens is None:
+                # 默认把每个batch都视为长度=seq_len_q
+                cu_seqlens = torch.arange(0, (batch_size + 1) * seq_len_q, step=seq_len_q, dtype=torch.int32, device=q.device)
+            if kv_cu_seqlens is None:
+                cu_seqlens_k = torch.arange(0, (batch_size + 1) * k.shape[2], step=k.shape[2], dtype=torch.int32, device=k.device)
+            else:
+                cu_seqlens_k = kv_cu_seqlens
+
+            # flatten [B, nH, T, d] -> [total_T, nH, d]
+            # 注意！FlashAttn2是 (total, nH, d)，不是 (nH, total, d)，和普通实现不一样
+            # 更安全的 flatten 方式
+            # [B, nH, T, d] -> [B, T, nH, d] -> [total_T, nH, d]
+            q_ = q.transpose(1, 2).contiguous().view(-1, self.num_heads, self.head_dim)
+            k_ = k.transpose(1, 2).contiguous().view(-1, self.num_heads, self.head_dim)
+            v_ = v.transpose(1, 2).contiguous().view(-1, self.num_heads, self.head_dim)
+            
+            max_seqlen_q = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
+            max_seqlen_k = (cu_seqlens_k[1:] - cu_seqlens_k[:-1]).max().item()
+
+            attn_output, _ = attention_interface(
+                self,
+                q_,
+                k_,
+                v_,
+                attention_mask=None,
+                scaling=self.scaling,
+                dropout=0.0 if not self.training else self.attention_dropout,
+                cu_seq_lens_q=cu_seqlens,
+                cu_seq_lens_k=cu_seqlens_k,
+                max_length_q=max_seqlen_q,
+                max_length_k=max_seqlen_k,
+                is_causal=self.is_causal,
+                **kwargs,
+            )
+            
+            # 更简洁的输出重构
+            # [total_q, nH, d] -> [B, seq_len_q, nH, d]
+            attn_output = attn_output.view(batch_size, seq_len_q, self.num_heads, self.head_dim).contiguous()
+        else:
+            # 普通实现，下游实现就是 [B, nH, T, d]
+            attn_output, _ = attention_interface(
+                self,
+                q, k, v,
+                attention_mask=attention_mask,
+                scaling=self.scaling,
+                dropout=0.0 if not self.training else self.attention_dropout,
+                is_causal=self.is_causal,
+                **kwargs,
+            )
+            # attn_output: [B, nH, seq_q, d]
+            attn_output = attn_output.transpose(1, 2).contiguous()  # [B, seq_q, nH, d]
+
+        attn_output = attn_output.reshape(batch_size, seq_len_q, self.dim)  # [B, seq_q, D]
+        attn_output = self.proj(attn_output)
+        if orig_2d:
+            attn_output = attn_output.squeeze(0)
+        return attn_output.contiguous()
+
+
+class ADCopilotCompareVisualEncoder(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.sequence_compare = getattr(config, "sequence_compare", True)
+        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 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)
+        self.encoder_norm4 = Qwen2RMSNorm(self.hidden_size, eps=1e-6)
+        self.encoder_mlp1 = Qwen2_5_VLMLP(config)
+        self.encoder_mlp2 = Qwen2_5_VLMLP(config)
+        
+        # Decoder 部分：Query 与编码特征交互
+        # 可学习的 Query Embeddings
+        self.query_embeddings = nn.Parameter(
+            torch.empty(self.token_size, self.hidden_size)
+        )
+        # 只保留 Cross Attention for queries to attend to encoded features
+        self.decoder_cross_attn = OptimizedCrossAttention(config, is_cross_attention=True)
+        
+        self.decoder_norm1 = Qwen2RMSNorm(self.hidden_size, eps=1e-6)
+        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 init_query_embeddings(self):
+        nn.init.normal_(self.query_embeddings, mean=0.0, std=0.02)
+
+    def forward(self, images_hidden_states: list) -> torch.Tensor:
+        """
+        Args:
+            images_hidden_states: List of tensor, each tensor has shape [seq_len, hidden_size]
+        
+        Returns:
+            Tensor of shape [total_images, token_size, hidden_size]
+        """
+        if not images_hidden_states:
+            return torch.empty(0, self.token_size, self.hidden_size)
+        
+        # 检查 query_embeddings 是否包含 NaN
+        if torch.isnan(self.query_embeddings).any():
+            print("警告：query_embeddings 包含 NaN 值")
+            # nn.init.normal_(self.query_embeddings, mean=0.0, std=0.02)
+        
+        # 获取每个图像的序列长度
+        seq_lengths = [state.size(0) for state in images_hidden_states]
+        max_seq_len = max(seq_lengths)
+        batch_size = len(images_hidden_states)
+        device = images_hidden_states[0].device
+        dtype = images_hidden_states[0].dtype
+        
+        # 将所有图像填充到相同长度并堆叠
+        padded_states = []
+        attention_masks = []
+        for state in images_hidden_states:
+            pad_len = max_seq_len - state.size(0)
+            if pad_len > 0:
+                # 填充序列
+                padded_state = F.pad(state, (0, 0, 0, pad_len), mode='constant', value=0)
+                # 创建注意力掩码
+                attention_mask = torch.ones(max_seq_len, dtype=torch.bool, device=device)
+                attention_mask[state.size(0):] = False
+            else:
+                padded_state = state
+                attention_mask = torch.ones(max_seq_len, dtype=torch.bool, device=device)
+            padded_states.append(padded_state)
+            attention_masks.append(attention_mask)
+        
+        # [batch_size, max_seq_len, hidden_size]
+        batched_states = torch.stack(padded_states)
+        # [batch_size, max_seq_len]
+        attention_masks = torch.stack(attention_masks)
+        
+        # 创建循环移位的状态用于对比
+        # 对于第一个图像，使用自身作为previous
+        previous_states = torch.roll(batched_states, shifts=1, dims=0)
+        previous_masks = torch.roll(attention_masks, shifts=1, dims=0)
+
+        if previous_states.size(0) > 1 and self.sequence_compare:
+            previous_states[0] = previous_states[1]
+            previous_masks[0] = previous_masks[1]
+        
+        # Encoder: 批量处理所有图像
+        encoded_features = self._encoder_forward(
+            batched_states,  # [batch_size, max_seq_len, hidden_size]
+            previous_states,  # [batch_size, max_seq_len, hidden_size]
+            attention_masks,  # [batch_size, max_seq_len]
+            previous_masks   # [batch_size, max_seq_len]
+        )
+        
+        # Decoder: 批量处理所有图像
+        # 扩展query_embeddings到batch维度
+        batch_queries = self.query_embeddings.unsqueeze(0).expand(batch_size, -1, -1)
+        # [batch_size, token_size, hidden_size]
+        compare_visual_embeds = self._decoder_forward(
+            batch_queries,
+            encoded_features,
+            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, out_hidden_size]
+    
+    def _encoder_forward(self, current_features, previous_features, current_mask=None, previous_mask=None):
+        """
+        Encoder: 双向图像特征交互
+        Args:
+            current_features: [batch_size, seq_len, hidden_size]
+            previous_features: [batch_size, seq_len, hidden_size]
+            current_mask: [batch_size, seq_len]
+            previous_mask: [batch_size, seq_len]
+        """
+        # 第一步：previous attend to current
+        residual = previous_features
+        
+        # Layer norm
+        previous_normed = self.encoder_norm1(previous_features)
+        current_normed1 = self.encoder_norm1(current_features)
+        
+        # Cross attention: previous attend to current
+        cross_attn_output1 = self.encoder_cross_attn1(
+            query_states=previous_normed,
+            key_value_states=current_normed1,
+            attention_mask=current_mask.unsqueeze(1).unsqueeze(2) if current_mask is not None else None
+        )
+        
+        # Residual connection
+        previous_features = residual + cross_attn_output1
+        
+        # MLP for previous features
+        residual = previous_features
+        mlp_input1 = self.encoder_norm2(previous_features)
+        mlp_output1 = self.encoder_mlp1(mlp_input1)
+        previous_features = residual + mlp_output1
+        
+        # 第二步：current attend to previous (enhanced)
+        residual = current_features
+        
+        # Layer norm
+        current_normed2 = self.encoder_norm3(current_features)
+        previous_normed2 = self.encoder_norm3(previous_features)
+        
+        # Cross attention: current attend to previous
+        cross_attn_output2 = self.encoder_cross_attn2(
+            query_states=current_normed2,
+            key_value_states=previous_normed2,
+            attention_mask=previous_mask.unsqueeze(1).unsqueeze(2) if previous_mask is not None else None
+        )
+        
+        # Residual connection
+        current_features = residual + cross_attn_output2
+        
+        # MLP for current features
+        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
+        return current_features
+    
+    def _decoder_forward(self, queries, encoded_features, query_mask=None, encoded_mask=None):
+        """
+        Decoder: Query 与编码特征交互
+        Args:
+            queries: [batch_size, token_size, hidden_size]
+            encoded_features: [batch_size, seq_len, hidden_size]
+            query_mask: [batch_size, token_size]
+            encoded_mask: [batch_size, seq_len]
+        """
+        # Cross attention: queries attend to encoded features
+        residual = queries
+        queries_normed = self.decoder_norm1(queries)
+        encoded_normed = self.decoder_norm1(encoded_features)
+        
+        cross_attn_output = self.decoder_cross_attn(
+            query_states=queries_normed,
+            key_value_states=encoded_normed,
+            attention_mask=encoded_mask.unsqueeze(1).unsqueeze(2) if encoded_mask is not None else None
+        )
+        
+        queries = residual + cross_attn_output
+        
+        # MLP
+        residual = queries
+        mlp_input = self.decoder_norm2(queries)
+        mlp_output = self.decoder_mlp(mlp_input)
+        queries = residual + mlp_output
+        
+        return queries  # [batch_size, token_size, hidden_size]
+
+
+# 先把组件继承出来方便修改
+class ADCopilotVisionTransformerPretrainedModel(Qwen2_5_VisionTransformerPretrainedModel):
+    def __init__(self, config, *inputs, **kwargs) -> None:
+        super().__init__(config, *inputs, **kwargs)
+        self.compare_visual_encoder = ADCopilotCompareVisualEncoder(config)
+        
+    def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs) -> torch.Tensor:
+        """
+        Args:
+            hidden_states (`torch.Tensor` of shape `(seq_len, hidden_size)`):
+                The final hidden states of the model.
+            grid_thw (`torch.Tensor` of shape `(num_images_or_videos, 3)`):
+                The temporal, height and width of feature shape of each image in LLM.
+
+        Returns:
+            `torch.Tensor`: hidden_states, compare_visual_embeds.
+        """
+        hidden_states = self.patch_embed(hidden_states)
+        rotary_pos_emb = self.rot_pos_emb(grid_thw)
+        window_index, cu_window_seqlens = self.get_window_index(grid_thw)
+        cu_window_seqlens = torch.tensor(
+            cu_window_seqlens,
+            device=hidden_states.device,
+            dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32,
+        )
+        cu_window_seqlens = torch.unique_consecutive(cu_window_seqlens)
+
+        seq_len, _ = hidden_states.size()
+        hidden_states = hidden_states.reshape(seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1)
+        hidden_states = hidden_states[window_index, :, :]
+        hidden_states = hidden_states.reshape(seq_len, -1)
+        rotary_pos_emb = rotary_pos_emb.reshape(seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1)
+        rotary_pos_emb = rotary_pos_emb[window_index, :, :]
+        rotary_pos_emb = rotary_pos_emb.reshape(seq_len, -1)
+        emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1)
+        position_embeddings = (emb.cos(), emb.sin())
+
+        cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum(
+            dim=0,
+            # Select dtype based on the following factors:
+            #  - FA2 requires that cu_seqlens_q must have dtype int32
+            #  - torch.onnx.export requires that cu_seqlens_q must have same dtype as grid_thw
+            # See https://github.com/huggingface/transformers/pull/34852 for more information
+            dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32,
+        )
+        cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0)
+
+        for layer_num, blk in enumerate(self.blocks):
+            if layer_num in self.fullatt_block_indexes:
+                cu_seqlens_now = cu_seqlens
+            else:
+                cu_seqlens_now = cu_window_seqlens
+
+            hidden_states = blk(
+                hidden_states,
+                cu_seqlens=cu_seqlens_now,
+                position_embeddings=position_embeddings,
+                **kwargs,
+            )
+
+        split_sizes = grid_thw.prod(-1).tolist()
+        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)
+
+        
+        hidden_states = self.merger(hidden_states)
+        reverse_indices = torch.argsort(window_index)
+        hidden_states = hidden_states[reverse_indices, :]
+
+        return hidden_states, compare_visual_embeds
+
+class ADCopilotVLModel(Qwen2_5_VLModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.visual = ADCopilotVisionTransformerPretrainedModel._from_config(config.vision_config)
+        self.compare_token_size = config.vision_config.compare_token_size
+        # self.learnable_image_embeddings = nn.Parameter(
+        #     torch.randn(100, config.hidden_size) * 0.02  # 使用小的初始化值
+        # )   
+        
+    def get_image_features(self, pixel_values: torch.FloatTensor, image_grid_thw: Optional[torch.LongTensor] = None):
+        """
+        Encodes images into continuous embeddings that can be forwarded to the language model.
+
+        Args:
+            pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`):
+                The tensors corresponding to the input images.
+            image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+                The temporal, height and width of feature shape of each image in LLM.
+        """
+        pixel_values = pixel_values.type(self.visual.dtype)
+        image_embeds, compare_visual_embeds = self.visual(pixel_values, grid_thw=image_grid_thw)
+        # 每个图像添加了对比感知token
+        split_sizes = (image_grid_thw.prod(-1) // self.visual.spatial_merge_size**2).tolist()
+        image_embeds = torch.split(image_embeds, split_sizes)
+
+        # 将图像嵌入和对比视觉嵌入拼接
+        enhanced_image_embeds = []
+        for i, embeds in enumerate(image_embeds):
+            # 确保 compare_visual_embeds[i] 与 embeds 在相同设备和数据类型
+            compare_embed = compare_visual_embeds[i].to(device=embeds.device, dtype=embeds.dtype)
+            enhanced_embeds = torch.cat([embeds, compare_embed], dim=0)
+            enhanced_image_embeds.append(enhanced_embeds)
+        
+        # image_embeds = torch.cat(enhanced_image_embeds, dim=0)
+        return enhanced_image_embeds
+    
+    def get_rope_index(self, input_ids: Optional[torch.LongTensor] = None, image_grid_thw: Optional[torch.LongTensor] = None, video_grid_thw: Optional[torch.LongTensor] = None, second_per_grid_ts: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, **kwargs) -> tuple[torch.Tensor, torch.Tensor]:
+        return self.get_rope_index_with_compare_token(input_ids, image_grid_thw, video_grid_thw, second_per_grid_ts, attention_mask)
+    
+    def get_rope_index_with_compare_token(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        image_grid_thw: Optional[torch.LongTensor] = None,
+        video_grid_thw: Optional[torch.LongTensor] = None,
+        second_per_grid_ts: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        spatial_merge_size = self.config.vision_config.spatial_merge_size
+        image_token_id = self.config.image_token_id
+        video_token_id = self.config.video_token_id
+        vision_start_token_id = self.config.vision_start_token_id
+        mrope_position_deltas = []
+        if input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None):
+            total_input_ids = input_ids
+            if attention_mask is None:
+                attention_mask = torch.ones_like(total_input_ids)
+            position_ids = torch.ones(
+                3,
+                input_ids.shape[0],
+                input_ids.shape[1],
+                dtype=input_ids.dtype,
+                device=input_ids.device,
+            )
+            image_index, video_index = 0, 0
+            attention_mask = attention_mask.to(total_input_ids.device)
+            for i, input_ids in enumerate(total_input_ids):
+                input_ids = input_ids[attention_mask[i] == 1]
+                image_nums, video_nums = 0, 0
+                vision_start_indices = torch.argwhere(input_ids == vision_start_token_id).squeeze(1)
+                vision_tokens = input_ids[vision_start_indices + 1]
+                image_nums = (vision_tokens == image_token_id).sum()
+                video_nums = (vision_tokens == video_token_id).sum()
+                input_tokens = input_ids.tolist()
+                llm_pos_ids_list: list = []
+                st = 0
+                remain_images, remain_videos = image_nums, video_nums
+                for vision_index in range(image_nums + video_nums):
+                    if image_token_id in input_tokens and remain_images > 0:
+                        ed_image = input_tokens.index(image_token_id, st)
+                    else:
+                        ed_image = len(input_tokens) + 1
+                    if video_token_id in input_tokens and remain_videos > 0:
+                        ed_video = input_tokens.index(video_token_id, st)
+                    else:
+                        ed_video = len(input_tokens) + 1
+                    if ed_image < ed_video:
+                        t, h, w = (
+                            image_grid_thw[image_index][0],
+                            image_grid_thw[image_index][1],
+                            image_grid_thw[image_index][2],
+                        )
+                        second_per_grid_t = 0
+                        image_index += 1
+                        remain_images -= 1
+                        ed = ed_image
+
+                    else:
+                        t, h, w = (
+                            video_grid_thw[video_index][0],
+                            video_grid_thw[video_index][1],
+                            video_grid_thw[video_index][2],
+                        )
+                        if second_per_grid_ts is not None:
+                            second_per_grid_t = second_per_grid_ts[video_index]
+                        else:
+                            second_per_grid_t = 1.0
+                        video_index += 1
+                        remain_videos -= 1
+                        ed = ed_video
+                    llm_grid_t, llm_grid_h, llm_grid_w = (
+                        t.item(),
+                        h.item() // spatial_merge_size,
+                        w.item() // spatial_merge_size,
+                    )
+                    text_len = ed - st
+
+                    st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+                    llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+
+                    range_tensor = torch.arange(llm_grid_t).view(-1, 1)
+                    expanded_range = range_tensor.expand(-1, llm_grid_h * llm_grid_w)
+
+                    ## normalize type, send to device.
+                    second_per_grid_t = torch.as_tensor(
+                        second_per_grid_t, dtype=range_tensor.dtype, device=range_tensor.device
+                    )
+
+                    time_tensor = expanded_range * second_per_grid_t * self.config.vision_config.tokens_per_second
+
+                    time_tensor_long = time_tensor.long()
+                    t_index = time_tensor_long.flatten()
+
+                    h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten()
+                    w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten()
+                    llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + text_len + st_idx)
+                    st = ed + llm_grid_t * llm_grid_h * llm_grid_w
+                    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 = 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
+
+                if st < len(input_tokens):
+                    st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+                    text_len = len(input_tokens) - st
+                    llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+
+                llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
+                position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(position_ids.device)
+                mrope_position_deltas.append(llm_positions.max() + 1 - len(total_input_ids[i]))
+            mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1)
+            return position_ids, mrope_position_deltas
+        else:
+            if attention_mask is not None:
+                position_ids = attention_mask.long().cumsum(-1) - 1
+                position_ids.masked_fill_(attention_mask == 0, 1)
+                position_ids = position_ids.unsqueeze(0).expand(3, -1, -1).to(attention_mask.device)
+                max_position_ids = position_ids.max(0, keepdim=False)[0].max(-1, keepdim=True)[0]
+                mrope_position_deltas = max_position_ids + 1 - attention_mask.shape[-1]
+            else:
+                position_ids = (
+                    torch.arange(input_ids.shape[1], device=input_ids.device)
+                    .view(1, 1, -1)
+                    .expand(3, input_ids.shape[0], -1)
+                )
+                mrope_position_deltas = torch.zeros(
+                    [input_ids.shape[0], 1],
+                    device=input_ids.device,
+                    dtype=input_ids.dtype,
+                )
+
+            return position_ids, mrope_position_deltas
+
+class ADCopilotVLForConditionalGeneration(Qwen2_5_VLForConditionalGeneration):    
+    config_class = ADCopilotConfig
+    
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = ADCopilotVLModel(config)