腾讯开源"全知视觉"模型 WeVisOne-VL：一个模型干翻 9 大 CV 任务，2B 参数吊打 9B 巨头

# wevisone_vl_core.py
"""
WeVisOne-VL 核心架构骨架（基于公开资料整理）
核心思想：VLM 主干输出 <OBJ> token 的隐藏状态 → 像素解码器 → 实例级像素输出
"""

import torch
import torch.nn as nn
import torch.nn.functional as F

# ============================================================
# 0. 特殊 Token 定义
# ============================================================
OBJ_TOKEN_ID = 152678          # <OBJ> token ID
COORD_TOKEN_RANGE = (0, 1000)  # 坐标 token <0>...<1000> 表示归一化框/点


# ============================================================
# 1. Visual Feature Adapter
#    作用：将 Qwen3-VL 多层级 ViT 特征适配到解码器输入空间
# ============================================================
class VisualFeatureAdapter(nn.Module):
    """
    对 Qwen3-VL 多层级 ViT 特征进行投影/融合，
    输出统一通道数的视觉特征，供后续可变形注意力采样使用。
    """
    def __init__(self, in_channels_list, hidden_dim=256):
        super().__init__()
        self.proj_layers = nn.ModuleList([
            nn.Conv2d(in_c, hidden_dim, kernel_size=1)
            for in_c in in_channels_list
        ])
        self.fusion = nn.Sequential(
            nn.Conv2d(hidden_dim * len(in_channels_list), hidden_dim, 1),
            nn.GroupNorm(32, hidden_dim),
            nn.GELU(),
        )

    def forward(self, multi_level_feats):
        """
        multi_level_feats: List[Tensor]，每个 Tensor 形状 (B, C_i, H_i, W_i)
        返回: (B, hidden_dim, H, W) 的统一特征图（以最高分辨率层为基准插值对齐）
        """
        target_size = multi_level_feats[0].shape[-2:]
        proj_feats = []
        for feat, proj in zip(multi_level_feats, self.proj_layers):
            x = proj(feat)
            if x.shape[-2:] != target_size:
                x = F.interpolate(x, size=target_size, mode="bilinear", align_corners=False)
            proj_feats.append(x)
        return self.fusion(torch.cat(proj_feats, dim=1))


# ============================================================
# 2. Box-Guided Query Decoder（核心）
#    6 层可变形交叉注意力，query 由 <OBJ> 隐藏状态初始化，
#    并结合 Fourier 编码的 box prior 进行空间引导
# ============================================================
class DeformableCrossAttentionBlock(nn.Module):
    """
    简化版可变形交叉注意力块。
    实际实现中应包含 offset / attention weight 预测分支与 grid_sample。
    此处保留核心结构语义，省略底层采样细节。
    """
    def __init__(self, dim=256, num_heads=8):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(dim, num_heads, batch_first=True)
        self.cross_attn = nn.MultiheadAttention(dim, num_heads, batch_first=True)
        self.norm1 = nn.LayerNorm(dim)
        self.norm2 = nn.LayerNorm(dim)
        self.norm3 = nn.LayerNorm(dim)
        self.ffn = nn.Sequential(
            nn.Linear(dim, dim * 4),
            nn.GELU(),
            nn.Linear(dim * 4, dim),
        )

    def forward(self, query, visual_feat_flat, box_prior=None):
        """
        query: (B, N_q, dim)  —— N_q 为 <OBJ> 数量
        visual_feat_flat: (B, H*W, dim)
        box_prior: (B, N_q, dim) 或 None，Fourier 编码后的 box 先验
        """
        # 自注意力：让同一答案内的多个 <OBJ> 互相感知
        q = self.norm1(query)
        query = query + self.self_attn(q, q, q)[0]

        # 交叉注意力：query 从视觉特征中采样
        q = self.norm2(query)
        if box_prior is not None:
            q = q + box_prior
        query = query + self.cross_attn(q, visual_feat_flat, visual_feat_flat)[0]

        # FFN
        query = query + self.ffn(self.norm3(query))
        return query


