fount-vlm-nell-02 0.3.11__py3-none-any.whl

mlx_vlm/models/molmo/language.py

@@ -0,0 +1,208 @@
+from typing import Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from ..base import (
+    LanguageModelOutput,
+    create_attention_mask,
+    scaled_dot_product_attention,
+)
+from ..cache import KVCache
+from .config import TextConfig
+
+
+class SwiGLU(nn.Module):
+    def __call__(self, x: mx.array) -> mx.array:
+        x, gate = mx.split(x, 2, axis=-1)
+        return nn.silu(gate) * x
+
+
+class MolmoBlock(nn.Module):
+    def __init__(self, config: TextConfig):
+        super().__init__()
+        self.attn_out = nn.Linear(config.d_model, config.d_model, bias=False)
+        self.ff_out = nn.Linear(
+            int(config.act_output_multiplier * config.mlp_hidden_size),
+            config.d_model,
+            bias=False,
+        )
+        self.attn_norm = nn.RMSNorm(config.d_model, eps=config.layer_norm_eps)
+        self.ff_norm = nn.RMSNorm(config.d_model, eps=config.layer_norm_eps)
+        self.ff_proj = nn.Linear(config.d_model, config.mlp_hidden_size, bias=False)
+        head_dim = config.d_model // config.n_heads
+        self.rotary_emb = nn.RoPE(head_dim, base=config.rope_theta)
+        self.scale = head_dim**-0.5
+        self.n_heads = config.n_heads
+        self.n_kv_heads = config.n_kv_heads
+        self.fused_dims = (
+            config.d_model,
+            config.n_kv_heads * head_dim,
+            config.n_kv_heads * head_dim,
+        )
+        self.att_proj = nn.Linear(config.d_model, sum(self.fused_dims), bias=True)
+        self.act = SwiGLU()
+
+    def __call__(self, x, mask=None, cache=None):
+        batch_size, seq_len, D = x.shape
+        attn_in = self.attn_norm(x)
+
+        qkv = self.att_proj(attn_in)
+
+        q, k, v = mx.split(
+            qkv, [self.fused_dims[0], self.fused_dims[0] + self.fused_dims[1]], axis=-1
+        )
+
+        q = q.reshape(batch_size, seq_len, self.n_heads, D // self.n_heads).transpose(
+            0, 2, 1, 3
+        )
+        k = k.reshape(
+            batch_size, seq_len, self.n_kv_heads, D // self.n_heads
+        ).transpose(0, 2, 1, 3)
+        v = v.reshape(
+            batch_size, seq_len, self.n_kv_heads, D // self.n_heads
+        ).transpose(0, 2, 1, 3)
+
+        if cache is not None:
+            q = self.rotary_emb(q, offset=cache.offset)
+            k = self.rotary_emb(k, offset=cache.offset)
+            k, v = cache.update_and_fetch(k, v)
+        else:
+            q = self.rotary_emb(q)
+            k = self.rotary_emb(k)
+
+        # Perform attention
+        att = scaled_dot_product_attention(q, k, v, cache, scale=self.scale, mask=mask)
+        att = att.transpose(0, 2, 1, 3).reshape(batch_size, seq_len, D)
+        att = self.attn_out(att)
+
+        # Add attention scores
+        # shape: (batch_size, seq_len, d_model)
+        x = x + att
+
+        # Feed-forward layer
+        og_x = x
+        x = self.ff_norm(x)
+        x = self.ff_proj(x)
+        x = self.act(x)
+        x = self.ff_out(x)
+        x = og_x + x
+
+        return x
+
+
+class Embedding(nn.Module):
+    def __init__(
+        self,
+        num_embeddings: int,
+        num_new_embeddings: int,
+        features: int,
+        initializer_range: float = 0.02,
+        new_embed_initializer_range: float = 0.02,
+    ):
+        super().__init__()
+        self.initializer_range = initializer_range
+        self.new_embed_initializer_range = new_embed_initializer_range
+
+        # Initialize embeddings
+        self.embedding = mx.random.normal(
+            (num_embeddings, features), scale=self.initializer_range
+        )
+        self.new_embedding = mx.random.normal(
+            (num_new_embeddings, features), scale=self.new_embed_initializer_range
+        )
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return mx.concat([self.embedding, self.new_embedding], axis=0)[x]
+
+
+class Molmo(nn.Module):
+    def __init__(self, config: TextConfig):
+        super().__init__()
+        self.config = config
+
+        self.wte = Embedding(
+            config.embedding_size, config.additional_vocab_size, config.d_model
+        )
+
+        self.blocks = [MolmoBlock(config) for _ in range(config.n_layers)]
+
+        self.ln_f = nn.RMSNorm(config.d_model, eps=config.layer_norm_eps)
+
+        if not config.weight_tying:
+            self.ff_out = nn.Linear(config.d_model, config.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        input_ids: mx.array,
+        inputs_embeds: Optional[mx.array] = None,
+        mask: Optional[mx.array] = None,
+        cache: Optional[KVCache] = None,
+    ) -> LanguageModelOutput:
+
+        if inputs_embeds is None:
+            h = self.wte(input_ids)
+        else:
+            h = inputs_embeds
+
+        if cache is None:
+            cache = [None] * self.config.n_layers
+
+        if mask is None:
+            mask = create_attention_mask(h, cache)
+        else:
+            seq_len = h.shape[-2]
+            causal_mask = mx.tril(mx.ones((seq_len, seq_len), dtype=mx.bool_))
+            mask = (mask & causal_mask) > 0
+
+        for block, c in zip(self.blocks, cache):
+            h = block(h, mask, c)
+
+        h = self.ln_f(h)
+
+        if self.config.weight_tying:
+            logits = mx.matmul(h, self.wte.weight.T)
+        else:
+            logits = self.ff_out(h)
+
+        return LanguageModelOutput(logits=logits)
+
+
+class LanguageModel(nn.Module):
+    def __init__(self, config: TextConfig):
+        super().__init__()
+        self.config = config
+        self.model_type = config.model_type
+        if self.model_type != "molmo":
+            raise ValueError(
+                f"Model type {self.model_type} not supported. Currently only 'molmo' is supported"
+            )
+        self.model = Molmo(config)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        inputs_embeds: Optional[mx.array] = None,
+        mask: Optional[mx.array] = None,
+        cache: Optional[KVCache] = None,
+        **kwargs,
+    ) -> LanguageModelOutput:
+        outputs = self.model(inputs, inputs_embeds, mask, cache)
+        return outputs
+
+    @staticmethod
+    def sanitize(weights):
+        # Remove unused precomputed rotary freqs
+        return {k: v for k, v in weights.items() if "rotary_emb.inv_freq" not in k}
+
+    @property
+    def layers(self):
+        return self.model.blocks
+
+    @property
+    def head_dim(self):
+        return self.config.d_model // self.config.n_heads
+
+    @property
+    def n_kv_heads(self):
+        return self.config.n_kv_heads

