Upload Moondream

config.json

@@ -1,4 +1,5 @@
 {
+  "_name_or_path": "vikhyatk/moondream2",
   "architectures": [
     "Moondream"
   ],
@@ -10,6 +11,20 @@
   "phi_config": {
     "model_type": "phi"
   },
+  "text_config": {
+    "architectures": [
+      "Moondream"
+    ],
+    "auto_map": {
+      "AutoConfig": "configuration_moondream.MoondreamConfig",
+      "AutoModelForCausalLM": "moondream.Moondream"
+    },
+    "model_type": "phi",
+    "phi_config": {
+      "model_type": "phi"
+    },
+    "torch_dtype": "float16"
+  },
   "torch_dtype": "float16",
   "transformers_version": "4.38.2"
 }

configuration_moondream.py

@@ -94,5 +94,5 @@ class MoondreamConfig(PretrainedConfig):
     model_type = "moondream1"
 
     def __init__(self, **kwargs):
-        self.phi_config = PhiConfig(**kwargs)
+        self.text_config = PhiConfig(**kwargs)
         super().__init__(**kwargs)

modeling_phi.py

@@ -400,40 +400,10 @@ class PhiAttention(nn.Module):
         key_states = repeat_kv(key_states, self.num_key_value_groups)
         value_states = repeat_kv(value_states, self.num_key_value_groups)
 
-        # Queries and keys upcast to fp32 is required by Phi-2 to avoid overflow
-        attn_weights = torch.matmul(
-            query_states.to(torch.float32), key_states.to(torch.float32).transpose(2, 3)
-        ) / math.sqrt(self.head_dim)
-
-        if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
-            raise ValueError(
-                f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"
-                f" {attn_weights.size()}"
-            )
-
-        if attention_mask is not None:
-            if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
-                raise ValueError(
-                    f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
-                )
-            attn_weights = attn_weights + attention_mask
-
-        # upcast attention to fp32
-        attn_weights = nn.functional.softmax(
-            attn_weights, dim=-1, dtype=torch.float32
-        ).to(value_states.dtype)
-        attn_weights = nn.functional.dropout(
-            attn_weights, p=self.attention_dropout, training=self.training
+        attn_output = torch.nn.functional.scaled_dot_product_attention(
+            query_states, key_states, value_states, attn_mask=attention_mask
         )
 
-        attn_output = torch.matmul(attn_weights, value_states)
-
-        if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
-            raise ValueError(
-                f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
-                f" {attn_output.size()}"
-            )
-
         attn_output = attn_output.transpose(1, 2).contiguous()
         attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
 
@@ -1115,7 +1085,6 @@ class PhiForCausalLM(PhiPreTrainedModel):
 
         hidden_states = outputs[0]
         logits = self.lm_head(hidden_states)
-        logits = logits.float()
 
         loss = None
         if labels is not None:

moondream.py

@@ -12,14 +12,16 @@ class Moondream(PreTrainedModel):
 
     def __init__(self, config):
         super().__init__(config)
-        self.vision_encoder = VisionEncoder()
+        self.vision_encoder = VisionEncoder(
+            use_flash_attn=config._attn_implementation == "flash_attention_2"
+        )
 
-        if type(config.phi_config) == dict:
+        if type(config.text_config) == dict:
             phi_config = PhiConfig(
-                **config.phi_config, attn_implementation=config._attn_implementation
+                **config.text_config, attn_implementation=config._attn_implementation
             )
         else:
-            phi_config = config.phi_config
+            phi_config = config.text_config
         self.text_model = PhiForCausalLM(phi_config)
 
     @property

vision_encoder.py

@@ -10,10 +10,20 @@ from torchvision.transforms.v2 import (
     ToDtype,
     Normalize,
 )
+from transformers.utils import is_flash_attn_2_available
+
+try:
+    if is_flash_attn_2_available():
+        from flash_attn.modules.mha import FlashSelfAttention
+    else:
+        FlashSelfAttention = None
+except ImportError:
+    FlashSelfAttention = None
 
 
 class Attention(nn.Module):
-    def __init__(self, dim, num_heads=16):
+
+    def __init__(self, dim, num_heads=16, use_flash_attn=False):
         super().__init__()
         assert dim % num_heads == 0, "dim should be divisible by num_heads"
 
@@ -23,6 +33,11 @@ class Attention(nn.Module):
         self.qkv = nn.Linear(dim, dim * 3)
         self.proj = nn.Linear(dim, dim)
 
+        if use_flash_attn and FlashSelfAttention is not None:
+            self.flash_attn = FlashSelfAttention()
+        else:
+            self.flash_attn = None
+
         torch.nn.init.kaiming_normal_(
             self.qkv.weight, mode="fan_in", nonlinearity="relu"
         )
@@ -31,25 +46,36 @@ class Attention(nn.Module):
         )
 
     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
