"""SigLIP2-base-patch16-512 vision encoder, adapted from lusxvr/nanoVLM (MIT). Loads real pretrained weights from google/siglip2-base-patch16-512. """ import math import torch import torch.nn as nn import torch.nn.functional as F VIT_HIDDEN_DIM = 768 VIT_INTER_DIM = 3072 VIT_PATCH_SIZE = 16 VIT_IMG_SIZE = 512 VIT_N_HEADS = 12 VIT_N_BLOCKS = 12 VIT_LN_EPS = 1e-6 VIT_MODEL_TYPE = "google/siglip2-base-patch16-512" class ViTPatchEmbeddings(nn.Module): def __init__(self): super().__init__() self.num_patches = (VIT_IMG_SIZE // VIT_PATCH_SIZE) ** 2 self.conv = nn.Conv2d(3, VIT_HIDDEN_DIM, kernel_size=VIT_PATCH_SIZE, stride=VIT_PATCH_SIZE, padding="valid") self.position_embedding = nn.Parameter(torch.rand(1, self.num_patches, VIT_HIDDEN_DIM)) def forward(self, x): x = self.conv(x) x = x.flatten(2).transpose(1, 2) x = x + self.position_embedding return x class ViTMultiHeadAttention(nn.Module): def __init__(self): super().__init__() self.n_heads = VIT_N_HEADS self.head_dim = VIT_HIDDEN_DIM // VIT_N_HEADS self.qkv_proj = nn.Linear(VIT_HIDDEN_DIM, 3 * VIT_HIDDEN_DIM, bias=True) self.out_proj = nn.Linear(VIT_HIDDEN_DIM, VIT_HIDDEN_DIM, bias=True) def forward(self, x): B, T, C = x.size() qkv = self.qkv_proj(x) q, k, v = qkv.split(C, dim=2) q = q.view(B, T, self.n_heads, self.head_dim).transpose(1, 2) k = k.view(B, T, self.n_heads, self.head_dim).transpose(1, 2) v = v.view(B, T, self.n_heads, self.head_dim).transpose(1, 2) y = F.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=False) y = y.transpose(1, 2).contiguous().view(B, T, C) return self.out_proj(y) class ViTMLP(nn.Module): def __init__(self): super().__init__() self.fc1 = nn.Linear(VIT_HIDDEN_DIM, VIT_INTER_DIM) self.fc2 = nn.Linear(VIT_INTER_DIM, VIT_HIDDEN_DIM) self.act = nn.GELU(approximate="tanh") def forward(self, x): return self.fc2(self.act(self.fc1(x))) class ViTBlock(nn.Module): def __init__(self): super().__init__() self.ln1 = nn.LayerNorm(VIT_HIDDEN_DIM, eps=VIT_LN_EPS) self.attn = ViTMultiHeadAttention() self.ln2 = nn.LayerNorm(VIT_HIDDEN_DIM, eps=VIT_LN_EPS) self.mlp = ViTMLP() def forward(self, x): x = x + self.attn(self.ln1(x)) x = x + self.mlp(self.ln2(x)) return x class ViT(nn.Module): def __init__(self): super().__init__() self.patch_embedding = ViTPatchEmbeddings() self.blocks = nn.ModuleList([ViTBlock() for _ in range(VIT_N_BLOCKS)]) self.layer_norm = nn.LayerNorm(VIT_HIDDEN_DIM, eps=VIT_LN_EPS) def forward(self, x): x = self.patch_embedding(x) for block in self.blocks: x = block(x) return self.layer_norm(x) @classmethod def from_pretrained(cls): from huggingface_hub import hf_hub_download import safetensors model = cls() safetensors_file = hf_hub_download(repo_id=VIT_MODEL_TYPE, filename="model.safetensors") sd = model.state_dict() mapping = { "vision_model.embeddings.patch_embedding.weight": "patch_embedding.conv.weight", "vision_model.embeddings.patch_embedding.bias": "patch_embedding.conv.bias", "vision_model.embeddings.position_embedding.weight": "patch_embedding.position_embedding", "vision_model.post_layernorm.weight": "layer_norm.weight", "vision_model.post_layernorm.bias": "layer_norm.bias", } for i in range(VIT_N_BLOCKS): mapping[f"vision_model.encoder.layers.{i}.layer_norm1.weight"] = f"blocks.{i}.ln1.weight" mapping[f"vision_model.encoder.layers.{i}.layer_norm1.bias"] = f"blocks.{i}.ln1.bias" mapping[f"vision_model.encoder.layers.{i}.layer_norm2.weight"] = f"blocks.{i}.ln2.weight" mapping[f"vision_model.encoder.layers.{i}.layer_norm2.bias"] = f"blocks.{i}.ln2.bias" mapping[f"vision_model.encoder.layers.{i}.mlp.fc1.weight"] = f"blocks.{i}.mlp.fc1.weight" mapping[f"vision_model.encoder.layers.{i}.mlp.fc1.bias"] = f"blocks.{i}.mlp.fc1.bias" mapping[f"vision_model.encoder.layers.{i}.mlp.fc2.weight"] = f"blocks.{i}.mlp.fc2.weight" mapping[f"vision_model.encoder.layers.{i}.mlp.fc2.bias"] = f"blocks.{i}.mlp.fc2.bias" mapping[f"vision_model.encoder.layers.{i}.self_attn.out_proj.weight"] = f"blocks.{i}.attn.out_proj.weight" mapping[f"vision_model.encoder.layers.{i}.self_attn.out_proj.bias"] = f"blocks.{i}.attn.out_proj.bias" with safetensors.safe_open(filename=safetensors_file, framework="pt", device="cpu") as f: for hf_key, our_key in mapping.items(): tensor = f.get_tensor(hf_key) if tensor.shape == sd[our_key].shape: sd[our_key].copy_(tensor) elif "position_embedding" in hf_key: sd[our_key].copy_(tensor.unsqueeze(0)) else: raise ValueError(f"shape mismatch {hf_key}: {tensor.shape} vs {sd[our_key].shape}") for i in range(VIT_N_BLOCKS): q = f.get_tensor(f"vision_model.encoder.layers.{i}.self_attn.q_proj.weight") k = f.get_tensor(f"vision_model.encoder.layers.{i}.self_attn.k_proj.weight") v = f.get_tensor(f"vision_model.encoder.layers.{i}.self_attn.v_proj.weight") sd[f"blocks.{i}.attn.qkv_proj.weight"].copy_(torch.cat([q, k, v], dim=0)) qb = f.get_tensor(f"vision_model.encoder.layers.{i}.self_attn.q_proj.bias") kb = f.get_tensor(f"vision_model.encoder.layers.{i}.self_attn.k_proj.bias") vb = f.get_tensor(f"vision_model.encoder.layers.{i}.self_attn.v_proj.bias") sd[f"blocks.{i}.attn.qkv_proj.bias"].copy_(torch.cat([qb, kb, vb], dim=0)) model.load_state_dict(sd) n_params = sum(p.numel() for p in model.parameters()) print(f"Loaded {VIT_MODEL_TYPE}: {n_params:,} params") return model class ModalityProjector(nn.Module): """Pixel-shuffle (x4) + linear projection into the target LLM's hidden size.""" def __init__(self, llm_hidden_size, scale_factor=4): super().__init__() self.scale_factor = scale_factor self.input_dim = VIT_HIDDEN_DIM * (scale_factor ** 2) self.output_dim = llm_hidden_size self.proj = nn.Linear(self.input_dim, self.output_dim, bias=False) nn.init.normal_(self.proj.weight, mean=0.0, std=0.02) def pixel_shuffle(self, x): bsz, seq, embed_dim = x.size() seq_root = int(seq ** 0.5) assert seq_root ** 2 == seq assert seq_root % self.scale_factor == 0 h = w = seq_root x = x.view(bsz, h, w, embed_dim) h_out, w_out = h // self.scale_factor, w // self.scale_factor x = x.reshape(bsz, h_out, self.scale_factor, w_out, self.scale_factor, embed_dim) x = x.permute(0, 1, 3, 2, 4, 5).contiguous() x = x.reshape(bsz, h_out * w_out, embed_dim * self.scale_factor ** 2) return x def forward(self, x): x = x.to(self.proj.weight.dtype) return self.proj(self.pixel_shuffle(x))