nexaai 1.0.16rc10__cp310-cp310-macosx_13_0_x86_64.whl → 1.0.16rc12__cp310-cp310-macosx_13_0_x86_64.whl

nexaai/__init__.py

@@ -21,6 +21,9 @@ except ImportError:
 # Import common configuration classes first (no external dependencies)
 from .common import ModelConfig, GenerationConfig, ChatMessage, SamplerConfig, PluginID
 
+# Import logging functionality
+from .log import set_logger, get_error_message
+
 # Create alias for PluginID to be accessible as plugin_id
 plugin_id = PluginID
 
@@ -45,6 +48,10 @@ __all__ = [
     "EmbeddingConfig",
     "PluginID",
     "plugin_id",
+    
+    # Logging functionality
+    "set_logger",
+    "get_error_message",
 
     "LLM",
     "Embedder",

nexaai/_version.py

@@ -1,4 +1,4 @@
 # This file is generated by CMake from _version.py.in
 # Do not modify this file manually - it will be overwritten
 
-__version__ = "1.0.16-rc10" 
+__version__ = "1.0.16-rc12" 

nexaai/log.py

@@ -0,0 +1,92 @@
+"""
+Logging configuration for NexaAI bridge.
+
+This module provides a minimal API to configure bridge-wide logging
+to route into Python's logging system.
+"""
+
+import logging
+import threading
+from enum import IntEnum
+from typing import Optional
+
+from nexaai.binds import common_bind
+from nexaai.runtime import is_initialized
+
+
+class LogLevel(IntEnum):
+    """Log levels matching ml_LogLevel from ml.h"""
+    TRACE = 0
+    DEBUG = 1
+    INFO = 2
+    WARN = 3
+    ERROR = 4
+
+
+# Module-level state
+_config_lock = threading.Lock()
+_current_logger: Optional[logging.Logger] = None
+
+
+def set_logger(logger: Optional[logging.Logger] = None, *, strict: bool = True) -> None:
+    """
+    Set the process-wide bridge logger.
+    
+    Args:
+        logger: Python logger to receive bridge logs. If None, uses "nexaai.ml" logger.
+        strict: If True, raises if called after runtime initialization.
+                If False, attempts to set anyway (best-effort).
+    
+    Raises:
+        RuntimeError: If strict=True and runtime is already initialized.
+    """
+    global _current_logger
+    
+    with _config_lock:
+        # Check initialization state if strict mode
+        if strict and is_initialized():
+            raise RuntimeError(
+                "Cannot configure logging after runtime initialization. "
+                "Call set_logger() before creating any models, or use strict=False for best-effort."
+            )
+        
+        # Use default logger if none provided
+        if logger is None:
+            logger = logging.getLogger("nexaai.ml")
+        
+        _current_logger = logger
+        
+        # Set the C callback
+        common_bind.ml_set_log(_log_callback)
+
+
+def _log_callback(level: int, message: str) -> None:
+    """Internal callback that forwards bridge logs to Python logger."""
+    if _current_logger is None:
+        return
+    
+    # Map bridge log levels to Python logging levels
+    if level == LogLevel.TRACE or level == LogLevel.DEBUG:
+        _current_logger.debug(message)
+    elif level == LogLevel.INFO:
+        _current_logger.info(message)
+    elif level == LogLevel.WARN:
+        _current_logger.warning(message)
+    elif level == LogLevel.ERROR:
+        _current_logger.error(message)
+    else:
+        # Fallback for unknown levels
+        _current_logger.info(f"[Level {level}] {message}")
+
+
+def get_error_message(error_code: int) -> str:
+    """
+    Get error message string for error code.
+    
+    Args:
+        error_code: ML error code (typically negative)
+        
+    Returns:
+        Human-readable error message
+    """
+    return common_bind.ml_get_error_message(error_code)

nexaai/mlx_backend/image_gen/__init__.py

@@ -0,0 +1 @@
+# Image generation module for MLX backend

nexaai/mlx_backend/image_gen/generate_sd.py

@@ -0,0 +1,244 @@
+from __future__ import annotations
+
+from typing import (
+    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 .stable_diffusion import StableDiffusion, StableDiffusionXL
+
+
+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 = 4,  # SDXL Turbo typically uses fewer steps
+        guidance_scale: float = 0.0,  # SDXL Turbo works well with no guidance
+        eta: float = 0.0,
+        seed: int = -1,
+    ) -> None:
+        """Initialize sampler configuration optimized for SDXL Turbo"""
+        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
+
+
+class ImageGen:
+    def __init__(
+        self,
+        model_path: str,
+        scheduler_config_path: Optional[str] = None,
+        device: Optional[str] = None,
+        float16: bool = True,
+        quantize: bool = False,
+    ) -> None:
+        """Initialize the image generation model for SDXL Turbo"""
+        self.model_path = model_path
+        self.scheduler_config_path = scheduler_config_path
+        self.float16 = float16
+        self.quantize = quantize
+        self.model = None
+
+    @staticmethod
+    def load_model(model_path: str, 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 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:
+            self.model = self.load_model(self.model_path)
+            if not self.model:
+                raise RuntimeError("Model not loaded")
+
+        sampler_config = config.sampler_config
+
+        negative_prompt = ""
+        if config.negative_prompts:
+            negative_prompt = config.negative_prompts if isinstance(
+                config.negative_prompts, str) else config.negative_prompts[0]
+
+        # 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
+        image_array = np.array(decoded_image.squeeze(0))
+
+        if clear_cache:
+            mx.clear_cache()
+
+        return Image.from_numpy(image_array)
+
+    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 using SDXL Turbo"""
+        if not self.model:
+            self.model = self.load_model(self.model_path)
+            if not self.model:
+                raise RuntimeError("Model not loaded")
+
+        sampler_config = config.sampler_config
+
+        negative_prompt = ""
+        if config.negative_prompts:
+            negative_prompt = config.negative_prompts if isinstance(
+                config.negative_prompts, str) else config.negative_prompts[0]
+
+        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)

nexaai/mlx_backend/image_gen/interface.py

@@ -0,0 +1,82 @@
+from __future__ import annotations
+import os
+from typing import Optional
+
+from ml import ImageGenCreateInput, ImageGenerationConfig, ImageGenImg2ImgInput, ImageGenTxt2ImgInput, ImageGenOutput
+from profiling import ProfilingMixin, StopReason
+
+from .generate_sd import ImageGen as SDImageGen, Image, ImageGenerationConfig as SDImageGenerationConfig, ImageSamplerConfig
+
+
+class ImageGen(ProfilingMixin):
+    sd_gen: Optional[SDImageGen] = None
+
+    def __init__(self, input: ImageGenCreateInput):
+        """Initialize the image generation model"""
+        self.sd_gen = SDImageGen(model_path=input.model_path)
+
+    def destroy(self) -> None:
+        """Clean up resources"""
+        self.sd_gen = None
+
+    def txt2img(self, input: ImageGenTxt2ImgInput) -> ImageGenOutput:
+        """Generate an image from a text prompt - public interface"""
+        height = input.config.height
+        width = input.config.width
+        assert height % 16 == 0, f"Height must be divisible by 16 ({height}/16={height/16})"
+        assert width % 16 == 0, f"Width must be divisible by 16 ({width}/16={width/16})"
+
+        internal_config = SDImageGenerationConfig(
+            prompts=input.prompt,
+            negative_prompts=input.config.negative_prompts,
+            height=height,
+            width=width,
+            sampler_config=ImageSamplerConfig(
+                steps=input.config.sampler_config.steps,
+                guidance_scale=input.config.sampler_config.guidance_scale,
+                seed=input.config.sampler_config.seed
+            ),
+            strength=input.config.strength
+        )
+
+        result_image = self.sd_gen.txt2img(input.prompt, internal_config)
+
+        parent_dir = os.path.dirname(input.output_path)
+        if parent_dir:
+            os.makedirs(parent_dir, exist_ok=True)
+        result_image.to_pil().save(input.output_path)
+
+        return ImageGenOutput(output_image_path=input.output_path)
+
+    def img2img(self, input: ImageGenImg2ImgInput) -> ImageGenOutput:
+        """Generate an image from an initial image and a text prompt - public interface"""
+        height = input.config.height
+        width = input.config.width
+        assert height % 16 == 0, f"Height must be divisible by 16 ({height}/16={height/16})"
+        assert width % 16 == 0, f"Width must be divisible by 16 ({width}/16={width/16})"
+
+        init_image = Image.from_pil(input.init_image_path)
+
+        internal_config = SDImageGenerationConfig(
+            prompts=input.prompt,
+            negative_prompts=input.config.negative_prompts,
+            height=height,
+            width=width,
+            sampler_config=ImageSamplerConfig(
+                steps=input.config.sampler_config.steps,
+                guidance_scale=input.config.sampler_config.guidance_scale,
+                seed=input.config.sampler_config.seed
+            ),
+            init_image=init_image,
+            strength=input.config.strength
+        )
+
+        result_image = self.sd_gen.img2img(
+            init_image, input.prompt, internal_config)
+
+        parent_dir = os.path.dirname(input.output_path)
+        if parent_dir:
+            os.makedirs(parent_dir, exist_ok=True)
+        result_image.to_pil().save(input.output_path)
+
+        return ImageGenOutput(output_image_path=input.output_path)