nexaai 1.0.29__cp310-cp310-macosx_14_0_universal2.whl

nexaai/image_gen.py

@@ -0,0 +1,141 @@
+from typing import List, Optional, Union
+from abc import abstractmethod
+from dataclasses import dataclass
+
+from nexaai.base import BaseModel
+from nexaai.common import PluginID
+
+
+@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: Union[PluginID, str] = PluginID.LLAMA_CPP,
+                   device_id: Optional[str] = None,
+                   float16: bool = True,
+                   quantize: bool = False,
+                   **kwargs
+        ) -> 'ImageGen':
+        """Load image generation model from local path, routing to appropriate implementation."""
+        # Check plugin_id value for routing - handle both enum and string
+        plugin_value = plugin_id.value if isinstance(plugin_id, PluginID) else plugin_id
+        
+        if plugin_value == "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

nexaai/image_gen_impl/mlx_image_gen_impl.py

@@ -0,0 +1,292 @@
+# Note: This code is generated by Cursor, not tested yet.
+
+from typing import List, Optional, Union
+import os
+
+from nexaai.common import PluginID
+from nexaai.image_gen import ImageGen, ImageGenerationConfig, ImageSamplerConfig, SchedulerConfig, Image
+from nexaai.mlx_backend.sd.interface import ImageGen as MLXImageGenInterface
+
+
+class MLXImageGenImpl(ImageGen):
+    def __init__(self):
+        """Initialize MLX Image Generation implementation."""
+        super().__init__()
+        self._mlx_image_gen = None
+
+    @classmethod
+    def _load_from(cls,
+                   model_path: str,
+                   scheduler_config_path: str = "",
+                   plugin_id: Union[PluginID, str] = PluginID.MLX,
+                   device_id: Optional[str] = None,
+                   float16: bool = True,
+                   quantize: bool = False
+        ) -> 'MLXImageGenImpl':
+        """Load image generation model from local path using MLX backend."""
+        try:
+            # MLX Image Generation interface is already imported
+            
+            # Create instance and load MLX image generation model
+            instance = cls()
+            instance._mlx_image_gen = MLXImageGenInterface(
+                model_path=model_path,
+                scheduler_config_path=scheduler_config_path,
+                device=device_id,
+                float16=float16,
+                quantize=quantize
+            )
+            
+            return instance
+        except Exception as e:
+            raise RuntimeError(f"Failed to load MLX Image Generation: {str(e)}")
+
+    def eject(self):
+        """Destroy the model and free resources."""
+        if self._mlx_image_gen:
+            self._mlx_image_gen.destroy()
+            self._mlx_image_gen = None
+
+    def load_model(self, model_path: str, extra_data: Optional[str] = None) -> bool:
+        """Load model from path."""
+        if not self._mlx_image_gen:
+            raise RuntimeError("MLX Image Generator not initialized")
+        
+        try:
+            return self._mlx_image_gen.load_model(model_path, extra_data)
+        except Exception as e:
+            raise RuntimeError(f"Failed to load image generation model: {str(e)}")
+
+    def set_scheduler(self, config: SchedulerConfig) -> None:
+        """Set scheduler configuration."""
+        if not self._mlx_image_gen:
+            raise RuntimeError("MLX Image Generator not loaded")
+        
+        try:
+            # Convert our config to MLX format
+            import sys
+            mlx_path = self._get_mlx_path()
+            if mlx_path not in sys.path:
+                sys.path.insert(0, mlx_path)
+            
+            # MLX ImageGen uses a simpler interface, we'll store the config for use in generation
+            self._scheduler_config = config
+            self._mlx_image_gen.set_scheduler(config)
+        except Exception as e:
+            raise RuntimeError(f"Failed to set scheduler: {str(e)}")
+
+    def set_sampler(self, config: ImageSamplerConfig) -> None:
+        """Set sampler configuration."""
+        if not self._mlx_image_gen:
+            raise RuntimeError("MLX Image Generator not loaded")
+        
+        try:
+            # Convert our config to MLX format
+            import sys
+            mlx_path = self._get_mlx_path()
+            if mlx_path not in sys.path:
+                sys.path.insert(0, mlx_path)
+            from sd.interface import ImageSamplerConfig as MLXImageSamplerConfig
+            
+            mlx_config = MLXImageSamplerConfig(
+                method=config.method,
+                steps=config.steps,
+                guidance_scale=config.guidance_scale,
+                eta=config.eta,
+                seed=config.seed
+            )
+            
+            self._mlx_image_gen.set_sampler(mlx_config)
+        except Exception as e:
+            raise RuntimeError(f"Failed to set sampler: {str(e)}")
+
+    def reset_sampler(self) -> None:
+        """Reset sampler to default configuration."""
+        if not self._mlx_image_gen:
+            raise RuntimeError("MLX Image Generator not loaded")
+        
+        try:
+            self._mlx_image_gen.reset_sampler()
+        except Exception as e:
+            raise RuntimeError(f"Failed to reset sampler: {str(e)}")
+
+    def txt2img(self, prompt: str, config: ImageGenerationConfig) -> Image:
+        """Generate image from text prompt."""
+        if not self._mlx_image_gen:
+            raise RuntimeError("MLX Image Generator not loaded")
+        
+        try:
+            # Convert our config to MLX format
+            import sys
+            mlx_path = self._get_mlx_path()
+            if mlx_path not in sys.path:
+                sys.path.insert(0, mlx_path)
+            from sd.interface import ImageGenerationConfig as MLXImageGenerationConfig
+            
+            mlx_config = self._convert_generation_config(config)
+            
+            # Use MLX text-to-image generation
+            result = self._mlx_image_gen.txt2img(prompt, mlx_config)
+            
+            # Convert MLX result to our format
+            return self._convert_image_result(result)
+            
+        except Exception as e:
+            raise RuntimeError(f"Failed to generate image from text: {str(e)}")
+
+    def img2img(self, init_image: Image, prompt: str, config: ImageGenerationConfig) -> Image:
+        """Generate image from initial image and text prompt."""
+        if not self._mlx_image_gen:
+            raise RuntimeError("MLX Image Generator not loaded")
+        
+        try:
+            # Convert our config to MLX format
+            mlx_config = self._convert_generation_config(config)
+            
+            # Convert our image to MLX format
+            import sys
+            mlx_path = self._get_mlx_path()
+            if mlx_path not in sys.path:
+                sys.path.insert(0, mlx_path)
+            from sd.interface import Image as MLXImage
+            
+            mlx_init_image = MLXImage(
+                data=init_image.data,
+                width=init_image.width,
+                height=init_image.height,
+                channels=init_image.channels
+            )
+            
+            # Use MLX image-to-image generation
+            result = self._mlx_image_gen.img2img(mlx_init_image, prompt, mlx_config)
+            
+            # Convert MLX result to our format
+            return self._convert_image_result(result)
+            
+        except Exception as e:
+            raise RuntimeError(f"Failed to generate image from image: {str(e)}")
+
+    def generate(self, config: ImageGenerationConfig) -> Image:
+        """Generate image from configuration."""
+        if not self._mlx_image_gen:
+            raise RuntimeError("MLX Image Generator not loaded")
+        
+        try:
+            # Convert our config to MLX format
+            mlx_config = self._convert_generation_config(config)
+            
+            # Use MLX generation
+            result = self._mlx_image_gen.generate(mlx_config)
+            
+            # Convert MLX result to our format
+            return self._convert_image_result(result)
+            
+        except Exception as e:
+            raise RuntimeError(f"Failed to generate image: {str(e)}")
+
+    def set_lora(self, lora_id: int) -> None:
+        """Set active LoRA adapter."""
+        if not self._mlx_image_gen:
+            raise RuntimeError("MLX Image Generator not loaded")
+        
+        try:
+            self._mlx_image_gen.set_lora(lora_id)
+        except Exception as e:
+            raise RuntimeError(f"Failed to set LoRA: {str(e)}")
+
+    def add_lora(self, lora_path: str) -> int:
+        """Add LoRA adapter and return its ID."""
+        if not self._mlx_image_gen:
+            raise RuntimeError("MLX Image Generator not loaded")
+        
+        try:
+            return self._mlx_image_gen.add_lora(lora_path)
+        except Exception as e:
+            raise RuntimeError(f"Failed to add LoRA: {str(e)}")
+
+    def remove_lora(self, lora_id: int) -> None:
+        """Remove LoRA adapter."""
+        if not self._mlx_image_gen:
+            raise RuntimeError("MLX Image Generator not loaded")
+        
+        try:
+            self._mlx_image_gen.remove_lora(lora_id)
+        except Exception as e:
+            raise RuntimeError(f"Failed to remove LoRA: {str(e)}")
+
+    def list_loras(self) -> List[int]:
+        """List available LoRA adapters."""
+        if not self._mlx_image_gen:
+            raise RuntimeError("MLX Image Generator not loaded")
+        
+        try:
+            return self._mlx_image_gen.list_loras()
+        except Exception as e:
+            raise RuntimeError(f"Failed to list LoRAs: {str(e)}")
+
+    def _convert_generation_config(self, config: ImageGenerationConfig):
+        """Convert our config to MLX format."""
+        import sys
+        mlx_path = self._get_mlx_path()
+        if mlx_path not in sys.path:
+            sys.path.insert(0, mlx_path)
+        from sd.interface import ImageGenerationConfig as MLXImageGenerationConfig, ImageSamplerConfig as MLXImageSamplerConfig
+        
+        mlx_sampler_config = None
+        if config.sampler_config:
+            mlx_sampler_config = MLXImageSamplerConfig(
+                method=config.sampler_config.method,
+                steps=config.sampler_config.steps,
+                guidance_scale=config.sampler_config.guidance_scale,
+                eta=config.sampler_config.eta,
+                seed=config.sampler_config.seed
+            )
+        
+        mlx_init_image = None
+        if config.init_image:
+            from sd.interface import Image as MLXImage
+            mlx_init_image = MLXImage(
+                data=config.init_image.data,
+                width=config.init_image.width,
+                height=config.init_image.height,
+                channels=config.init_image.channels
+            )
+        
+        return MLXImageGenerationConfig(
+            prompts=config.prompts,
+            negative_prompts=config.negative_prompts,
+            height=config.height,
+            width=config.width,
+            sampler_config=mlx_sampler_config,
+            lora_id=config.lora_id,
+            init_image=mlx_init_image,
+            strength=config.strength,
+            n_images=config.n_images,
+            n_rows=config.n_rows,
+            decoding_batch_size=config.decoding_batch_size
+        )
+
+    def _convert_image_result(self, mlx_image):
+        """Convert MLX image to our format."""
+        return Image(
+            data=mlx_image.data,
+            width=mlx_image.width,
+            height=mlx_image.height,
+            channels=mlx_image.channels
+        )
+
+    @staticmethod
+    def _get_mlx_path():
+        """Get the MLX py-lib path."""
+        from pathlib import Path
+        current_dir = Path(__file__).parent.parent.parent.parent  # Go back to nexa root
+        possible_paths = [
+            current_dir / "plugins" / "mlx" / "py-lib",
+            Path(os.environ.get("NEXA_MLX_PATH", "")) if os.environ.get("NEXA_MLX_PATH") else None,
+        ]
+        
+        for path in possible_paths:
+            if path and path.exists() and (path / "sd" / "interface.py").exists():
+                return str(path.absolute())
+        
+        raise ImportError("Could not find MLX py-lib directory")

