nexaai 1.0.16rc10__cp310-cp310-macosx_14_0_universal2.whl → 1.0.16rc11__cp310-cp310-macosx_14_0_universal2.whl

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()

nexaai/mlx_backend/image_gen/stable_diffusion/__init__.py

@@ -0,0 +1,306 @@
+# Copyright © 2023-2024 Apple Inc.
+
+import time
+from typing import Optional, Tuple
+
+import mlx.core as mx
+
+from .model_io import (
+    _DEFAULT_MODEL,
+    load_autoencoder,
+    load_diffusion_config,
+    load_text_encoder,
+    load_tokenizer,
+    load_unet,
+)
+from .sampler import SimpleEulerAncestralSampler, SimpleEulerSampler
+
+
+class StableDiffusion:
+    def __init__(self, model: str = _DEFAULT_MODEL, float16: bool = False):
+        self.dtype = mx.float16 if float16 else mx.float32
+        self.diffusion_config = load_diffusion_config(model)
+        self.unet = load_unet(model, float16)
+        self.text_encoder = load_text_encoder(model, float16)
+        self.autoencoder = load_autoencoder(model, False)
+        self.sampler = SimpleEulerSampler(self.diffusion_config)
+        self.tokenizer = load_tokenizer(model)
+
+    def ensure_models_are_loaded(self):
+        mx.eval(self.unet.parameters())
+        mx.eval(self.text_encoder.parameters())
+        mx.eval(self.autoencoder.parameters())
+
+    def _tokenize(self, tokenizer, text: str, negative_text: Optional[str] = None):
+        # Tokenize the text
+        tokens = [tokenizer.tokenize(text)]
+        if negative_text is not None:
+            tokens += [tokenizer.tokenize(negative_text)]
+        lengths = [len(t) for t in tokens]
+        N = max(lengths)
+        tokens = [t + [0] * (N - len(t)) for t in tokens]
+        tokens = mx.array(tokens)
+
+        return tokens
+
+    def _get_text_conditioning(
+        self,
+        text: str,
+        n_images: int = 1,
+        cfg_weight: float = 7.5,
+        negative_text: str = "",
+    ):
+        # Tokenize the text
+        tokens = self._tokenize(
+            self.tokenizer, text, (negative_text if cfg_weight > 1 else None)
+        )
+
+        # Compute the features
+        conditioning = self.text_encoder(tokens).last_hidden_state
+
+        # Repeat the conditioning for each of the generated images
+        if n_images > 1:
+            conditioning = mx.repeat(conditioning, n_images, axis=0)
+
+        return conditioning
+
+    def _denoising_step(
+        self, x_t, t, t_prev, conditioning, cfg_weight: float = 7.5, text_time=None
+    ):
+        x_t_unet = mx.concatenate([x_t] * 2, axis=0) if cfg_weight > 1 else x_t
+        t_unet = mx.broadcast_to(t, [len(x_t_unet)])
+        eps_pred = self.unet(
+            x_t_unet, t_unet, encoder_x=conditioning, text_time=text_time
+        )
+
+        if cfg_weight > 1:
+            eps_text, eps_neg = eps_pred.split(2)
+            eps_pred = eps_neg + cfg_weight * (eps_text - eps_neg)
+
+        x_t_prev = self.sampler.step(eps_pred, x_t, t, t_prev)
+
+        return x_t_prev
+
+    def _denoising_loop(
+        self,
+        x_T,
+        T,
+        conditioning,
+        num_steps: int = 50,
+        cfg_weight: float = 7.5,
+        text_time=None,
+    ):
+        x_t = x_T
+        for t, t_prev in self.sampler.timesteps(
+            num_steps, start_time=T, dtype=self.dtype
+        ):
+            x_t = self._denoising_step(
+                x_t, t, t_prev, conditioning, cfg_weight, text_time
+            )
+            yield x_t
+
+    def generate_latents(
+        self,
+        text: str,
+        n_images: int = 1,
+        num_steps: int = 50,
+        cfg_weight: float = 7.5,
+        negative_text: str = "",
+        latent_size: Tuple[int] = (64, 64),
+        seed=None,
+    ):
+        # Set the PRNG state
+        seed = int(time.time()) if seed is None else seed
+        mx.random.seed(seed)
+
+        # Get the text conditioning
+        conditioning = self._get_text_conditioning(
+            text, n_images, cfg_weight, negative_text
+        )
+
+        # Create the latent variables
+        x_T = self.sampler.sample_prior(
+            (n_images, *latent_size, self.autoencoder.latent_channels), dtype=self.dtype
+        )
+
+        # Perform the denoising loop
+        yield from self._denoising_loop(
+            x_T, self.sampler.max_time, conditioning, num_steps, cfg_weight
+        )
+
+    def generate_latents_from_image(
+        self,
+        image,
+        text: str,
+        n_images: int = 1,
+        strength: float = 0.8,
+        num_steps: int = 50,
+        cfg_weight: float = 7.5,
+        negative_text: str = "",
+        seed=None,
+    ):
+        # Set the PRNG state
+        seed = int(time.time()) if seed is None else seed
+        mx.random.seed(seed)
+
+        # Define the num steps and start step
+        start_step = self.sampler.max_time * strength
+        num_steps = int(num_steps * strength)
+
+        # Get the text conditioning
+        conditioning = self._get_text_conditioning(
+            text, n_images, cfg_weight, negative_text
+        )
+
+        # Get the latents from the input image and add noise according to the
+        # start time.
+        x_0, _ = self.autoencoder.encode(image[None])
+        x_0 = mx.broadcast_to(x_0, (n_images,) + x_0.shape[1:])
+        x_T = self.sampler.add_noise(x_0, mx.array(start_step))
+
+        # Perform the denoising loop
+        yield from self._denoising_loop(
+            x_T, start_step, conditioning, num_steps, cfg_weight
+        )
+
+    def decode(self, x_t):
+        x = self.autoencoder.decode(x_t)
+        x = mx.clip(x / 2 + 0.5, 0, 1)
+        return x
+
+
+class StableDiffusionXL(StableDiffusion):
+    def __init__(self, model: str = _DEFAULT_MODEL, float16: bool = False):
+        super().__init__(model, float16)
+
+        self.sampler = SimpleEulerAncestralSampler(self.diffusion_config)
+
+        self.text_encoder_1 = self.text_encoder
+        self.tokenizer_1 = self.tokenizer
+        del self.tokenizer, self.text_encoder
+
+        self.text_encoder_2 = load_text_encoder(
+            model,
+            float16,
+            model_key="text_encoder_2",
+        )
+        self.tokenizer_2 = load_tokenizer(
+            model,
+            merges_key="tokenizer_2_merges",
+            vocab_key="tokenizer_2_vocab",
+        )
+
+    def ensure_models_are_loaded(self):
+        mx.eval(self.unet.parameters())
+        mx.eval(self.text_encoder_1.parameters())
+        mx.eval(self.text_encoder_2.parameters())
+        mx.eval(self.autoencoder.parameters())
+
+    def _get_text_conditioning(
+        self,
+        text: str,
+        n_images: int = 1,
+        cfg_weight: float = 7.5,
+        negative_text: str = "",
+    ):
+        tokens_1 = self._tokenize(
+            self.tokenizer_1,
+            text,
+            (negative_text if cfg_weight > 1 else None),
+        )
+        tokens_2 = self._tokenize(
+            self.tokenizer_2,
+            text,
+            (negative_text if cfg_weight > 1 else None),
+        )
+
+        conditioning_1 = self.text_encoder_1(tokens_1)
+        conditioning_2 = self.text_encoder_2(tokens_2)
+        conditioning = mx.concatenate(
+            [conditioning_1.hidden_states[-2], conditioning_2.hidden_states[-2]],
+            axis=-1,
+        )
+        pooled_conditioning = conditioning_2.pooled_output
+
+        if n_images > 1:
+            conditioning = mx.repeat(conditioning, n_images, axis=0)
+            pooled_conditioning = mx.repeat(pooled_conditioning, n_images, axis=0)
+
+        return conditioning, pooled_conditioning
+
+    def generate_latents(
+        self,
+        text: str,
+        n_images: int = 1,
+        num_steps: int = 2,
+        cfg_weight: float = 0.0,
+        negative_text: str = "",
+        latent_size: Tuple[int] = (64, 64),
+        seed=None,
+    ):
+        # Set the PRNG state
+        seed = int(time.time()) if seed is None else seed
+        mx.random.seed(seed)
+
+        # Get the text conditioning
+        conditioning, pooled_conditioning = self._get_text_conditioning(
+            text, n_images, cfg_weight, negative_text
+        )
+        text_time = (
+            pooled_conditioning,
+            mx.array([[512, 512, 0, 0, 512, 512.0]] * len(pooled_conditioning)),
+        )
+
+        # Create the latent variables
+        x_T = self.sampler.sample_prior(
+            (n_images, *latent_size, self.autoencoder.latent_channels), dtype=self.dtype
+        )
+
+        # Perform the denoising loop
+        yield from self._denoising_loop(
+            x_T,
+            self.sampler.max_time,
+            conditioning,
+            num_steps,
+            cfg_weight,
+            text_time=text_time,
+        )
+
+    def generate_latents_from_image(
+        self,
+        image,
+        text: str,
+        n_images: int = 1,
+        strength: float = 0.8,
+        num_steps: int = 2,
+        cfg_weight: float = 0.0,
+        negative_text: str = "",
+        seed=None,
+    ):
+        # Set the PRNG state
+        seed = seed or int(time.time())
+        mx.random.seed(seed)
+
+        # Define the num steps and start step
+        start_step = self.sampler.max_time * strength
+        num_steps = int(num_steps * strength)
+
+        # Get the text conditioning
+        conditioning, pooled_conditioning = self._get_text_conditioning(
+            text, n_images, cfg_weight, negative_text
+        )
+        text_time = (
+            pooled_conditioning,
+            mx.array([[512, 512, 0, 0, 512, 512.0]] * len(pooled_conditioning)),
+        )
+
+        # Get the latents from the input image and add noise according to the
+        # start time.
+        x_0, _ = self.autoencoder.encode(image[None])
+        x_0 = mx.broadcast_to(x_0, (n_images,) + x_0.shape[1:])
+        x_T = self.sampler.add_noise(x_0, mx.array(start_step))
+
+        # Perform the denoising loop
+        yield from self._denoising_loop(
+            x_T, start_step, conditioning, num_steps, cfg_weight, text_time=text_time
+        )

