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

mlx_vlm/models/molmo2/language.py

@@ -0,0 +1,206 @@
+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 ModelConfig, TextConfig
+
+
+class Molmo2Embedding(nn.Module):
+    def __init__(
+        self,
+        num_embeddings: int,
+        num_new_embeddings: int,
+        features: int,
+    ):
+        super().__init__()
+        self.embedding = mx.zeros((num_embeddings, features))
+        self.new_embedding = mx.zeros((num_new_embeddings, features))
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return mx.concatenate([self.embedding, self.new_embedding], axis=0)[x]
+
+
+class LanguageModelMLP(nn.Module):
+    def __init__(self, input_dim: int, intermediate_size: int):
+        super().__init__()
+        self.ff_proj = nn.Linear(input_dim, intermediate_size * 2, bias=False)
+        self.ff_out = nn.Linear(intermediate_size, input_dim, bias=False)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        x = self.ff_proj(x)
+        x, gate = mx.split(x, 2, axis=-1)
+        x = nn.silu(gate) * x
+        return self.ff_out(x)
+
+
+class Molmo2Attention(nn.Module):
+    def __init__(self, config: TextConfig):
+        super().__init__()
+        self.config = config
+        self.num_heads = config.num_attention_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        self.head_dim = config.head_dim
+        self.scale = self.head_dim**-0.5
+
+        self.fused_dims = (
+            config.num_attention_heads * config.head_dim,
+            config.head_dim * config.num_key_value_heads,
+            config.head_dim * config.num_key_value_heads,
+        )
+
+        self.att_proj = nn.Linear(
+            config.hidden_size,
+            sum(self.fused_dims),
+            bias=config.qkv_bias,
+        )
+
+        self.q_norm = nn.RMSNorm(dims=config.head_dim, eps=config.layer_norm_eps)
+        self.k_norm = nn.RMSNorm(dims=config.head_dim, eps=config.layer_norm_eps)
+
+        self.attn_out = nn.Linear(
+            config.head_dim * config.num_attention_heads,
+            config.hidden_size,
+            bias=False,
+        )
+
+        self.rotary_emb = nn.RoPE(self.head_dim, base=config.rope_theta)
+
+    def __call__(
+        self,
+        hidden_states: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[KVCache] = None,
+    ) -> mx.array:
+        batch_size, seq_len, _ = hidden_states.shape
+
+        qkv = self.att_proj(hidden_states)
+        q, k, v = mx.split(
+            qkv,
+            [self.fused_dims[0], self.fused_dims[0] + self.fused_dims[1]],
+            axis=-1,
+        )
+
+        q = self.q_norm(q.reshape(batch_size, seq_len, self.num_heads, self.head_dim))
+        k = self.k_norm(
+            k.reshape(batch_size, seq_len, self.num_key_value_heads, self.head_dim)
+        )
+        v = v.reshape(batch_size, seq_len, self.num_key_value_heads, self.head_dim)
+
+        q = q.transpose(0, 2, 1, 3)
+        k = k.transpose(0, 2, 1, 3)
+        v = v.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)
+
+        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, -1)
+        return self.attn_out(att)
+
+
+class Molmo2DecoderLayer(nn.Module):
+    def __init__(self, config: TextConfig):
+        super().__init__()
+        self.self_attn = Molmo2Attention(config)
+        self.attn_norm = nn.RMSNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.ff_norm = nn.RMSNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.mlp = LanguageModelMLP(config.hidden_size, config.intermediate_size)
+
+    def __call__(
+        self,
+        hidden_states: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[KVCache] = None,
+    ) -> mx.array:
+        residual = hidden_states
+        hidden_states = self.attn_norm(hidden_states)
+        hidden_states = residual + self.self_attn(hidden_states, mask, cache)
+
+        residual = hidden_states
+        hidden_states = self.ff_norm(hidden_states)
+        hidden_states = residual + self.mlp(hidden_states)
+        return hidden_states
+
+
+class Molmo2Transformer(nn.Module):
+    def __init__(self, config: TextConfig):
+        super().__init__()
+        self.config = config
+
+        self.wte = Molmo2Embedding(
+            config.vocab_size, config.additional_vocab_size, config.hidden_size
+        )
+        self.blocks = [
+            Molmo2DecoderLayer(config) for _ in range(config.num_hidden_layers)
+        ]
+        self.ln_f = nn.RMSNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+        self.emb_drop = nn.Dropout(config.embedding_dropout)
+
+    def __call__(
+        self,
+        input_ids: mx.array,
+        inputs_embeds: Optional[mx.array] = None,
+        mask: Optional[mx.array] = None,
+        cache: Optional[list[KVCache]] = None,
+    ) -> mx.array:
+        if inputs_embeds is None:
+            hidden_states = self.wte(input_ids)
+        else:
+            hidden_states = inputs_embeds
+
+        if cache is None:
+            cache = [None] * len(self.blocks)
+
+        if mask is None:
+            mask = create_attention_mask(hidden_states, cache)
+
+        hidden_states = self.emb_drop(hidden_states)
+
+        for block, c in zip(self.blocks, cache):
+            hidden_states = block(hidden_states, mask, c)
+
+        return self.ln_f(hidden_states)
+
+
+class LanguageModel(nn.Module):
+    def __init__(self, args: TextConfig, config: ModelConfig = None):
+        super().__init__()
+        self.args = args
+        self.config = config
+        self.model_type = args.model_type
+        self.model = Molmo2Transformer(args)
+        self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        inputs_embeds: Optional[mx.array] = None,
+        mask: Optional[mx.array] = None,
+        cache: Optional[list[KVCache]] = None,
+        **kwargs,
+    ) -> LanguageModelOutput:
+        hidden_states = self.model(inputs, inputs_embeds, mask, cache)
+        logits = self.lm_head(hidden_states)
+        return LanguageModelOutput(logits=logits)
+
+    @staticmethod
+    def sanitize(weights):
+        # Remove unused precomputed rotary freqs if present.
+        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

