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

mlx_vlm/models/qwen3_vl/qwen3_vl.py

@@ -0,0 +1,166 @@
+from typing import Optional
+
+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 Model(nn.Module):
+    def __init__(self, config: ModelConfig):
+        super().__init__()
+        self.config = config
+        self.vision_tower = VisionModel(config.vision_config)
+        self.language_model = LanguageModel(config.text_config, config)
+
+    def get_input_embeddings(
+        self,
+        input_ids: Optional[mx.array] = None,
+        pixel_values: Optional[mx.array] = None,
+        **kwargs,
+    ):
+        image_grid_thw = kwargs.get("image_grid_thw", None)
+        video_grid_thw = kwargs.get("video_grid_thw", None)
+        mask = kwargs.get("mask", None)
+        grid_thw = image_grid_thw if image_grid_thw is not None else video_grid_thw
+
+        if pixel_values is None:
+            # Reset position state for text-only generation
+            self.language_model._position_ids = None
+            self.language_model._rope_deltas = None
+            return InputEmbeddingsFeatures(
+                inputs_embeds=self.language_model.model.embed_tokens(input_ids)
+            )
+
+        dtype = self.vision_tower.patch_embed.proj.weight.dtype
+        pixel_values = pixel_values.astype(dtype)
+
+        # Get the input embeddings from the language model
+        inputs_embeds = self.language_model.model.embed_tokens(input_ids)
+
+        # Get the ouptut hidden states from the vision model
+        hidden_states, deepstack_image_embeds = self.vision_tower(
+            pixel_values, grid_thw
+        )
+
+        visual_pos_masks = None
+        deepstack_visual_embeds = None
+
+        # Insert special image tokens in the input_ids
+        inputs_embeds, image_mask = self.merge_input_ids_with_image_features(
+            hidden_states,
+            inputs_embeds,
+            input_ids,
+            self.config.image_token_index,
+            self.config.video_token_index,
+        )
+
+        image_mask = image_mask[..., 0]
+        visual_pos_masks = image_mask
+        deepstack_visual_embeds = mx.eval(deepstack_image_embeds)
+
+        # Pre-calculate position_ids for chunked prefill
+        if image_grid_thw is not None or video_grid_thw is not None:
+            position_ids, rope_deltas = self.language_model.get_rope_index(
+                input_ids, image_grid_thw, video_grid_thw, mask
+            )
+            self.language_model._position_ids = position_ids
+            self.language_model._rope_deltas = rope_deltas
+
+        return InputEmbeddingsFeatures(
+            inputs_embeds=inputs_embeds,
+            visual_pos_masks=visual_pos_masks,
+            deepstack_visual_embeds=deepstack_visual_embeds,
+        )
+
+    @staticmethod
+    def merge_input_ids_with_image_features(
+        image_features, inputs_embeds, input_ids, image_token_index, video_token_index
+    ):
+        special_image_mask = input_ids == image_token_index
+        special_video_mask = input_ids == video_token_index
+        special_image_mask = special_image_mask | special_video_mask
+        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, special_image_mask
+
+    @property
+    def layers(self):
+        return self.language_model.model.layers
+
+    def __call__(
+        self,
+        input_ids: mx.array,
+        pixel_values: Optional[mx.array] = None,
+        mask: Optional[mx.array] = None,
+        cache=None,
+        **kwargs,
+    ):
+
+        input_embeddings_features = self.get_input_embeddings(
+            input_ids, pixel_values, **kwargs
+        )
+
+        kwargs.update(
+            {
+                "pixel_values": pixel_values,
+                **input_embeddings_features.to_dict(),
+            }
+        )
+
+        logits = self.language_model(input_ids, mask=mask, cache=cache, **kwargs)
+        return logits
+
+    def sanitize(self, weights):
+        sanitized_weights = {}
+        for key, value in weights.items():
+            if "model" in key:
+                if "model.language_model" in key:
+                    key = key.replace("model.language_model", "language_model.model")
+
+                elif "model.visual" in key:
+                    key = key.replace("model.visual", "vision_tower")
+            elif "lm_head" in key:
+                key = key.replace("lm_head", "language_model.lm_head")
+
+            sanitized_weights[key] = value
+
+        return sanitized_weights

