Spaces:
Running on Zero
Running on Zero
File size: 11,238 Bytes
87608ea | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 | """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
|