Upload vision_encoder.py with huggingface_hub

vision_encoder.py

@@ -0,0 +1,167 @@
+"""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))