• Stars
    star
    16,002
  • Rank 1,682 (Top 0.04 %)
  • Language
    Python
  • License
    Other
  • Created over 1 year ago
  • Updated about 2 months ago

Reviews

There are no reviews yet. Be the first to send feedback to the community and the maintainers!

Repository Details

Stable Diffusion with Core ML on Apple Silicon

Core ML Stable Diffusion

Run Stable Diffusion on Apple Silicon with Core ML

[Blog Post] [BibTeX]

This repository comprises:

  • python_coreml_stable_diffusion, a Python package for converting PyTorch models to Core ML format and performing image generation with Hugging Face diffusers in Python
  • StableDiffusion, a Swift package that developers can add to their Xcode projects as a dependency to deploy image generation capabilities in their apps. The Swift package relies on the Core ML model files generated by python_coreml_stable_diffusion

If you run into issues during installation or runtime, please refer to the FAQ section. Please refer to the System Requirements section before getting started.

System Requirements

Model Conversion:

macOS Python coremltools
13.1 3.8 7.0

Project Build:

macOS Xcode Swift
13.1 14.3 5.8

Target Device Runtime:

macOS iPadOS, iOS
13.1 16.2

Target Device Runtime (With Memory Improvements):

macOS iPadOS, iOS
14.0 17.0

Target Device Hardware Generation:

Mac iPad iPhone
M1 M1 A14

Performance Benchmarks

stabilityai/stable-diffusion-2-1-base Benchmark:

Device --compute-unit --attention-implementation End-to-End Latency (s) Diffusion Speed (iter/s)
iPhone 12 Mini CPU_AND_NE SPLIT_EINSUM_V2 20 1.3
iPhone 12 Pro Max CPU_AND_NE SPLIT_EINSUM_V2 17 1.4
iPhone 13 CPU_AND_NE SPLIT_EINSUM_V2 15 1.7
iPhone 13 Pro Max CPU_AND_NE SPLIT_EINSUM_V2 12 1.8
iPhone 14 CPU_AND_NE SPLIT_EINSUM_V2 13 1.8
iPhone 14 Pro Max CPU_AND_NE SPLIT_EINSUM_V2 9 2.3
iPad Pro (M1) CPU_AND_NE SPLIT_EINSUM_V2 11 2.1
iPad Pro (M2) CPU_AND_NE SPLIT_EINSUM_V2 8 2.9
Mac Studio (M1 Ultra) CPU_AND_GPU ORIGINAL 4 6.3
Mac Studio (M2 Ultra) CPU_AND_GPU ORIGINAL 3 7.6
Details (Click to expand)
  • This benchmark was conducted by Apple using public beta versions of iOS 17.0, iPadOS 17.0 and macOS 14.0 in June 2023.
  • The performance data was collected by running the StableDiffusion Swift pipeline.
  • Swift code is not fully optimized, introducing up to ~10% overhead unrelated to Core ML model execution.
  • The median latency value across 3 back-to-back end-to-end executions are reported
  • The image generation procedure follows the standard configuration: 20 inference steps, 512x512 output image resolution, 77 text token sequence length, classifier-free guidance (batch size of 2 for unet).
  • The actual prompt length does not impact performance because the Core ML model is converted with a static shape that computes the forward pass for all of the 77 elements (tokenizer.model_max_length) in the text token sequence regardless of the actual length of the input text.
  • Weights are compressed to 6 bit precision. Please refer to this section for details.
  • Activations are in float16 precision for both the GPU and the Neural Engine.
  • For iPhone 12 mini and iPhone 13, we enable the reduceMemory option to load and unload models just-in-time to avoid memory issues. This added up to 2 seconds to the end-to-end latency.
  • In the benchmark table, we report the best performing --compute-unit and --attention-implementation values per device. The former does not modify the Core ML model and can be applied during runtime. The latter modifies the Core ML model. Note that the best performing compute unit is model version and hardware-specific.
  • Note that the performance optimizations in this repository (e.g. --attention-implementation) are generally applicable to Transformers and not customized to Stable Diffusion. Better performance may be observed upon custom kernel tuning. Therefore, these numbers do not represent peak HW capability.
  • Performance may vary across different versions of Stable Diffusion due to architecture changes in the model itself. Each reported number is specific to the model version mentioned in that context.
  • Performance may vary due to factors like increased system load from other applications or suboptimal device thermal state. Given these factors, we do not report sub-second variance in latency.

