model/layers.py

import math
import inspect
from contextlib import nullcontext

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Tuple


def conv_layer(in_dim, out_dim, kernel_size=1, padding=0, stride=1):
    return nn.Sequential(
        nn.Conv2d(in_dim, out_dim, kernel_size, stride, padding, bias=False),
        nn.BatchNorm2d(out_dim), nn.ReLU(True))


def linear_layer(in_dim, out_dim, bias=False):
    return nn.Sequential(nn.Linear(in_dim, out_dim, bias),
                         nn.BatchNorm1d(out_dim), nn.ReLU(True))


def _build_cabs_mlp(depth, hidden_size, output_hidden_size):
    layers = [nn.Linear(hidden_size, output_hidden_size)]
    for _ in range(1, depth):
        layers.append(nn.SiLU())
        layers.append(nn.Linear(output_hidden_size, output_hidden_size))
    return nn.Sequential(*layers)


class LayerNorm2d(nn.Module):
    def __init__(self, num_channels, eps=1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(num_channels))
        self.bias = nn.Parameter(torch.zeros(num_channels))
        self.eps = eps

    def forward(self, x):
        mean = x.mean(dim=1, keepdim=True)
        var = (x - mean).pow(2).mean(dim=1, keepdim=True)
        x = (x - mean) * torch.rsqrt(var + self.eps)
        return x * self.weight[:, None, None] + self.bias[:, None, None]


class CAbstractorBlock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.shortcut = (
            nn.Identity()
            if in_channels == out_channels
            else nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False)
        )
        self.net = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False),
            LayerNorm2d(out_channels),
            nn.SiLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False),
            LayerNorm2d(out_channels),
            nn.SiLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=1, bias=False),
            LayerNorm2d(out_channels),
        )
        self.act = nn.SiLU(inplace=True)

    def forward(self, x):
        return self.act(self.net(x) + self.shortcut(x))


class CAbstractorAdapter(nn.Module):
    """
    Local C-Abstractor adapter for v2_C.

    It mirrors Honeybee's C-Abstractor shape contract without importing Honeybee:
    2D visual features are optionally position encoded, processed by convolutional
    abstractor blocks, adaptively pooled to query tokens, then projected to the
    language hidden size.
    """
    def __init__(
        self,
        in_channels,
        hidden_channels,
        out_channels,
        depth=3,
        mlp_depth=2,
        num_query_tokens=144,
        num_input_tokens=None,
        pos_emb=True,
        prenorm=False,
    ):
        super().__init__()
        self.in_channels = in_channels
        self.num_query_tokens = num_query_tokens
        self.input_hw = None
        self.output_hw = None

        if num_query_tokens is not None and num_query_tokens > 0:
            query_hw = int(math.sqrt(num_query_tokens))
            if query_hw * query_hw != num_query_tokens:
                raise ValueError("num_query_tokens must be a square number")
            self.output_hw = (query_hw, query_hw)

        self.prenorm = nn.LayerNorm(in_channels) if prenorm else None
        if pos_emb and num_input_tokens is not None:
            self.pos_emb = nn.Parameter(torch.zeros(1, num_input_tokens, in_channels))
            if hasattr(nn.init, "trunc_normal_"):
                nn.init.trunc_normal_(self.pos_emb, mean=0.0, std=0.02)
            else:
                nn.init.normal_(self.pos_emb, mean=0.0, std=0.02)
        else:
            self.pos_emb = None

        if depth > 0:
            first_stage = self._make_stage(depth, in_channels, hidden_channels)
            second_stage = self._make_stage(depth, hidden_channels, hidden_channels)
            self.net = nn.ModuleList([first_stage, second_stage])
            self.readout = _build_cabs_mlp(mlp_depth, hidden_channels, out_channels)
        else:
            self.net = nn.ModuleList()
            self.readout = _build_cabs_mlp(mlp_depth, in_channels, out_channels)

    @staticmethod
    def _make_stage(depth, in_channels, out_channels):
        blocks = []
        for idx in range(depth):
            blocks.append(
                CAbstractorBlock(
                    in_channels if idx == 0 else out_channels,
                    out_channels,
                )
            )
        return nn.Sequential(*blocks)

    def _pos_emb_for_shape(self, h, w):
        if self.pos_emb is None:
            return None

        if self.pos_emb.size(1) == h * w:
            return self.pos_emb

        src_hw = int(math.sqrt(self.pos_emb.size(1)))
        if src_hw * src_hw != self.pos_emb.size(1):
            raise ValueError("Cannot interpolate non-square C-Abstractor positional embeddings")

        pos = self.pos_emb.transpose(1, 2).reshape(1, self.in_channels, src_hw, src_hw)
        pos = F.interpolate(pos, size=(h, w), mode="bicubic", align_corners=False)
        return pos.flatten(2).transpose(1, 2)

    def forward(self, x):
        b, c, h, w = x.shape
        tokens = x.flatten(2).transpose(1, 2)

        if self.prenorm is not None:
            tokens = self.prenorm(tokens)

        pos_emb = self._pos_emb_for_shape(h, w)
        if pos_emb is not None:
            tokens = tokens + pos_emb

        x = tokens.transpose(1, 2).reshape(b, c, h, w)
        if self.net:
            x = self.net[0](x)

        target_hw = self.output_hw if self.output_hw is not None else (h, w)
        x = F.adaptive_avg_pool2d(x, target_hw)

        if self.net:
            x = self.net[1](x)

        x = x.flatten(2).transpose(1, 2)
        x = self.readout(x)
        return x, (h, w)


