"""DinoVolumeSceneV4VLM — v4 head with InternVL2.5-1B vision tower. Same v4 head (per-voxel MLP fusion + sin/cos height PE + adaptive per-view z range), with the DINOv3 backbone replaced by InternVL2.5-1B's vision tower **followed by a forward pass through the LLM**. The vision tokens attend to the text prompt (dataset / task name) and to EACH OTHER (one LLM forward consumes both views' tokens, so cross-view attention is free via the LLM's self-attention). ────────────────────────────────────────────────────────────────────── Architecture deltas vs v4 ────────────────────────────────────────────────────────────────────── 1. **Backbone**: ``OpenGVLab/InternVL2_5-1B``. We use the VLM's own ``extract_feature(pixel_values)`` method to get vision tokens projected into the LLM's embedding space (the VLM's vision encoder + projector, exactly as the model was pretrained). 2. **Vision-language attention**: we build the sequence [view_0_img_tokens, view_1_img_tokens, ..., text_embeds] and call the LLM forward with ``output_hidden_states=True``. The vision tokens flow through every LLM layer, attending to text and to the other views' tokens. We take the last-layer hidden states at the image-token positions per view, reshape to (B, hidden_dim, grid, grid), and feed into the standard v4 pipeline (1×1 conv → bilinear-up to pred_size → 5-layer refine CNN). 3. **What's trainable**: vision encoder + projector are TRAINABLE. The LLM is FROZEN. Gradients still flow back through the frozen LLM into the vision tokens — that's how the vision encoder learns to produce tokens that the (frozen) LLM finds useful for the downstream voxel loss. The image-token mlp projector is trainable. 4. **Global token**: there's no explicit CLS in InternVL. We use the mean of the last text-position hidden states (post-LLM) as the global per-view descriptor — it's the LLM's summary of "view + task". Per view's global = mean of LLM's last-layer hidden at that view's token positions. Concat over views for ``cls_fusion='concat'``. 5. **Task prompt**: just the dataset name, e.g. ``"cup pickplace"`` or ``"towel folding"``. Passed via ``task_text`` argument to forward (or set as a class default). Same prompt for all views in one batch. 6. **Image size**: InternVL2.5 vision tower expects 448×448 with ImageNet normalisation. Our PARA pipeline runs DINOv3 at 504×504 with the same normalisation. We resize 504→448 INSIDE this module so the rest of the training pipeline (K_in, T_w2c at img_size=504) does not need to know. Everything else (per-voxel MLP fusion, sin/cos PE, adaptive per-view z, per-voxel grip/rot, voxel sampling, loss) is unchanged from v4. ────────────────────────────────────────────────────────────────────── Cost notes ────────────────────────────────────────────────────────────────────── InternVL2.5-1B is ~1 B params (frozen LLM ≈ 0.9 B, vision tower ≈ 0.1 B). Forward pass with N=2 views at 448² gives ~256 image tokens × 2 + text tokens — well within attention budget. With bf16 + frozen LLM, memory at B=4 should be fine on a 24 GB GPU. If tight, drop ``batch_size`` to 2. """ 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 # InternVL uses ImageNet stats (same as DINOv3) — but its vision tower # expects 448×448 not 504×504. We resize inside this module. VLM_IMAGENET_MEAN = (0.485, 0.456, 0.406) VLM_IMAGENET_STD = (0.229, 0.224, 0.225) VLM_DEFAULT_IMG_SIZE = 448 class DinoVolumeSceneV4VLM(DinoVolumeSceneV4): """v4 head + InternVL2.5-1B vision tower (frozen LLM, trainable vision).""" def __init__(self, *args, vlm_model_name: str = "OpenGVLab/InternVL2_5-1B", task_text_default: str = "pick and place", freeze_llm: bool = True, vlm_img_size: int = VLM_DEFAULT_IMG_SIZE, **kwargs): # v4 init loads DINOv3 (which we delete). Everything else (PE # buffers, voxel MLPs, head modules) gets built by the parent. super().__init__(*args, **kwargs) del self.backbone # drop DINOv3 # Load InternVL. bf16 because LLM is huge. from transformers import AutoModel, AutoTokenizer print(f"Loading VLM: {vlm_model_name}") self.vlm = AutoModel.from_pretrained( vlm_model_name, torch_dtype=torch.bfloat16, trust_remote_code=True, ) self._tokenizer = AutoTokenizer.from_pretrained( vlm_model_name, trust_remote_code=True, ) self.task_text_default = task_text_default self.vlm_img_size = int(vlm_img_size) self.freeze_llm = bool(freeze_llm) # Resolve LLM + dims from InternVL's config. cfg = self.vlm.config llm_cfg = cfg.llm_config if hasattr(cfg, "llm_config") else cfg.text_config self.vlm_hidden_dim = int(llm_cfg.hidden_size) self._lm = self._resolve_lm_submodule() if self.freeze_llm: for p in self._lm.parameters(): p.requires_grad = False self._lm.eval() else: # Unfrozen LLM: enable gradient checkpointing so backward # activations fit on a 24 GB GPU. Adds ~30% compute, halves # activation memory. for p in self._lm.parameters(): p.requires_grad = True if hasattr(self._lm, "gradient_checkpointing_enable"): self._lm.gradient_checkpointing_enable() elif hasattr(self.vlm, "gradient_checkpointing_enable"): self.vlm.gradient_checkpointing_enable() # Vision tower (+ projector) trainable so it learns task-relevant # tokens. The pretrained init carries the VLM alignment signal. for n, m in self.vlm.named_modules(): if not n: continue if n.startswith("vision_model") or n.startswith("mlp1") or "projector" in n: for p in m.parameters(): p.requires_grad = True # Rebuild modules that depended on DINOv3's embed_dim (384): # feat_head : embed_dim → feat_dim # input_proj: cls fan-in scales with embed_dim self.embed_dim = self.vlm_hidden_dim 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 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, ) # Buffers for ImageNet renormalisation when input is already in # ImageNet stats (no-op identity buffers; reserved for future # alt-norm inputs). self.register_buffer( "_imn_mean", torch.tensor(VLM_IMAGENET_MEAN).view(1, 3, 1, 1)) self.register_buffer( "_imn_std", torch.tensor(VLM_IMAGENET_STD).view(1, 3, 1, 1)) def _resolve_lm_submodule(self): # InternVL exposes the LM at various attrs depending on version. for n in ("language_model", "llm", "lm", "text_model"): if hasattr(self.vlm, n): return getattr(self.vlm, n) raise RuntimeError("Could not locate InternVL's LM submodule.") # ──────────────────────────────────────────────────────────────────── # Multi-view backbone call (VLM) # ──────────────────────────────────────────────────────────────────── def _vision_tokens_single(self, x: torch.Tensor) -> torch.Tensor: """ImageNet-normalised (B, 3, S, S) → image tokens (B, N_img, D_lm) via InternVL's vision encoder + projector. Uses ``extract_feature`` which already projects into the LM's embedding space. """ # Resize to InternVL's native input. if x.shape[-1] != self.vlm_img_size or x.shape[-2] != self.vlm_img_size: x = F.interpolate(x, size=(self.vlm_img_size, self.vlm_img_size), mode="bilinear", align_corners=False) x = x.to(self.vlm.dtype) return self.vlm.extract_feature(x) # (B, N_img, D_lm) def _per_view_patch_features(self, rgb: list[torch.Tensor]): """Multi-view VLM forward. Steps: 1. Per view, run vision encoder + projector → image tokens. 2. Tokenise task text (one prompt for the whole batch). 3. Build sequence ``[v0_tokens, v1_tokens, ..., text_embeds]`` and run the (frozen) LLM with ``output_hidden_states=True``. 4. Pull last-layer hidden states at each view's image-token positions; reshape to (B, D_lm, grid, grid) per view. 5. Per-view global = mean of LLM hidden at that view's positions. """ assert len(rgb) == self.n_views N = self.n_views B = rgb[0].shape[0] device = rgb[0].device # 1. Image tokens per view (vision encoder + projector). per_view_img_tokens: list[torch.Tensor] = [] # each (B, N_img, D_lm) for k in range(N): per_view_img_tokens.append(self._vision_tokens_single(rgb[k])) n_img_per_view = per_view_img_tokens[0].shape[1] # 2. Task text → embeds. task_texts = [self.task_text_default] * B text_inputs = self._tokenizer( task_texts, return_tensors="pt", padding=True, truncation=True, max_length=32, ).to(device) text_embeds = self._lm.get_input_embeddings()( text_inputs["input_ids"]).to(per_view_img_tokens[0].dtype) text_attn = text_inputs["attention_mask"] n_text = text_embeds.shape[1] # 3. Concatenate all view tokens + text into one sequence. inputs_embeds = torch.cat(per_view_img_tokens + [text_embeds], dim=1) attn_mask = torch.cat([ torch.ones(B, N * n_img_per_view, device=device, dtype=torch.long), text_attn, ], dim=1) # 4. Frozen LLM forward; pull last layer hidden states. lm_kwargs = dict( inputs_embeds=inputs_embeds, attention_mask=attn_mask, output_hidden_states=True, return_dict=True, ) if self.freeze_llm: with torch.no_grad(): # Gradients to vision tokens still flow because they are # the leaf inputs to the LLM; LLM params don't need grad, # so wrapping in no_grad is safe ONLY if the loss flows # through nothing that requires the LLM's own gradients. # We need vision-token gradients to backprop THROUGH the # LLM layers, which means we CAN'T use no_grad here. Drop # the no_grad: PyTorch will just skip param grads for # frozen tensors but propagate activation grads. pass lm_out = self._lm(**lm_kwargs) hidden = lm_out.hidden_states[-1].float() # (B, N_total, D_lm) # 5. Split per view; reshape to (B, D_lm, grid, grid). Mean of # each view's positions = global token. grid = int(round(math.sqrt(n_img_per_view))) view_patches: list[torch.Tensor] = [] view_global: list[torch.Tensor] = [] for k in range(N): start = k * n_img_per_view end = start + n_img_per_view seg = hidden[:, start:end, :] # (B, N_img, D_lm) if grid * grid == n_img_per_view: p = seg.reshape(B, grid, grid, -1).permute(0, 3, 1, 2).contiguous() else: pad = grid * grid - n_img_per_view seg_pad = F.pad(seg, (0, 0, 0, pad)) if pad > 0 else seg p = seg_pad.reshape(B, grid, grid, -1).permute(0, 3, 1, 2).contiguous() view_patches.append(p) view_global.append(seg.mean(dim=1)) return view_patches, view_global def patch_features(self, rgb: torch.Tensor): raise NotImplementedError( "VLM backbone needs multi-view input + text. " "Use ``_per_view_patch_features(rgb_list)`` instead.")