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

mlx_vlm/models/internvl_chat/vision.py

@@ -0,0 +1,265 @@
+from typing import Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+import numpy as np
+
+from ..base import interpolate
+from .config import VisionConfig
+
+
+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
+
+
+class Attention(nn.Module):
+    def __init__(self, config: VisionConfig):
+        super().__init__()
+
+        if (config.hidden_size % config.num_attention_heads) != 0:
+            raise ValueError(
+                "The input feature dimensions should be divisible by the "
+                f"number of heads ({config.hidden_size} % {config.num_attention_heads}) != 0"
+            )
+
+        self.dims = dims = config.hidden_size
+
+        self.num_heads = config.num_attention_heads
+        head_dim = config.hidden_size // config.num_attention_heads
+        self.scale = head_dim**-0.5
+        self.qkv_bias = config.qkv_bias
+
+        self.qkv = nn.Linear(dims, 3 * dims, bias=config.qkv_bias)
+        self.proj = nn.Linear(dims, dims)
+
+        self.qk_normalization = config.qk_normalization
+
+        if self.qk_normalization:
+            self.q_norm = nn.RMSNorm(dims, eps=config.layer_norm_eps)
+            self.k_norm = nn.RMSNorm(dims, eps=config.layer_norm_eps)
+
+    def __call__(self, x, mask=None):
+        B, L, C = x.shape
+        qkv = self.qkv(x).reshape(B, L, 3, self.num_heads, C // self.num_heads)
+        qkv = qkv.transpose(2, 0, 3, 1, 4)
+        queries, keys, values = (
+            qkv[0],
+            qkv[1],
+            qkv[2],
+        )  # Each has shape (B, groups, N, C//groups)
+
+        if self.qk_normalization:
+            B_, H_, N_, D_ = queries.shape
+            queries = (
+                self.q_norm(queries.transpose(0, 2, 1, 3).flatten(-2, -1))
+                .reshape(B_, N_, H_, D_)
+                .transpose(0, 2, 1, 3)
+            )
+            keys = (
+                self.k_norm(keys.transpose(0, 2, 1, 3).flatten(-2, -1))
+                .reshape(B_, N_, H_, D_)
+                .transpose(0, 2, 1, 3)
+            )
+
+        output = mx.fast.scaled_dot_product_attention(
+            queries, keys, values, scale=self.scale, mask=mask
+        )
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.proj(output)
+
+
+class MLP(nn.Module):
+    def __init__(self, config: VisionConfig):
+        super().__init__()
+        self.activation_fn = nn.GELU(approx="precise")
+        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.fc1(x)
+        x = self.activation_fn(x)
+        x = self.fc2(x)
+        return x
+
+
+class EncoderLayer(nn.Module):
+    def __init__(self, config: VisionConfig, drop_path_rate: float = 0.0):
+        super().__init__()
+        self.embed_dim = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.norm_type = getattr(config, "norm_type", "layer_norm")
+
+        self.attn = Attention(config)
+        self.mlp = MLP(config)
+
+        if self.norm_type == "layer_norm":
+            self.norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
+            self.norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
+        elif self.norm_type == "rms_norm":
+            self.norm1 = nn.RMSNorm(self.embed_dim, eps=config.layer_norm_eps)
+            self.norm2 = nn.RMSNorm(self.embed_dim, eps=config.layer_norm_eps)
+        else:
+            raise ValueError(f"Unsupported normalization type: {self.norm_type}")
+
+        self.ls1 = mx.ones((self.embed_dim,))
+        self.ls2 = mx.ones((self.embed_dim,))
+
+        self.drop_path1 = (
+            nn.Dropout(drop_path_rate) if drop_path_rate > 0.0 else nn.Identity()
+        )
+        self.drop_path2 = (
+            nn.Dropout(drop_path_rate) if drop_path_rate > 0.0 else nn.Identity()
+        )
+
+    def __call__(self, x: mx.array, mask: Optional[mx.array] = None) -> mx.array:
+        dtype = x.dtype
+        x = x + self.drop_path1(self.attn(self.norm1(x).astype(dtype)) * self.ls1)
+
+        x = x + self.drop_path2(self.mlp(self.norm2(x).astype(dtype)) * self.ls2)
+
+        return x.astype(dtype)
+
+
+class Encoder(nn.Module):
+    def __init__(self, config: VisionConfig):
+        super().__init__()
+        dpr = [
+            mx.array(x)
+            for x in np.linspace(0, config.drop_path_rate, config.num_hidden_layers)
+        ]
+        self.layers = [
+            EncoderLayer(config, dpr[i]) for i in range(config.num_hidden_layers)
+        ]
+
+    def __call__(
+        self,
+        x: mx.array,
+        output_hidden_states: Optional[bool] = None,
+        mask: Optional[mx.array] = None,
+    ) -> mx.array:
+        encoder_states = (x,) if output_hidden_states else None
+        h = x
+        for l in self.layers:
+            x = l(x, mask=mask)
+            if output_hidden_states:
+                encoder_states = encoder_states + (x,)
+
+            h = x
+
+        return (h, encoder_states)
+
+
+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.random.normal((1, 1, self.embed_dim))
+
+        self.patch_embedding = nn.Conv2d(
+            in_channels=3,
+            out_channels=self.embed_dim,
+            kernel_size=self.patch_size,
+            stride=self.patch_size,
+        )
+
+        self.num_patches = (self.image_size // self.patch_size) ** 2
+        self.num_positions = self.num_patches + 1
+
+        self.position_embedding = mx.random.normal(
+            (1, self.num_positions, self.embed_dim)
+        )
+
+    def _get_pos_embed(self, pos_embed, H, W):
+        target_dtype = pos_embed.dtype
+        pos_embed = pos_embed.reshape(
+            1,
+            self.image_size // self.patch_size,
+            self.image_size // self.patch_size,
+            -1,
+        ).transpose(0, 3, 1, 2)
+        pos_embed = interpolate(pos_embed, (H, W))
+        pos_embed = (
+            pos_embed.reshape(1, -1, H * W).transpose(0, 2, 1).astype(target_dtype)
+        )
+        return pos_embed
+
+    def __call__(self, x: mx.array) -> mx.array:
+        target_dtype = self.patch_embedding.weight.dtype
+        patch_embeds = self.patch_embedding(x).transpose(
+            0, 3, 1, 2
+        )  # shape = [*, channel, width, height]
+        batch_size, _, height, width = patch_embeds.shape
+        patch_embeds = mx.flatten(patch_embeds, start_axis=2).transpose(0, 2, 1)
+        class_embeds = mx.broadcast_to(
+            self.class_embedding, (batch_size, 1, self.embed_dim)
+        ).astype(target_dtype)
+        embeddings = mx.concatenate([class_embeds, patch_embeds], axis=1)
+        position_embedding = mx.concatenate(
+            [
+                self.position_embedding[:, :1, :],
+                self._get_pos_embed(self.position_embedding[:, 1:, :], height, width),
+            ],
+            axis=1,
+        )
+        embeddings = embeddings + position_embedding.astype(target_dtype)
+
+        return embeddings
+
+
+class VisionModel(nn.Module):
+    def __init__(self, config: VisionConfig):
+        super().__init__()
+        self.model_type = config.model_type
+        if self.model_type not in ["siglip_vision_model", "intern_vit_6b"]:
+            raise ValueError(f"Unsupported model type: {self.model_type}")
+
+        self.embeddings = VisionEmbeddings(config)
+        self.encoder = Encoder(config)
+
+    def __call__(
+        self,
+        x: mx.array,
+        output_hidden_states: Optional[bool] = None,
+    ) -> mx.array:
+        x = self.embeddings(x)
+        last_hidden_state, encoder_outputs = self.encoder(
+            x=x, output_hidden_states=output_hidden_states, mask=None
+        )
+        pooler_output = last_hidden_state[:, 0, :]
+        return last_hidden_state, pooler_output, encoder_outputs[1:]
+
+    def sanitize(self, 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]
+                if check_array_shape(v):
+                    sanitized_weights[k] = v
+                else:
+                    sanitized_weights[k] = v.transpose(0, 2, 3, 1)
+            else:
+                sanitized_weights[k] = v
+
+        return sanitized_weights

mlx_vlm/models/interpolate.py

@@ -0,0 +1,183 @@
+import mlx.core as mx
+import numpy as np
+
+
+def gaussian_blur_axis(image, sigma, axis):
+    """
+    Applies a 1D Gaussian blur along the given axis.
+    This version works for arrays with any number of dimensions.
+    """
+    radius = int(3 * sigma)
+    if radius < 1:
+        return image
+    x = mx.arange(-radius, radius + 1)
+    kernel = mx.exp(-(x**2) / (2 * sigma**2))
+    kernel = kernel / mx.sum(kernel)
+
+    # MLX doesn't have a direct apply_along_axis equivalent,
+    # so we'll implement the convolution differently based on the axis
+
+    # Helper function to apply 1D convolution along specific axis
+    def conv_1d(array, kernel, axis):
+        # Reshape kernel to broadcast along the right dimensions
+        kernel_shape = [1] * image.ndim
+        kernel_shape[axis] = len(kernel)
+        kernel_reshaped = kernel.reshape(kernel_shape)
+
+        # Pad the array
+        pad_width = [(0, 0)] * image.ndim
+        pad_width[axis] = (radius, radius)
+        padded = mx.pad(array, pad_width, mode="edge")
+
+        # Perform convolution via sliding window sum
+        result = mx.zeros_like(array)
+        slices = [slice(None)] * padded.ndim
+
+        for i in range(2 * radius + 1):
+            slices[axis] = slice(i, i + array.shape[axis])
+            result = result + padded[tuple(slices)] * kernel_reshaped
+
+        return result
+
+    return conv_1d(image, kernel, axis)
+
+
+def bilinear_interpolate(image, new_height, new_width, align_corners=False):
+    """
+    Performs bilinear interpolation on an array whose spatial dimensions are the first two.
+    It supports extra dimensions (e.g. channels or batch dimensions that have been moved to the trailing axes).
+    """
+    # image is assumed to have shape (H, W, ...) where H and W are spatial dimensions.
+    H_in, W_in = image.shape[0], image.shape[1]
+
+    # Compute sampling positions in the input image.
+    if new_height == 1:
+        row_positions = mx.array([0.0])
+    else:
+        if align_corners:
+            row_positions = mx.linspace(0, H_in - 1, new_height)
+        else:
+            row_positions = (mx.arange(new_height) + 0.5) * H_in / new_height - 0.5
+
+    if new_width == 1:
+        col_positions = mx.array([0.0])
+    else:
+        if align_corners:
+            col_positions = mx.linspace(0, W_in - 1, new_width)
+        else:
+            col_positions = (mx.arange(new_width) + 0.5) * W_in / new_width - 0.5
+
+    # Compute floor and ceil indices.
+    row_floor = mx.floor(row_positions).astype(mx.int32)
+    col_floor = mx.floor(col_positions).astype(mx.int32)
+    row_ceil = row_floor + 1
+    col_ceil = col_floor + 1
+
+    row_floor = mx.clip(row_floor, 0, H_in - 1)
+    row_ceil = mx.clip(row_ceil, 0, H_in - 1)
+    col_floor = mx.clip(col_floor, 0, W_in - 1)
+    col_ceil = mx.clip(col_ceil, 0, W_in - 1)
+
+    row_weight = row_positions - row_floor  # shape (new_height,)
+    col_weight = col_positions - col_floor  # shape (new_width,)
+
+    # Use advanced indexing for gather operations
+    # Create meshgrid for coordinates
+    row_floor_grid, col_floor_grid = mx.meshgrid(row_floor, col_floor, indexing="ij")
+    row_ceil_grid, col_floor_grid = mx.meshgrid(row_ceil, col_floor, indexing="ij")
+    row_floor_grid, col_ceil_grid = mx.meshgrid(row_floor, col_ceil, indexing="ij")
+    row_ceil_grid, col_ceil_grid = mx.meshgrid(row_ceil, col_ceil, indexing="ij")
+
+    # Gather the four surrounding pixels using take_along_axis
+    # For higher dimensional arrays, we'll need to reshape and broadcast
+    extra_dims = image.ndim - 2
+
+    def gather_pixels(row_indices, col_indices):
+        # Flatten the spatial dimensions for gathering
+        flat_indices = row_indices * W_in + col_indices
+        flat_image = mx.reshape(image, (-1,) + image.shape[2:])
+        # Gather and reshape back
+        gathered = mx.take(flat_image, flat_indices.reshape(-1), axis=0)
+        return mx.reshape(gathered, (new_height, new_width) + image.shape[2:])
+
+    top_left = gather_pixels(row_floor_grid, col_floor_grid)
+    top_right = gather_pixels(row_floor_grid, col_ceil_grid)
+    bottom_left = gather_pixels(row_ceil_grid, col_floor_grid)
+    bottom_right = gather_pixels(row_ceil_grid, col_ceil_grid)
+
+    # Expand the weights to have shape (new_height, new_width, *[1]*extra_dims)
+    r_weight = row_weight.reshape(new_height, 1, *([1] * extra_dims))
+    c_weight = col_weight.reshape(1, new_width, *([1] * extra_dims))
+
+    # Perform bilinear interpolation.
+    result = (
+        (1 - r_weight) * (1 - c_weight) * top_left
+        + (1 - r_weight) * c_weight * top_right
+        + r_weight * (1 - c_weight) * bottom_left
+        + r_weight * c_weight * bottom_right
+    )
+    return result
+
+
+def resize_bilinear(image, new_size, align_corners=False, antialias=True):
+    """
+    Resizes an image (or embedding tensor) to new_size=(new_height, new_width)
+    using bilinear interpolation with MLX.
+
+    Supports:
+      - 2D: (H, W)
+      - 3D: (H, W, C)
+      - 4D: (B, C, H, W)  (assumed for typical image batches)
+    """
+    new_height, new_width = new_size
+
+    # Convert numpy arrays to MLX arrays if needed
+    if isinstance(image, np.ndarray):
+        image = mx.array(image)
+
+    if image.ndim == 2 or image.ndim == 3:
+        # Assume spatial dims are the first two.
+        resized = image
+        H_in, W_in = image.shape[:2]
+        if antialias:
+            if new_height < H_in:
+                scale_y = new_height / H_in
+                sigma_y = (1 / scale_y - 1) / 2.0  # heuristic
+                if sigma_y > 0:
+                    resized = gaussian_blur_axis(resized, sigma_y, axis=0)
+            if new_width < W_in:
+                scale_x = new_width / W_in
+                sigma_x = (1 / scale_x - 1) / 2.0
+                if sigma_x > 0:
+                    resized = gaussian_blur_axis(resized, sigma_x, axis=1)
+        resized = bilinear_interpolate(
+            resized, new_height, new_width, align_corners=align_corners
+        )
+        return resized
+
+    elif image.ndim == 4:
+        # Assume shape is (B, C, H, W) (typical PyTorch/MLX format).
+        B, C, H_in, W_in = image.shape
+        # Permute to bring spatial dims to the front: (H, W, B, C)
+        image_perm = mx.transpose(image, (2, 3, 0, 1))
+        resized = image_perm
+        if antialias:
+            if new_height < H_in:
+                scale_y = new_height / H_in
+                sigma_y = (1 / scale_y - 1) / 2.0
+                if sigma_y > 0:
+                    resized = gaussian_blur_axis(resized, sigma_y, axis=0)
+            if new_width < W_in:
+                scale_x = new_width / W_in
+                sigma_x = (1 / scale_x - 1) / 2.0
+                if sigma_x > 0:
+                    resized = gaussian_blur_axis(resized, sigma_x, axis=1)
+        resized = bilinear_interpolate(
+            resized, new_height, new_width, align_corners=align_corners
+        )
+        # Permute back to (B, C, new_height, new_width)
+        resized = mx.transpose(resized, (2, 3, 0, 1))
+        return resized
+
+    else:
+        raise ValueError("Unsupported image dimensions.")

mlx_vlm/models/jina_vlm/__init__.py

@@ -0,0 +1,3 @@
+from .config import ModelConfig, TextConfig, VisionConfig
+from .image_processor import ImageProcessor
+from .jina_vlm import JinaVLMProcessor, LanguageModel, Model, VisionModel

mlx_vlm/models/jina_vlm/config.py

@@ -0,0 +1,142 @@
+from dataclasses import dataclass, field
+from typing import Tuple
+
+from ..base import BaseModelConfig
+
+
+@dataclass
+class VisionConfig(BaseModelConfig):
+    """Vision encoder configuration for Jina VLM."""
+
+    model_type: str = "jina_vlm"
+    hidden_size: int = 1152
+    num_hidden_layers: int = 27
+    num_attention_heads: int = 16
+    head_dim: int = 72
+    patch_size: int = 14
+    image_size: int = 378
+    num_channels: int = 3
+    intermediate_size: int = 4304
+    layer_norm_eps: float = 1e-6
+    use_bias: bool = True
+    use_cls_token: bool = False
+    post_layer_norm: bool = True
+    activation: str = "gelu_pytorch_tanh"
+    vit_layers: Tuple[int, ...] = (-4, -10)
+    output_size: int = 2048
+    # Connector config
+    pooling_h: int = 2
+    pooling_w: int = 2
+    connector_hidden_size: int = 6144
+
+
+@dataclass
+class TextConfig(BaseModelConfig):
+    """Text decoder configuration for Jina VLM."""
+
+    model_type: str = "jina_vlm"
+    hidden_size: int = 2048
+    num_hidden_layers: int = 28
+    num_attention_heads: int = 16
+    num_key_value_heads: int = 8
+    head_dim: int = 128
+    vocab_size: int = 151936
+    additional_vocab_size: int = 128
+    intermediate_size: int = 6144
+    rms_norm_eps: float = 1e-6
+    rope_theta: float = 1000000.0
+    max_position_embeddings: int = 40960
+    use_qk_norm: bool = True
+    tie_word_embeddings: bool = False
+
+
+@dataclass
+class ModelConfig(BaseModelConfig):
+    """Full Jina VLM configuration."""
+
+    text_config: TextConfig = field(default_factory=TextConfig)
+    vision_config: VisionConfig = field(default_factory=VisionConfig)
+    model_type: str = "jina_vlm"
+    vocab_size: int = 151936
+    bos_token_id: int = 151643
+    eos_token_id: int = 151643
+    pad_token_id: int = 151643
+    image_token_index: int = 151940  # <|image|>
+    image_token_id: int = 151940  # <|image|>
+    image_start_token_id: int = 151666  # <im_start>
+    image_end_token_id: int = 151667  # <im_end>
+    image_patch_token_id: int = 151665  # <im_patch>
+    image_column_token_id: int = 151668  # <im_col>
+    ignore_index: int = -100
+    tie_word_embeddings: bool = False
+
+    @classmethod
+    def from_dict(cls, params):
+        # Parse vision config
+        vision_cfg = params.get("vision_config", {})
+        vision_block = vision_cfg.get("block_config", {})
+        vision_attn = vision_block.get("attn_config", {})
+        vision_ffn = vision_block.get("ffn_config", {})
+        vl_connector = vision_cfg.get("vl_connector_config", {})
+        connector_mlp = vl_connector.get("mlp_projector_config", {})
+
+        vision_config = VisionConfig(
+            hidden_size=vision_cfg.get("hidden_size", 1152),
+            num_hidden_layers=vision_cfg.get("n_layers", 27),
+            num_attention_heads=vision_attn.get("n_heads", 16),
+            head_dim=vision_attn.get("head_dim", 72),
+            patch_size=vision_cfg.get("patch_size", 14),
+            image_size=(
+                vision_cfg.get("input_size", [378, 378])[0]
+                if isinstance(vision_cfg.get("input_size"), list)
+                else 378
+            ),
+            num_channels=vision_cfg.get("n_channels", 3),
+            intermediate_size=vision_ffn.get("size", 4304),
+            use_bias=vision_attn.get("q_bias", True),
+            use_cls_token=vision_cfg.get("use_cls_token", False),
+            post_layer_norm=vision_cfg.get("post_lnorm", True),
+            activation=vision_ffn.get("activation_type", "gelu_pytorch_tanh"),
+            vit_layers=tuple(vision_cfg.get("vit_layers", [-4, -10])),
+            output_size=vision_cfg.get("output_size", 2048),
+            pooling_h=vl_connector.get("pooling_h", 2),
+            pooling_w=vl_connector.get("pooling_w", 2),
+            connector_hidden_size=connector_mlp.get("size", 6144),
+        )
+
+        # Parse text config
+        text_cfg = params.get("text_config", {})
+        text_block = text_cfg.get("block_config", {})
+        text_attn = text_block.get("attn_config", {})
+        text_ffn = text_block.get("ffn_config", {})
+        text_lnorm = text_block.get("lnorm_config", {})
+
+        text_config = TextConfig(
+            hidden_size=text_cfg.get("hidden_size", 2048),
+            num_hidden_layers=text_cfg.get(
+                "n_layers", text_cfg.get("num_hidden_layers", 28)
+            ),
+            num_attention_heads=text_attn.get("n_heads", 16),
+            num_key_value_heads=text_attn.get("n_kv_heads", 8),
+            head_dim=text_attn.get("head_dim", 128),
+            vocab_size=text_cfg.get("vocab_size", 151936),
+            additional_vocab_size=text_cfg.get("additional_vocab_size", 128),
+            intermediate_size=text_ffn.get("size", 6144),
+            rms_norm_eps=text_lnorm.get("eps", 1e-6),
+            rope_theta=text_cfg.get("rope_theta", 1000000.0),
+            max_position_embeddings=text_cfg.get("max_sequence_length", 40960),
+            use_qk_norm=text_attn.get("q_lnorm", True),
+            tie_word_embeddings=text_cfg.get("tie_word_embeddings", False),
+        )
+
+        return cls(
+            text_config=text_config,
+            vision_config=vision_config,
+            model_type=params.get("model_type", "jina_vlm"),
+            vocab_size=params.get("vocab_size", text_config.vocab_size),
+            bos_token_id=params.get("bos_token_id", 151643),
+            eos_token_id=params.get("eos_token_id", 151643),
+            pad_token_id=params.get("pad_token_id", 151643),
+            image_token_index=params.get("image_token_index", 151940),
+            tie_word_embeddings=params.get("tie_word_embeddings", False),
+        )