nexaai/mlx_backend/image_gen/stable_diffusion/clip.py

@@ -0,0 +1,116 @@
+# Copyright © 2023-2024 Apple Inc.
+
+from dataclasses import dataclass
+from typing import List, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .config import CLIPTextModelConfig
+
+_ACTIVATIONS = {"quick_gelu": nn.gelu_fast_approx, "gelu": nn.gelu}
+
+
+@dataclass
+class CLIPOutput:
+    # The last_hidden_state indexed at the EOS token and possibly projected if
+    # the model has a projection layer
+    pooled_output: Optional[mx.array] = None
+
+    # The full sequence output of the transformer after the final layernorm
+    last_hidden_state: Optional[mx.array] = None
+
+    # A list of hidden states corresponding to the outputs of the transformer layers
+    hidden_states: Optional[List[mx.array]] = None
+
+
+class CLIPEncoderLayer(nn.Module):
+    """The transformer encoder layer from CLIP."""
+
+    def __init__(self, model_dims: int, num_heads: int, activation: str):
+        super().__init__()
+
+        self.layer_norm1 = nn.LayerNorm(model_dims)
+        self.layer_norm2 = nn.LayerNorm(model_dims)
+
+        self.attention = nn.MultiHeadAttention(model_dims, num_heads)
+        # Add biases to the attention projections to match CLIP
+        self.attention.query_proj.bias = mx.zeros(model_dims)
+        self.attention.key_proj.bias = mx.zeros(model_dims)
+        self.attention.value_proj.bias = mx.zeros(model_dims)
+        self.attention.out_proj.bias = mx.zeros(model_dims)
+
+        self.linear1 = nn.Linear(model_dims, 4 * model_dims)
+        self.linear2 = nn.Linear(4 * model_dims, model_dims)
+
+        self.act = _ACTIVATIONS[activation]
+
+    def __call__(self, x, attn_mask=None):
+        y = self.layer_norm1(x)
+        y = self.attention(y, y, y, attn_mask)
+        x = y + x
+
+        y = self.layer_norm2(x)
+        y = self.linear1(y)
+        y = self.act(y)
+        y = self.linear2(y)
+        x = y + x
+
+        return x
+
+
+class CLIPTextModel(nn.Module):
+    """Implements the text encoder transformer from CLIP."""
+
+    def __init__(self, config: CLIPTextModelConfig):
+        super().__init__()
+
+        self.token_embedding = nn.Embedding(config.vocab_size, config.model_dims)
+        self.position_embedding = nn.Embedding(config.max_length, config.model_dims)
+        self.layers = [
+            CLIPEncoderLayer(config.model_dims, config.num_heads, config.hidden_act)
+            for i in range(config.num_layers)
+        ]
+        self.final_layer_norm = nn.LayerNorm(config.model_dims)
+
+        if config.projection_dim is not None:
+            self.text_projection = nn.Linear(
+                config.model_dims, config.projection_dim, bias=False
+            )
+
+    def _get_mask(self, N, dtype):
+        indices = mx.arange(N)
+        mask = indices[:, None] < indices[None]
+        mask = mask.astype(dtype) * (-6e4 if dtype == mx.float16 else -1e9)
+        return mask
+
+    def __call__(self, x):
+        # Extract some shapes
+        B, N = x.shape
+        eos_tokens = x.argmax(-1)
+
+        # Compute the embeddings
+        x = self.token_embedding(x)
+        x = x + self.position_embedding.weight[:N]
+
+        # Compute the features from the transformer
+        mask = self._get_mask(N, x.dtype)
+        hidden_states = []
+        for l in self.layers:
+            x = l(x, mask)
+            hidden_states.append(x)
+
+        # Apply the final layernorm and return
+        x = self.final_layer_norm(x)
+        last_hidden_state = x
+
+        # Select the EOS token
+        pooled_output = x[mx.arange(len(x)), eos_tokens]
+        if "text_projection" in self:
+            pooled_output = self.text_projection(pooled_output)
+
+        return CLIPOutput(
+            pooled_output=pooled_output,
+            last_hidden_state=last_hidden_state,
+            hidden_states=hidden_states,
+        )