class BoxGuidedQueryDecoder(nn.Module):
    def __init__(self, dim=256, num_layers=6, num_heads=8):
        super().__init__()
        self.layers = nn.ModuleList([
            DeformableCrossAttentionBlock(dim, num_heads)
            for _ in range(num_layers)
        ])
        self.box_prior_proj = nn.Linear(dim, dim)  # Fourier 编码 → 模型维度

    def forward(self, obj_hidden, visual_feat, boxes_fourier):
        """
        obj_hidden: (B, N_obj, D_vlm)  —— 来自 VLM 的 <OBJ> 隐藏状态
        visual_feat: (B, H*W, D)       —— Adapter 输出展平
        boxes_fourier: (B, N_obj, D)   —— 每个 <OBJ> 对应框的 Fourier 编码
        """
        query = obj_hidden
        box_prior = self.box_prior_proj(boxes_fourier)
        for layer in self.layers:
            query = layer(query, visual_feat, box_prior)
        return query


# ============================================================
# 3. Instance-Conditioned Dense Prediction Head
#    FullCanvasAlphaCascadeHead：在 S/4 → S/2 → S/1 上逐级细化
# ============================================================
class FullCanvasAlphaCascadeHead(nn.Module):
    """
    从解码后的实例 query 生成全分辨率 alpha matte / mask。
    级联结构：低分辨率粗预测 → 上采样 + 特征融合 → 高分辨率精修。
    """
    def __init__(self, dim=256, num_classes=1):
        super().__init__()
        # S/4 阶段
        self.stage_s4 = nn.Sequential(
            nn.Conv2d(dim, dim, 3, padding=1),
            nn.GroupNorm(32, dim),
            nn.GELU(),
            nn.Conv2d(dim, num_classes, 1),
        )
        # S/2 阶段：融合 S/4 上采样结果与图像特征
        self.stage_s2 = nn.Sequential(
            nn.Conv2d(dim + num_classes, dim, 3, padding=1),
            nn.GroupNorm(32, dim),
            nn.GELU(),
            nn.Conv2d(dim, num_classes, 1),
        )
        # S/1 阶段
        self.stage_s1 = nn.Sequential(
            nn.Conv2d(dim + num_classes, dim, 3, padding=1),
            nn.GroupNorm(32, dim),
            nn.GELU(),
            nn.Conv2d(dim, num_classes, 1),
        )

    def forward(self, instance_feat, image_feat_pyramid):
        """
        instance_feat: (B*N, dim, H/4, W/4) —— 解码器 query 重塑为空间形式
        image_feat_pyramid: Dict[str, Tensor]，包含 "s4", "s2", "s1" 的图像特征
        """
        # S/4
        pred_s4 = self.stage_s4(instance_feat)

        # S/2：上采样 S/4 预测 + 融合
        pred_s4_up = F.interpolate(pred_s4, scale_factor=2, mode="bilinear", align_corners=False)
        feat_s2 = torch.cat([image_feat_pyramid["s2"], pred_s4_up], dim=1)
        pred_s2 = self.stage_s2(feat_s2)

        # S/1：再上采样 + 融合，输出全分辨率
        pred_s2_up = F.interpolate(pred_s2, scale_factor=2, mode="bilinear", align_corners=False)
        feat_s1 = torch.cat([image_feat_pyramid["s1"], pred_s2_up], dim=1)
        pred_s1 = self.stage_s1(feat_s1)

        return pred_s1  # (B*N, 1, H, W)，stride 1


