Release Project Norn V20: 9.55B Unified Holographic Titan (Titans-HRR + Sliding Window Attention)
102d648 verified | """ | |
| 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 | |
| def hidden_size(self) -> int: | |
| return self._hidden_size | |
| def num_patches(self) -> int: | |
| return self._num_patches | |
| def grid_size(self) -> int: | |
| return self._grid_size | |
| 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) | |
| 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) | |