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/kernels.py

@@ -1,307 +0,0 @@
-import mlx.core as mx
-
-
-def nearest_interpolate(x, size=None, scale_factor=None):
-    """
-    Nearest neighbor interpolation that exactly matches PyTorch's behavior.
-    """
-    # Get input dimensions
-    batch_size, channels, in_h, in_w = x.shape
-
-    # Calculate output dimensions
-    if size is not None:
-        out_h, out_w = size
-    elif scale_factor is not None:
-        if isinstance(scale_factor, (int, float)):
-            scale_h = scale_w = scale_factor
-        else:
-            scale_h, scale_w = scale_factor
-        out_h, out_w = int(in_h * scale_h), int(in_w * scale_w)
-    else:
-        raise ValueError("Either size or scale_factor must be specified")
-
-    # Create dimensions tensor
-    dims = mx.array([batch_size, channels, in_h, in_w, out_h, out_w], dtype=mx.int32)
-
-    # Reshape input tensor to 1D for kernel processing
-    x_flat = x.reshape(-1)
-    input_dtype = x.dtype
-    if input_dtype != mx.float32:
-        x_flat = x_flat.astype(mx.float32)
-
-    # Metal kernel source that matches PyTorch's coordinate calculation
-    source = """
-        uint x_out = thread_position_in_grid.x;
-        uint y_out = thread_position_in_grid.y;
-        uint bc_idx = thread_position_in_grid.z;
-
-        int batch_size = dims[0];
-        int channels = dims[1];
-        int in_h = dims[2];
-        int in_w = dims[3];
-        int out_h = dims[4];
-        int out_w = dims[5];
-
-        if (x_out >= (uint)out_w || y_out >= (uint)out_h || bc_idx >= (uint)(batch_size * channels))
-            return;
-
-        int c = bc_idx % channels;
-        int b = bc_idx / channels;
-
-        // PyTorch's coordinate calculation for nearest neighbor
-        // This matches: torch.nn.functional.interpolate(..., mode='nearest')
-        float scale_h = float(in_h) / float(out_h);
-        float scale_w = float(in_w) / float(out_w);
-
-        // PyTorch uses floor for nearest neighbor coordinate mapping
-        int y_in = int(floor(float(y_out) * scale_h));
-        int x_in = int(floor(float(x_out) * scale_w));
-
-        // Clamp to bounds
-        y_in = max(0, min(y_in, in_h - 1));
-        x_in = max(0, min(x_in, in_w - 1));
-
-        int input_offset = ((b * channels + c) * in_h + y_in) * in_w + x_in;
-        int output_offset = ((b * channels + c) * out_h + y_out) * out_w + x_out;
-
-        output[output_offset] = input[input_offset];
-    """
-
-    # Create and run kernel
-    kernel = mx.fast.metal_kernel(
-        name="nearest_interpolation",
-        input_names=["input", "dims"],
-        output_names=["output"],
-        source=source,
-    )
-
-    threadgroup = get_optimal_threadgroup(out_w, out_h)
-    outputs = kernel(
-        inputs=[x_flat, dims],
-        grid=(out_w, out_h, batch_size * channels),
-        threadgroup=threadgroup,
-        output_shapes=[(batch_size * channels * out_h * out_w,)],
-        output_dtypes=[mx.float32],
-    )
-
-    result = outputs[0].reshape(batch_size, channels, out_h, out_w)
-    if input_dtype != mx.float32:
-        result = result.astype(input_dtype)
-
-    return result
-
-
-def bicubic_interpolate(x, size=None, scale_factor=None, align_corners=False):
-    """
-    Bicubic interpolation using MLX's built-in interpolate function.
-
-    Args:
-        x: MLX tensor of shape [B, C, H, W]
-        size: Tuple of (out_h, out_w) or None
-        scale_factor: Float or tuple of (scale_h, scale_w) or None
-        align_corners: Whether to align corners
-
-    Returns:
-        Interpolated MLX tensor
-    """
-    # Get input dimensions
-    batch_size, channels, in_h, in_w = x.shape
-
-    # Calculate output dimensions
-    if size is not None:
-        out_h, out_w = size
-        scale_h, scale_w = out_h / in_h, out_w / in_w
-    elif scale_factor is not None:
-        if isinstance(scale_factor, (int, float)):
-            scale_h = scale_w = scale_factor
-        else:
-            scale_h, scale_w = scale_factor
-        out_h, out_w = int(in_h * scale_h), int(in_w * scale_w)
-    else:
-        raise ValueError("Either size or scale_factor must be specified")
-
-    # Create scale and align_corners parameters tensor
-    params = mx.array(
-        [scale_h, scale_w, 1.0 if align_corners else 0.0], dtype=mx.float32
-    )
-
-    # Create dimensions tensor
-    dims = mx.array([batch_size, channels, in_h, in_w, out_h, out_w], dtype=mx.int32)
-
-    # Reshape input tensor to 1D for kernel processing
-    x_flat = x.reshape(-1)
-
-    # Convert to float32 for processing if needed
-    input_dtype = x.dtype
-    if input_dtype != mx.float32:
-        x_flat = x_flat.astype(mx.float32)
-
-    # Metal kernel source code
-    source = """
-        // Get thread position
-        uint x_out = thread_position_in_grid.x;
-        uint y_out = thread_position_in_grid.y;
-        uint bc_idx = thread_position_in_grid.z;
-
-        // Extract dimensions from dims
-        int batch_size = dims[0];
-        int channels = dims[1];
-        int in_h = dims[2];
-        int in_w = dims[3];
-        int out_h = dims[4];
-        int out_w = dims[5];
-
-        // Extract scales and flags
-        float scale_h = params[0];
-        float scale_w = params[1];
-        bool align_corners = params[2] > 0.5;
-
-        // Check bounds
-        if (x_out >= (uint)out_w || y_out >= (uint)out_h || bc_idx >= (uint)(batch_size * channels))
-            return;
-
-        // Calculate batch and channel indices
-        int c = bc_idx % channels;
-        int b = bc_idx / channels;
-
-        // Calculate input coordinates based on output position
-        float x_in, y_in;
-
-        if (align_corners && out_w > 1 && out_h > 1) {
-            x_in = float(x_out) * (in_w - 1) / (out_w - 1);
-            y_in = float(y_out) * (in_h - 1) / (out_h - 1);
-        } else {
-            // Fix the alignment calculation to ensure consistent mapping across thread boundaries
-            x_in = ((float(x_out) + 0.5f) / float(out_w)) * float(in_w) - 0.5f;
-            y_in = ((float(y_out) + 0.5f) / float(out_h)) * float(in_h) - 0.5f;
-        }
-
-        // Get integer and fractional parts
-        int x0 = int(floor(x_in));
-        int y0 = int(floor(y_in));
-        float x_frac = x_in - x0;
-        float y_frac = y_in - y0;
-
-        // Improved cubic kernel function for better continuity
-        auto cubic_kernel = [](float x) -> float {
-            float absx = fabs(x);
-            float absx2 = absx * absx;
-            float absx3 = absx2 * absx;
-
-            // Use a=-0.5 for smoother interpolation
-            const float a = -0.5f;
-
-            if (absx <= 1.0f) {
-                return (a+2.0f)*absx3 - (a+3.0f)*absx2 + 1.0f;
-            } else if (absx < 2.0f) {
-                return a*absx3 - 5.0f*a*absx2 + 8.0f*a*absx - 4.0f*a;
-            }
-            return 0.0f;
-        };
-
-        // Perform bicubic interpolation with improved boundary handling
-        float result = 0.0f;
-        float weight_sum = 0.0f;  // Track weight sum for normalization
-
-        for (int i = -1; i <= 2; i++) {
-            int y_pos = y0 + i;
-            // Clamp y coordinate to valid range
-            y_pos = max(0, min(y_pos, in_h - 1));
-            float wy = cubic_kernel(y_frac - i);
-
-            for (int j = -1; j <= 2; j++) {
-                int x_pos = x0 + j;
-                // Clamp x coordinate to valid range
-                x_pos = max(0, min(x_pos, in_w - 1));
-                float wx = cubic_kernel(x_frac - j);
-                float weight = wy * wx;
-
-                // Calculate input tensor offset
-                int input_offset = ((b * channels + c) * in_h + y_pos) * in_w + x_pos;
-
-                // Add weighted contribution
-                result += input[input_offset] * weight;
-                weight_sum += weight;
-            }
-        }
-
-        // Normalize by weight sum to ensure consistent intensity
-        if (weight_sum > 0.0f) {
-            result /= weight_sum;
-        }
-
-        // Calculate output tensor offset
-        int output_offset = ((b * channels + c) * out_h + y_out) * out_w + x_out;
-
-        // Assign the result to output
-        output[output_offset] = (float)result;
-    """
-
-    # Create the kernel
-    kernel = mx.fast.metal_kernel(
-        name="bicubic_interpolation",
-        input_names=["input", "dims", "params"],
-        output_names=["output"],
-        source=source,
-    )
-
-    # Run the kernel
-    threadgroup = get_optimal_threadgroup(out_w, out_h)
-    outputs = kernel(
-        inputs=[x_flat, dims, params],
-        grid=(out_w, out_h, batch_size * channels),
-        threadgroup=threadgroup,
-        output_shapes=[(batch_size * channels * out_h * out_w,)],
-        output_dtypes=[mx.float32],  # Always use float32 for kernel output
-    )
-
-    # Reshape output back to 4D tensor and convert back to original dtype
-    result = outputs[0].reshape(batch_size, channels, out_h, out_w)
-    if input_dtype != mx.float32:
-        result = result.astype(input_dtype)
-
-    return result
-
-
-def get_optimal_threadgroup(out_w, out_h):
-    # Calculate optimal threadgroup dimensions based on output dimensions
-
-    # Maximum threadgroup size for most Metal GPUs
-    # This could be made more dynamic with Metal API queries if needed
-    MAX_THREADS_PER_GROUP = 1024
-    MAX_THREADS_PER_DIM = 1024
-
-    # Start with a reasonable default size for 2D workloads
-    default_threadgroup = (32, 32, 1)
-
-    try:
-        # Don't create threadgroups larger than the work dimensions
-        max_width = min(MAX_THREADS_PER_DIM, out_w)
-        max_height = min(MAX_THREADS_PER_DIM, out_h)
-
-        # Find largest power of 2 that fits within our dimensions
-        width = 2 ** (max_width.bit_length() - 1)
-        if width > max_width:
-            width = width // 2
-
-        height = 2 ** (max_height.bit_length() - 1)
-        if height > max_height:
-            height = height // 2
-
-        # Ensure we don't exceed maximum threads per threadgroup
-        while width * height > MAX_THREADS_PER_GROUP:
-            # Reduce the larger dimension first
-            if width >= height:
-                width = width // 2
-            else:
-                height = height // 2
-
-        # Ensure minimum size for efficiency
-        width = max(8, width)
-        height = max(8, height)
-
-        return (width, height, 1)
-
-    except Exception:
-        # Return safe defaults if calculation fails
-        return default_threadgroup

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

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

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

