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

nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/florence2/vision.py

@@ -1,591 +0,0 @@
-import inspect
-import math
-from collections import OrderedDict
-from dataclasses import dataclass, field
-from typing import List, Optional, Tuple, Union
-
-import mlx.core as mx
-import mlx.nn as nn
-
-
-@dataclass
-class VisionConfig:
-    """Configuration class for Florence2 Vision model (DaViT)."""
-
-    model_type: str = "davit"
-    in_chans: int = 3
-    num_classes: int = 1000
-    depths: List[int] = field(default_factory=lambda: [1, 1, 9, 1])
-    dim_embed: List[int] = field(default_factory=lambda: [128, 256, 512, 1024])
-    num_heads: List[int] = field(default_factory=lambda: [4, 8, 16, 32])
-    num_groups: List[int] = field(default_factory=lambda: [4, 8, 16, 32])
-    window_size: int = 12
-    mlp_ratio: float = 4.0
-    drop_path_rate: float = 0.1
-    patch_size: List[int] = field(default_factory=lambda: [7, 3, 3, 3])
-    patch_stride: List[int] = field(default_factory=lambda: [4, 2, 2, 2])
-    patch_padding: List[int] = field(default_factory=lambda: [3, 1, 1, 1])
-    patch_prenorm: List[bool] = field(
-        default_factory=lambda: [False, False, False, False]
-    )
-    qkv_bias: bool = True
-    conv_at_attn: bool = True
-    conv_at_ffn: bool = True
-    hidden_size: int = 768
-    image_size: Tuple[int, int] = (768, 768)
-
-    @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 MlpFC(nn.Module):
-    """MLP FC module"""
-
-    def __init__(
-        self,
-        in_features: int,
-        hidden_features: Optional[int] = None,
-        out_features: Optional[int] = None,
-    ):
-        super().__init__()
-        self.fc1 = nn.Linear(in_features, hidden_features)
-        self.fc2 = nn.Linear(hidden_features, out_features)
-        self.gelu = nn.GELU()
-
-    def __call__(self, x):
-        return self.fc2(self.gelu(self.fc1(x)))
-
-
-class Mlp(nn.Module):
-    """MLP module"""
-
-    def __init__(
-        self,
-        in_features: int,
-        hidden_features: Optional[int] = None,
-        out_features: Optional[int] = None,
-    ):
-        super().__init__()
-        out_features = out_features or in_features
-        hidden_features = hidden_features or in_features
-
-        self.net = MlpFC(in_features, hidden_features, out_features)
-
-    def __call__(self, x, size):
-        return self.net(x), size
-
-
-class DepthWiseConv2d(nn.Module):
-    """Depthwise Convolution"""
-
-    def __init__(
-        self,
-        dim_in: int,
-        kernel_size: int,
-        padding: int,
-        stride: int,
-        bias: bool = True,
-    ):
-        super().__init__()
-
-        self.dw = nn.Conv2d(
-            dim_in,
-            dim_in,
-            kernel_size=kernel_size,
-            padding=padding,
-            stride=stride,
-            bias=bias,
-            groups=dim_in,
-        )
-
-    def __call__(self, x, size):
-        B, N, C = x.shape
-        H, W = size
-        assert N == H * W
-
-        x = self.dw(x.reshape(B, H, W, C))
-
-        x = x.transpose(0, 3, 1, 2)
-
-        size = (x.shape[-2], x.shape[-1])
-        x = x.flatten(2).transpose(0, 2, 1)
-        return x, size
-
-
-class ConvEmbed(nn.Module):
-    """Image to Patch Embedding"""
-
-    def __init__(
-        self,
-        patch_size: int = 7,
-        in_chans: int = 3,
-        embed_dim: int = 64,
-        stride: int = 4,
-        padding: int = 2,
-        norm_layer: Optional[nn.Module] = None,
-        pre_norm: bool = True,
-    ):
-        super().__init__()
-        self.proj = nn.Conv2d(
-            in_chans, embed_dim, kernel_size=patch_size, stride=stride, padding=padding
-        )
-
-        if norm_layer and pre_norm:
-            self.norm = norm_layer(in_chans)
-        elif norm_layer:
-            self.norm = norm_layer(embed_dim)
-        else:
-            self.norm = None
-
-        self.pre_norm = pre_norm
-
-    def __call__(self, x, size):
-        H, W = size
-        if len(x.shape) == 3:
-
-            if self.norm and self.pre_norm:
-                x = self.norm(x)
-
-            x = x.reshape(-1, H, W, x.shape[-1])
-        else:
-            x = x.transpose(0, 2, 3, 1)
-
-        x = self.proj(x)
-
-        B, H, W, C = x.shape
-
-        x = x.reshape(B, H * W, C)
-
-        if self.norm and not self.pre_norm:
-            x = self.norm(x)
-
-        return x, (H, W)
-
-
-class ChannelAttention(nn.Module):
-    """Channel Attention module"""
-
-    def __init__(self, dim: int, groups: int = 8, qkv_bias: bool = True):
-        super().__init__()
-        self.groups = groups
-        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
-        self.proj = nn.Linear(dim, dim)
-
-    def __call__(self, x, size):
-        B, N, C = x.shape
-
-        qkv = self.qkv(x).reshape(B, N, 3, self.groups, C // self.groups)
-        qkv = qkv.transpose(2, 0, 3, 1, 4)
-        q, k, v = qkv[0], qkv[1], qkv[2]  # Each has shape (B, groups, N, C//groups)
-
-        q = q * (float(N) ** -0.5)
-
-        # For multi-head attention, we need to keep the groups dimension
-        attention = mx.matmul(q.transpose(0, 1, 3, 2), k)  # (B, groups, N, N)
-        attention = mx.softmax(attention, axis=-1)
-
-        x = mx.matmul(attention, v.transpose(0, 1, 3, 2)).transpose(
-            0, 1, 3, 2
-        )  # (B, groups, N, C//groups)
-        x = x.transpose(0, 2, 1, 3).reshape(B, N, C)
-        x = self.proj(x)
-
-        return x, size
-
-
-def window_partition(x: mx.array, window_size: int):
-    """Partition into non-overlapping windows"""
-    B, H, W, C = x.shape
-    x = mx.reshape(
-        x, (B, H // window_size, window_size, W // window_size, window_size, C)
-    )
-    windows = mx.reshape(
-        mx.transpose(x, (0, 1, 3, 2, 4, 5)), (-1, window_size, window_size, C)
-    )
-    return windows
-
-
-def window_reverse(
-    windows: mx.array, batch_size: int, window_size: int, H: int, W: int
-):
-    """Merge windows back to feature map"""
-    B = batch_size
-    x = mx.reshape(
-        windows, (B, H // window_size, W // window_size, window_size, window_size, -1)
-    )
-    x = mx.reshape(mx.transpose(x, (0, 1, 3, 2, 4, 5)), (B, H, W, -1))
-    return x
-
-
-class WindowAttention(nn.Module):
-    """Window based multi-head self attention module"""
-
-    def __init__(
-        self, dim: int, num_heads: int, window_size: int, qkv_bias: bool = True
-    ):
-        super().__init__()
-        self.dim = dim
-        self.window_size = window_size
-        self.num_heads = num_heads
-        head_dim = dim // num_heads
-        self.scale = float(head_dim) ** -0.5
-
-        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
-        self.proj = nn.Linear(dim, dim)
-
-    def __call__(self, x, size):
-        H, W = size
-        B, L, C = x.shape
-
-        assert L == H * W, f"input feature has wrong size {L} == {H * W}"
-
-        x = mx.reshape(x, (B, H, W, C))
-
-        # Calculate padding
-        pad_l = pad_t = 0
-        pad_r = (self.window_size - W % self.window_size) % self.window_size
-        pad_b = (self.window_size - H % self.window_size) % self.window_size
-
-        # MLX padding
-        x = mx.pad(x, [(0, 0), (pad_t, pad_b), (pad_l, pad_r), (0, 0)])
-
-        _, Hp, Wp, _ = x.shape
-
-        # Window partition
-        x = window_partition(x, self.window_size)
-        x = mx.reshape(x, (-1, self.window_size * self.window_size, C))
-
-        # Multi-head self attention
-        B_, N, C = x.shape
-        qkv = mx.reshape(self.qkv(x), (B_, N, 3, self.num_heads, C // self.num_heads))
-        qkv = mx.transpose(qkv, (2, 0, 3, 1, 4))
-        q, k, v = qkv[0], qkv[1], qkv[2]
-
-        # Scaled dot-product attention
-        q = q * self.scale
-        attn = mx.matmul(q, mx.transpose(k, (0, 1, 3, 2)))
-        attn = mx.softmax(attn, axis=-1)
-
-        x = mx.reshape(mx.transpose(mx.matmul(attn, v), (0, 2, 1, 3)), (B_, N, C))
-        x = self.proj(x)
-
-        # Merge windows
-        x = mx.reshape(x, (-1, self.window_size, self.window_size, C))
-        x = window_reverse(x, B, self.window_size, Hp, Wp)
-
-        if pad_r > 0 or pad_b > 0:
-            x = x[:, :H, :W, :]
-
-        x = mx.reshape(x, (B, H * W, C))
-        return x, size
-
-
-class PreNorm(nn.Module):
-    """Pre-normalization module"""
-
-    def __init__(self, norm, fn, drop_path=None):
-        super().__init__()
-        self.norm = norm
-        self.fn = fn
-        self.drop_path = drop_path
-
-    def __call__(self, x, size):
-        shortcut = x
-        if self.norm is not None:
-            x = self.norm(x)
-        x, size = self.fn(x, size)
-
-        if self.drop_path is not None:
-            x = self.drop_path(x)
-
-        x = shortcut + x
-        return x, size
-
-
-class DropPath(nn.Module):
-    """Drop paths (Stochastic Depth) per sample."""
-
-    def __init__(self, drop_prob: float = 0.0):
-        super().__init__()
-        self.drop_prob = drop_prob
-
-    def __call__(self, x):
-        if self.drop_prob == 0.0 or not self.training:
-            return x
-
-        keep_prob = 1 - self.drop_prob
-        shape = (x.shape[0],) + (1,) * (x.ndim - 1)
-        random_tensor = keep_prob + mx.random.uniform(shape)
-        random_tensor = mx.floor(random_tensor)
-        output = x * random_tensor / keep_prob
-        return output
-
-
-class SpatialBlock(nn.Module):
-    """Spatial attention block"""
-
-    def __init__(
-        self,
-        dim: int,
-        num_heads: int,
-        window_size: int,
-        mlp_ratio: float = 4.0,
-        qkv_bias: bool = True,
-        drop_path_rate: float = 0.0,
-        conv_at_attn: bool = True,
-        conv_at_ffn: bool = True,
-    ):
-        super().__init__()
-
-        drop_path = DropPath(drop_path_rate) if drop_path_rate > 0.0 else None
-
-        self.conv1 = (
-            PreNorm(None, DepthWiseConv2d(dim, 3, 1, 1), None) if conv_at_attn else None
-        )
-
-        self.window_attn = PreNorm(
-            nn.LayerNorm(dim),
-            WindowAttention(dim, num_heads, window_size, qkv_bias=qkv_bias),
-            drop_path,
-        )
-
-        self.conv2 = (
-            PreNorm(None, DepthWiseConv2d(dim, 3, 1, 1), None) if conv_at_ffn else None
-        )
-
-        self.ffn = PreNorm(
-            nn.LayerNorm(dim),
-            Mlp(
-                in_features=dim, hidden_features=int(dim * mlp_ratio), out_features=dim
-            ),
-            drop_path,
-        )
-
-    def __call__(self, x, size):
-        if self.conv1 is not None:
-            x, size = self.conv1(x, size)
-        x, size = self.window_attn(x, size)
-
-        if self.conv2 is not None:
-            x, size = self.conv2(x, size)
-        x, size = self.ffn(x, size)
-        return x, size
-
-
-class ChannelBlock(nn.Module):
-    """Channel attention block"""
-
-    def __init__(
-        self,
-        dim: int,
-        groups: int,
-        mlp_ratio: float = 4.0,
-        qkv_bias: bool = True,
-        drop_path_rate: float = 0.0,
-        conv_at_attn: bool = True,
-        conv_at_ffn: bool = True,
-    ):
-        super().__init__()
-
-        drop_path = DropPath(drop_path_rate) if drop_path_rate > 0.0 else None
-
-        self.conv1 = (
-            PreNorm(None, DepthWiseConv2d(dim, 3, 1, 1), None) if conv_at_attn else None
-        )
-
-        self.channel_attn = PreNorm(
-            nn.LayerNorm(dim),
-            ChannelAttention(dim, groups=groups, qkv_bias=qkv_bias),
-            drop_path,
-        )
-
-        self.conv2 = (
-            PreNorm(None, DepthWiseConv2d(dim, 3, 1, 1), None) if conv_at_ffn else None
-        )
-
-        self.ffn = PreNorm(
-            nn.LayerNorm(dim),
-            Mlp(
-                in_features=dim, hidden_features=int(dim * mlp_ratio), out_features=dim
-            ),
-            drop_path,
-        )
-
-    def __call__(self, x, size):
-        if self.conv1 is not None:
-            x, size = self.conv1(x, size)
-        x, size = self.channel_attn(x, size)
-
-        if self.conv2 is not None:
-            x, size = self.conv2(x, size)
-        x, size = self.ffn(x, size)
-
-        return x, size
-
-
-class Block(nn.Module):
-    def __init__(
-        self,
-        dim: int,
-        num_heads: int,
-        num_groups: int,
-        window_size: int,
-        mlp_ratio: float = 4.0,
-        qkv_bias: bool = True,
-        drop_path_rate: Tuple[float, float] = (0.0, 0.0),
-        conv_at_attn: bool = True,
-        conv_at_ffn: bool = True,
-    ):
-        super().__init__()
-        self.spatial_block = SpatialBlock(
-            dim,
-            num_heads,
-            window_size,
-            drop_path_rate=drop_path_rate[0],
-            qkv_bias=qkv_bias,
-            mlp_ratio=mlp_ratio,
-            conv_at_attn=conv_at_attn,
-            conv_at_ffn=conv_at_ffn,
-        )
-        self.channel_block = ChannelBlock(
-            dim,
-            num_groups,
-            drop_path_rate=drop_path_rate[1],
-            qkv_bias=qkv_bias,
-            mlp_ratio=mlp_ratio,
-            conv_at_attn=conv_at_attn,
-            conv_at_ffn=conv_at_ffn,
-        )
-
-    def __call__(self, x, size):
-        x, size = self.spatial_block(x, size)
-        x, size = self.channel_block(x, size)
-        return x, size
-
-
-class VisionModel(nn.Module):
-    """DaViT: Dual Attention Vision Transformer"""
-
-    def __init__(self, config: VisionConfig):
-        super().__init__()
-
-        self.num_classes = config.num_classes
-        self.model_type = config.model_type
-        self.dim_embed = config.dim_embed
-        self.num_heads = config.num_heads
-        self.num_groups = config.num_groups
-        self.num_stages = len(self.dim_embed)
-        assert self.num_stages == len(self.num_heads) == len(self.num_groups)
-
-        if self.model_type not in ["davit", ""]:
-            raise ValueError(
-                f"Model type {self.model_type} not supported. Currently only 'davit' is supported"
-            )
-
-        # Convert PyTorch's linspace to MLX equivalent
-        total_blocks = sum(config.depths) * 2
-        dpr = [
-            i * config.drop_path_rate / (total_blocks - 1) for i in range(total_blocks)
-        ]
-
-        depth_offset = 0
-        self.convs = []
-        self.blocks = []
-
-        for i in range(self.num_stages):
-
-            conv_embed = ConvEmbed(
-                patch_size=config.patch_size[i],
-                stride=config.patch_stride[i],
-                padding=config.patch_padding[i],
-                in_chans=config.in_chans if i == 0 else self.dim_embed[i - 1],
-                embed_dim=self.dim_embed[i],
-                norm_layer=nn.LayerNorm,
-                pre_norm=config.patch_prenorm[i],
-            )
-            self.convs.append(conv_embed)
-
-            block = []
-            for j in range(config.depths[i]):
-                block.append(
-                    Block(
-                        self.dim_embed[i],
-                        config.num_heads[i],
-                        config.num_groups[i],
-                        config.window_size,
-                        config.mlp_ratio,
-                        config.qkv_bias,
-                        (dpr[depth_offset + j * 2], dpr[depth_offset + j * 2 + 1]),
-                        config.conv_at_attn,
-                        config.conv_at_ffn,
-                    )
-                )
-
-            self.blocks.append(block)
-
-            depth_offset += config.depths[i] * 2
-
-    def __call__(self, x):
-        input_size = x.shape[2:]
-
-        # Process through stages
-        for conv, blks in zip(self.convs, self.blocks):
-            x, input_size = conv(x, input_size)
-            for blk in blks:
-                x, input_size = blk(x, input_size)
-
-        return x
-
-    @staticmethod
-    def sanitize(weights):
-        sanitized_weights = {}
-        for k, v in weights.items():
-            if "position_ids" in k:
-                # Remove unused position_ids
-                continue
-            elif "convs" in k:
-                if "proj.weight" in k:
-                    # PyTorch conv2d weight tensors have shape:
-                    #   [out_channels, in_channels, kH, KW]
-                    # MLX conv2d expects the weight be of shape:
-                    #   [out_channels, kH, KW, in_channels]
-                    if check_array_shape(v):
-                        sanitized_weights[k] = v
-                    else:
-                        sanitized_weights[k] = v.transpose(0, 2, 3, 1)
-                else:
-                    sanitized_weights[k] = v
-            elif "blocks" in k:
-                if "dw.weight" in k:
-                    sanitized_weights[k] = (
-                        v.transpose(0, 2, 3, 1) if v.shape[1] < v.shape[-1] else v
-                    )
-                else:
-                    sanitized_weights[k] = v
-            else:
-                sanitized_weights[k] = v
-
-        return sanitized_weights

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

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