"""DinoVolumeSceneV4DynaFLIP — v4 head with DynaFLIP-base backbone. Same v4 head (per-voxel MLP fusion + sin/cos height PE + adaptive per-view z range), with the DINOv3 backbone replaced by **DynaFLIP-base** from HuggingFace (``jlee-larr/dynaflip-base``). DynaFLIP is a FLIP-derived ~0.2 B-param vision encoder; the README advertises stronger robotics transfer than vanilla CLIP/DINO. We use just the vision branch. ────────────────────────────────────────────────────────────────────── Architecture deltas vs v4 ────────────────────────────────────────────────────────────────────── 1. **Backbone**: ``jlee-larr/dynaflip-base`` loaded via ``AutoModel.from_pretrained(..., trust_remote_code=True)``. Output is ``v.last_hidden_state`` of shape ``(B, num_patches, 768)`` and ``v.pooler_output`` of shape ``(B, 1536)`` (CLS + mean(patches)). 2. **Per-view loop, no backbone-level cross-view attention**. Same as the default DINOv3 path — cross-view fusion happens later in the per-voxel MLP. If you want cross-view at the backbone level, set ``cross_view_layers > 0`` to enable the v4 ``CrossViewAttnStack``. 3. **Per-pixel features**: 768-d patch tokens → 1×1 conv (feat_head) → bilinear-up to ``pred_size`` → 5-layer refine CNN (the standard v4 pipeline; no DPT, no extra projection). 4. **Global descriptor**: the pooler output is 1536-d (twice embed_dim). We keep it as the per-view global token; ``input_proj`` is rebuilt to accept ``N · 1536 + feat_dim + past_fan`` instead of v4's ``N · embed_dim + feat_dim + past_fan``. 5. **Image size + normalisation**: DynaFLIP's processor is consulted at init to read its native input size and mean/std. Inputs from the PARA pipeline arrive ImageNet-normalised at 504×504; we denormalise, resize to DynaFLIP's native size (typically 224 or 336), then renormalise with DynaFLIP's stats. 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 # Standard ImageNet stats — the PARA pipeline normalises with these # before handing tensors to the model. IMAGENET_MEAN = (0.485, 0.456, 0.406) IMAGENET_STD = (0.229, 0.224, 0.225) class DinoVolumeSceneV4DynaFLIP(DinoVolumeSceneV4): """v4 head + DynaFLIP-base vision backbone (single-view loop).""" def __init__(self, *args, dynaflip_model_name: str = "jlee-larr/dynaflip-base", **kwargs): super().__init__(*args, **kwargs) # v4 already built the DINOv3 backbone; drop it before loading. del self.backbone from transformers import AutoModel, AutoProcessor print(f"Loading DynaFLIP backbone: {dynaflip_model_name}") self.dynaflip = AutoModel.from_pretrained( dynaflip_model_name, trust_remote_code=True, ) processor = AutoProcessor.from_pretrained( dynaflip_model_name, trust_remote_code=True, ) # Resolve patch tokens' embed_dim from a tiny dummy forward. # Some FLIP variants expose .config.vision_config.hidden_size; the # safest path is to read the README's stated 768 but verify here. df_embed_dim = 768 try: cfg = self.dynaflip.config for attr in ("vision_config", "vision_model_config", "model_config"): if hasattr(cfg, attr): sub = getattr(cfg, attr) if hasattr(sub, "hidden_size"): df_embed_dim = int(sub.hidden_size) break except Exception: pass # Native input + normalisation from the processor (fall back to # CLIP-style defaults if the processor's structure isn't standard). ip = getattr(processor, "image_processor", processor) df_size = 224 df_mean = (0.48145466, 0.4578275, 0.40821073) df_std = (0.26862954, 0.26130258, 0.27577711) try: if hasattr(ip, "size"): sz = ip.size if isinstance(sz, dict): df_size = int(sz.get("height", sz.get("shortest_edge", df_size))) elif isinstance(sz, int): df_size = int(sz) if getattr(ip, "image_mean", None) is not None: df_mean = tuple(float(x) for x in ip.image_mean) if getattr(ip, "image_std", None) is not None: df_std = tuple(float(x) for x in ip.image_std) except Exception: pass self.df_img_size = int(df_size) self.register_buffer( "_df_mean", torch.tensor(df_mean).view(1, 3, 1, 1)) self.register_buffer( "_df_std", torch.tensor(df_std).view(1, 3, 1, 1)) self.register_buffer( "_imn_mean", torch.tensor(IMAGENET_MEAN).view(1, 3, 1, 1)) self.register_buffer( "_imn_std", torch.tensor(IMAGENET_STD).view(1, 3, 1, 1)) print(f" DynaFLIP: input {self.df_img_size}², " f"embed_dim={df_embed_dim}, " f"mean={df_mean[:1]}…, std={df_std[:1]}…") # Pooler dim — confirmed from a dummy forward (handles both the # 1536-d README claim and any future variant). with torch.no_grad(): dummy = torch.zeros( 1, 3, self.df_img_size, self.df_img_size, dtype=next(self.dynaflip.parameters()).dtype, device=next(self.dynaflip.parameters()).device, ) v = self.dynaflip.vision_outputs(dummy) patch_tok = v.last_hidden_state # (1, N, D_patch) pooled = getattr(v, "pooler_output", None) self._df_pool_dim = int(pooled.shape[-1]) if pooled is not None \ else int(patch_tok.shape[-1]) df_embed_dim = int(patch_tok.shape[-1]) self._df_n_patches = int(patch_tok.shape[1]) # Rebuild head modules sized for DynaFLIP's dims. # ``self.embed_dim`` controls the CLS fan-in used by input_proj; # we set it to the *pooler* dim so the global descriptor flows in. self.embed_dim = self._df_pool_dim self._df_embed_dim = df_embed_dim feat_dim = self.feat_dim d_model = self.d_model self.feat_head = nn.Conv2d(df_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) # Rebuild the optional cross-view attention stack at the new # patch-token embed_dim (note: NOT the pooler dim). 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=df_embed_dim, n_heads=n_heads, ) # ──────────────────────────────────────────────────────────────────── # Per-view backbone forward # ──────────────────────────────────────────────────────────────────── def _to_dynaflip_norm(self, x: torch.Tensor) -> torch.Tensor: """Convert (B, 3, S, S) from ImageNet stats → DynaFLIP stats, resized to DynaFLIP's native input.""" # Denormalise from ImageNet → [0, 1]. x01 = x * self._imn_std + self._imn_mean # Resize to native. if x01.shape[-1] != self.df_img_size or x01.shape[-2] != self.df_img_size: x01 = F.interpolate( x01, size=(self.df_img_size, self.df_img_size), mode="bilinear", align_corners=False) # Renormalise with DynaFLIP stats. return (x01 - self._df_mean) / self._df_std def _per_view_patch_features(self, rgb: list[torch.Tensor]): """Per-view DynaFLIP forward. Returns: view_patches: list of N tensors of shape (B, D_patch, ph, pw) view_global: list of N tensors of shape (B, D_pool) """ assert len(rgb) == self.n_views view_patches: list[torch.Tensor] = [] view_global: list[torch.Tensor] = [] # Skip CLS token if its presence pushes the patch count above a # perfect square. Patch grid side inferred from the count returned # by the dummy forward at init. n_tokens = self._df_n_patches side = int(round(math.sqrt(n_tokens))) has_cls = (side * side != n_tokens) n_patches_grid = side * side for k in range(self.n_views): x = self._to_dynaflip_norm(rgb[k]) v = self.dynaflip.vision_outputs(x) patches = v.last_hidden_state # (B, n_tokens, D_patch) if has_cls and patches.shape[1] - 1 == n_patches_grid: patches = patches[:, 1:, :] patches = patches[:, :n_patches_grid, :] # safe slice B = patches.shape[0] patch_grid = patches.transpose(1, 2).reshape( B, self._df_embed_dim, side, side).contiguous() view_patches.append(patch_grid) pooled = getattr(v, "pooler_output", None) if pooled is None: pooled = patches.mean(dim=1) # fall back view_global.append(pooled) return view_patches, view_global def patch_features(self, rgb: torch.Tensor): raise NotImplementedError( "DynaFLIP backbone uses the multi-view entry point. " "Use ``_per_view_patch_features(rgb_list)`` instead.")