stabilityai/stable-diffusion-xl-base-1.0 Benchmark:

Device --compute-unit --attention-implementation End-to-End Latency (s) Diffusion Speed (iter/s)
MacBook Pro (M1 Max) CPU_AND_GPU ORIGINAL 46 0.46
MacBook Pro (M2 Max) CPU_AND_GPU ORIGINAL 37 0.57
Mac Studio (M1 Ultra) CPU_AND_GPU ORIGINAL 25 0.89
Mac Studio (M2 Ultra) CPU_AND_GPU ORIGINAL 20 1.11
Details (Click to expand)
  • This benchmark was conducted by Apple and Hugging Face using public beta versions of iOS 17.0, iPadOS 17.0 and macOS 14.0 in July 2023.
  • The performance data was collected by running the StableDiffusion Swift pipeline.
  • The median latency value across 3 back-to-back end-to-end executions are reported
  • The image generation procedure follows the standard configuration: 20 inference steps, 1024x1024 output image resolution, classifier-free guidance (batch size of 2 for unet).
  • Weights and activations are in float16 precision
  • Performance may vary across different versions of Stable Diffusion due to architecture changes in the model itself. Each reported number is specific to the model version mentioned in that context.
  • Performance may vary due to factors like increased system load from other applications or suboptimal device thermal state. Given these factors, we do not report sub-second variance in latency.

Weight Compression

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:

MBP: Post-Training Mixed-Bit Palettization

Details (Click to expand)
3.41-bit 4.50-bit 6.55-bit 16-bit (original)

This is an advanced compression technique that picks a suitable number of bits (among 1, 2, 4, 6 and 8) in order to achieve the desired signal strength as measured by end-to-end PSNR. The compression rate depends on the model version as well as the tolerance for signal loss (drop in PSNR).

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.

Using Stable Diffusion XL

Details (Click to expand)

Model Conversion

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 --bundle-resources-for-swift-cli --attention-implementation ORIGINAL -o <output-dir>
  • --xl-version: Additional argument to pass to the conversion script when specifying an XL model
  • --attention-implementation ORIGINAL (recommended for cpuAndGPU)
  • Due to known float16 overflow issues in the VAE, it runs in float32 precision for now

Inference

e.g.

swift run StableDiffusionSample <prompt> --resource-path <output-mlpackages-directory/Resources> --output-path <output-dir> --compute-units cpuAndGPU --xl
  • Only --compute-units cpuAndGPU is supported for now
  • Only the base model is supported, refiner model is not yet supported
  • ControlNet for XL is not yet supported

Using ControlNet

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.

Using the System Multilingual Text Encoder

Details (Click to expand)

With iOS 17 and macOS 14, NaturalLanguage framework introduced the NLContextualEmbedding which provides Transformer-based textual embeddings for Latin (20 languages), Cyrillic (4 langauges) 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:

Using Ready-made Core ML Models from Hugging Face Hub

Click to expand

🤗 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.

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.

Converting Models to Core ML

Click to expand

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

  • --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.

  • --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 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.

Image Generation with Python

Click to expand

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 <output-mlpackages-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.
  • --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.

Image Generation with Swift

Click to expand

Example CLI Usage

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.

Example Library Usage

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

Swift Package Details

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

Example Swift App

Click to expand

🤗 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.

FAQ

Click to expand
Q1: ERROR: Failed building wheel for tokenizers or error: can't find Rust compiler

A1: Please review this potential solution.

