nexaai 1.0.19rc7__cp310-cp310-macosx_14_0_universal2.whl → 1.0.19rc8__cp310-cp310-macosx_14_0_universal2.whl

nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3vl_moe/qwen3vl_moe.py

@@ -1,1309 +0,0 @@
-# Copyright © 2023-2024 Apple Inc.
-
-from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Tuple, Union
-
-import mlx.core as mx
-import mlx.nn as nn
-import math
-import numpy as np
-
-import os
-import sys
-
-curr_dir = os.path.dirname(os.path.abspath(__file__))
-llm_common_dir = os.path.join(curr_dir, "llm_common")
-sys.path.append(llm_common_dir)
-
-# Try relative imports first, fallback to sys.path approach for Nuitka compatibility
-try:
-    from .llm_common.base import (
-        BaseModelArgs,
-        create_attention_mask,
-        scaled_dot_product_attention,
-    )
-    from .llm_common.rope_utils import initialize_rope
-except ImportError:
-    # Fallback for Nuitka compiled environment
-    from llm_common.base import (
-        BaseModelArgs,
-        create_attention_mask,
-        scaled_dot_product_attention,
-    )
-    from llm_common.rope_utils import initialize_rope
-from switch_layers import SwitchGLU
-
-
-@dataclass
-class VisionConfig:
-    hidden_size: int = 1152
-    intermediate_size: int = 4304
-    num_heads: int = 16
-    num_hidden_layers: int = 27
-    patch_size: int = 16
-    temporal_patch_size: int = 2
-    in_channels: int = 3
-    hidden_act: str = "gelu_pytorch_tanh"
-    spatial_merge_size: int = 2
-    out_hidden_size: int = 2048
-    num_position_embeddings: int = 2304
-    deepstack_visual_indexes: List[int] = None
-
-    def __post_init__(self):
-        if self.deepstack_visual_indexes is None:
-            self.deepstack_visual_indexes = [8, 16, 24]
-
-
-@dataclass
-class TextConfig(BaseModelArgs):
-    model_type: str = "qwen3_vl_moe_text"
-    hidden_size: int = 2048
-    num_hidden_layers: int = 48
-    intermediate_size: int = 6144
-    num_attention_heads: int = 32
-    num_key_value_heads: int = 4
-    rms_norm_eps: float = 1e-6
-    vocab_size: int = 152064
-    max_position_embeddings: int = 128000
-    rope_theta: float = 1000000.0
-    head_dim: int = 128
-    tie_word_embeddings: bool = False
-    attention_bias: bool = False
-    attention_dropout: float = 0.0
-    rope_scaling: Optional[Dict[str, Union[float, str]]] = None
-    # MoE specific parameters
-    num_experts: int = 128
-    num_experts_per_tok: int = 8
-    moe_intermediate_size: int = 768
-    shared_expert_intermediate_size: int = 0
-    norm_topk_prob: bool = True
-    decoder_sparse_step: int = 1
-    max_window_layers: int = 48
-    sliding_window: int = 32768
-    mlp_only_layers: List[int] = None
-    use_qk_norm: bool = True
-    layer_types: List[str] = None
-
-    def __post_init__(self):
-        if self.rope_scaling is None:
-            self.rope_scaling = {
-                "mrope_interleaved": True,
-                "mrope_section": [24, 20, 20],
-                "rope_type": "default"
-            }
-        if self.mlp_only_layers is None:
-            self.mlp_only_layers = []
-        if self.layer_types is None:
-            # This would need to be populated based on the actual model architecture
-            self.layer_types = []
-
-
-@dataclass
-class ModelArgs(BaseModelArgs):
-    vision_config: VisionConfig = None
-    text_config: TextConfig = None
-    image_token_id: int = 151655
-    vision_start_token_id: int = 151652
-    vision_end_token_id: int = 151653
-
-    def __post_init__(self):
-        if self.vision_config is None:
-            self.vision_config = VisionConfig()
-        if self.text_config is None:
-            self.text_config = TextConfig()
-
-
-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(q, k, cos, sin):
-    cos = mx.expand_dims(cos, axis=-2)
-    sin = mx.expand_dims(sin, axis=-2)
-    q_embed = (q * cos) + (rotate_half(q) * sin)
-    k_embed = (k * cos) + (rotate_half(k) * sin)
-    return q_embed, k_embed
-
-
-def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
-    cos = mx.expand_dims(cos, axis=unsqueeze_dim)
-    sin = mx.expand_dims(sin, axis=unsqueeze_dim)
-    q_embed = (q * cos) + (rotate_half(q) * sin)
-    k_embed = (k * cos) + (rotate_half(k) * sin)
-    return q_embed, k_embed
-
-
-class VisionMLP(nn.Module):
-    def __init__(self, config: VisionConfig):
-        super().__init__()
-        self.hidden_size = config.hidden_size
-        self.intermediate_size = config.intermediate_size
-        self.linear_fc1 = nn.Linear(self.hidden_size, self.intermediate_size, bias=True)
-        self.linear_fc2 = nn.Linear(self.intermediate_size, self.hidden_size, bias=True)
-
-    def __call__(self, hidden_state):
-        return self.linear_fc2(nn.gelu(self.linear_fc1(hidden_state)))
-
-
-class VisionPatchEmbed(nn.Module):
-    def __init__(self, config: VisionConfig):
-        super().__init__()
-        self.patch_size = config.patch_size
-        self.temporal_patch_size = config.temporal_patch_size
-        self.in_channels = config.in_channels
-        self.embed_dim = config.hidden_size
-
-        kernel_size = [self.temporal_patch_size, self.patch_size, self.patch_size]
-        self.proj = nn.Conv3d(
-            self.in_channels, 
-            self.embed_dim, 
-            kernel_size=kernel_size, 
-            stride=kernel_size, 
-            bias=True
-        )
-
-    def __call__(self, hidden_states: mx.array) -> mx.array:
-        target_dtype = self.proj.weight.dtype
-        
-        # Reshape to 5D: [batch, channels, temporal, height, width] (PyTorch format)
-        # This matches the PyTorch ground truth exactly
-        hidden_states = hidden_states.reshape(
-            -1, self.in_channels, self.temporal_patch_size, self.patch_size, self.patch_size
-        )
-        
-        # Convert to MLX format: [batch, temporal, height, width, channels]
-        hidden_states = hidden_states.transpose(0, 2, 3, 4, 1)
-        
-        # Apply conv3d with target dtype and reshape to match PyTorch output
-        hidden_states = self.proj(hidden_states.astype(target_dtype)).reshape(-1, self.embed_dim)
-        
-        return hidden_states
-
-
-class VisionRotaryEmbedding(nn.Module):
-    def __init__(self, dim: int, theta: float = 10000.0):
-        super().__init__()
-        # Don't store inv_freq as a parameter since it causes loading issues
-        self.dim = dim
-        self.theta = theta
-
-    def __call__(self, seqlen: int) -> mx.array:
-        # Compute inv_freq on the fly
-        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 VisionPatchMerger(nn.Module):
-    def __init__(self, config: VisionConfig, use_postshuffle_norm=False):
-        super().__init__()
-        self.hidden_size = config.hidden_size * (config.spatial_merge_size ** 2)
-        self.use_postshuffle_norm = use_postshuffle_norm
-        
-        norm_size = self.hidden_size if use_postshuffle_norm else config.hidden_size
-        self.ln_q = nn.LayerNorm(norm_size, eps=1e-6)
-        self.linear_fc1 = nn.Linear(self.hidden_size, self.hidden_size)
-        self.linear_fc2 = nn.Linear(self.hidden_size, config.out_hidden_size)
-
-    def __call__(self, x: mx.array) -> mx.array:
-        if self.use_postshuffle_norm:
-            x = self.ln_q(x.reshape(-1, self.hidden_size)).reshape(-1, self.hidden_size)
-        else:
-            x = self.ln_q(x).reshape(-1, self.hidden_size)
-        
-        x = self.linear_fc2(nn.gelu(self.linear_fc1(x)))
-        return x
-
-
-class VisionAttention(nn.Module):
-    def __init__(self, config: VisionConfig):
-        super().__init__()
-        self.dim = config.hidden_size
-        self.num_heads = config.num_heads
-        self.head_dim = self.dim // self.num_heads
-        self.scaling = self.head_dim ** -0.5
-        
-        self.qkv = nn.Linear(self.dim, self.dim * 3, bias=True)
-        self.proj = nn.Linear(self.dim, self.dim)
-
-    def __call__(
-        self,
-        hidden_states: mx.array,
-        cu_seqlens: mx.array,
-        rotary_pos_emb: Optional[mx.array] = None,
-        position_embeddings: Optional[Tuple[mx.array, mx.array]] = None,
-        **kwargs,
-    ) -> mx.array:
-        seq_length = hidden_states.shape[0]
-        qkv = self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1)
-        qkv = qkv.transpose(1, 0, 2, 3)
-        query_states, key_states, value_states = qkv[0], qkv[1], qkv[2]
-        
-        cos, sin = position_embeddings
-        query_states, key_states = apply_rotary_pos_emb_vision(
-            query_states, key_states, cos, sin
-        )
-
-        query_states = query_states.transpose(1, 0, 2)
-        key_states = key_states.transpose(1, 0, 2)
-        value_states = value_states.transpose(1, 0, 2)
-        
-        query_states = mx.expand_dims(query_states, axis=0)
-        key_states = mx.expand_dims(key_states, axis=0)
-        value_states = mx.expand_dims(value_states, axis=0)
-        
-        lengths = cu_seqlens[1:] - cu_seqlens[:-1]
-        
-        split_indices = []
-        cumsum = 0
-        for length in lengths[:-1]:
-            cumsum += int(length)
-            split_indices.append(cumsum)
-        
-        if split_indices:
-            q_splits = mx.split(query_states, split_indices, axis=1)
-            k_splits = mx.split(key_states, split_indices, axis=1)
-            v_splits = mx.split(value_states, split_indices, axis=1)
-        else:
-            q_splits = [query_states]
-            k_splits = [key_states] 
-            v_splits = [value_states]
-        
-        attn_outputs = []
-        for q, k, v in zip(q_splits, k_splits, v_splits):
-            attn_out = scaled_dot_product_attention(
-                q, k, v, 
-                scale=self.scaling, mask=None, cache=None
-            )
-            attn_outputs.append(attn_out)
-        
-        attn_output = mx.concatenate(attn_outputs, axis=1)
-        
-        attn_output = attn_output[0].transpose(1, 0, 2)
-        attn_output = attn_output.reshape(seq_length, -1)
-        attn_output = self.proj(attn_output)
-        
-        return attn_output
-
-
-class VisionBlock(nn.Module):
-    def __init__(self, config: VisionConfig):
-        super().__init__()
-        self.norm1 = nn.LayerNorm(config.hidden_size, eps=1e-6)
-        self.norm2 = nn.LayerNorm(config.hidden_size, eps=1e-6)
-        self.attn = VisionAttention(config)
-        self.mlp = VisionMLP(config)
-
-    def __call__(
-        self,
-        hidden_states: mx.array,
-        cu_seqlens: mx.array,
-        position_embeddings: Tuple[mx.array, mx.array],
-    ) -> mx.array:
-        hidden_states = hidden_states + self.attn(
-            self.norm1(hidden_states),
-            cu_seqlens=cu_seqlens,
-            position_embeddings=position_embeddings,
-        )
-        hidden_states = hidden_states + self.mlp(self.norm2(hidden_states))
-        return hidden_states
-
-
-class VisionModel(nn.Module):
-    def __init__(self, config: VisionConfig):
-        super().__init__()
-        self.config = config
-        self.spatial_merge_size = config.spatial_merge_size
-        self.patch_size = config.patch_size
-
-        self.patch_embed = VisionPatchEmbed(config)
-        self.pos_embed = nn.Embedding(config.num_position_embeddings, config.hidden_size)
-        self.num_grid_per_side = int(config.num_position_embeddings ** 0.5)
-
-        head_dim = config.hidden_size // config.num_heads
-        self.rotary_pos_emb = VisionRotaryEmbedding(head_dim // 2)
-
-        self.blocks = [VisionBlock(config) for _ in range(config.num_hidden_layers)]
-        self.merger = VisionPatchMerger(config, use_postshuffle_norm=False)
-
-        self.deepstack_visual_indexes = config.deepstack_visual_indexes
-        self.deepstack_merger_list = [
-            VisionPatchMerger(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(grid_thw[:, 1:].max().item())
-        freq_table = self.rotary_pos_emb(max_hw)  # (max_hw, dim // 2)
-
-        pos_ids_parts = []
-
-        for i in range(grid_thw.shape[0]):
-            num_frames = int(grid_thw[i, 0].item())
-            height = int(grid_thw[i, 1].item())
-            width = int(grid_thw[i, 2].item())
-            
-            merged_h, merged_w = height // merge_size, width // merge_size
-
-            block_rows = mx.arange(merged_h)  # block row indices
-            block_cols = mx.arange(merged_w)  # block col indices
-            intra_row = mx.arange(merge_size)  # intra-block row offsets
-            intra_col = mx.arange(merge_size)  # intra-block col offsets
-
-            # Compute full-resolution positions using broadcasting
-            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_parts.append(coords)
-
-        # Concatenate all coordinate parts
-        pos_ids = mx.concatenate(pos_ids_parts, axis=0)
-        
-        embeddings = freq_table[pos_ids]  # lookup rotary embeddings
-        embeddings = embeddings.reshape(embeddings.shape[0], -1)
-        return embeddings
-
-    def fast_pos_embed_interpolate(self, grid_thw: mx.array):
-        patch_pos_embeds = []
-        
-        for i in range(grid_thw.shape[0]):
-            t = int(grid_thw[i, 0].item())
-            h = int(grid_thw[i, 1].item())
-            w = int(grid_thw[i, 2].item())
-            
-            # Simple position embedding interpolation
-            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 = mx.floor(h_idxs).astype(mx.int32)
-            w_idxs_floor = mx.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
-
-            # Compute bilinear interpolation indices and weights
-            indices_tl = (base_h[:, None] + w_idxs_floor[None, :]).reshape(-1)
-            indices_tr = (base_h[:, None] + w_idxs_ceil[None, :]).reshape(-1)
-            indices_bl = (base_h_ceil[:, None] + w_idxs_floor[None, :]).reshape(-1)
-            indices_br = (base_h_ceil[:, None] + w_idxs_ceil[None, :]).reshape(-1)
-
-            weights_tl = ((1 - dh)[:, None] * (1 - dw)[None, :]).reshape(-1)
-            weights_tr = ((1 - dh)[:, None] * dw[None, :]).reshape(-1)
-            weights_bl = (dh[:, None] * (1 - dw)[None, :]).reshape(-1)
-            weights_br = (dh[:, None] * dw[None, :]).reshape(-1)
-
-            # Get embeddings and interpolate
-            pos_embed_tl = self.pos_embed(indices_tl) * weights_tl[:, None]
-            pos_embed_tr = self.pos_embed(indices_tr) * weights_tr[:, None]
-            pos_embed_bl = self.pos_embed(indices_bl) * weights_bl[:, None]
-            pos_embed_br = self.pos_embed(indices_br) * weights_br[:, None]
-
-            pos_embed = pos_embed_tl + pos_embed_tr + pos_embed_bl + pos_embed_br
-
-            # Repeat for temporal dimension and apply spatial merging
-            pos_embed = mx.tile(pos_embed, (t, 1))
-            
-            # Apply spatial merging pattern
-            merge_size = self.config.spatial_merge_size
-            pos_embed = pos_embed.reshape(t, h // merge_size, merge_size, w // merge_size, merge_size, -1)
-            pos_embed = mx.transpose(pos_embed, (0, 1, 3, 2, 4, 5))
-            pos_embed = pos_embed.reshape(-1, pos_embed.shape[-1])
-            
-            patch_pos_embeds.append(pos_embed)
-        
-        return mx.concatenate(patch_pos_embeds, axis=0)
-
-    def __call__(self, hidden_states: mx.array, grid_thw: mx.array) -> Tuple[mx.array, List[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]
-        
-        emb = mx.concatenate([rotary_pos_emb, rotary_pos_emb], axis=-1)
-        position_embeddings = (mx.cos(emb), mx.sin(emb))
-
-            # Create cumulative sequence lengths (following HuggingFace implementation)
-        # torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0])
-        seq_lens_per_image = grid_thw[:, 1] * grid_thw[:, 2]  # h * w for each image
-        seq_lens = []
-        for i, (seq_len, repeats) in enumerate(zip(seq_lens_per_image, grid_thw[:, 0])):
-            seq_lens.extend([seq_len] * int(repeats))
-        seq_lens = mx.array(seq_lens)
-        
-        # Then compute cumulative sum
-        cu_seqlens = mx.cumsum(seq_lens)
-        # Pad with 0 at the beginning
-        cu_seqlens = mx.concatenate([mx.array([0]), cu_seqlens])
-
-        deepstack_feature_lists = []
-        for layer_num, blk in enumerate(self.blocks):
-            hidden_states = blk(
-                hidden_states,
-                cu_seqlens=cu_seqlens,
-                position_embeddings=position_embeddings,
-            )
-            if layer_num in self.deepstack_visual_indexes:
-                idx = self.deepstack_visual_indexes.index(layer_num)
-                deepstack_feature = self.deepstack_merger_list[idx](hidden_states)
-                deepstack_feature_lists.append(deepstack_feature)
-
-        hidden_states = self.merger(hidden_states)
-        return hidden_states, deepstack_feature_lists
-
-
-class TextRotaryEmbedding(nn.Module):
-    def __init__(self, config: TextConfig):
-        super().__init__()
-        self.config = config
-        self.max_seq_len_cached = config.max_position_embeddings
-        self.original_max_seq_len = config.max_position_embeddings
-        
-        # MRoPE configuration
-        if hasattr(config, "rope_scaling") and config.rope_scaling is not None:
-            self.rope_type = config.rope_scaling.get("rope_type", "default")
-            self.mrope_section = config.rope_scaling.get("mrope_section", [24, 20, 20])
-        else:
-            self.rope_type = "default"
-            self.mrope_section = [24, 20, 20]
-        
-        # Store parameters for computing inv_freq on the fly
-        self.head_dim = config.head_dim
-        self.theta = config.rope_theta
-        
-        # Attention scaling (simplified - may need adjustment based on actual config)
-        self.attention_scaling = 1.0
-
-    def _get_inv_freq(self):
-        """Compute inverse frequencies on the fly"""
-        inv_freq = 1.0 / (self.theta ** (mx.arange(0, self.head_dim, 2).astype(mx.float32) / self.head_dim))
-        # Expand for 3 dimensions (T, H, W)
-        return mx.broadcast_to(inv_freq[None, :], (3, len(inv_freq)))
-
-    def apply_interleaved_mrope(self, freqs, mrope_section):
-        """Apply interleaved MRoPE to 3D rotary embeddings.
-        Reorganizes frequency layout from chunked [TTT...HHH...WWW] to
-        interleaved [THTHWHTHW...TT], preserving frequency continuity.
-        args:
-            x: (3, bs, seq_len, head_dim // 2)
-            mrope_section: (3,)
-        returns:
-            x_t: (bs, seq_len, head_dim // 2)
-        """
-        freqs_t = freqs[0]  # just overwrite the first dimension T
-        for dim, offset in enumerate((1, 2), start=1):  # H, W
-            length = mrope_section[dim] * 3
-            idx = slice(offset, length, 3)
-            freqs_t[..., idx] = freqs[dim, ..., idx]
-        return freqs_t
-
-    def __call__(self, x: mx.array, position_ids: mx.array) -> mx.array:
-        """
-        Args:
-            x: Input tensor for dtype reference
-            position_ids: Position indices, shape (3, batch_size, seq_len) for MRoPE
-            
-        Returns:
-            cos, sin: Cosine and sine embeddings
-        """
-        # Handle 2D position_ids by expanding to 3D for MRoPE
-        if position_ids.ndim == 2:
-            position_ids = mx.broadcast_to(position_ids[None, ...], (3, position_ids.shape[0], position_ids.shape[1]))
-        
-        batch_size, seq_len = position_ids.shape[1], position_ids.shape[2]
-        
-        # Expand inverse frequencies: (3, 1, 1, dim//2) -> (3, batch_size, 1, dim//2)
-        inv_freq_expanded = mx.broadcast_to(
-            self._get_inv_freq()[:, None, None, :], 
-            (3, batch_size, 1, self._get_inv_freq().shape[-1])
-        )
-        
-        # Expand position ids: (3, batch_size, seq_len) -> (3, batch_size, seq_len, 1)
-        position_ids_expanded = position_ids[..., None].astype(mx.float32)
-        
-        # Compute frequencies: (3, batch_size, seq_len, dim//2)
-        freqs = inv_freq_expanded * position_ids_expanded
-        
-        # Apply interleaved MRoPE
-        freqs = self.apply_interleaved_mrope(freqs, self.mrope_section)
-        
-        # Create embeddings
-        emb = mx.concatenate([freqs, freqs], axis=-1)  # (batch_size, seq_len, head_dim)
-        cos = mx.cos(emb) * self.attention_scaling
-        sin = mx.sin(emb) * self.attention_scaling
-        
-        return cos.astype(x.dtype), sin.astype(x.dtype)
-
-
-class TextAttention(nn.Module):
-    def __init__(self, config: TextConfig, layer_idx: int):
-        super().__init__()
-        self.config = config
-        self.layer_idx = layer_idx
-        
-        dim = config.hidden_size
-        self.n_heads = config.num_attention_heads
-        self.n_kv_heads = config.num_key_value_heads
-        self.head_dim = config.head_dim
-        self.scale = self.head_dim ** -0.5
-
-        self.q_proj = nn.Linear(dim, self.n_heads * self.head_dim, bias=config.attention_bias)
-        self.k_proj = nn.Linear(dim, self.n_kv_heads * self.head_dim, bias=config.attention_bias)
-        self.v_proj = nn.Linear(dim, self.n_kv_heads * self.head_dim, bias=config.attention_bias)
-        self.o_proj = nn.Linear(self.n_heads * self.head_dim, dim, bias=config.attention_bias)
-
-        self.q_norm = nn.RMSNorm(self.head_dim, eps=config.rms_norm_eps)
-        self.k_norm = nn.RMSNorm(self.head_dim, eps=config.rms_norm_eps)
-        
-        # Initialize rope directly
-        self.rope = initialize_rope(
-            config.head_dim,
-            base=config.rope_theta,
-            traditional=False,
-            scaling_config=config.rope_scaling,
-            max_position_embeddings=config.max_position_embeddings,
-        )
-
-    def __call__(
-        self,
-        hidden_states: mx.array,
-        attention_mask: Optional[mx.array] = None,
-        cache: Optional[Any] = None,
-        cos: Optional[mx.array] = None,
-        sin: Optional[mx.array] = None,
-        rope_deltas: Optional[mx.array] = None,
-    ) -> Tuple[mx.array, Optional[mx.array]]:
-        B, L, D = hidden_states.shape
-
-        queries = self.q_proj(hidden_states).reshape(B, L, self.n_heads, -1)
-        keys = self.k_proj(hidden_states).reshape(B, L, self.n_kv_heads, -1)
-        values = self.v_proj(hidden_states).reshape(B, L, self.n_kv_heads, -1)
-
-        queries = self.q_norm(queries).transpose(0, 2, 1, 3)
-        keys = self.k_norm(keys).transpose(0, 2, 1, 3)
-        values = values.transpose(0, 2, 1, 3)
-
-        # Apply rope directly to queries and keys
-        if cos is not None and sin is not None:
-            queries, keys = apply_rotary_pos_emb(queries, keys, cos, sin)
-            if cache is not None:
-                keys, values = cache.update_and_fetch(keys, values)
-        else:
-            if cache is not None:
-                queries = self.rope(queries, offset=cache.offset+rope_deltas)
-                keys = self.rope(keys, offset=cache.offset+rope_deltas)
-                keys, values = cache.update_and_fetch(keys, values)
-            else:
-                queries = self.rope(queries)
-                keys = self.rope(keys)
-
-        output = scaled_dot_product_attention(
-            queries, keys, values, cache=cache, scale=self.scale, mask=attention_mask
-        )
-        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
-        return self.o_proj(output), None
-
-
-class TextMLP(nn.Module):
-    def __init__(self, config: TextConfig):
-        super().__init__()
-        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
-        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
-        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
-
-    def __call__(self, x):
-        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
-
-
-# Add this custom MoE implementation to replace SwitchGLU usage
-
-class TextMoEExperts(nn.Module):
-    def __init__(self, config: TextConfig):
-        super().__init__()
-        # Use the optimized SwitchGLU implementation for efficient expert computation
-        self.switch_glu = SwitchGLU(
-            input_dims=config.hidden_size,
-            hidden_dims=config.moe_intermediate_size,
-            num_experts=config.num_experts,
-            activation=nn.SiLU(),
-            bias=False
-        )
-
-    def __call__(self, hidden_states: mx.array, routing_weights: mx.array, router_indices: mx.array) -> mx.array:
-        # Use the efficient SwitchGLU implementation
-        # SwitchGLU handles the expert routing internally and is highly optimized
-        expert_output = self.switch_glu(hidden_states, router_indices)
-        
-        # Apply routing weights and sum over experts (top_k dimension)
-        weighted_output = expert_output * mx.expand_dims(routing_weights, -1)
-        final_output = mx.sum(weighted_output, axis=-2)
-        
-        return final_output
-
-class TextSparseMoeBlock(nn.Module):
-    def __init__(self, config: TextConfig):
-        super().__init__()
-        self.hidden_size = config.hidden_size
-        self.num_experts = config.num_experts
-        self.top_k = config.num_experts_per_tok
-        self.norm_topk_prob = config.norm_topk_prob
-        
-        self.gate = nn.Linear(self.hidden_size, self.num_experts, bias=False)
-        self.experts = TextMoEExperts(config)
-
-    def __call__(self, x: mx.array) -> mx.array:
-        batch_size, sequence_length, hidden_dim = x.shape
-        x_flat = x.reshape(-1, hidden_dim)
-        
-        # Router computation
-        router_logits = self.gate(x_flat)
-        routing_weights = mx.softmax(router_logits, axis=-1, precise=True)
-        
-        # Top-k selection
-        router_indices = mx.argpartition(-routing_weights, kth=self.top_k - 1, axis=-1)[..., :self.top_k]
-        routing_weights = mx.take_along_axis(routing_weights, router_indices, axis=-1)
-        
-        if self.norm_topk_prob:
-            routing_weights = routing_weights / mx.sum(routing_weights, axis=-1, keepdims=True)
-        
-        # Expert computation
-        final_hidden_states = self.experts(x, routing_weights, router_indices)
-        
-        return final_hidden_states
-
-
-class TextDecoderLayer(nn.Module):
-    def __init__(self, config: TextConfig, layer_idx: int):
-        super().__init__()
-        self.hidden_size = config.hidden_size
-        self.self_attn = TextAttention(config, layer_idx)
-        
-        # Determine if this layer should use MoE
-        use_moe = (
-            layer_idx not in config.mlp_only_layers and
-            config.num_experts > 0 and 
-            (layer_idx + 1) % config.decoder_sparse_step == 0
-        )
-        
-        if use_moe:
-            self.mlp = TextSparseMoeBlock(config)
-        else:
-            self.mlp = TextMLP(config)
-            
-        self.input_layernorm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.post_attention_layernorm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-
-    def __call__(
-        self,
-        hidden_states: mx.array,
-        attention_mask: Optional[mx.array] = None,
-        cache: Optional[Any] = None,
-        cos: Optional[mx.array] = None,
-        sin: Optional[mx.array] = None,
-        rope_deltas: Optional[mx.array] = None,
-    ) -> mx.array:
-        residual = hidden_states
-        hidden_states = self.input_layernorm(hidden_states)
-        
-        hidden_states, _ = self.self_attn(
-            hidden_states=hidden_states,
-            attention_mask=attention_mask,
-            cache=cache,
-            cos=cos,
-            sin=sin,
-            rope_deltas=rope_deltas,
-        )
-        hidden_states = residual + hidden_states
-        residual = hidden_states
-        hidden_states = self.post_attention_layernorm(hidden_states)
-        hidden_states = self.mlp(hidden_states)
-        hidden_states = residual + hidden_states
-        return hidden_states
-
-
-class TextModel(nn.Module):
-    def __init__(self, config: TextConfig):
-        super().__init__()
-        self.config = config
-        self.vocab_size = config.vocab_size
-        
-        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
-        self.layers = [
-            TextDecoderLayer(config, layer_idx) 
-            for layer_idx in range(config.num_hidden_layers)
-        ]
-        self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.rotary_emb = TextRotaryEmbedding(config)
-
-    def _deepstack_process(
-        self,
-        hidden_states: mx.array,
-        visual_pos_masks: mx.array,
-        deepstack_visual_embeds: mx.array,
-    ) -> mx.array:
-        if visual_pos_masks is None or deepstack_visual_embeds is None:
-            return hidden_states
-        B, L, D = hidden_states.shape
-        mask_flat = visual_pos_masks.astype(mx.int32).reshape(-1)
-        idx_flat = mx.cumsum(mask_flat, axis=0) - 1
-        N = deepstack_visual_embeds.shape[0]
-        idx_flat = mx.maximum(idx_flat, 0)
-        eq = (idx_flat[:, None] == mx.arange(N)[None, :]).astype(hidden_states.dtype)
-        add_flat = eq @ deepstack_visual_embeds.astype(hidden_states.dtype)
-        add_flat = add_flat * mask_flat[:, None].astype(hidden_states.dtype)
-        add = add_flat.reshape(B, L, D)
-        return hidden_states + add
-
-    def __call__(
-        self,
-        input_ids: Optional[mx.array] = None,
-        inputs_embeds: Optional[mx.array] = None,
-        attention_mask: Optional[mx.array] = None,
-        cache=None,
-        visual_pos_masks: Optional[mx.array] = None,
-        deepstack_visual_embeds: Optional[List[mx.array]] = None,
-        cos: Optional[mx.array] = None,
-        sin: Optional[mx.array] = None,
-        rope_deltas: Optional[mx.array] = None,
-    ):
-        if inputs_embeds is None:
-            inputs_embeds = self.embed_tokens(input_ids)
-
-        hidden_states = inputs_embeds
-
-        if attention_mask is None:
-            attention_mask = create_attention_mask(hidden_states, cache, return_array=True)
-
-        if cache is None:
-            cache = [None] * len(self.layers)
-
-        for layer_idx, (decoder_layer, c) in enumerate(zip(self.layers, cache)):
-            hidden_states = decoder_layer(
-                hidden_states,
-                attention_mask=attention_mask,
-                cache=c,
-                cos=cos,
-                sin=sin,
-                rope_deltas=rope_deltas,
-            )
-            if deepstack_visual_embeds is not None and layer_idx < len(deepstack_visual_embeds):
-                hidden_states = self._deepstack_process(hidden_states, visual_pos_masks, deepstack_visual_embeds[layer_idx])
-        hidden_states = self.norm(hidden_states)
-        return hidden_states
-
-
-# Standalone Vision Model
-class VEGModel(nn.Module):
-    def __init__(self, vision_config: VisionConfig):
-        super().__init__()
-        self.config = vision_config
-        self.visual = VisionModel(vision_config)
-    
-    def __call__(self, pixel_values: mx.array, image_grid_thw: mx.array):
-        return self.visual(pixel_values, image_grid_thw)
-    
-    def sanitize(self, weights):
-        sanitized = {}
-        for k, v in weights.items():
-            if 'visual.' in k:
-                # Remove prefixes to match our model structure
-                clean_key = k.replace('model.visual.', '').replace('visual.', '')
-                sanitized[f'visual.{clean_key}'] = v
-        return sanitized
-
-
-# Pure LLM Model (no vision components)
-class LLMModel(nn.Module):
-    def __init__(self, text_config: TextConfig):
-        super().__init__()
-        self.args = text_config
-        self.config = text_config
-        self.language_model = TextModel(text_config)
-        if not text_config.tie_word_embeddings:
-            self.lm_head = nn.Linear(text_config.hidden_size, text_config.vocab_size, bias=False)
-            
-    def get_rope_index(
-            self,
-            input_ids: Optional[mx.array] = None,
-            image_grid_thw: Optional[mx.array] = None,
-            attention_mask: Optional[mx.array] = None,
-        ) -> Tuple[mx.array, mx.array]:
-            """Simplified version for images only (no video support)."""
-
-            spatial_merge_size = 2
-            image_token_id = 151655
-            vision_start_token_id = 151652
-            mrope_position_deltas = []
-            
-            if input_ids is not None and image_grid_thw is not None:
-                total_input_ids = input_ids
-                if attention_mask is None:
-                    attention_mask = mx.ones_like(total_input_ids)
-                
-                batch_size, seq_len = input_ids.shape
-                position_ids_list = []
-                image_index = 0
-                
-                for i in range(batch_size):
-                    input_ids_seq = total_input_ids[i]
-                    mask_seq = attention_mask[i]
-                    
-                    # Use mask to get valid length
-                    valid_length = int(mx.sum(mask_seq).item())
-                    input_ids_seq = input_ids_seq[:valid_length]
-                    
-                    image_nums = 0
-                    # Find vision start tokens by iterating through the sequence
-                    vision_start_positions = []
-                    for pos in range(input_ids_seq.shape[0]):
-                        if input_ids_seq[pos].item() == vision_start_token_id:
-                            vision_start_positions.append(pos)
-                    
-                    if len(vision_start_positions) > 0:
-                        for pos in vision_start_positions:
-                            if pos + 1 < input_ids_seq.shape[0]:
-                                if input_ids_seq[pos + 1].item() == image_token_id:
-                                    image_nums += 1
-                    
-                    input_tokens = input_ids_seq.tolist()
-                    llm_pos_ids_list = []
-                    st = 0
-                    remain_images = image_nums
-                    
-                    for _ in range(image_nums):
-                        ed_image = input_tokens.index(image_token_id, st)
-                        
-                        t = image_grid_thw[image_index, 0].item()
-                        h = image_grid_thw[image_index, 1].item()
-                        w = image_grid_thw[image_index, 2].item()
-                        image_index += 1
-                        remain_images -= 1
-                        ed = ed_image
-                        
-                        llm_grid_t = int(t)
-                        llm_grid_h = int(h) // spatial_merge_size
-                        llm_grid_w = int(w) // spatial_merge_size
-                        text_len = ed - st
-
-                        st_idx = llm_pos_ids_list[-1].max().item() + 1 if len(llm_pos_ids_list) > 0 else 0
-                        text_pos = mx.arange(text_len).reshape(1, -1)
-                        text_pos = mx.broadcast_to(text_pos, (3, text_len)) + st_idx
-                        llm_pos_ids_list.append(text_pos)
-
-                        # t_index is always 0 because llm_grid_t is always 1 for images
-                        t_index = mx.arange(llm_grid_t).reshape(-1, 1)
-                        t_index = mx.broadcast_to(t_index, (llm_grid_t, llm_grid_h * llm_grid_w)).reshape(-1)
-                        
-                        h_index = mx.arange(llm_grid_h).reshape(1, -1, 1)
-                        h_index = mx.broadcast_to(h_index, (llm_grid_t, llm_grid_h, llm_grid_w)).reshape(-1)
-                        
-                        w_index = mx.arange(llm_grid_w).reshape(1, 1, -1)
-                        w_index = mx.broadcast_to(w_index, (llm_grid_t, llm_grid_h, llm_grid_w)).reshape(-1)
-                        
-                        vision_pos = mx.stack([t_index, h_index, w_index]) + text_len + st_idx
-                        llm_pos_ids_list.append(vision_pos)
-                        st = ed + llm_grid_t * llm_grid_h * llm_grid_w
-
-                    if st < len(input_tokens):
-                        st_idx = llm_pos_ids_list[-1].max().item() + 1 if len(llm_pos_ids_list) > 0 else 0
-                        text_len = len(input_tokens) - st
-                        text_pos = mx.arange(text_len).reshape(1, -1)
-                        text_pos = mx.broadcast_to(text_pos, (3, text_len)) + st_idx
-                        llm_pos_ids_list.append(text_pos)
-
-                    llm_positions = mx.concatenate(llm_pos_ids_list, axis=1).reshape(3, -1)
-                    
-                    # Create position_ids for this batch item, pad to seq_len
-                    batch_position_ids = mx.ones((3, seq_len), dtype=input_ids.dtype)
-                    valid_length = min(seq_len, llm_positions.shape[1])
-                    
-                    # Create new arrays for each dimension
-                    pos_dim0 = mx.concatenate([llm_positions[0, :valid_length], 
-                                            mx.ones(seq_len - valid_length, dtype=input_ids.dtype)])
-                    pos_dim1 = mx.concatenate([llm_positions[1, :valid_length], 
-                                            mx.ones(seq_len - valid_length, dtype=input_ids.dtype)])
-                    pos_dim2 = mx.concatenate([llm_positions[2, :valid_length], 
-                                            mx.ones(seq_len - valid_length, dtype=input_ids.dtype)])
-                    
-                    batch_position_ids = mx.stack([pos_dim0, pos_dim1, pos_dim2])
-                    position_ids_list.append(batch_position_ids)
-                    
-                    mrope_position_deltas.append(llm_positions.max().item() + 1 - len(total_input_ids[i]))
-                
-                # Stack all batch position_ids
-                position_ids = mx.stack(position_ids_list, axis=1)  # Shape: (3, batch_size, seq_len)
-                mrope_position_deltas = mx.array(mrope_position_deltas).reshape(-1, 1)
-                return position_ids, mrope_position_deltas
-            else:
-                if attention_mask is not None:
-                    position_ids = mx.cumsum(attention_mask.astype(mx.int32), axis=-1) - 1
-                    position_ids = mx.where(attention_mask == 0, 1, position_ids)
-                    position_ids = mx.expand_dims(position_ids, axis=0)
-                    position_ids = mx.broadcast_to(position_ids, (3, position_ids.shape[1], position_ids.shape[2]))
-                    max_position_ids = mx.max(mx.max(position_ids, axis=0, keepdims=False), axis=-1, keepdims=True)
-                    mrope_position_deltas = max_position_ids + 1 - attention_mask.shape[-1]
-                else:
-                    seq_len = input_ids.shape[1]
-                    batch_size = input_ids.shape[0]
-                    position_ids = mx.arange(seq_len).reshape(1, 1, -1)
-                    position_ids = mx.broadcast_to(position_ids, (3, batch_size, seq_len))
-                    mrope_position_deltas = mx.zeros((batch_size, 1), dtype=input_ids.dtype)
-
-                return position_ids, mrope_position_deltas
-    
-    def __call__(
-        self,
-        inputs: mx.array = None,
-        mask: mx.array = None,
-        cache=None,
-        inputs_embeds: Optional[mx.array] = None,
-        visual_pos_masks: Optional[mx.array] = None,
-        deepstack_visual_embeds: Optional[List[mx.array]] = None,
-        cos: Optional[mx.array] = None,
-        sin: Optional[mx.array] = None,
-        rope_deltas: Optional[mx.array] = None,
-    ):
-        out = self.language_model(
-            input_ids=inputs,
-            inputs_embeds=inputs_embeds,
-            attention_mask=mask,
-            cache=cache,
-            visual_pos_masks=visual_pos_masks,
-            deepstack_visual_embeds=deepstack_visual_embeds,
-            cos=cos,
-            sin=sin,
-            rope_deltas=rope_deltas,
-        )
-        if self.args.tie_word_embeddings:
-            return self.language_model.embed_tokens.as_linear(out)
-        else:
-            return self.lm_head(out)
-    
-    def sanitize(self, weights):
-        sanitized = {}
-        for k, v in weights.items():
-            if not ('visual.' in k):
-                # Handle key mapping from combined model to LLM-only model
-                clean_key = k
-                
-                # Remove model. prefix if present
-                if clean_key.startswith('model.'):
-                    clean_key = clean_key[6:]  # Remove 'model.'
-                
-                # Map language_ prefixed keys to language_model structure
-                if clean_key.startswith('language_'):
-                    if clean_key.startswith('language_layers.'):
-                        clean_key = 'language_model.layers.' + clean_key[16:]  # Map to language_model.layers.
-                    elif clean_key.startswith('language_embed_tokens.'):
-                        clean_key = 'language_model.embed_tokens.' + clean_key[22:]  # Map to language_model.embed_tokens.
-                    elif clean_key.startswith('language_norm.'):
-                        clean_key = 'language_model.norm.' + clean_key[14:]  # Map to language_model.norm.
-                
-                sanitized[clean_key] = v
-        
-        # Handle tied embeddings - remove lm_head if using tied embeddings
-        if self.args.tie_word_embeddings:
-            sanitized.pop("lm_head.weight", None)
-            
-        return sanitized
-    
-    @property
-    def layers(self):
-        return self.language_model.layers
-
-
-# Combined Model (for compatibility and utility functions)
-class Qwen3VLModel(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        self.args = args
-        self.config = args
-        self.visual = VisionModel(args.vision_config)
-        self.language_model = TextModel(args.text_config)
-    
-    def sanitize(self, weights):
-        # Map weights to match the combined model structure
-        sanitized = {}
-        for k, v in weights.items():
-            # Remove 'model.' prefix if present to match our structure
-            clean_key = k.replace('model.', '') if k.startswith('model.') else k
-            sanitized[clean_key] = v
-        return sanitized
-
-    def get_image_features(
-        self, 
-        pixel_values: mx.array, 
-        image_grid_thw: Optional[mx.array] = None
-    ):
-        image_embeds, deepstack_visual_embeds = self.visual(pixel_values, image_grid_thw)
-        # Split based on grid dimensions
-        if image_grid_thw is not None:
-            split_sizes = (mx.prod(image_grid_thw, axis=-1) // (self.visual.spatial_merge_size ** 2)).tolist()
-            # Convert sizes to indices for mx.split (cumulative sum, excluding the last)
-            split_indices = []
-            cumsum = 0
-            for size in split_sizes[:-1]:  # Exclude last element
-                cumsum += size
-                split_indices.append(cumsum)
-            
-            if split_indices:  # Only split if we have indices
-                image_embeds = mx.split(image_embeds, split_indices)
-            else:
-                image_embeds = [image_embeds]  # Single image case
-        return image_embeds, deepstack_visual_embeds
-
-
-    def __call__(
-        self,
-        input_ids: mx.array = None,
-        attention_mask: Optional[mx.array] = None,
-        inputs_embeds: Optional[mx.array] = None,
-        pixel_values: Optional[mx.array] = None,
-        image_grid_thw: Optional[mx.array] = None,
-        cache=None,
-        visual_pos_masks: Optional[mx.array] = None,
-        deepstack_visual_embeds: Optional[List[mx.array]] = None,
-        cos: Optional[mx.array] = None,
-        sin: Optional[mx.array] = None,
-        rope_deltas: Optional[mx.array] = None,
-    ):
-        if inputs_embeds is None:
-            inputs_embeds = self.language_model.embed_tokens(input_ids)
-
-        # Process images
-        
-        if pixel_values is not None:
-            image_embeds, deepstack_visual_embeds = self.get_image_features(
-                pixel_values, image_grid_thw
-            )
-            
-            # Create masks and embed visual features
-            if isinstance(image_embeds, list):
-                image_embeds = mx.concatenate(image_embeds, axis=0)
-            
-            # Find image token positions and replace with visual embeddings
-            image_mask = (input_ids == self.args.image_token_id)
-            visual_pos_masks = image_mask
-            
-            # Replace image tokens with visual embeddings
-            inputs_embeds = inputs_embeds.at[image_mask].set(
-                image_embeds.astype(inputs_embeds.dtype)
-            )
-
-
-        outputs = self.language_model(
-            inputs_embeds=inputs_embeds,
-            attention_mask=attention_mask,
-            cache=cache,
-            visual_pos_masks=visual_pos_masks,
-            deepstack_visual_embeds=deepstack_visual_embeds,
-            cos=cos,
-            sin=sin,
-            rope_deltas=rope_deltas,
-        )
-
-        return outputs
-
-
-def handle_multimodal_embeds(vision_model, llm_model, input_ids, pixel_values, image_grid_thw):
-    """
-    Handle the processing of multimodal embeddings including image features and position encoding.
-    
-    This function processes vision and text inputs to create unified embeddings that can be fed
-    into the language model. It handles:
-    - Vision feature extraction from pixel values
-    - Deepstack visual embedding collection
-    - Image token replacement in text embeddings
-    - Position encoding setup for MRoPE (Multi-dimensional RoPE)
-    
-    Args:
-        vision_model: The vision encoder model (VEGModel instance)
-        llm_model: The language model (LLMModel instance) 
-        input_ids: Tokenized text input with image token placeholders [batch_size, seq_len]
-        pixel_values: Preprocessed image pixel data [num_patches, feature_dim]
-        image_grid_thw: Grid dimensions for each image [num_images, 3] (time, height, width)
-        
-    Returns:
-        tuple: (inputs_embeds, deepstack_visual_embeds, visual_pos_masks, cos, sin, rope_deltas)
-            - inputs_embeds: Combined text and image embeddings [batch_size, seq_len, hidden_size]
-            - deepstack_visual_embeds: Multi-layer visual features for deepstack processing
-            - visual_pos_masks: Boolean mask indicating image token positions
-            - cos: Cosine values for rotary position encoding
-            - sin: Sine values for rotary position encoding  
-            - rope_deltas: Position offset deltas for rope computation
-    """
-    inputs_embeds = llm_model.language_model.embed_tokens(input_ids.squeeze(0))
-    deepstack_visual_embeds = None
-    visual_pos_masks = None
-    cos = None
-    sin = None
-    rope_deltas = 0
-    
-    if pixel_values is not None:
-        if pixel_values.ndim == 4:
-            pixel_values = mx.expand_dims(pixel_values, axis=2)
-        
-        # Process each image individually to prevent feature mixing
-        image_embeds_list = []
-        all_deepstack_embeds = []
-        
-        # Calculate cumulative indices for each image
-        cumulative_patches = 0
-        
-        for i in range(image_grid_thw.shape[0]):
-            # Calculate number of patches for current image
-            current_patches = int(image_grid_thw[i, 1] * image_grid_thw[i, 2])
-            start_idx = cumulative_patches
-            end_idx = cumulative_patches + current_patches
-            cumulative_patches += current_patches
-            
-            single_pixel_values = pixel_values[start_idx:end_idx]
-            single_grid_thw = image_grid_thw[i:i+1]
-            
-            # Use vision model directly
-            single_embeds, single_deepstack = vision_model(single_pixel_values, single_grid_thw)
-            
-            # Split based on grid dimensions
-            if single_grid_thw is not None:
-                split_sizes = (mx.prod(single_grid_thw, axis=-1) // (vision_model.visual.spatial_merge_size ** 2)).tolist()
-                split_indices = []
-                cumsum = 0
-                for size in split_sizes[:-1]:
-                    cumsum += size
-                    split_indices.append(cumsum)
-                
-                if split_indices:
-                    single_embeds = mx.split(single_embeds, split_indices)
-                else:
-                    single_embeds = [single_embeds]
-            
-            image_embeds_list.extend(single_embeds)
-            
-            # Collect deepstack embeddings
-            if i == 0:
-                all_deepstack_embeds = single_deepstack
-            else:
-                # Concatenate deepstack embeddings from different images
-                for j in range(len(all_deepstack_embeds)):
-                    all_deepstack_embeds[j] = mx.concatenate([all_deepstack_embeds[j], single_deepstack[j]], axis=0)
-        
-        deepstack_visual_embeds = all_deepstack_embeds
-        
-        # Concatenate all image embeddings for processing
-        image_embeds = mx.concatenate(image_embeds_list, axis=0)
-        
-        # Find all image token positions
-        image_token_id = 151655  # Default image token ID
-        image_mask = (input_ids.squeeze(0) == image_token_id)
-        image_mask_np = np.array(image_mask)
-        image_token_positions = np.where(image_mask_np)[0]
-        
-        # Verify we have the correct number of image tokens
-        expected_total_tokens = sum(embed.shape[0] for embed in image_embeds_list)
-        assert len(image_token_positions) == expected_total_tokens, f"Expected {expected_total_tokens} image tokens, got {len(image_token_positions)}"
-        
-        # Replace image tokens with image embeddings
-        seq_len = inputs_embeds.shape[0]
-        result = inputs_embeds
-        
-        # Replace image tokens with image embeddings sequentially
-        embed_idx = 0
-        for img_embed in image_embeds_list:
-            for patch_idx in range(img_embed.shape[0]):
-                token_pos = image_token_positions[embed_idx]
-                pos_mask = mx.arange(seq_len) == token_pos
-                result = mx.where(
-                    mx.expand_dims(pos_mask, axis=-1),
-                    mx.expand_dims(img_embed[patch_idx], axis=0).astype(inputs_embeds.dtype),
-                    result
-                )
-                embed_idx += 1
-
-        inputs_embeds = result
-        position_ids, rope_deltas = llm_model.get_rope_index(input_ids, image_grid_thw)
-        cos, sin = llm_model.language_model.rotary_emb(inputs_embeds, position_ids)
-        if inputs_embeds.ndim == 2:
-            inputs_embeds = mx.expand_dims(inputs_embeds, axis=0)
-        
-        if image_mask is not None:
-            visual_pos_masks = image_mask
-    
-    return inputs_embeds, deepstack_visual_embeds, visual_pos_masks, cos, sin, rope_deltas
-
-
-# Legacy Model wrapper (for backward compatibility)
-class Model(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        self.args = args
-        self.model = Qwen3VLModel(args)
-        if not args.text_config.tie_word_embeddings:
-            self.lm_head = nn.Linear(args.text_config.hidden_size, args.text_config.vocab_size, bias=False)
-
-    def __call__(
-        self,
-        inputs: mx.array = None,
-        mask: mx.array = None,
-        cache=None,
-        inputs_embeds: Optional[mx.array] = None,
-        pixel_values: Optional[mx.array] = None,        
-        image_grid_thw: Optional[mx.array] = None,
-        visual_pos_masks: Optional[mx.array] = None,
-        deepstack_visual_embeds: Optional[List[mx.array]] = None,
-        cos: Optional[mx.array] = None,
-        sin: Optional[mx.array] = None,
-        rope_deltas: Optional[mx.array] = None,
-    ):
-        out = self.model(
-            input_ids=inputs,
-            inputs_embeds=inputs_embeds,
-            attention_mask=mask,
-            cache=cache,
-            pixel_values=pixel_values,
-            image_grid_thw=image_grid_thw,
-            visual_pos_masks=visual_pos_masks,
-            deepstack_visual_embeds=deepstack_visual_embeds,
-            cos=cos,
-            sin=sin,
-            rope_deltas=rope_deltas,
-        )
-        if self.args.text_config.tie_word_embeddings:
-            return self.model.language_model.embed_tokens.as_linear(out)
-        else:
-            return self.lm_head(out)
-
-    def sanitize(self, weights):
-        # Remove any unnecessary weights
-        sanitized = {}
-        for k, v in weights.items():
-            sanitized[k] = v
-        
-        # Handle tied embeddings - remove lm_head if using tied embeddings
-        if self.args.text_config.tie_word_embeddings:
-            sanitized.pop("lm_head.weight", None)
-            
-        return sanitized
-
-    @property
-    def layers(self):
-        return self.model.language_model.layers