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

nexaai/image_gen_impl/mlx_image_gen_impl.py

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