Core ML Stable Diffusion

Model Conversion:

Project Build:

Target Device Runtime:

Target Device Runtime (With Memory Improvements):

Target Device Hardware Generation:

stabilityai/stable-diffusion-2-1-base (512x512)

stabilityai/stable-diffusion-xl-base-1.0-ios (768x768)

stabilityai/stable-diffusion-xl-base-1.0 (1024x1024)

coremltools-7.0 supports advanced weight compression techniques for pruning, palettization and linear 8-bit quantization. For these techniques, coremltools.optimize.torch.* includes APIs that require fine-tuning to maintain accuracy at higher compression rates whereas coremltools.optimize.coreml.* includes APIs that are applied post-training and are data-free.

We demonstrate how data-free post-training palettization implemented in coremltools.optimize.coreml.palettize_weights enables us to achieve greatly improved performance for Stable Diffusion on mobile devices. This API implements the Fast Exact k-Means algorithm for optimal weight clustering which yields more accurate palettes. Using --quantize-nbits {2,4,6,8} during conversion is going to apply this compression to the unet and text_encoder models.

For best results, we recommend training-time palettization: coremltools.optimize.torch.palettization.DKMPalettizer if fine-tuning your model is feasible. This API implements the Differentiable k-Means (DKM) learned palettization algorithm. In this exercise, we stick to post-training palettization for the sake of simplicity and ease of reproducibility.

The Neural Engine is capable of accelerating models with low-bit palettization: 1, 2, 4, 6 or 8 bits. With iOS 17 and macOS 14, compressed weights for Core ML models can be just-in-time decompressed during runtime (as opposed to ahead-of-time decompression upon load) to match the precision of activation tensors. This yields significant memory savings and enables models to run on devices with smaller RAM (e.g. iPhone 12 Mini). In addition, compressed weights are faster to fetch from memory which reduces the latency of memory bandwidth-bound layers. The just-in-time decompression behavior depends on the compute unit, layer type and hardware generation.

Weight Precision	--compute-unit	stabilityai/stable-diffusion-2-1-base generating "a high quality photo of a surfing dog"
6-bit	cpuAndNeuralEngine
16-bit	cpuAndNeuralEngine
16-bit	cpuAndGPU

Note that there are minor differences across 16-bit (float16) and 6-bit results. These differences are comparable to the differences across float16 and float32 or differences across compute units as exemplified above. We recommend a minimum of 6 bits for palettizing Stable Diffusion. Smaller number of bits (1, 2 and 4) will require either fine-tuning or advanced palettization techniques such as MBP.

Resources:

• Core ML Tools Docs: Optimizing Models
• WWDC23 Session Video: Use Core ML Tools for machine learning model compression

This section describes an advanced compression algorithm called Mixed-Bit Palettization (MBP) built on top of the Post-Training Weight Palettization tools and using the Weights Metadata API from coremltools.

MBP builds a per-layer "palettization recipe" by picking a suitable number of bits among the Neural Engine supported bit-widths of 1, 2, 4, 6 and 8 in order to achieve the minimum average bit-width while maintaining a desired level of signal strength. The signal strength is measured by comparing the compressed model's output to that of the original float16 model. Given the same random seed and text prompts, PSNR between denoised latents is computed. The compression rate will depend on the model version as well as the tolerance for signal loss (drop in PSNR) since this algorithm is adaptive.

3.41-bit	4.50-bit	6.55-bit	16-bit (original)

For example, the original float16 stabilityai/stable-diffusion-xl-base-1.0 model has an ~82 dB signal strength. Naively applying linear 8-bit quantization to the Unet model drops the signal to ~65 dB. Instead, applying MBP yields an average of 2.81-bits quantization while maintaining a signal strength of ~67 dB. This technique generally yields better results compared to using --quantize-nbits during model conversion but requires a "pre-analysis" run that takes up to a few hours on a single GPU (mps or cuda).

Here is the signal strength (PSNR in dB) versus model size reduction (% of float16 size) for stabilityai/stable-diffusion-xl-base-1.0. The {1,2,4,6,8}-bit curves are generated by progresssively palettizing more layers using a palette with fixed number of bits. The layers were ordered in ascending order of their isolated impact to end-to-end signal strength so the cumulative compression's impact is delayed as much as possible. The mixed-bit curve is based on falling back to a higher number of bits as soon as a layer's isolated impact to end-to-end signal integrity drops below a threshold. Note that all curves based on palettization outperform linear 8-bit quantization at the same model size except for 1-bit.

