# Copyright © 2024 Apple Inc. import inspect import math from dataclasses import dataclass from typing import Optional import mlx.core as mx import mlx.nn as nn @dataclass class VisionConfig: model_type: str num_hidden_layers: int = 24 hidden_size: int = 1024 intermediate_size: int = 4096 num_attention_heads: int = 16 image_size: int = 336 patch_size: int = 14 projection_dim: int = 768 vocab_size: int = 32000 num_channels: int = 3 layer_norm_eps: float = 1e-5 @classmethod def from_dict(cls, params): return cls( **{ k: v for k, v in params.items() if k in inspect.signature(cls).parameters } ) class Attention(nn.Module): def __init__( self, dims: int, num_heads: int, query_input_dims: Optional[int] = None, key_input_dims: Optional[int] = None, value_input_dims: Optional[int] = None, value_dims: Optional[int] = None, value_output_dims: Optional[int] = None, bias: bool = False, ): super().__init__() if (dims % num_heads) != 0: raise ValueError( "The input feature dimensions should be divisible by the " f"number of heads ({dims} % {num_heads}) != 0" ) query_input_dims = query_input_dims or dims key_input_dims = key_input_dims or dims value_input_dims = value_input_dims or key_input_dims value_dims = value_dims or dims value_output_dims = value_output_dims or dims self.num_heads = num_heads self.q_proj = nn.Linear(query_input_dims, dims, bias=bias) self.k_proj = nn.Linear(key_input_dims, dims, bias=bias) self.v_proj = nn.Linear(value_input_dims, value_dims, bias=bias) self.out_proj = nn.Linear(value_dims, value_output_dims, bias=bias) def __call__(self, queries, keys, values, mask=None): queries = self.q_proj(queries) keys = self.k_proj(keys) values = self.v_proj(values) num_heads = self.num_heads B, L, D = queries.shape _, S, _ = keys.shape queries = queries.reshape(B, L, num_heads, -1).transpose(0, 2, 1, 3) keys = keys.reshape(B, S, num_heads, -1).transpose(0, 2, 3, 1) values = values.reshape(B, S, num_heads, -1).transpose(0, 2, 1, 3) scale = math.sqrt(1 / queries.shape[-1]) scores = (queries * scale) @ keys if mask is not None: scores = scores + mask.astype(scores.dtype) scores = mx.softmax(scores, axis=-1) values_hat = (scores @ values).transpose(0, 2, 1, 3).reshape(B, L, -1) return self.out_proj(values_hat) class MLP(nn.Module): def __init__(self, config: VisionConfig): super().__init__() self.activation_fn = nn.GELU(approx="fast") self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size) self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size) def __call__(self, x: mx.array) -> mx.array: x = self.activation_fn(self.fc1(x)) x = self.fc2(x) return x class EncoderLayer(nn.Module): def __init__(self, config: VisionConfig): super().__init__() self.embed_dim = config.hidden_size self.self_attn = Attention( config.hidden_size, config.num_attention_heads, bias=True ) self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) self.mlp = MLP(config) self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) def __call__(self, x: mx.array, mask: Optional[mx.array] = None) -> mx.array: y = self.layer_norm1(x) y = self.self_attn(y, y, y, mask) x = x + y y = self.layer_norm2(x) y = self.mlp(y) return x + y class Encoder(nn.Module): def __init__(self, config: VisionConfig): super().__init__() self.layers = [EncoderLayer(config) for _ in range(config.num_hidden_layers)] class VisionEmbeddings(nn.Module): def __init__(self, config: VisionConfig): super().__init__() self.config = config self.embed_dim = config.hidden_size self.image_size = config.image_size self.patch_size = config.patch_size self.class_embedding = mx.zeros((config.hidden_size,)) 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=False, ) self.num_patches = (self.image_size // self.patch_size) ** 2 self.num_positions = self.num_patches + 1 self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim) def __call__(self, x: mx.array) -> mx.array: batch_size = x.shape[0] patch_embeddings = self.patch_embedding(x) patch_embeddings = mx.flatten(patch_embeddings, start_axis=1, end_axis=2) embed_dim = patch_embeddings.shape[-1] cls_embeddings = mx.broadcast_to( self.class_embedding, (batch_size, 1, embed_dim) ) embeddings = mx.concatenate((cls_embeddings, patch_embeddings), axis=1) embeddings += self.position_embedding.weight return embeddings class ClipVisionModel(nn.Module): def __init__(self, config: VisionConfig): super().__init__() self.embeddings = VisionEmbeddings(config) self.pre_layrnorm = nn.LayerNorm(config.hidden_size) self.encoder = Encoder(config) self.post_layernorm = nn.LayerNorm(config.hidden_size) def __call__( self, x: mx.array, output_hidden_states: Optional[bool] = None, ) -> mx.array: x = self.embeddings(x) x = self.pre_layrnorm(x) encoder_states = (x,) if output_hidden_states else None for l in self.encoder.layers: x = l(x, mask=None) if output_hidden_states: encoder_states = encoder_states + (x,) pooler_output = self.post_layernorm(x[:, 0, :]) return pooler_output, x, encoder_states class VisionModel(nn.Module): def __init__(self, config: VisionConfig): super().__init__() self.model_type = config.model_type if self.model_type != "clip_vision_model": raise ValueError(f"Unsupported model type: {self.model_type}") self.vision_model = ClipVisionModel(config) def __call__( self, x: mx.array, output_hidden_states: Optional[bool] = None ) -> mx.array: return self.vision_model(x, output_hidden_states) @staticmethod def sanitize(weights): sanitized_weights = {} for k, v in weights.items(): if "position_ids" in k: # Remove unused position_ids continue elif "patch_embedding.weight" in k: # PyTorch conv2d weight tensors have shape: # [out_channels, in_channels, kH, KW] # MLX conv2d expects the weight be of shape: # [out_channels, kH, KW, in_channels] sanitized_weights[k] = v.transpose(0, 2, 3, 1) else: sanitized_weights[k] = v return sanitized_weights