class DAbstractorAdapter(nn.Module):
    """
    Local deformable-attention abstractor for v2_D.

    This follows Honeybee's D-Abstractor interface without importing Honeybee:
    visual features are projected to a Deformable DETR decoder memory, learned
    query tokens attend to the visual grid with deformable cross-attention, and
    the resulting query states are projected to the language hidden size.
    """
    def __init__(
        self,
        in_channels,
        out_channels,
        d_model=512,
        decoder_layers=2,
        num_heads=8,
        ffn_dim=1024,
        dropout=0.1,
        attention_dropout=0.0,
        num_query_tokens=144,
        num_input_tokens=None,
        pos_emb=True,
        pooled_v_target="query",
        learnable_ref_points=True,
        disable_custom_kernels=True,
        force_fp32=False,
        clamp_reference_points=False,
        reference_point_eps=1e-4,
        output_norm=False,
        output_scale=1.0,
        output_clip=None,
        sanitize_output=False,
    ):
        super().__init__()
        try:
            from transformers.models.deformable_detr import DeformableDetrConfig
            from transformers.models.deformable_detr.modeling_deformable_detr import (
                DeformableDetrDecoderLayer,
            )
        except ImportError as exc:
            raise ImportError(
                "DAbstractorAdapter requires transformers with Deformable DETR support."
            ) from exc

        if d_model % num_heads != 0:
            raise ValueError("d_abs_d_model must be divisible by d_abs_num_heads")

        query_hw = int(math.sqrt(num_query_tokens))
        if query_hw * query_hw != num_query_tokens:
            raise ValueError("d_abs_num_query_tokens must be a square number")

        if isinstance(in_channels, (list, tuple)):
            self.in_channels_per_level = [int(channels) for channels in in_channels]
        else:
            self.in_channels_per_level = [int(in_channels)]

        self.in_channels = self.in_channels_per_level[-1]
        self.num_feature_levels = len(self.in_channels_per_level)
        self.d_model = d_model
        self.num_query_tokens = num_query_tokens
        self.pooled_v_target = pooled_v_target
        self.force_fp32 = force_fp32
        self.clamp_reference_points = clamp_reference_points
        self.reference_point_eps = reference_point_eps
        self.output_scale = output_scale
        self.output_clip = output_clip
        self.sanitize_output = sanitize_output

        self.config = DeformableDetrConfig(
            d_model=d_model,
            decoder_layers=decoder_layers,
            decoder_attention_heads=num_heads,
            decoder_ffn_dim=ffn_dim,
            dropout=dropout,
            attention_dropout=attention_dropout,
            num_feature_levels=self.num_feature_levels,
            num_queries=num_query_tokens,
            disable_custom_kernels=disable_custom_kernels,
            use_pretrained_backbone=False,
        )
        self.layers = nn.ModuleList(
            [DeformableDetrDecoderLayer(self.config) for _ in range(decoder_layers)]
        )
        decoder_layer_params = inspect.signature(DeformableDetrDecoderLayer.forward).parameters
        self._uses_spatial_shapes_list = (
            "spatial_shapes_list" in decoder_layer_params
        )

        self.input_proj = nn.ModuleList(
            [
                nn.Linear(channels, d_model) if channels != d_model else nn.Identity()
                for channels in self.in_channels_per_level
            ]
        )
        if pos_emb and num_input_tokens is not None:
            num_input_tokens_per_level = self._normalize_num_input_tokens(num_input_tokens)
            pos_embs = []
            for tokens in num_input_tokens_per_level:
                pos_emb = nn.Parameter(torch.zeros(1, tokens, d_model))
                if hasattr(nn.init, "trunc_normal_"):
                    nn.init.trunc_normal_(pos_emb, mean=0.0, std=0.02)
                else:
                    nn.init.normal_(pos_emb, mean=0.0, std=0.02)
                pos_embs.append(pos_emb)
            self.pos_emb = (
                pos_embs[0]
                if self.num_feature_levels == 1
                else nn.ParameterList(pos_embs)
            )
        else:
            self.pos_emb = None

        if self.num_feature_levels > 1:
            self.level_emb = nn.Parameter(
                torch.empty(1, self.num_feature_levels, 1, d_model)
            )
            nn.init.normal_(self.level_emb)
        else:
            self.level_emb = None

        if pooled_v_target == "none":
            self.query_position_embeddings = nn.Embedding(num_query_tokens, d_model * 2)
            self.downsampler = None
        else:
            self.query_position_embeddings = nn.Embedding(num_query_tokens, d_model)
            self.downsampler = nn.AdaptiveAvgPool2d((query_hw, query_hw))

        reference_points = self._init_reference_points(num_query_tokens, num_input_tokens)
        if learnable_ref_points:
            self.reference_points = nn.Parameter(reference_points)
        else:
            self.register_buffer("reference_points", reference_points)

        self.output_proj = (
            nn.Linear(d_model, out_channels) if d_model != out_channels else nn.Identity()
        )
        self.output_norm = (
            nn.LayerNorm(out_channels, elementwise_affine=False)
            if output_norm
            else None
        )

    def _normalize_num_input_tokens(self, num_input_tokens):
        if isinstance(num_input_tokens, (list, tuple)):
            if len(num_input_tokens) != self.num_feature_levels:
                raise ValueError(
                    "num_input_tokens must have one value per D-Abstractor feature level"
                )
            return [int(tokens) for tokens in num_input_tokens]

        if self.num_feature_levels > 1:
            raise ValueError(
                "multi-scale DAbstractorAdapter requires num_input_tokens per level"
            )

        return [int(num_input_tokens)]

    @staticmethod
    def _init_reference_points(num_query_tokens, num_input_tokens):
        query_hw = int(math.sqrt(num_query_tokens))
        memory_hw = query_hw

        if isinstance(num_input_tokens, (list, tuple)):
            for tokens in num_input_tokens:
                inferred_hw = int(math.sqrt(tokens))
                if inferred_hw * inferred_hw == tokens:
                    memory_hw = max(memory_hw, inferred_hw)
        elif num_input_tokens is not None:
            inferred_hw = int(math.sqrt(num_input_tokens))
            if inferred_hw * inferred_hw == num_input_tokens:
                memory_hw = inferred_hw

        ref_y = torch.linspace(0.5, memory_hw - 0.5, query_hw, dtype=torch.float32)
        ref_x = torch.linspace(0.5, memory_hw - 0.5, query_hw, dtype=torch.float32)
        ref_y, ref_x = torch.meshgrid(ref_y, ref_x, indexing="ij")
        ref_y = ref_y.reshape(-1) / memory_hw
        ref_x = ref_x.reshape(-1) / memory_hw
        return torch.stack((ref_x, ref_y), dim=-1).unsqueeze(0)

    def _pos_emb_for_shape(self, h, w, level=0):
        if self.pos_emb is None:
            return None

        pos_emb = self.pos_emb[level] if isinstance(self.pos_emb, nn.ParameterList) else self.pos_emb

        if pos_emb.size(1) == h * w:
            return pos_emb

        src_hw = int(math.sqrt(pos_emb.size(1)))
        if src_hw * src_hw != pos_emb.size(1):
            raise ValueError("Cannot interpolate non-square D-Abstractor positional embeddings")

        pos = pos_emb.transpose(1, 2).reshape(1, self.d_model, src_hw, src_hw)
        pos = F.interpolate(pos, size=(h, w), mode="bicubic", align_corners=False)
        return pos.flatten(2).transpose(1, 2)

    @staticmethod
    def _prepare_deformable_inputs(batch_size, spatial_hw, dtype, device):
        valid_ratios = torch.ones(batch_size, len(spatial_hw), 2, dtype=dtype, device=device)
        spatial_shapes = torch.tensor(spatial_hw, dtype=torch.long, device=device)
        level_lengths = spatial_shapes[:, 0] * spatial_shapes[:, 1]
        level_start_index = torch.cat(
            [
                torch.zeros(1, dtype=torch.long, device=device),
                level_lengths.cumsum(0)[:-1],
            ],
            dim=0,
        )
        return valid_ratios, spatial_shapes, level_start_index

    def _make_pooled_queries(self, visual_tokens, h, w):
        batch_size = visual_tokens.size(0)
        query_embeds = self.query_position_embeddings.weight

        if self.pooled_v_target == "none":
            query_embed, target = torch.split(query_embeds, self.d_model, dim=1)
            query_embed = query_embed.unsqueeze(0).expand(batch_size, -1, -1)
            target = target.unsqueeze(0).expand(batch_size, -1, -1)
            return query_embed, target

        pooled = visual_tokens.transpose(1, 2).reshape(batch_size, self.d_model, h, w)
        pooled = self.downsampler(pooled).flatten(2).transpose(1, 2)
        learned_queries = query_embeds.unsqueeze(0).expand(batch_size, -1, -1)

        if self.pooled_v_target == "tgt":
            return learned_queries, pooled
        return pooled, learned_queries

    def forward(self, x):
        feature_maps = list(x) if isinstance(x, (list, tuple)) else [x]
        if len(feature_maps) != self.num_feature_levels:
            raise ValueError(
                f"Expected {self.num_feature_levels} D-Abstractor feature levels, "
                f"got {len(feature_maps)}"
            )

        original_dtype = feature_maps[-1].dtype
        autocast_context = (
            torch.cuda.amp.autocast(enabled=False)
            if self.force_fp32 and feature_maps[-1].is_cuda
            else nullcontext()
        )

        with autocast_context:
            batch_size = feature_maps[0].size(0)
            visual_levels = []
            spatial_hw = []
            query_source = None
            query_h, query_w = None, None

            for level, feature_map in enumerate(feature_maps):
                _, _, h, w = feature_map.shape
                spatial_hw.append((int(h), int(w)))

                visual_tokens = feature_map.flatten(2).transpose(1, 2)
                if self.force_fp32:
                    visual_tokens = visual_tokens.float()
                visual_tokens = self.input_proj[level](visual_tokens)

                pos_emb = self._pos_emb_for_shape(h, w, level)
                if pos_emb is not None:
                    visual_tokens = visual_tokens + pos_emb.to(
                        dtype=visual_tokens.dtype, device=visual_tokens.device
                    )
                if self.level_emb is not None:
                    visual_tokens = visual_tokens + self.level_emb[:, level].to(
                        dtype=visual_tokens.dtype, device=visual_tokens.device
                    )

                visual_levels.append(visual_tokens)
                query_source = visual_tokens
                query_h, query_w = h, w

            visual_tokens = torch.cat(visual_levels, dim=1)
            query_embed, target = self._make_pooled_queries(query_source, query_h, query_w)
            reference_points = self.reference_points.to(
                dtype=visual_tokens.dtype, device=visual_tokens.device
            )
            if self.clamp_reference_points:
                reference_points = reference_points.clamp(
                    self.reference_point_eps,
                    1.0 - self.reference_point_eps,
                )
            reference_points = reference_points.expand(batch_size, -1, -1)
            valid_ratios, spatial_shapes, level_start_index = self._prepare_deformable_inputs(
                batch_size,
                spatial_hw,
                visual_tokens.dtype,
                visual_tokens.device,
            )
            spatial_shapes_list = spatial_hw

            hidden_states = target
            for decoder_layer in self.layers:
                reference_points_input = reference_points[:, :, None] * valid_ratios[:, None]
                layer_kwargs = {
                    "position_embeddings": query_embed,
                    "encoder_hidden_states": visual_tokens,
                    "reference_points": reference_points_input,
                    "spatial_shapes": spatial_shapes,
                    "level_start_index": level_start_index,
                    "encoder_attention_mask": None,
                    "output_attentions": False,
                }
                if self._uses_spatial_shapes_list:
                    layer_kwargs["spatial_shapes_list"] = spatial_shapes_list
                layer_outputs = decoder_layer(hidden_states, **layer_kwargs)
                hidden_states = layer_outputs[0]

            hidden_states = self.output_proj(hidden_states)
            if self.output_norm is not None:
                hidden_states = self.output_norm(hidden_states)
                hidden_states = hidden_states * self.output_scale
            if self.output_clip is not None:
                hidden_states = hidden_states.clamp(-self.output_clip, self.output_clip)
            if self.sanitize_output:
                hidden_states = torch.nan_to_num(
                    hidden_states,
                    nan=0.0,
                    posinf=self.output_clip if self.output_clip is not None else 1.0,
                    neginf=-(self.output_clip if self.output_clip is not None else 1.0),
                )
        return hidden_states.to(dtype=original_dtype), (query_h, query_w)




