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/sd/main.py

@@ -1,286 +0,0 @@
-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

@@ -1,306 +0,0 @@
-# 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

@@ -1,116 +0,0 @@
-# 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,
-        )