nexaai 1.0.29__cp310-cp310-macosx_14_0_universal2.whl

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)

nexaai/mlx_backend/image_gen/main.py

@@ -0,0 +1,281 @@
+#!/usr/bin/env python3
+"""
+Command line interface for text-to-image generation using MLX backend.
+"""
+
+import argparse
+import sys
+import os
+from pathlib import Path
+from typing import Optional
+
+# Add the parent directory to the path to import the interface
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from interface import ImageGen, ImageSamplerConfig
+from ml import (
+    ImageGenCreateInput, 
+    ImageGenTxt2ImgInput, 
+    ImageGenerationConfig, 
+    ImageSamplerConfig as MLImageSamplerConfig,
+    SchedulerConfig,
+    ModelConfig
+)
+
+
+def create_default_config() -> ImageGenerationConfig:
+    """Create a default image generation configuration."""
+    sampler_config = MLImageSamplerConfig(
+        method="ddim",
+        steps=4,  # SDXL Turbo optimized
+        guidance_scale=0.0,  # SDXL Turbo works well with no guidance
+        eta=0.0,
+        seed=-1
+    )
+    
+    scheduler_config = SchedulerConfig(
+        type="ddim",
+        num_train_timesteps=1000,
+        steps_offset=0,
+        beta_start=0.00085,
+        beta_end=0.012,
+        beta_schedule="scaled_linear",
+        prediction_type="epsilon",
+        timestep_type="discrete",
+        timestep_spacing="linspace",
+        interpolation_type="linear"
+    )
+    
+    return ImageGenerationConfig(
+        prompts=[""],  # Will be set by user input
+        sampler_config=sampler_config,
+        scheduler_config=scheduler_config,
+        strength=1.0,
+        negative_prompts=None,
+        height=512,
+        width=512
+    )
+
+
+def parse_arguments():
+    """Parse command line arguments."""
+    parser = argparse.ArgumentParser(
+        description="Generate images from text prompts using MLX backend",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  python main.py "a beautiful sunset over mountains"
+  python main.py "a cat sitting on a chair" --output output.png --width 1024 --height 1024
+  python main.py "a futuristic city" --model-path ./models/sdxl-turbo --steps 8 --seed 42
+        """
+    )
+    
+    # Required arguments
+    parser.add_argument(
+        "prompt",
+        help="Text prompt for image generation"
+    )
+    
+    # Optional arguments
+    parser.add_argument(
+        "--output", "-o",
+        type=str,
+        help="Output image path (default: generated_image.png)"
+    )
+    
+    parser.add_argument(
+        "--model-path", "-m",
+        type=str,
+        default="stabilityai/sdxl-turbo",
+        help="Path to the model or HuggingFace model name (default: stabilityai/sdxl-turbo)"
+    )
+    
+    parser.add_argument(
+        "--width", "-w",
+        type=int,
+        default=512,
+        help="Image width (must be divisible by 16, default: 512)"
+    )
+    
+    parser.add_argument(
+        "--height", "-h",
+        type=int,
+        default=512,
+        help="Image height (must be divisible by 16, default: 512)"
+    )
+    
+    parser.add_argument(
+        "--steps", "-s",
+        type=int,
+        default=4,
+        help="Number of denoising steps (default: 4 for SDXL Turbo)"
+    )
+    
+    parser.add_argument(
+        "--guidance-scale", "-g",
+        type=float,
+        default=0.0,
+        help="Guidance scale (default: 0.0 for SDXL Turbo)"
+    )
+    
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=-1,
+        help="Random seed (-1 for random, default: -1)"
+    )
+    
+    parser.add_argument(
+        "--negative-prompt", "-n",
+        type=str,
+        help="Negative prompt to avoid certain elements"
+    )
+    
+    parser.add_argument(
+        "--device-id",
+        type=str,
+        help="Device ID to use (default: auto-detect)"
+    )
+    
+    parser.add_argument(
+        "--verbose", "-v",
+        action="store_true",
+        help="Enable verbose logging"
+    )
+    
+    return parser.parse_args()
+
+
+def validate_arguments(args):
+    """Validate command line arguments."""
+    # Check dimensions are divisible by 16
+    if args.width % 16 != 0:
+        raise ValueError(f"Width must be divisible by 16, got {args.width}")
+    if args.height % 16 != 0:
+        raise ValueError(f"Height must be divisible by 16, got {args.height}")
+    
+    # Check steps is positive
+    if args.steps <= 0:
+        raise ValueError(f"Steps must be positive, got {args.steps}")
+    
+    # Check guidance scale is non-negative
+    if args.guidance_scale < 0:
+        raise ValueError(f"Guidance scale must be non-negative, got {args.guidance_scale}")
+
+
+def main():
+    """Main function for command line interface."""
+    try:
+        # Parse arguments
+        args = parse_arguments()
+        
+        # Validate arguments
+        validate_arguments(args)
+        
+        # Set up output path
+        if args.output:
+            output_path = Path(args.output)
+        else:
+            output_path = Path("generated_image.png")
+        
+        # Ensure output directory exists
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+        
+        if args.verbose:
+            print(f"Initializing image generation...")
+            print(f"Model: {args.model_path}")
+            print(f"Prompt: {args.prompt}")
+            print(f"Output: {output_path}")
+            print(f"Dimensions: {args.width}x{args.height}")
+            print(f"Steps: {args.steps}")
+            print(f"Guidance scale: {args.guidance_scale}")
+            print(f"Seed: {args.seed}")
+            if args.negative_prompt:
+                print(f"Negative prompt: {args.negative_prompt}")
+        
+        # Create model configuration
+        model_config = ModelConfig(
+            name="sdxl-turbo",
+            version="1.0",
+            description="SDXL Turbo model for fast image generation"
+        )
+        
+        # Create image generator
+        create_input = ImageGenCreateInput(
+            model_name="sdxl-turbo",
+            model_path=args.model_path,
+            config=model_config,
+            scheduler_config_path="",  # Not used for SDXL Turbo
+            plugin_id="mlx",
+            device_id=args.device_id
+        )
+        
+        image_gen = ImageGen(create_input)
+        
+        # Create generation configuration
+        sampler_config = MLImageSamplerConfig(
+            method="ddim",
+            steps=args.steps,
+            guidance_scale=args.guidance_scale,
+            eta=0.0,
+            seed=args.seed
+        )
+        
+        scheduler_config = SchedulerConfig(
+            type="ddim",
+            num_train_timesteps=1000,
+            steps_offset=0,
+            beta_start=0.00085,
+            beta_end=0.012,
+            beta_schedule="scaled_linear",
+            prediction_type="epsilon",
+            timestep_type="discrete",
+            timestep_spacing="linspace",
+            interpolation_type="linear"
+        )
+        
+        generation_config = ImageGenerationConfig(
+            prompts=[args.prompt],
+            sampler_config=sampler_config,
+            scheduler_config=scheduler_config,
+            strength=1.0,
+            negative_prompts=[args.negative_prompt] if args.negative_prompt else None,
+            height=args.height,
+            width=args.width
+        )
+        
+        # Create text-to-image input
+        txt2img_input = ImageGenTxt2ImgInput(
+            prompt=args.prompt,
+            config=generation_config,
+            output_path=str(output_path)
+        )
+        
+        if args.verbose:
+            print("Generating image...")
+        
+        # Generate image
+        result = image_gen.txt2img(txt2img_input)
+        
+        if args.verbose:
+            print(f"Image generated successfully!")
+            print(f"Saved to: {result.output_image_path}")
+        else:
+            print(f"Image saved to: {result.output_image_path}")
+        
+        # Clean up
+        image_gen.close()
+        
+    except KeyboardInterrupt:
+        print("\nGeneration cancelled by user.")
+        sys.exit(1)
+    except Exception as e:
+        print(f"Error: {e}", file=sys.stderr)
+        if args.verbose:
+            import traceback
+            traceback.print_exc()
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()