class CoordConv(nn.Module):
    def __init__(self,
                 in_channels,
                 out_channels,
                 kernel_size=3,
                 padding=1,
                 stride=1):
        super().__init__()
        self.conv1 = conv_layer(in_channels + 2, out_channels, kernel_size,
                                padding, stride)

    def add_coord(self, input):
        b, _, h, w = input.size()
        x_range = torch.linspace(-1, 1, w, device=input.device)
        y_range = torch.linspace(-1, 1, h, device=input.device)
        y, x = torch.meshgrid(y_range, x_range)
        y = y.expand([b, 1, -1, -1])
        x = x.expand([b, 1, -1, -1])
        coord_feat = torch.cat([x, y], 1)
        input = torch.cat([input, coord_feat], 1)
        return input

    def forward(self, x):
        x = self.add_coord(x)
        x = self.conv1(x)
        return x


class MultiTaskProjector(nn.Module):
    def __init__(self, word_dim=1024, in_dim=256, kernel_size=3):
        super().__init__()
        self.in_dim = in_dim
        self.kernel_size = kernel_size
        # visual projector
        # 视觉特征的reshape 最终输出的是 (B, 256*5, 104, 104)，表示 5 个任务的 feature map（5×256）
        self.vis = nn.Sequential(  # os16 -> os4
            nn.Upsample(scale_factor=2, mode='bilinear'),          # (26x26) → (52x52)  
            conv_layer(in_dim * 2, in_dim * 2, 3, padding=1),      # 512 → 512
            nn.Upsample(scale_factor=2, mode='bilinear'),          # (52x52) → (104x104)
            conv_layer(in_dim * 2, in_dim, 3, padding=1),          # 512 → 256
            nn.Conv2d(in_dim, in_dim*5, 1))                        # 256 → 256*5

        # textual projector
        out_dim = 1 * in_dim * kernel_size * kernel_size + 1
        self.txt = nn.Linear(word_dim, out_dim)

    def forward(self, x, word):
        '''
            x: b, 512, 26, 26
            word: b, 512
        '''
        x = self.vis(x)
        x = torch.tensor_split(x, 5, dim=1) # no tensor_split api in torch 1.7, please use it in higher version
        # x = torch.chunk(x, 5, dim=1)

        mask_x = x[0]
        grasp_qua_x = x[1]
        grasp_sin_x = x[2]
        grasp_cos_x = x[3]
        grasp_wid_x = x[4]

        B, C, H, W = mask_x.size()


        # 1, b*256, 104, 104
        mask_x = mask_x.reshape(1, B * C, H, W)
        grasp_qua_x = grasp_qua_x.reshape(1, B * C, H, W)
        grasp_sin_x = grasp_sin_x.reshape(1, B * C, H, W)
        grasp_cos_x = grasp_cos_x.reshape(1, B * C, H, W)
        grasp_wid_x = grasp_wid_x.reshape(1, B * C, H, W)


        # txt: b, (256*3*3 + 1) -> b, 256, 3, 3 / b
        word = self.txt(word)
        weight, bias = word[:, :-1], word[:, -1]
        weight = weight.reshape(B, C, self.kernel_size, self.kernel_size)
        # Conv2d - 1, b*256, 104, 104 -> 1, b, 104, 104
        mask_out = F.conv2d(mask_x,
                       weight,
                       padding=self.kernel_size // 2,
                       groups=weight.size(0),
                       bias=bias)
        
        grasp_qua_out = F.conv2d(grasp_qua_x,
                            weight,
                            padding=self.kernel_size // 2,
                            groups=weight.size(0),
                            bias=bias)
        
        grasp_sin_out = F.conv2d(grasp_sin_x,
                            weight,
                            padding=self.kernel_size // 2,
                            groups=weight.size(0),
                            bias=bias)

        grasp_cos_out = F.conv2d(grasp_cos_x,
                            weight,
                            padding=self.kernel_size // 2,
                            groups=weight.size(0),
                            bias=bias)
        
        grasp_wid_out = F.conv2d(grasp_wid_x,
                            weight,
                            padding=self.kernel_size // 2,
                            groups=weight.size(0),
                            bias=bias)
            
        mask_out = mask_out.transpose(0, 1)
        grasp_qua_out = grasp_qua_out.transpose(0, 1)
        grasp_sin_out = grasp_sin_out.transpose(0, 1)
        grasp_cos_out = grasp_cos_out.transpose(0, 1)
        grasp_wid_out = grasp_wid_out.transpose(0, 1)
        # b, 1, 104, 104

        return mask_out, grasp_qua_out, grasp_sin_out, grasp_cos_out, grasp_wid_out