mlx_vlm/models/molmo2/molmo2.py

@@ -0,0 +1,330 @@
+from typing import Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+import numpy as np
+
+from ..base import InputEmbeddingsFeatures, LanguageModelOutput
+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, config)
+        self.vision_tower = VisionModel(config.vision_config)
+
+    @property
+    def layers(self):
+        return self.language_model.layers
+
+    def build_batched_images(
+        self,
+        input_ids: mx.array,
+        pixel_values: mx.array,
+        image_token_pooling: mx.array,
+        image_grids: mx.array,
+        image_num_crops: mx.array,
+    ) -> tuple[mx.array, mx.array]:
+        raw_counts = (input_ids == self.config.image_end_token_id).sum(axis=1)
+        counts = raw_counts // 2
+        batch_size = counts.shape[0]
+
+        num_images = int(counts.sum().item())
+
+        if image_grids.shape[0] != num_images:
+            raise ValueError(
+                f"Expected {num_images} image grids, got {image_grids.shape[0]}"
+            )
+        if image_num_crops.shape[0] != num_images:
+            raise ValueError(
+                f"Expected {num_images} image crop counts, got {image_num_crops.shape[0]}"
+            )
+
+        num_pooled_patches_per_image = (
+            (image_grids[:, :2].prod(axis=1) + image_grids[:, 2:].prod(axis=1))
+            .astype(image_num_crops.dtype)
+            .reshape(-1)
+        )
+
+        n_crops, n_patches, pixels_per_patch = pixel_values.shape
+
+        example_ids_for_image = mx.array(
+            np.repeat(
+                np.arange(batch_size), np.array(counts).astype(np.int32).tolist()
+            ),
+            dtype=mx.int32,
+        )
+
+        crops_per_example = mx.zeros((batch_size,), dtype=image_num_crops.dtype)
+        pooled_per_example = mx.zeros(
+            (batch_size,), dtype=num_pooled_patches_per_image.dtype
+        )
+        for image_idx in range(num_images):
+            ex = int(example_ids_for_image[image_idx].item())
+            crops_per_example[ex] = crops_per_example[ex] + image_num_crops[image_idx]
+            pooled_per_example[ex] = (
+                pooled_per_example[ex] + num_pooled_patches_per_image[image_idx]
+            )
+
+        total_crops = int(crops_per_example.sum().item())
+        if total_crops != n_crops:
+            raise ValueError(f"Expected {total_crops} crops, got {n_crops}")
+
+        total_pooled = int(pooled_per_example.sum().item())
+        if total_pooled != image_token_pooling.shape[0]:
+            raise ValueError(
+                f"Expected {total_pooled} pooled patches, got {image_token_pooling.shape[0]}"
+            )
+
+        max_crops = int(crops_per_example.max().item())
+        images = mx.full(
+            (batch_size, max_crops, n_patches, pixels_per_patch),
+            vals=-1,
+            dtype=pixel_values.dtype,
+        )
+
+        offset_crop = 0
+        for i in range(batch_size):
+            num = int(crops_per_example[i].item())
+            images[i, :num] = pixel_values[offset_crop : offset_crop + num]
+            offset_crop += num
+
+        max_pooled = int(pooled_per_example.max().item())
+        token_dim = image_token_pooling.shape[1]
+        new_token_pooling = mx.full(
+            (batch_size, max_pooled, token_dim),
+            vals=-1,
+            dtype=image_token_pooling.dtype,
+        )
+
+        patches_per_image = image_num_crops * n_patches
+        counts_list = counts.tolist()
+        image_idx = 0
+        pooled_offset = 0
+        patch_offset = 0
+        for ex, c in enumerate(counts_list):
+            num_pooled = int(pooled_per_example[ex].item())
+            cur = mx.array(
+                image_token_pooling[pooled_offset : pooled_offset + num_pooled]
+            )
+
+            per_img_patches = patches_per_image[image_idx : image_idx + c]
+            index_offsets = [0] + np.cumsum(per_img_patches.tolist()).tolist()[:-1]
+            per_img_pooled = num_pooled_patches_per_image[
+                image_idx : image_idx + c
+            ].tolist()
+
+            offset = 0
+            for j in range(c):
+                n = int(per_img_pooled[j])
+                idx_off = int(index_offsets[j])
+                cur_slice = cur[offset : offset + n]
+                cur[offset : offset + n] = mx.where(
+                    cur_slice >= 0,
+                    cur_slice + idx_off,
+                    cur_slice,
+                )
+                offset += n
+
+            new_token_pooling[ex, :num_pooled] = cur
+            pooled_offset += num_pooled
+            image_idx += c
+            patch_offset += num_pooled
+
+        return images, new_token_pooling
+
+    def build_batched_videos(
+        self,
+        input_ids: mx.array,
+        pixel_values_videos: mx.array,
+        video_token_pooling: mx.array,
+        video_grids: mx.array,
+    ) -> tuple[mx.array, mx.array]:
+        end_token_id = (
+            self.config.frame_end_token_id
+            if self.config.use_frame_special_tokens
+            else self.config.image_end_token_id
+        )
+        counts = mx.any(input_ids == end_token_id, axis=1).astype(mx.int32)
+        batch_size = counts.shape[0]
+        num_videos = int(counts.sum().item())
+
+        if video_grids.shape[0] != num_videos:
+            raise ValueError(
+                f"Expected {num_videos} videos, got {video_grids.shape[0]}"
+            )
+
+        num_pooled_patches_per_video = (video_grids[:, 1] * video_grids[:, 2]).astype(
+            video_token_pooling.dtype
+        )
+
+        n_frames, n_patches, pixels_per_patch = pixel_values_videos.shape
+
+        frames_per_example = mx.zeros((batch_size,), dtype=mx.int32)
+        pooled_per_example = mx.zeros((batch_size,), dtype=video_token_pooling.dtype)
+
+        video_index = 0
+        for i in range(batch_size):
+            if counts[i].item() == 1:
+                frames_per_example[i] = int(video_grids[video_index][0].item())
+                pooled_per_example[i] = num_pooled_patches_per_video[video_index]
+                video_index += 1
+
+        max_frames = int(frames_per_example.max().item()) if num_videos else 0
+        videos = mx.full(
+            (batch_size, max_frames, n_patches, pixels_per_patch),
+            vals=-1,
+            dtype=pixel_values_videos.dtype,
+        )
+
+        offset = 0
+        for i in range(batch_size):
+            num = int(frames_per_example[i].item())
+            if num > 0:
+                videos[i, :num] = pixel_values_videos[offset : offset + num]
+                offset += num
+
+        max_pooled = int(pooled_per_example.max().item()) if num_videos else 0
+        token_dim = video_token_pooling.shape[1]
+        new_token_pooling = mx.full(
+            (batch_size, max_pooled, token_dim),
+            vals=-1,
+            dtype=video_token_pooling.dtype,
+        )
+
+        pooled_offset = 0
+        for i in range(batch_size):
+            num = int(pooled_per_example[i].item())
+            if num > 0:
+                new_token_pooling[i, :num] = video_token_pooling[
+                    pooled_offset : pooled_offset + num
+                ]
+                pooled_offset += num
+
+        if offset != n_frames:
+            raise ValueError(f"Expected {n_frames} frames, got {offset}")
+        if pooled_offset != video_token_pooling.shape[0]:
+            raise ValueError(
+                f"Expected {video_token_pooling.shape[0]} pooled tokens, got {pooled_offset}"
+            )
+
+        return videos, new_token_pooling
+
+    def merge_visual_inputs(
+        self,
+        *,
+        input_ids: mx.array,
+        pixel_values: Optional[mx.array] = None,
+        image_token_pooling: Optional[mx.array] = None,
+        image_grids: Optional[mx.array] = None,
+        image_num_crops: Optional[mx.array] = None,
+        video_token_pooling: Optional[mx.array] = None,
+        video_grids: Optional[mx.array] = None,
+    ) -> tuple[Optional[mx.array], Optional[mx.array]]:
+        if pixel_values is None:
+            return None, None
+
+        if video_token_pooling is not None or video_grids is not None:
+            if video_token_pooling is None or video_grids is None:
+                raise ValueError(
+                    "video_token_pooling and video_grids are required for videos"
+                )
+            return self.build_batched_videos(
+                input_ids=input_ids,
+                pixel_values_videos=pixel_values,
+                video_token_pooling=video_token_pooling,
+                video_grids=video_grids,
+            )
+
+        if (
+            image_token_pooling is None
+            or image_grids is None
+            or image_num_crops is None
+        ):
+            raise ValueError(
+                "image_token_pooling, image_grids, and image_num_crops are required for images"
+            )
+
+        return self.build_batched_images(
+            input_ids=input_ids,
+            pixel_values=pixel_values,
+            image_token_pooling=image_token_pooling,
+            image_grids=image_grids,
+            image_num_crops=image_num_crops,
+        )
+
+    def get_input_embeddings(
+        self,
+        input_ids: mx.array,
+        pixel_values: Optional[mx.array] = None,
+        **kwargs,
+    ) -> mx.array:
+        input_ids = input_ids * (input_ids != -1).astype(input_ids.dtype)
+        x = self.language_model.model.wte(input_ids)
+
+        if pixel_values is not None:
+
+            pixel_values, token_pooling = self.merge_visual_inputs(
+                input_ids=input_ids,
+                pixel_values=pixel_values,
+                image_token_pooling=kwargs.get("image_token_pooling", None),
+                image_grids=kwargs.get("image_grids", None),
+                image_num_crops=kwargs.get("image_num_crops", None),
+                video_token_pooling=kwargs.get("video_token_pooling", None),
+                video_grids=kwargs.get("video_grids", None),
+            )
+
+            dtype = self.vision_tower.image_vit.patch_embedding.weight.dtype
+            pixel_values = pixel_values.astype(dtype)
+            image_features = self.vision_tower(pixel_values, token_pooling)
+            is_image_patch = mx.reshape(input_ids, (-1,)) == self.config.image_patch_id
+            if int(is_image_patch.sum().item()) != image_features.shape[0]:
+                raise ValueError(
+                    f"Expected {int(is_image_patch.sum().item())} image features, got {image_features.shape[0]}"
+                )
+            flat_x = mx.reshape(x, (-1, x.shape[-1]))
+            positions = mx.array(np.where(np.array(is_image_patch))[0], dtype=mx.uint32)
+            flat_x[positions] = flat_x[positions] + image_features
+            x = flat_x.reshape(x.shape)
+
+        return InputEmbeddingsFeatures(inputs_embeds=x)
+
+    def __call__(
+        self,
+        input_ids: mx.array,
+        pixel_values: Optional[mx.array] = None,
+        mask: Optional[mx.array] = None,
+        cache=None,
+        **kwargs,
+    ) -> LanguageModelOutput:
+        if input_ids.ndim == 1:
+            input_ids = input_ids[None, :]
+
+        input_embeddings_features = self.get_input_embeddings(
+            input_ids=input_ids, pixel_values=pixel_values, **kwargs
+        )
+
+        return self.language_model(
+            input_ids,
+            inputs_embeds=input_embeddings_features.inputs_embeds,
+            mask=mask,
+            cache=cache,
+        )
+
+    def sanitize(self, weights):
+        def transform_key(key: str) -> str:
+            if key.startswith("model.transformer."):
+                key = key.replace("model.transformer.", "language_model.model.", 1)
+            if key.startswith("model.vision_backbone."):
+                key = key.replace("model.vision_backbone.", "vision_tower.", 1)
+            if key.startswith("lm_head."):
+                key = key.replace("lm_head.", "language_model.lm_head.", 1)
+            # Vision transformer uses list not named submodule
+            key = key.replace(".transformer.resblocks.", ".transformer.")
+            return key
+
+        return {transform_key(k): v for k, v in weights.items()}