vit_insert_merger.py

import torch
from torch import nn
from typing import Any, Optional, Tuple, Union, List
import torch.nn.functional as F
from .modeling_navit_siglip_fast import SiglipAttention, SiglipFlashAttention2, SiglipMLP

from transformers.activations import ACT2FN
from transformers.utils import (
    is_flash_attn_2_available,
    logging,
)

logger = logging.get_logger(__name__)

if is_flash_attn_2_available():
    from flash_attn import flash_attn_func, flash_attn_varlen_func
    from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input  # noqa

def get_vit_insert_merger(hidden_size, intermediate_size, vpm, insert_layer_id):
    return ViTWindowAttentionMerger(vpm, insert_layer_id, downsample='ViTmlp', foreach=True)

class ViTWindowAttentionMerger(nn.Module):
    def __init__(self, vpm, insert_layer_id, downsample=None, foreach=False):
        super().__init__()
        self.window_kernel_size = (2, 2)
        assert downsample in ['average', 'ViTmlp', 'ViTmlp_only', 'self_attention_ViTmlp', None], f"Unknown downsample way: {downsample}"
        self.downsample = downsample
        self.foreach = foreach

        self.embed_dim = vpm.config.hidden_size
        self._use_flash_attention_2 = vpm.config._attn_implementation == "flash_attention_2"
        self.self_attn = (
            SiglipAttention(vpm.config)
            if not self._use_flash_attention_2
            else SiglipFlashAttention2(vpm.config)
        )
        self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=vpm.config.layer_norm_eps)
        self.mlp = SiglipMLP(vpm.config)
        self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=vpm.config.layer_norm_eps)

        if self.downsample in ['ViTmlp', 'ViTmlp_only', 'self_attention_ViTmlp']:
            self.hidden_size = (
                self.embed_dim
                * self.window_kernel_size[0]
                * self.window_kernel_size[1]
            )

            self.intermediate_size = (
                vpm.config.intermediate_size
                * self.window_kernel_size[0]
                * self.window_kernel_size[1]
            )

            self.pre_norm = torch.nn.LayerNorm(self.hidden_size, eps=1e-6)
            self.linear_1 = nn.Linear(self.hidden_size, self.intermediate_size, bias=True)
            self.act = ACT2FN["gelu_pytorch_tanh"]
            self.linear_2 = nn.Linear(
                self.intermediate_size, self.embed_dim, bias=True
            )

        self._init_weight(vpm, insert_layer_id)

    def _init_weight(self, vpm, insert_layer_id):
        copy_block = vpm.encoder.layers[insert_layer_id]

        # 拷贝 self-attention
        with torch.no_grad():
            for target_module, src_module in [
                (self.self_attn, copy_block.self_attn),
                (self.layer_norm1, copy_block.layer_norm1),
                (self.layer_norm2, copy_block.layer_norm2),
                (self.mlp, copy_block.mlp),
            ]:
                target_state = target_module.state_dict()
                src_state = src_module.state_dict()
                for k, v in src_state.items():
                    if k in target_state and v.shape == target_state[k].shape:
                        target_state[k].copy_(v)
                target_module.load_state_dict(target_state)

            if self.downsample in ['ViTmlp', 'ViTmlp_only', 'self_attention_ViTmlp']:
                fc1_old = copy_block.mlp.fc1       # [intermediate, hidden]
                fc2_old = copy_block.mlp.fc2       # [hidden, intermediate]

                # -------------------------------
                # fc1_new: Linear(hidden*4, inter*4)
                # 形成 block diagonal (4个 fc1)
                # -------------------------------
                hidden = fc1_old.weight.shape[1]
                inter = fc1_old.weight.shape[0]

                fc1_blocks = [fc1_old.weight.data] * 4  # list of [inter, hidden]
                w_fc1_new = torch.zeros(inter * 4, hidden * 4, device=fc1_old.weight.device)
                for i in range(4):
                    w_fc1_new[i*inter:(i+1)*inter, i*hidden:(i+1)*hidden] = fc1_blocks[i]
                
                b_fc1_new = fc1_old.bias.data.repeat(4)  # [inter*4]

                self.linear_1.weight.copy_(w_fc1_new)
                self.linear_1.bias.copy_(b_fc1_new)

                # -------------------------------
                # fc2_new: Linear(inter*4, hidden)
                # 水平拼接 4 份 + 除以 4
                # -------------------------------
                w_fc2_new = torch.cat([fc2_old.weight.data] * 4, dim=1) / 4.0  # [hidden, inter*4]
                b_fc2_new = fc2_old.bias.data  # [hidden]

                self.linear_2.weight.copy_(w_fc2_new)
                self.linear_2.bias.copy_(b_fc2_new)

                # -------------------------------
                # pre_norm
                # -------------------------------
                self.pre_norm.weight.data.copy_(
                    copy_block.layer_norm2.weight.data.repeat(4)
                )
                self.pre_norm.bias.data.copy_(
                    copy_block.layer_norm2.bias.data.repeat(4)
                )

    def get_window_index(self, tgt_sizes):
        """
        tgt_sizes: list or tensor of (H, W)
        return:
            window_index: Tensor[total_tokens] -> 按 window 顺序排列的 token 索引
            cu_seqlens: Tensor[num_windows + 1]
        """
        window_h, window_w = self.window_kernel_size  # e.g. (2, 2)
        max_seqlens = window_h * window_w

        window_index_list = []
        cu_seqlens = [0]
        token_offset = 0

        for (H, W) in tgt_sizes:
            assert H % window_h == 0 and W % window_w == 0, \
                f"H={H}, W={W} must be divisible by window size ({window_h}, {window_w})"

            index = torch.arange(H * W).reshape(H, W)

            num_windows_h = H // window_h
            num_windows_w = W // window_w
            num_windows = num_windows_h * num_windows_w

            index = index.reshape(num_windows_h, window_h, num_windows_w, window_w)
            index = index.permute(0, 2, 1, 3).reshape(num_windows, window_h * window_w)

            index_flat = index.reshape(-1) + token_offset
            window_index_list.append(index_flat)

            window_token_count = window_h * window_w
            cu_this = torch.arange(1, num_windows + 1) * window_token_count + cu_seqlens[-1]
            cu_seqlens.extend(cu_this.tolist())

            token_offset += H * W      # 偏移

        window_index = torch.cat(window_index_list)
        cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int32)

        return window_index, cu_seqlens, max_seqlens

    def forward(
        self,
        hidden_states: torch.Tensor,
        tgt_sizes: torch.IntTensor, 
        attention_mask: torch.Tensor,
        cu_seqlens: torch.Tensor = None,
        max_seqlens: torch.Tensor = None,
    ) -> Tuple[torch.FloatTensor]:
        
        if self.downsample == 'self_attention_ViTmlp':
            residual = hidden_states
            hidden_states = self.layer_norm1(hidden_states)

            hidden_states, attn_weights = self.self_attn(
                hidden_states=hidden_states,
                attention_mask=attention_mask,
                cu_seqlens=cu_seqlens,
                max_seqlens=max_seqlens,
                tgt_sizes=tgt_sizes,
            )
            hidden_states = residual + hidden_states
        elif self.downsample != 'ViTmlp_only':
            residual = hidden_states
            hidden_states = self.layer_norm1(hidden_states)
            device = hidden_states.device

            if self.foreach:
                all_pixel_values = []
                batch_size, _ = tgt_sizes.shape
                for batch_idx in range(batch_size):
                    hidden_state = hidden_states[0, cu_seqlens[batch_idx]:cu_seqlens[batch_idx+1], :].unsqueeze(0)
                    tgt_size = tgt_sizes[batch_idx].unsqueeze(0)
                    
                    window_index, window_cu_seqlens, window_max_seqlens = self.get_window_index(tgt_size)
                    hidden_state = hidden_state[:, window_index, :]

                    hidden_state, _ = self.self_attn(
                        hidden_states=hidden_state,
                        attention_mask=attention_mask,
                        cu_seqlens=window_cu_seqlens.to(device),
                        max_seqlens=window_max_seqlens,
                        tgt_sizes=tgt_size,
                    )
                    
                    all_pixel_values.append(hidden_state[:, torch.argsort(window_index), :])    # 生成反向索引

                hidden_states = torch.concat(all_pixel_values, dim=1)
                hidden_states = residual + hidden_states

                # # 判断两种方法得到的结果相不相等
                # print("hidden_states1: ", hidden_states1)
                # print("hidden_states2: ", hidden_states2)
                # if torch.allclose(hidden_states1, hidden_states2, rtol=1e-4, atol=1e-6):
                #     print("✅ 两个张量近似相等")
                # else:
                #     print("❌ 两个张量不相等")
            else:
                window_index, window_cu_seqlens, window_max_seqlens = self.get_window_index(tgt_sizes)
                hidden_states = hidden_states[:, window_index, :]
                hidden_states, attn_weights = self.self_attn(
                    hidden_states=hidden_states,
                    attention_mask=attention_mask,
                    cu_seqlens=window_cu_seqlens.to(device),
                    max_seqlens=window_max_seqlens,
                    tgt_sizes=tgt_sizes,
                )
                hidden_states = hidden_states[:, torch.argsort(window_index), :]    # 生成反向索引
                hidden_states = residual + hidden_states

        if self.downsample == 'average':
            batch_size, _ = tgt_sizes.shape
            all_pixel_values = []
            new_tgt_sizes = torch.zeros_like(tgt_sizes, dtype=tgt_sizes.dtype, device=tgt_sizes.device)

            m1, m2 = self.window_kernel_size
            for batch_idx in range(batch_size):
                h, w = tgt_sizes[batch_idx]
                assert h % 2 == 0 and w % 2 == 0, "patch尺寸不能被2整除, 无法拼接4个相邻patch"
                from einops import rearrange
                hidden_state = rearrange(hidden_states[0, cu_seqlens[batch_idx]:cu_seqlens[batch_idx+1], :].squeeze(0), "(h p1 w p2) d -> (h w) (p1 p2) d", h=h // m1, p1=m1, w=w // m2, p2=m2)

                hidden_state = hidden_state.mean(dim=1)  # (num_windows, D)
                
                all_pixel_values.append(hidden_state)
                new_tgt_sizes[batch_idx, :2] = torch.tensor([h // 2, w // 2], device=new_tgt_sizes.device, dtype=new_tgt_sizes.dtype)

            new_hidden_states = torch.concat(all_pixel_values, dim=0).unsqueeze(0)
            new_cu_seqlens = F.pad(torch.cumsum(new_tgt_sizes[:, 0] * new_tgt_sizes[:, 1], dim=0, dtype=torch.int32).cuda(), (1, 0))
            assert max_seqlens % 4 == 0
            new_max_seqlens = max_seqlens // 4

            return new_hidden_states, new_tgt_sizes, attention_mask, new_cu_seqlens, new_max_seqlens   # attention_mask=None
        elif self.downsample in ['ViTmlp', 'ViTmlp_only', 'self_attention_ViTmlp']:
            batch_size, _ = tgt_sizes.shape
            all_pixel_values = []
            new_tgt_sizes = torch.zeros_like(tgt_sizes, dtype=tgt_sizes.dtype, device=tgt_sizes.device)

            m1, m2 = self.window_kernel_size
            for batch_idx in range(batch_size):
                h, w = tgt_sizes[batch_idx]
                assert h % 2 == 0 and w % 2 == 0, "patch尺寸不能被2整除, 无法拼接4个相邻patch"
                from einops import rearrange
                hidden_state = rearrange(hidden_states[0, cu_seqlens[batch_idx]:cu_seqlens[batch_idx+1], :].squeeze(0), "(h p1 w p2) d -> (h w) (p1 p2 d)", h=h // m1, p1=m1, w=w // m2, p2=m2)
                
                residual = rearrange(hidden_states[0, cu_seqlens[batch_idx]:cu_seqlens[batch_idx+1], :].squeeze(0), "(h p1 w p2) d -> (h w) (p1 p2) d", h=h // m1, p1=m1, w=w // m2, p2=m2).mean(dim=1)   # [H/2*W/2, D]
                
                hidden_state = self.pre_norm(hidden_state)
                hidden_state = self.linear_1(hidden_state)
                hidden_state = self.act(hidden_state)
                hidden_state = self.linear_2(hidden_state)
                
                all_pixel_values.append(hidden_state + residual)
                new_tgt_sizes[batch_idx, :2] = torch.tensor([h // 2, w // 2], device=new_tgt_sizes.device, dtype=new_tgt_sizes.dtype)

            new_hidden_states = torch.concat(all_pixel_values, dim=0).unsqueeze(0)
            new_cu_seqlens = F.pad(torch.cumsum(new_tgt_sizes[:, 0] * new_tgt_sizes[:, 1], dim=0, dtype=torch.int32).cuda(), (1, 0))
            assert max_seqlens % 4 == 0
            new_max_seqlens = max_seqlens // 4

            return new_hidden_states, new_tgt_sizes, attention_mask, new_cu_seqlens, new_max_seqlens   # attention_mask=None
        else:
            residual = hidden_states
            hidden_states = self.layer_norm2(hidden_states)
            hidden_states = self.mlp(hidden_states)
            hidden_states = residual + hidden_states

            return hidden_states, tgt_sizes, attention_mask, cu_seqlens, max_seqlens   # attention_mask=None