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

mlx_vlm/models/qwen3_omni_moe/vision.py

@@ -0,0 +1,419 @@
+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
+
+    if len(shape) not in [4, 5]:
+        return False
+
+    B, out_channels, kH, KW, t = shape
+
+    if t == 3:
+        return True
+
+    if (out_channels >= kH) and (out_channels >= KW) and (kH == KW):
+        return True
+    else:
+        return False
+
+
+def rotate_half(x):
+    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()
+
+    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 not in ["qwen3_vl_moe", "qwen3_omni_moe_vision_encoder"]:
+            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
+
+        max_hw = int(mx.max(grid_thw[:, 1:]).item())
+        freq_table = self.rotary_pos_emb(max_hw)
+
+        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
+
+            block_rows = mx.arange(merged_h)
+            block_cols = mx.arange(merged_w)
+
+            intra_row = mx.arange(merge_size)
+            intra_col = mx.arange(merge_size)
+
+            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, :]
+            )
+
+            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)
+
+            coords = mx.stack([row_idx, col_idx], axis=-1)
+
+            if num_frames > 1:
+                coords = mx.tile(coords, (num_frames, 1))
+
+            pos_ids.append(coords)
+
+        pos_ids = mx.concatenate(pos_ids, axis=0)
+
+        h_embeddings = freq_table[pos_ids[:, 0]]
+        w_embeddings = freq_table[pos_ids[:, 1]]
+
+        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)
+
+        batch_size = grid_thw.shape[0]
+
+        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]))
+
+        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:
+                continue
+            elif "patch_embed.proj.weight" in k:
+                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/__init__.py

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

mlx_vlm/models/qwen3_vl/config.py

@@ -0,0 +1,103 @@
+import inspect
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional, Union
+
+from ..base import BaseModelConfig
+
+
+@dataclass
+class VisionConfig(BaseModelConfig):
+    model_type: str = "qwen3_vl"
+    depth: int = 32
+    hidden_size: int = 1280
+    intermediate_size: int = 3420
+    out_hidden_size: int = 1536
+    num_heads: int = 16
+    image_size: int = 384
+    patch_size: int = 14
+    vocab_size: int = 32000
+    mlp_ratio: float = 4.0
+    in_channels: int = 3
+    layer_norm_eps: float = 1e-6
+    spatial_patch_size: int = 14
+    spatial_merge_size: int = 2
+    tokens_per_second: int = 2
+    temporal_patch_size: int = 2
+    num_position_embeddings: int = 2304
+    window_size: int = 112
+    fullatt_block_indexes: list[int] = field(default_factory=lambda: [7, 15, 23, 31])
+    deepstack_visual_indexes: list[int] = field(default_factory=list)
+
+
+@dataclass
+class TextConfig(BaseModelConfig):
+    model_type: str
+    num_hidden_layers: int
+    hidden_size: int
+    intermediate_size: int
+    num_attention_heads: int
+    rms_norm_eps: float
+    vocab_size: int
+    num_key_value_heads: Optional[int]
+    head_dim: int
+    rope_theta: float
+    max_position_embeddings: int
+    norm_topk_prob: bool = True
+    rope_scaling: Optional[Dict[str, Union[float, str, bool, List[int]]]] = field(
+        default_factory=lambda: {"type": "default", "mrope_section": [24, 20, 20]}
+    )
+    tie_word_embeddings: bool = False
+    attention_bias: bool = False
+    hidden_act: str = "silu"
+
+    def __post_init__(self):
+        if self.num_key_value_heads is None:
+            self.num_key_value_heads = self.num_attention_heads
+
+        if self.rope_scaling:
+            # Normalize rope_scaling keys (accept both 'rope_type' and 'type')
+            if "type" not in self.rope_scaling and "rope_type" in self.rope_scaling:
+                self.rope_scaling["type"] = self.rope_scaling.pop("rope_type")
+
+            required_keys = {"mrope_section", "type"}
+            if not all(key in self.rope_scaling for key in required_keys):
+                raise ValueError(f"rope_scaling must contain keys {required_keys}")
+
+            if not self.rope_scaling["type"] in ["mrope", "default"]:
+                raise ValueError(f"rope_scaling type must be 'mrope' or 'default'")
+
+
+@dataclass
+class ModelConfig(BaseModelConfig):
+    text_config: TextConfig
+    vision_config: VisionConfig
+    model_type: str
+    ignore_index: int = -100
+    image_token_id: int = 151655
+    video_token_id: int = 151656
+    image_token_index: Optional[int] = None
+    video_token_index: Optional[int] = None
+    vision_start_token_id: int = 151652
+    vision_end_token_id: int = 151653
+    vision_token_id: int = 151654
+    vision_feature_select_strategy: str = "default"
+    vision_feature_layer: int = -2
+    vocab_size: int = 32000
+    eos_token_id: Optional[List[int]] = None
+
+    def __post_init__(self):
+        if self.image_token_index is None:
+            self.image_token_index = self.image_token_id
+        if self.video_token_index is None:
+            self.video_token_index = self.video_token_id
+
+    @classmethod
+    def from_dict(cls, params):
+
+        return cls(
+            **{
+                k: v
+                for k, v in params.items()
+                if k in inspect.signature(cls).parameters
+            }
+        )