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

nexaai/binds/nexa_mlx/py-lib/sd/main.py

@@ -0,0 +1,286 @@
+from interface import ImageGen, ImageGenerationConfig, ImageSamplerConfig, Image
+import numpy as np
+from PIL import Image as PILImage
+import mlx.core as mx
+
+
+def test_txt2image(
+    prompt="A photo of an astronaut riding a horse on Mars.",
+    model="sdxl",
+    local_model_path="",
+    n_images=1,
+    steps=None,
+    cfg=None,
+    negative_prompt="",
+    n_rows=1,
+    decoding_batch_size=1,
+    float16=True,
+    quantize=False,
+    preload_models=False,
+    output="out_txt2img.png",
+    seed=None,
+    verbose=False,
+    width=512,
+    height=512,
+):
+    """Generate images from text prompt using high-level interface"""
+
+    # Determine model path based on model type
+    if model == "sdxl":
+        model_path = local_model_path or "stabilityai/sdxl-turbo"
+        default_cfg = 0.0
+        default_steps = 2
+    else:
+        model_path = local_model_path or "stabilityai/stable-diffusion-2-1-base"
+        default_cfg = 7.5
+        default_steps = 50
+
+    # Use provided values or defaults
+    cfg = cfg or default_cfg
+    steps = steps or default_steps
+
+    # Create ImageGen instance with proper parameters
+    image_gen = ImageGen(model_path, "", device=None, float16=float16, quantize=quantize)
+
+    # Load the model
+    if not image_gen.load_model(model_path):
+        print(f"Failed to load model: {model_path}")
+        return None
+
+    # Create sampler configuration
+    sampler_config = ImageSamplerConfig(
+        method="ddim",
+        steps=steps,
+        guidance_scale=cfg,
+        seed=seed if seed is not None else -1,
+    )
+
+    # Create generation configuration with all parameters
+    gen_config = ImageGenerationConfig(
+        prompts=prompt,
+        negative_prompts=negative_prompt,
+        height=height,
+        width=width,
+        sampler_config=sampler_config,
+        n_images=n_images,
+        n_rows=n_rows,
+        decoding_batch_size=decoding_batch_size,
+    )
+
+    if verbose:
+        print(f"Generating {n_images} image(s) with prompt: '{prompt}'")
+        print(f"Model: {model_path}, Steps: {steps}, CFG: {cfg}")
+        print(f"Float16: {float16}, Quantize: {quantize}")
+
+    # Generate image using txt2img
+    result_image = image_gen.txt2img(prompt, gen_config)
+
+    # Free memory by deleting model components (following main_duplicate.py pattern)
+    if image_gen.model:
+        if model == "sdxl":
+            if hasattr(image_gen.model, "text_encoder_1"):
+                del image_gen.model.text_encoder_1
+            if hasattr(image_gen.model, "text_encoder_2"):
+                del image_gen.model.text_encoder_2
+        else:
+            if hasattr(image_gen.model, "text_encoder"):
+                del image_gen.model.text_encoder
+
+        if hasattr(image_gen.model, "unet"):
+            del image_gen.model.unet
+        if hasattr(image_gen.model, "sampler"):
+            del image_gen.model.sampler
+
+    # Get peak memory usage
+    peak_mem_unet = mx.metal.get_peak_memory() / 1024**3
+
+    # Convert to PIL and save
+    image_np = result_image.to_numpy()
+    image_pil = PILImage.fromarray((image_np * 255).astype(np.uint8))
+    image_pil.save(output)
+
+    print(f"Text-to-image output saved to: {output}")
+
+    # Get final peak memory usage
+    peak_mem_overall = mx.metal.get_peak_memory() / 1024**3
+
+    # Report memory usage
+    if verbose:
+        print(f"Peak memory used for unet: {peak_mem_unet:.3f}GB")
+        print(f"Peak memory used overall: {peak_mem_overall:.3f}GB")
+
+    # Clean up
+    image_gen.close()
+
+    return output
+
+
+def test_image2image(
+    prompt="A lit fireplace",
+    model="sdxl",
+    strength=0.5,
+    local_model_path="",
+    n_images=1,
+    steps=None,
+    cfg=None,
+    negative_prompt="",
+    n_rows=1,
+    decoding_batch_size=1,
+    quantize=False,
+    float16=True,
+    preload_models=False,
+    init_image_path="out_txt2img.png",
+    output="out_img2img.png",
+    verbose=False,
+    seed=None,
+    width=256,
+    height=256,
+):
+    """Generate images from image and text prompt using high-level interface"""
+
+    # Determine model path based on model type
+    if model == "sdxl":
+        model_path = local_model_path or "stabilityai/sdxl-turbo"
+        default_cfg = 0.0
+        default_steps = 2
+    else:
+        model_path = local_model_path or "stabilityai/stable-diffusion-2-1-base"
+        default_cfg = 7.5
+        default_steps = 50
+
+    # Use provided values or defaults
+    cfg = cfg or default_cfg
+    steps = steps or default_steps
+
+    # Load and process input image
+    try:
+        pil_img = PILImage.open(init_image_path)
+        # Ensure RGB format
+        if pil_img.mode != "RGB":
+            pil_img = pil_img.convert("RGB")
+
+        # Convert to numpy array and then to our Image class
+        img_np = np.array(pil_img).astype(np.float32) / 255.0  # Normalize to [0,1]
+        init_image = Image.from_numpy(img_np)
+
+    except FileNotFoundError:
+        print(f"Error: Image file '{init_image_path}' not found.")
+        return None
+    except Exception as e:
+        print(f"Error loading image: {e}")
+        return None
+
+    # Create ImageGen instance
+    image_gen = ImageGen(model_path, "", device=None)
+
+    # Load the model
+    if not image_gen.load_model(model_path):
+        print(f"Failed to load model: {model_path}")
+        return None
+
+    # Create sampler configuration
+    sampler_config = ImageSamplerConfig(
+        method="ddim",
+        steps=steps,
+        guidance_scale=cfg,
+        seed=seed if seed is not None else -1,
+    )
+
+    # Create generation configuration
+    gen_config = ImageGenerationConfig(
+        prompts=prompt,
+        negative_prompts=negative_prompt,
+        height=height,
+        width=width,
+        sampler_config=sampler_config,
+        init_image=init_image,
+        strength=strength,
+    )
+
+    if verbose:
+        print(f"Generating image with prompt: '{prompt}' and strength: {strength}")
+        print(f"Model: {model_path}, Steps: {steps}, CFG: {cfg}")
+
+    # Generate image using img2img
+    result_image = image_gen.img2img(init_image, prompt, gen_config)
+
+    # Free memory by deleting model components (following main_duplicate.py pattern)
+    if image_gen.model:
+        if model == "sdxl":
+            if hasattr(image_gen.model, "text_encoder_1"):
+                del image_gen.model.text_encoder_1
+            if hasattr(image_gen.model, "text_encoder_2"):
+                del image_gen.model.text_encoder_2
+        else:
+            if hasattr(image_gen.model, "text_encoder"):
+                del image_gen.model.text_encoder
+
+        if hasattr(image_gen.model, "unet"):
+            del image_gen.model.unet
+        if hasattr(image_gen.model, "sampler"):
+            del image_gen.model.sampler
+
+    # Get peak memory usage
+    peak_mem_unet = mx.metal.get_peak_memory() / 1024**3
+
+    # Convert to PIL and save
+    image_np = result_image.to_numpy()
+    image_pil = PILImage.fromarray((image_np * 255).astype(np.uint8))
+    image_pil.save(output)
+
+    print(f"Image-to-image output saved to: {output}")
+
+    # Get final peak memory usage
+    peak_mem_overall = mx.metal.get_peak_memory() / 1024**3
+
+    # Report memory usage
+    if verbose:
+        print(f"Peak memory used for unet: {peak_mem_unet:.3f}GB")
+        print(f"Peak memory used overall: {peak_mem_overall:.3f}GB")
+
+    # Clean up
+    image_gen.close()
+
+    return output
+
+
+if __name__ == "__main__":
+    # Text-to-image parameters
+    txt2img_params = {
+        "prompt": "A photo of an astronaut riding a horse on Mars.",
+        "model": "sdxl",
+        "n_images": 1,
+        "n_rows": 1,
+        "output": "out_txt2img.png",
+        "verbose": True,
+        "width": 256,
+        "height": 256,
+    }
+
+    # Image-to-image parameters
+    img2img_params = {
+        "prompt": "A lit fireplace",
+        "model": "sdxl",
+        "strength": 0.5,
+        "n_images": 1,
+        "n_rows": 1,
+        "init_image_path": "out_txt2img.png",
+        "output": "out_img2img.png",
+        "verbose": True,
+        "width": 512,
+        "height": 512,
+    }
+
+    print("Running text-to-image generation...")
+    generated_image = test_txt2image(**txt2img_params)
+
+    if generated_image:
+        print(f"\nRunning image-to-image generation using: {generated_image}")
+        img2img_params["init_image_path"] = generated_image
+        test_image2image(**img2img_params)
+
+        print(f"\nPipeline complete!")
+        print(f"Text-to-image result: {txt2img_params['output']}")
+        print(f"Image-to-image result: {img2img_params['output']}")
+    else:
+        print("Failed to generate initial image, skipping img2img test")

nexaai/binds/nexa_mlx/py-lib/sd/modeling/__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/binds/nexa_mlx/py-lib/sd/modeling/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,
+        )