Q2: RuntimeError: {NSLocalizedDescription = "Error computing NN outputs."

A2: There are many potential causes for this error. In this context, it is highly likely to be encountered when your system is under increased memory pressure from other applications. Reducing memory utilization of other applications is likely to help alleviate the issue.

Q3: My Mac has 8GB RAM and I am converting models to Core ML using the example command. The process is getting killed because of memory issues. How do I fix this issue?

A3: In order to minimize the memory impact of the model conversion process, please execute the following command instead:

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

If you need --chunk-unet, you may do so in yet another independent command which will reuse the previously exported Unet model and simply chunk it in place:

python -m python_coreml_stable_diffusion.torch2coreml --convert-unet --chunk-unet -o <output-mlpackages-directory>
Q4: My Mac has 8GB RAM, should image generation work on my machine?

A4: Yes! Especially the --compute-unit CPU_AND_NE option should work under reasonable system load from other applications. Note that part of the Example Results were generated using an M2 MacBook Air with 8GB RAM.

Q5: Every time I generate an image using the Python pipeline, loading all the Core ML models takes 2-3 minutes. Is this expected?

A5: Yes and using the Swift library reduces this to just a few seconds. The reason is that coremltools loads Core ML models (.mlpackage) and each model is compiled to be run on the requested compute unit during load time. Because of the size and number of operations of the unet model, it takes around 2-3 minutes to compile it for Neural Engine execution. Other models should take at most a few seconds. Note that coremltools does not cache the compiled model for later loads so each load takes equally long. In order to benefit from compilation caching, StableDiffusion Swift package by default relies on compiled Core ML models (.mlmodelc) which will be compiled down for the requested compute unit upon first load but then the cache will be reused on subsequent loads until it is purged due to lack of use.

If you intend to use the Python pipeline in an application, we recommend initializing the pipeline once so that the load time is only incurred once. Afterwards, generating images using different prompts and random seeds will not incur the load time for the current session of your application.

Q6: I want to deploy StableDiffusion, the Swift package, in my mobile app. What should I be aware of?

A6: The Image Generation with Swift section describes the minimum SDK and OS versions as well as the device models supported by this package. We recommend carefully testing the package on the device with the least amount of RAM available among your deployment targets.

The image generation process in StableDiffusion can yield over 2 GB of peak memory during runtime depending on the compute units selected. On iPadOS, we recommend using .cpuAndNeuralEngine in your configuration and the reduceMemory option when constructing a StableDiffusionPipeline to minimize memory pressure.

If your app crashes during image generation, consider adding the Increased Memory Limit capability to inform the system that some of your app’s core features may perform better by exceeding the default app memory limit on supported devices.

On iOS, depending on the iPhone model, Stable Diffusion model versions, selected compute units, system load and design of your app, this may still not be sufficient to keep your apps peak memory under the limit. Please remember, because the device shares memory between apps and iOS processes, one app using too much memory can compromise the user experience across the whole device.

We strongly recommend quantizing your models with --quantize-nbits 6 for iOS deployment. This reduces the peak RAM usage by 1GB or more depending on the model version and robustly enables inference even on iPhone 12 Mini.

Q7: How do I generate images with different resolutions using the same Core ML models?

A7: The current version of python_coreml_stable_diffusion does not support single-model multi-resolution out of the box. However, developers may fork this project and leverage the flexible shapes support from coremltools to extend the torch2coreml script by using coremltools.EnumeratedShapes. Note that, while the text_encoder is agnostic to the image resolution, the inputs and outputs of vae_decoder and unet models are dependent on the desired image resolution.

Q8: Are the Core ML and PyTorch generated images going to be identical?

A8: If desired, the generated images across PyTorch and Core ML can be made approximately identical. However, it is not guaranteed by default. There are several factors that might lead to different images across PyTorch and Core ML:

1. Random Number Generator Behavior

The main source of potentially different results across PyTorch and Core ML is the Random Number Generator (RNG) behavior. PyTorch and Numpy have different sources of randomness. python_coreml_stable_diffusion generally relies on Numpy for RNG (e.g. latents initialization) and StableDiffusion Swift Library reproduces this RNG behavior by default. However, PyTorch-based pipelines such as Hugging Face diffusers relies on PyTorch's RNG behavior. Thanks to @liuliu's contribution, one can match the PyTorch (CPU) RNG behavior in Swift by specifying --rng torch which selects the torchRNG mode.

2. PyTorch

"Completely reproducible results are not guaranteed across PyTorch releases, individual commits, or different platforms. Furthermore, results may not be reproducible between CPU and GPU executions, even when using identical seeds." (source).

3. Model Function Drift During Conversion

The difference in outputs across corresponding PyTorch and Core ML models is a potential cause. The signal integrity is tested during the conversion process (enabled via --check-output-correctness argument to python_coreml_stable_diffusion.torch2coreml) and it is verified to be above a minimum PSNR value as tested on random inputs. Note that this is simply a sanity check and does not guarantee this minimum PSNR across all possible inputs. Furthermore, the results are not guaranteed to be identical when executing the same Core ML models across different compute units. This is not expected to be a major source of difference as the sample visual results indicate in this section.

4. Weights and Activations Data Type

When quantizing models from float32 to lower-precision data types such as float16, the generated images are known to vary slightly in semantics even when using the same PyTorch model. Core ML models generated by coremltools have float16 weights and activations by default unless explicitly overridden. This is not expected to be a major source of difference.

Q9: The model files are very large, how do I avoid a large binary for my App?

A9: The recommended option is to prompt the user to download these assets upon first launch of the app. This keeps the app binary size independent of the Core ML models being deployed. Disclosing the size of the download to the user is extremely important as there could be data charges or storage impact that the user might not be comfortable with.

Q10: `Could not initialize NNPACK! Reason: Unsupported hardware`

A10: This warning is safe to ignore in the context of this repository.

Q11: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect

A11: This warning is safe to ignore in the context of this repository.

Q12: UserWarning: resource_tracker: There appear to be 1 leaked semaphore objects to clean up at shutdown

A12: If this warning is printed right after zsh: killed python -m python_coreml_stable_diffusion.torch2coreml ... , then it is highly likely that your Mac has run out of memory while converting models to Core ML. Please see Q3 from above for the solution.

BibTeX Reference

@misc{stable-diffusion-coreml-apple-silicon,
title = {Stable Diffusion with Core ML on Apple Silicon},
author = {Atila Orhon and Michael Siracusa and Aseem Wadhwa},
year = {2022},
URL = {null}
}

More Repositories

1

swift

The Swift Programming Language
C++
65,899
star
2

swift-evolution

This maintains proposals for changes and user-visible enhancements to the Swift Programming Language.
Markdown
15,013
star
3

foundationdb

FoundationDB - the open source, distributed, transactional key-value store
C++
13,947
star
4

turicreate

Turi Create simplifies the development of custom machine learning models.
C++
11,153
star
5

darwin-xnu

The Darwin Kernel (mirror). This repository is a pure mirror and contributions are currently not accepted via pull-requests, please submit your contributions via https://developer.apple.com/bug-reporting/
C
10,558
star
6

swift-package-manager

The Package Manager for the Swift Programming Language
Swift
9,592
star
7

ml-ferret

Python
7,576
star
8

swift-nio

Event-driven network application framework for high performance protocol servers & clients, non-blocking.
Swift
7,274
star
9

swift-algorithms

Commonly used sequence and collection algorithms for Swift
Swift
5,622
star
10

swift-corelibs-foundation

The Foundation Project, providing core utilities, internationalization, and OS independence
Swift
5,189
star
11

swift-protobuf

Plugin and runtime library for using protobuf with Swift
Swift
4,446
star
12

password-manager-resources

A place for creators and users of password managers to collaborate on resources to make password management better.
JavaScript
4,010
star
13

coremltools

Core ML tools contain supporting tools for Core ML model conversion, editing, and validation.
Python
3,974
star
14

ml-mgie

Python
3,682
star
15

tensorflow_macos

TensorFlow for macOS 11.0+ accelerated using Apple's ML Compute framework.
Shell
3,643
star
16

swift-collections

Commonly used data structures for Swift
Swift
3,434
star
17

pkl

A configuration as code language with rich validation and tooling.
Java
3,360
star
18

swift-argument-parser

Straightforward, type-safe argument parsing for Swift
Swift
3,163
star
19

sourcekit-lsp

Language Server Protocol implementation for Swift and C-based languages
Swift
3,110
star
20

swift-log

A Logging API for Swift
Swift
2,931
star
21

swift-syntax

A set of Swift libraries for parsing, inspecting, generating, and transforming Swift source code.
Swift
2,887
star
22

swift-async-algorithms

Async Algorithms for Swift
Swift
2,695
star
23

swift-markdown

A Swift package for parsing, building, editing, and analyzing Markdown documents.
Swift
2,586
star
24

HomeKitADK

C
2,456
star
25

ml-ane-transformers

Reference implementation of the Transformer architecture optimized for Apple Neural Engine (ANE)
Python
2,431
star
26

swift-corelibs-libdispatch

The libdispatch Project, (a.k.a. Grand Central Dispatch), for concurrency on multicore hardware
C
2,420
star
27

swift-format

Formatting technology for Swift source code
Swift
2,261
star
28

homebrew-apple

Ruby
2,227
star
29

swift-foundation

The Foundation project
Swift
2,088
star
30

cups

Apple CUPS Sources
C
1,828
star
31

sample-food-truck

SwiftUI sample code from WWDC22
Swift
1,695
star
32

ml-fastvit

This repository contains the official implementation of the research paper, "FastViT: A Fast Hybrid Vision Transformer using Structural Reparameterization" ICCV 2023
Python
1,693
star
33

ml-cvnets

CVNets: A library for training computer vision networks
Python
1,664
star
34

swift-book

The Swift Programming Language book
Markdown
1,616
star
35

swift-numerics

Advanced mathematical types and functions for Swift
Swift
1,602
star
36

ml-hypersim

Hypersim: A Photorealistic Synthetic Dataset for Holistic Indoor Scene Understanding
Python
1,495
star
37

swift-crypto

Open-source implementation of a substantial portion of the API of Apple CryptoKit suitable for use on Linux platforms.
C
1,400
star
38

swift-docker

Docker Official Image packaging for Swift
Dockerfile
1,331
star
39

ml-neuman

Official repository of NeuMan: Neural Human Radiance Field from a Single Video (ECCV 2022)
Python
1,233
star
40

swift-openapi-generator

Generate Swift client and server code from an OpenAPI document.
Swift
1,142
star
41

swift-system

Low-level system calls and types for Swift
Swift
1,137
star
42

swift-corelibs-xctest

The XCTest Project, A Swift core library for providing unit test support
Swift
1,120
star
43

swift-docc

Documentation compiler that produces rich API reference documentation and interactive tutorials for your Swift framework or package.
Swift
1,093
star
44

swift-llbuild

A low-level build system, used by Xcode and the Swift Package Manager
C++
1,067
star
45

swift-atomics

Low-level atomic operations for Swift
Swift
1,004
star
46

swift-testing

Swift
981
star
47

servicetalk

A networking framework that evolves with your application
Java
881
star
48

swift-http-types

Version-independent HTTP currency types for Swift
Swift
815
star
49

swift-llvm

LLVM
815
star
50

swift-driver

Swift compiler driver reimplementation in Swift
Swift
764
star
51

swift-protobuf-plugin

Moved to apple/swift-protobuf
757
star
52

swift-lldb

This is the version of LLDB that supports the Swift programming language & REPL.
C++
673
star
53

swift-clang

C++
673
star
54

unityplugins

C#
645
star
55

ml-mobileone

This repository contains the official implementation of the research paper, "An Improved One millisecond Mobile Backbone".
Swift
641
star
56

ml-gaudi

602
star
57

ml-aim

This repository provides the code and model checkpoints of the research paper: Scalable Pre-training of Large Autoregressive Image Models
Python
602
star
58

swift-metrics

Metrics API for Swift
Swift
602
star
59

axlearn

Python
564
star
60

swift-distributed-actors

Peer-to-peer cluster implementation for Swift Distributed Actors
Swift
562
star
61

ARKitScenes

This repo accompanies the research paper, ARKitScenes - A Diverse Real-World Dataset for 3D Indoor Scene Understanding Using Mobile RGB-D Data and contains the data, scripts to visualize and process assets, and training code described in our paper.
Python
552
star
62

sample-backyard-birds

Swift
506
star
63

device-management

Device management schema data for MDM.
506
star
64

ccs-calendarserver

The Calendar and Contacts Server.
Python
470
star
65

ml-facelit

Official repository of FaceLit: Neural 3D Relightable Faces (CVPR 2023)
Python
457
star
66

swift-3-api-guidelines-review

Swift
452
star
67

swift-org-website

Swift.org website
SCSS
438
star
68

GCGC

Jupyter Notebook
436
star
69

swift-nio-http2

HTTP/2 support for SwiftNIO
Swift
405
star
70

swift-tools-support-core

Contains common infrastructural code for both SwiftPM and llbuild.
Swift
390
star
71

swift-nio-ssh

SwiftNIO SSH is a programmatic implementation of SSH using SwiftNIO
Swift
364
star
72

swift-nio-ssl

TLS Support for SwiftNIO, based on BoringSSL.
C
345
star
73

ml-gmpi

Official PyTorch implementation of GMPI (ECCV 2022, Oral Presentation)
Python
329
star
74

example-package-dealer

Example package for use with the Swift Package Manager
Swift
319
star
75

swift-collections-benchmark

A benchmarking tool for Swift Collection algorithms
Swift
316
star
76

example-package-playingcard

Example package for use with the Swift Package Manager
Swift
308
star
77

swift-docc-render

Web renderer for Swift-DocC documentation.
JavaScript
300
star
78

indexstore-db

Index database library for use with sourcekit-lsp
C++
299
star
79

swift-playdate-examples

A technical demonstration of Embedded Swift running on Playdate by Panic
Swift
295
star
80

swift-docc-plugin

Swift Package Manager command plugin for Swift-DocC
Swift
295
star
81

ml-hierarchical-confusion-matrix

Neo: Hierarchical Confusion Matrix Visualization (CHI 2022)
TypeScript
292
star
82

ml-gsn

Python
284
star
83

swift-llbuild2

A fresh take on a low-level build system API.
Swift
280
star
84

swift-source-compat-suite

The infrastructure and project index comprising the Swift source compatibility suite.
Python
278
star
85

sample-cloudkit-sharing

Swift
275
star
86

swift-xcode-playground-support

Logging and communication to allow Swift toolchains to communicate with Xcode.
Swift
270
star
87

swift-experimental-string-processing

An early experimental general-purpose pattern matching engine for Swift.
Swift
263
star
88

ml-sigma-reparam

Python
255
star
89

swift-standard-library-preview

Swift
253
star
90

swift-nio-transport-services

Extensions for SwiftNIO to support Apple platforms as first-class citizens.
Swift
252
star
91

swift-stress-tester

Stress testing utilities for Swift's tooling
Swift
207
star
92

swift-service-discovery

A service discovery API for Swift.
Swift
203
star
93

swift-certificates

An implementation of X.509 for Swift
Swift
195
star
94

swift-nio-examples

examples of how to use swift-nio
Swift
195
star
95

swift-cluster-membership

Distributed Membership Protocol implementations in Swift
Swift
191
star
96

swift-aoc-starter-example

Swift starter project for solving Advent of Code challenges.
Swift
189
star
97

sample-cloudkit-coredatasync

Swift
187
star
98

swift-distributed-tracing

Instrumentation library for Swift server applications
Swift
186
star
99

pfl-research

Simulation framework for accelerating research in Private Federated Learning
Python
186
star
100

swift-internals

HTML
182
star