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/llama4/vision.py

@@ -1,560 +0,0 @@
-import inspect
-import math
-from dataclasses import dataclass
-from typing import Optional, Tuple
-
-import mlx.core as mx
-import mlx.nn as nn
-
-from ..base import pixel_shuffle
-
-
-@dataclass
-class VisionConfig:
-    model_type: str
-    hidden_size: int
-    image_size: int
-    initializer_range: float
-    intermediate_size: int
-    norm_eps: float
-    num_attention_heads: int
-    num_channels: int
-    num_hidden_layers: int
-    patch_size: int
-    pixel_shuffle_ratio: float
-    projector_dropout: float
-    projector_input_dim: int
-    projector_output_dim: int
-    rope_theta: float
-    vision_feature_layer: int
-    vision_feature_select_strategy: str
-    vision_output_dim: int
-
-    @classmethod
-    def from_dict(cls, params):
-        return cls(
-            **{
-                k: v
-                for k, v in params.items()
-                if k in inspect.signature(cls).parameters
-            }
-        )
-
-
-def check_array_shape(arr):
-    shape = arr.shape
-
-    # Check if the shape has 4 dimensions
-    if len(shape) != 4:
-        return False
-
-    out_channels, kH, KW, _ = shape
-
-    # Check if out_channels is the largest, and kH and KW are the same
-    if (out_channels >= kH) and (out_channels >= KW) and (kH == KW):
-        return True
-    else:
-        return False
-
-
-class Llama4MultiModalProjector(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.linear_1 = nn.Linear(
-            config.vision_config.vision_output_dim,
-            config.text_config.hidden_size,
-            bias=False,
-        )
-
-    def __call__(self, image_features):
-        hidden_states = self.linear_1(image_features)
-        return hidden_states
-
-
-class Llama4VisionPixelShuffleMLP(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.pixel_shuffle_ratio = config.pixel_shuffle_ratio
-        self.inner_dim = int(
-            config.projector_input_dim // (self.pixel_shuffle_ratio**2)
-        )
-        self.output_dim = config.projector_output_dim
-        self.mlp = Llama4VisionMLP(config, bias=False, is_projector=True)
-
-    def __call__(self, encoded_patches: mx.array) -> mx.array:
-        encoded_patches = pixel_shuffle(encoded_patches, self.pixel_shuffle_ratio)
-        return self.mlp(encoded_patches)
-
-
-# TODO there is a different RoPE for vision encoder, defined as below
-def reshape_for_broadcast(freqs_ci: mx.array, query: mx.array):
-    ndim = query.ndim
-    shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(query.shape)]
-    return freqs_ci.reshape(*shape)
-
-
-def view_as_complex(x):
-    """
-    Convert a tensor with shape (..., 2) to a complex tensor with shape (...).
-
-    Args:
-        x: A real tensor with last dimension of size 2.
-
-    Returns:
-        A complex tensor with size one less than the input.
-    """
-    # Ensure the last dimension is size 2
-    assert x.shape[-1] == 2, f"Last dimension must be 2, got {x.shape[-1]}"
-
-    # Get real and imaginary parts
-    real, imag = x[..., 0], x[..., 1]
-
-    # Create complex tensor
-    return real + 1j * imag
-
-
-def view_as_real(x):
-    """
-    Convert a complex tensor with shape (...) to a real tensor with shape (..., 2).
-
-    Args:
-        x: A complex tensor.
-
-    Returns:
-        A real tensor with an extra dimension of size 2.
-    """
-    # Get real and imaginary parts
-    real = mx.real(x)
-    imag = mx.imag(x)
-
-    # Combine into a tensor with last dimension 2
-    return mx.stack([real, imag], axis=-1)
-
-
-def vision_apply_rotary_emb(
-    query: mx.array,
-    key: mx.array,
-    freqs_ci: mx.array,
-) -> Tuple[mx.array, mx.array]:
-
-    query_ = view_as_complex(query.astype(mx.float32).reshape(*query.shape[:-1], -1, 2))
-    key_ = view_as_complex(key.astype(mx.float32).reshape(*key.shape[:-1], -1, 2))
-    freqs_ci = reshape_for_broadcast(freqs_ci=freqs_ci, query=query_)
-    query_out = view_as_real(query_ * freqs_ci).flatten(3)
-    key_out = view_as_real(key_ * freqs_ci).flatten(3)
-    return query_out.astype(query.dtype), key_out.astype(key.dtype)
-
-
-class Llama4VisionAttention(nn.Module):
-    def __init__(self, config: VisionConfig):
-        super().__init__()
-        self.config = config
-        self.embed_dim = config.hidden_size
-        self.num_heads = config.num_attention_heads
-        self.head_dim = config.hidden_size // config.num_attention_heads
-        self.num_key_value_groups = 1
-        self.scale = self.head_dim**-0.5
-
-        self.q_proj = nn.Linear(
-            self.embed_dim, self.num_heads * self.head_dim, bias=True
-        )
-        self.k_proj = nn.Linear(
-            self.embed_dim, self.num_heads * self.head_dim, bias=True
-        )
-        self.v_proj = nn.Linear(
-            self.embed_dim, self.num_heads * self.head_dim, bias=True
-        )
-        self.o_proj = nn.Linear(
-            self.num_heads * self.head_dim, self.embed_dim, bias=True
-        )
-
-    def __call__(
-        self,
-        hidden_states: mx.array,
-        freqs_ci: mx.array,
-        mask: Optional[mx.array] = None,
-        cache: Optional[mx.array] = None,
-    ):
-        B, L, D = hidden_states.shape
-
-        query_states = self.q_proj(hidden_states).reshape(B, L, self.num_heads, -1)
-        key_states = self.k_proj(hidden_states).reshape(B, L, self.num_heads, -1)
-        value_states = self.v_proj(hidden_states).reshape(B, L, self.num_heads, -1)
-
-        query_states, key_states = vision_apply_rotary_emb(
-            query_states, key_states, freqs_ci=freqs_ci
-        )
-
-        query_states = query_states.transpose(0, 2, 1, 3)
-        key_states = key_states.transpose(0, 2, 1, 3)
-        value_states = value_states.transpose(0, 2, 1, 3)
-
-        attn_output = mx.fast.scaled_dot_product_attention(
-            query_states, key_states, value_states, scale=self.scale
-        )
-
-        attn_output = attn_output.transpose(0, 2, 1, 3).reshape(B, L, -1)
-        attn_output = self.o_proj(attn_output)
-        return attn_output
-
-
-class Llama4VisionMLP(nn.Module):
-    def __init__(self, config, bias=True, is_projector=False):
-        super().__init__()
-        self.config = config
-        self.activation_fn = nn.GELU(approx="fast")  # ACT2FN[config.hidden_act]
-        self.is_projector = is_projector
-        self.hidden_size = config.hidden_size
-        self.intermediate_size = config.intermediate_size
-
-        # Determine dimensions for first linear layer based on whether this is a projector
-        fc1_input_dim = self.intermediate_size if is_projector else self.hidden_size
-        fc1_output_dim = (
-            config.projector_input_dim if is_projector else self.intermediate_size
-        )
-
-        self.fc1 = nn.Linear(fc1_input_dim, fc1_output_dim, bias=bias)
-
-        # Determine dimensions for second linear layer
-        fc2_input_dim = (
-            config.projector_output_dim if is_projector else self.intermediate_size
-        )
-        fc2_output_dim = (
-            config.projector_output_dim if is_projector else self.hidden_size
-        )
-
-        self.fc2 = nn.Linear(fc2_input_dim, fc2_output_dim, bias=bias)
-
-        self.is_projector = is_projector
-
-    def __call__(self, hidden_states: mx.array) -> mx.array:
-        hidden_states = self.fc1(hidden_states)
-        hidden_states = self.activation_fn(hidden_states)
-
-        if self.is_projector:
-            return self.activation_fn(self.fc2(hidden_states))
-
-        return self.fc2(hidden_states)
-
-
-class Llama4VisionEncoderLayer(nn.Module):
-    def __init__(self, config: VisionConfig):
-        super().__init__()
-        self.hidden_size = config.hidden_size
-
-        self.self_attn = Llama4VisionAttention(config)
-        self.mlp = Llama4VisionMLP(config)
-
-        self.input_layernorm = nn.LayerNorm(config.hidden_size)
-        self.post_attention_layernorm = nn.LayerNorm(config.hidden_size)
-
-    def __call__(
-        self,
-        hidden_state: mx.array,
-        freqs_ci: mx.array,
-        mask: Optional[mx.array] = None,
-    ):
-        # Self Attention
-        residual = hidden_state
-
-        hidden_state = self.input_layernorm(hidden_state)
-
-        hidden_state = self.self_attn(
-            hidden_state,
-            freqs_ci=freqs_ci,
-            mask=mask,
-        )
-        hidden_state = residual + hidden_state
-
-        # Feed forward
-        residual = hidden_state
-        hidden_state = self.post_attention_layernorm(hidden_state)
-        hidden_state = self.mlp(hidden_state)
-        hidden_state = residual + hidden_state
-        return hidden_state
-
-
-class Llama4VisionEncoder(nn.Module):
-    """
-    Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a
-    [`Llama4VisionEncoderLayer`].
-
-    Args:
-        config: VisionConfig
-    """
-
-    def __init__(self, config: VisionConfig):
-        super().__init__()
-        self.config = config
-        self.layers = [
-            Llama4VisionEncoderLayer(config) for _ in range(config.num_hidden_layers)
-        ]
-        self.config = config
-
-    def __call__(
-        self,
-        hidden_states: mx.array,
-        freqs_ci: mx.array,  # TODO move this to an attribute instead of keeping it around
-        mask: Optional[mx.array] = None,
-    ):
-
-        for i, encoder_layer in enumerate(self.layers):
-            hidden_states = encoder_layer(
-                hidden_state=hidden_states,
-                mask=mask,
-                freqs_ci=freqs_ci,
-            )
-
-        return hidden_states
-
-
-class Llama4UnfoldConvolution(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        kernel_size = config.patch_size
-        if isinstance(kernel_size, int):
-            kernel_size = (kernel_size, kernel_size)
-        self.kernel_size = kernel_size
-        self.stride = config.patch_size
-        self.linear = nn.Linear(
-            config.num_channels * kernel_size[0] * kernel_size[1],
-            config.hidden_size,
-            bias=False,
-        )
-
-    def _pair(self, x):
-        """Convert input to a pair of values."""
-        if isinstance(x, (list, tuple)):
-            return tuple(x)
-        return (x, x)
-
-    def unfold(self, input_tensor):
-        """
-        Extract sliding local blocks from a batched input tensor (MLX implementation).
-
-        This is equivalent to PyTorch's nn.functional.unfold or im2col operation.
-
-        Args:
-            input_tensor: Input tensor of shape (B, C, H, W)
-
-        Returns:
-            Unfolded tensor of shape (B, C*kernel_height*kernel_width, L)
-            where L is the number of blocks
-        """
-        # Convert to pairs
-        kernel_size = self._pair(self.kernel_size)
-        stride = self._pair(self.stride)
-        padding = (0, 0)  # No padding in the original code
-        dilation = (1, 1)  # Default dilation
-
-        # Input shape
-        batch_size, channels, height, width = input_tensor.shape
-
-        # Calculate output dimensions
-        height_out = (
-            height + 2 * padding[0] - dilation[0] * (kernel_size[0] - 1) - 1
-        ) // stride[0] + 1
-        width_out = (
-            width + 2 * padding[1] - dilation[1] * (kernel_size[1] - 1) - 1
-        ) // stride[1] + 1
-
-        # Initialize output arrays
-        blocks = []
-
-        # Extract blocks
-        for i in range(0, height - kernel_size[0] * dilation[0] + 1, stride[0]):
-            for j in range(0, width - kernel_size[1] * dilation[1] + 1, stride[1]):
-                # Extract the block for all channels
-                block = []
-                for di in range(kernel_size[0]):
-                    for dj in range(kernel_size[1]):
-                        h_idx = i + di * dilation[0]
-                        w_idx = j + dj * dilation[1]
-                        # Get the block for all channels and add to our list
-                        block.append(input_tensor[:, :, h_idx, w_idx])
-
-                # Stack the channel-blocks
-                block = mx.stack(block, axis=1)  # Shape: (B, k*k, C)
-                block = mx.transpose(block, [0, 2, 1])  # Shape: (B, C, k*k)
-                blocks.append(block)
-
-        # Stack all blocks together
-        result = mx.stack(blocks, axis=-1)  # Shape: (B, C, k*k, L)
-
-        # Reshape to match PyTorch's unfold output format: (B, C*k*k, L)
-        result = mx.reshape(
-            result,
-            (
-                batch_size,
-                channels * kernel_size[0] * kernel_size[1],
-                height_out * width_out,
-            ),
-        )
-
-        return result
-
-    def __call__(self, hidden_states: mx.array) -> mx.array:
-        hidden_states = self.unfold(hidden_states)
-        hidden_states = hidden_states.swapaxes(1, 2)
-        hidden_states = self.linear(hidden_states)
-        return hidden_states
-
-
-class Llama4VisionRotaryEmbedding:
-    def __init__(self, config):
-        super().__init__()
-        idx = config.image_size // config.patch_size
-        img_idx = mx.arange(idx**2, dtype=mx.int32).reshape(idx**2, 1)
-        img_idx = mx.concatenate([img_idx, img_idx[:1]], axis=0)
-        img_idx[-1, -1] = -2  # ID_CLS_TOKEN
-        frequencies_x = img_idx % idx  # get the coordinates of the 2d matrix along x
-        frequencies_y = img_idx // idx  # get the coordinates of the 2d matrix along y
-        freq_dim = config.hidden_size // config.num_attention_heads // 2
-        rope_freq = 1.0 / (
-            config.rope_theta
-            ** (
-                mx.arange(0, freq_dim, 2, dtype=mx.float32)[: (freq_dim // 2)]
-                / freq_dim
-            )
-        )
-
-        # Expand dimensions for frequencies_x and frequencies_y
-        freqs_x_expanded = (frequencies_x + 1)[..., None] * rope_freq[None, None, :]
-        freqs_y_expanded = (frequencies_y + 1)[..., None] * rope_freq[None, None, :]
-
-        def repeat_interleave(tensor, repeats, dim=-1):
-            # Get the shape
-            shape = list(tensor.shape)
-
-            # Reshape to add an extra dimension for repeating
-            tensor = mx.reshape(tensor, shape[:-1] + [shape[-1], 1])
-
-            # Repeat along the new dimension
-            tensor = mx.repeat(tensor, repeats, axis=-1)
-
-            # Reshape to flatten the last two dimensions
-            return mx.reshape(tensor, shape[:-1] + [shape[-1] * repeats])
-
-        # Apply interleaving
-        freqs_x = repeat_interleave(freqs_x_expanded, 2)
-        freqs_y = repeat_interleave(freqs_y_expanded, 2)
-        freqs = mx.concatenate([freqs_x, freqs_y], axis=-1).astype(mx.float32)[..., ::2]
-        # Replaced masked_fill with where
-        mask = img_idx.reshape(-1, 1, 1) < 0
-        freqs = mx.where(mask, mx.zeros_like(freqs), freqs)
-        freq_cis = mx.stack([mx.cos(freqs), mx.sin(freqs)], axis=-1)
-        freq_cis = view_as_complex(freq_cis)
-        self.freqs_ci = freq_cis  # idx**2, idx**2, idx * 2
-
-    def __call__(self, hidden_states):
-        return self.freqs_ci
-
-
-class VisionModel(nn.Module):
-    def __init__(self, config: VisionConfig):
-        super().__init__()
-        self.image_size = config.image_size
-        self.patch_size = config.patch_size
-        self.hidden_size = config.hidden_size
-        self.num_channels = config.num_channels
-        self.model_type = config.model_type
-        if self.model_type not in ["llama4", "llama4_vision_model"]:
-            raise ValueError(f"Model type {self.model_type} not supported")
-
-        self.num_patches = (self.image_size // self.patch_size) ** 2 + 1
-        self.scale = config.hidden_size**-0.5
-
-        self.class_embedding = self.scale * mx.random.normal((self.hidden_size,))
-        self.positional_embedding_vlm = self.scale * mx.random.normal(
-            (self.num_patches, self.hidden_size)
-        )
-
-        self.patch_embedding = Llama4UnfoldConvolution(config)
-
-        self.rotary_embedding = Llama4VisionRotaryEmbedding(config)
-
-        # layer norms
-        self.layernorm_pre = nn.LayerNorm(self.hidden_size)
-        self.layernorm_post = nn.LayerNorm(self.hidden_size)
-
-        # encoders
-        self.model = Llama4VisionEncoder(config)
-        self.vision_adapter = Llama4VisionPixelShuffleMLP(config)
-
-    def get_input_embeddings(self):
-        """
-        This function is used to fetch the first embedding layer to activate grads on inputs.
-        """
-        return self.patch_embedding
-
-    def __call__(
-        self,
-        pixel_values: mx.array,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        capture_activations: Optional[bool] = True,
-    ):
-
-        batch_size_times_num_tiles, num_channels, height, width = pixel_values.shape
-        num_concurrent_media = 1
-        num_chunks = 1
-
-        hidden_state = self.patch_embedding(pixel_values)
-
-        _, num_patches, hidden_dim = hidden_state.shape
-
-        # Add cls token
-        hidden_state = hidden_state.reshape(
-            batch_size_times_num_tiles * num_concurrent_media * num_chunks,
-            num_patches,
-            hidden_dim,
-        )
-
-        class_embedding = mx.broadcast_to(
-            self.class_embedding, (hidden_state.shape[0], 1, hidden_state.shape[-1])
-        )
-        hidden_state = mx.concatenate([hidden_state, class_embedding], axis=1)
-        num_patches += 1
-
-        # Position embeddings
-        hidden_state = hidden_state.reshape(
-            batch_size_times_num_tiles * num_concurrent_media,
-            num_chunks,
-            num_patches,
-            hidden_dim,
-        )
-
-        positional_embedding = self.positional_embedding_vlm
-        hidden_state = hidden_state + positional_embedding
-
-        hidden_state = self.layernorm_pre(hidden_state)
-
-        hidden_state = hidden_state.reshape(batch_size_times_num_tiles, -1, hidden_dim)
-        freqs_ci = self.rotary_embedding(pixel_values)
-
-        hidden_state = self.model(
-            hidden_state,
-            mask=None,
-            freqs_ci=freqs_ci,
-        )
-
-        hidden_state = self.layernorm_post(hidden_state)
-
-        hidden_state = hidden_state[:, :-1, :]
-
-        # now, we use Llama4VisionPixelShuffle + mlp to project embeddings
-        final_hidden_state = self.vision_adapter(hidden_state)
-
-        # Return only the final state
-        return final_hidden_state
-
-    def sanitize(self, weights):
-        sanitized_weights = {}
-        for k, v in weights.items():
-            if "position_ids" in k:
-                # Remove unused position_ids
-                continue
-            else:
-                sanitized_weights[k] = v
-
-        return sanitized_weights

nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/llava/__init__.py

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