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

mlx_vlm/models/hunyuan_vl/vision.py

@@ -0,0 +1,322 @@
+import mlx.core as mx
+import mlx.nn as nn
+
+from ..base import chunked_attention
+from .config import VisionConfig
+
+
+class VisionMLP(nn.Module):
+    def __init__(self, config: VisionConfig):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.dense_h_to_4h = nn.Linear(
+            self.hidden_size, self.intermediate_size, bias=True
+        )
+        self.dense_4h_to_h = nn.Linear(
+            self.intermediate_size, self.hidden_size, bias=True
+        )
+
+    def __call__(self, x: mx.array) -> mx.array:
+        x = self.dense_h_to_4h(x)
+        x = nn.gelu(x)
+        x = self.dense_4h_to_h(x)
+        return x
+
+
+class VisionAttention(nn.Module):
+
+    def __init__(self, config: VisionConfig):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = self.hidden_size // self.num_heads
+        self.scale = self.head_dim**-0.5
+
+        self.q_proj = nn.Linear(
+            config.hidden_size, self.num_heads * self.head_dim, bias=True
+        )
+        self.k_proj = nn.Linear(
+            config.hidden_size, self.num_heads * self.head_dim, bias=True
+        )
+        self.v_proj = nn.Linear(
+            config.hidden_size, self.num_heads * self.head_dim, bias=True
+        )
+        self.o_proj = nn.Linear(
+            self.num_heads * self.head_dim, config.hidden_size, bias=True
+        )
+
+    def __call__(self, x: mx.array, chunk_size: int = 1024) -> mx.array:
+        B, L, _ = x.shape
+
+        queries = self.q_proj(x)
+        keys = self.k_proj(x)
+        values = self.v_proj(x)
+
+        # Reshape to (B, n_heads, L, head_dim)
+        queries = queries.reshape(B, L, self.num_heads, self.head_dim).transpose(
+            0, 2, 1, 3
+        )
+        keys = keys.reshape(B, L, self.num_heads, self.head_dim).transpose(0, 2, 1, 3)
+        values = values.reshape(B, L, self.num_heads, self.head_dim).transpose(
+            0, 2, 1, 3
+        )
+
+        output = chunked_attention(
+            queries,
+            keys,
+            values,
+            scale=self.scale,
+            chunk_size=chunk_size,
+        )
+
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.o_proj(output)
+
+
+class VisionBlock(nn.Module):
+
+    def __init__(self, config: VisionConfig):
+        super().__init__()
+        self.input_layernorm = nn.LayerNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = nn.LayerNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+        self.self_attn = VisionAttention(config)
+        self.mlp = VisionMLP(config)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        # Self-attention with residual
+        h = x + self.self_attn(self.input_layernorm(x))
+        # MLP with residual
+        out = h + self.mlp(self.post_attention_layernorm(h))
+        return out
+
+
+class PatchEmbed(nn.Module):
+
+    def __init__(self, config: VisionConfig):
+        super().__init__()
+        self.config = config
+        self.embed_dim = config.hidden_size
+        self.patch_size = config.patch_size
+        self.num_channels = config.num_channels
+        self.spatial_merge_size = config.spatial_merge_size
+        self.interpolate_mode = config.interpolate_mode
+
+        self.patch_embedding = nn.Conv2d(
+            in_channels=config.num_channels,
+            out_channels=self.embed_dim,
+            kernel_size=self.patch_size,
+            stride=self.patch_size,
+            bias=True,
+        )
+
+        self.max_num_patches = (config.max_image_size // self.patch_size) ** 2
+        self.num_positions = self.max_num_patches + 1
+        self.position_edge = int(self.num_positions**0.5)
+        self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim)
+
+    def __call__(self, pixel_values: mx.array, grid_thw: list) -> mx.array:
+        num_patches = pixel_values.shape[0]
+        # Reshape: (num_patches, C*P*P) -> (num_patches, C, P, P) -> (num_patches, P, P, C) for MLX conv
+        pixel_values = pixel_values.reshape(
+            num_patches, self.num_channels, self.patch_size, self.patch_size
+        )
+        pixel_values = pixel_values.transpose(0, 2, 3, 1)  # NCHW -> NHWC for MLX
+
+        # Apply patch embedding
+        patch_embeds = self.patch_embedding(pixel_values)  # (N, 1, 1, embed_dim)
+        patch_embeds = patch_embeds.reshape(1, num_patches, self.embed_dim)
+
+        # Get position embeddings and interpolate for each grid
+        pos_embed_weights = self.position_embedding.weight[1:, :]  # Skip cls token
+        base_pos_embed = pos_embed_weights.reshape(
+            1, self.position_edge, self.position_edge, self.embed_dim
+        )
+
+        patch_pos_embed_list = []
+        for grid in grid_thw:
+            t, h, w = grid
+            h_float = float(h) + 0.1
+            w_float = float(w) + 0.1
+
+            target_h = int(h)
+            target_w = int(w)
+
+            # Simple bilinear interpolation
+            pos_embed = self._interpolate_pos_embed(base_pos_embed, target_h, target_w)
+            pos_embed = pos_embed.reshape(1, -1, self.embed_dim)
+            patch_pos_embed_list.append(pos_embed)
+
+        patch_pos_embed = mx.concatenate(patch_pos_embed_list, axis=1)
+        embeddings = patch_embeds + patch_pos_embed
+
+        return embeddings
+
+    def _interpolate_pos_embed(
+        self, pos_embed: mx.array, target_h: int, target_w: int
+    ) -> mx.array:
+        dtype = pos_embed.dtype
+        src_h, src_w = pos_embed.shape[1], pos_embed.shape[2]
+
+        if src_h == target_h and src_w == target_w:
+            return pos_embed
+
+        # Create coordinate grids
+        h_scale = src_h / (target_h + 0.1)
+        w_scale = src_w / (target_w + 0.1)
+        h_coords = (mx.arange(target_h) + 0.5) * h_scale - 0.5
+        w_coords = (mx.arange(target_w) + 0.5) * w_scale - 0.5
+
+        i0 = h_coords.astype(mx.int32)
+        j0 = w_coords.astype(mx.int32)
+        i1 = mx.minimum(i0 + 1, src_h - 1)
+        j1 = mx.minimum(j0 + 1, src_w - 1)
+
+        di = (h_coords - i0.astype(mx.float32))[:, None, None]
+        dj = (w_coords - j0.astype(mx.float32))[None, :, None]
+
+        # Gather corners and interpolate
+        p00 = pos_embed[0, i0][:, j0]
+        p01 = pos_embed[0, i0][:, j1]
+        p10 = pos_embed[0, i1][:, j0]
+        p11 = pos_embed[0, i1][:, j1]
+
+        result = (
+            (1 - di) * (1 - dj) * p00
+            + (1 - di) * dj * p01
+            + di * (1 - dj) * p10
+            + di * dj * p11
+        )
+
+        return result[None].astype(dtype)
+
+
+class PatchMerger(nn.Module):
+    def __init__(
+        self,
+        config: VisionConfig,
+    ):
+        super().__init__()
+        self.config = config
+        self.spatial_merge_size = config.spatial_merge_size
+        self.hidden_size = config.hidden_size
+        self.out_hidden_size = config.out_hidden_size
+
+        merge_hidden = config.hidden_size * 2  # 2304
+        final_hidden = config.hidden_size * 4  # 4608
+
+        self.before_rms = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.after_rms = nn.RMSNorm(config.out_hidden_size, eps=config.rms_norm_eps)
+
+        self.proj = [
+            nn.Conv2d(
+                config.hidden_size,
+                merge_hidden,
+                kernel_size=config.spatial_merge_size,
+                stride=config.spatial_merge_size,
+                bias=True,
+            ),
+            nn.GELU(),
+            nn.Conv2d(merge_hidden, final_hidden, kernel_size=1, bias=True),
+        ]
+
+        self.mlp = nn.Linear(final_hidden, config.out_hidden_size, bias=True)
+
+        self.image_newline = mx.zeros((final_hidden,))
+        self.image_begin = mx.zeros((config.out_hidden_size,))
+        self.image_end = mx.zeros((config.out_hidden_size,))
+        self.image_sep = mx.zeros((config.out_hidden_size,))
+
+    def __call__(self, hidden_states: mx.array, grid_h: int, grid_w: int) -> mx.array:
+
+        B = hidden_states.shape[0]
+        final_hidden = self.config.hidden_size * 4  # 4608
+
+        x = self.before_rms(hidden_states)
+
+        x = x.reshape(B, grid_h, grid_w, self.hidden_size)
+
+        for layer in self.proj:
+            x = layer(x)
+
+        merged_h = grid_h // self.spatial_merge_size
+        merged_w = grid_w // self.spatial_merge_size
+
+        x = x.reshape(B, merged_h, merged_w, final_hidden)
+
+        newlines = mx.broadcast_to(
+            self.image_newline[None, None, None, :], (B, merged_h, 1, final_hidden)
+        )
+
+        x = mx.concatenate(
+            [x, newlines], axis=2
+        )  # (B, merged_h, merged_w+1, final_hidden)
+        x = x.reshape(B, merged_h * (merged_w + 1), final_hidden)
+
+        x = self.mlp(x)
+
+        begin = mx.broadcast_to(
+            self.image_begin[None, None, :], (B, 1, self.out_hidden_size)
+        )
+        end = mx.broadcast_to(
+            self.image_end[None, None, :], (B, 1, self.out_hidden_size)
+        )
+
+        x = mx.concatenate([begin, x, end], axis=1)
+
+        x = self.after_rms(x)
+
+        return x
+
+
+class VisionModel(nn.Module):
+    def __init__(self, config: VisionConfig):
+        super().__init__()
+        self.config = config
+        self.model_type = config.model_type
+        if self.model_type != "hunyuan_vl":
+            raise ValueError(f"Unsupported model type: {self.model_type}")
+        self.embeddings = PatchEmbed(config)
+        self.layers = [VisionBlock(config) for _ in range(config.num_hidden_layers)]
+        self.perceive = PatchMerger(
+            config=config,
+        )
+
+    def __call__(
+        self,
+        pixel_values: mx.array,
+        grid_thw: list,
+    ) -> mx.array:
+        """
+        Args:
+            pixel_values: Flattened pixel values of shape (total_patches, C*P*P)
+            grid_thw: List of [t, h, w] for each image
+
+        Returns:
+            Image features of shape (1, total_tokens, text_hidden_size)
+        """
+        hidden_states = self.embeddings(pixel_values, grid_thw)
+
+        for layer in self.layers:
+            hidden_states = layer(hidden_states)
+
+        # Calculate cumulative sequence lengths
+        cu_seqlens = [0]
+        for t, h, w in grid_thw:
+            cu_seqlens.append(int(h * w))
+        cu_seqlens = mx.cumsum(mx.array(cu_seqlens, dtype=mx.int32))
+
+        # Split and process each image
+        processed_items = []
+        for i, grid in enumerate(grid_thw):
+            t, h, w = grid
+            start_idx = int(cu_seqlens[i])
+            end_idx = int(cu_seqlens[i + 1])
+            item = hidden_states[:, start_idx:end_idx, :]
+            processed = self.perceive(item, int(h), int(w))
+            processed_items.append(processed)
+
+        hidden_states = mx.concatenate(processed_items, axis=1)
+        return hidden_states

