pxdepth/model/Pixel_Space_Depth_Predictor.py

"""Pixel-Space Depth Predictor that preserves the dense image lattice.

Normalized RGB is embedded with a 1x1 projection and processed by shared
CM-PiT trunk blocks before branching into depth and finite-mask predictors.
Linear patch compression is used only within transformer blocks, after which
features are expanded back to per-pixel tokens for dense output heads.
"""

from typing import Iterable, Optional, Tuple

import torch
import torch.nn as nn
from torch.utils.checkpoint import checkpoint

from ..registry import PREDICTORS
from .CM_PiT import CMPiTBlock
from .RoPE import PositionGetter, RotaryPositionEmbedding2D
from .precision import full_precision, reduced_precision


@PREDICTORS.register()
class PixelSpaceDepthPredictor(nn.Module):
    """Pixel-Space Depth Predictor built from cascaded CM-PiT blocks.

    A 1x1 projection first embeds normalized RGB into dense pixel features.
    Shared trunk blocks refine those features, after which independent depth
    and validity branches predict normalized log-depth and finite-depth logits.
    No convolution larger than 1x1 is applied to the pixel representation.
    """

    def __init__(
        self,
        in_channels: int = 3,
        dim_ctx: int = 1024,
        attn_dim: int = 1536,
        ctx_patch_size: int = 14,
        dim_pix: int = 16,
        trunk_patch_size: int = 14,
        depth_patch_size: int = 7,
        mask_patch_size: int = 14,
        num_heads: int = 24,
        trunk_depth: int = 4,
        depth_depth: int = 4,
        mask_depth: int = 2,
        mlp_ratio: float = 4.0,
        qk_norm: bool = True,
        rope_frequency: float = 100.0,
        eps: float = 1e-6,
        gradient_checkpointing: bool = True,
    ) -> None:
        """Construct the shared trunk and two prediction branches.

        Args:
            in_channels: Number of image channels, equal to three for RGB.
            dim_ctx: Global context-token channel count ``C_ctx``.
            attn_dim: Channel count ``D`` after linear patch compression.
            ctx_patch_size: Encoder patch size ``P_ctx`` in image pixels.
            dim_pix: Channel count ``C_pix`` of the dense pixel feature map.
            trunk_patch_size: Attention patch size used by shared trunk blocks.
            depth_patch_size: Attention patch size used by depth blocks.
            mask_patch_size: Attention patch size used by validity-mask blocks.
            num_heads: Number of gated-attention heads.
            trunk_depth: Number of shared CM-PiT blocks.
            depth_depth: Number of depth-branch CM-PiT blocks.
            mask_depth: Number of validity-branch CM-PiT blocks.
            mlp_ratio: SwiGLU expansion ratio inside every block.
            qk_norm: Enable FP32 RMSNorm for attention queries and keys.
            rope_frequency: Base frequency of the shared 2D RoPE module.
            eps: Numerical epsilon for normalization layers.
            gradient_checkpointing: Recompute CM-PiT blocks during backward to
                reduce activation memory.

        Returns:
            ``None``. The complete pixel predictor is registered on the module.
        """
        super().__init__()
        if attn_dim % num_heads != 0:
            raise ValueError(f"attn_dim ({attn_dim}) must be divisible by num_heads ({num_heads})")
        for name, patch_size in {
            "trunk_patch_size": trunk_patch_size,
            "depth_patch_size": depth_patch_size,
            "mask_patch_size": mask_patch_size,
        }.items():
            if patch_size <= 0 or ctx_patch_size % patch_size != 0:
                raise ValueError(
                    f"{name} ({patch_size}) must be positive and divide ctx_patch_size ({ctx_patch_size})"
                )

        self.dim_ctx = dim_ctx
        self.dim_pix = dim_pix
        self.attn_dim = attn_dim
        self.ctx_patch_size = ctx_patch_size
        self.trunk_patch_size = trunk_patch_size
        self.depth_patch_size = depth_patch_size
        self.mask_patch_size = mask_patch_size
        self.gradient_checkpointing = gradient_checkpointing

        self.pos = PositionGetter()
        self.rope = RotaryPositionEmbedding2D(frequency=rope_frequency)
        self.input_proj = nn.Conv2d(in_channels, dim_pix, kernel_size=1, bias=True)

        block_args = dict(
            dim_ctx=dim_ctx,
            ctx_patch_size=ctx_patch_size,
            dim_pix=dim_pix,
            attn_dim=attn_dim,
            num_heads=num_heads,
            mlp_ratio=mlp_ratio,
            qk_norm=qk_norm,
            rope=self.rope,
            eps=eps,
        )
        self.trunk_blocks = nn.ModuleList(
            CMPiTBlock(patch_size=trunk_patch_size, **block_args) for _ in range(trunk_depth)
        )
        self.depth_blocks = nn.ModuleList(
            CMPiTBlock(patch_size=depth_patch_size, **block_args) for _ in range(depth_depth)
        )
        self.mask_blocks = nn.ModuleList(
            CMPiTBlock(patch_size=mask_patch_size, **block_args) for _ in range(mask_depth)
        )
        self.depth_head = nn.Conv2d(dim_pix, 1, kernel_size=1, bias=True)
        self.mask_head = nn.Conv2d(dim_pix, 1, kernel_size=1, bias=True)
        self.reset_parameters()

    def _blocks(self) -> Iterable[CMPiTBlock]:
        """Iterate over every CM-PiT block in execution-independent order.

        Returns:
            Iterable containing shared trunk, depth, and mask blocks. The method
            takes no tensor inputs and is used for parameter initialization.
        """
        return (*self.trunk_blocks, *self.depth_blocks, *self.mask_blocks)

    def reset_parameters(self) -> None:
        """Initialize projections and start adaptive modulation at identity.

        Linear and 1x1 convolution weights use Xavier uniform initialization.
        Normalization scales start at one. The final adaptive-normalization
        projections are zeroed so every CM-PiT residual branch initially has
        zero modulation and zero gate.

        Returns:
            ``None``. Parameters are modified in place.
        """
        def init(module: nn.Module) -> None:
            """Initialize one child module visited by :meth:`nn.Module.apply`.

            Args:
                module: Child ``nn.Module`` to initialize in place.

            Returns:
                ``None``.
            """
            if isinstance(module, (nn.Linear, nn.Conv2d)):
                nn.init.xavier_uniform_(module.weight)
                if module.bias is not None:
                    nn.init.zeros_(module.bias)
            elif isinstance(module, (nn.LayerNorm, nn.RMSNorm)):
                if module.weight is not None:
                    nn.init.ones_(module.weight)
                if getattr(module, "bias", None) is not None:
                    nn.init.zeros_(module.bias)

        self.apply(init)
        for block in self._blocks():
            nn.init.zeros_(block.ada_norm.proj[-1].weight)
            nn.init.zeros_(block.ada_norm.proj[-1].bias)

    def enable_gradient_checkpointing(self) -> None:
        """Enable activation recomputation for CM-PiT blocks.

        The flag is consulted only while the module is in training mode.

        Returns:
            ``None``. The runtime flag is changed in place.
        """
        self.gradient_checkpointing = True

    def disable_gradient_checkpointing(self) -> None:
        """Disable activation recomputation for CM-PiT blocks.

        Subsequent training forwards retain block activations for backward.

        Returns:
            ``None``. The runtime flag is changed in place.
        """
        self.gradient_checkpointing = False

    def _position(self, batch: int, height: int, width: int, patch_size: int, device: torch.device):
        """Create cached 2D coordinates for one decoder patch grid.

        Args:
            batch: Batch size ``B``.
            height: Dense image-feature height ``H``.
            width: Dense image-feature width ``W``.
            patch_size: Block patch side length ``P``.
            device: Device on which coordinates are allocated.

        Returns:
            Integer position tensor ``[B, (H/P)*(W/P), 2]``.
        """
        return self.pos(batch, height // patch_size, width // patch_size, device=device).to(device)

    def _run(
        self,
        x: torch.Tensor,
        blocks: nn.ModuleList,
        ctx: torch.Tensor,
        pos: torch.Tensor,
    ) -> torch.Tensor:
        """Run a sequence of CM-PiT blocks with optional checkpointing.

        Args:
            x: Dense pixel features ``[B, C_pix, H, W]``.
            blocks: Ordered CM-PiT block collection for one branch.
            ctx: Global context tokens ``[B, N_ctx, C_ctx]``.
            pos: 2D positions ``[B, N, 2]`` matching the blocks' patch grid.

        Returns:
            Refined dense features ``[B, C_pix, H, W]``.
        """
        for block in blocks:
            if self.training and self.gradient_checkpointing:
                x = checkpoint(block, x, ctx, pos, use_reentrant=False)
            else:
                x = block(x, ctx, pos)
        return x

    def forward(
        self,
        image: torch.Tensor,
        ctx: torch.Tensor,
        autocast_dtype: Optional[torch.dtype] = torch.bfloat16,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Predict normalized log-depth and finite-depth logits from RGB.

        Args:
            image: ImageNet-normalized RGB tensor ``[B, 3, H, W]``.
            ctx: Global context tokens ``[B, (H/P_ctx)*(W/P_ctx), C_ctx]``.
            autocast_dtype: Decoder attention dtype. Use ``None`` for FP32,
                ``torch.float16`` for FP16, or ``torch.bfloat16`` for BF16.

        Returns:
            depth: Raw normalized log-depth tensor ``[B, 1, H, W]``.
            mask: Raw finite-depth logit tensor ``[B, 1, H, W]``.
        """
        batch, _, height, width = image.shape
        p_ctx = self.ctx_patch_size
        if height % p_ctx != 0 or width % p_ctx != 0:
            raise ValueError(f"Input resolution ({height}, {width}) must be divisible by {p_ctx}")
        expected = (height // p_ctx) * (width // p_ctx)
        if tuple(ctx.shape) != (batch, expected, self.dim_ctx):
            raise ValueError(
                f"Context shape {tuple(ctx.shape)} does not match ({batch}, {expected}, {self.dim_ctx})"
            )

        with full_precision(image.device):
            pix = self.input_proj(image.float())

        with reduced_precision(image.device, autocast_dtype):
            trunk_pos = self._position(batch, height, width, self.trunk_patch_size, image.device)
            pix = self._run(pix, self.trunk_blocks, ctx, trunk_pos)

            depth_pos = self._position(batch, height, width, self.depth_patch_size, image.device)
            depth_feat = self._run(pix, self.depth_blocks, ctx, depth_pos)

            mask_pos = self._position(batch, height, width, self.mask_patch_size, image.device)
            mask_feat = self._run(pix, self.mask_blocks, ctx, mask_pos)

        with full_precision(image.device):
            depth = self.depth_head(depth_feat.float())
            mask = self.mask_head(mask_feat.float())
        return depth, mask