class Projector(nn.Module):
    def __init__(self, word_dim=1024, in_dim=256, kernel_size=3):
        super().__init__()
        self.in_dim = in_dim
        self.kernel_size = kernel_size
        # visual projector
        self.vis = nn.Sequential(  # os16 -> os4
            nn.Upsample(scale_factor=2, mode='bilinear'),           # 26→52
            conv_layer(in_dim * 2, in_dim * 2, 3, padding=1),       # 卷积层 → 提取局部模式
            nn.Upsample(scale_factor=2, mode='bilinear'),           # 52→104
            conv_layer(in_dim * 2, in_dim, 3, padding=1),           # 通道压缩 + 特征压缩
            nn.Conv2d(in_dim, in_dim, 1))                           # 最后一个线性变换
        
        # textual projector
        # 将文本向量 word: [B, 1024] 投影为卷积核权重和偏置 输出 B x (256×3×3 + 1) 的向量
        out_dim = 1 * in_dim * kernel_size * kernel_size + 1
        self.txt = nn.Linear(word_dim, out_dim)

    def forward(self, x, word):
        '''
            x: b, 512, 26, 26
            word: b, 512
        '''
        x = self.vis(x)         # → shape: [B, 256, 104, 104]
        B, C, H, W = x.size()  
        # 1, b*256, 104, 104
        x = x.reshape(1, B * C, H, W)    # → [1, B×256, 104, 104] 是为了支持后面 groups=B 的逐样本动态卷积
        # txt: b, (256*3*3 + 1) -> b, 256, 3, 3 / b
        word = self.txt(word)   # [B, 2305]
        weight, bias = word[:, :-1], word[:, -1]
        weight = weight.reshape(B, C, self.kernel_size, self.kernel_size)   # [B, 256, 3, 3]
        # Conv2d - 1, b*256, 104, 104 -> 1, b, 104, 104
        out = F.conv2d(x,
                       weight,
                       padding=self.kernel_size // 2,
                       groups=weight.size(0),
                       bias=bias)
        out = out.transpose(0, 1)
        # b, 1, 104, 104
        return out


