nexaai 1.0.29__cp310-cp310-macosx_14_0_universal2.whl

nexaai/mlx_backend/vlm/modeling/models/gemma3n/vision.py

@@ -0,0 +1,1022 @@
+import inspect
+from collections.abc import Sequence
+from dataclasses import dataclass
+from math import sqrt
+from typing import Dict, List, Optional, Tuple, Type
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .config import VisionConfig
+
+from ..base import check_array_shape
+from ..kernels import bicubic_interpolate, nearest_interpolate
+
+
+# https://github.com/huggingface/new-model-addition-timm-gemma3p5-non-fork/blob/mobilenet-gemma3n-rw/timm/models/mobilenetv5.py#L24
+class MobileNetV5MultiScaleFusionAdapter(nn.Module):
+    """Multi-layer fusion token adapter.
+    Attributes:
+      out_filters: The number of output filters.
+      output_resolution: The output resolution.
+      activation: The activation function.
+      expansion_ratio: The expansion ratio.
+      upsampling_interpolation: The upsampling interpolation.
+      use_layer_scale: Whether to use layer scale.
+      layer_scale_init_value: The initial value of the layer scale.
+      skip_projection: Whether to skip the projection.
+      name: The name of the module.
+      upsize: The upsampling fn.
+      downsize: The downsampling fn.
+    """
+
+    def __init__(
+        self,
+        in_chs: List[int],
+        out_chs: int,
+        output_resolution: int,
+        expansion_ratio: float = 2.0,
+        interpolation_mode: str = "nearest",
+        use_layer_scale: bool = False,
+        layer_scale_init_value: float = 1e-5,
+        noskip: bool = True,
+    ):
+        super().__init__()
+        self.in_channels = sum(in_chs) if isinstance(in_chs, Sequence) else in_chs
+        self.out_channels = out_chs
+        self.output_resolution = to_2tuple(output_resolution)
+        self.expansion_ratio = expansion_ratio
+        self.interpolation_mode = interpolation_mode
+        self.use_layer_scale = use_layer_scale
+        self.layer_scale_init_value = layer_scale_init_value
+        self.noskip = noskip
+
+        norm_layer = RMSNormAct2d
+        self.ffn = UniversalInvertedResidual(
+            in_chs=self.in_channels,
+            out_chs=self.out_channels,
+            dw_kernel_size_mid=0,
+            exp_ratio=self.expansion_ratio,
+            norm_layer=norm_layer,
+            noskip=self.noskip,
+            layer_scale_init_value=(
+                self.layer_scale_init_value if self.use_layer_scale else None
+            ),
+        )
+
+        self.norm = norm_layer(self.out_channels, eps=1e-6, apply_act=False)
+
+    def __call__(self, inputs: list[mx.array]) -> mx.array:
+        inputs = [i.transpose(0, 3, 1, 2) for i in inputs]
+        high_resolution = inputs[0].shape[
+            -2:
+        ]  # Assuming the first input is the highest resolution.
+        resized_inputs = []
+
+        for _, img in enumerate(inputs):
+            if any([r < hr for r, hr in zip(img.shape[-2:], high_resolution)]):
+                img = nearest_interpolate(img, size=high_resolution)
+
+            resized_inputs.append(img)
+
+        channel_cat_imgs = mx.concatenate(
+            resized_inputs, axis=1
+        )  # Cat on channel dim, must equal self.in_channels
+        img = self.ffn(channel_cat_imgs.swapaxes(1, 3)).swapaxes(1, 3)
+
+        if any([ro != rh for ro, rh in zip(high_resolution, self.output_resolution)]):
+            if (
+                high_resolution[0] % self.output_resolution[0] != 0
+                or high_resolution[1] % self.output_resolution[1] != 0
+            ):
+                img = bicubic_interpolate(img, self.output_resolution)
+            else:
+                h_strides = high_resolution[0] // self.output_resolution[0]
+                w_strides = high_resolution[1] // self.output_resolution[1]
+
+                img = nn.AvgPool2d(
+                    kernel_size=(h_strides, w_strides),
+                    stride=(h_strides, w_strides),
+                )(img.swapaxes(1, 3))
+
+            img = self.norm(img) if self.noskip else img
+
+        return img
+
+
+# https://github.com/huggingface/new-model-addition-timm-gemma3p5-non-fork/blob/mobilenet-gemma3n-rw/timm/layers/layer_scale.py#L22
+class LayerScale2d(nn.Module):
+    def __init__(self, dim: int, init_values: float = 1e-5, inplace: bool = False):
+        super().__init__()
+        self.inplace = inplace
+        self.gamma = init_values * mx.ones((dim,))
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return x.mul_(self.gamma) if self.inplace else x * self.gamma
+
+
+def rms_norm2d(
+    x: mx.array,
+    normalized_shape: List[int],
+    weight: Optional[mx.array] = None,
+    eps: float = 1e-5,
+):
+    assert len(normalized_shape) == 1
+    dtype = x.dtype
+    v = mx.power(x, 2)
+    v = mx.mean(v, axis=1, keepdims=True)
+    x = x * mx.rsqrt(v + eps)
+    if weight is not None:
+        x = x.astype(dtype) * weight.reshape(1, -1, 1, 1)
+    return x
+
+
+# https://github.com/huggingface/new-model-addition-timm-gemma3p5-non-fork/blob/mobilenet-gemma3n-rw/timm/layers/norm_act.py#L504
+class RMSNormAct2d(nn.RMSNorm):
+    def __init__(
+        self,
+        num_channels,
+        eps=1e-6,
+        apply_act: bool = True,
+    ):
+        super().__init__(dims=num_channels, eps=eps)
+        self.normalized_shape = [num_channels]
+        self.drop = nn.Identity()
+        self.act = nn.GELU() if apply_act else nn.Identity()
+
+    def __call__(self, x: mx.array) -> mx.array:
+
+        x = x.transpose(0, 3, 1, 2)  # Convert from NHWC to NCHW
+        x = rms_norm2d(x, self.normalized_shape, self.weight, self.eps)
+        x = self.drop(x)
+        x = self.act(x)
+        x = x.transpose(0, 2, 3, 1)  # Convert back to NHWC
+        return x
+
+
+# https://github.com/huggingface/new-model-addition-timm-gemma3p5-non-fork/blob/mobilenet-gemma3n-rw/timm/models/_efficientnet_blocks.py#L310
+class UniversalInvertedResidual(nn.Module):
+    def __init__(
+        self,
+        in_chs: int,
+        out_chs: int,
+        dw_kernel_size_start: int = 0,
+        dw_kernel_size_mid: int = 3,
+        dw_kernel_size_end: int = 0,
+        stride: int = 1,
+        dilation: int = 1,
+        group_size: int = 1,
+        pad_type: str = "",
+        noskip: bool = False,
+        exp_ratio: float = 1.0,
+        norm_layer=RMSNormAct2d,
+        conv_kwargs: Optional[Dict] = None,
+        drop_path_rate: float = 0.0,
+        layer_scale_init_value: Optional[float] = 1e-5,
+    ):
+        super().__init__()
+        self.has_skip = (in_chs == out_chs and stride == 1) and not noskip
+        if stride > 1:
+            assert dw_kernel_size_start or dw_kernel_size_mid or dw_kernel_size_end
+
+        if dw_kernel_size_start:
+            dw_start_stride = stride if not dw_kernel_size_mid else 1
+            dw_start_groups = num_groups(group_size, in_chs)
+            self.dw_start = ConvNormAct(
+                nn.Conv2d,
+                in_chs,
+                in_chs,
+                kernel_size=dw_kernel_size_start,
+                stride=dw_start_stride,
+                padding=(dw_kernel_size_start - 1) // 2,
+                dilation=dilation,
+                groups=dw_start_groups,
+                bias=False,
+                apply_act=False,
+                eps=1e-05,
+            )
+        else:
+            self.dw_start = nn.Identity()
+
+        mid_chs = make_divisible(in_chs * exp_ratio)
+        self.pw_exp = ConvNormAct(
+            nn.Conv2d,
+            in_chs,
+            mid_chs,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            groups=1,
+            bias=False,
+            eps=1e-05,
+        )
+
+        if dw_kernel_size_mid:
+            dw_mid_groups = num_groups(group_size, mid_chs)
+            self.dw_mid = ConvNormAct(
+                Conv2dSame,
+                mid_chs,
+                mid_chs,
+                kernel_size=dw_kernel_size_mid,
+                stride=stride,
+                padding=0,
+                dilation=dilation,
+                groups=dw_mid_groups,
+                bias=False,
+                eps=1e-05,
+            )
+        else:
+            self.dw_mid = nn.Identity()
+
+        self.pw_proj = ConvNormAct(
+            nn.Conv2d,
+            mid_chs,
+            out_chs,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+            groups=1,
+            bias=False,
+            apply_act=False,
+            eps=1e-05,
+        )
+        if layer_scale_init_value is not None:
+            self.layer_scale = LayerScale2d(out_chs, layer_scale_init_value)
+        else:
+            self.layer_scale = nn.Identity()
+
+    def __call__(self, x: mx.array) -> mx.array:
+        shortcut = x
+        x = self.dw_start(x)
+        x = self.pw_exp(x)
+        x = self.dw_mid(x)
+        x = self.pw_proj(x)
+        x = self.layer_scale(x)
+        if self.has_skip:
+            x = x + shortcut
+        return x
+
+
+# https://github.com/huggingface/new-model-addition-timm-gemma3p5-non-fork/blob/mobilenet-gemma3n-rw/timm/layers/conv_bn_act.py#L15
+class ConvNormAct(nn.Module):
+    def __init__(
+        self,
+        conv_cls,
+        in_chs: int,
+        out_chs: int,
+        kernel_size: int = 3,
+        stride: int = 1,
+        padding: int = 0,
+        dilation: int = 1,
+        groups: int = 1,
+        bias: bool = False,
+        apply_act: bool = True,
+        eps: float = 1e-6,
+    ):
+        super().__init__()
+        self.out_chs = out_chs
+        self.conv = conv_cls(
+            in_chs,
+            out_chs,
+            kernel_size,
+            stride,
+            padding,
+            (dilation, dilation),
+            groups,
+            bias,
+        )
+        self.bn = RMSNormAct2d(out_chs, eps=eps, apply_act=apply_act)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        c = self.conv(x)
+        r = self.bn(c)
+        return r
+
+
+def pad_same(
+    x,
+    kernel_size: List[int],
+    stride: List[int],
+    dilation: List[int] = (1, 1),
+    value: float = 0,
+):
+    """
+    Input should be in MLX format
+    """
+    ih, iw = x.shape[1:3]
+    pad_h = get_same_padding(ih, kernel_size[0], stride[0], dilation[0])
+    pad_w = get_same_padding(iw, kernel_size[1], stride[1], dilation[1])
+
+    # MLX pad format: [(low, high), (low, high), ...] for each axis
+    # Padding order is reversed compared to PyTorch F.pad
+    pad_widths = [
+        (0, 0),  # No padding for batch dimension
+        (pad_h // 2, pad_h - pad_h // 2),  # Height padding
+        (pad_w // 2, pad_w - pad_w // 2),  # Width padding
+        (0, 0),  # No padding for channel dimension
+    ]
+
+    x = mx.pad(x, pad_widths, constant_values=value)
+    return x
+
+
+def get_padding_value(padding, kernel_size, **kwargs) -> Tuple[Tuple, bool]:
+    dynamic = False
+    if isinstance(padding, str):
+        # for any string padding, the padding will be calculated for you, one of three ways
+        padding = padding.lower()
+        if padding == "same":
+            # TF compatible 'SAME' padding, has a performance and GPU memory allocation impact
+            if is_static_pad(kernel_size, **kwargs):
+                # static case, no extra overhead
+                padding = get_padding(kernel_size, **kwargs)
+            else:
+                # dynamic 'SAME' padding, has runtime/GPU memory overhead
+                padding = 0
+                dynamic = True
+        elif padding == "valid":
+            # 'VALID' padding, same as padding=0
+            padding = 0
+        else:
+            # Default to PyTorch style 'same'-ish symmetric padding
+            padding = get_padding(kernel_size, **kwargs)
+    return padding, dynamic
+
+
+def get_same_padding(
+    input_size: int, kernel_size: int, stride: int, dilation: int = 1
+) -> int:
+    """Calculate padding needed for 'same' output size."""
+    effective_kernel_size = dilation * (kernel_size - 1) + 1
+    output_size = (input_size + stride - 1) // stride
+    total_padding = max(
+        0, (output_size - 1) * stride + effective_kernel_size - input_size
+    )
+    return total_padding
+
+
+def get_padding(kernel_size, stride=1, dilation=1, **_):
+    """Get symmetric padding for given kernel size."""
+    if isinstance(kernel_size, int):
+        kernel_size = [kernel_size, kernel_size]
+    if isinstance(stride, int):
+        stride = [stride, stride]
+    if isinstance(dilation, int):
+        dilation = [dilation, dilation]
+
+    padding = []
+    for k, d in zip(kernel_size, dilation):
+        effective_k = d * (k - 1) + 1
+        pad_total = effective_k - 1
+        padding.append(pad_total // 2)
+    return tuple(padding)
+
+
+def is_static_pad(kernel_size, stride=1, dilation=1, **_):
+    """Check if padding can be calculated statically."""
+    if isinstance(kernel_size, int):
+        kernel_size = [kernel_size, kernel_size]
+    if isinstance(stride, int):
+        stride = [stride, stride]
+    if isinstance(dilation, int):
+        dilation = [dilation, dilation]
+
+    # Static padding is possible when stride is 1 for all dimensions
+    return all(s == 1 for s in stride)
+
+
+class Conv2dSame(nn.Conv2d):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.kernel_size = self.weight.shape[1:3]
+
+    def __call__(self, x: mx.array) -> mx.array:
+        x = pad_same(x, self.kernel_size, self.stride, self.dilation)
+        y = mx.conv2d(
+            x, self.weight, self.stride, self.padding, self.dilation, self.groups
+        )
+        if "bias" in self:
+            y = y + self.bias
+        return y
+
+
+# https://github.com/huggingface/new-model-addition-timm-gemma3p5-non-fork/blob/mobilenet-gemma3n-rw/timm/models/_efficientnet_blocks.py#L629
+class EdgeResidual(nn.Module):
+    def __init__(
+        self,
+        in_chs: int,
+        out_chs: int,
+        exp_kernel_size: int = 3,
+        stride: int = 1,
+        dilation: int = 1,
+        group_size: int = 0,
+        pad_type: str = "",
+        force_in_chs: int = 0,
+        noskip: bool = False,
+        expand_ratio: float = 1.0,
+        pw_kernel_size: int = 1,
+        norm_layer=RMSNormAct2d,
+    ):
+        super().__init__()
+
+        if force_in_chs > 0:
+            mid_chs = make_divisible(force_in_chs * expand_ratio)
+        else:
+            mid_chs = make_divisible(in_chs * expand_ratio)
+
+        groups = num_groups(group_size, mid_chs)
+
+        self.has_skip = (in_chs == out_chs and stride == 1) and not noskip
+
+        self.conv_exp = Conv2dSame(
+            in_chs,
+            mid_chs,
+            kernel_size=exp_kernel_size,
+            stride=stride,
+            padding=0,
+            dilation=(dilation, dilation),
+            groups=groups,
+            bias=False,
+        )
+
+        self.bn1 = norm_layer(mid_chs, eps=1e-05) if norm_layer else nn.Identity()
+
+        # Point-wise linear projection
+        padding_pwl = (pw_kernel_size - 1) // 2
+        self.conv_pwl = nn.Conv2d(
+            mid_chs,
+            out_chs,
+            kernel_size=pw_kernel_size,
+            padding=padding_pwl,
+            bias=False,
+        )
+
+        self.bn2 = (
+            norm_layer(out_chs, eps=1e-05, apply_act=False)
+            if norm_layer
+            else nn.Identity()
+        )
+
+    def __call__(self, x: mx.array) -> mx.array:
+        shortcut = x
+        x = self.conv_exp(x)
+        x = self.bn1(x)
+        x = self.conv_pwl(x)
+        x = self.bn2(x)
+        if self.has_skip:
+            x = x + shortcut
+        return x
+
+
+# https://github.com/huggingface/new-model-addition-timm-gemma3p5-non-fork/blob/mobilenet-gemma3n-rw/timm/models/_efficientnet_blocks.py#L449
+class MobileAttention(nn.Module):
+    def __init__(
+        self,
+        in_chs: int,
+        out_chs: int,
+        stride: int = 1,
+        dw_kernel_size: int = 3,
+        dilation: int = 1,
+        group_size: int = 1,
+        pad_type: str = "",
+        num_heads: int = 8,
+        key_dim: int = 64,
+        value_dim: int = 64,
+        use_multi_query: bool = True,
+        query_strides: Tuple[int, int] = (1, 1),
+        kv_stride: int = 1,
+        cpe_dw_kernel_size: int = 3,
+        noskip: bool = False,
+        act_layer=nn.GELU,
+        aa_layer=None,
+        drop_path_rate: float = 0.0,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+        layer_scale_init_value: Optional[float] = 1e-5,
+        use_bias: bool = False,
+    ):
+        super().__init__()
+        self.has_skip = (stride == 1 and in_chs == out_chs) and not noskip
+        self.query_strides = to_2tuple(query_strides)
+        self.kv_stride = kv_stride
+        self.has_query_stride = any([s > 1 for s in self.query_strides])
+
+        # Normalization layer
+        self.norm = RMSNormAct2d(
+            in_chs,
+            eps=1e-05,
+            apply_act=False,
+        )
+        # Determine number of heads if not provided
+        if num_heads is None:
+            assert in_chs % key_dim == 0
+            num_heads = in_chs // key_dim
+
+        # Attention layer
+        if use_multi_query:
+            self.attn = MultiQueryAttention2d(
+                in_chs,
+                dim_out=out_chs,
+                num_heads=num_heads,
+                key_dim=key_dim,
+                value_dim=value_dim,
+                query_strides=query_strides,
+                kv_stride=kv_stride,
+                dilation=dilation,
+                padding=pad_type,
+                dw_kernel_size=dw_kernel_size,
+                attn_drop=attn_drop,
+                proj_drop=proj_drop,
+            )
+        else:
+            raise NotImplementedError("attention not implemented")
+
+        # Layer scaling
+        if layer_scale_init_value is not None:
+            self.layer_scale = LayerScale2d(out_chs, layer_scale_init_value)
+        else:
+            self.layer_scale = nn.Identity()
+
+        # Drop path for residual connection
+        self.drop_path = nn.Identity()
+
+    def __call__(self, x: mx.array) -> mx.array:
+        shortcut = x
+        x = self.norm(x)
+        x = self.attn(x)
+        x = self.layer_scale(x)
+
+        # Apply skip connection if available
+        if self.has_skip:
+            x = self.drop_path(x) + shortcut
+        return x
+
+
+def create_conv2d(
+    in_channels,
+    out_channels,
+    kernel_size,
+    stride=1,
+    dilation=1,
+    depthwise=False,
+    bias=False,
+    **kwargs,
+):
+    """Helper function to create a 2D convolution with common parameters"""
+    if depthwise:
+        # Depthwise convolution
+        return nn.Conv2d(
+            in_channels,
+            out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=(kernel_size - 1) // 2 * dilation,
+            dilation=dilation,
+            groups=in_channels,
+            bias=bias,
+        )
+    else:
+        # Regular convolution
+        return nn.Conv2d(
+            in_channels,
+            out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=(kernel_size - 1) // 2 * dilation,
+            dilation=dilation,
+            bias=bias,
+        )
+
+
+def to_2tuple(x):
+    """Convert input to 2-tuple"""
+    if isinstance(x, tuple):
+        return x
+    return (x, x)
+
+
+class NamedSequential(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self._order = []
+
+    def add_module(self, name, module):
+        setattr(self, name, module)
+        self._order.append(name)
+
+    def __call__(self, x):
+        for name in self._order:
+            x = getattr(self, name)(x)
+        return x
+
+
+# https://github.com/huggingface/new-model-addition-timm-gemma3p5-non-fork/blob/mobilenet-gemma3n-rw/timm/layers/attention2d.py#L82
+class MultiQueryAttention2d(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        dim_out: Optional[int] = None,
+        num_heads: int = 8,
+        key_dim: int = 64,
+        value_dim: int = 64,
+        query_strides: Tuple[int, int] = (1, 1),
+        kv_stride: int = 1,
+        dilation: int = 1,
+        padding: str = "",
+        dw_kernel_size: int = 3,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+    ):
+        super().__init__()
+        dim_out = dim_out or dim
+        self.num_heads = num_heads
+        self.query_strides = to_2tuple(query_strides)
+        self.kv_stride = kv_stride
+        self.fused_attn = True
+        self.key_dim = key_dim
+        self.value_dim = value_dim
+        head_dim = key_dim
+        self.scale = head_dim**-0.5
+
+        self.query = NamedSequential()
+        self.query.add_module(
+            "proj",
+            create_conv2d(
+                dim,
+                self.num_heads * self.key_dim,
+                kernel_size=1,
+            ),
+        )
+        self.key = NamedSequential()
+        if kv_stride > 1:
+            self.key.add_module(
+                "down_conv",
+                create_conv2d(
+                    dim,
+                    dim,
+                    kernel_size=dw_kernel_size,
+                    stride=kv_stride,
+                    dilation=dilation,
+                    padding=padding,
+                    depthwise=True,
+                ),
+            )
+            self.key.add_module("norm", RMSNormAct2d(dim, eps=1e-6, apply_act=False))
+        self.key.add_module(
+            "proj", create_conv2d(dim, key_dim, kernel_size=1, bias=False)
+        )
+
+        self.value = NamedSequential()
+        if kv_stride > 1:
+            self.value.add_module(
+                "down_conv",
+                create_conv2d(
+                    dim,
+                    dim,
+                    kernel_size=dw_kernel_size,
+                    stride=kv_stride,
+                    dilation=dilation,
+                    padding=padding,
+                    depthwise=True,
+                ),
+            )
+            self.value.add_module("norm", RMSNormAct2d(dim, eps=1e-6, apply_act=False))
+        self.value.add_module(
+            "proj", create_conv2d(dim, value_dim, kernel_size=1, bias=False)
+        )
+
+        # Attention dropout
+        self.attn_drop = nn.Dropout(attn_drop) if attn_drop > 0 else nn.Identity()
+
+        # Output projection
+        self.output = NamedSequential()
+        self.output.add_module(
+            "proj",
+            create_conv2d(
+                value_dim * num_heads,
+                dim_out,
+                kernel_size=1,
+                stride=1,
+                bias=False,
+            ),
+        )
+        self.proj_drop = nn.Dropout(proj_drop) if proj_drop > 0 else nn.Identity()
+
+    def _reshape_input(self, t: mx.array):
+        """
+        Input shape MLX: [B, H, W, C]
+        Input shape PyTorch: [B, C, H, W]
+
+        PyTorch Reshape: [B, C, H, W] -> [B, C, -1] -> [B, -1, C] -> [B, 1, -1, C] -> SDPA
+        MLX Reshape: [B, H, W, C] -> [B, -1, C] -> [B, 1, -1, C] -> SDPA
+        """
+        s = t.shape
+        t = t.reshape(s[0], -1, s[3])[:, None, :, :]
+
+        return t
+
+    def _reshape_projected_query(self, t: mx.array, num_heads: int, key_dim: int):
+        """
+        Input shape MLX: [B, H, W, C] where C = num_heads * key_dim
+        """
+        B, H, W, C = t.shape
+        # t = t.reshape(B, H, W, num_heads, key_dim)
+        t = t.reshape(B, H * W, num_heads, key_dim)
+        return t.transpose(0, 2, 1, 3)
+
+    def _reshape_output(self, t: mx.array, num_heads: int, h_px: int, w_px: int):
+        """
+        Input shape: [B, NH, L, D] where L = h_px * w_px
+        Output shape MLX: [B, H, W, C] where C = NH * D
+        """
+        B, NH, L, D = t.shape
+        # First transpose to [B, L, NH, D]
+        t = t.transpose(0, 2, 1, 3)
+        # Then reshape to [B, H, W, NH*D]
+        t = t.reshape(B, h_px, w_px, NH * D)
+        return t
+
+    def __call__(self, x: mx.array, attn_mask: Optional[mx.array] = None) -> mx.array:
+        B, H, W, C = x.shape
+        q = self.query(x)
+        q = self._reshape_projected_query(q, self.num_heads, self.key_dim)
+
+        k = self.key(x)
+        k = self._reshape_input(k)
+
+        v = self.value(x)
+        v = self._reshape_input(v)
+
+        if self.fused_attn:
+            o = mx.fast.scaled_dot_product_attention(
+                q,
+                k,
+                v,
+                scale=1.0 / sqrt(q.shape[-1]),
+            )
+        else:
+            raise NotImplementedError("unfused attention not implemented")
+
+        o = self._reshape_output(
+            o, self.num_heads, H // self.query_strides[0], W // self.query_strides[1]
+        )
+        x = self.output(o)
+        return x
+
+
+def num_groups(group_size: Optional[int], channels: int) -> int:
+    if not group_size:  # 0 or None
+        return 1  # normal conv with 1 group
+    else:
+        # NOTE group_size == 1 -> depthwise conv
+        assert channels % group_size == 0
+        return channels // group_size
+
+
+def make_divisible(v, divisor: int = 8, min_value=None, round_limit: float = 0.9):
+    min_value = min_value or divisor
+    new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
+    # Make sure that round down does not go down by more than 10%.
+    if new_v < round_limit * v:
+        new_v += divisor
+    return new_v
+
+
+@dataclass(frozen=True)
+class EdgeResidualConfig:
+    kernel_size: int = 3
+    filters: int = 32
+    strides: int = 1
+    expand_ratio: float = 4.0
+    is_multiscale: bool = False
+
+
+def _er(kernel_size, filters, strides=1, expand_ratio=4.0, is_multiscale=False):
+    return EdgeResidualConfig(
+        kernel_size=kernel_size,
+        filters=filters,
+        strides=strides,
+        expand_ratio=expand_ratio,
+        is_multiscale=is_multiscale,
+    )
+
+
+@dataclass(frozen=True)
+class UniversalInvertedResidualConfig:
+    start_dw_kernel_size: int = 0  # Zero size means no conv
+    mid_dw_kernel_size: int = 0  # Zero size means no conv
+    filters: int = 32
+    strides: int = 1
+    expand_ratio: float = 4.0
+    is_multiscale: bool = False
+
+
+def _uir(
+    start_dw_kernel_size,
+    mid_dw_kernel_size,
+    filters,
+    strides=1,
+    expand_ratio=4.0,
+    is_multiscale=False,
+):
+    return UniversalInvertedResidualConfig(
+        start_dw_kernel_size=start_dw_kernel_size,
+        mid_dw_kernel_size=mid_dw_kernel_size,
+        filters=filters,
+        strides=strides,
+        expand_ratio=expand_ratio,
+        is_multiscale=is_multiscale,
+    )
+
+
+@dataclass(frozen=True)
+class MultiQueryAttentionBlockConfig:
+    num_heads: int = 8
+    kv_dim: int = 16
+    kv_strides: int = 1
+    mmqa_avg_pool_kv: bool = False
+    mmqa_dropout: float = 0.0
+    mmqa_dw_kernel_size: int = 3
+    is_multiscale: bool = False
+
+
+def _mmqa(
+    num_heads,
+    kv_dim,
+    kv_strides,
+    mmqa_avg_pool_kv=False,
+    is_multiscale=False,
+):
+    conf = MultiQueryAttentionBlockConfig(
+        num_heads=num_heads,
+        kv_dim=kv_dim,
+        kv_strides=kv_strides,
+        mmqa_avg_pool_kv=mmqa_avg_pool_kv,
+        is_multiscale=is_multiscale,
+    )
+    return conf
+
+
+# https://github.com/huggingface/new-model-addition-timm-gemma3p5-non-fork/blob/mobilenet-gemma3n-rw/timm/models/mobilenetv5.py#L596
+def gemma3n_mobilenet_def():
+    return [
+        # Stage 1: Edge Residuals
+        [_er(3, 128, 2)] + [_er(3, 128, 1)] * 2,
+        # Stage 2: Universal Inverted Residuals
+        [_uir(3, 5, 256, 2, 6.0)] + [_uir(k, 0, 256) for k in [5, 3, 5, 3]],
+        # Stage 3: Universal Inverted Residuals with Multi-Query Attention
+        [_uir(5, 5, 640, 2, 6.0)]
+        + [_uir(5, 0, 640)] * 7
+        + [_uir(0, 0, 640, 1, 1.0)]
+        + [_mmqa(12, 64, 2), _uir(0, 0, 640, 1, 2.0)] * 13
+        + [_mmqa(12, 64, 2), _uir(0, 0, 640, 1, 2.0, is_multiscale=True)],
+        # Stage 4: Universal Inverted Residuals with Multi-Query Attention
+        [_uir(5, 5, 1280, 2, 6.0)]
+        + [_mmqa(16, 96, 1), _uir(0, 0, 1280, 1, 2.0)] * 18
+        + [_mmqa(16, 96, 1), _uir(0, 0, 1280, 1, 2.0, is_multiscale=True)],
+    ]
+
+
+class VisionTower(nn.Module):
+    def __init__(self, config: VisionConfig):
+        super().__init__()
+        self.conv_stem = ConvNormAct(
+            Conv2dSame,
+            in_chs=3,
+            out_chs=64,
+            kernel_size=3,
+            stride=2,
+            padding=0,
+            eps=1e-05,
+            bias=True,
+        )
+        msfa_indices = (3, 4)
+        msfa_output_resolution = (16, 16)
+
+        (num_features, self.blocks) = self.build()
+        self.num_features = self.head_hidden_size = (
+            num_features  # output of msfa is output of forward_features()
+        )
+        self.msfa_indices = msfa_indices
+        self.msfa_output_resolution = msfa_output_resolution
+
+        self.msfa = MobileNetV5MultiScaleFusionAdapter(
+            in_chs=[1920],
+            out_chs=2048,
+            output_resolution=self.msfa_output_resolution,
+        )
+
+    def build(self):
+        blocks = []
+        in_chs = self.conv_stem.out_chs
+        for stage, block_config in enumerate(gemma3n_mobilenet_def()):
+            block_group = []
+            for config in block_config:
+                match config:
+                    case EdgeResidualConfig(
+                        kernel_size, filters, strides, expand_ratio, is_multiscale
+                    ):
+                        x = EdgeResidual(
+                            exp_kernel_size=kernel_size,
+                            in_chs=in_chs,
+                            out_chs=filters,
+                            stride=strides,
+                            expand_ratio=expand_ratio,
+                        )
+                        in_chs = filters  # in_chs of next is out_chs of prev
+                        block_group.append(x)
+                    case UniversalInvertedResidualConfig(
+                        start_dw_kernel_size,
+                        mid_dw_kernel_size,
+                        filters,
+                        strides,
+                        expand_ratio,
+                        is_multiscale,
+                    ):
+                        x = UniversalInvertedResidual(
+                            in_chs=in_chs,
+                            out_chs=filters,
+                            dw_kernel_size_start=start_dw_kernel_size,
+                            dw_kernel_size_mid=mid_dw_kernel_size,
+                            stride=strides,
+                            exp_ratio=expand_ratio,
+                        )
+                        in_chs = filters
+                        block_group.append(x)
+                    case MultiQueryAttentionBlockConfig(
+                        num_heads,
+                        kv_dim,
+                        kv_strides,
+                        mmqa_avg_pool_kv,
+                        is_multiscale,
+                    ):
+                        x = MobileAttention(
+                            in_chs=in_chs,
+                            out_chs=in_chs,
+                            stride=1,
+                            num_heads=num_heads,
+                            key_dim=kv_dim,
+                            value_dim=kv_dim,
+                            kv_stride=kv_strides,
+                            act_layer=None,
+                        )
+                        block_group.append(x)
+                    case _:
+                        continue
+            blocks.append(block_group)
+        return (in_chs, blocks)
+
+    def __call__(
+        self, x: mx.array, output_hidden_states: Optional[bool] = None
+    ) -> mx.array:
+        feat_idx = 0
+        x = x.transpose(0, 2, 3, 1)  # Convert from NCHW to NHWC
+        x = self.conv_stem(x)
+        intermediates = []
+
+        if feat_idx in self.msfa_indices:
+            intermediates.append(x)
+
+        # MBV5 is constructed of 4 stages, each stage is a group of blocks.
+        for block_group in self.blocks:
+            feat_idx += 1
+            for block in block_group:
+                x = block(x)
+
+            if feat_idx in self.msfa_indices:
+                intermediates.append(x)
+
+        x = self.msfa(intermediates)
+        return x
+
+
+class VisionModel(nn.Module):
+    def __init__(self, config: VisionConfig):
+        super().__init__()
+        self.model_type = config.model_type
+        if self.model_type not in ["gemma3", "gemma3_vision", "gemma3n_vision"]:
+            raise ValueError(f"Unsupported model type: {self.model_type}")
+
+        self.timm_model = VisionTower(config)
+
+    def __call__(
+        self, x: mx.array, output_hidden_states: Optional[bool] = None
+    ) -> mx.array:
+        return self.timm_model(x, output_hidden_states)
+
+    def sanitize(self, weights):
+        sanitized_weights = {}
+        skip_transpose = False
+        _, H, _, C = weights["vision_tower.timm_model.blocks.0.0.conv_exp.weight"].shape
+        if C > H:
+            skip_transpose = True
+
+        for k, v in weights.items():
+            # PyTorch conv2d weight: [out_channels, in_channels, kH, kW]
+            # MLX conv2d weight: [out_channels, kH, KW, in_channels]
+            if ("conv" in k and "weight" in k) or ("attn" and "proj.weight") in k:
+                if len(v.shape) == 4 and not skip_transpose:
+                    v = v.transpose(0, 2, 3, 1)
+            sanitized_weights[k] = v
+
+        return sanitized_weights