nexaai 1.0.21rc5__cp313-cp313-win_arm64.whl → 1.0.21rc16__cp313-cp313-win_arm64.whl

nexaai/image_gen_impl/pybind_image_gen_impl.py

@@ -1,85 +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")
+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

@@ -1,98 +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
+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