@@ -1,143 +0,0 @@
-import glob
-import inspect
-import json
-import re
-from dataclasses import dataclass
-from pathlib import Path
-from typing import List, Optional
-
-import mlx.core as mx
-import mlx.nn as nn
-import numpy as np
-from huggingface_hub import snapshot_download
-from transformers import AutoConfig
-
-from .language import LanguageModel, TextConfig
-from .vision import VisionConfig, VisionModel
-
-
-@dataclass
-class ModelConfig:
-    text_config: TextConfig
-    vision_config: VisionConfig
-    model_type: str
-    ignore_index: int = -100
-    vocab_size: int = 128259
-    scale_factor: int = 2
-    media_placeholder_token_id: int = 163606
-    image_token_index: Optional[int] = None
-    eos_token_id: Optional[List[int]] = None
-
-    def __post_init__(self):
-        if self.image_token_index is None:
-            self.image_token_index = self.media_placeholder_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
-            }
-        )
-
-
-class KimiVLMultiModalProjector(nn.Module):
-
-    def __init__(self, config: ModelConfig):
-        super().__init__()
-
-        self.hidden_size = (
-            config.vision_config.hidden_size
-            * config.vision_config.merge_kernel_size[0]
-            * config.vision_config.merge_kernel_size[1]
-        )
-
-        self.pre_norm = nn.LayerNorm(config.vision_config.hidden_size, eps=1e-05)
-        self.linear_1 = nn.Linear(self.hidden_size, self.hidden_size, bias=True)
-        self.act = nn.GELU()
-        self.linear_2 = nn.Linear(
-            self.hidden_size, config.text_config.hidden_size, bias=True
-        )
-
-    def __call__(self, image_features: list[mx.array]) -> mx.array:
-        image_features = mx.concatenate(image_features, axis=0)
-        h = self.pre_norm(image_features).reshape(-1, self.hidden_size)
-        h = self.linear_1(h)
-        h = self.act(h)
-        h = self.linear_2(h)
-        return h
-
-
-class Model(nn.Module):
-    def __init__(self, config: ModelConfig):
-        super().__init__()
-        self.model_type = config.model_type
-        self.config = config
-
-        self.vision_tower = VisionModel(config.vision_config)
-        self.language_model = LanguageModel(config.text_config)
-        self.multi_modal_projector = KimiVLMultiModalProjector(config)
-
-    def get_input_embeddings(
-        self,
-        input_ids: Optional[mx.array] = None,
-        pixel_values: Optional[mx.array] = None,
-        grid_thw: Optional[mx.array] = None,
-    ):
-        if pixel_values is None:
-            return self.language_model.embed_tokens(input_ids)
-
-        inputs_embeds = self.language_model.embed_tokens(input_ids)
-
-        hidden_state = self.vision_tower(
-            pixel_values.transpose(0, 2, 3, 1),
-            output_hidden_states=True,
-            grid_thw=grid_thw,
-        )
-
-        image_features = self.multi_modal_projector(hidden_state)
-
-        final_inputs_embeds = self._prepare_inputs_for_multimodal(
-            image_features, inputs_embeds, input_ids
-        )
-        return final_inputs_embeds
-
-    def _prepare_inputs_for_multimodal(self, image_features, inputs_embeds, input_ids):
-        image_token_index = self.config.image_token_index
-
-        # Positions of <image> tokens in input_ids, assuming batch size is 1
-        image_positions = np.where(input_ids == image_token_index)[1].tolist()
-
-        inputs_embeds[:, image_positions, :] = image_features
-
-        return inputs_embeds
-
-    @property
-    def layers(self):
-        return self.language_model.model.layers
-
-    def __call__(
-        self,
-        input_ids: mx.array,
-        pixel_values: mx.array,
-        cache=None,
-        **kwargs,
-    ):
-        image_grid_thw = kwargs.pop("image_grid_hws", None)
-        video_grid_thw = kwargs.pop("video_grid_hws", None)
-        grid_thw = image_grid_thw if image_grid_thw is not None else video_grid_thw
-        input_embeddings = self.get_input_embeddings(
-            input_ids, pixel_values, grid_thw=grid_thw
-        )
-        logits = self.language_model(
-            inputs=input_ids, cache=cache, inputs_embeds=input_embeddings
-        )
-        return logits
-
-    def sanitize(self, weights):
-        return {
-            k.replace("encoder.", "") if "vision_tower" in k else k: v
-            for k, v in weights.items()
-        }