class TransformerDecoder(nn.Module):
    def __init__(self,
                 num_layers,
                 d_model,
                 nhead,
                 dim_ffn,
                 dropout,
                 return_intermediate=False):
        super().__init__()
        self.layers = nn.ModuleList([
            TransformerDecoderLayer(d_model=d_model,
                                    nhead=nhead,
                                    dim_feedforward=dim_ffn,
                                    dropout=dropout) for _ in range(num_layers)
        ])
        self.num_layers = num_layers
        self.norm = nn.LayerNorm(d_model)
        # 是否返回中间层输出
        self.return_intermediate = return_intermediate

    @staticmethod
    def pos1d(d_model, length):
        """
        :param d_model: dimension of the model
        :param length: length of positions
        :return: length*d_model position matrix
        位置编码
        """
        if d_model % 2 != 0:
            raise ValueError("Cannot use sin/cos positional encoding with "
                             "odd dim (got dim={:d})".format(d_model))
        pe = torch.zeros(length, d_model)
        position = torch.arange(0, length).unsqueeze(1)
        div_term = torch.exp((torch.arange(0, d_model, 2, dtype=torch.float) *
                              -(math.log(10000.0) / d_model)))
        pe[:, 0::2] = torch.sin(position.float() * div_term)
        pe[:, 1::2] = torch.cos(position.float() * div_term)

        return pe.unsqueeze(1)  # n, 1, 512

    @staticmethod
    def pos2d(d_model, height, width):
        """
        :param d_model: dimension of the model
        :param height: height of the positions
        :param width: width of the positions
        :return: d_model*height*width position matrix
        """
        if d_model % 4 != 0:
            raise ValueError("Cannot use sin/cos positional encoding with "
                             "odd dimension (got dim={:d})".format(d_model))
        pe = torch.zeros(d_model, height, width)
        # Each dimension use half of d_model
        d_model = int(d_model / 2)
        div_term = torch.exp(
            torch.arange(0., d_model, 2) * -(math.log(10000.0) / d_model))
        pos_w = torch.arange(0., width).unsqueeze(1)
        pos_h = torch.arange(0., height).unsqueeze(1)
        pe[0:d_model:2, :, :] = torch.sin(pos_w * div_term).transpose(
            0, 1).unsqueeze(1).repeat(1, height, 1)
        pe[1:d_model:2, :, :] = torch.cos(pos_w * div_term).transpose(
            0, 1).unsqueeze(1).repeat(1, height, 1)
        pe[d_model::2, :, :] = torch.sin(pos_h * div_term).transpose(
            0, 1).unsqueeze(2).repeat(1, 1, width)
        pe[d_model + 1::2, :, :] = torch.cos(pos_h * div_term).transpose(
            0, 1).unsqueeze(2).repeat(1, 1, width)

        return pe.reshape(-1, 1, height * width).permute(2, 1, 0)  # hw, 1, 512

    def forward(self, vis, txt, pad_mask):
        '''
            vis: b, 512, h, w
            txt: b, L, 512
            pad_mask: b, L
        '''
        B, C, H, W = vis.size()
        _, L, D = txt.size()
        # position encoding
        vis_pos = self.pos2d(C, H, W)
        txt_pos = self.pos1d(D, L)
        # reshape & permute
        vis = vis.reshape(B, C, -1).permute(2, 0, 1)
        txt = txt.permute(1, 0, 2)
        # forward
        output = vis
        intermediate = []
        for layer in self.layers:
            output = layer(output, txt, vis_pos, txt_pos, pad_mask)
            if self.return_intermediate:
                # HW, b, 512 -> b, 512, HW
                intermediate.append(self.norm(output).permute(1, 2, 0))

        if self.norm is not None:
            # HW, b, 512 -> b, 512, HW
            output = self.norm(output).permute(1, 2, 0)
            if self.return_intermediate:
                intermediate.pop()
                intermediate.append(output)
                # [output1, output2, ..., output_n]
                return intermediate
            else:
                # b, 512, HW
                return output
        return output