mlx_vlm/models/idefics2/__init__.py

@@ -0,0 +1,2 @@
+from .config import ModelConfig, PerceiverConfig, TextConfig, VisionConfig
+from .idefics2 import LanguageModel, Model, VisionModel

mlx_vlm/models/idefics2/config.py

@@ -0,0 +1,65 @@
+from dataclasses import dataclass
+from typing import List, Optional
+
+from ..base import BaseModelConfig
+
+
+@dataclass
+class VisionConfig(BaseModelConfig):
+    model_type: str
+    hidden_size: int = 4096
+    intermediate_size: int = 14336
+    num_hidden_layers: int = 32
+    num_attention_heads: int = 32
+    num_key_value_heads: int = 8
+    num_channels: int = 3
+    image_size: int = 224
+    patch_size: int = 32
+    layer_norm_eps: float = 1e-6
+
+
+@dataclass
+class TextConfig(BaseModelConfig):
+    model_type: str
+    hidden_size: int = 4096
+    intermediate_size: int = 14336
+    num_hidden_layers: int = 32
+    num_attention_heads: int = 32
+    num_key_value_heads: int = 8
+    rms_norm_eps: float = 1e-5
+    vocab_size: int = 32003
+    rope_theta: float = 1000000.0
+    rope_traditional: bool = False
+    max_position_embeddings: int = 32768
+    tie_word_embeddings: bool = False
+
+    def __post_init__(self):
+        if self.num_key_value_heads is None:
+            self.num_key_value_heads = self.num_attention_heads
+
+
+@dataclass
+class PerceiverConfig(BaseModelConfig):
+    model_type: str
+    num_key_value_heads: int = 4
+    resampler_depth: int = 3
+    resampler_head_dim: int = 96
+    resampler_n_heads: int = 16
+    resampler_n_latents: int = 64
+
+
+@dataclass
+class ModelConfig(BaseModelConfig):
+    text_config: TextConfig
+    vision_config: VisionConfig
+    perceiver_config: PerceiverConfig
+    model_type: str
+    ignore_index: int = -100
+    image_token_id: int = 32001
+    vocab_size: int = 151936
+    image_token_index: Optional[int] = None
+    eos_token_id: Optional[List[int]] = None
+
+    def __post_init__(self):
+        if self.image_token_index is None:
+            self.image_token_index = self.image_token_id

