"""High-resolution DINOv3 PARA models — bilinear up + 5-layer CNN refinement. Two new variants of the DINOv3 PARA family. Both replace the production ``feat_head + bilinear up + 2-layer residual MLP`` per-pixel refinement with **bilinear up to pred_size=256** followed by a **5-layer CNN** (plain 3×3 conv + GELU stack, last layer ungated). The bigger pred grid + deeper refinement aims to push the spatial resolution ceiling we hit with the 128×128 factorized model. - ``DinoVolumeSceneHires``: factorized scoring (same shape as the production ``DinoVolumeScene``). Defaults: ``pred_size=256``, ``n_height_bins=128``, ``feat_dim=48``. - ``DinoVolumeScenePerVoxelHires``: per-voxel scoring. CNN refines at ``feat_dim=16`` channels (cheap), then a **final 1×1 conv expands to ``feat_dim * n_height_bins`` channels** and we reshape to a per-voxel feature volume. Defaults: ``pred_size=256``, ``n_height_bins=64``, ``feat_dim=16``. Both share ``DINOv3 ViT-S/16+`` backbone + ``AdaLN-Zero`` query stack + ``sin/cos`` PEs from ``lib.model``. Volume scoring uses raw dot products (no L2-norm, no temperature; see ``vault/.../known-bugs.md`` #11). """ from __future__ import annotations import math import os import sys from pathlib import Path import torch import torch.nn as nn import torch.nn.functional as F sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) from lib.model import ( # noqa: E402 AdaLNZeroBlock, sin_pe, DEFAULT_DINO_REPO, DEFAULT_DINO_WEIGHTS, ) class FiveLayerCNN(nn.Module): """Modest 5-layer 3×3 CNN, same in/out channels, GELU between layers. Last conv ungated (no activation) so the output is unbounded — matches the production ``F_scene`` distribution that the factored scoring expects. Plain (non-residual) since 5 layers is shallow enough that vanishing gradients aren't a concern. """ def __init__(self, channels: int, kernel: int = 3): super().__init__() pad = kernel // 2 self.net = nn.Sequential( nn.Conv2d(channels, channels, kernel, padding=pad), nn.GELU(), nn.Conv2d(channels, channels, kernel, padding=pad), nn.GELU(), nn.Conv2d(channels, channels, kernel, padding=pad), nn.GELU(), nn.Conv2d(channels, channels, kernel, padding=pad), nn.GELU(), nn.Conv2d(channels, channels, kernel, padding=pad), ) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.net(x) class _DinoHiresBase(nn.Module): """Shared DINOv3 backbone + 1×1 conv to ``feat_dim`` + bilinear up to ``pred_size`` + 5-layer CNN refinement. Subclasses add the head-specific final stage (factorized scoring vs per-voxel expansion + scoring). """ def __init__( self, n_window: int, n_height_bins: int, pred_size: int, feat_dim: int, d_model: int, d_cond: int, d_sin_z: int, d_sin_t: int, n_blocks: int, n_gripper_bins: int, n_rot_clusters: int, dino_repo: str, dino_weights: str, ): super().__init__() self.n_window = n_window self.n_z = n_height_bins self.pred_size = pred_size self.feat_dim = feat_dim self.d_model = d_model self.d_sin_z = d_sin_z self.d_sin_t = d_sin_t self.n_gripper_bins = n_gripper_bins self.n_rot_clusters = n_rot_clusters self.backbone = torch.hub.load( os.environ.get("DINO_REPO_DIR", dino_repo), "dinov3_vits16plus", source="local", weights=os.environ.get("DINO_WEIGHTS_PATH", dino_weights), ) embed_dim = self.backbone.embed_dim self.embed_dim = embed_dim # Per-pixel refinement stage: 1×1 conv to feat_dim → bilinear up # to pred_size → 5-layer CNN. (Bilinear up + CNN replaces the # production ``ResidualMLP2`` and bumps pred_size 128→256.) self.feat_head = nn.Conv2d(embed_dim, feat_dim, 1) self.refine_cnn = FiveLayerCNN(feat_dim, kernel=3) # Sin/cos PEs (unchanged from production). self.register_buffer("z_sin", sin_pe(n_height_bins, d_sin_z)) self.register_buffer("t_sin", sin_pe(n_window, d_sin_t)) # Input projection: (eef_feat ⊕ cls) → d_model. self.input_proj = nn.Linear(feat_dim + embed_dim, d_model) self.t_cond_proj = nn.Linear(d_sin_t, d_cond) self.blocks = nn.ModuleList([ AdaLNZeroBlock(d_model, d_cond, mlp_ratio=4) for _ in range(n_blocks) ]) self.final_norm = nn.LayerNorm(d_model) # Subclasses override d_q logic + add scoring stage. self.grip_head = nn.Linear(d_model, n_gripper_bins) self.rot_head = nn.Linear(d_model, n_rot_clusters) def patch_features(self, rgb: torch.Tensor): """rgb (B, 3, S, S) → (patch_grid (B, embed, P_dino, P_dino), cls (B, embed)).""" out = self.backbone.forward_features(rgb) cls = out["x_norm_clstoken"] toks = out["x_norm_patchtokens"] P = int(round(math.sqrt(toks.shape[1]))) return toks.transpose(1, 2).reshape(rgb.size(0), -1, P, P), cls def _pixel_features(self, rgb: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """rgb → (F_scene (B, feat_dim, P, P), patch (raw DINOv3 grid), cls).""" patch, cls = self.patch_features(rgb) F0 = self.feat_head(patch) F1 = F.interpolate(F0, size=(self.pred_size, self.pred_size), mode="bilinear", align_corners=False) F_scene = self.refine_cnn(F1) # (B, feat_dim, P, P) return F_scene, patch, cls class DinoVolumeSceneHires(_DinoHiresBase): """Factorized YX × Z × T scoring at pred_size=256, n_height_bins=128. Same factored scoring as production ``DinoVolumeScene`` — only the per-pixel refinement and resolution change. """ def __init__( self, n_window: int = 16, n_height_bins: int = 128, pred_size: int = 256, feat_dim: int = 48, d_model: int = 256, d_cond: int = 128, d_sin_z: int = 48, d_sin_t: int = 48, n_blocks: int = 5, mlp_hidden: int = 128, # kept for CLI parity, unused n_gripper_bins: int = 32, n_rot_clusters: int = 64, dino_repo: str = DEFAULT_DINO_REPO, dino_weights: str = DEFAULT_DINO_WEIGHTS, ): super().__init__(n_window, n_height_bins, pred_size, feat_dim, d_model, d_cond, d_sin_z, d_sin_t, n_blocks, n_gripper_bins, n_rot_clusters, dino_repo, dino_weights) del mlp_hidden d_q = feat_dim + d_sin_z + d_sin_t self.q_head = nn.Linear(d_model, d_q) def forward(self, rgb: torch.Tensor, start_pix: torch.Tensor) -> dict: B, _, S, _ = rgb.shape T, Z, P = self.n_window, self.n_z, self.pred_size F_scene, patch, cls = self._pixel_features(rgb) scale = P / S sx = (start_pix[..., 0] * scale).long().clamp(0, P - 1) sy = (start_pix[..., 1] * scale).long().clamp(0, P - 1) b_idx = torch.arange(B, device=rgb.device) eef_feat = F_scene[b_idx, :, sy, sx] q_in = self.input_proj(torch.cat([eef_feat, cls], dim=-1)) q_in_bt = q_in.unsqueeze(1).expand(B, T, self.d_model).reshape(B * T, -1) cond_t = self.t_cond_proj(self.t_sin) cond_bt = cond_t.unsqueeze(0).expand(B, T, -1).reshape(B * T, -1) h = q_in_bt for blk in self.blocks: h = blk(h, cond_bt) h = self.final_norm(h) penult = h.view(B, T, self.d_model) q_spatial = self.q_head(penult) q_F = q_spatial[..., :self.feat_dim] q_z = q_spatial[..., self.feat_dim:self.feat_dim + self.d_sin_z] q_t = q_spatial[..., self.feat_dim + self.d_sin_z:] # Raw dot products — same as DinoVolumeScene (no F.normalize, no temperature). score_yx = torch.einsum("btc,bchw->bthw", q_F, F_scene) score_z = torch.einsum("btc,zc->btz", q_z, self.z_sin) score_t = torch.einsum("btc,tc->bt", q_t, self.t_sin) vol = (score_yx[:, :, None, :, :] + score_z[:, :, :, None, None] + score_t[:, :, None, None, None]) # (B, T, Z, P, P) return { "volume_logits": vol, "grip_logits": self.grip_head(penult), "rot_logits": self.rot_head(penult), "F_refined": F_scene, "patch_features": patch, } class DinoVolumeScenePerVoxelHires(_DinoHiresBase): """Per-voxel scoring at pred_size=256, n_height_bins=64, feat_dim=16. The 5-layer CNN refines a thin ``feat_dim=16`` channel stack at pred_size=256 — cheap on both params and memory. **After** refinement a single ``1×1 conv`` expands to ``feat_dim * n_z`` channels which are reshaped to the per-voxel feature volume ``(B, feat_dim, Z, P, P)``. Per-T query is a ``feat_dim``-d vector dotted against every voxel. """ def __init__( self, n_window: int = 16, n_height_bins: int = 64, pred_size: int = 256, feat_dim: int = 16, d_model: int = 256, d_cond: int = 128, d_sin_z: int = 48, # unused here; kept for CLI parity d_sin_t: int = 48, n_blocks: int = 5, mlp_hidden: int = 128, # kept for CLI parity n_gripper_bins: int = 32, n_rot_clusters: int = 64, dino_repo: str = DEFAULT_DINO_REPO, dino_weights: str = DEFAULT_DINO_WEIGHTS, ): super().__init__(n_window, n_height_bins, pred_size, feat_dim, d_model, d_cond, d_sin_z, d_sin_t, n_blocks, n_gripper_bins, n_rot_clusters, dino_repo, dino_weights) del mlp_hidden, d_sin_z # Final expansion: feat_dim → feat_dim * n_z channels per pixel. # Reshape gives the per-voxel feature volume (B, feat_dim, Z, P, P). self.feat_expand = nn.Conv2d(feat_dim, feat_dim * n_height_bins, kernel_size=1) # Per-T query is a single feat_dim-d vector (no z/t split — z # information lives entirely in the voxel features themselves). self.q_head = nn.Linear(d_model, feat_dim) def forward(self, rgb: torch.Tensor, start_pix: torch.Tensor) -> dict: del start_pix # per-voxel uses CLS only B, _, _, _ = rgb.shape T, Z, P, C = self.n_window, self.n_z, self.pred_size, self.feat_dim # Refine at thin C=feat_dim channels, then expand to C*Z and reshape. F_refined, patch, cls = self._pixel_features(rgb) F_full = self.feat_expand(F_refined) # (B, C*Z, P, P) F_vol = F_full.reshape(B, C, Z, P, P) # per-voxel volume # Per-T query from CLS only (no eef-sampling for per-voxel — matches # the existing DinoVolumeScenePerVoxel design). # input_proj expects (eef_feat ⊕ cls); we feed zeros for eef_feat to # reuse the same Linear shape without retraining a separate input layer. zero_eef = torch.zeros(B, C, device=rgb.device, dtype=cls.dtype) q_in = self.input_proj(torch.cat([zero_eef, cls], dim=-1)) q_in_bt = q_in.unsqueeze(1).expand(B, T, self.d_model).reshape(B * T, -1) cond_t = self.t_cond_proj(self.t_sin) cond_bt = cond_t.unsqueeze(0).expand(B, T, -1).reshape(B * T, -1) h = q_in_bt for blk in self.blocks: h = blk(h, cond_bt) h = self.final_norm(h) penult = h.view(B, T, self.d_model) # (B, T, d_model) q = self.q_head(penult) # (B, T, C) vol = torch.einsum("btc,bczhw->btzhw", q, F_vol) # (B, T, Z, P, P) return { "volume_logits": vol, "grip_logits": self.grip_head(penult), "rot_logits": self.rot_head(penult), # PCA viz operates on the pre-expansion 16-ch feature map (not # F_full's 1024 channels). SVD on (P*P, 1024) at inference was # adding ~1.5 s lag between chunks; (P*P, 16) is ~20× faster. "F_refined": F_refined, "patch_features": patch, } class DinoVolumeSceneHiresFlat(DinoVolumeSceneHires): """Variant of DinoVolumeSceneHires that predicts the WHOLE T-step trajectory as a single flat MLP output, instead of T independent per-token AdaLN-Zero forwards. Motivation: AdaLN-Zero processes T tokens independently (no cross-T mixing, only FiLM scalar from time sin/cos differentiates them), so temporal coherence has to be learned implicitly from data. A single MLP whose output dim is ``T × per_T_dim`` produces the entire trajectory in one forward — by construction the whole trajectory is one coherent prediction. Roughly 32× fewer MLP forwards too. Temporal positional structure is *not lost* — it's preserved at the scoring stage via the same ``q_z · z_sin`` / ``q_t · t_sin`` einsums the parent uses. """ def __init__(self, *args, mlp_layers: int = 3, mlp_hidden_mult: float = 2.0, **kwargs): super().__init__(*args, **kwargs) T = self.n_window n_gripper_bins = self.n_gripper_bins n_rot_clusters = self.n_rot_clusters d_model = self.d_model d_q = self.feat_dim + self.d_sin_z + self.d_sin_t # Replace the AdaLN-Zero stack with a plain MLP on (B, d_model). hidden = int(d_model * mlp_hidden_mult) layers: list[nn.Module] = [nn.Linear(d_model, hidden), nn.GELU()] for _ in range(mlp_layers - 1): layers.extend([nn.Linear(hidden, hidden), nn.GELU()]) self.global_mlp = nn.Sequential(*layers) # Output heads now produce the WHOLE flattened trajectory at once. # Reshape happens in forward. self.q_head = nn.Linear(hidden, T * d_q) self.grip_head = nn.Linear(hidden, T * n_gripper_bins) self.rot_head = nn.Linear(hidden, T * n_rot_clusters) # The parent's ``self.blocks``, ``self.final_norm`` and # ``self.t_cond_proj`` are now unused; leaving them in the # state_dict for backward init compatibility but they won't be # called in forward. def forward(self, rgb: torch.Tensor, start_pix: torch.Tensor) -> dict: B, _, S, _ = rgb.shape T, Z, P = self.n_window, self.n_z, self.pred_size F_scene, patch, cls = self._pixel_features(rgb) # Sample F_scene at start_pix. scale = P / S sx = (start_pix[..., 0] * scale).long().clamp(0, P - 1) sy = (start_pix[..., 1] * scale).long().clamp(0, P - 1) b_idx = torch.arange(B, device=rgb.device) eef_feat = F_scene[b_idx, :, sy, sx] # (B, feat_dim) q_in = self.input_proj(torch.cat([eef_feat, cls], dim=-1)) # (B, d_model) h = self.global_mlp(q_in) # (B, hidden) — ONE forward d_q = self.feat_dim + self.d_sin_z + self.d_sin_t q_flat = self.q_head(h).reshape(B, T, d_q) q_F = q_flat[..., :self.feat_dim] q_z = q_flat[..., self.feat_dim:self.feat_dim + self.d_sin_z] q_t = q_flat[..., self.feat_dim + self.d_sin_z:] # Raw dot products — same as the parent (no L2-norm, no temperature). score_yx = torch.einsum("btc,bchw->bthw", q_F, F_scene) score_z = torch.einsum("btc,zc->btz", q_z, self.z_sin) score_t = torch.einsum("btc,tc->bt", q_t, self.t_sin) vol = (score_yx[:, :, None, :, :] + score_z[:, :, :, None, None] + score_t[:, :, None, None, None]) grip_logits = self.grip_head(h).reshape(B, T, self.n_gripper_bins) rot_logits = self.rot_head(h).reshape(B, T, self.n_rot_clusters) return { "volume_logits": vol, "grip_logits": grip_logits, "rot_logits": rot_logits, "F_refined": F_scene, "patch_features": patch, } class _ResidualMLP(nn.Module): """N-layer MLP with skip connections every block; GELU activations. Layout: ``Linear(in_dim → hidden) GELU`` × 1 (no res, dim change), then ``(x + GELU(Linear(hidden → hidden)))`` × (n_layers-2), then ``Linear(hidden → out_dim)``. Skip-connections kick in once dims align, so even at n_layers=3 we get one residual block. """ def __init__(self, in_dim: int, hidden: int, out_dim: int, n_layers: int): super().__init__() assert n_layers >= 2 self.fc_in = nn.Linear(in_dim, hidden) self.mid = nn.ModuleList( [nn.Linear(hidden, hidden) for _ in range(n_layers - 2)]) self.fc_out = nn.Linear(hidden, out_dim) self.act = nn.GELU() def forward(self, x: torch.Tensor) -> torch.Tensor: h = self.act(self.fc_in(x)) for fc in self.mid: h = h + self.act(fc(h)) return self.fc_out(h) class DinoVolumeSceneHiresAblate(DinoVolumeSceneHires): """Ablation variant of the factorized hires model with optional past- trajectory encoding, configurable global MLP, and optional multi-layer DINOv3 patch-feature fusion. All ablation knobs default to "unchanged" so ``DinoVolumeSceneHiresAblate()`` is functionally identical to ``DinoVolumeSceneHires()``. Past encoding (concatenated with eef_feat + cls before the global MLP): - ``past_n``: number of past timesteps to consume (0 = none). - ``past_use_z`` / ``past_use_grip``: per-step (1+1 = 2) dims. - ``past_use_uv``: per-step (2) dims; px coords normalized to [-1, 1] in ``img_size`` space. - ``past_uv_delta``: subtract current ``start_pix`` from past UVs. - ``past_use_xyz``: per-step (3) dims for absolute world position. - ``past_xyz_delta``: also concatenate (past_xyz − current_xyz). Global MLP (replaces the original 1-Linear ``input_proj``): - ``mlp_layers``: depth (1 = plain Linear, ≥2 = MLP). - ``mlp_residual``: use ``_ResidualMLP`` (skip every block) for ≥3 layers. - ``mlp_hidden_mult``: hidden = d_model × this (only when mlp_layers > 1). Backbone fusion: - ``backbone_layer_fusion``: concat patch tokens from DINOv3 layers [3, 6, 9, 11] before ``feat_head`` 1×1 conv (4·embed_dim per patch). Same 31×31 grid; richer per-pixel info than the last layer alone. """ def __init__( self, n_window: int = 16, n_height_bins: int = 128, pred_size: int = 256, feat_dim: int = 48, d_model: int = 256, d_cond: int = 128, d_sin_z: int = 48, d_sin_t: int = 48, n_blocks: int = 5, mlp_hidden: int = 128, # unused; kept for CLI parity n_gripper_bins: int = 32, n_rot_clusters: int = 64, dino_repo: str = DEFAULT_DINO_REPO, dino_weights: str = DEFAULT_DINO_WEIGHTS, # ── ablation knobs ────────────────────────────────────────────── img_size: int = 504, past_n: int = 0, past_use_z: bool = True, past_z_delta: bool = False, past_use_grip: bool = True, past_use_uv: bool = False, past_uv_delta: bool = False, past_use_xyz: bool = False, past_xyz_delta: bool = False, mlp_layers: int = 1, mlp_residual: bool = False, mlp_hidden_mult: float = 2.0, backbone_layer_fusion: bool = False, ): super().__init__( n_window=n_window, n_height_bins=n_height_bins, pred_size=pred_size, feat_dim=feat_dim, d_model=d_model, d_cond=d_cond, d_sin_z=d_sin_z, d_sin_t=d_sin_t, n_blocks=n_blocks, mlp_hidden=mlp_hidden, n_gripper_bins=n_gripper_bins, n_rot_clusters=n_rot_clusters, dino_repo=dino_repo, dino_weights=dino_weights, ) # Stash flags (gate on past_n > 0). self.past_n = past_n self.past_use_z = past_use_z and past_n > 0 self.past_z_delta = past_z_delta self.past_use_grip = past_use_grip and past_n > 0 self.past_use_uv = past_use_uv and past_n > 0 self.past_uv_delta = past_uv_delta self.past_use_xyz = past_use_xyz and past_n > 0 self.past_xyz_delta = past_xyz_delta self.backbone_layer_fusion = backbone_layer_fusion self.img_size_model = img_size # Compute past-vector dimension. past_dim = 0 if past_n > 0: if self.past_use_z: past_dim += past_n if self.past_use_grip: past_dim += past_n if self.past_use_uv: past_dim += past_n * 2 if self.past_use_xyz: past_dim += past_n * 3 if self.past_xyz_delta: past_dim += past_n * 3 self.past_dim = past_dim # Rebuild input_proj as the requested MLP. Old super().__init__ already # built a plain Linear(feat_dim+embed_dim → d_model); we replace. in_dim = self.feat_dim + self.embed_dim + past_dim if mlp_layers == 1: self.input_proj = nn.Linear(in_dim, self.d_model) else: hidden = int(self.d_model * mlp_hidden_mult) if mlp_residual and mlp_layers >= 3: self.input_proj = _ResidualMLP(in_dim, hidden, self.d_model, mlp_layers) else: layers: list[nn.Module] = [nn.Linear(in_dim, hidden), nn.GELU()] for _ in range(mlp_layers - 2): layers.extend([nn.Linear(hidden, hidden), nn.GELU()]) layers.append(nn.Linear(hidden, self.d_model)) self.input_proj = nn.Sequential(*layers) # Backbone fusion: rebuild feat_head to take 4*embed_dim if fusing. if self.backbone_layer_fusion: self.feat_head = nn.Conv2d(4 * self.embed_dim, self.feat_dim, 1) def patch_features(self, rgb: torch.Tensor): """Standard path returns last-layer patch tokens at 31×31. With ``backbone_layer_fusion``, return concat of patch tokens from layers [3, 6, 9, 11] at the same 31×31 grid (giving 4×embed_dim per patch).""" if not self.backbone_layer_fusion: return super().patch_features(rgb) # DINOv3 vits16plus has 12 blocks (indices 0..11). Layers we want: # [3, 6, 9, 11] — early-mid-late + final, balanced sampling. layer_idx = [3, 6, 9, 11] feats = self.backbone.get_intermediate_layers(rgb, n=layer_idx, norm=True) # Each element: (B, N, embed_dim). Concat on channel dim. concat = torch.cat(list(feats), dim=-1) # (B, N, 4*embed) out = self.backbone.forward_features(rgb) cls = out["x_norm_clstoken"] P = int(round(math.sqrt(concat.shape[1]))) return concat.transpose(1, 2).reshape(rgb.size(0), -1, P, P), cls def forward( self, rgb: torch.Tensor, start_pix: torch.Tensor, past_z: torch.Tensor = None, past_grip: torch.Tensor = None, past_uv: torch.Tensor = None, past_xyz: torch.Tensor = None, current_xyz: torch.Tensor = None, ) -> dict: B, _, S, _ = rgb.shape T, Z, P = self.n_window, self.n_z, self.pred_size F_scene, patch, cls = self._pixel_features(rgb) # Sample eef_feat at start_pix. scale = P / S sx = (start_pix[..., 0] * scale).long().clamp(0, P - 1) sy = (start_pix[..., 1] * scale).long().clamp(0, P - 1) b_idx = torch.arange(B, device=rgb.device) eef_feat = F_scene[b_idx, :, sy, sx] # Build the global input vector. Normalize past features inline # using fixed scales so the Linear's first layer doesn't need to # absorb large absolute-value differences across modalities. parts = [eef_feat, cls] if self.past_n > 0: if self.past_use_z and past_z is not None: if self.past_z_delta and current_xyz is not None: # Delta from current EE z, scaled — deltas are small # (cm scale), multiply by 10 to give a comparable signal # to other normalized inputs. parts.append((past_z - current_xyz[:, 2:3]) * 10.0) else: # z in meters: workspace ≈ [0.05, 0.35]; center 0.2 scale 5. parts.append((past_z - 0.2) * 5.0) # (B, past_n) if self.past_use_grip and past_grip is not None: # grip in [-0.2, 0.8]: center 0.3 scale 2. parts.append((past_grip - 0.3) * 2.0) if self.past_use_uv and past_uv is not None: uv_n = (past_uv / self.img_size_model - 0.5) * 2.0 # (B, past_n, 2) if self.past_uv_delta: sp_n = (start_pix.float() / self.img_size_model - 0.5) * 2.0 uv_n = uv_n - sp_n.unsqueeze(1) parts.append(uv_n.reshape(B, -1)) if self.past_use_xyz and past_xyz is not None: # Workspace center ~0.3m, scale 2 → roughly [-1, 1]. parts.append(((past_xyz - 0.3) * 2.0).reshape(B, -1)) if self.past_xyz_delta and current_xyz is not None: # ±10cm delta → ±0.5 signal. delta = (past_xyz - current_xyz.unsqueeze(1)) * 5.0 parts.append(delta.reshape(B, -1)) q_in = self.input_proj(torch.cat(parts, dim=-1)) # (B, d_model) q_in_bt = q_in.unsqueeze(1).expand(B, T, self.d_model).reshape(B * T, -1) cond_t = self.t_cond_proj(self.t_sin) cond_bt = cond_t.unsqueeze(0).expand(B, T, -1).reshape(B * T, -1) h = q_in_bt for blk in self.blocks: h = blk(h, cond_bt) h = self.final_norm(h) penult = h.view(B, T, self.d_model) q_spatial = self.q_head(penult) q_F = q_spatial[..., :self.feat_dim] q_z = q_spatial[..., self.feat_dim:self.feat_dim + self.d_sin_z] q_t = q_spatial[..., self.feat_dim + self.d_sin_z:] score_yx = torch.einsum("btc,bchw->bthw", q_F, F_scene) score_z = torch.einsum("btc,zc->btz", q_z, self.z_sin) score_t = torch.einsum("btc,tc->bt", q_t, self.t_sin) vol = (score_yx[:, :, None, :, :] + score_z[:, :, :, None, None] + score_t[:, :, None, None, None]) return { "volume_logits": vol, "grip_logits": self.grip_head(penult), "rot_logits": self.rot_head(penult), "F_refined": F_scene, "patch_features": patch, } if __name__ == "__main__": import time device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"device = {device}") for cls, kw in [(DinoVolumeSceneHires, dict(n_height_bins=128, feat_dim=48)), (DinoVolumeScenePerVoxelHires, dict(n_height_bins=64, feat_dim=16))]: m = cls(n_window=16, pred_size=256, d_model=256, n_blocks=5, **kw).to(device).eval() n_t = sum(p.numel() for p in m.parameters() if p.requires_grad) print(f"\n{cls.__name__}: trainable {n_t / 1e6:.1f}M") rgb = torch.rand(1, 3, 504, 504, device=device) sp = torch.tensor([[252.0, 252.0]], device=device) with torch.no_grad(): t0 = time.time() out = m(rgb, sp) t1 = time.time() for k, v in out.items(): if hasattr(v, "shape"): print(f" {k}: {tuple(v.shape)}") print(f" forward = {(t1 - t0) * 1000:.1f} ms") if device.type == "cuda": print(f" peak: {torch.cuda.max_memory_allocated() / 1e9:.2f} GB") torch.cuda.reset_peak_memory_stats()