class TransformerDecoderLayer(nn.Module):
    def __init__(self,
                 d_model=512,
                 nhead=9,
                 dim_feedforward=2048,
                 dropout=0.1):
        super().__init__()
        # Normalization Layer
        self.self_attn_norm = nn.LayerNorm(d_model)
        self.cross_attn_norm = nn.LayerNorm(d_model)
        # Attention Layer
        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
        self.multihead_attn = nn.MultiheadAttention(d_model,
                                                    nhead,
                                                    dropout=dropout,
                                                    kdim=d_model,
                                                    vdim=d_model)
        # FFN
        self.ffn = nn.Sequential(nn.Linear(d_model, dim_feedforward),
                                 nn.ReLU(True), nn.Dropout(dropout),
                                 nn.LayerNorm(dim_feedforward),
                                 nn.Linear(dim_feedforward, d_model))
        # LayerNorm & Dropout
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
        self.dropout3 = nn.Dropout(dropout)

    def with_pos_embed(self, tensor, pos):
        return tensor if pos is None else tensor + pos.to(tensor.device)

    def forward(self, vis, txt, vis_pos, txt_pos, pad_mask):
        '''
            vis: 26*26, b, 512
            txt: L, b, 512
            vis_pos: 26*26, 1, 512
            txt_pos: L, 1, 512
            pad_mask: b, L
        '''
        # Self-Attention
        vis2 = self.norm1(vis)
        q = k = self.with_pos_embed(vis2, vis_pos)
        vis2 = self.self_attn(q, k, value=vis2)[0]
        vis2 = self.self_attn_norm(vis2)
        vis = vis + self.dropout1(vis2)
        # Cross-Attention
        vis2 = self.norm2(vis)
        vis2 = self.multihead_attn(query=self.with_pos_embed(vis2, vis_pos),
                                   key=self.with_pos_embed(txt, txt_pos),
                                   value=txt,
                                   key_padding_mask=pad_mask)[0]
        vis2 = self.cross_attn_norm(vis2)
        vis = vis + self.dropout2(vis2)
        # FFN
        vis2 = self.norm3(vis)
        vis2 = self.ffn(vis2)
        vis = vis + self.dropout3(vis2)
        return vis


class FPN(nn.Module):
    def __init__(self,
                 in_channels=[512, 1024, 1024],
                 out_channels=[256, 512, 1024]):
        super(FPN, self).__init__()
        # text projection 添加一个可学习的投影层
        self.txt_proj = linear_layer(in_channels[2], out_channels[2])
        # fusion 1: v5 & seq -> f_5: b, 1024, 13, 13
        self.f1_v_proj = conv_layer(in_channels[2], out_channels[2], 1, 0)
        self.norm_layer = nn.Sequential(nn.BatchNorm2d(out_channels[2]),
                                        nn.ReLU(True))
        # fusion 2: v4 & fm -> f_4: b, 512, 26, 26
        self.f2_v_proj = conv_layer(in_channels[1], out_channels[1], 3, 1)
        self.f2_cat = conv_layer(out_channels[2] + out_channels[1],
                                 out_channels[1], 1, 0)
        # fusion 3: v3 & fm_mid -> f_3: b, 512, 52, 52
        self.f3_v_proj = conv_layer(in_channels[0], out_channels[0], 3, 1)
        self.f3_cat = conv_layer(out_channels[0] + out_channels[1],
                                 out_channels[1], 1, 0)
        # fusion 4: f_3 & f_4 & f_5 -> fq: b, 256, 26, 26
        self.f4_proj5 = conv_layer(out_channels[2], out_channels[1], 3, 1)
        self.f4_proj4 = conv_layer(out_channels[1], out_channels[1], 3, 1)
        self.f4_proj3 = conv_layer(out_channels[1], out_channels[1], 3, 1)
        # aggregation
        self.aggr = conv_layer(3 * out_channels[1], out_channels[1], 1, 0)
        self.coordconv = nn.Sequential(
            CoordConv(out_channels[1], out_channels[1], 3, 1),
            conv_layer(out_channels[1], out_channels[1], 3, 1))

    def forward(self, imgs, state):
        # v3, v4, v5: 256, 52, 52 / 512, 26, 26 / 1024, 13, 13
        v3, v4, v5 = imgs
        # fusion 1: b, 1024, 13, 13
        # text projection: [B, 1024, 1, 1]
        state = self.txt_proj(state).unsqueeze(-1).unsqueeze(
            -1)  
        
        # Fusion 1（v5 和文本 state 融合）：
        f5 = self.f1_v_proj(v5)                 # Conv1x1 → [B, 1024, 13, 13]
        f5 = self.norm_layer(f5 * state)        # 广播乘文本，增强目标相关性    

        # fusion 2: v4 + 上采样的 f5
        f4 = self.f2_v_proj(v4)                 # Conv3x3 → [B, 512, 26, 26]
        f5_ = F.interpolate(f5, scale_factor=2, mode='bilinear')    # 上采样
        f4 = self.f2_cat(torch.cat([f4, f5_], dim=1))   # concat + Conv1x1 → [B, 512, 26, 26]


        # fusion 3: v3 + avg pool 的 f4
        f3 = self.f3_v_proj(v3)                 # Conv3x3 → [B, 256, 52, 52]
        f3 = F.avg_pool2d(f3, 2, 2)             # 降采样 → [B, 256, 26, 26]
        f3 = self.f3_cat(torch.cat([f3, f4], dim=1))    # concat + Conv1x1 → [B, 512, 26, 26]

        # fusion 4: 融合 f3/f4/f5
        fq5 = self.f4_proj5(f5)                 # → [B, 512, 13, 13]
        fq4 = self.f4_proj4(f4)                 # → [B, 512, 26, 26]
        fq3 = self.f4_proj3(f3)                 # → [B, 512, 26, 26]
        fq5 = F.interpolate(fq5, scale_factor=2, mode='bilinear')

        # query
        fq = torch.cat([fq3, fq4, fq5], dim=1)  # → [B, 1536, 26, 26]
        fq = self.aggr(fq)                      # Conv1x1 聚合通道 → [B, 512, 26, 26]
        fq = self.coordconv(fq)                 # CoordConv + Conv3x3 → [B, 512, 26, 26]
        # b, 512, 26, 26
        return fq

