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

nexaai/binds/nexa_mlx/py-lib/rerank/modeling/nexa_jina_rerank.py

@@ -0,0 +1,330 @@
+# 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

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

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

@@ -0,0 +1,362 @@
+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)