Upload Moondream

config.json

@@ -10,6 +10,6 @@
   "phi_config": {
     "model_type": "phi-msft"
   },
-  "torch_dtype": "float32",
+  "torch_dtype": "float16",
   "transformers_version": "4.36.2"
 }

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:892e51df302d98a83974761c4f386caddbad2edd0e84f228d9935b4aed33ee25
+size 3715037856

moondream.py

@@ -1,10 +1,12 @@
 import torch
+from torch import nn
 from .vision_encoder import VisionEncoder
-from .text_model import TextModel
 from .configuration_moondream import MoondreamConfig
 from transformers import PreTrainedModel
 import re
 
+from .modeling_phi import PhiForCausalLM
+from .configuration_moondream import PhiConfig
 
 class Moondream(PreTrainedModel):
     config_class = MoondreamConfig
@@ -12,11 +14,16 @@ class Moondream(PreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
         self.vision_encoder = VisionEncoder()
-        self.text_model = TextModel(config)
+
+        if type(config.phi_config) == dict:
+            phi_config = PhiConfig(**config.phi_config)
+        else:
+            phi_config = config.phi_config
+        self.text_model = PhiForCausalLM(phi_config)
 
     @property
     def device(self):
-        return self.text_model.model.device
+        return self.text_model.device
 
     def encode_image(self, image):
         return self.vision_encoder(image)
@@ -27,22 +34,22 @@ class Moondream(PreTrainedModel):
                 txt, return_tensors="pt", add_special_tokens=False
             ).input_ids.to(self.device)
 
+        text_emb = self.text_model.get_input_embeddings()
+
         # Add BOS token
         embeds = []
         embeds.append(
-            self.text_model.text_emb(
-                (torch.tensor([[tokenizer.bos_token_id]], device=self.device))
-            )
+            text_emb((torch.tensor([[tokenizer.bos_token_id]], device=self.device)))
         )
 
         if "<image>" not in prompt:
-            embeds.append(self.text_model.text_emb(_tokenize(prompt)))
+            embeds.append(text_emb(_tokenize(prompt)))
         else:
             assert prompt.count("<image>") == 1
             before, after = prompt.split("<image>")
-            embeds.append(self.text_model.text_emb(_tokenize(f"{before}<image>")))
+            embeds.append(text_emb(_tokenize(f"{before}<image>")))
             embeds.append(image_embeds.to(self.device))
-            embeds.append(self.text_model.text_emb(_tokenize(f"</image>{after}")))
+            embeds.append(text_emb(_tokenize(f"</image>{after}")))
 
         return torch.cat(embeds, dim=1)
 
@@ -67,7 +74,7 @@ class Moondream(PreTrainedModel):
 
         with torch.no_grad():
             inputs_embeds = self.input_embeds(prompt, image_embeds, tokenizer)
-            output_ids = self.text_model.model.generate(
+            output_ids = self.text_model.generate(
                 inputs_embeds=inputs_embeds, **generate_config
             )
 

vision_encoder.py

@@ -80,23 +80,18 @@ class VisionProjection(nn.Module):
         model_dim = 2048
         hidden_dim = model_dim * 4
 
-        self.mlp1 = MLP(image_embedding_dim, hidden_dim, model_dim)
-        self.mlp2 = MLP(model_dim, hidden_dim, model_dim)
-        self.ln = nn.LayerNorm(model_dim)
+        self.mlp = MLP(image_embedding_dim, hidden_dim, model_dim)
 
     @property
     def device(self):
-        return self.mlp1.fc1.weight.device
+        return self.mlp.fc1.weight.device
 
     def forward(self, x):
-        x = self.mlp1(x)
-        x = self.ln(x)
-        x = x + self.mlp2(x)
-        return x
+        return self.mlp(x)
 
 
-class VisionTower(nn.Module):
-    def __init__(self):
+class VisionEncoder(nn.Module):
+    def __init__(self) -> None:
         super().__init__()
 
         self.encoder = ModelHolder(
@@ -109,17 +104,6 @@ class VisionTower(nn.Module):
 
         self.projection = VisionProjection()
 
-    def forward(self, x):
-        x = self.encoder(x)
-        x = self.projection(x)
-        return x
-
-
-class VisionEncoder(nn.Module):
-    def __init__(self) -> None:
-        super().__init__()
-
-        self.model = VisionTower()
         self.preprocess = Compose(
             [
                 Resize(size=(378, 378), interpolation=InterpolationMode.BICUBIC),
@@ -131,20 +115,22 @@ class VisionEncoder(nn.Module):
 
     @property
     def device(self):
-        return self.model.projection.mlp1.fc1.weight.device
+        return self.projection.mlp.fc1.weight.device
 
     @property
     def dtype(self):
-        return self.model.projection.mlp1.fc1.weight.dtype
+        return self.projection.mlp.fc1.weight.dtype
 
     def __call__(self, image: Image) -> torch.Tensor:
         with torch.no_grad():
-            image_vec = (
+            x = (
                 self.preprocess(image.convert("RGB"))
                 .unsqueeze(0)
                 .to(self.device, dtype=self.dtype)
             )
-            image_vec = rearrange(
-                image_vec, "b c (h p1) (w p2) -> b (h w) (c p1 p2)", p1=14, p2=14
-            )
-            return self.model(image_vec)
+            x = rearrange(x, "b c (h p1) (w p2) -> b (h w) (c p1 p2)", p1=14, p2=14)
+
+            x = self.encoder(x)
+            x = self.projection(x)
+
+            return x