mlx_vlm/models/idefics2/idefics2.py

@@ -0,0 +1,321 @@
+import re
+from typing import Optional, Tuple
+
+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
+
+
+def masked_scatter(
+    final_embedding: mx.array,
+    image_mask_expanded: mx.array,
+    scaled_image_features: mx.array,
+):
+    # Reshape the tensors to 1D
+    final_embedding_shape = final_embedding.shape
+    scaled_image_features_flattened = mx.flatten(scaled_image_features)
+    final_embedding_flattened = mx.flatten(final_embedding)
+    image_mask_expanded_flattened = mx.flatten(image_mask_expanded)
+
+    # Scatter the scaled image features into the special image token positions
+    image_positions = mx.array(np.where(image_mask_expanded_flattened)[0], mx.uint32)
+    final_embedding_flattened[image_positions] = scaled_image_features_flattened
+
+    # Reshape back to the original shape
+    final_embedding = mx.reshape(final_embedding_flattened, final_embedding_shape)
+
+    return final_embedding
+
+
+class Idefics2PerceiverAttention(nn.Module):
+    def __init__(self, config: ModelConfig):
+        super().__init__()
+
+        dim = config.text_config.hidden_size
+        self.n_heads = n_heads = config.perceiver_config.resampler_n_heads
+        self.n_kv_heads = n_kv_heads = config.perceiver_config.num_key_value_heads
+
+        head_dim = config.perceiver_config.resampler_head_dim
+        self.scale = head_dim**-0.5
+
+        self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=False)
+        self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False)
+        self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False)
+        self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=False)
+
+    def __call__(
+        self,
+        x: mx.array,
+        kv: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Tuple[mx.array, mx.array]] = None,
+    ) -> mx.array:
+        B, L, D = x.shape
+        kv_seq_len = L + kv.shape[1]
+        hidden_states = mx.concatenate([kv, x], axis=-2)
+
+        queries = self.q_proj(x)
+        keys = self.k_proj(hidden_states)
+        values = self.v_proj(hidden_states)
+
+        # Prepare the queries, keys and values for the attention computation
+        queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
+        keys = keys.reshape(B, kv_seq_len, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+        values = values.reshape(B, kv_seq_len, self.n_kv_heads, -1).transpose(
+            0, 2, 1, 3
+        )
+
+        if cache is not None:
+            key_cache, value_cache = cache
+            keys = mx.concatenate([key_cache, keys], axis=2)
+            values = mx.concatenate([value_cache, values], axis=2)
+
+        output = mx.fast.scaled_dot_product_attention(
+            queries, keys, values, scale=self.scale
+        )
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.o_proj(output)
+
+
+class Idefics2PerceiverLayer(nn.Module):
+    def __init__(self, config: ModelConfig):
+        super().__init__()
+        self.hidden_size = config.text_config.hidden_size
+        self.n_latents = config.perceiver_config.resampler_n_latents
+        self.depth = config.perceiver_config.resampler_depth
+        self.rms_norm_eps = config.text_config.rms_norm_eps
+
+        self.input_latents_norm = nn.RMSNorm(self.hidden_size, eps=self.rms_norm_eps)
+        self.input_context_norm = nn.RMSNorm(self.hidden_size, eps=self.rms_norm_eps)
+        self.self_attn = Idefics2PerceiverAttention(config)
+        self.post_attention_layernorm = nn.RMSNorm(
+            self.hidden_size, eps=self.rms_norm_eps
+        )
+        self.mlp = MLP(self.hidden_size, self.hidden_size * 4, self.hidden_size)
+
+    def __call__(
+        self,
+        x: mx.array,
+        hidden_states: mx.array,
+        mask: Optional[mx.array] = None,
+    ) -> mx.array:
+        latents = self.input_latents_norm(x)
+        context = self.input_context_norm(hidden_states)
+
+        latents = self.self_attn(latents, context, mask=mask)
+
+        latents = x + latents
+        r = latents
+
+        latents = self.post_attention_layernorm(latents)
+        latents = self.mlp(latents)
+        latents = r + latents
+        return latents
+
+
+class Idefics2PerceiverResampler(nn.Module):
+    def __init__(self, config: ModelConfig):
+        super().__init__()
+        self.hidden_size = config.text_config.hidden_size
+        self.n_latents = config.perceiver_config.resampler_n_latents
+
+        self.latents = mx.ones((self.n_latents, self.hidden_size))
+        self.layers = [
+            Idefics2PerceiverLayer(config)
+            for _ in range(config.perceiver_config.resampler_depth)
+        ]
+        self.norm = nn.RMSNorm(self.hidden_size, eps=config.text_config.rms_norm_eps)
+
+    def __call__(self, x: mx.array, mask: Optional[mx.array] = None):
+
+        h = mx.expand_dims(self.latents, axis=0)
+        h = mx.repeat(h, x.shape[0], axis=0)
+
+        for layer in self.layers:
+            h = layer(h, x, mask=mask)
+
+        return self.norm(h)
+
+
+class MLP(nn.Module):
+    def __init__(self, dim, hidden_dim, output_size):
+        super().__init__()
+        self.gate_proj = nn.Linear(dim, hidden_dim, bias=False)
+        self.down_proj = nn.Linear(hidden_dim, output_size, bias=False)
+        self.up_proj = nn.Linear(dim, hidden_dim, bias=False)
+
+    def __call__(self, x) -> mx.array:
+        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class Idefics2Connector(nn.Module):
+    def __init__(self, config: ModelConfig):
+        super().__init__()
+        self.modality_projection = MLP(
+            config.vision_config.hidden_size,
+            config.text_config.intermediate_size,
+            config.text_config.hidden_size,
+        )
+
+        self.perceiver_resampler = Idefics2PerceiverResampler(config)
+
+    def __call__(self, x: mx.array, mask=None) -> mx.array:
+        x = self.modality_projection(x)
+        x = self.perceiver_resampler(x, mask=mask)
+        return x
+
+
+class Model(nn.Module):
+    def __init__(self, config: ModelConfig):
+        super().__init__()
+        self.model_type = config.model_type
+        self.config = config
+
+        self.vision_model = VisionModel(config.vision_config)
+        self.language_model = LanguageModel(config.text_config)
+        self.connector = Idefics2Connector(config)
+
+    def get_input_embeddings(
+        self,
+        input_ids: Optional[mx.array] = None,
+        pixel_values: Optional[mx.array] = None,
+        **kwargs,
+    ):
+        pixel_attention_mask = kwargs.get("pixel_attention_mask", None)
+
+        if pixel_values is None:
+            return InputEmbeddingsFeatures(
+                inputs_embeds=self.language_model.embed_tokens(input_ids)
+            )
+
+        inputs_embeds = self.language_model.embed_tokens(input_ids)
+
+        batch_size, num_images, num_channels, height, width = pixel_values.shape
+        pixel_values = pixel_values.reshape(
+            batch_size * num_images, num_channels, height, width
+        )
+
+        # Remove padding images - padding image are full 0.
+        nb_values_per_image = np.prod(pixel_values.shape[1:])
+        real_images_mask = (pixel_values == 0.0).sum(
+            axis=(-1, -2, -3)
+        ) != nb_values_per_image
+        real_images_inds = np.where(real_images_mask)[0].tolist()
+        pixel_values = pixel_values[real_images_inds, ...]
+
+        if pixel_attention_mask is None:
+            pixel_attention_mask = mx.ones(
+                (pixel_values.size(0), pixel_values.size(2), pixel_values.size(3)),
+                dtype=mx.bool,
+            )
+        else:
+            # Remove padding images from the mask
+            pixel_attention_mask = pixel_attention_mask.reshape(
+                batch_size * num_images, height, width
+            )
+            pixel_attention_mask = pixel_attention_mask[real_images_inds]
+
+        patch_size = self.config.vision_config.patch_size
+        batch_size, height, width = pixel_attention_mask.shape
+
+        # Calculate number of patches
+        patches_h = height // patch_size
+        patches_w = width // patch_size
+
+        # Reshape to extract patches
+        reshaped = pixel_attention_mask[
+            :, : patches_h * patch_size, : patches_w * patch_size
+        ]
+        reshaped = reshaped.reshape(
+            batch_size, patches_h, patch_size, patches_w, patch_size
+        )
+        reshaped = reshaped.transpose(
+            0, 1, 3, 2, 4
+        )  # (batch, patches_h, patches_w, patch_size, patch_size)
+
+        # Sum over patch dimensions and check if any pixels are active
+        patch_attention_mask = reshaped.sum(axis=(-1, -2)) > 0
+
+        pooler_output, *_ = self.vision_model(
+            pixel_values.transpose(0, 2, 3, 1),
+            patch_attention_mask=patch_attention_mask,
+            output_hidden_states=True,
+        )
+
+        image_features = pooler_output.astype(pixel_values.dtype)
+        image_features = self.connector(image_features)
+
+        final_inputs_embeds = self._prepare_inputs_for_multimodal(
+            image_features, inputs_embeds, input_ids
+        )
+        return InputEmbeddingsFeatures(inputs_embeds=final_inputs_embeds)
+
+    def _prepare_inputs_for_multimodal(self, image_features, inputs_embeds, input_ids):
+        special_image_mask = input_ids == self.config.image_token_index
+        n_image_tokens = special_image_mask.sum()
+        special_image_mask = special_image_mask[..., None]
+        special_image_mask = mx.broadcast_to(special_image_mask, inputs_embeds.shape)
+
+        n_image_features = image_features.shape[0]
+        n_image_mask_elements = special_image_mask.sum()
+        if n_image_mask_elements != image_features.size:
+            raise ValueError(
+                f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}"
+            )
+
+        inputs_embeds = masked_scatter(
+            inputs_embeds, special_image_mask, image_features
+        )
+
+        return inputs_embeds
+
+    @property
+    def layers(self):
+        return self.language_model.model.layers
+
+    def __call__(
+        self,
+        input_ids: mx.array,
+        pixel_values: mx.array,
+        mask: mx.array,
+        cache=None,
+        **kwargs,
+    ):
+        input_embeddings_features = self.get_input_embeddings(
+            input_ids, pixel_values, **kwargs
+        )
+        logits = self.language_model(
+            inputs=input_ids,
+            cache=cache,
+            inputs_embeds=input_embeddings_features.inputs_embeds,
+        )
+        return logits
+
+    def sanitize(self, weights):
+        weights = {
+            (
+                f"{k.split('.', 1)[1]}"
+                if re.match(r"^model\.", k)
+                else (f"language_model.{k}" if re.match(r"^lm_head\.", k) else k)
+            ): v
+            for k, v in weights.items()
+        }
+
+        weights = {
+            (
+                f"language_model.{k.split('.', 1)[1]}"
+                if re.match(
+                    r"^text_model\.",
+                    k,
+                )
+                else k
+            ): v
+            for k, v in weights.items()
+        }
+
+        return weights