""" Vision Encoder Wrapper for Project Norn V20 ============================================= Wraps Google SigLIP-Base-patch16-384 as a frozen nn.Module for extracting patch-level visual features from input images. Output: [B, 576, 768] patch embeddings (24x24 grid x 768 dims) The encoder is always frozen during both training stages. Weights are downloaded from HuggingFace on first use and cached locally. """ import torch import torch.nn as nn from typing import Optional, Tuple from PIL import Image class NornVisionEncoder(nn.Module): """ Frozen SigLIP vision encoder that extracts patch embeddings from images. Architecture: SigLIP-Base-patch16-384 - Input: [B, 3, 384, 384] pixel values (normalized) - Output: [B, 576, 768] patch embeddings (24x24 grid) The encoder is always frozen and never trained. """ def __init__( self, model_name: str = "google/siglip-base-patch16-384", cache_dir: Optional[str] = None, device: Optional[torch.device] = None, ): super().__init__() self.model_name = model_name self.device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu") self._hidden_size = 768 self._num_patches = 576 # 24 x 24 self._image_size = 384 self._patch_size = 16 self._grid_size = 24 # 384 / 16 # Lazy load - encoder is created on first forward pass or explicit init self._encoder = None self._processor = None self._cache_dir = cache_dir @property def hidden_size(self) -> int: return self._hidden_size @property def num_patches(self) -> int: return self._num_patches @property def grid_size(self) -> int: return self._grid_size @property def image_size(self) -> int: return self._image_size def to(self, *args, **kwargs): super().to(*args, **kwargs) # Parse device if provided try: device = torch._C._nn._parse_to(*args, **kwargs)[0] if device is not None: self.device = device if self._encoder is not None: self._encoder.to(device) except Exception: pass return self def _ensure_loaded(self, target_device: Optional[torch.device] = None): """Lazy-load the SigLIP model and processor on first use.""" if target_device is not None: self.device = target_device if self._encoder is not None: if self._encoder.device != self.device: self._encoder.to(self.device) return try: from transformers import SiglipVisionModel, SiglipImageProcessor print(f"[VisionEncoder] Loading {self.model_name}...") self._processor = SiglipImageProcessor.from_pretrained( self.model_name, cache_dir=self._cache_dir ) self._encoder = SiglipVisionModel.from_pretrained( self.model_name, cache_dir=self._cache_dir ).to(self.device) # Freeze all parameters for param in self._encoder.parameters(): param.requires_grad = False self._encoder.eval() total_params = sum(p.numel() for p in self._encoder.parameters()) print(f"[VisionEncoder] Loaded: {total_params / 1e6:.1f}M params (frozen)") except ImportError: raise ImportError( "SigLIP requires transformers >= 4.37.0. " "Install with: pip install transformers>=4.37.0" ) def preprocess_images(self, images: list) -> torch.Tensor: """ Preprocess PIL images or file paths into normalized pixel tensors. Args: images: List of PIL.Image objects or file path strings Returns: pixel_values: [B, 3, 384, 384] normalized tensor """ self._ensure_loaded() pil_images = [] for img in images: if isinstance(img, str): pil_images.append(Image.open(img).convert("RGB")) elif isinstance(img, Image.Image): pil_images.append(img.convert("RGB")) else: raise ValueError(f"Expected PIL.Image or file path, got {type(img)}") processed = self._processor( images=pil_images, return_tensors="pt", ) return processed["pixel_values"].to(self.device) @torch.no_grad() def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: """ Extract patch embeddings from preprocessed pixel values. Args: pixel_values: [B, 3, 384, 384] normalized pixel tensor Returns: patch_embeddings: [B, 576, 768] - one embedding per 16x16 patch """ self.device = pixel_values.device self._ensure_loaded(target_device=pixel_values.device) outputs = self._encoder( pixel_values=pixel_values.to(self.device), output_hidden_states=False, ) patch_embeds = outputs.last_hidden_state # [B, 576, 768] # SigLIP-Base has no CLS token, so all 576 patches are spatial assert patch_embeds.shape[1] == self._num_patches, ( f"Expected {self._num_patches} patches, got {patch_embeds.shape[1]}. " f"Check model variant." ) return patch_embeds.to(pixel_values.device) def get_patch_grid_positions(self) -> Tuple[torch.Tensor, torch.Tensor]: """ Returns the (x, y) normalized grid positions for each of the 576 patches. Positions are in [0, 1] range, suitable for HRR spatial encoding. Returns: x_coords: [576] - horizontal positions y_coords: [576] - vertical positions """ grid_y, grid_x = torch.meshgrid( torch.linspace(0, 1, self._grid_size, device=self.device), torch.linspace(0, 1, self._grid_size, device=self.device), indexing="ij", ) return grid_x.reshape(-1), grid_y.reshape(-1) # --------------------------------------------------------------------------- # Self-Test # --------------------------------------------------------------------------- if __name__ == "__main__": print("=" * 60) print(" NornVisionEncoder Self-Test") print("=" * 60) encoder = NornVisionEncoder() # Test 1: Grid positions x, y = encoder.get_patch_grid_positions() assert x.shape == (576,), f"x shape: {x.shape}" assert y.shape == (576,), f"y shape: {y.shape}" assert x.min() >= 0.0 and x.max() <= 1.0 assert y.min() >= 0.0 and y.max() <= 1.0 print(f" [PASS] Grid positions: {x.shape}, range [{x.min():.2f}, {x.max():.2f}]") # Test 2: Properties print(f" [INFO] Hidden size: {encoder.hidden_size}") print(f" [INFO] Num patches: {encoder.num_patches}") print(f" [INFO] Grid size: {encoder.grid_size}x{encoder.grid_size}") print(f" [INFO] Image size: {encoder.image_size}") # Test 3: Load and run SigLIP (requires internet + GPU) try: encoder._ensure_loaded() dummy_pixels = torch.randn(1, 3, 384, 384, device=encoder.device) patches = encoder(dummy_pixels) assert patches.shape == (1, 576, 768), f"Unexpected shape: {patches.shape}" print(f" [PASS] Forward pass: {patches.shape}, dtype={patches.dtype}") print(f" [PASS] Patch stats: mean={patches.mean():.4f}, std={patches.std():.4f}") except Exception as e: print(f" [SKIP] SigLIP forward test: {e}") print("=" * 60) print(" All vision encoder tests passed!") print("=" * 60)