nexaai 1.0.4rc13__cp310-cp310-win_amd64.whl

nexaai/embedder.py

@@ -0,0 +1,68 @@
+from typing import List, Union
+from dataclasses import dataclass
+from abc import abstractmethod
+import numpy as np
+
+from nexaai.base import BaseModel
+
+
+@dataclass
+class EmbeddingConfig:
+    batch_size: int = 32
+    normalize: bool = True
+    normalize_method: str = "l2"
+
+
+class Embedder(BaseModel):
+    def __init__(self):
+        """
+        Internal initializer
+        """
+        pass
+
+    @classmethod
+    def _load_from(cls, model_path: str, tokenizer_file: str = "tokenizer.json", plugin_id: str = "llama_cpp"):
+        """
+        Load an embedder from model files, routing to appropriate implementation.
+
+        Args:
+            model_path: Path to the model file
+            tokenizer_file: Path to the tokenizer file (default: "tokenizer.json")
+            plugin_id: Plugin ID to use for the model (default: "llama_cpp")
+
+        Returns:
+            Embedder instance
+        """
+        if plugin_id == "mlx":
+            from nexaai.embedder_impl.mlx_embedder_impl import MLXEmbedderImpl
+            return MLXEmbedderImpl._load_from(model_path, tokenizer_file, plugin_id)
+        else:
+            from nexaai.embedder_impl.pybind_embedder_impl import PyBindEmbedderImpl
+            return PyBindEmbedderImpl._load_from(model_path, tokenizer_file, plugin_id)
+
+    @abstractmethod
+    def generate(self, texts: Union[List[str], str] = None, config: EmbeddingConfig = EmbeddingConfig(), input_ids: Union[List[int], List[List[int]]] = None) -> np.ndarray:
+        """
+        Generate embeddings for the given texts or input_ids.
+
+        Args:
+            texts: List of strings or single string to embed
+            input_ids: Pre-tokenized input as:
+                      - Single sequence: list of integers [1, 2, 3, 4]
+                      - Multiple sequences: list of lists [[1, 2, 3], [4, 5, 6]]
+            config: Configuration for embedding generation
+
+        Returns:
+            numpy array of embeddings with shape (num_sequences, embedding_dim)
+        """
+        pass
+
+    @abstractmethod
+    def get_embedding_dim(self) -> int:
+        """
+        Get the embedding dimension of the model
+
+        Returns:
+            The embedding dimension in int
+        """
+        pass

nexaai/embedder_impl/mlx_embedder_impl.py

@@ -0,0 +1,114 @@
+from typing import List, Union
+import numpy as np
+
+from nexaai.embedder import Embedder, EmbeddingConfig
+from nexaai.mlx_backend.embedding.interface import Embedder as MLXEmbedderInterface
+from nexaai.mlx_backend.ml import ModelConfig as MLXModelConfig, SamplerConfig as MLXSamplerConfig, GenerationConfig as MLXGenerationConfig, EmbeddingConfig
+
+
+class MLXEmbedderImpl(Embedder):
+    def __init__(self):
+        """Initialize MLX Embedder implementation."""
+        super().__init__()
+        self._mlx_embedder = None
+
+    @classmethod
+    def _load_from(cls, model_path: str, tokenizer_file: str = "tokenizer.json", plugin_id: str = "mlx"):
+        """
+        Load an embedder from model files using MLX backend.
+
+        Args:
+            model_path: Path to the model file
+            tokenizer_file: Path to the tokenizer file (default: "tokenizer.json")
+            plugin_id: Plugin ID to use for the model (default: "mlx")
+
+        Returns:
+            MLXEmbedderImpl instance
+        """
+        try:
+            # MLX interface is already imported
+            
+            # Create instance and load MLX embedder
+            instance = cls()
+            instance._mlx_embedder = MLXEmbedderInterface(
+                model_path=model_path,
+                tokenizer_path=tokenizer_file
+            )
+            
+            # Load the model
+            success = instance._mlx_embedder.load_model(model_path)
+            if not success:
+                raise RuntimeError("Failed to load MLX embedder model")
+            
+            return instance
+        except Exception as e:
+            raise RuntimeError(f"Failed to load MLX Embedder: {str(e)}")
+
+    def eject(self):
+        """
+        Clean up resources and destroy the embedder
+        """
+        if self._mlx_embedder:
+            self._mlx_embedder.destroy()
+            self._mlx_embedder = None
+
+    def generate(self, texts: Union[List[str], str] = None, config: EmbeddingConfig = EmbeddingConfig(), input_ids: Union[List[int], List[List[int]]] = None) -> np.ndarray:
+        """
+        Generate embeddings for the given texts or input_ids.
+
+        Args:
+            texts: List of strings or single string to embed
+            input_ids: Pre-tokenized input as:
+                      - Single sequence: list of integers [1, 2, 3, 4]
+                      - Multiple sequences: list of lists [[1, 2, 3], [4, 5, 6]]
+            config: Configuration for embedding generation
+
+        Returns:
+            numpy array of embeddings with shape (num_sequences, embedding_dim)
+        """
+        if not self._mlx_embedder:
+            raise RuntimeError("MLX Embedder not loaded")
+        
+        if texts is None and input_ids is None:
+            raise ValueError("Either texts or input_ids must be provided")
+        
+        # MLX embedder currently only supports text input, not pre-tokenized input_ids
+        if input_ids is not None:
+            raise NotImplementedError("MLX embedder does not support input_ids, only text input")
+        
+        try:
+            # Convert single string to list if needed
+            if isinstance(texts, str):
+                texts = [texts]
+            
+            # MLX config classes are already imported
+            
+            # Convert our config to MLX config
+            mlx_config = EmbeddingConfig()
+            mlx_config.batch_size = config.batch_size
+            mlx_config.normalize = config.normalize
+            mlx_config.normalize_method = config.normalize_method
+            
+            # Generate embeddings using MLX
+            embeddings = self._mlx_embedder.embed(texts, mlx_config)
+            
+            # Convert to numpy array
+            return np.array(embeddings, dtype=np.float32)
+            
+        except Exception as e:
+            raise RuntimeError(f"Failed to generate embeddings: {str(e)}")
+
+    def get_embedding_dim(self) -> int:
+        """
+        Get the embedding dimension of the model
+
+        Returns:
+            The embedding dimension in int
+        """
+        if not self._mlx_embedder:
+            raise RuntimeError("MLX Embedder not loaded")
+        
+        try:
+            return self._mlx_embedder.embedding_dim()
+        except Exception as e:
+            raise RuntimeError(f"Failed to get embedding dimension: {str(e)}")

