"""DinoVolumeSceneV4PaliGemma — v4 head with PaliGemma backbone. Mirrors :mod:`lib.model_volumetric_v4_vlm` but swaps InternVL2.5-1B for **PaliGemma** (default ``google/paligemma-3b-pt-448``). PaliGemma pairs a SigLIP-So400m/14 vision tower with a Gemma-2B LLM; we use the vision tower + multi-modal projector + a forward pass through the LLM with the task prompt, exactly like v4_vlm. ────────────────────────────────────────────────────────────────────── Architecture deltas vs v4 ────────────────────────────────────────────────────────────────────── 1. **Backbone**: ``PaliGemmaForConditionalGeneration.from_pretrained(...)``. We rely on ``get_image_features(pixel_values)`` to produce projected image tokens of shape ``(B, n_img_per_view, 2048)`` already in the Gemma hidden dim. SigLIP-So400m/14 at 448² gives n_img_per_view=1024. 2. **Vision-language attention**: identical to v4_vlm — build the sequence ``[v0_tokens, v1_tokens, …, text_embeds]``, run the LLM with ``output_hidden_states=True``, take last-layer hidden states at each view's image-token positions, reshape to ``(B, D_lm, grid, grid)``. 3. **Trainability**: vision tower + multi-modal projector are trainable. The LLM is **frozen by default** (``--no-vlm_freeze_llm`` to unfreeze with gradient checkpointing). Unfreezing is recommended on the 96 GB GPUs but optional on 24 GB. 4. **Global token**: per-view mean of the LLM's last-layer hidden states at that view's image-token positions, same as v4_vlm. 5. **Image size + normalisation**: SigLIP/PaliGemma's processor uses either ImageNet stats (mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) or variant-specific stats. We resolve from the processor; default is ``(0.5, 0.5, 0.5)`` for both mean and std (SigLIP convention). ────────────────────────────────────────────────────────────────────── Cost notes ────────────────────────────────────────────────────────────────────── PaliGemma-3B at 448² is roughly: - vision_tower: ~0.4 B params (SigLIP-So400m/14) - multi_modal_projector: 1152×2048 ≈ 2.4 M - language_model: ~2.6 B params (Gemma-2B) - Total: ~3.0 B params Frozen LLM: trainable ≈ 0.4 B vision + projector + v4 head. Unfrozen LLM with gradient checkpointing: still tight on 24 GB at B=2 even at bf16 — run on the 96 GB yukon GPUs if unfreezing. """ 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 # SigLIP normalisation (matches the PaliGemma processor default). SIGLIP_MEAN = (0.5, 0.5, 0.5) SIGLIP_STD = (0.5, 0.5, 0.5) # PARA inputs arrive ImageNet-normalised. IMAGENET_MEAN = (0.485, 0.456, 0.406) IMAGENET_STD = (0.229, 0.224, 0.225) class DinoVolumeSceneV4PaliGemma(DinoVolumeSceneV4): """v4 head + PaliGemma vision tower + Gemma LLM (frozen by default).""" def __init__(self, *args, paligemma_model_name: str = "google/paligemma-3b-pt-448", task_text_default: str = "pick and place", freeze_llm: bool = True, **kwargs): super().__init__(*args, **kwargs) del self.backbone # drop DINOv3 from transformers import ( PaliGemmaForConditionalGeneration, AutoProcessor) print(f"Loading PaliGemma: {paligemma_model_name}") self.vlm = PaliGemmaForConditionalGeneration.from_pretrained( paligemma_model_name, torch_dtype=torch.bfloat16, ) processor = AutoProcessor.from_pretrained(paligemma_model_name) self._tokenizer = processor.tokenizer self.task_text_default = task_text_default self.freeze_llm = bool(freeze_llm) # Resolve native image size + normalisation from the processor. ip = getattr(processor, "image_processor", processor) pg_size = 448 pg_mean = SIGLIP_MEAN pg_std = SIGLIP_STD try: if hasattr(ip, "size"): sz = ip.size if isinstance(sz, dict): pg_size = int(sz.get("height", pg_size)) elif isinstance(sz, int): pg_size = int(sz) if getattr(ip, "image_mean", None) is not None: pg_mean = tuple(float(x) for x in ip.image_mean) if getattr(ip, "image_std", None) is not None: pg_std = tuple(float(x) for x in ip.image_std) except Exception: pass self.pg_img_size = int(pg_size) self.register_buffer( "_pg_mean", torch.tensor(pg_mean).view(1, 3, 1, 1)) self.register_buffer( "_pg_std", torch.tensor(pg_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)) # LLM submodule + hidden dim. self._lm = self._resolve_lm_submodule() try: self.vlm_hidden_dim = int(self._lm.config.hidden_size) except Exception: self.vlm_hidden_dim = int( getattr(self.vlm.config, "hidden_size", 2048)) print(f" PaliGemma: input {self.pg_img_size}², " f"LLM hidden={self.vlm_hidden_dim}, " f"mean={pg_mean[:1]}…, std={pg_std[:1]}…, " f"freeze_llm={self.freeze_llm}") if self.freeze_llm: for p in self._lm.parameters(): p.requires_grad = False self._lm.eval() else: for p in self._lm.parameters(): p.requires_grad = True if hasattr(self._lm, "gradient_checkpointing_enable"): self._lm.gradient_checkpointing_enable() # Vision tower + multi-modal projector trainable. for n, m in self.vlm.named_modules(): if not n: continue if (n.startswith("vision_tower") or n.startswith("multi_modal_projector")): for p in m.parameters(): p.requires_grad = True # Rebuild head modules sized to Gemma hidden_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, ) def _resolve_lm_submodule(self): for n in ("language_model", "llm", "lm", "text_model"): if hasattr(self.vlm, n): return getattr(self.vlm, n) raise RuntimeError("Could not locate PaliGemma's LM submodule.") # ──────────────────────────────────────────────────────────────────── # Multi-view forward (PaliGemma) # ──────────────────────────────────────────────────────────────────── def _to_paligemma_norm(self, x: torch.Tensor) -> torch.Tensor: """ImageNet stats (PARA) → SigLIP stats (PaliGemma), at native res.""" x01 = x * self._imn_std + self._imn_mean if (x01.shape[-1] != self.pg_img_size or x01.shape[-2] != self.pg_img_size): x01 = F.interpolate( x01, size=(self.pg_img_size, self.pg_img_size), mode="bilinear", align_corners=False) return (x01 - self._pg_mean) / self._pg_std def _vision_tokens_single(self, x: torch.Tensor) -> torch.Tensor: """Run vision tower + projector for one view. Returns image tokens already in the LLM embedding space, shape (B, n_img, hidden_dim).""" x = self._to_paligemma_norm(x).to(self.vlm.dtype) # PaliGemma exposes a helper that does vision_tower → projector. if hasattr(self.vlm, "get_image_features"): return self.vlm.get_image_features(x) # Fall-back: do it manually. v_out = self.vlm.vision_tower(x) if hasattr(v_out, "last_hidden_state"): feats = v_out.last_hidden_state else: feats = v_out return self.vlm.multi_modal_projector(feats) def _per_view_patch_features(self, rgb: list[torch.Tensor]): """Same pipeline as v4_vlm: per-view vision tokens, concat with text embeds, LLM forward, split per view + reshape to (B, D, g, g). """ assert len(rgb) == self.n_views N = self.n_views B = rgb[0].shape[0] device = rgb[0].device # 1. Image tokens per view. per_view_img_tokens: list[torch.Tensor] = [] 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 views + text. 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. LLM forward → last-layer hidden states. lm_out = self._lm( inputs_embeds=inputs_embeds, attention_mask=attn_mask, output_hidden_states=True, return_dict=True, ) hidden = lm_out.hidden_states[-1].float() # (B, N_total, D_lm) # 5. Per-view split → (B, D_lm, grid, grid). 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, :] 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( "PaliGemma backbone uses the multi-view entry point. " "Use ``_per_view_patch_features(rgb_list)`` instead.")