nexaai 1.0.29__cp310-cp310-macosx_14_0_universal2.whl

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

@@ -0,0 +1,522 @@
+import inspect
+from dataclasses import dataclass
+from typing import List, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from ..kernels import bicubic_interpolate
+
+
+@dataclass
+class VisionConfig:
+    model_type: str = "moonvit"
+    depth: int = 27
+    embed_dim: int = 1152
+    hidden_size: int = 1152
+    num_heads: int = 16
+    image_size: int = 384
+    patch_size: int = 14
+    vocab_size: int = 32000
+    mlp_ratio: float = 4.0
+    num_channels: int = 3
+    layer_norm_eps: float = 1e-6
+    intermediate_size: int = 4304
+    init_pos_emb_height: int = 64
+    init_pos_emb_width: int = 64
+    spatial_patch_size: int = 14
+    spatial_merge_size: int = 2
+    temporal_patch_size: int = 2
+    merge_kernel_size: list[int, int] = None
+
+    def __post_init__(self):
+        if self.merge_kernel_size is None:
+            self.merge_kernel_size = (self.spatial_merge_size, self.spatial_merge_size)
+
+    @classmethod
+    def from_dict(cls, params):
+        return cls(
+            **{
+                k: v
+                for k, v in params.items()
+                if k in inspect.signature(cls).parameters
+            }
+        )
+
+
+def check_array_shape(arr):
+    shape = arr.shape
+
+    # Check if the shape has 4 dimensions
+    if len(shape) != 4:
+        return False
+
+    out_channels, kH, KW, _ = shape
+
+    # Check if out_channels is the largest, and kH and KW are the same
+    if (out_channels >= kH) and (out_channels >= KW) and (kH == KW):
+        return True
+    else:
+        return False
+
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return mx.concatenate([-x2, x1], axis=-1)
+
+
+def apply_rotary_pos_emb_vision(tensor, freqs) -> mx.array:
+    orig_dtype = tensor.dtype
+
+    cos = mx.cos(freqs)
+    sin = mx.sin(freqs)
+
+    cos = mx.expand_dims(cos, axis=1)  # Equivalent to unsqueeze(1)
+    cos = mx.tile(cos, (1, 1, 2))  # Equivalent to repeat(1, 1, 2)
+    cos = mx.expand_dims(cos, axis=0)  # Equivalent to [None, ...]
+
+    sin = mx.expand_dims(sin, axis=1)  # Equivalent to unsqueeze(1)
+    sin = mx.tile(sin, (1, 1, 2))  # Equivalent to repeat(1, 1, 2)
+    sin = mx.expand_dims(sin, axis=0)  # Equivalent to [None, ...]
+
+    output = (tensor * cos) + (rotate_half(tensor) * sin)
+    return output.astype(orig_dtype)
+
+
+class VisionRotaryEmbedding(nn.Module):
+    def __init__(self, dim: int, theta: float = 10000.0) -> None:
+        super().__init__()
+        self.dim = dim
+        self.theta = theta
+
+    def __call__(self, seqlen: int) -> mx.array:
+        inv_freq = 1.0 / (
+            self.theta ** (mx.arange(0, self.dim, 2, dtype=mx.float32) / self.dim)
+        )
+        seq = mx.arange(seqlen.tolist(), dtype=inv_freq.dtype)
+        freqs = mx.outer(seq, inv_freq)
+        return freqs
+
+
+class Learnable2DInterpPosEmb(nn.Module):
+    def __init__(
+        self, height: int, width: int, dim: int, interpolation_mode: str = "bicubic"
+    ) -> None:
+        super().__init__()
+        self.height = height
+        self.width = width
+        self.interpolation_mode = interpolation_mode
+        self.weight = mx.ones((height, width, dim))
+
+    def __call__(self, x: mx.array, grid_hws: mx.array) -> mx.array:
+        pos_embs = []
+        for shape in grid_hws.tolist():
+            if shape == self.weight.shape[:-1]:
+                pos_embs.append(self.weight.flatten(end_axis=1))
+            else:
+                result = (
+                    bicubic_interpolate(
+                        mx.expand_dims(self.weight.transpose(2, 0, 1), axis=0),
+                        size=shape,
+                    )
+                    .squeeze(0)
+                    .transpose(1, 2, 0)
+                    .flatten(end_axis=1)
+                )
+
+                pos_embs.append(result)
+
+        out = x + mx.concatenate(pos_embs).astype(x.dtype)
+        return out
+
+
+class PatchEmbed(nn.Module):
+    def __init__(
+        self,
+        patch_size: int = 14,
+        num_channels: int = 3,
+        embed_dim: int = 1152,
+        init_pos_emb_height: int = 64,
+    ) -> None:
+        super().__init__()
+        self.patch_size = patch_size
+        self.num_channels = num_channels
+        self.embed_dim = embed_dim
+        self.init_pos_emb_height = init_pos_emb_height
+
+        self.proj = nn.Conv2d(
+            num_channels,
+            embed_dim,
+            kernel_size=patch_size,
+            stride=patch_size,
+            bias=True,
+        )
+        self.pos_emb = Learnable2DInterpPosEmb(
+            height=init_pos_emb_height, width=init_pos_emb_height, dim=embed_dim
+        )
+
+    def __call__(self, hidden_states: mx.array, grid_thw: mx.array) -> mx.array:
+        hidden_states = self.proj(hidden_states).swapaxes(1, 3)
+        hidden_states = hidden_states.reshape(hidden_states.shape[0], -1)
+        hidden_states = self.pos_emb(hidden_states, grid_thw)
+        return hidden_states
+
+
+def _apply_rope_input_validation(x, freqs_cis):
+    assert x.ndim == freqs_cis.ndim + 1, (x.shape, freqs_cis.shape)
+    assert x.shape[:-2] == freqs_cis.shape[:-1], (x.shape, freqs_cis.shape)
+    assert x.shape[-1] == 2 * freqs_cis.shape[-1], (x.shape, freqs_cis.shape)
+    assert freqs_cis.dtype == mx.complex64, freqs_cis.dtype
+
+
+def view_as_complex(x):
+    """
+    Convert a tensor with shape (..., 2) to a complex tensor with shape (...).
+    """
+    # Get real and imaginary parts
+    real, imag = x[..., 0], x[..., 1]
+    # Create complex tensor
+    return real + 1j * imag
+
+
+def view_as_real(x):
+    """
+    Convert a complex tensor with shape (...) to a real tensor with shape (..., 2).
+    """
+    # Get real and imaginary parts
+    real = mx.real(x)
+    imag = mx.imag(x)
+    # Combine into a tensor with last dimension 2
+    return mx.stack([real, imag], axis=-1)
+
+
+def apply_rope(
+    q: mx.array, k: mx.array, freqs_cis: mx.array
+) -> tuple[mx.array, mx.array]:
+    """
+    Args: (The leading dimensions of all inputs should be the same)
+        q: query, array of shape (..., num_heads, head_dim)
+        k: key, array of shape (..., num_heads, head_dim)
+        freqs_cis: array of shape (..., head_dim/2), dtype=mx.complex64. It contains the precomputed cis(freqs) for each position in the 2D grid.
+    Returns:
+        xq_out, xk_out: arrays of shape (..., num_heads, head_dim)
+    """
+    _apply_rope_input_validation(q, freqs_cis)
+    _apply_rope_input_validation(k, freqs_cis)
+
+    freqs_cis = mx.expand_dims(freqs_cis, axis=-2)  # ..., 1, head_dim/2
+    # ..., num_heads, head_dim/2
+    q_ = view_as_complex(q.astype(mx.float32).reshape(*q.shape[:-1], -1, 2))
+    k_ = view_as_complex(k.astype(mx.float32).reshape(*k.shape[:-1], -1, 2))
+    q_out = view_as_real(q_ * freqs_cis).flatten(-2)  # ..., num_heads, head_dim
+    k_out = view_as_real(k_ * freqs_cis).flatten(-2)  # ..., num_heads, head_dim
+    return q_out.astype(q.dtype), k_out.astype(k.dtype)
+
+
+class Attention(nn.Module):
+    def __init__(self, dim: int, num_heads: int = 16) -> None:
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = head_dim = dim // num_heads
+        self.scale = head_dim**-0.5
+        self.wqkv = nn.Linear(dim, dim * 3, bias=True)
+        self.wo = nn.Linear(dim, dim, bias=True)
+
+    def __call__(
+        self, x: mx.array, cu_seqlens: mx.array, rotary_pos_emb: mx.array = None
+    ) -> mx.array:
+        seq_length = x.shape[0]
+        qkv = self.wqkv(x)
+
+        qkv_shape = qkv.shape[:-1] + (
+            3,
+            self.num_heads,
+            self.head_dim,
+        )
+        # xqkv: (batch_size, seqlen, 3, nheads, headdim)
+        qkv = qkv.reshape(*qkv_shape)
+
+        q, k, v = mx.split(qkv, 3, axis=1)
+        q = q.squeeze(1)
+        k = k.squeeze(1)
+        v = v.squeeze(1)
+
+        q, k = apply_rope(q, k, rotary_pos_emb)
+
+        attention_mask = mx.zeros((1, seq_length, seq_length), dtype=x.dtype)
+
+        # Create attention mask for each sequence in the batch
+        for i in range(1, len(cu_seqlens)):
+            start = int(cu_seqlens[i - 1])
+            end = int(cu_seqlens[i])
+            attention_mask[..., start:end, start:end] = 1
+
+        q = q.transpose(1, 0, 2)
+        k = k.transpose(1, 0, 2)
+        v = v.transpose(1, 0, 2)
+
+        attn_weight = q @ k.swapaxes(-2, -1) / mx.sqrt(q.shape[-1])
+        attn_weight += attention_mask
+        attn_weight = mx.softmax(attn_weight, axis=-1).astype(q.dtype)
+
+        attn_output = attn_weight @ v
+        attn_output = attn_output.transpose(1, 0, 2)
+        attn_output = attn_output.reshape(seq_length, -1)
+        return self.wo(attn_output)
+
+
+class MLP(nn.Module):
+    def __init__(self, dim, hidden_dim):
+        super().__init__()
+        self.activation_fn = nn.GELU()
+        self.fc0 = nn.Linear(dim, hidden_dim)
+        self.fc1 = nn.Linear(hidden_dim, dim)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        x = self.activation_fn(self.fc0(x))
+        x = self.fc1(x)
+        return x
+
+
+class Qwen2VLVisionBlock(nn.Module):
+    def __init__(self, config: VisionConfig) -> None:
+        super().__init__()
+        self.norm0 = nn.LayerNorm(config.embed_dim, eps=1e-6)
+        self.norm1 = nn.LayerNorm(config.embed_dim, eps=1e-6)
+
+        self.attn = Attention(dim=config.embed_dim, num_heads=config.num_heads)
+        self.mlp = MLP(dim=config.embed_dim, hidden_dim=config.intermediate_size)
+
+    def __call__(self, hidden_states, cu_seqlens, rotary_pos_emb) -> mx.array:
+        hidden_states = hidden_states + self.attn(
+            self.norm0(hidden_states),
+            cu_seqlens=cu_seqlens,
+            rotary_pos_emb=rotary_pos_emb,
+        )
+        hidden_states = hidden_states + self.mlp(self.norm1(hidden_states))
+        return hidden_states
+
+
+class Rope2DPosEmb(nn.Module):
+    """2D rotary position embedding with multi-resolution support.
+
+    This class is intended to be used in the following way:
+    1. Before training, create an instance of Rope2DPosEmb. This instance will hold the precomputed cis.
+    2. Before each forward pass, call `get_freqs_cis_by_*` to get the `freqs_cis` tensor for this iteration.
+    3. During the forward pass, pass the `freqs_cis` tensor to each attention layer, and call `apply` just before each attention operation.
+        The rope is shared across all attention layers and all heads.
+
+    Refs:
+    - RoFormer: https://arxiv.org/abs/2104.09864
+    - VisionLLaMA: https://arxiv.org/abs/2403.00522
+    - https://github.com/Meituan-AutoML/VisionLLaMA/blob/main/dit/models.py
+
+    Args:
+        dim (int): usually the multi-head attention dimension, should be divisible by 4 (TODO: relax this constraint if needed)
+        max_height (int): the maximum height of the 2D grid
+        max_width (int): the maximum width of the 2D grid
+        theta_base (float): the base of the theta
+    """
+
+    def __init__(self, dim: int, max_height: int, max_width: int, theta_base=10000):
+        super().__init__()
+        self.dim = dim
+        assert self.dim % 4 == 0, "dim must be divisible by 4"
+        self.max_height = max_height
+        self.max_width = max_width
+        self.theta_base = theta_base
+
+        self._freqs_cis = None
+
+    def extra_repr(self):
+        return f"dim={self.dim}, max_height={self.max_height}, max_width={self.max_width}, theta_base={self.theta_base}"
+
+    def _precompute_freqs_cis(self) -> mx.array:
+        """Calculate the cis(freqs) for each position in the 2D grid.
+
+        Return: complex array of shape (max_height, max_width, dim//2) and value:
+            height axis: ret[h, w, 2*i] = cis(h * theta_base**(-4*i/dim))
+            weight axis: ret[h, w, 2*i+1] = cis(w * theta_base**(-4*i/dim))   with (i in [0, dim//4))
+            note: `cis` is a mathematical notation defined by cis x = cos x + i sin x,
+        """
+        N = self.max_height * self.max_width
+        flat_pos = mx.arange(0, N, dtype=mx.float32)
+        x_pos = flat_pos % self.max_width
+        y_pos = flat_pos // self.max_width
+        dim_range = mx.arange(0, self.dim, 4)[: (self.dim // 4)].astype(
+            mx.float32
+        )  # C/4
+        freqs = 1.0 / (self.theta_base ** (dim_range / self.dim))
+        x_freqs = mx.outer(x_pos, freqs)  # N, C/4
+        y_freqs = mx.outer(y_pos, freqs)  # N, C/4
+
+        # Create complex numbers using cos and sin
+        x_cos = mx.cos(x_freqs)
+        x_sin = mx.sin(x_freqs)
+        y_cos = mx.cos(y_freqs)
+        y_sin = mx.sin(y_freqs)
+
+        # Create complex numbers
+        x_cis = x_cos + 1j * x_sin  # N, C/4
+        y_cis = y_cos + 1j * y_sin  # N, C/4
+
+        # N, C/4, 2
+        freqs_cis = mx.stack([x_cis, y_cis], axis=-1)
+
+        # max_height, max_width, C/2
+        freqs_cis = freqs_cis.reshape(self.max_height, self.max_width, -1)
+        return freqs_cis
+
+    def get_freqs_cis(self, grid_hws: mx.array) -> mx.array:
+        """
+        Args:
+            grid_hws (mx.array): grid height and width
+
+        Returns:
+            freqs_cis: array of shape (sum(t * height * width), dim//2)
+        """
+        if self._freqs_cis is None:
+            self._freqs_cis = self._precompute_freqs_cis()
+
+        shapes = grid_hws.tolist()
+        assert all(
+            1 <= h <= self.max_height and 1 <= w <= self.max_width for h, w in shapes
+        ), (
+            shapes,
+            self.max_height,
+            self.max_width,
+        )
+
+        freqs_cis_list = []
+        for h, w in shapes:
+            # Get the slice of precomputed frequencies for this shape
+            shape_freqs = self._freqs_cis[:h, :w]
+            # Reshape to flatten the spatial dimensions
+            shape_freqs = shape_freqs.reshape(-1, self.dim // 2)
+            freqs_cis_list.append(shape_freqs)
+
+        freqs_cis = mx.concatenate(freqs_cis_list, axis=0)
+        return freqs_cis
+
+
+def patch_merger(
+    x: mx.array,
+    grid_hws: mx.array,
+    merge_kernel_size: list[int, int] = (2, 2),
+) -> List[mx.array]:
+    d_model = x.shape[-1]
+
+    outputs = []
+    pre_sum = 0
+    for x_shape in grid_hws.tolist():
+        height, width = x_shape[0], x_shape[1]
+        # Get the current sequence
+        seq = x[pre_sum : pre_sum + height * width]
+        # Reshape along self.merge_kernel_size and concat to the last dimension
+        kernel_height, kernel_width = merge_kernel_size
+        new_height, new_width = height // kernel_height, width // kernel_width
+        reshaped_seq = seq.reshape(
+            new_height, kernel_height, new_width, kernel_width, d_model
+        )
+        reshaped_seq = mx.transpose(reshaped_seq, (0, 2, 1, 3, 4))
+        padded_seq = reshaped_seq.reshape(
+            new_height * new_width, kernel_height * kernel_width, -1
+        )
+        outputs.append(padded_seq)
+        pre_sum += height * width
+
+    return outputs
+
+
+class VisionModel(nn.Module):
+
+    def __init__(self, config: VisionConfig) -> None:
+        super().__init__()
+        self.config = config
+        self.model_type = config.model_type
+        if self.model_type not in ["qwen2_vl", "moonvit"]:
+            raise ValueError(f"Unsupported model type: {self.model_type}")
+        self.spatial_merge_size = config.spatial_merge_size
+        self.merge_kernel_size = config.merge_kernel_size
+
+        self.patch_embed = PatchEmbed(
+            patch_size=config.patch_size,
+            num_channels=config.num_channels,
+            embed_dim=config.embed_dim,
+            init_pos_emb_height=config.init_pos_emb_height,
+        )
+
+        head_dim = config.embed_dim // config.num_heads
+        self.rotary_pos_emb = VisionRotaryEmbedding(head_dim // 2)
+
+        self.blocks = [Qwen2VLVisionBlock(config) for _ in range(config.depth)]
+        self.final_layernorm = nn.LayerNorm(config.hidden_size, eps=1e-6)
+        self.rope_pos_emb = Rope2DPosEmb(head_dim, 512, 512)
+
+    def __call__(
+        self,
+        hidden_states: mx.array,
+        grid_thw: mx.array,
+        output_hidden_states: Optional[bool] = None,
+    ) -> mx.array:
+
+        hidden_states = self.patch_embed(hidden_states, grid_thw)
+        rotary_pos_emb = self.rope_pos_emb.get_freqs_cis(grid_thw)
+
+        # Assuming grid_thw has shape (batch_size, 3)
+        batch_size = grid_thw.shape[0]
+
+        # Calculate cu_seqlens for each item in the batch
+        lengths = mx.concatenate(
+            (
+                mx.zeros((1,), dtype=grid_thw.dtype),
+                grid_thw[:, 0] * grid_thw[:, 1],
+            )
+        )
+        cu_seqlens = mx.cumsum(lengths.astype(mx.int32), axis=0)
+
+        encoder_states = (hidden_states,) if output_hidden_states else None
+
+        for blk in self.blocks:
+            hidden_states = blk(
+                hidden_states, cu_seqlens=cu_seqlens, rotary_pos_emb=rotary_pos_emb
+            )
+            if output_hidden_states:
+                encoder_states = encoder_states + (hidden_states,)
+
+        hidden_states = self.final_layernorm(hidden_states)
+
+        hidden_states = patch_merger(
+            hidden_states, grid_thw, merge_kernel_size=self.merge_kernel_size
+        )
+
+        return hidden_states
+
+    def sanitize(self, weights):
+        sanitized_weights = {}
+        for k, v in weights.items():
+            if "position_ids" in k:
+                # Remove unused position_ids
+                continue
+            elif "patch_embed.proj.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]
+                if check_array_shape(v):
+                    sanitized_weights[k] = v
+                else:
+                    sanitized_weights[k] = v.transpose(0, 2, 3, 1)
+
+            elif "vision_tower.blocks" in k:
+                if "attn" not in k and ("wqkv" in k or "wo" in k):
+                    new_key = k.replace("wqkv", "attn.wqkv").replace("wo", "attn.wo")
+                    sanitized_weights[new_key] = v
+                else:
+                    sanitized_weights[k] = v
+            else:
+                sanitized_weights[k] = v
+
+        return sanitized_weights

nexaai/mlx_backend/vlm/modeling/models/llama4/__init__.py

@@ -0,0 +1,8 @@
+from .llama4 import (
+    LanguageModel,
+    Model,
+    ModelConfig,
+    TextConfig,
+    VisionConfig,
+    VisionModel,
+)