Here are the steps for applying this technique on another model version:

Step 1: Run the pre-analysis script to generate "recipes" with varying signal strength:

python -m python_coreml_stable_diffusion.mixed_bit_compression_pre_analysis --model-version <model-version> -o <output-dir>

For popular base models, you may find the pre-computed pre-analysis results here. Fine-tuned models models are likely to honor the recipes of their corresponding base models but this is untested.

Step 2: The resulting JSON file from Step 1 will list "baselines", e.g.:

{
  "model_version": "stabilityai/stable-diffusion-xl-base-1.0",
  "baselines": {
    "original": 82.2,
    "linear_8bit": 66.025,
    "recipe_6.55_bit_mixedpalette": 79.9,
    "recipe_5.52_bit_mixedpalette": 78.2,
    "recipe_4.89_bit_mixedpalette": 76.8,
    "recipe_4.41_bit_mixedpalette": 75.5,
    "recipe_4.04_bit_mixedpalette": 73.2,
    "recipe_3.67_bit_mixedpalette": 72.2,
    "recipe_3.32_bit_mixedpalette": 71.4,
    "recipe_3.19_bit_mixedpalette": 70.4,
    "recipe_3.08_bit_mixedpalette": 69.6,
    "recipe_2.98_bit_mixedpalette": 68.6,
    "recipe_2.90_bit_mixedpalette": 67.8,
    "recipe_2.83_bit_mixedpalette": 67.0,
    "recipe_2.71_bit_mixedpalette": 66.3
  },
}

Among these baselines, select a recipe based on your desired signal strength. We recommend palettizing to ~4 bits depending on the use case even if the signal integrity for lower bit values are higher than the linear 8-bit quantization baseline.

Finally, apply the selected recipe to the float16 Core ML model as follows:

python -m python_coreml_stable_diffusion.mixed_bit_compression_apply --mlpackage-path <path-to-float16-unet-mlpackage> -o <output-dir> --pre-analysis-json-path <path-to--pre-analysis-json> --selected-recipe <selected-recipe-string-key>

An example <selected-recipe-string-key> would be "recipe_4.50_bit_mixedpalette" which achieves an average of 4.50-bits compression (compressed from ~5.2GB to ~1.46GB for SDXL). Please note that signal strength does not directly map to image-text alignment. Always verify that your MBP-compressed model variant is accurately generating images for your test prompts.

e.g.:

python -m python_coreml_stable_diffusion.torch2coreml --convert-unet --convert-vae-decoder --convert-text-encoder --xl-version --model-version stabilityai/stable-diffusion-xl-base-1.0 --refiner-version stabilityai/stable-diffusion-xl-refiner-1.0 --bundle-resources-for-swift-cli --attention-implementation {ORIGINAL,SPLIT_EINSUM} -o <output-dir>

• --xl-version: Additional argument to pass to the conversion script when specifying an XL model
• --refiner-version: Additional argument to pass to the conversion script when specifying an XL refiner model, required for "Ensemble of Expert Denoisers" inference.
• --attention-implementation: ORIGINAL is recommended for cpuAndGPU for deployment on Mac
• --attention-implementation: SPLIT_EINSUM is recommended for cpuAndNeuralEngine for deployment on iPhone & iPad
• --attention-implementation: SPLIT_EINSUM_V2 is not recommended for Stable Diffusion XL because of prohibitively long compilation time
• Tip: Adding --latent-h 96 --latent-w 96 is recommended for iOS and iPadOS deployment which leads to 768x768 generation as opposed to the default 1024x1024.
• Tip: Due to known float16 overflow issues in the original Stable Diffusion XL VAE, the model conversion script enforces float32 precision. Using a custom VAE version such as madebyollin/sdxl-vae-fp16-fix by @madebyollin via --custom-vae-version madebyollin/sdxl-vae-fp16-fix will restore the default float16 precision for VAE.

swift run StableDiffusionSample <prompt> --resource-path <output-mlpackages-directory/Resources> --output-path <output-dir> --compute-units {cpuAndGPU,cpuAndNeuralEngine} --xl

• Only the base model is required, refiner model is optional and will be used by default if provided in the resource directory
• ControlNet for XL is not yet supported