mlx_vlm/models/qwen3_vl/vision.py

@@ -0,0 +1,441 @@
+from itertools import accumulate
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .config import VisionConfig
+
+
+def check_array_shape(arr):
+    shape = arr.shape
+
+    # Check if the shape has 4 or 5 dimensions
+    if len(shape) not in [4, 5]:
+        return False
+
+    B, out_channels, kH, KW, t = shape
+
+    if t == 3:
+        return True
+
+    # 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)
+    cos = mx.tile(cos, (1, 1, 2))
+    cos = mx.expand_dims(cos, axis=0)
+
+    sin = mx.expand_dims(sin, axis=1)
+    sin = mx.tile(sin, (1, 1, 2))
+    sin = mx.expand_dims(sin, axis=0)
+
+    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, dtype=inv_freq.dtype)
+        freqs = mx.outer(seq, inv_freq)
+        return freqs
+
+
+class PatchEmbed(nn.Module):
+    def __init__(
+        self,
+        patch_size: int = 14,
+        temporal_patch_size: int = 2,
+        in_channels: int = 3,
+        hidden_size: int = 1152,
+    ) -> None:
+        super().__init__()
+        self.patch_size = patch_size
+        self.temporal_patch_size = temporal_patch_size
+        self.in_channels = in_channels
+        self.hidden_size = hidden_size
+
+        kernel_size = [temporal_patch_size, patch_size, patch_size]
+        self.proj = nn.Conv3d(
+            in_channels,
+            hidden_size,
+            kernel_size=kernel_size,
+            stride=kernel_size,
+            bias=True,
+        )
+
+    def __call__(self, hidden_states: mx.array) -> mx.array:
+        hidden_states = hidden_states.reshape(
+            -1,
+            self.in_channels,
+            self.temporal_patch_size,
+            self.patch_size,
+            self.patch_size,
+        ).moveaxis(1, 4)
+
+        hidden_states = self.proj(hidden_states)
+        hidden_states = hidden_states.reshape(-1, self.hidden_size)
+        return hidden_states
+
+
+class PatchMerger(nn.Module):
+    def __init__(self, config: VisionConfig, use_postshuffle_norm=False) -> None:
+        super().__init__()
+        self.hidden_size = config.hidden_size * (config.spatial_merge_size**2)
+        self.use_postshuffle_norm = use_postshuffle_norm
+        self.norm = nn.LayerNorm(
+            self.hidden_size if use_postshuffle_norm else config.hidden_size, eps=1e-6
+        )
+        self.linear_fc1 = nn.Linear(self.hidden_size, self.hidden_size)
+        self.act_fn = nn.GELU()
+        self.linear_fc2 = nn.Linear(self.hidden_size, config.out_hidden_size)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        x = self.norm(
+            x.reshape(-1, self.hidden_size) if self.use_postshuffle_norm else x
+        ).reshape(-1, self.hidden_size)
+        x = self.linear_fc2(self.act_fn(self.linear_fc1(x)))
+        return x
+
+
+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.qkv = nn.Linear(dim, dim * 3, bias=True)
+        self.proj = nn.Linear(dim, dim)
+
+    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.qkv(x).reshape(seq_length, 3, self.num_heads, -1).transpose(1, 0, 2, 3)
+        )
+        q, k, v = mx.split(qkv, 3)
+
+        q = apply_rotary_pos_emb_vision(mx.expand_dims(q, 0), rotary_pos_emb)[0]
+        k = apply_rotary_pos_emb_vision(mx.expand_dims(k, 0), rotary_pos_emb)[0]
+
+        q = q.transpose(0, 2, 1, 3)
+        k = k.transpose(0, 2, 1, 3)
+        v = v.transpose(0, 2, 1, 3)
+
+        splits = [
+            mx.split(tensor, cu_seqlens[1:-1].tolist(), axis=2) for tensor in (q, k, v)
+        ]
+
+        attn_outputs = []
+        for q, k, v in zip(*splits):
+            output = mx.fast.scaled_dot_product_attention(
+                q, k, v, scale=self.scale, mask=None
+            )
+            attn_outputs.append(output)
+
+        output = mx.concatenate(attn_outputs, axis=2)
+        output = output.transpose(0, 2, 1, 3).reshape(seq_length, -1)
+        return self.proj(output)
+
+
+class MLP(nn.Module):
+    def __init__(self, dim, hidden_dim):
+        super().__init__()
+        self.linear_fc1 = nn.Linear(dim, hidden_dim, bias=True)
+        self.linear_fc2 = nn.Linear(hidden_dim, dim, bias=True)
+        self.act_fn = nn.GELU(approx="tanh")
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return self.linear_fc2(self.act_fn(self.linear_fc1(x)))
+
+
+class Qwen3VLMoEVisionBlock(nn.Module):
+    def __init__(self, config: VisionConfig) -> None:
+        super().__init__()
+        self.norm1 = nn.LayerNorm(config.hidden_size, eps=1e-6)
+        self.norm2 = nn.LayerNorm(config.hidden_size, eps=1e-6)
+
+        self.attn = Attention(dim=config.hidden_size, num_heads=config.num_heads)
+        self.mlp = MLP(dim=config.hidden_size, hidden_dim=config.intermediate_size)
+
+    def __call__(self, hidden_states, cu_seqlens, rotary_pos_emb) -> mx.array:
+        hidden_states = hidden_states + self.attn(
+            self.norm1(hidden_states),
+            cu_seqlens=cu_seqlens,
+            rotary_pos_emb=rotary_pos_emb,
+        )
+        hidden_states = hidden_states + self.mlp(self.norm2(hidden_states))
+        return hidden_states
+
+
+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 != "qwen3_vl":
+            raise ValueError(f"Unsupported model type: {self.model_type}")
+
+        self.spatial_merge_size = config.spatial_merge_size
+
+        self.patch_embed = PatchEmbed(
+            patch_size=config.patch_size,
+            temporal_patch_size=config.temporal_patch_size,
+            in_channels=config.in_channels,
+            hidden_size=config.hidden_size,
+        )
+
+        head_dim = config.hidden_size // config.num_heads
+        self.rotary_pos_emb = VisionRotaryEmbedding(head_dim // 2)
+
+        self.pos_embed = nn.Embedding(
+            config.num_position_embeddings, config.hidden_size
+        )
+        self.num_grid_per_side = int(config.num_position_embeddings**0.5)
+
+        self.blocks = [Qwen3VLMoEVisionBlock(config) for _ in range(config.depth)]
+        self.merger = PatchMerger(config=config, use_postshuffle_norm=False)
+
+        self.deepstack_visual_indexes = config.deepstack_visual_indexes
+        self.deepstack_merger_list = [
+            PatchMerger(
+                config=config,
+                use_postshuffle_norm=True,
+            )
+            for _ in range(len(config.deepstack_visual_indexes))
+        ]
+
+    def rot_pos_emb(self, grid_thw: mx.array) -> mx.array:
+        merge_size = self.spatial_merge_size
+
+        # Get max grid size for frequency table
+        max_hw = int(mx.max(grid_thw[:, 1:]).item())
+        freq_table = self.rotary_pos_emb(max_hw)  # Shape: (max_hw, dim // 2)
+
+        pos_ids = []
+
+        for num_frames, height, width in grid_thw.tolist():
+            num_frames, height, width = int(num_frames), int(height), int(width)
+            merged_h, merged_w = height // merge_size, width // merge_size
+
+            # Create block indices
+            block_rows = mx.arange(merged_h)
+            block_cols = mx.arange(merged_w)
+
+            # Create intra-block indices
+            intra_row = mx.arange(merge_size)
+            intra_col = mx.arange(merge_size)
+
+            # Compute full-resolution positions
+            row_idx = (
+                block_rows[:, None, None, None] * merge_size
+                + intra_row[None, None, :, None]
+            )
+            col_idx = (
+                block_cols[None, :, None, None] * merge_size
+                + intra_col[None, None, None, :]
+            )
+
+            # Broadcast and flatten
+            row_idx = mx.broadcast_to(
+                row_idx, (merged_h, merged_w, merge_size, merge_size)
+            ).reshape(-1)
+            col_idx = mx.broadcast_to(
+                col_idx, (merged_h, merged_w, merge_size, merge_size)
+            ).reshape(-1)
+
+            # Stack into coordinate pairs
+            coords = mx.stack([row_idx, col_idx], axis=-1)
+
+            # Repeat for temporal dimension
+            if num_frames > 1:
+                coords = mx.tile(coords, (num_frames, 1))
+
+            pos_ids.append(coords)
+
+        # Concatenate all position IDs - shape: (total_tokens, 2)
+        pos_ids = mx.concatenate(pos_ids, axis=0)
+
+        # Lookup embeddings: freq_table[h_pos] and freq_table[w_pos]
+        # pos_ids[:, 0] = height positions, pos_ids[:, 1] = width positions
+        h_embeddings = freq_table[pos_ids[:, 0]]  # (total_tokens, dim // 2)
+        w_embeddings = freq_table[pos_ids[:, 1]]  # (total_tokens, dim // 2)
+
+        # Concatenate height and width embeddings
+        embeddings = mx.concatenate([h_embeddings, w_embeddings], axis=-1)
+
+        return embeddings
+
+    def fast_pos_embed_interpolate(self, grid_thw):
+        grid_thw_list = grid_thw.tolist()
+        idx_list = [[] for _ in range(4)]
+        weight_list = [[] for _ in range(4)]
+
+        for t, h, w in grid_thw_list:
+            h = int(h)
+            w = int(w)
+            t = int(t)
+
+            h_idxs = mx.linspace(0, self.num_grid_per_side - 1, h)
+            w_idxs = mx.linspace(0, self.num_grid_per_side - 1, w)
+
+            h_idxs_floor = h_idxs.astype(mx.int32)
+            w_idxs_floor = w_idxs.astype(mx.int32)
+            h_idxs_ceil = mx.minimum(h_idxs_floor + 1, self.num_grid_per_side - 1)
+            w_idxs_ceil = mx.minimum(w_idxs_floor + 1, self.num_grid_per_side - 1)
+
+            dh = h_idxs - h_idxs_floor.astype(mx.float32)
+            dw = w_idxs - w_idxs_floor.astype(mx.float32)
+
+            base_h = h_idxs_floor * self.num_grid_per_side
+            base_h_ceil = h_idxs_ceil * self.num_grid_per_side
+
+            indices = [
+                (base_h[:, None] + w_idxs_floor[None, :]).flatten(),
+                (base_h[:, None] + w_idxs_ceil[None, :]).flatten(),
+                (base_h_ceil[:, None] + w_idxs_floor[None, :]).flatten(),
+                (base_h_ceil[:, None] + w_idxs_ceil[None, :]).flatten(),
+            ]
+
+            weights = [
+                ((1 - dh)[:, None] * (1 - dw)[None, :]).flatten(),
+                ((1 - dh)[:, None] * dw[None, :]).flatten(),
+                (dh[:, None] * (1 - dw)[None, :]).flatten(),
+                (dh[:, None] * dw[None, :]).flatten(),
+            ]
+
+            for i in range(4):
+                idx_list[i].extend(indices[i].tolist())
+                weight_list[i].extend(weights[i].tolist())
+
+        idx_tensor = mx.array(idx_list, dtype=mx.int32)
+        weight_tensor = mx.array(weight_list, dtype=self.pos_embed.weight.dtype)
+
+        pos_embeds = self.pos_embed(idx_tensor) * weight_tensor[:, :, None]
+        patch_pos_embeds = pos_embeds[0] + pos_embeds[1] + pos_embeds[2] + pos_embeds[3]
+
+        split_sizes = [int(h * w) for t, h, w in grid_thw_list]
+        if len(split_sizes) > 1:
+            split_indices = list(accumulate(split_sizes[:-1]))
+            patch_pos_embeds_split = mx.split(patch_pos_embeds, split_indices, axis=0)
+        else:
+            patch_pos_embeds_split = [patch_pos_embeds]
+
+        patch_pos_embeds_permute = []
+        merge_size = self.config.spatial_merge_size
+
+        for pos_embed, (t, h, w) in zip(patch_pos_embeds_split, grid_thw_list):
+            t, h, w = int(t), int(h), int(w)
+            feature_dim = pos_embed.shape[-1]
+            pos_embed = mx.tile(pos_embed, (t, 1))
+            pos_embed = pos_embed.reshape(t, h, w, feature_dim)
+            pos_embed = (
+                pos_embed.reshape(
+                    t,
+                    h // merge_size,
+                    merge_size,
+                    w // merge_size,
+                    merge_size,
+                    feature_dim,
+                )
+                .transpose(0, 1, 3, 2, 4, 5)
+                .reshape(-1, feature_dim)
+            )
+            patch_pos_embeds_permute.append(pos_embed)
+
+        patch_pos_embeds = mx.concatenate(patch_pos_embeds_permute)
+        return patch_pos_embeds
+
+    def __call__(
+        self,
+        hidden_states: mx.array,
+        grid_thw: mx.array,
+        **kwargs,
+    ) -> mx.array:
+
+        hidden_states = self.patch_embed(hidden_states)
+        pos_embeds = self.fast_pos_embed_interpolate(grid_thw)
+        hidden_states = hidden_states + pos_embeds
+        rotary_pos_emb = self.rot_pos_emb(grid_thw)
+
+        seq_len = hidden_states.shape[0]
+        hidden_states = hidden_states.reshape(seq_len, -1)
+        rotary_pos_emb = rotary_pos_emb.reshape(seq_len, -1)
+
+        # Assuming grid_thw has shape (batch_size, 3)
+        batch_size = grid_thw.shape[0]
+
+        # Calculate cu_seqlens for each item in the batch
+        cu_seqlens = []
+        for i in range(batch_size):
+            seq_len = grid_thw[i, 1] * grid_thw[i, 2]
+            cu_seqlens.append(mx.repeat(seq_len, grid_thw[i, 0]))
+
+        # Concatenate the cu_seqlens for all items in the batch
+        cu_seqlens = mx.concatenate(cu_seqlens)
+
+        cu_seqlens = mx.cumsum(cu_seqlens.astype(mx.int32), axis=0)
+        cu_seqlens = mx.pad(cu_seqlens, (1, 0), mode="constant", constant_values=0)
+
+        deepstack_feature_lists = []
+        for layer_num, blk in enumerate(self.blocks):
+            hidden_states = blk(
+                hidden_states,
+                cu_seqlens=cu_seqlens,
+                rotary_pos_emb=rotary_pos_emb,
+            )
+            if layer_num in self.deepstack_visual_indexes:
+                deepstack_feature = self.deepstack_merger_list[
+                    self.deepstack_visual_indexes.index(layer_num)
+                ](hidden_states)
+                deepstack_feature_lists.append(deepstack_feature)
+
+        hidden_states = self.merger(hidden_states)
+
+        return hidden_states, deepstack_feature_lists
+
+    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, 4, 1)
+            else:
+                sanitized_weights[k] = v
+
+        return sanitized_weights

mlx_vlm/models/qwen3_vl_moe/__init__.py

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