mlx_vlm/models/molmo/molmo.py

@@ -0,0 +1,108 @@
+from typing import Dict, List, Tuple, Union
+
+import mlx.core as mx
+import mlx.nn as nn
+import numpy as np
+
+from ..base import InputEmbeddingsFeatures
+from .config import ModelConfig
+from .language import LanguageModel
+from .vision import VisionModel
+
+
+class Model(nn.Module):
+    def __init__(self, config: ModelConfig):
+        super().__init__()
+        self.config = config
+        self.language_model = LanguageModel(config.text_config)
+        self.vision_tower = VisionModel(config.vision_config)
+
+    @property
+    def layers(self):
+        return self.language_model.model.layers
+
+    def get_input_embeddings(
+        self,
+        input_ids: mx.array,
+        pixel_values: mx.array = None,
+        **kwargs,
+    ):
+        """Get input embeddings with image features merged in."""
+        if input_ids.ndim == 1:
+            input_ids = input_ids[None, :]
+
+        batch_size, seq_len = input_ids.shape
+
+        image_input_idx = kwargs.get("image_input_idx", None)
+        image_masks = kwargs.get("image_masks", None)
+
+        if pixel_values is None:
+            inputs_embeds = self.language_model.model.wte(input_ids)
+            return InputEmbeddingsFeatures(inputs_embeds=inputs_embeds)
+
+        # Process images
+        dtype = self.vision_tower.image_vit.patch_embedding.weight.dtype
+        pixel_values = pixel_values.astype(dtype)
+
+        if pixel_values.ndim == 3:
+            pixel_values = mx.expand_dims(pixel_values, 0)
+            image_masks = (
+                mx.expand_dims(image_masks, 0) if image_masks is not None else None
+            )
+            image_input_idx = (
+                mx.expand_dims(image_input_idx, 0)
+                if image_input_idx is not None
+                else None
+            )
+
+        image_features, cls_embed = self.vision_tower(pixel_values, image_masks)
+
+        # Insert image features into the input embeddings
+        num_image, num_patch = image_features.shape[1:3]
+
+        # Reshape for merging
+        image_features = image_features.reshape(batch_size, num_image * num_patch, -1)
+        image_input_idx = image_input_idx.reshape(batch_size, num_image * num_patch)
+
+        valid = np.where(image_input_idx >= 0)[0].tolist()
+        batch_idx = mx.arange(batch_size)
+        batch_idx = mx.tile(batch_idx[:, None], [1, image_features.shape[1]])
+
+        input_embeddings = self.language_model.model.wte(input_ids)
+        input_embeddings[batch_idx[valid], image_input_idx[valid]] += image_features[
+            valid
+        ]
+
+        return InputEmbeddingsFeatures(inputs_embeds=input_embeddings)
+
+    def __call__(
+        self,
+        input_ids: mx.array,
+        pixel_values: mx.array,
+        mask: mx.array,
+        cache=None,
+        **kwargs,
+    ) -> Dict[str, Union[mx.array, List[Tuple[mx.array, mx.array]]]]:
+        # Get input embeddings with image features merged
+        embedding_output = self.get_input_embeddings(input_ids, pixel_values, **kwargs)
+        input_embeddings = embedding_output.inputs_embeds
+
+        # Forward pass through the language model
+        logits = self.language_model(
+            input_ids,
+            inputs_embeds=input_embeddings,
+            mask=mask,
+            cache=cache,
+        )
+
+        return logits
+
+    def sanitize(self, weights):
+        def transform_key(key):
+            if "model.transformer" in key:
+                key = key.replace("model.transformer", "language_model.model")
+            if "model.vision_backbone" in key:
+                key = key.replace("model.vision_backbone", "vision_tower")
+            return key
+
+        return {transform_key(k): v for k, v in weights.items()}