"""DinoVolumeSceneV4Radio — v4 head with NVIDIA C-RADIOv2-B backbone. Same v4 head (per-voxel MLP fusion + sin/cos height PE + adaptive per-view z range), with the DINOv3 backbone replaced by **C-RADIOv2-B** from NVLabs (``nvidia/C-RADIOv2-B``). RADIO is a distillation of DINOv2, SigLIP, SAM, and DFN into a single ViT-B/16 backbone; the spatial features it emits are effectively a "multi-teacher fused DINO+CLIP+SAM" representation. ────────────────────────────────────────────────────────────────────── Architecture deltas vs v4 ────────────────────────────────────────────────────────────────────── 1. **Backbone**: ``nvidia/C-RADIOv2-B`` loaded via ``AutoModel.from_pretrained(..., trust_remote_code=True)``. Forward returns ``(summary, spatial_features)`` where ``summary`` is the global CLS-like descriptor (B, D_summary) and ``spatial_features`` is (B, N, D_spatial). RADIO accepts arbitrary input resolutions (both dims must be multiples of 16); we feed the same 504² PARA input so the resulting patch grid is 31×31 — matching DINO's grid exactly for fair comparison. 2. **Per-view loop, no backbone-level cross-view attention**. Same as the default DINOv3 path. 3. **Per-pixel features**: D_spatial → 1×1 conv (feat_head) → bilinear-up to ``pred_size`` → 5-layer refine CNN (standard v4 pipeline; no DPT, no extra projection). 4. **Global descriptor**: RADIO's summary token (typically D_summary = D_spatial); flows into ``input_proj`` as the per-view CLS. 5. **Image size + normalisation**: RADIO ships its own input conditioner that internally denormalises → renormalises with its own stats, so we feed ImageNet-normalised tensors directly (matches the PARA pipeline default). Input res: 504×504 (same as DINO). 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 math import torch import torch.nn as nn import torch.nn.functional as F from .model_volumetric_v4 import DinoVolumeSceneV4 class DinoVolumeSceneV4Radio(DinoVolumeSceneV4): """v4 head + NVIDIA C-RADIOv2-B backbone (single-view loop).""" def __init__(self, *args, radio_model_name: str = "nvidia/C-RADIOv2-B", radio_input_size: int = 496, # 31*16 → 31² patches, matches DINO **kwargs): super().__init__(*args, **kwargs) # v4 already built the DINOv3 backbone; drop it before loading. del self.backbone from transformers import AutoModel print(f"Loading RADIO backbone: {radio_model_name}") self.radio = AutoModel.from_pretrained( radio_model_name, trust_remote_code=True, ) # RADIO requires input dims to be multiples of 16 — the input # conditioner will pad-and-truncate otherwise. Round to nearest # multiple of 16 that matches the caller's requested size. assert radio_input_size % 16 == 0, ( f"radio_input_size={radio_input_size} must be a multiple of 16") self.radio_img_size = int(radio_input_size) # Resolve embed dims + patch grid via a tiny dummy forward. RADIO's # config has `.hidden_size` but not always exposed; dummy fwd is # the safe path. with torch.no_grad(): dummy = torch.zeros( 1, 3, self.radio_img_size, self.radio_img_size, dtype=next(self.radio.parameters()).dtype, device=next(self.radio.parameters()).device, ) summary, spatial = self.radio(dummy) # summary: (B, D_summary); spatial: (B, N, D_spatial) n_tokens = int(spatial.shape[1]) self._radio_spatial_dim = int(spatial.shape[-1]) self._radio_summary_dim = int(summary.shape[-1]) self._radio_n_patches = n_tokens side = int(round(math.sqrt(n_tokens))) assert side * side == n_tokens, ( f"RADIO patch count {n_tokens} is not a perfect square; " f"input {self.radio_img_size} may not be a clean multiple " f"of the ViT patch size (16).") self._radio_side = side print(f" RADIO: input {self.radio_img_size}², " f"spatial_dim={self._radio_spatial_dim}, " f"summary_dim={self._radio_summary_dim}, " f"patch_grid={side}²") # Rebuild head modules sized for RADIO's dims. self.embed_dim = self._radio_summary_dim feat_dim = self.feat_dim d_model = self.d_model self.feat_head = nn.Conv2d(self._radio_spatial_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) # Rebuild the optional cross-view attention stack at the spatial # embed_dim (NOT the summary dim — CVA operates on patch tokens). 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._radio_spatial_dim, n_heads=n_heads, ) # ──────────────────────────────────────────────────────────────────── # Per-view backbone forward # ──────────────────────────────────────────────────────────────────── def _resize_for_radio(self, x: torch.Tensor) -> torch.Tensor: """Resize (B, 3, S, S) to RADIO's expected input size. RADIO's input conditioner handles normalisation internally, so we keep the ImageNet-normalised tensor as-is.""" if x.shape[-1] != self.radio_img_size or x.shape[-2] != self.radio_img_size: x = F.interpolate( x, size=(self.radio_img_size, self.radio_img_size), mode="bilinear", align_corners=False) return x def _per_view_patch_features(self, rgb: list[torch.Tensor]): """Per-view RADIO forward. Returns: view_patches: list of N tensors of shape (B, D_spatial, ph, pw) view_global: list of N tensors of shape (B, D_summary) """ assert len(rgb) == self.n_views view_patches: list[torch.Tensor] = [] view_global: list[torch.Tensor] = [] side = self._radio_side for k in range(self.n_views): x = self._resize_for_radio(rgb[k]) summary, spatial = self.radio(x) B = spatial.shape[0] patch_grid = spatial.transpose(1, 2).reshape( B, self._radio_spatial_dim, side, side).contiguous() view_patches.append(patch_grid) view_global.append(summary) return view_patches, view_global def patch_features(self, rgb: torch.Tensor): raise NotImplementedError( "RADIO backbone uses the multi-view entry point. " "Use ``_per_view_patch_features(rgb_list)`` instead.")