pxdepth/model/CM_PiT.py

"""Context-Modulated Pixel Transformer (CM-PiT) building blocks.

CM-PiT compresses local dense pixel features into attention tokens, processes
them with gated self-attention and SwiGLU, and expands them back without losing
the original pixel lattice. Global encoder tokens generate adaptive shift,
scale, and residual gates that condition both transformer sublayers.
"""

from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange

from .Gated_Attention import GatedAttention
from .RoPE import RotaryPositionEmbedding2D
from .precision import full_precision


def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
    """Apply affine context modulation without changing the tensor layout.

    Args:
        x: Normalized pixel tokens with shape ``[B, N, P, C]``.
        shift: Context-predicted additive offsets with shape ``[B, N, P, C]``.
        scale: Context-predicted residual scales with shape ``[B, N, P, C]``.

    Returns:
        Modulated tokens ``x * (1 + scale) + shift`` with shape
        ``[B, N, P, C]``.
    """
    return x * (1.0 + scale) + shift


class SwiGLU(nn.Module):
    """SwiGLU feed-forward layer operating independently on every pixel token.

    The first projection creates value and gate branches, SiLU activates the
    gate, and the second projection returns to the pixel-channel dimension.
    Spatial and patch axes are preserved throughout the module.
    """

    def __init__(self, dim: int, hidden_dim: int) -> None:
        """Construct the gated feed-forward projections.

        Args:
            dim: Input and output channel count ``C``.
            hidden_dim: Width of each hidden value/gate branch.

        Returns:
            ``None``. Learnable linear layers are registered on the module.
        """
        super().__init__()
        self.fc1 = nn.Linear(dim, hidden_dim * 2)
        self.fc2 = nn.Linear(hidden_dim, dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Transform pixel tokens with a SiLU-gated hidden representation.

        Args:
            x: Floating tensor with arbitrary leading dimensions and final
                channel dimension ``C=dim``. CM-PiT supplies ``[B,N,P,C]``.

        Returns:
            Tensor with the same shape and dtype as ``x``.
        """
        value, gate = self.fc1(x).chunk(2, dim=-1)
        return self.fc2(value * F.silu(gate))


class ContextAdaNorm(nn.Module):
    """Predict Context-Guided Adaptive Normalization parameters.

    Each global context token produces shift, scale, and residual-gate values
    for both the attention and MLP sublayers over every pixel represented by
    that encoder token. The six parameter groups are unpacked by
    :class:`CMPiTBlock`.
    """

    def __init__(self, dim_ctx: int, patch_size: int, dim_pix: int) -> None:
        """Create the context-to-modulation projection.

        Args:
            dim_ctx: Channel count of each Global Context Encoder token.
            patch_size: Encoder patch side length ``P_ctx`` in image pixels.
            dim_pix: Pixel-feature channel count ``C_pix``.

        Returns:
            ``None``. The projection outputs ``6 * P_ctx^2 * C_pix`` values per
            context token.
        """
        super().__init__()
        self.proj = nn.Sequential(
            nn.SiLU(),
            nn.Linear(dim_ctx, 6 * patch_size * patch_size * dim_pix),
        )

    def forward(self, ctx: torch.Tensor) -> torch.Tensor:
        """Project context tokens into six dense pixel-wise parameter fields.

        Args:
            ctx: Context token tensor ``[B, N_ctx, C_ctx]``.

        Returns:
            Modulation tensor ``[B, N_ctx, 6 * P_ctx^2 * C_pix]``.
        """
        return self.proj(ctx)


class CMPiTBlock(nn.Module):
    """Context-Modulated Pixel Transformer block.

    The block groups a dense pixel feature map into local patches, linearly
    compresses every patch to an attention token, applies gated global
    self-attention, expands the token back to pixel features, and follows it
    with a per-pixel SwiGLU MLP. Both residual branches use Context-Guided
    Adaptive Normalization generated from DINO context tokens.
    """

    def __init__(
        self,
        dim_ctx: int,
        ctx_patch_size: int,
        dim_pix: int,
        patch_size: int,
        attn_dim: int,
        num_heads: int,
        mlp_ratio: float = 4.0,
        qk_norm: bool = True,
        rope: Optional[RotaryPositionEmbedding2D] = None,
        eps: float = 1e-6,
    ) -> None:
        """Configure one CM-PiT block.

        Args:
            dim_ctx: Context-token channel count ``C_ctx``.
            ctx_patch_size: Image patch size ``P_ctx`` represented by one
                context token.
            dim_pix: Dense pixel-feature channel count ``C_pix``.
            patch_size: Side length ``P`` grouped into one attention token.
                It must divide ``ctx_patch_size``.
            attn_dim: Compressed attention-token channel count ``D``.
            num_heads: Number of attention heads. ``D`` must be divisible by it.
            mlp_ratio: Expansion ratio controlling the SwiGLU hidden width.
            qk_norm: Whether to apply FP32 RMSNorm to each query/key head.
            rope: Optional 2D rotary position embedding shared by decoder blocks.
            eps: Numerical epsilon used by RMSNorm layers.

        Returns:
            ``None``. Attention, modulation, MLP, and projection layers are
            registered on the block.
        """
        super().__init__()
        if ctx_patch_size % patch_size != 0:
            raise ValueError(
                f"ctx_patch_size ({ctx_patch_size}) must be divisible by patch_size ({patch_size})"
            )

        self.dim_ctx = dim_ctx
        self.dim_pix = dim_pix
        self.ctx_patch_size = ctx_patch_size
        self.patch_size = patch_size
        patch_dim = patch_size * patch_size * dim_pix

        self.norm1 = nn.RMSNorm(dim_pix, eps=eps)
        self.linear_compress = nn.Linear(patch_dim, attn_dim)
        self.attn = GatedAttention(attn_dim, num_heads, qk_norm=qk_norm, rope=rope, eps=eps)
        self.linear_expand = nn.Linear(attn_dim, patch_dim)
        self.norm2 = nn.RMSNorm(dim_pix, eps=eps)
        hidden_dim = max(1, int(round(dim_pix * mlp_ratio * 2.0 / 3.0)))
        self.mlp = SwiGLU(dim_pix, hidden_dim)
        self.ada_norm = ContextAdaNorm(dim_ctx, ctx_patch_size, dim_pix)

    @staticmethod
    def _norm(norm: nn.Module, x: torch.Tensor) -> torch.Tensor:
        """Evaluate a normalization layer in FP32 and restore input dtype.

        Args:
            norm: Normalization module acting on the final channel dimension.
            x: Pixel tokens ``[B, N, P^2, C_pix]`` in the active model dtype.

        Returns:
            Normalized tensor with the same shape and dtype as ``x``.
        """
        dtype = x.dtype
        with full_precision(x.device):
            out = norm(x.float())
        return out.to(dtype)

    def _modulation(self, ctx: torch.Tensor, height: int, width: int) -> torch.Tensor:
        """Align context modulation fields with the block's pixel patches.

        Args:
            ctx: Global context tokens ``[B, H_ctx*W_ctx, C_ctx]``.
            height: Dense pixel-map height ``H``.
            width: Dense pixel-map width ``W``.

        Returns:
            Six modulation groups with shape
            ``[B, (H/P)*(W/P), 6, P^2, C_pix]``. Rearrangement is exact and
            contains no interpolation.
        """
        batch = ctx.shape[0]
        p_ctx, p = self.ctx_patch_size, self.patch_size
        ctx_h, ctx_w = height // p_ctx, width // p_ctx
        if ctx.shape[1] != ctx_h * ctx_w:
            raise ValueError(
                f"Context token count ({ctx.shape[1]}) does not match grid ({ctx_h}x{ctx_w})"
            )

        mod = self.ada_norm(ctx).view(batch, ctx_h, ctx_w, 6, p_ctx, p_ctx, self.dim_pix)
        if p_ctx == p:
            return rearrange(mod, "b h w m ph pw c -> b (h w) m (ph pw) c")

        ratio = p_ctx // p
        return rearrange(
            mod,
            "b h w m (rh ph) (rw pw) c -> b (h rh w rw) m (ph pw) c",
            rh=ratio,
            rw=ratio,
            ph=p,
            pw=p,
        )

    def forward(self, x: torch.Tensor, ctx: torch.Tensor, pos: torch.Tensor) -> torch.Tensor:
        """Apply context-modulated attention and MLP residual updates.

        Args:
            x: Dense pixel features ``[B, C_pix, H, W]``.
            ctx: Global context tokens ``[B, (H/P_ctx)*(W/P_ctx), C_ctx]``.
            pos: Integer 2D token positions ``[B, (H/P)*(W/P), 2]`` used by
                rotary position embedding in self-attention.

        Returns:
            Updated dense pixel features ``[B, C_pix, H, W]``.
        """
        batch, _, height, width = x.shape
        p = self.patch_size
        pix = rearrange(x, "b c (h ph) (w pw) -> b (h w) (ph pw) c", ph=p, pw=p)
        shift_attn, scale_attn, gate_attn, shift_mlp, scale_mlp, gate_mlp = self._modulation(
            ctx, height, width
        ).unbind(dim=2)

        out = modulate(self._norm(self.norm1, pix), shift_attn, scale_attn)
        out = self.linear_compress(out.flatten(2))
        out = self.attn(out, pos=pos)
        out = self.linear_expand(out).view(batch, -1, p * p, self.dim_pix)
        pix = pix + gate_attn * out

        out = modulate(self._norm(self.norm2, pix), shift_mlp, scale_mlp)
        pix = pix + gate_mlp * self.mlp(out)
        return rearrange(
            pix,
            "b (h w) (ph pw) c -> b c (h ph) (w pw)",
            h=height // p,
            w=width // p,
            ph=p,
            pw=p,
        )