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

Details (Click to expand)

Model Conversion:

Project Build:

Target Device Runtime:

Target Device Runtime (With Memory Improvements):

Target Device Hardware Generation:

Performance Benchmarks

Details (Click to expand)

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

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 Seed 8 in August 2023.
• The performance data was collected using the benchmark branch of the Diffusers app
• Swift code is not fully optimized, introducing up to ~10% overhead unrelated to Core ML model execution.
• The median latency value across 5 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.
• * indicates that the reduceMemory option was enabled which loads and unloads models just-in-time to avoid memory shortage. 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.

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

Details (Click to expand)

• This benchmark was conducted by Apple and Hugging Face using iOS 17.0.2 and iPadOS 17.0.2 in September 2023.
• The performance data was collected using the benchmark branch of the Diffusers app
• The median latency value across 5 back-to-back end-to-end executions are reported
• The image generation procedure follows this configuration: 20 inference steps, 768x768 output image resolution, 77 text token sequence length, classifier-free guidance (batch size of 2 for unet).
• Unet.mlmodelc is compressed to 4.04 bit precision following the Mixed-Bit Palettization algorithm recipe published here
• All models except for Unet.mlmodelc are compressed to 16 bit precision
• madebyollin/sdxl-vae-fp16-fix by @madebyollin was used as the source PyTorch model for VAEDecoder.mlmodelc in order to enable float16 weight and activation quantization for the VAE model.
• --attention-implementation SPLIT_EINSUM is chosen in lieu of SPLIT_EINSUM_V2 due to the prohibitively long compilation time of the latter
• * indicates that the reduceMemory option was enabled which loads and unloads models just-in-time to avoid memory shortage. This added significant overhead to the end-to-end latency. Note that end-to-end latency difference between iPad Pro (M1) and iPhone 13 Pro Max despite identical diffusion speed.
• 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.
• 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.

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

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 (6-bits and higher)

Details (Click to expand)

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 o