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/rerank/modeling/nexa_jina_rerank.py

@@ -1,330 +0,0 @@
-# Copyright © 2023-2024 Apple Inc.
-
-import math
-from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Union
-
-import mlx.core as mx
-import mlx.nn as nn
-
-from mlx_lm.models.base import (
-    BaseModelArgs,
-    scaled_dot_product_attention,
-)
-
-
-@dataclass
-class ModelArgs(BaseModelArgs):
-    model_type: str = "xlm_roberta"
-    vocab_size: int = 250002
-    hidden_size: int = 768
-    num_hidden_layers: int = 12
-    num_attention_heads: int = 12
-    intermediate_size: int = 3072
-    hidden_act: str = "gelu"
-    hidden_dropout_prob: float = 0.1
-    attention_probs_dropout_prob: float = 0.1
-    max_position_embeddings: int = 1026
-    type_vocab_size: int = 1
-    initializer_range: float = 0.02
-    layer_norm_eps: float = 1e-05
-    pad_token_id: int = 1
-    bos_token_id: int = 0
-    eos_token_id: int = 2
-    position_embedding_type: str = "absolute"
-    use_cache: bool = True
-    classifier_dropout: Optional[float] = None
-    num_labels: int = 1
-
-
-class XLMRobertaEmbeddings(nn.Module):
-    def __init__(self, config: ModelArgs):
-        super().__init__()
-        self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size)
-        self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
-        self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)
-
-    def __call__(
-        self,
-        input_ids: Optional[mx.array] = None,
-        position_ids: Optional[mx.array] = None,
-        token_type_ids: Optional[mx.array] = None,
-    ) -> mx.array:
-        if token_type_ids is None:
-            token_type_ids = mx.zeros_like(input_ids)
-
-        inputs_embeds = self.word_embeddings(input_ids)
-        position_embeddings = self.position_embeddings(position_ids)
-        token_type_embeddings = self.token_type_embeddings(token_type_ids)
-
-        embeddings = inputs_embeds + position_embeddings + token_type_embeddings
-        return embeddings
-
-
-class SelfAttention(nn.Module):
-    def __init__(self, config: ModelArgs):
-        super().__init__()
-        self.num_attention_heads = config.num_attention_heads
-        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
-        self.all_head_size = self.num_attention_heads * self.attention_head_size
-
-    def __call__(
-        self,
-        qkv: mx.array,
-        key_padding_mask: Optional[mx.array] = None,
-    ) -> mx.array:
-        # qkv shape: [batch, seqlen, 3, num_heads, head_dim]
-        batch_size, seqlen = qkv.shape[0], qkv.shape[1]
-        q, k, v = mx.split(qkv, 3, axis=2)  # Each: [batch, seqlen, 1, num_heads, head_dim]
-        q = mx.squeeze(q, axis=2)  # [batch, seqlen, num_heads, head_dim]
-        k = mx.squeeze(k, axis=2)
-        v = mx.squeeze(v, axis=2)
-
-        # Transpose for attention computation: [batch, num_heads, seqlen, head_dim]
-        q = mx.transpose(q, (0, 2, 1, 3))
-        k = mx.transpose(k, (0, 2, 1, 3))
-        v = mx.transpose(v, (0, 2, 1, 3))
-
-        scale = 1.0 / math.sqrt(self.attention_head_size)
-
-        mask = None
-        if key_padding_mask is not None:
-            # key_padding_mask: [batch, seqlen] where True means keep, False means mask
-            # Convert to attention mask: [batch, 1, 1, seqlen]
-            mask = mx.expand_dims(mx.expand_dims(key_padding_mask, axis=1), axis=1)
-            # Use the same dtype as the query tensor to match model dtype
-            target_dtype = q.dtype
-            mask = (1.0 - mask.astype(target_dtype)) * -10000.0
-
-        context = scaled_dot_product_attention(q, k, v, cache=None, scale=scale, mask=mask)
-
-        # Transpose back and reshape: [batch, seqlen, hidden_size]
-        context = mx.transpose(context, (0, 2, 1, 3))
-        new_context_shape = context.shape[:-2] + (self.all_head_size,)
-        context = mx.reshape(context, new_context_shape)
-        return context
-
-
-class MHA(nn.Module):
-    def __init__(self, config: ModelArgs):
-        super().__init__()
-        self.embed_dim = config.hidden_size
-        self.num_heads = config.num_attention_heads
-        self.head_dim = self.embed_dim // self.num_heads
-
-        # QKV projection
-        qkv_dim = self.head_dim * (self.num_heads + 2 * self.num_heads)  # q + k + v
-        self.Wqkv = nn.Linear(self.embed_dim, qkv_dim, bias=True)
-
-        # Self attention
-        self.inner_attn = SelfAttention(config)
-
-        # Output projection
-        self.out_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True)
-
-    def __call__(
-        self,
-        x: mx.array,
-        key_padding_mask: Optional[mx.array] = None,
-    ) -> tuple:
-        residual = x
-        qkv = self.Wqkv(x)
-
-        # Reshape to [batch, seqlen, 3, num_heads, head_dim]
-        batch, seqlen = qkv.shape[0], qkv.shape[1]
-        qkv = mx.reshape(qkv, (batch, seqlen, 3, self.num_heads, self.head_dim))
-
-        context = self.inner_attn(qkv, key_padding_mask=key_padding_mask)
-        out = self.out_proj(context)
-
-        return out, residual
-
-
-class Mlp(nn.Module):
-    def __init__(self, config: ModelArgs):
-        super().__init__()
-        self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size, bias=True)
-        self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size, bias=True)
-
-    def __call__(self, x: mx.array) -> tuple:
-        residual = x
-        y = self.fc1(x)
-        y = nn.gelu(y)
-        y = self.fc2(y)
-        return y, residual
-
-
-class Block(nn.Module):
-    def __init__(self, config: ModelArgs):
-        super().__init__()
-        self.mixer = MHA(config)
-        self.norm1 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
-        self.mlp = Mlp(config)
-        self.norm2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
-
-    def __call__(
-        self,
-        hidden_states: mx.array,
-        mixer_kwargs: Optional[dict] = None,
-    ) -> mx.array:
-        mixer_kwargs = mixer_kwargs or {}
-
-        # Attention block
-        mixer_out, residual = self.mixer(hidden_states, **mixer_kwargs)
-        hidden_states = self.norm1(mixer_out + residual)
-
-        # MLP block
-        mlp_out, residual = self.mlp(hidden_states)
-        hidden_states = self.norm2(mlp_out + residual)
-
-        return hidden_states
-
-
-class XLMRobertaEncoder(nn.Module):
-    def __init__(self, config: ModelArgs):
-        super().__init__()
-        # Create layers list to match torch naming
-        self.layers = [Block(config) for _ in range(config.num_hidden_layers)]
-
-    def __call__(
-        self,
-        hidden_states: mx.array,
-        key_padding_mask: Optional[mx.array] = None,
-    ) -> mx.array:
-        mixer_kwargs = None
-        if key_padding_mask is not None:
-            mixer_kwargs = {"key_padding_mask": key_padding_mask}
-
-        # Access layers from the list
-        for layer_module in self.layers:
-            hidden_states = layer_module(hidden_states, mixer_kwargs=mixer_kwargs)
-
-        return hidden_states
-
-
-class XLMRobertaModel(nn.Module):
-    def __init__(self, config: ModelArgs):
-        super().__init__()
-        self.config = config
-        self.embeddings = XLMRobertaEmbeddings(config)
-        self.emb_ln = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
-        self.encoder = XLMRobertaEncoder(config)
-
-    def __call__(
-        self,
-        input_ids: mx.array,
-        attention_mask: Optional[mx.array] = None,
-        token_type_ids: Optional[mx.array] = None,
-        position_ids: Optional[mx.array] = None,
-    ) -> mx.array:
-        hidden_states = self.embeddings(
-            input_ids=input_ids,
-            position_ids=position_ids,
-            token_type_ids=token_type_ids,
-        )
-
-        hidden_states = self.emb_ln(hidden_states)
-
-        # Convert attention_mask for padding (True=keep, False=mask)
-        key_padding_mask = attention_mask
-
-        sequence_output = self.encoder(hidden_states, key_padding_mask=key_padding_mask)
-
-        return sequence_output
-
-
-class XLMRobertaClassificationHead(nn.Module):
-    def __init__(self, config: ModelArgs):
-        super().__init__()
-        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
-        self.out_proj = nn.Linear(config.hidden_size, config.num_labels)
-
-    def __call__(self, features: mx.array) -> mx.array:
-        x = features[:, 0, :]  # take first token (equivalent to [CLS])
-        x = self.dense(x)
-        x = mx.tanh(x)
-        x = self.out_proj(x)
-        return x
-
-
-class XLMRobertaForSequenceClassification(nn.Module):
-    def __init__(self, config: ModelArgs):
-        super().__init__()
-        self.num_labels = config.num_labels
-        self.config = config
-        self.roberta = XLMRobertaModel(config)
-        self.classifier = XLMRobertaClassificationHead(config)
-
-    def __call__(
-        self,
-        input_ids: mx.array,
-        attention_mask: Optional[mx.array] = None,
-        token_type_ids: Optional[mx.array] = None,
-        position_ids: Optional[mx.array] = None,
-    ) -> mx.array:
-        sequence_output = self.roberta(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            token_type_ids=token_type_ids,
-            position_ids=position_ids,
-        )
-        logits = self.classifier(sequence_output)
-        return logits
-
-    def nexa_forward(
-        self,
-        input_ids: mx.array,
-        attention_mask: mx.array,
-        token_type_ids: mx.array,
-        position_ids: mx.array,
-    ) -> mx.array:
-        return self(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            token_type_ids=token_type_ids,
-            position_ids=position_ids,
-        )
-
-
-class Model(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        self.args = args
-        self.model_type = args.model_type
-        self.model = XLMRobertaForSequenceClassification(args)
-
-    def __call__(
-        self,
-        input_ids: mx.array,
-        attention_mask: Optional[mx.array] = None,
-        token_type_ids: Optional[mx.array] = None,
-        position_ids: Optional[mx.array] = None,
-    ) -> mx.array:
-        return self.model(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            token_type_ids=token_type_ids,
-            position_ids=position_ids,
-        )
-
-    def nexa_forward(
-        self,
-        input_ids: mx.array,
-        attention_mask: mx.array,
-        token_type_ids: mx.array,
-        position_ids: mx.array,
-    ) -> mx.array:
-        return self.model.nexa_forward(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            token_type_ids=token_type_ids,
-            position_ids=position_ids,
-        )
-
-    def sanitize(self, weights):
-        """Remove parameters that don't exist in our model"""
-        return weights
-
-    @property
-    def layers(self):
-        return self.model.roberta.encoder.layers

nexaai/binds/nexa_mlx/py-lib/sd/__init__.py

@@ -1 +0,0 @@
-"""Stable Diffusion MLX interface package"""

nexaai/binds/nexa_mlx/py-lib/sd/interface.py

@@ -1,362 +0,0 @@
-from __future__ import annotations
-
-import time
-from typing import (
-    Any,
-    Callable,
-    List,
-    Optional,
-)
-
-import mlx.core as mx
-import numpy as np
-from PIL import Image as PILImage
-import mlx.nn as nn
-import os
-
-from .modeling import StableDiffusion, StableDiffusionXL
-
-# --------------------------------------------------------------------------------------
-# Core aliases & callback protocols
-# --------------------------------------------------------------------------------------
-
-Path = str
-LogCallback = Callable[[str], None]
-
-# --------------------------------------------------------------------------------------
-# Core module functions
-# --------------------------------------------------------------------------------------
-
-def init() -> None: 
-    """Initialize the stable diffusion module"""
-    pass
-
-def deinit() -> None: 
-    """Deinitialize the stable diffusion module"""
-    pass
-
-def set_log(callback: LogCallback) -> None: 
-    """Set the logging callback"""
-    pass
-
-def log(message: str) -> None: 
-    """Log a message"""
-    print(message)
-
-# --------------------------------------------------------------------------------------
-# Basic data structures
-# --------------------------------------------------------------------------------------
-
-class Image:
-    def __init__(self, data: List[float], width: int, height: int, channels: int) -> None:
-        """Initialize an image with pixel data"""
-        self.data = data
-        self.width = width
-        self.height = height
-        self.channels = channels
-
-    @classmethod
-    def from_numpy(cls, array: np.ndarray) -> 'Image':
-        """Create Image from numpy array (H, W, C)"""
-        height, width, channels = array.shape
-        data = array.flatten().tolist()
-        return cls(data, width, height, channels)
-    
-    @classmethod
-    def from_pil(cls, pil_image: PILImage.Image) -> 'Image':
-        """Create Image from PIL Image"""
-        array = np.array(pil_image).astype(np.float32) / 255.0
-        return cls.from_numpy(array)
-    
-    def to_numpy(self) -> np.ndarray:
-        """Convert to numpy array (H, W, C)"""
-        return np.array(self.data).reshape(self.height, self.width, self.channels)
-    
-    def to_pil(self) -> PILImage.Image:
-        """Convert to PIL Image"""
-        array = (self.to_numpy() * 255).astype(np.uint8)
-        return PILImage.fromarray(array)
-
-class ImageSamplerConfig:
-    def __init__(
-        self,
-        method: str = "ddim",
-        steps: int = 20,
-        guidance_scale: float = 7.5,
-        eta: float = 0.0,
-        seed: int = -1,
-    ) -> None:
-        """Initialize sampler configuration"""
-        self.method = method
-        self.steps = steps
-        self.guidance_scale = guidance_scale
-        self.eta = eta
-        self.seed = seed
-
-class ImageGenerationConfig:
-    def __init__(
-        self,
-        prompts: str | List[str],
-        negative_prompts: str | List[str] | None = None,
-        height: int = 512,
-        width: int = 512,
-        sampler_config: Optional[ImageSamplerConfig] = None,
-        lora_id: int = -1,  # Not used but kept for compatibility
-        init_image: Optional[Image] = None,
-        strength: float = 1.0,
-        n_images: int = 1,
-        n_rows: int = 1,
-        decoding_batch_size: int = 1,
-    ) -> None:
-        """Initialize image generation configuration"""
-        self.prompts = prompts
-        self.negative_prompts = negative_prompts or ""
-        self.height = height
-        self.width = width
-        self.sampler_config = sampler_config or ImageSamplerConfig()
-        self.lora_id = lora_id
-        self.init_image = init_image
-        self.strength = strength
-        self.n_images = n_images
-        self.n_rows = n_rows
-        self.decoding_batch_size = decoding_batch_size
-
-# --------------------------------------------------------------------------------------
-# Helper functions - following txt2img.py pattern
-# --------------------------------------------------------------------------------------
-
-def load_model(model_path: Path, float16: bool = True, quantize: bool = False) -> StableDiffusion:
-    """Load a model from the given path - following txt2img.py pattern"""
-    
-    # Check if it's a local path or HuggingFace repo
-    # If it contains path separators or exists as a file/directory, treat as local
-    is_local_path = ('/' in model_path or '\\' in model_path or os.path.exists(model_path))
-    
-    if is_local_path:
-        # For local paths, determine model type from the path or model files
-        if "xl" in model_path.lower() or "turbo" in model_path.lower():
-            model = StableDiffusionXL(model_path, float16=float16)
-        else:
-            model = StableDiffusion(model_path, float16=float16)
-    else:
-        # For HuggingFace repo names, use the original logic
-        if "xl" in model_path.lower() or "turbo" in model_path.lower():
-            model = StableDiffusionXL(model_path, float16=float16)
-        else:
-            model = StableDiffusion(model_path, float16=float16)
-    
-    # Apply quantization if requested - same as txt2img.py
-    if quantize:
-        if "xl" in model_path.lower() or "turbo" in model_path.lower():
-            nn.quantize(
-                model.text_encoder_1, class_predicate=lambda _, m: isinstance(m, nn.Linear)
-            )
-            nn.quantize(
-                model.text_encoder_2, class_predicate=lambda _, m: isinstance(m, nn.Linear)
-            )
-        else:
-            nn.quantize(
-                model.text_encoder, class_predicate=lambda _, m: isinstance(m, nn.Linear)
-            )
-        nn.quantize(model.unet, group_size=32, bits=8)
-    
-    return model
-
-def _prepare_image_for_sd(image: Image, target_width: int, target_height: int) -> mx.array:
-    """Prepare image for stable diffusion processing - simplified"""
-    # Convert to PIL and resize
-    pil_img = image.to_pil()
-    pil_img = pil_img.resize((target_width, target_height), PILImage.LANCZOS)
-    
-    # Convert to array and normalize to [0,1] range (following txt2img.py pattern)
-    img_array = np.array(pil_img).astype(np.float32)[:, :, :3]  # Ensure RGB
-    img_tensor = mx.array(img_array / 255.0)
-    
-    return img_tensor
-
-# --------------------------------------------------------------------------------------
-# Image generation
-# --------------------------------------------------------------------------------------
-
-class ImageGen:
-    def __init__(
-        self,
-        model_path: Path,
-        scheduler_config_path: Path = "",  # Make optional
-        device: Optional[str] = None,
-        float16: bool = True,
-        quantize: bool = False,
-    ) -> None:
-        """Initialize the image generation model"""
-        self.model_path = model_path
-        self.scheduler_config_path = scheduler_config_path  # Store for compatibility
-        self.float16 = float16
-        self.quantize = quantize
-        self.model = None
-    
-    def destroy(self) -> None:
-        """Clean up resources"""
-        self.model = None
-    
-    def load_model(self, model_path: Path, extra_data: Any = None) -> bool:
-        """Load the model from a file"""
-        try:
-            if os.path.isfile(model_path):
-                model_path = os.path.dirname(model_path)
-            
-            self.model_path = model_path
-            self.model = load_model(model_path, self.float16, self.quantize)
-            self.model.ensure_models_are_loaded()
-            return True
-        except Exception as e:
-            log(f"Failed to load model: {e}")
-            return False
-    
-    def close(self) -> None:
-        """Close the model"""
-        self.destroy()
-    
-    def set_scheduler(self, config: Any) -> None:
-        """Set scheduler configuration (placeholder for compatibility)"""
-        log("Warning: set_scheduler not implemented")
-        pass
-    
-    def set_sampler(self, config: ImageSamplerConfig) -> None:
-        """Set sampler configuration (placeholder for compatibility)"""
-        log("Warning: set_sampler not implemented")
-        pass
-    
-    def reset_sampler(self) -> None:
-        """Reset sampler configuration (placeholder for compatibility)"""
-        log("Warning: reset_sampler not implemented")
-        pass
-    
-    def set_lora(self, lora_id: int) -> None:
-        """Set LoRA (placeholder for compatibility)"""
-        log("Warning: LoRA management not implemented")
-        pass
-    
-    def add_lora(self, lora_path: Path) -> int:
-        """Add LoRA (placeholder for compatibility)"""
-        log("Warning: LoRA management not implemented")
-        return -1
-    
-    def remove_lora(self, lora_id: int) -> None:
-        """Remove LoRA (placeholder for compatibility)"""
-        log("Warning: LoRA management not implemented")
-        pass
-    
-    def list_loras(self) -> List[int]:
-        """List LoRAs (placeholder for compatibility)"""
-        log("Warning: LoRA management not implemented")
-        return []
-    
-    def txt2img(self, prompt: str, config: ImageGenerationConfig, clear_cache: bool = True) -> Image:
-        """Generate an image from a text prompt - following txt2img.py pattern"""
-        if not self.model and not self.load_model(self.model_path):
-            raise RuntimeError("Model not loaded")
-        
-        sampler_config = config.sampler_config
-        
-        # Extract prompts
-        negative_prompt = ""
-        if config.negative_prompts:
-            negative_prompt = config.negative_prompts if isinstance(config.negative_prompts, str) else config.negative_prompts[0]
-        
-        try:
-            # Generate latents - following txt2img.py approach
-            latents_generator = self.model.generate_latents(
-                prompt,
-                n_images=1,
-                num_steps=sampler_config.steps,
-                cfg_weight=sampler_config.guidance_scale,
-                negative_text=negative_prompt,
-                seed=sampler_config.seed if sampler_config.seed >= 0 else None
-            )
-            
-            # Get final latents - following txt2img.py pattern
-            final_latents = None
-            for latents in latents_generator:
-                final_latents = latents
-                mx.eval(final_latents)
-            
-            if final_latents is None:
-                raise RuntimeError("No latents generated")
-            
-            # Decode to image - following txt2img.py pattern
-            decoded_image = self.model.decode(final_latents)
-            mx.eval(decoded_image)
-            
-            # Convert to numpy array - following txt2img.py pattern
-            image_array = np.array(decoded_image.squeeze(0))
-            
-            if clear_cache:
-                mx.clear_cache()
-            
-            return Image.from_numpy(image_array)
-            
-        except Exception as e:
-            log(f"Generation failed: {e}")
-            raise e
-    
-    def img2img(self, init_image: Image, prompt: str, config: ImageGenerationConfig, clear_cache: bool = True) -> Image:
-        """Generate an image from an initial image and a text prompt"""
-        if not self.model and not self.load_model(self.model_path):
-            raise RuntimeError("Model not loaded")
-        
-        sampler_config = config.sampler_config
-        
-        # Extract prompts
-        negative_prompt = ""
-        if config.negative_prompts:
-            negative_prompt = config.negative_prompts if isinstance(config.negative_prompts, str) else config.negative_prompts[0]
-        
-        try:
-            # Prepare image for SD processing
-            img_tensor = _prepare_image_for_sd(init_image, config.width, config.height)
-            
-            # Generate latents from image
-            latents_generator = self.model.generate_latents_from_image(
-                img_tensor,
-                prompt,
-                n_images=1,
-                strength=config.strength,
-                num_steps=sampler_config.steps,
-                cfg_weight=sampler_config.guidance_scale,
-                negative_text=negative_prompt,
-                seed=sampler_config.seed if sampler_config.seed >= 0 else None
-            )
-            
-            # Get final latents
-            final_latents = None
-            for latents in latents_generator:
-                final_latents = latents
-                mx.eval(final_latents)
-            
-            if final_latents is None:
-                raise RuntimeError("No latents generated")
-            
-            # Decode to image
-            decoded_image = self.model.decode(final_latents)
-            mx.eval(decoded_image)
-            
-            # Convert to numpy array
-            image_array = np.array(decoded_image.squeeze(0))
-            
-            if clear_cache:
-                mx.clear_cache()
-            
-            return Image.from_numpy(image_array)
-            
-        except Exception as e:
-            log(f"Generation failed: {e}")
-            raise e
-    
-    def generate(self, config: ImageGenerationConfig) -> Image:
-        """Generate an image from configuration"""
-        if config.init_image:
-            prompt = config.prompts if isinstance(config.prompts, str) else config.prompts[0]
-            return self.img2img(config.init_image, prompt, config)
-        else:
-            prompt = config.prompts if isinstance(config.prompts, str) else config.prompts[0]
-            return self.txt2img(prompt, config)