nexaai/embedder_impl/pybind_embedder_impl.py

@@ -0,0 +1,91 @@
+from typing import List, Union
+import numpy as np
+
+from nexaai.embedder import Embedder, EmbeddingConfig
+from nexaai.binds import embedder_bind
+from nexaai.runtime import _ensure_runtime
+
+
+class PyBindEmbedderImpl(Embedder):
+    def __init__(self, _handle_ptr):
+        """
+        Internal initializer
+        """
+        super().__init__()
+        self._handle = _handle_ptr
+
+    @classmethod
+    def _load_from(cls, model_path: str, tokenizer_file: str = "tokenizer.json", plugin_id: str = "llama_cpp"):
+        """
+        Load an embedder from model files
+
+        Args:
+            model_path: Path to the model file
+            tokenizer_file: Path to the tokenizer file (default: "tokenizer.json")
+            plugin_id: Plugin ID to use for the model (default: "llama_cpp")
+
+        Returns:
+            PyBindEmbedderImpl instance
+        """
+        _ensure_runtime()
+        handle = embedder_bind.ml_embedder_create(model_path, tokenizer_file, plugin_id)
+        return cls(handle)
+
+    def eject(self):
+        """
+        Clean up resources and destroy the embedder
+        """
+        # Destructor of the handle will unload the model correctly
+        del self._handle
+        self._handle = None
+
+    def generate(self, texts: Union[List[str], str] = None, config: EmbeddingConfig = EmbeddingConfig(), input_ids: Union[List[int], List[List[int]]] = None) -> np.ndarray:
+        """
+        Generate embeddings for the given texts or input_ids.
+
+        Args:
+            texts: List of strings or single string to embed
+            input_ids: Pre-tokenized input as:
+                      - Single sequence: list of integers [1, 2, 3, 4]
+                      - Multiple sequences: list of lists [[1, 2, 3], [4, 5, 6]]
+            config: Configuration for embedding generation
+
+        Returns:
+            numpy array of embeddings with shape (num_sequences, embedding_dim)
+        """
+        if texts is None and input_ids is None:
+            raise ValueError("Either texts or input_ids must be provided")
+
+        # Create bind config
+        bind_config = embedder_bind.EmbeddingConfig()
+        bind_config.batch_size = config.batch_size
+        bind_config.normalize = config.normalize
+        bind_config.normalize_method = config.normalize_method
+
+        # Convert single string to list if needed
+        if isinstance(texts, str):
+            texts = [texts]
+
+        # Convert input_ids to 2D format if needed
+        processed_input_ids = None
+        if input_ids is not None:
+            if len(input_ids) > 0 and isinstance(input_ids[0], int):
+                # Single sequence: convert [1, 2, 3] to [[1, 2, 3]]
+                processed_input_ids = [input_ids]
+            else:
+                # Multiple sequences: already in correct format [[1, 2], [3, 4]]
+                processed_input_ids = input_ids
+
+        # Pass both parameters, let the ABI handle validation
+        embeddings = embedder_bind.ml_embedder_embed(self._handle, bind_config, texts, processed_input_ids)
+
+        return embeddings
+
+    def get_embedding_dim(self) -> int:
+        """
+        Get the embedding dimension of the model
+
+        Returns:
+            The embedding dimension in int
+        """
+        return embedder_bind.ml_embedder_embedding_dim(self._handle)