python -m python_coreml_stable_diffusion.pipeline --prompt <prompt> --compute-unit {CPU_AND_GPU,CPU_AND_NE} -o <output-dir> -i <output-mlpackages-directory/Resources> --model-version stabilityai/stable-diffusion-xl-base-1.0

• refiner model is not yet supported
• ControlNet for XL is not yet supported

Example results using the prompt "a high quality photo of a surfing dog" conditioned on the scribble (leftmost):

ControlNet allows users to condition image generation with Stable Diffusion on signals such as edge maps, depth maps, segmentation maps, scribbles and pose. Thanks to @ryu38's contribution, both the Python CLI and the Swift package support ControlNet models. Please refer to this section for details on setting up Stable Diffusion with ControlNet.

Note that ControlNet is not yet supported for Stable Diffusion XL.

With iOS 17 and macOS 14, NaturalLanguage framework introduced the NLContextualEmbedding which provides Transformer-based textual embeddings for Latin (20 languages), Cyrillic (4 languages) and CJK (3 languages) scripts. The WWDC23 session titled Explore Natural Language multilingual models demonstrated how this powerful new model can be used by developers to train downstream tasks such as multilingual image generation with Stable Diffusion.

The code to reproduce this demo workflow is made available in this repository. There are several ways in which this workflow can be implemented. Here is an example:

Step 1: Curate an image-text dataset with the desired languages.

Step 2: Pre-compute the NLContextualEmbedding values and replace the text strings with these embedding vectors in your dataset.

Step 3: Fine-tune a base model from Hugging Face Hub that is compatible with the StableDiffusionPipeline by using your new dataset and replacing the default text_encoder with your pre-computed NLContextualEmbedding values.

Step 4: In order to be able to swap the text_encoder of a base model without training new layers, the base model's text_encoder.hidden_size must match that of NLContextualEmbedding. If it doesn't, you will need to train a linear projection layer to map between the two dimensionalities. After fine-tuning, this linear layer should be converted to CoreML as follows:

python -m python_coreml_stable_diffusion.multilingual_projection --input-path <path-to-projection-torchscript> --output-dir <output-dir>

The command above will yield a MultilingualTextEncoderProjection.mlmodelc file under --output-dir and this should be colocated with the rest of the Core ML model assets that were generated through --bundle-resources-for-swift-cli.

Step 5: The multilingual system text encoder can now be invoked by setting useMultilingualTextEncoder to true when initializing a pipeline or setting --use-multilingual-text-encoder in the CLI. Note that the model assets are distributed over-the-air so the first invocation will trigger asset downloads which is less than 100MB.

Resources:

• WWDC23 Session Video: Explore Natural Language multilingual models
• NLContextualEmbedding API Documentation

🤗 Hugging Face ran the conversion procedure on the following models and made the Core ML weights publicly available on the Hub. If you would like to convert a version of Stable Diffusion that is not already available on the Hub, please refer to the Converting Models to Core ML.

• 6-bit quantized models (suitable for iOS 17 and macOS 14):

  • CompVis/stable-diffusion-v1-4
  • runwayml/stable-diffusion-v1-5
  • stabilityai/stable-diffusion-2-base
  • stabilityai/stable-diffusion-2-1-base

• Mixed-bit quantized models

• stabilityai/stable-diffusion-xl-base-1.0
• stabilityai/stable-diffusion-xl-base-1.0-ios

• Uncompressed models:
  • CompVis/stable-diffusion-v1-4
  • runwayml/stable-diffusion-v1-5
  • stabilityai/stable-diffusion-2-base
  • stabilityai/stable-diffusion-2-1-base
  • stabilityai/stable-diffusion-xl-base-1.0
  • stabilityai/stable-diffusion-xl-{base+refiner}-1.0

If you want to use any of those models you may download the weights and proceed to generate images with Python or Swift.

There are several variants in each model repository. You may clone the whole repos using git and git lfs to download all variants, or selectively download the ones you need.

To clone the repos using git, please follow this process:

Step 1: Install the git lfs extension for your system.

git lfs stores large files outside the main git repo, and it downloads them from the appropriate server after you clone or checkout. It is available in most package managers, check the installation page for details.

Step 2: Enable git lfs by running this command once:

git lfs install

Step 3: Use git clone to download a copy of the repo that includes all model variants. For Stable Diffusion version 1.4, you'd issue the following command in your terminal:

git clone https://huggingface.co/apple/coreml-stable-diffusion-v1-4