nexaai/image_gen_impl/pybind_image_gen_impl.py

@@ -0,0 +1,85 @@
+from typing import List, Optional, Union
+
+from nexaai.common import PluginID
+from nexaai.image_gen import ImageGen, ImageGenerationConfig, ImageSamplerConfig, SchedulerConfig, Image
+
+
+class PyBindImageGenImpl(ImageGen):
+    def __init__(self):
+        """Initialize PyBind Image Generation implementation."""
+        super().__init__()
+        # TODO: Add PyBind-specific initialization
+
+    @classmethod
+    def _load_from(cls,
+                   model_path: str,
+                   scheduler_config_path: str = "",
+                   plugin_id: Union[PluginID, str] = PluginID.LLAMA_CPP,
+                   device_id: Optional[str] = None,
+                   float16: bool = True,
+                   quantize: bool = False
+        ) -> 'PyBindImageGenImpl':
+        """Load image generation model from local path using PyBind backend."""
+        # TODO: Implement PyBind image generation loading
+        instance = cls()
+        return instance
+
+    def eject(self):
+        """Destroy the model and free resources."""
+        # TODO: Implement PyBind image generation cleanup
+        pass
+
+    def load_model(self, model_path: str, extra_data: Optional[str] = None) -> bool:
+        """Load model from path."""
+        # TODO: Implement PyBind image generation model loading
+        raise NotImplementedError("PyBind image generation model loading not yet implemented")
+
+    def set_scheduler(self, config: SchedulerConfig) -> None:
+        """Set scheduler configuration."""
+        # TODO: Implement PyBind scheduler setting
+        raise NotImplementedError("PyBind scheduler setting not yet implemented")
+
+    def set_sampler(self, config: ImageSamplerConfig) -> None:
+        """Set sampler configuration."""
+        # TODO: Implement PyBind sampler setting
+        raise NotImplementedError("PyBind sampler setting not yet implemented")
+
+    def reset_sampler(self) -> None:
+        """Reset sampler to default configuration."""
+        # TODO: Implement PyBind sampler reset
+        raise NotImplementedError("PyBind sampler reset not yet implemented")
+
+    def txt2img(self, prompt: str, config: ImageGenerationConfig) -> Image:
+        """Generate image from text prompt."""
+        # TODO: Implement PyBind text-to-image
+        raise NotImplementedError("PyBind text-to-image not yet implemented")
+
+    def img2img(self, init_image: Image, prompt: str, config: ImageGenerationConfig) -> Image:
+        """Generate image from initial image and text prompt."""
+        # TODO: Implement PyBind image-to-image
+        raise NotImplementedError("PyBind image-to-image not yet implemented")
+
+    def generate(self, config: ImageGenerationConfig) -> Image:
+        """Generate image from configuration."""
+        # TODO: Implement PyBind image generation
+        raise NotImplementedError("PyBind image generation not yet implemented")
+
+    def set_lora(self, lora_id: int) -> None:
+        """Set active LoRA adapter."""
+        # TODO: Implement PyBind LoRA setting
+        raise NotImplementedError("PyBind LoRA setting not yet implemented")
+
+    def add_lora(self, lora_path: str) -> int:
+        """Add LoRA adapter and return its ID."""
+        # TODO: Implement PyBind LoRA addition
+        raise NotImplementedError("PyBind LoRA addition not yet implemented")
+
+    def remove_lora(self, lora_id: int) -> None:
+        """Remove LoRA adapter."""
+        # TODO: Implement PyBind LoRA removal
+        raise NotImplementedError("PyBind LoRA removal not yet implemented")
+
+    def list_loras(self) -> List[int]:
+        """List available LoRA adapters."""
+        # TODO: Implement PyBind LoRA listing
+        raise NotImplementedError("PyBind LoRA listing not yet implemented")

