jolia_multimodal_msa.py

# Vendored verbatim from the internal `raidium.rd.models` library for the
# self-contained Hugging Face release. Only imports were rewritten (raidium
# hub base classes -> jolia_shim; sibling modules -> jolia_* names).
# Do not edit by hand: regenerate with scripts/build_hf_jolia.py.

import math
from typing import List, Optional, Tuple, Union

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.checkpoint as checkpoint
from einops import rearrange
from timm.models.layers import DropPath, LayerNorm2d
from timm.models.vision_transformer import Mlp


class CrossWindowAttention(nn.Module):
    """Cross-window attention where queries come from a separate input."""

    def __init__(
        self,
        dim: int,
        num_heads: int = 8,
        qkv_bias: bool = False,
        qk_scale: Optional[float] = None,
        attn_drop: float = 0.0,
        proj_drop: float = 0.0,
    ):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.scale = qk_scale or self.head_dim**-0.5

        # Separate Q projection for query input
        self.q = nn.Linear(dim, dim, bias=qkv_bias)
        # KV projection for context input
        self.kv = nn.Linear(dim, dim * 2, bias=qkv_bias)

        # Output projection
        self.proj = nn.Linear(dim, dim)

        # Dropouts
        self.attn_drop = attn_drop
        self.proj_drop = nn.Dropout(proj_drop)

    def forward(self, x_q: torch.Tensor, x_kv: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor:
        """
        Args:
            x_q: Query input tensor (B, Nq, C)
            x_kv: Key-value input tensor (B, Nkv, C)
            mask: Optional attention mask
        """
        B, Nq, C = x_q.shape
        _, Nkv, _ = x_kv.shape

        # Generate Q from x_q
        q = self.q(x_q).reshape(B, Nq, self.num_heads, self.head_dim).permute(0, 2, 1, 3)

        # Generate K,V from x_kv
        kv = self.kv(x_kv).reshape(B, Nkv, 2, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
        k, v = kv.unbind(0)  # Each shape: (B, num_heads, Nkv, head_dim)

        x = F.scaled_dot_product_attention(q, k, v, dropout_p=self.attn_drop)

        x = x.transpose(1, 2).reshape(B, Nq, C)
        x = self.proj_drop(self.proj(x))
        return x

    def run_attn(self, q, k, v, mask=None):
        B, H, Nq, D = q.shape
        C = H * D
        x = F.scaled_dot_product_attention(q, k, v, dropout_p=self.attn_drop)
        x = x.transpose(1, 2).reshape(B, Nq, C)
        x = self.proj_drop(self.proj(x))
        return x

    def get_qkv(self, x_q, x_kv):
        B, Nq, C = x_q.shape
        _, Nkv, _ = x_kv.shape
        q = self.q(x_q).reshape(B, Nq, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
        kv = self.kv(x_kv).reshape(B, Nkv, 2, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
        k, v = kv.unbind(0)
        return q, k, v

    def get_q(self, x):
        B, Nq, C = x.shape
        q = self.q(x).reshape(B, Nq, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
        return q

    def get_kv(self, x):
        B, Nkv, C = x.shape
        kv = self.kv(x).reshape(B, Nkv, 2, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
        k, v = kv.unbind(0)
        return [k, v]


class CrossWindowBlock(nn.Module):
    """Transformer block with cross-window attention and MLP."""

    def __init__(
        self,
        dim: int,
        num_heads: int = 8,
        mlp_ratio: float = 4.0,
        qkv_bias: bool = False,
        qk_scale: Optional[float] = None,
        drop: float = 0.0,
        attn_drop: float = 0.0,
        drop_path: float = 0.0,
        act_layer: nn.Module = nn.GELU,
        norm_layer: nn.Module = nn.LayerNorm,
    ):
        super().__init__()
        # Cross window attention
        self.norm1_q = norm_layer(dim)
        self.norm1_kv = norm_layer(dim)
        self.attn = CrossWindowAttention(
            dim=dim,
            num_heads=num_heads,
            qkv_bias=qkv_bias,
            qk_scale=qk_scale,
            attn_drop=attn_drop,
            proj_drop=drop,
        )

        # MLP
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(
            in_features=dim,
            hidden_features=mlp_hidden_dim,
            act_layer=act_layer,
            drop=drop,
        )

        # Drop path
        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()

    def forward(self, x_q: torch.Tensor, x_kv: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor:
        """
        Args:
            x_q: Query input tensor
            x_kv: Key-value input tensor
            mask: Optional attention mask
        """
        # Cross window attention with residual
        x = x_q + self.drop_path(self.attn(self.norm1_q(x_q), self.norm1_kv(x_kv), mask))

        # MLP with residual
        x = x + self.drop_path(self.mlp(self.norm2(x)))

        return x

    def get_qkv(self, x_q, x_kv=None):
        if x_kv is None:
            x_kv = x_q
        x_q = self.norm1_q(x_q)
        x_kv = self.norm1_kv(x_kv)
        q, k, v = self.attn.get_qkv(x_q, x_kv)
        return q, k, v

    def get_qkv_tokens(self, x, key="q"):
        if key == "q":
            return self.attn.get_q(self.norm1_q(x))
        if key == "kv":
            return self.attn.get_kv(self.norm1_kv(x))

    def xattn_qkv(self, q, k, v, mask=None):
        x = self.attn.run_attn(q, k, v, mask)
        return x

    def mlp_residual(self, x):
        x = x + self.drop_path(self.mlp(self.norm2(x)))
        return x

    def skip_with_drop(self, x, skip):
        x = x + self.drop_path(skip)
        return x


class RelativePosEmb(nn.Module):
    """
    Learnable relative positional embedding for 3D grids, supporting linear or conv projections.
    Coordinate tables are pre-computed at init time as registered buffers to avoid graph breaks.
    """

    def __init__(
        self,
        dim: int,
        rank: int = 2,
        conv: bool = False,
        modality_to_grid_size: dict[str, list[int]] | None = None,
    ):
        super().__init__()
        self.rank = rank
        self.conv = conv
        if not conv:
            self.cpb_mlp = nn.Sequential(
                nn.Linear(rank, 512, bias=True),
                nn.ReLU(),
                nn.Linear(512, dim, bias=False),
            )
        else:
            self.cpb_mlp = nn.Sequential(
                nn.Conv1d(rank, 512, 1, bias=True),
                nn.ReLU(),
                nn.Conv1d(512, dim, 1, bias=False),
            )

        if modality_to_grid_size:
            for modality, grid_size in modality_to_grid_size.items():
                self.register_buffer(
                    f"_coord_{modality}",
                    self._build_coord_table(grid_size, conv=conv),
                )

    @staticmethod
    def _build_coord_table(grid_size: list[int], conv: bool = False) -> torch.Tensor:
        h, w, d = grid_size
        table = (
            torch.stack(
                torch.meshgrid(
                    torch.arange(h, dtype=torch.float32),
                    torch.arange(w, dtype=torch.float32),
                    torch.arange(d, dtype=torch.float32),
                    indexing="ij",
                )
            )
            .contiguous()
            .unsqueeze(0)
        )  # [1, 3, h, w, d]
        if h > 1:
            table[0, 0] -= h // 2
            table[0, 0] /= h // 2
        if w > 1:
            table[0, 1] -= w // 2
            table[0, 1] /= w // 2
        if d > 1:
            table[0, 2] -= d // 2
            table[0, 2] /= d // 2

        if not conv:
            return table.permute(0, 2, 3, 4, 1).reshape(1, h * w * d, 3)
        return table.squeeze(0).reshape(3, -1)

    def forward(
        self,
        x: torch.Tensor,
        grid_size: list[int] | None = None,
        modality: str = "default",
    ) -> torch.Tensor:
        coord_table = getattr(self, f"_coord_{modality}", None)
        if coord_table is None:
            coord_table = self._build_coord_table(grid_size, conv=self.conv).to(x.device)
        return x + self.cpb_mlp(coord_table)


class MultiScaleAttentionBlock(nn.Module):
    """
    MultiScaleAttentionBlock: Implements multi-scale attention with various communication protocols.
    """

    def __init__(
        self,
        dim,
        num_heads,
        mlp_ratio=4.0,
        qkv_bias=True,
        qk_norm=False,
        drop=0.0,
        attn_drop=0.0,
        init_values=None,
        drop_path=0.0,
        act_layer=nn.GELU,
        norm_layer=nn.LayerNorm,
        pool_op="max",
        window_dims=4,
        weight_share=True,
        ignore_registers=False,
        accumulate_window_summary=True,
        num_scales=None,
        local2global_per_modality=None,
        posemb_grid_sizes=None,
        **kwargs,
    ):
        super().__init__()
        self._posemb_grid_sizes = posemb_grid_sizes
        self._init_basic_config(
            dim,
            num_heads,
            drop,
            attn_drop,
            qkv_bias,
            mlp_ratio,
            drop_path,
            window_dims,
            init_values,
            norm_layer,
            weight_share,
            num_scales,
            local2global_per_modality,
        )

        self._init_multiscale_attention()
        self._init_multiscale_position_embeddings()

    def _init_basic_config(
        self,
        dim,
        num_heads,
        drop,
        attn_drop,
        qkv_bias,
        mlp_ratio,
        drop_path,
        window_dims,
        init_values,
        norm_layer,
        weight_share,
        num_scales,
        local2global_per_modality=None,
    ):
        """Initialize basic configuration parameters."""
        self.dim = dim
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.scale = self.head_dim**0.5
        self.window_dims = window_dims
        self.init_values = init_values
        self.norm_layer = norm_layer
        self.mlp_ratio = mlp_ratio
        self.drop_path = drop_path

        # Dropout configurations
        self.attn_drop_p = attn_drop
        self.drop = drop
        self.proj_drop = nn.Dropout(drop)

        # Component configurations
        self.qkv_bias = qkv_bias
        self.additional_scale = None
        self.communication_protocol = "all2all_sattn__sequential"
        self.aggregation_protocol = "one2one_xattn"
        self.num_scales = num_scales
        self.weight_share = weight_share

        # Pre-create bottom-up pools per modality
        if local2global_per_modality:
            self.bottom_up_pools = nn.ModuleDict(
                {mod: nn.MaxPool3d(kernel_size=l2g) for mod, l2g in local2global_per_modality.items()}
            )

    def _init_multiscale_attention(self):
        """Initialize multiscale attention components, with one x-attn block per window."""
        self.blocks = nn.ModuleList(
            [
                CrossWindowBlock(
                    dim=self.dim,
                    num_heads=self.num_heads,
                    mlp_ratio=self.mlp_ratio,
                    qkv_bias=self.qkv_bias,
                    drop=self.drop,
                    attn_drop=self.attn_drop_p,
                    drop_path=self.drop_path,
                    norm_layer=self.norm_layer,
                )
                for scale_idx in range(self.num_scales)
            ]
        )

    def _init_multiscale_position_embeddings(self):
        """Initialize position embeddings with pre-computed coordinate tables."""
        self.posemb = nn.ModuleList(
            [
                RelativePosEmb(
                    self.dim,
                    rank=3,
                    modality_to_grid_size=self._posemb_grid_sizes[scale_idx] if self._posemb_grid_sizes else None,
                )
                for scale_idx in range(self.num_scales)
            ]
        )

    def propagate_bottom_up(
        self,
        stages: List[torch.Tensor],
        grid_sizes: List[Tuple[int, int]],
        merge_ratio: list,
        modality: str,
    ) -> List[torch.Tensor]:
        """
        Propagate information from local to global representations in bottom-up pass.
        """
        downscaling_op = self.bottom_up_pools[modality]

        for i in range(len(stages) - 1):
            current_stage = stages[i]
            current_grid_size = grid_sizes[i]
            nw = math.prod(current_grid_size)

            current_stage = rearrange(
                current_stage,
                "bnw (m0 m1 m2) c -> bnw c m0 m1 m2",
                m0=merge_ratio[0],
                m1=merge_ratio[1],
                m2=merge_ratio[2],
            )
            current_stage = downscaling_op(current_stage)

            current_stage = rearrange(current_stage, "(b nw) c m0 m1 m2 -> b nw m0 m1 m2 c", nw=nw)
            current_stage = rearrange(
                current_stage,
                "b (d h w) m0 m1 m2 c -> b (d m0) (h m1) (w m2) c",
                h=current_grid_size[0],
                w=current_grid_size[1],
                d=current_grid_size[2],
            )

            d, h, w = current_stage.shape[1:4]
            if d == merge_ratio[0] and h == merge_ratio[1] and w == merge_ratio[2]:
                propagated = rearrange(current_stage, "b d h w c -> b (d h w) c")
            elif d >= merge_ratio[0] and h >= merge_ratio[1] and w >= merge_ratio[2]:
                propagated = rearrange(
                    current_stage,
                    "b (d m0) (h m1) (w m2) c -> (b d h w) (m0 m1 m2) c",
                    m0=merge_ratio[0],
                    m1=merge_ratio[1],
                    m2=merge_ratio[2],
                )
            else:
                propagated = rearrange(current_stage, "b d h w c -> b (d h w) c")

            stages[i + 1] = stages[i + 1] + propagated

        return stages

    def forward(
        self,
        scales,
        grid_sizes=None,
        multiscale_layout=None,
        merge_ratio=None,
        local2global=None,
        modality: str = "chest_cxr_single_view",
    ):
        if "sequential" in self.communication_protocol:
            return self.forward_sequential(
                scales,
                grid_sizes,
                multiscale_layout,
                merge_ratio,
                local2global,
                modality=modality,
            )
        else:
            raise NotImplementedError

    def forward_sequential(
        self,
        scales: List[torch.Tensor],
        grid_sizes: Optional[List[Tuple[int, int]]] = None,
        multiscale_layout=None,
        merge_ratio=None,
        local2global=None,
        modality: str = "chest_xray_single_view",
    ) -> List[torch.Tensor]:
        """
        Implements communication protocol for sequential processing of scales.
        """
        num_scales = len(scales)

        for idx in range(num_scales):
            scales[idx] = self.posemb[idx](
                scales[idx],
                grid_size=multiscale_layout[idx]["window_dims"],
                modality=modality,
            )

        scales = self.propagate_bottom_up(scales, grid_sizes, merge_ratio, modality)

        # List-based out_scales (no dict mutation)
        out_scales: list[Optional[torch.Tensor]] = [None] * num_scales

        # Top-down: message passing from higher to lower level scales
        for S in range(num_scales - 1, -1, -1):
            out_scales[S] = self._process_all2all_sattn(scales[S], S, out_scales)

        # Bottom-up: aggregate from lower to higher level scales
        for S in range(1, num_scales):
            out_scales[S] = self._aggregate_one2one_xattn(S, out_scales, multiscale_layout)

        return out_scales

    def _aggregate_one2one_xattn(
        self,
        S: int,
        out_scales: list[Optional[torch.Tensor]],
        multiscale_layout=None,
    ) -> torch.Tensor:
        """Aggregate cross-attention from scale S-1 into scale S."""
        x_S = out_scales[S]
        x_Sm1 = out_scales[S - 1]

        q_S = self.blocks[S].get_qkv_tokens(x_S, "q")
        k_Sm1, v_Sm1 = self.blocks[S - 1].get_qkv_tokens(x_Sm1, "kv")

        kH, kW, kD = multiscale_layout[S]["grid_size"]
        mH, mW, mD = multiscale_layout[S]["window_dims"]
        q_S = rearrange(
            q_S,
            "(b kD kH kW) h (mD mH mW) c -> b h (kD mD) (kH mH) (kW mW) c",
            kD=kD,
            kH=kH,
            kW=kW,
            mD=mD,
            mH=mH,
            mW=mW,
        )
        mH, mW, mD = multiscale_layout[S]["window_dims"]
        sH, sW, sD = multiscale_layout[S - 1]["grid_size"]

        q_S = rearrange(
            q_S,
            "b h (sD mD) (sH mH) (sW mW) c -> (b sD sH sW) h mD mH mW c",
            sD=sD,
            sH=sH,
            sW=sW,
        )

        m0, m1, m2 = q_S.shape[2:5]
        q_S = rearrange(q_S, "b h m0 m1 m2 c -> b h (m0 m1 m2) c", m0=m0, m1=m1, m2=m2)

        xattn_l2g = self.blocks[S].xattn_qkv(q_S, k_Sm1, v_Sm1)
        xattn_l2g = rearrange(
            xattn_l2g,
            "(b sD sH sW) (m0 m1 m2) c -> b (sD m0) (sH m1) (sW m2) c",
            sD=sD,
            sH=sH,
            sW=sW,
            m0=m0,
            m1=m1,
            m2=m2,
        )
        xattn_l2g = rearrange(
            xattn_l2g,
            "b (kD m0) (kH m1) (kW m2) c -> (b kD kH kW) (m0 m1 m2) c",
            kD=kD,
            kH=kH,
            kW=kW,
        )

        x_S = self.blocks[S].skip_with_drop(x_S, xattn_l2g)
        x_S = self.blocks[S].mlp_residual(x_S)

        return x_S

    def _process_all2all_sattn(
        self,
        x_S: torch.Tensor,
        S: int,
        out_scales: list[Optional[torch.Tensor]],
    ) -> torch.Tensor:
        """Process scale S with all-to-all self-attention across already-processed scales."""
        q_S, k_S, v_S = self.blocks[S].get_qkv(x_S)

        k_list, v_list = [k_S], [v_S]
        for T in range(len(out_scales)):
            if out_scales[T] is None:
                continue
            x_t = out_scales[T]
            num_repeats = x_S.shape[0] // x_t.shape[0]
            k_t, v_t = self.blocks[T].get_qkv_tokens(x_t, "kv")
            k_list.append(k_t.repeat_interleave(num_repeats, dim=0))
            v_list.append(v_t.repeat_interleave(num_repeats, dim=0))

        k_cat = torch.cat(k_list, dim=2)
        v_cat = torch.cat(v_list, dim=2)

        x_S = self.blocks[S].skip_with_drop(x_S, self.blocks[S].xattn_qkv(q_S, k_cat, v_cat))
        return self.blocks[S].mlp_residual(x_S)


class AtlasStage(nn.Module):
    """
    AtlasStage: A single stage of the AtlasMultiScale architecture that processes
    input features through multiple attention blocks with window-based operations.
    """

    def __init__(
        self,
        dim: int,
        depth: int,
        num_heads: int,
        mlp_ratio: float = 4.0,
        qkv_bias: bool = True,
        qk_scale: Optional[float] = None,
        drop: float = 0.0,
        attn_drop: float = 0.0,
        drop_path: Union[float, List[float]] = 0.0,
        num_scales=None,
        activation_checkpointing: bool = False,
        local2global_per_modality=None,
        posemb_grid_sizes=None,
        **kwargs,
    ):
        super().__init__()

        drop_path_rates = [0.0] if depth == 1 else [i * (drop_path / (depth - 1)) for i in range(depth)]
        self.activation_checkpointing = activation_checkpointing

        self.blocks = nn.ModuleList(
            [
                MultiScaleAttentionBlock(
                    dim=dim,
                    num_heads=num_heads,
                    mlp_ratio=mlp_ratio,
                    qkv_bias=qkv_bias,
                    drop=drop,
                    attn_drop=attn_drop,
                    drop_path=drop_path_rates[i],
                    weight_share=False,
                    num_scales=num_scales,
                    local2global_per_modality=local2global_per_modality,
                    posemb_grid_sizes=posemb_grid_sizes,
                )
                for i in range(depth)
            ]
        )

    def forward(
        self,
        x: torch.Tensor,
        grid_sizes: List[Tuple[int, int]],
        multiscale_layout=None,
        merge_ratio=None,
        local2global=None,
        modality: str = "chest_xray_single_view",
    ) -> torch.Tensor:
        """Forward pass for the Atlas Stages.

        Args:
            x: Input tensor
            grid_sizes: List of grid sizes for multi-scale processing

        Returns:
            Processed tensor after attention blocks
        """
        # Process through attention blocks
        for block in self.blocks:
            if self.activation_checkpointing and self.training:

                def _run_block(*scales):
                    out = block(
                        list(scales),
                        grid_sizes,
                        multiscale_layout,
                        merge_ratio,
                        local2global,
                        modality=modality,
                    )
                    return tuple(out)

                x = list(checkpoint.checkpoint(_run_block, *x, use_reentrant=False))
            else:
                x = block(
                    x,
                    grid_sizes,
                    multiscale_layout,
                    merge_ratio,
                    local2global,
                    modality=modality,
                )

        return x