If you prefer to download specific variants instead of cloning the repos, you can use the huggingface_hub Python library. For example, to do generation in Python using the ORIGINAL attention implementation (read this section for details), you could use the following helper code:

from huggingface_hub import snapshot_download
from pathlib import Path

repo_id = "apple/coreml-stable-diffusion-v1-4"
variant = "original/packages"

model_path = Path("./models") / (repo_id.split("/")[-1] + "_" + variant.replace("/", "_"))
snapshot_download(repo_id, allow_patterns=f"{variant}/*", local_dir=model_path, local_dir_use_symlinks=False)
print(f"Model downloaded at {model_path}")

model_path would be the path in your local filesystem where the checkpoint was saved. Please, refer to this post for additional details.

Step 1: Create a Python environment and install dependencies:

conda create -n coreml_stable_diffusion python=3.8 -y
conda activate coreml_stable_diffusion
cd /path/to/cloned/ml-stable-diffusion/repository
pip install -e .

Step 2: Log in to or register for your Hugging Face account, generate a User Access Token and use this token to set up Hugging Face API access by running huggingface-cli login in a Terminal window.

Step 3: Navigate to the version of Stable Diffusion that you would like to use on Hugging Face Hub and accept its Terms of Use. The default model version is CompVis/stable-diffusion-v1-4. The model version may be changed by the user as described in the next step.

Step 4: Execute the following command from the Terminal to generate Core ML model files (.mlpackage)

python -m python_coreml_stable_diffusion.torch2coreml --convert-unet --convert-text-encoder --convert-vae-decoder --convert-safety-checker --model-version <model-version-string-from-hub> -o <output-mlpackages-directory>

WARNING: This command will download several GB worth of PyTorch checkpoints from Hugging Face. Please ensure that you are on Wi-Fi and have enough disk space.

This generally takes 15-20 minutes on an M1 MacBook Pro. Upon successful execution, the 4 neural network models that comprise Stable Diffusion will have been converted from PyTorch to Core ML (.mlpackage) and saved into the specified <output-mlpackages-directory>. Some additional notable arguments:

• --model-version: The model version name as published on the Hugging Face Hub

• --refiner-version: The refiner version name as published on the Hugging Face Hub. This is optional and if specified, this argument will convert and bundle the refiner unet alongside the model unet.

• --bundle-resources-for-swift-cli: Compiles all 4 models and bundles them along with necessary resources for text tokenization into <output-mlpackages-directory>/Resources which should provided as input to the Swift package. This flag is not necessary for the diffusers-based Python pipeline. However using these compiled models in Python will significantly speed up inference.