nexaai/llm.py

@@ -0,0 +1,98 @@
+from typing import Generator, Optional, Union
+from abc import abstractmethod
+import queue
+import threading
+
+from nexaai.common import ModelConfig, GenerationConfig, ChatMessage, PluginID
+from nexaai.base import BaseModel, ProfilingData
+
+class LLM(BaseModel):
+    def __init__(self, m_cfg: ModelConfig = ModelConfig()):
+        """Initialize base LLM class."""
+        self._m_cfg = m_cfg
+        self._cancel_event = threading.Event()  # New attribute to control cancellation
+
+    @classmethod
+    def _load_from(cls,
+                   local_path: str,
+                   model_name: Optional[str] = None,
+                   tokenizer_path: Optional[str] = None,
+                   m_cfg: ModelConfig = ModelConfig(),
+                   plugin_id: Union[PluginID, str] = PluginID.LLAMA_CPP,
+                   device_id: Optional[str] = None,
+                   **kwargs
+        ) -> 'LLM':
+        """Load model from local path, routing to appropriate implementation."""
+        # Check plugin_id value for routing - handle both enum and string
+        plugin_value = plugin_id.value if isinstance(plugin_id, PluginID) else plugin_id
+        
+        if plugin_value == "mlx":
+            from nexaai.llm_impl.mlx_llm_impl import MLXLLMImpl
+            return MLXLLMImpl._load_from(local_path, model_name, tokenizer_path, m_cfg, plugin_id, device_id)
+        else:
+            from nexaai.llm_impl.pybind_llm_impl import PyBindLLMImpl
+            return PyBindLLMImpl._load_from(local_path, model_name, tokenizer_path, m_cfg, plugin_id, device_id)
+
+    def cancel_generation(self):
+        """Signal to cancel any ongoing stream generation."""
+        self._cancel_event.set()
+
+    def reset_cancel(self):
+        """Reset the cancel event. Call before starting a new generation if needed."""
+        self._cancel_event.clear()
+
+    @abstractmethod
+    def apply_chat_template(self, messages: list[ChatMessage], tools: Optional[str] = None, enable_thinking: bool = True, add_generation_prompt: bool = True) -> str:
+        """Apply the chat template to messages."""
+        pass
+
+    @abstractmethod
+    def generate_stream(self, prompt: str, g_cfg: GenerationConfig = GenerationConfig()) -> Generator[str, None, None]:
+        """Generate text with streaming."""
+        pass
+
+    @abstractmethod
+    def generate(self, prompt: str, g_cfg: GenerationConfig = GenerationConfig()) -> str:
+        """
+        Generate text without streaming.
+
+        Args:
+            prompt (str): The prompt to generate text from. For chat models, this is the chat messages after chat template is applied.
+            g_cfg (GenerationConfig): Generation configuration.
+
+        Returns:
+            str: The generated text.
+        """
+        pass
+
+    def get_profiling_data(self) -> Optional[ProfilingData]:
+        """Get profiling data from the last generation."""  
+        pass
+
+    @abstractmethod
+    def save_kv_cache(self, path: str):
+        """
+        Save the key-value cache to the file.
+
+        Args:
+            path (str): The path to the file.
+        """
+        pass
+
+    @abstractmethod
+    def load_kv_cache(self, path: str):
+        """
+        Load the key-value cache from the file.
+
+        Args:
+            path (str): The path to the file.
+        """
+        pass
+
+    @abstractmethod
+    def reset(self):
+        """
+        Reset the LLM model context and KV cache. If not reset, the model will skip the number of evaluated tokens and treat tokens after those as the new incremental tokens.
+        If your past chat history changed, or you are starting a new chat, you should always reset the model before running generate.
+        """
+        pass