+        if self.flash_attn is not None:
+            qkv = self.qkv(x)
+            qkv = rearrange(
+                qkv, "... (three h d) -> ... three h d", three=3, h=self.num_heads
+            )
+            attn_output = self.flash_attn(qkv)
+            output = rearrange(attn_output, "... h d -> ... (h d)")
+            output = self.proj(output)
+            return output
+        else:
+            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):
+
+    def __init__(self, embed_dim, use_flash_attn=False):
         super().__init__()
-        self.attn = Attention(embed_dim)
+        self.attn = Attention(embed_dim, use_flash_attn=use_flash_attn)
         self.mlp = MLP(embed_dim, 4304)
         self.norm1 = nn.LayerNorm(embed_dim)
         self.norm2 = nn.LayerNorm(embed_dim)
@@ -62,7 +88,7 @@ class VitBlock(nn.Module):
 
 class VisionTransformer(nn.Module):
 
-    def __init__(self):
+    def __init__(self, use_flash_attn=False):
         super().__init__()
 
         embed_len = 729
@@ -70,7 +96,9 @@ class VisionTransformer(nn.Module):
 
         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.blocks = nn.Sequential(
+            *[VitBlock(embed_dim, use_flash_attn=use_flash_attn) for _ in range(27)]
+        )
         self.norm = nn.LayerNorm(embed_dim)
 
     def forward(self, x):
@@ -83,9 +111,9 @@ class VisionTransformer(nn.Module):
 
 class EncoderWrapper(nn.Module):
 
-    def __init__(self):
+    def __init__(self, use_flash_attn=False):
         super().__init__()
-        self.model = nn.ModuleDict({"visual": VisionTransformer()})
+        self.model = nn.ModuleDict({"visual": VisionTransformer(use_flash_attn)})
 
     def forward(self, x):
         return self.model["visual"](x)
@@ -98,6 +126,13 @@ class LinearPatchEmbedding(nn.Module):
         self.linear = nn.Linear(588, 1152)
 
     def forward(self, x):
+        b, c, hp1, wp2 = x.shape
+        p1, p2 = 14, 14
+        h, w = hp1 // p1, wp2 // p2
+        x = x.reshape(b, c, h, p1, w, p2)
+        x = x.permute(0, 2, 4, 1, 3, 5)
+        x = x.reshape(b, h * w, c * p1 * p2)
+
         return self.linear(x)
 
 
@@ -148,10 +183,11 @@ class VisionProjection(nn.Module):
 
 
 class VisionEncoder(nn.Module):
-    def __init__(self) -> None:
+
+    def __init__(self, use_flash_attn=False):
         super().__init__()
 
-        self.encoder = EncoderWrapper()
+        self.encoder = EncoderWrapper(use_flash_attn)
         self.projection = VisionProjection()
 
         self.preprocess = Compose(
@@ -172,16 +208,20 @@ class VisionEncoder(nn.Module):
         return self.projection.mlp.fc1.weight.dtype
 
     def __call__(self, images) -> torch.Tensor:
-        if not isinstance(images, list):
+        if not isinstance(images, list) and not isinstance(images, torch.Tensor):
             images = [images]
 
         with torch.no_grad():
-            x = torch.stack(
-                [self.preprocess(image.convert("RGB")) for image in images]
-            ).to(self.device, dtype=self.dtype)
+            # Skip preprocess if images are already tensors
+            if not isinstance(images, torch.Tensor) and not isinstance(
+                images[0], torch.Tensor
+            ):
+                images = [self.preprocess(image.convert("RGB")) for image in images]
 
-            x = rearrange(x, "b c (h p1) (w p2) -> b (h w) (c p1 p2)", p1=14, p2=14)
+            if isinstance(images, list):
+                images = torch.stack(images)
 
+            x = images.to(self.device, dtype=self.dtype)
             x = self.encoder(x)
             x = self.projection(x)