nexaai 1.0.19rc7__cp310-cp310-macosx_14_0_universal2.whl → 1.0.19rc9__cp310-cp310-macosx_14_0_universal2.whl

nexaai/binds/nexa_mlx/py-lib/image_gen/generate_sd.py

@@ -1,244 +0,0 @@
-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/binds/nexa_mlx/py-lib/image_gen/interface.py

@@ -1,82 +0,0 @@
-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/binds/nexa_mlx/py-lib/image_gen/main.py

@@ -1,281 +0,0 @@
-#!/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()