src/mavt/model/backbone.py

"""Stage 2: Hybrid Transformer-RGAT4D Backbone.

Block layout (12 blocks):
  0-3:  StandardTransformerBlock  (SigLIP2 init)
  4:    RGAT4DBlock               (zero-init output)
  5-7:  StandardTransformerBlock  (SigLIP2 init)
  8:    RGAT4DBlock               (zero-init output)
  9-11: StandardTransformerBlock  (SigLIP2 init)

Adjacency masks are precomputed once per (modality, resolution) and cached.
"""

from __future__ import annotations
from typing import Dict, List, Optional, Tuple

import torch
import torch.nn as nn
import torch.utils.checkpoint as cp

from mavt.model.transformer import StandardTransformerBlock
from mavt.model.rgat import RGAT4DBlock, build_adjacency


RGAT_POSITIONS = {4, 8}   # which block indices are RGAT4D


class HybridBackbone(nn.Module):
    """12-block hybrid Transformer-RGAT backbone."""

    def __init__(
        self,
        dim: int = 1152,
        num_heads: int = 16,
        num_blocks: int = 12,
        mlp_ratio: float = 4.0,
        dropout: float = 0.0,
        r_s: int = 2,
        r_t: int = 1,
        use_gradient_checkpointing: bool = False,
    ):
        super().__init__()
        self.r_s = r_s
        self.r_t = r_t
        self.use_gradient_checkpointing = use_gradient_checkpointing

        self.blocks = nn.ModuleList()
        for i in range(num_blocks):
            if i in RGAT_POSITIONS:
                self.blocks.append(RGAT4DBlock(dim, num_heads,
                                               mlp_ratio=mlp_ratio, dropout=dropout))
            else:
                self.blocks.append(StandardTransformerBlock(dim, num_heads,
                                                             mlp_ratio=mlp_ratio, dropout=dropout))

        # Mask cache: key = (modality, N) → (adj_mask, edge_type_masks)
        self._mask_cache: Dict[Tuple[str, int], Tuple[torch.Tensor, List[torch.Tensor]]] = {}

    # ------------------------------------------------------------------ #

    def _get_masks(
        self,
        positions: torch.Tensor,  # (N, 4)
        plane_ids: torch.Tensor,  # (N,)
        modality: str,
    ) -> Tuple[torch.Tensor, List[torch.Tensor]]:
        key = (modality, positions.shape[0])
        if key not in self._mask_cache:
            adj, etype_masks = build_adjacency(positions, plane_ids, modality,
                                               self.r_s, self.r_t)
            self._mask_cache[key] = (adj, etype_masks)
        return self._mask_cache[key]

    # ------------------------------------------------------------------ #

    def _run_transformer(self, block: StandardTransformerBlock, x: torch.Tensor) -> torch.Tensor:
        if self.use_gradient_checkpointing and self.training:
            return cp.checkpoint(block, x, use_reentrant=False)
        return block(x)

    def _run_rgat(
        self,
        block: RGAT4DBlock,
        x: torch.Tensor,
        adj_mask: torch.Tensor,
        edge_type_masks: List[torch.Tensor],
    ) -> torch.Tensor:
        if self.use_gradient_checkpointing and self.training:
            return cp.checkpoint(block, x, adj_mask, edge_type_masks, use_reentrant=False)
        return block(x, adj_mask, edge_type_masks)

    # ------------------------------------------------------------------ #

    def forward(
        self,
        x: torch.Tensor,          # (B, N, D)
        positions: torch.Tensor,  # (N, 4)
        plane_ids: torch.Tensor,  # (N,)
        modality: str,
    ) -> torch.Tensor:
        adj_mask, edge_type_masks = self._get_masks(positions, plane_ids, modality)
        # Move cached masks to current device if needed
        adj_mask = adj_mask.to(x.device)
        edge_type_masks = [m.to(x.device) for m in edge_type_masks]

        for i, block in enumerate(self.blocks):
            if i in RGAT_POSITIONS:
                x = self._run_rgat(block, x, adj_mask, edge_type_masks)
            else:
                x = self._run_transformer(block, x)
        return x

    # ------------------------------------------------------------------ #
    #  SigLIP2 weight loading utility                                     #
    # ------------------------------------------------------------------ #

    def load_siglip2_weights(self, model_name: str = "google/siglip2-base-patch16-224",
                              freeze_stages: int = 0) -> None:
        """Load SigLIP2 backbone weights into Transformer blocks (best-effort).

        freeze_stages: number of initial Transformer blocks to freeze (stage 1: all,
        stage 2: leave last 4 unfrozen, stage 3: none frozen).
        """
        try:
            from transformers import AutoModel
            import re
            siglip = AutoModel.from_pretrained(model_name)
            siglip_blocks = siglip.vision_model.encoder.layers

            transformer_idx = 0  # index into siglip_blocks
            for block_idx, block in enumerate(self.blocks):
                if block_idx in RGAT_POSITIONS:
                    continue
                if transformer_idx >= len(siglip_blocks):
                    break
                src = siglip_blocks[transformer_idx]
                _copy_siglip2_block(src, block)
                transformer_idx += 1

            # Freeze early blocks
            frozen = 0
            for block_idx, block in enumerate(self.blocks):
                if block_idx in RGAT_POSITIONS:
                    continue
                if frozen < freeze_stages:
                    for p in block.parameters():
                        p.requires_grad_(False)
                    frozen += 1

        except Exception as exc:  # noqa: BLE001
            print(f"[backbone] SigLIP2 weight loading skipped: {exc}")


def _copy_siglip2_block(src: nn.Module, dst: StandardTransformerBlock) -> None:
    """Best-effort copy from a SigLIP2 encoder layer to our StandardTransformerBlock."""
    state = dst.state_dict()
    # SigLIP2 uses self_attn.{q,k,v,out}_proj; we use fused qkv + out_proj
    try:
        Q = src.self_attn.q_proj.weight.data
        K = src.self_attn.k_proj.weight.data
        V = src.self_attn.v_proj.weight.data
        state['qkv.weight'] = torch.cat([Q, K, V], dim=0)
        if src.self_attn.q_proj.bias is not None:
            Qb = src.self_attn.q_proj.bias.data
            Kb = src.self_attn.k_proj.bias.data
            Vb = src.self_attn.v_proj.bias.data
            state['qkv.bias'] = torch.cat([Qb, Kb, Vb], dim=0)
        state['out_proj.weight'] = src.self_attn.out_proj.weight.data
        if src.self_attn.out_proj.bias is not None:
            state['out_proj.bias'] = src.self_attn.out_proj.bias.data
        # LayerNorm
        state['norm1.weight'] = src.layer_norm1.weight.data
        state['norm1.bias']   = src.layer_norm1.bias.data
        state['norm2.weight'] = src.layer_norm2.weight.data
        state['norm2.bias']   = src.layer_norm2.bias.data
        # MLP
        state['mlp.0.weight'] = src.mlp.fc1.weight.data
        state['mlp.0.bias']   = src.mlp.fc1.bias.data
        state['mlp.3.weight'] = src.mlp.fc2.weight.data
        state['mlp.3.bias']   = src.mlp.fc2.bias.data
        dst.load_state_dict(state)
    except (AttributeError, RuntimeError):
        pass  # dimension mismatch or different naming — skip silently