nexaai/image_gen.py

@@ -0,0 +1,136 @@
+from typing import List, Optional, Union
+from abc import abstractmethod
+from dataclasses import dataclass
+
+from nexaai.base import BaseModel
+
+
+@dataclass
+class Image:
+    """Image data structure."""
+    data: List[float]  # width × height × channels
+    width: int
+    height: int
+    channels: int  # 3 = RGB, 4 = RGBA
+
+
+@dataclass
+class ImageSamplerConfig:
+    """Configuration for image sampling."""
+    method: str = "ddim"
+    steps: int = 20
+    guidance_scale: float = 7.5
+    eta: float = 0.0
+    seed: int = -1  # –1 for random
+
+
+@dataclass
+class ImageGenerationConfig:
+    """Configuration for image generation."""
+    prompts: Union[str, List[str]]
+    negative_prompts: Optional[Union[str, List[str]]] = None
+    height: int = 512
+    width: int = 512
+    sampler_config: Optional[ImageSamplerConfig] = None
+    lora_id: int = -1  # –1 for none
+    init_image: Optional[Image] = None
+    strength: float = 1.0
+    n_images: int = 1
+    n_rows: int = 1
+    decoding_batch_size: int = 1
+
+
+@dataclass
+class SchedulerConfig:
+    """Configuration for diffusion scheduler."""
+    type: str = "ddim"
+    num_train_timesteps: int = 1000
+    steps_offset: int = 0  # An offset added to the inference steps
+    beta_start: float = 0.00085
+    beta_end: float = 0.012
+    beta_schedule: str = "scaled_linear"
+    prediction_type: str = "epsilon"
+    timestep_type: str = "discrete"
+    timestep_spacing: str = "linspace"
+    interpolation_type: str = "linear"
+    config_path: Optional[str] = None
+
+
+class ImageGen(BaseModel):
+    """Abstract base class for image generation models."""
+
+    def __init__(self):
+        """Initialize base image generation class."""
+        pass
+
+    @classmethod
+    def _load_from(cls,
+                   model_path: str,
+                   scheduler_config_path: str = "",
+                   plugin_id: str = "llama_cpp",
+                   device_id: Optional[str] = None,
+                   float16: bool = True,
+                   quantize: bool = False
+        ) -> 'ImageGen':
+        """Load image generation model from local path, routing to appropriate implementation."""
+        if plugin_id == "mlx":
+            from nexaai.image_gen_impl.mlx_image_gen_impl import MLXImageGenImpl
+            return MLXImageGenImpl._load_from(model_path, scheduler_config_path, plugin_id, device_id, float16, quantize)
+        else:
+            from nexaai.image_gen_impl.pybind_image_gen_impl import PyBindImageGenImpl
+            return PyBindImageGenImpl._load_from(model_path, scheduler_config_path, plugin_id, device_id, float16, quantize)
+
+    @abstractmethod
+    def load_model(self, model_path: str, extra_data: Optional[str] = None) -> bool:
+        """Load model from path."""
+        pass
+
+    @abstractmethod
+    def set_scheduler(self, config: SchedulerConfig) -> None:
+        """Set scheduler configuration."""
+        pass
+
+    @abstractmethod
+    def set_sampler(self, config: ImageSamplerConfig) -> None:
+        """Set sampler configuration."""
+        pass
+
+    @abstractmethod
+    def reset_sampler(self) -> None:
+        """Reset sampler to default configuration."""
+        pass
+
+    @abstractmethod
+    def txt2img(self, prompt: str, config: ImageGenerationConfig) -> Image:
+        """Generate image from text prompt."""
+        pass
+
+    @abstractmethod
+    def img2img(self, init_image: Image, prompt: str, config: ImageGenerationConfig) -> Image:
+        """Generate image from initial image and text prompt."""
+        pass
+
+    @abstractmethod
+    def generate(self, config: ImageGenerationConfig) -> Image:
+        """Generate image from configuration."""
+        pass
+
+    @abstractmethod
+    def set_lora(self, lora_id: int) -> None:
+        """Set active LoRA adapter."""
+        pass
+
+    @abstractmethod
+    def add_lora(self, lora_path: str) -> int:
+        """Add LoRA adapter and return its ID."""
+        pass
+
+    @abstractmethod
+    def remove_lora(self, lora_id: int) -> None:
+        """Remove LoRA adapter."""
+        pass
+
+    @abstractmethod
+    def list_loras(self) -> List[int]:
+        """List available LoRA adapters."""
+        pass