"""DinoVolumeSceneV4DA3 — v4 head with Depth-Anything-3 backbone. Identical to ``DinoVolumeSceneV4`` (per-voxel MLP fusion + sin/cos height PE + adaptive per-view z range), but the DINOv3 backbone is swapped for DA3's multi-view ViT. We use DA3's built-in cross-view attention and the camera-token global descriptor — and we DELIBERATELY skip the DPT head because in past tests its gradients were poorly conditioned for our volumetric scoring objective. ────────────────────────────────────────────────────────────────────── Architecture deltas vs v4 ────────────────────────────────────────────────────────────────────── 1. **Backbone**: DA3 ViT in place of DINOv3 ViT-S/16+. DA3 is a multi-view ViT pretrained on monocular depth + camera-pose objectives. Its blocks alternate between local (within-view) attention and global (cross-view) attention starting from a configurable ``alt_start`` block index. The cross-view attention IS the cross-view fusion — no need for our own ``CrossViewAttnStack`` on top. 2. **Multi-view forward**: where v4 loops per view through DINOv3, v4_da3 stacks the views into ``(B, S=N_views, 3, H, W)`` and makes a single DA3 backbone call. Cross-view attention fires naturally at the right blocks. Per-view patches are split out of the backbone output for the downstream 1×1 conv + bilinear-up + 5-layer refine CNN (identical to v4). 3. **Global token**: DA3's camera token replaces DINOv3's CLS. The camera token is the geometric global descriptor DA3 was trained against (its own pose-decoder head consumes it), so it carries strictly more useful signal than DINOv3's SSL-CLS. 4. **No DPT head**: we tested it; gradients are noisy. We take the *raw* patch features at the last backbone layer (with ``cat_token=True`` so the per-pixel dim is ``2 × embed_dim`` = 768 for da3-small) and feed them through OUR 1×1 conv → bilinear-up → 5-layer refine CNN, exactly as v4 does on DINOv3. No DPT projects, no DPT resize layers, no DPT refinenets. 5. **Patch grid resolution**: DA3 uses patch_size=14, so at img_size=504 we get 36×36 patches (vs DINOv3's 31×31 at patch_size=16). The 5-layer refine CNN bilinear-ups both to ``pred_size`` so downstream code doesn't care. Everything else (per-voxel MLP fusion, sin/cos PE, adaptive per-view z, per-voxel grip/rot, voxel sampling, loss) is unchanged from v4. """ from __future__ import annotations import os import sys import types import torch import torch.nn as nn from .model_volumetric_v4 import DinoVolumeSceneV4 # ──────────────────────────────────────────────────────────────────────── # DA3 repo + weight setup (mirrors lib/model_da3.py) # ──────────────────────────────────────────────────────────────────────── DA3_REPO_DIR = os.environ.get( "DA3_REPO_DIR", os.path.expanduser("~/lab/da3_repo/src") if os.path.exists(os.path.expanduser("~/lab/da3_repo/src")) else "/data/cameron/da3_repo/src", ) DA3_WEIGHTS_DEFAULT = os.environ.get( "DA3_WEIGHTS_PATH", os.path.expanduser("~/lab/da3_weights") if os.path.exists(os.path.expanduser("~/lab/da3_weights")) else "/data/cameron/da3_weights", ) if DA3_REPO_DIR not in sys.path: sys.path.insert(0, DA3_REPO_DIR) # Stub optional utils we don't need at train/inference time. for _n in ("depth_anything_3.utils.export", "depth_anything_3.utils.pose_align"): if _n not in sys.modules: sys.modules[_n] = types.ModuleType(_n) sys.modules["depth_anything_3.utils.export"].export = lambda *a, **k: None sys.modules["depth_anything_3.utils.pose_align"].align_poses_umeyama = ( lambda *a, **k: None) sys.modules["depth_anything_3.utils.pose_align"].batch_align_poses_umeyama = ( lambda *a, **k: None) from depth_anything_3.api import DepthAnything3 # noqa: E402 class DinoVolumeSceneV4DA3(DinoVolumeSceneV4): """v4 head + DA3 backbone (no DPT, cam_token as global descriptor).""" def __init__(self, *args, da3_weights: str = DA3_WEIGHTS_DEFAULT, da3_patch_size: int = 14, **kwargs): # v4 init loads DINOv3 (which we'll immediately delete). We accept # this small wasted load because every other v4 attribute we need # (PE buffers, MLPs, head modules) is constructed by the parent. super().__init__(*args, **kwargs) # Drop DINOv3, load DA3. del self.backbone full = DepthAnything3.from_pretrained(da3_weights) self.backbone = full.model.backbone # cat_token=True default makes the per-pixel dim 2 × embed_dim. # da3-small: 384 × 2 = 768. embed_dim_eff = full.model.head.norm.normalized_shape[0] del full # drop cam_dec, gs head, DPT, etc. self.embed_dim = int(embed_dim_eff) self.da3_patch_size = int(da3_patch_size) # Rebuild modules that depended on the old DINOv3 embed_dim (384): # feat_head : embed_dim → feat_dim (1×1 conv) # input_proj: includes cls fan-in = N · embed_dim (when "concat") feat_dim = self.feat_dim d_model = self.d_model self.feat_head = nn.Conv2d(self.embed_dim, feat_dim, 1) cls_fan = (self.n_views * self.embed_dim if self.cls_fusion == "concat" else self.embed_dim) past_fan = self.past_enc_dim if self.past_n > 0 else 0 self.input_proj = nn.Linear(feat_dim + cls_fan + past_fan, d_model) # If the parent built a CrossViewAttnStack on top of DINOv3 patches, # rebuild it for the larger embed_dim. (Usually disabled for v4_da3 # since DA3 already does cross-view in the backbone.) if self.cross_view_attn is not None: from .model_volumetric_v4 import CrossViewAttnStack n_heads = self.cross_view_attn.blocks[0].attn.num_heads self.cross_view_attn = CrossViewAttnStack( n_layers=self.cross_view_layers, embed_dim=self.embed_dim, n_heads=n_heads, ) # ──────────────────────────────────────────────────────────────────── # Multi-view backbone call # ──────────────────────────────────────────────────────────────────── def _per_view_patch_features(self, rgb: list[torch.Tensor]): """Stack views, run DA3 backbone with cross-view attention enabled, split per-view patches + cam_tokens. Backbone runs in bf16 autocast. Returns: view_patches: list of N tensors, each (B, embed_dim, ph, pw) view_global: list of N tensors, each (B, embed_dim) — cam_token per view """ assert len(rgb) == self.n_views # Stack to (B, S=N, 3, H, W). x = torch.stack(rgb, dim=1) # (B, N, 3, H, W) B, N, _, H, W = x.shape ac_dtype = (torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16) with torch.autocast(device_type=x.device.type, dtype=ac_dtype): feats_tuples, _ = self.backbone( x, cam_token=None, export_feat_layers=[], ref_view_strategy="saddle_balanced", ) last = feats_tuples[-1] # ((B, N, n_patches, embed), (B, N, embed)) patches_all = last[0].float() # (B, N, n_patches, embed_dim) cam_tokens = last[1].float() # (B, N, embed_dim) ph = H // self.da3_patch_size pw = W // self.da3_patch_size view_patches: list[torch.Tensor] = [] view_global: list[torch.Tensor] = [] for k in range(N): p = patches_all[:, k] # (B, n_patches, embed_dim) n_patches = p.shape[1] # DA3 patches may include leading non-image tokens (registers # / cam token) — when the count exceeds ph·pw, drop the prefix. if n_patches > ph * pw: p = p[:, n_patches - ph * pw:] elif n_patches != ph * pw: # Backbone reshape mismatch — defensive fall-back to a # square grid inferred from the token count. side = int(round(n_patches ** 0.5)) ph_eff = pw_eff = side view_patches.append( p.permute(0, 2, 1).reshape(B, self.embed_dim, ph_eff, pw_eff)) view_global.append(cam_tokens[:, k]) continue view_patches.append( p.permute(0, 2, 1).reshape(B, self.embed_dim, ph, pw)) view_global.append(cam_tokens[:, k]) return view_patches, view_global # patch_features is the legacy single-view entry point inherited from # _DinoHiresBase via v4. Override it to fail loudly if anyone calls it # on this subclass — they should use _per_view_patch_features instead. def patch_features(self, rgb: torch.Tensor): raise NotImplementedError( "DA3 backbone needs multi-view input for cross-attn. " "Use ``_per_view_patch_features(rgb_list)`` instead.")