class GeoPriorGen(nn.Module):
    def __init__(self, embed_dim, num_heads, initial_value, heads_range):
        super().__init__()
        angle = 1.0 / (10000 ** torch.linspace(0, 1, embed_dim // num_heads // 2))
        angle = angle.unsqueeze(-1).repeat(1, 2).flatten()
        self.weight = nn.Parameter(torch.ones(2, 1, 1, 1), requires_grad=True)
        decay = torch.log(
            1 - 2 ** (-initial_value - heads_range * torch.arange(num_heads, dtype=torch.float) / num_heads)
        )
        self.register_buffer("angle", angle)
        self.register_buffer("decay", decay)

    def generate_depth_decay(self, H: int, W: int, depth_grid):
        """
        generate 2d decay mask, the result is (HW)*(HW)
        H, W are the numbers of patches at each column and row
        """
        B, _, H, W = depth_grid.shape
        grid_d = depth_grid.reshape(B, H * W, 1)
        mask_d = grid_d[:, :, None, :] - grid_d[:, None, :, :]
        mask_d = (mask_d.abs()).sum(dim=-1)
        mask_d = mask_d.unsqueeze(1) * self.decay[None, :, None, None]
        return mask_d

    def generate_pos_decay(self, H: int, W: int):
        """
        generate 2d decay mask, the result is (HW)*(HW)
        H, W are the numbers of patches at each column and row
        """
        index_h = torch.arange(H).to(self.decay)
        index_w = torch.arange(W).to(self.decay)
        grid = torch.meshgrid([index_h, index_w])
        grid = torch.stack(grid, dim=-1).reshape(H * W, 2)
        mask = grid[:, None, :] - grid[None, :, :]
        mask = (mask.abs()).sum(dim=-1)
        mask = mask * self.decay[:, None, None]
        return mask

    def generate_1d_depth_decay(self, H, W, depth_grid):
        """
        generate 1d depth decay mask, the result is l*l
        """
        mask = depth_grid[:, :, :, :, None] - depth_grid[:, :, :, None, :]
        mask = mask.abs()
        mask = mask * self.decay[:, None, None, None]
        assert mask.shape[2:] == (W, H, H)
        return mask

    def generate_1d_decay(self, l: int):
        """
        generate 1d decay mask, the result is l*l
        """
        index = torch.arange(l).to(self.decay)
        mask = index[:, None] - index[None, :]
        mask = mask.abs()
        mask = mask * self.decay[:, None, None]
        return mask

    def forward(self, HW_tuple: Tuple[int], depth_map, split_or_not=False):
        """
        depth_map: depth patches
        HW_tuple: (H, W)
        H * W == l
        """
        depth_map = F.interpolate(depth_map, size=HW_tuple, mode="bilinear", align_corners=False)

        if split_or_not:
            index = torch.arange(HW_tuple[0] * HW_tuple[1]).to(self.decay)
            sin = torch.sin(index[:, None] * self.angle[None, :])
            sin = sin.reshape(HW_tuple[0], HW_tuple[1], -1)
            cos = torch.cos(index[:, None] * self.angle[None, :])
            cos = cos.reshape(HW_tuple[0], HW_tuple[1], -1)

            mask_d_h = self.generate_1d_depth_decay(HW_tuple[0], HW_tuple[1], depth_map.transpose(-2, -1))
            mask_d_w = self.generate_1d_depth_decay(HW_tuple[1], HW_tuple[0], depth_map)

            mask_h = self.generate_1d_decay(HW_tuple[0])
            mask_w = self.generate_1d_decay(HW_tuple[1])

            mask_h = self.weight[0] * mask_h.unsqueeze(0).unsqueeze(2) + self.weight[1] * mask_d_h
            mask_w = self.weight[0] * mask_w.unsqueeze(0).unsqueeze(2) + self.weight[1] * mask_d_w

            geo_prior = ((sin, cos), (mask_h, mask_w))

        else:
            index = torch.arange(HW_tuple[0] * HW_tuple[1]).to(self.decay)
            sin = torch.sin(index[:, None] * self.angle[None, :])
            sin = sin.reshape(HW_tuple[0], HW_tuple[1], -1)
            cos = torch.cos(index[:, None] * self.angle[None, :])
            cos = cos.reshape(HW_tuple[0], HW_tuple[1], -1)
            mask = self.generate_pos_decay(HW_tuple[0], HW_tuple[1])

            mask_d = self.generate_depth_decay(HW_tuple[0], HW_tuple[1], depth_map)
            mask = self.weight[0] * mask + self.weight[1] * mask_d

            geo_prior = ((sin, cos), mask)

        return geo_prior
    
    
class DWConv2d(nn.Module):
    def __init__(self, dim, kernel_size, stride, padding):
        super().__init__()
        self.dwconv = nn.Conv2d(dim, dim, kernel_size, stride, padding, groups=dim)

    def forward(self, x: torch.Tensor):
        """
        input (b h w c)
        """
        x = x.permute(0, 3, 1, 2)
        x = self.dwconv(x)
        x = x.permute(0, 2, 3, 1)
        return x
    
def angle_transform(x, sin, cos):
    x1 = x[:, :, :, :, ::2]
    x2 = x[:, :, :, :, 1::2]
    #print(x1.shape, x2.shape, sin.shape, cos.shape)
    return (x * cos) + (torch.stack([-x2, x1], dim=-1).flatten(-2) * sin)

class CrossGSA(nn.Module):
    def __init__(self, embed_dim, num_heads, value_factor=1):
        super().__init__()
        self.factor = value_factor
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = self.embed_dim * self.factor // num_heads
        self.key_dim = self.embed_dim // num_heads
        self.scaling = self.key_dim**-0.5

        self.q_proj = nn.Linear(embed_dim, embed_dim)
        self.k_proj = nn.Linear(embed_dim, embed_dim)
        self.v_proj = nn.Linear(embed_dim, embed_dim * self.factor)

        self.lepe = DWConv2d(embed_dim, 5, 1, 2)
        self.out_proj = nn.Linear(embed_dim * self.factor, embed_dim)
        
        self.layer_norm = nn.LayerNorm(self.embed_dim, eps=1e-6)
        self.reset_parameters()

    def forward(self, x: torch.Tensor, y: torch.Tensor=None, rel_pos=None):
        """
        x: image feature [b h w c]
        y: depth feature [b h w c]
        rel_pos: (sin, cos), mask from GeoPriorGen
        """
        bsz, h, w, _ = x.size()
        q = self.q_proj(x) 
        if y:          
            k = self.k_proj(y)      
            v = self.v_proj(y)  
        else:
            k = self.k_proj(x)
            v = self.v_proj(x)            
        lepe = self.lepe(v)

        k = k * self.scaling

        qr = q.view(bsz, h, w, self.num_heads, -1).permute(0, 3, 1, 2, 4)
        kr = k.view(bsz, h, w, self.num_heads, -1).permute(0, 3, 1, 2, 4)
        vr = v.reshape(bsz, h, w, self.num_heads, -1).permute(0, 3, 1, 2, 4)

        # Apply angle transform if provided
        if rel_pos is not None:
            sin, cos = rel_pos[0]
            qr = angle_transform(qr, sin, cos)
            kr = angle_transform(kr, sin, cos)

        qr = qr.flatten(2, 3)
        kr = kr.flatten(2, 3)
        vr = vr.flatten(2, 3)

        qk_mat = qr @ kr.transpose(-1, -2)
        

        if rel_pos is not None:
            _, mask = rel_pos
            qk_mat = qk_mat + mask
        qk_mat = torch.softmax(qk_mat, -1)
        output = torch.matmul(qk_mat, vr)

        output = output.transpose(1, 2).reshape(bsz, h, w, -1)
        output = output + lepe
        output = x + self.out_proj(output)
        
        output = output + self.layer_norm(output)
        
        return output

    def reset_parameters(self):
        nn.init.xavier_normal_(self.q_proj.weight, gain=2**-2.5)
        nn.init.xavier_normal_(self.k_proj.weight, gain=2**-2.5)
        nn.init.xavier_normal_(self.v_proj.weight, gain=2**-2.5)
        nn.init.xavier_normal_(self.out_proj.weight)
        nn.init.constant_(self.out_proj.bias, 0.0)

class DepthCoordConv(nn.Module):
    """CoordConv with one additional normalized depth channel."""

    def __init__(
        self,
        in_channels,
        out_channels,
        kernel_size=3,
        padding=1,
        stride=1,
        normalize_depth=True,
        eps=1e-6,
    ):
        super().__init__()
        self.normalize_depth = normalize_depth
        self.eps = eps
        self._depth = None
        self.conv1 = conv_layer(
            in_channels + 3,
            out_channels,
            kernel_size,
            padding,
            stride,
        )

    def set_depth(self, depth):
        self._depth = depth

    def _resize_depth(self, depth, x):
        b, _, h, w = x.shape
        if depth is None:
            return x.new_zeros((b, 1, h, w))
        if depth.dim() == 3:
            depth = depth.unsqueeze(1)
        depth = depth.to(device=x.device, dtype=x.dtype)
        if depth.shape[-2:] != (h, w):
            depth = F.interpolate(
                depth,
                size=(h, w),
                mode="bilinear",
                align_corners=False,
            )
        depth = torch.nan_to_num(depth, nan=0.0, posinf=1.0, neginf=0.0)
        if self.normalize_depth:
            valid = depth > 0
            valid_min = depth.masked_fill(~valid, float("inf")).amin(
                dim=(2, 3),
                keepdim=True,
            )
            valid_max = depth.masked_fill(~valid, float("-inf")).amax(
                dim=(2, 3),
                keepdim=True,
            )
            has_valid = valid.any(dim=(2, 3), keepdim=True)
            depth_min = torch.where(has_valid, valid_min, torch.zeros_like(valid_min))
            depth_max = torch.where(has_valid, valid_max, torch.ones_like(valid_max))
            depth = (depth - depth_min) / (depth_max - depth_min + self.eps)
            depth = torch.where(valid, depth, torch.zeros_like(depth))
        return depth.clamp(0.0, 1.0)

    def add_coord(self, x, depth=None):
        b, _, h, w = x.shape
        x_range = torch.linspace(-1, 1, w, device=x.device, dtype=x.dtype)
        y_range = torch.linspace(-1, 1, h, device=x.device, dtype=x.dtype)
        y_grid, x_grid = torch.meshgrid(y_range, x_range)
        x_grid = x_grid.expand(b, 1, h, w)
        y_grid = y_grid.expand(b, 1, h, w)
        depth = self._resize_depth(depth, x)
        return torch.cat([x, x_grid, y_grid, depth], dim=1)

    def forward(self, x, depth=None):
        if depth is None:
            depth = self._depth
        x = self.add_coord(x, depth)
        return self.conv1(x)