# ============================================================
# 4. 完整像素解码器：组合以上三个模块
# ============================================================
class WeVisOneMaskDecoder(nn.Module):
    """
    约 9.6M 参数的核心解码器（实际参数量取决于各层通道配置）。
    输入：VLM 的 <OBJ> 隐藏状态 + 多级 ViT 特征 + 框的 Fourier 编码
    输出：每个 <OBJ> 对应的全分辨率 alpha matte / mask
    """
    def __init__(self, vlm_dim, vit_channels_list, hidden_dim=256, num_queries_per_obj=1):
        super().__init__()
        self.num_queries_per_obj = num_queries_per_obj

        # <OBJ> 隐藏状态投影到解码器维度
        self.obj_proj = nn.Linear(vlm_dim, hidden_dim)

        # 1) 视觉特征适配器
        self.visual_adapter = VisualFeatureAdapter(vit_channels_list, hidden_dim)

        # 2) 框引导的 Query 解码器（6 层）
        self.query_decoder = BoxGuidedQueryDecoder(hidden_dim, num_layers=6)

        # 3) 实例条件密集预测头（级联）
        self.dense_head = FullCanvasAlphaCascadeHead(hidden_dim)

        # 将解码后的 query 重塑为空间特征的小型卷积
        self.query_to_spatial = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim * 4 * 4),
            nn.GELU(),
        )

    def forward(self, obj_hidden_states, multi_level_vit_feats, boxes, image_feat_pyramid):
        """
        obj_hidden_states: (B, N_obj, D_vlm)  —— 从 VLM 中提取的 <OBJ> 最后层隐藏状态
        multi_level_vit_feats: List[Tensor]  —— Qwen3-VL 多层级 pre-merging ViT 特征
        boxes: (B, N_obj, 4)  —— 归一化 xyxy
        image_feat_pyramid: Dict[str, Tensor]  —— "s4", "s2", "s1" 分辨率特征，用于级联头融合
        """
        B, N_obj, _ = obj_hidden_states.shape

        # 1) 视觉特征适配
        visual_feat = self.visual_adapter(multi_level_vit_feats)   # (B, C, H, W)
        Bv, Cv, Hv, Wv = visual_feat.shape
        visual_feat_flat = visual_feat.flatten(2).transpose(1, 2)  # (B, H*W, C)

        # 2) <OBJ> 隐藏状态投影
        query = self.obj_proj(obj_hidden_states)                   # (B, N_obj, hidden_dim)

        # 3) 框的 Fourier 编码
        boxes_fourier = fourier_encode_boxes(boxes)                # (B, N_obj, D_fourier)
        # 实际中 D_fourier 可能与 hidden_dim 不同，此处假设已对齐或通过 proj 处理

        # 4) 框引导的 Query 解码（6 层可变形交叉注意力）
        decoded_queries = self.query_decoder(query, visual_feat_flat, boxes_fourier)

        # 5) 将每个 <OBJ> 的 query 重塑为 4×4 的空间特征图
        spatial_feat = self.query_to_spatial(decoded_queries)      # (B, N_obj, C*16)
        spatial_feat = spatial_feat.view(B * N_obj, -1, 4, 4)      # (B*N, C, 4, 4)
        # 上采样到 S/4 尺度，与视觉特征图对齐
        spatial_feat = F.interpolate(spatial_feat, size=(Hv, Wv), mode="bilinear", align_corners=False)

        # 6) 级联密集预测头：输出全分辨率 alpha matte
        masks = self.dense_head(spatial_feat, image_feat_pyramid)  # (B*N, 1, H, W)

        return masks.view(B, N_obj, *masks.shape[1:])


def fourier_encode_boxes(boxes, num_bands=32):
    """
    对归一化框 (xyxy) 做 Fourier 特征编码，作为解码器的空间先验。
    简化实现：对 4 个坐标分别做 sin/cos 多频编码后拼接。
    """
    # boxes: (B, N, 4)
    B, N, _ = boxes.shape
    freqs = 2.0 ** torch.arange(num_bands, device=boxes.device).float() * torch.pi
    coords = boxes.unsqueeze(-1)                       # (B, N, 4, 1)
    encoded = torch.cat([torch.sin(coords * freqs), torch.cos(coords * freqs)], dim=-1)
    return encoded.view(B, N, -1)                      # (B, N, 4 * num_bands * 2)


# ============================================================
# 5. 与 VLM 主干的衔接（示意）
# ============================================================
def extract_obj_hidden_states(vlm_outputs, input_ids):
    """
    从 Qwen3-VL 的输出中提取所有 <OBJ> token 位置的隐藏状态。
    vlm_outputs.last_hidden_state: (B, L, D_vlm)
    """
    obj_positions = (input_ids == OBJ_TOKEN_ID)        # (B, L)
    # 实际实现需处理 batch 内不同数量的 <OBJ>，此处假设可 gather 出 (B, N_obj, D)
    obj_hidden = vlm_outputs.last_hidden_state[obj_positions].view(
        input_ids.shape[0], -1, vlm_outputs.last_hidden_state.shape[-1]
    )
    return obj_hidden