Upload Moondream

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:15daa0c9a135d2084520e47d517ef689e09ee71c63d606bd8b8ff209bacc3e34
+oid sha256:95f9c7ac62666e56e45a66d4c07a0d110be1c5cfbe4eef1ae86857ef2787ce19
 size 3715037856

moondream.py

@@ -2,7 +2,6 @@ import torch
 from .vision_encoder import VisionEncoder
 from .configuration_moondream import MoondreamConfig
 from transformers import PreTrainedModel
-import re
 
 from .modeling_phi import PhiForCausalLM
 from .configuration_moondream import PhiConfig
@@ -62,16 +61,13 @@ class Moondream(PreTrainedModel):
         image_embeds,
         prompt,
         tokenizer,
-        eos_text="<END>",
         max_new_tokens=128,
         **kwargs,
     ):
-        eos_tokens = tokenizer(eos_text, add_special_tokens=False)[0].ids
-
         generate_config = {
-            "eos_token_id": eos_tokens,
+            "eos_token_id": tokenizer.eos_token_id,
             "bos_token_id": tokenizer.bos_token_id,
-            "pad_token_id": tokenizer.eos_token_id,
+            "pad_token_id": tokenizer.bos_token_id,
             "max_new_tokens": max_new_tokens,
             **kwargs,
         }
@@ -97,12 +93,11 @@ class Moondream(PreTrainedModel):
         answer = self.generate(
             image_embeds,
             prompt,
-            eos_text="<END>",
             tokenizer=tokenizer,
             max_new_tokens=512,
             **kwargs,
         )[0]
-        cleaned_answer = re.sub("<$|<END$", "", answer).strip()
+        cleaned_answer = answer.strip()
 
         # Use the result_queue to pass the result if it is provided
         if result_queue:
@@ -117,8 +112,6 @@ class Moondream(PreTrainedModel):
         tokenizer,
         **kwargs,
     ):
-        eos_tokens = tokenizer("<END>", add_special_tokens=False)[0].ids
-
         image_embeds = self.encode_image(images)
 
         templated_prompts = [
@@ -159,9 +152,9 @@ class Moondream(PreTrainedModel):
         )
 
         generate_config = {
-            "eos_token_id": eos_tokens,
+            "eos_token_id": tokenizer.eos_token_id,
             "bos_token_id": tokenizer.bos_token_id,
-            "pad_token_id": tokenizer.eos_token_id,
+            "pad_token_id": tokenizer.bos_token_id,
             "max_new_tokens": 512,
             **kwargs,
         }
@@ -174,6 +167,6 @@ class Moondream(PreTrainedModel):
             )
 
         return [
-            re.sub("<$|<END$", "", x).strip()
+            x.strip()
             for x in tokenizer.batch_decode(output_ids, skip_special_tokens=True)
         ]

vision_encoder.py

@@ -1,6 +1,6 @@
 import torch
+import torch.nn.functional as F
 from torch import nn
-from PIL import Image
 from einops import rearrange
 from torchvision.transforms.v2 import (
     Compose,
@@ -10,34 +10,92 @@ from torchvision.transforms.v2 import (
     ToDtype,
     Normalize,
 )
-import timm
 
 
-class VisualHolder(nn.Module):
-    def __init__(self, model):
+class Attention(nn.Module):
+    def __init__(self, dim, num_heads=16):
         super().__init__()
-        self.visual = model
+        assert dim % num_heads == 0, "dim should be divisible by num_heads"
+
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+
+        self.qkv = nn.Linear(dim, dim * 3)
+        self.proj = nn.Linear(dim, dim)
+
+        torch.nn.init.kaiming_normal_(
+            self.qkv.weight, mode="fan_in", nonlinearity="relu"
+        )
+        torch.nn.init.kaiming_normal_(
+            self.proj.weight, mode="fan_in", nonlinearity="relu"
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, N, C = x.shape
+        qkv = (
+            self.qkv(x)
+            .reshape(B, N, 3, self.num_heads, self.head_dim)
+            .permute(2, 0, 3, 1, 4)
+        )
+        q, k, v = qkv.unbind(0)
+
+        x = F.scaled_dot_product_attention(q, k, v)
+
+        x = x.transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        return x
+
+
+class VitBlock(nn.Module):
+    def __init__(self, embed_dim):
+        super().__init__()
+        self.attn = Attention(embed_dim)
+        self.mlp = MLP(embed_dim, 4304)
+        self.norm1 = nn.LayerNorm(embed_dim)
+        self.norm2 = nn.LayerNorm(embed_dim)
 
     def forward(self, x):
-        return self.visual(x)
+        x = x + self.attn(self.norm1(x))
+        x = x + self.mlp(self.norm2(x))
+        return x
 
 
-class ModelHolder(nn.Module):
-    def __init__(self, model):
+class VisionTransformer(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+
+        embed_len = 729
+        embed_dim = 1152
+
+        self.patch_embed = LinearPatchEmbedding()
+        self.pos_embed = nn.Parameter(torch.randn(1, embed_len, embed_dim) * 0.02)
+        self.blocks = nn.Sequential(*[VitBlock(embed_dim) for _ in range(27)])
+        self.norm = nn.LayerNorm(embed_dim)
+
+    def forward(self, x):
+        x = self.patch_embed(x)
+        x = x + self.pos_embed
+        for block in self.blocks:
+            x = block(x)
+        return self.norm(x)
+
+
+class EncoderWrapper(nn.Module):
+
+    def __init__(self):
         super().__init__()
-        self.model = model
+        self.model = nn.ModuleDict({"visual": VisionTransformer()})
 
     def forward(self, x):
-        return self.model(x)
+        return self.model["visual"](x)
 
 
 class LinearPatchEmbedding(nn.Module):
-    def __init__(self, conv):
+
+    def __init__(self):
         super().__init__()
         self.linear = nn.Linear(588, 1152)
-        self.linear.weight.data = conv.weight.data.view(1152, -1)
-        if conv.bias is not None:
-            self.linear.bias.data = conv.bias.data
 
     def forward(self, x):
         return self.linear(x)
@@ -49,13 +107,12 @@ class MLP(nn.Module):
         in_features: int,
         hidden_features: int = None,
         out_features: int = None,
-        act_layer: nn.Module = nn.GELU,
     ) -> None:
         super().__init__()
         out_features = out_features or in_features
         hidden_features = hidden_features or in_features
         self.fc1 = nn.Linear(in_features, hidden_features)
-        self.act = act_layer()
+        self.act = nn.GELU(approximate="tanh")
         self.fc2 = nn.Linear(hidden_features, out_features)
 
         torch.nn.init.kaiming_normal_(
@@ -94,14 +151,7 @@ class VisionEncoder(nn.Module):
     def __init__(self) -> None:
         super().__init__()
 
-        self.encoder = ModelHolder(
-            VisualHolder(timm.create_model("vit_so400m_patch14_siglip_384"))
-        )
-        self.encoder.model.visual.patch_embed = LinearPatchEmbedding(
-            self.encoder.model.visual.patch_embed.proj
-        )
-        self.encoder.model.visual.attn_pool = nn.Identity()
-
+        self.encoder = EncoderWrapper()
         self.projection = VisionProjection()
 
         self.preprocess = Compose(