• --quantize-nbits: Quantizes the weights of unet and text_encoder models down to 2, 4, 6 or 8 bits using a globally optimal k-means clustering algorithm. By default all models are weight-quantized to 16 bits even if this argument is not specified. Please refer to [this section](#compression-6-bits-and-higher for details and further guidance on weight compression.

• --chunk-unet: Splits the Unet model in two approximately equal chunks (each with less than 1GB of weights) for mobile-friendly deployment. This is required for Neural Engine deployment on iOS and iPadOS if weights are not quantized to 6-bits or less (--quantize-nbits {2,4,6}). This is not required for macOS. Swift CLI is able to consume both the chunked and regular versions of the Unet model but prioritizes the former. Note that chunked unet is not compatible with the Python pipeline because Python pipeline is intended for macOS only.

• --attention-implementation: Defaults to SPLIT_EINSUM which is the implementation described in Deploying Transformers on the Apple Neural Engine. --attention-implementation SPLIT_EINSUM_V2 yields 10-30% improvement for mobile devices, still targeting the Neural Engine. --attention-implementation ORIGINAL will switch to an alternative implementation that should be used for CPU or GPU deployment on some Mac devices. Please refer to the Performance Benchmark section for further guidance.

• --check-output-correctness: Compares original PyTorch model's outputs to final Core ML model's outputs. This flag increases RAM consumption significantly so it is recommended only for debugging purposes.

• --convert-controlnet: Converts ControlNet models specified after this option. This can also convert multiple models if you specify like --convert-controlnet lllyasviel/sd-controlnet-mlsd lllyasviel/sd-controlnet-depth.

• --unet-support-controlnet: enables a converted UNet model to receive additional inputs from ControlNet. This is required for generating image with using ControlNet and saved with a different name, *_control-unet.mlpackage, distinct from normal UNet. On the other hand, this UNet model can not work without ControlNet. Please use normal UNet for just txt2img.

• --convert-vae-encoder: not required for text-to-image applications. Required for image-to-image applications in order to map the input image to the latent space.

Run text-to-image generation using the example Python pipeline based on diffusers:

python -m python_coreml_stable_diffusion.pipeline --prompt "a photo of an astronaut riding a horse on mars" -i <core-ml-model-directory> -o </path/to/output/image> --compute-unit ALL --seed 93

Please refer to the help menu for all available arguments: python -m python_coreml_stable_diffusion.pipeline -h. Some notable arguments:

• -i: Should point to the -o directory from Step 4 of Converting Models to Core ML section from above. If you specified --bundle-resources-for-swift-cli during conversion, then use the resulting Resources folder (which holds the compiled .mlmodelc files). The compiled models load much faster after first use.
• --model-version: If you overrode the default model version while converting models to Core ML, you will need to specify the same model version here.
• --compute-unit: Note that the most performant compute unit for this particular implementation may differ across different hardware. CPU_AND_GPU or CPU_AND_NE may be faster than ALL. Please refer to the Performance Benchmark section for further guidance.
• --scheduler: If you would like to experiment with different schedulers, you may specify it here. For available options, please see the help menu. You may also specify a custom number of inference steps by --num-inference-steps which defaults to 50.
• --controlnet: ControlNet models specified with this option are used in image generation. Use this option in the format --controlnet lllyasviel/sd-controlnet-mlsd lllyasviel/sd-controlnet-depth and make sure to use --controlnet-inputs in conjunction.
• --controlnet-inputs: Image inputs corresponding to each ControlNet model. Please provide image paths in same order as models in --controlnet, for example: --controlnet-inputs image_mlsd image_depth.

swift run StableDiffusionSample "a photo of an astronaut riding a horse on mars" --resource-path <output-mlpackages-directory>/Resources/ --seed 93 --output-path </path/to/output/image>

The output will be named based on the prompt and random seed: e.g. </path/to/output/image>/a_photo_of_an_astronaut_riding_a_horse_on_mars.93.final.png

Please use the --help flag to learn about batched generation and more.

import StableDiffusion
...
let pipeline = try StableDiffusionPipeline(resourcesAt: resourceURL)
pipeline.loadResources()
let image = try pipeline.generateImages(prompt: prompt, seed: seed).first

On iOS, the reduceMemory option should be set to true when constructing StableDiffusionPipeline

This Swift package contains two products:

• StableDiffusion library
• StableDiffusionSample command-line tool

Both of these products require the Core ML models and tokenization resources to be supplied. When specifying resources via a directory path that directory must contain the following:

• TextEncoder.mlmodelc or `TextEncoder2.mlmodelc (text embedding model)
• Unet.mlmodelc or UnetChunk1.mlmodelc & UnetChunk2.mlmodelc (denoising autoencoder model)
• VAEDecoder.mlmodelc (image decoder model)
• vocab.json (tokenizer vocabulary file)
• merges.text (merges for byte pair encoding file)

Optionally, for image2image, in-painting, or similar:

• VAEEncoder.mlmodelc (image encoder model)

Optionally, it may also include the safety checker model that some versions of Stable Diffusion include:

• SafetyChecker.mlmodelc

Optionally, for the SDXL refiner:

• UnetRefiner.mlmodelc (refiner unet model)

Optionally, for ControlNet:

• ControlledUNet.mlmodelc or ControlledUnetChunk1.mlmodelc & ControlledUnetChunk2.mlmodelc (enabled to receive ControlNet values)
• controlnet/ (directory containing ControlNet models)
  • LllyasvielSdControlnetMlsd.mlmodelc (for example, from lllyasviel/sd-controlnet-mlsd)
  • LllyasvielSdControlnetDepth.mlmodelc (for example, from lllyasviel/sd-controlnet-depth)
  • Other models you converted

Note that the chunked version of Unet is checked for first. Only if it is not present will the full Unet.mlmodelc be loaded. Chunking is required for iOS and iPadOS and not necessary for macOS.

🤗 Hugging Face created an open-source demo app on top of this library. It's written in native Swift and Swift UI, and runs on macOS, iOS and iPadOS. You can use the code as a starting point for your app, or to see how to integrate this library in your own projects.

Hugging Face has made the app available in the Mac App Store.