Update pull_request_template.md after changing to MIT license

Changing repository license to MIT

.github/pull_request_template.md

@@ -1,8 +1 @@
-Thank you for your interest in contributing to Core ML Stable Diffusion! Please review [CONTRIBUTING.md](../CONTRIBUTING.md) first. If you would like to proceed with making a pull request, please indicate your agreement to the terms outlined in CONTRIBUTING.md by checking the box below. If not, please go ahead and fork this repo and make your updates.
-
-We appreciate your interest in the project!
-
-Do not erase the below when submitting your pull request:
-#########
-
-- [ ] I agree to the terms outlined in CONTRIBUTING.md 
+Thank you for your interest in contributing to Core ML Stable Diffusion! Please review [CONTRIBUTING.md](../CONTRIBUTING.md) first. We appreciate your interest in the project!

.gitignore

@@ -1,3 +1,5 @@
+*~
+
 # Swift Package
 .DS_Store
 /.build

CONTRIBUTING.md

@@ -1,14 +1,11 @@
 # Contribution Guide
 
-Thank you for your interest in contributing to Core ML Stable Diffusion! This project was released for system demonstration purposes and there are limited plans for future development of the repository. While we welcome new pull requests and issues please note that our response may be limited. Forks and out-of-tree improvements are strongly encouraged.
+Thank you for your interest in contributing to Core ML Stable Diffusion! This project was released for system demonstration purposes and there are limited plans for future development of the repository. While we welcome new pull requests and issues please note that our response may be limited.
+
 
 ## Submitting a Pull Request
 
-Under the project's [LICENSE](LICENSE.md), confirmation of the following is required to merge your changes:
-
-*"I agree that all information entered is original and owned by me, and I hereby provide an irrevocable, royalty-free license to Apple to use, modify, copy, publish, prepare derivate works of, distribute (including under the Apple Sample Code License), such information and all intellectual property therein in whole or part, in perpetuity and worldwide, without any attribution."*
-
-By submitting a pull request, you represent that you have the right to license your contribution to Apple and the community, and agree by submitting the patch that your contributions are licensed under the [LICENSE](LICENSE.md).
+The project is licensed under the MIT license. By submitting a pull request, you represent that you have the right to license your contribution to Apple and the community, and agree by submitting the patch that your contributions are licensed under the MIT license.
 
 ## Code of Conduct
 

LICENSE.md

@@ -1,39 +1,9 @@
-Copyright (C) 2022 Apple Inc. All Rights Reserved.
+MIT License
 
-IMPORTANT:  This Apple software is supplied to you by Apple
-Inc. ("Apple") in consideration of your agreement to the following
-terms, and your use, installation, modification or redistribution of
-this Apple software constitutes acceptance of these terms.  If you do
-not agree with these terms, please do not use, install, modify or
-redistribute this Apple software.
+Copyright (c) 2024 Apple Inc.
 
-In consideration of your agreement to abide by the following terms, and
-subject to these terms, Apple grants you a personal, non-exclusive
-license, under Apple's copyrights in this original Apple software (the
-"Apple Software"), to use, reproduce, modify and redistribute the Apple
-Software, with or without modifications, in source and/or binary forms;
-provided that if you redistribute the Apple Software in its entirety and
-without modifications, you must retain this notice and the following
-text and disclaimers in all such redistributions of the Apple Software.
-Neither the name, trademarks, service marks or logos of Apple Inc. may
-be used to endorse or promote products derived from the Apple Software
-without specific prior written permission from Apple.  Except as
-expressly stated in this notice, no other rights or licenses, express or
-implied, are granted by Apple herein, including but not limited to any
-patent rights that may be infringed by your derivative works or by other
-works in which the Apple Software may be incorporated.
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
 
-The Apple Software is provided by Apple on an "AS IS" basis.  APPLE
-MAKES NO WARRANTIES, EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION
-THE IMPLIED WARRANTIES OF NON-INFRINGEMENT, MERCHANTABILITY AND FITNESS
-FOR A PARTICULAR PURPOSE, REGARDING THE APPLE SOFTWARE OR ITS USE AND
-OPERATION ALONE OR IN COMBINATION WITH YOUR PRODUCTS.
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
 
-IN NO EVENT SHALL APPLE BE LIABLE FOR ANY SPECIAL, INDIRECT, INCIDENTAL
-OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
-SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
-INTERRUPTION) ARISING IN ANY WAY OUT OF THE USE, REPRODUCTION,
-MODIFICATION AND/OR DISTRIBUTION OF THE APPLE SOFTWARE, HOWEVER CAUSED
-AND WHETHER UNDER THEORY OF CONTRACT, TORT (INCLUDING NEGLIGENCE),
-STRICT LIABILITY OR OTHERWISE, EVEN IF APPLE HAS BEEN ADVISED OF THE
-POSSIBILITY OF SUCH DAMAGE.
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

Package.swift

@@ -18,7 +18,7 @@ let package = Package(
             targets: ["StableDiffusionCLI"])
     ],
     dependencies: [
-        .package(url: "https://github.com/apple/swift-argument-parser.git", exact: "1.2.0")
+        .package(url: "https://github.com/apple/swift-argument-parser.git", from: "1.2.3")
     ],
     targets: [
         .target(

README.md

@@ -16,6 +16,9 @@ If you run into issues during installation or runtime, please refer to the [FAQ]
 
 ## <a name="system-requirements"></a> System Requirements
 
+<details>
+  <summary> Details (Click to expand) </summary>
+
 Model Conversion:
 
  macOS  | Python | coremltools |
@@ -26,7 +29,7 @@ Project Build:
 
   macOS | Xcode | Swift |
 :------:|:-----:|:-----:|
-  14.0  | 15.0  |  5.8  |
+  13.1  | 14.3  |  5.8  |
 
 Target Device Runtime:
 
@@ -34,7 +37,7 @@ Target Device Runtime:
 :------:|:-----------:|
   13.1  |     16.2    |
 
-Target Device Runtime ([With Memory Improvements](#compression)):
+Target Device Runtime ([With Memory Improvements](#compression-6-bits-and-higher)):
 
   macOS | iPadOS, iOS |
 :------:|:-----------:|
@@ -46,47 +49,112 @@ Target Device Hardware Generation:
 :------:|:-------:|:-------:|
    M1   |   M1    |  A14    |
 
+</details>
+
 
 ## <a name="performance-benchmark"></a> Performance Benchmarks
 
-
-[`stabilityai/stable-diffusion-2-1-base`](https://huggingface.co/apple/coreml-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>
   <summary> Details (Click to expand) </summary>
 
-- 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.
+
+[`stabilityai/stable-diffusion-2-1-base`](https://huggingface.co/apple/coreml-stable-diffusion-2-1-base) (512x512)
+
+
+
+|        Device         | `--compute-unit`| `--attention-implementation` | End-to-End Latency (s) | Diffusion Speed (iter/s) |
+| --------------------- | --------------- | ---------------------------- | ---------------------- | ------------------------ |
+| iPhone 12 Mini        | `CPU_AND_NE`    |      `SPLIT_EINSUM_V2`       |      18.5*             |        1.44              |
+| iPhone 12 Pro Max     | `CPU_AND_NE`    |      `SPLIT_EINSUM_V2`       |      15.4              |        1.45              |
+| iPhone 13             | `CPU_AND_NE`    |      `SPLIT_EINSUM_V2`       |      10.8*             |        2.53              |
+| iPhone 13 Pro Max     | `CPU_AND_NE`    |      `SPLIT_EINSUM_V2`       |      10.4              |        2.55              |
+| iPhone 14             | `CPU_AND_NE`    |      `SPLIT_EINSUM_V2`       |      8.6               |        2.57              |
+| iPhone 14 Pro Max     | `CPU_AND_NE`    |      `SPLIT_EINSUM_V2`       |      7.9               |        2.69              |
+| iPad Pro (M1)         | `CPU_AND_NE`    |      `SPLIT_EINSUM_V2`       |      11.2              |        2.19              |
+| iPad Pro (M2)         | `CPU_AND_NE`    |      `SPLIT_EINSUM_V2`       |      7.0               |        3.07              |
+
+<details>
+  <summary> Details (Click to expand) </summary>
+
+- 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](https://github.com/huggingface/swift-coreml-diffusers)
 - 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 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](#compression) for details.
+- Weights are compressed to 6 bit precision. Please refer to [this section](#compression-6-bits-and-higher) 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](https://github.com/apple/ml-stable-diffusion/blob/main/swift/StableDiffusion/pipeline/StableDiffusionPipeline.swift#L65) 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.
+- `*` indicates that the [reduceMemory](https://github.com/apple/ml-stable-diffusion/blob/main/swift/StableDiffusion/pipeline/StableDiffusionPipeline.swift#L91) 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.
+
+</details>
+
+
+[`stabilityai/stable-diffusion-xl-base-1.0-ios`](https://huggingface.co/apple/coreml-stable-diffusion-xl-base-ios) (768x768)
+
+|        Device         | `--compute-unit`| `--attention-implementation` | End-to-End Latency (s) | Diffusion Speed (iter/s) |
+| --------------------- | --------------- | ---------------------------- | ---------------------- | ------------------------ |
+| iPhone 12 Pro         | `CPU_AND_NE`    |      `SPLIT_EINSUM`          |            116*        |        0.50              |
+| iPhone 13 Pro Max     | `CPU_AND_NE`    |      `SPLIT_EINSUM`          |            86*         |        0.68              |
+| iPhone 14 Pro Max     | `CPU_AND_NE`    |      `SPLIT_EINSUM`          |            77*         |        0.83              |
+| iPhone 15 Pro Max     | `CPU_AND_NE`    |      `SPLIT_EINSUM`          |            31          |        0.85              |
+| iPad Pro (M1)         | `CPU_AND_NE`    |      `SPLIT_EINSUM`          |            36          |        0.69              |
+| iPad Pro (M2)         | `CPU_AND_NE`    |      `SPLIT_EINSUM`          |            27          |        0.98              |
+
+<details>
+  <summary> Details (Click to expand) </summary>
+
+- 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](https://github.com/huggingface/swift-coreml-diffusers)
+- 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](#compression-lower-than-6-bits) algorithm recipe published [here](https://huggingface.co/apple/coreml-stable-diffusion-mixed-bit-palettization/blob/main/recipes/stabilityai-stable-diffusion-xl-base-1.0_palettization_recipe.json)
+- All models except for `Unet.mlmodelc` are compressed to 16 bit precision
+- [madebyollin/sdxl-vae-fp16-fix](https://huggingface.co/madebyollin/sdxl-vae-fp16-fix) by [@madebyollin](https://github.com/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](https://github.com/apple/ml-stable-diffusion/blob/main/swift/StableDiffusion/pipeline/StableDiffusionPipeline.swift#L91) 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.
+
+
+</details>
+
+
+
+[`stabilityai/stable-diffusion-xl-base-1.0`](https://huggingface.co/apple/coreml-stable-diffusion-xl-base) (1024x1024)
+
+|        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>
+  <summary> Details (Click to expand) </summary>
+
+- 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.
 
 </details>
+</details>
 
-## <a name="compression"></a> Weight Compression
+
+## <a name="compression-6-bits-and-higher"></a> Weight Compression (6-bits and higher)
+
+<details>
+  <summary> Details (Click to expand) </summary>
 
 coremltools-7.0 supports advanced weight compression techniques for [pruning](https://coremltools.readme.io/v7.0/docs/pruning), [palettization](https://coremltools.readme.io/v7.0/docs/palettization-overview) and [linear 8-bit quantization](https://coremltools.readme.io/v7.0/docs/quantization-aware-training). 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.
 
@@ -94,7 +162,7 @@ We demonstrate how data-free [post-training palettization](https://coremltools.r
 
 For best results, we recommend [training-time palettization](https://coremltools.readme.io/v7.0/docs/training-time-palettization): `coremltools.optimize.torch.palettization.DKMPalettizer` if fine-tuning your model is feasible. This API implements the [Differentiable k-Means (DKM)](https://machinelearning.apple.com/research/differentiable-k-means) 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: 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.
+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`](https://huggingface.co/apple/coreml-stable-diffusion-2-1-base) generating *"a high quality photo of a surfing dog"* |
 | :---------------:| :----------------: | ------------------------------------------------------  |
@@ -102,27 +170,142 @@ The Neural Engine is capable of accelerating models with low-bit palettization:
 | 16-bit           | cpuAndNeuralEngine | <img src="assets/float16_cpuandne_readmereel.png">  |
 | 16-bit           | cpuAndGPU          | <img src="assets/float16_gpu_readmereel.png"> |
 
-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 (2 and 4) will require fine-tuning to recover image generation quality as previously mentioned.
+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](#compression-lower-than-6-bits).
 
 Resources:
 - [Core ML Tools Docs: Optimizing Models](https://coremltools.readme.io/v7.0/docs/optimizing-models)
 - [WWDC23 Session Video: Use Core ML Tools for machine learning model compression](https://developer.apple.com/videos/play/wwdc2023/10047)
 
+</details>
+
+## <a name="compression-lower-than-6-bits"></a> Advanced Weight Compression (Lower than 6-bits)
+
+<details>
+  <summary> Details (Click to expand) </summary>
+
+This section describes an advanced compression algorithm called [Mixed-Bit Palettization (MBP)](https://huggingface.co/blog/stable-diffusion-xl-coreml#what-is-mixed-bit-palettization) built on top of the [Post-Training Weight Palettization tools](https://apple.github.io/coremltools/docs-guides/source/post-training-palettization.html) and using the [Weights Metadata API](https://apple.github.io/coremltools/docs-guides/source/mlmodel-utilities.html#get-weights-metadata) from [coremltools](https://github.com/apple/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) |
+| :-------:| :-------:| :-------:| :----------------:|
+| <img src="assets/mbp/a_high_quality_photo_of_a_surfing_dog.7667.final_3.41-bits.png"> | <img src="assets/mbp/a_high_quality_photo_of_a_surfing_dog.7667.final_4.50-bits.png">  | <img src="assets/mbp/a_high_quality_photo_of_a_surfing_dog.7667.final_6.55-bits.png"> | <img src="assets/mbp/a_high_quality_photo_of_a_surfing_dog.7667.final_float16_original.png"> |
+
+
+For example, the original float16 [stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0) model has an ~82 dB signal strength. Naively applying [linear 8-bit quantization](https://coremltools.readme.io/docs/data-free-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.
+
+<img src="assets/mbp/stabilityai_stable-diffusion-xl-base-1.0_psnr_vs_size.png" width="640">
+
+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
+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](https://huggingface.co/apple/coreml-stable-diffusion-mixed-bit-palettization/tree/main/recipes). 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.:
+
+```json
+{
+  "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
+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.
+
+</details>
+
+## <a name="using-stable-diffusion-xl"></a> Using Stable Diffusion XL
+
+<details>
+  <summary> Details (Click to expand) </summary>
+
+### Model Conversion
+
+e.g.:
+
+```bash
+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"](https://huggingface.co/docs/diffusers/main/en/api/pipelines/stable_diffusion/stable_diffusion_xl#1-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](https://github.com/apple/ml-stable-diffusion/blob/main/python_coreml_stable_diffusion/torch2coreml.py#L486). Using a custom VAE version such as [madebyollin/sdxl-vae-fp16-fix](https://huggingface.co/madebyollin/sdxl-vae-fp16-fix) by [@madebyollin](https://github.com/madebyollin) via `--custom-vae-version madebyollin/sdxl-vae-fp16-fix` will restore the default float16 precision for VAE.
+
+### Swift Inference
+
+```bash
+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 Inference
+
+```bash
+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
+
+</details>
+
 ## <a name="using-controlnet"></a> Using ControlNet
 
+<details>
+  <summary> Details (Click to expand) </summary>
+
 Example results using the prompt *"a high quality photo of a surfing dog"* conditioned on the scribble (leftmost):
 
 <img src="assets/controlnet_readme_reel.png">
 
 [ControlNet](https://huggingface.co/lllyasviel/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](https://github.com/apple/ml-stable-diffusion/pull/153), both the Python CLI and the Swift package support ControlNet models. Please refer to [this section](#converting-models-to-coreml) for details on setting up Stable Diffusion with ControlNet.
 
+Note that ControlNet is not yet supported for Stable Diffusion XL.
+
+</details>
 
 ## <a name="system-multilingual-text-encoder"></a> Using the System Multilingual Text Encoder
 
 <details>
   <summary> Details (Click to expand) </summary>
 
-With iOS 17 and macOS 14, `NaturalLanguage` framework introduced the [NLContextualEmbedding](https://developer.apple.com/documentation/naturallanguage/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](https://developer.apple.com/videos/play/wwdc2023/10042) demonstrated how this powerful new model can be used by developers to train downstream tasks such as multilingual image generation with Stable Diffusion.
+With iOS 17 and macOS 14, `NaturalLanguage` framework introduced the [NLContextualEmbedding](https://developer.apple.com/documentation/naturallanguage/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](https://developer.apple.com/videos/play/wwdc2023/10042) 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:
 
@@ -156,10 +339,23 @@ Resources:
 
 🤗 Hugging Face ran the [conversion procedure](#converting-models-to-coreml) 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](#converting-models-to-core-ml).
 
-* [`CompVis/stable-diffusion-v1-4`](https://huggingface.co/apple/coreml-stable-diffusion-v1-4)
-* [`runwayml/stable-diffusion-v1-5`](https://huggingface.co/apple/coreml-stable-diffusion-v1-5)
-* [`stabilityai/stable-diffusion-2-base`](https://huggingface.co/apple/coreml-stable-diffusion-2-base)
-* [`stabilityai/stable-diffusion-2-1-base`](https://huggingface.co/apple/coreml-stable-diffusion-2-1-base)
+* 6-bit quantized models (suitable for iOS 17 and macOS 14):
+  - [`CompVis/stable-diffusion-v1-4`](https://huggingface.co/apple/coreml-stable-diffusion-1-4-palettized)
+  - [`runwayml/stable-diffusion-v1-5`](https://huggingface.co/apple/coreml-stable-diffusion-v1-5-palettized)
+  - [`stabilityai/stable-diffusion-2-base`](https://huggingface.co/apple/coreml-stable-diffusion-2-base-palettized)
+  - [`stabilityai/stable-diffusion-2-1-base`](https://huggingface.co/apple/coreml-stable-diffusion-2-1-base-palettized)
+
+* Mixed-bit quantized models
+- [`stabilityai/stable-diffusion-xl-base-1.0`](https://huggingface.co/apple/coreml-stable-diffusion-mixed-bit-palettization)
+- [`stabilityai/stable-diffusion-xl-base-1.0-ios`](https://huggingface.co/apple/coreml-stable-diffusion-xl-base-ios)
+
+* Uncompressed models:
+  - [`CompVis/stable-diffusion-v1-4`](https://huggingface.co/apple/coreml-stable-diffusion-v1-4)
+  - [`runwayml/stable-diffusion-v1-5`](https://huggingface.co/apple/coreml-stable-diffusion-v1-5)
+  - [`stabilityai/stable-diffusion-2-base`](https://huggingface.co/apple/coreml-stable-diffusion-2-base)
+  - [`stabilityai/stable-diffusion-2-1-base`](https://huggingface.co/apple/coreml-stable-diffusion-2-1-base)
+  - [`stabilityai/stable-diffusion-xl-base-1.0`](https://huggingface.co/apple/coreml-stable-diffusion-xl-base)
+  - [`stabilityai/stable-diffusion-xl-{base+refiner}-1.0`](https://huggingface.co/apple/coreml-stable-diffusion-xl-base-with-refiner)
 
 If you want to use any of those models you may download the weights and proceed to [generate images with Python](#image-generation-with-python) or [Swift](#image-generation-with-swift).
 
@@ -222,19 +418,20 @@ pip install -e .
 **Step 4:** Execute the following command from the Terminal to generate Core ML model files (`.mlpackage`)
 
 ```shell
-python -m python_coreml_stable_diffusion.torch2coreml --convert-unet --convert-text-encoder --convert-vae-decoder --convert-safety-checker -o <output-mlpackages-directory>
+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 defaults to [CompVis/stable-diffusion-v1-4](https://huggingface.co/CompVis/stable-diffusion-v1-4). Developers may specify other versions that are available on [Hugging Face Hub](https://huggingface.co/models?search=stable-diffusion), e.g. [stabilityai/stable-diffusion-2-base](https://huggingface.co/stabilityai/stable-diffusion-2-base) & [runwayml/stable-diffusion-v1-5](https://huggingface.co/runwayml/stable-diffusion-v1-5).
+- `--model-version`: The model version name as published on the [Hugging Face Hub](https://huggingface.co/models?search=stable-diffusion)
 
+- `--refiner-version`: The refiner version name as published on the [Hugging Face Hub](https://huggingface.co/models?search=stable-diffusion). 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.
+- `--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](https://apple.github.io/coremltools/docs-guides/source/model-prediction.html#why-use-a-compiled-model).
 
-- `--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.
+- `--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.
 
@@ -246,6 +443,8 @@ This generally takes 15-20 minutes on an M1 MacBook Pro. Upon successful executi
 
 - `--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.
+
 </details>
 
 ## <a name="image-generation-with-python"></a> Image Generation with Python
@@ -256,11 +455,11 @@ This generally takes 15-20 minutes on an M1 MacBook Pro. Upon successful executi
 Run text-to-image generation using the example Python pipeline based on [diffusers](https://github.com/huggingface/diffusers):
 
 ```shell
-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
+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](#converting-models-to-coreml) section from above.
+- `-i`: Should point to the `-o` directory from Step 4 of [Converting Models to Core ML](#converting-models-to-coreml) 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](https://apple.github.io/coremltools/docs-guides/source/model-prediction.html#why-use-a-compiled-model).
 - `--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](#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.
@@ -303,7 +502,7 @@ This Swift package contains two products:
 
 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)
+- `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)
@@ -317,6 +516,10 @@ Optionally, it may also include the safety checker model that some versions of S
 
 - `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)
@@ -366,11 +569,11 @@ Hugging Face has made the app [available in the Mac App Store](https://apps.appl
 <b> A3: </b>  In order to minimize the memory impact of the model conversion process, please execute the following command instead:
 
 ```bash
-python -m python_coreml_stable_diffusion.torch2coreml --convert-vae-encoder -o <output-mlpackages-directory> && \
-python -m python_coreml_stable_diffusion.torch2coreml --convert-vae-decoder -o <output-mlpackages-directory> && \
-python -m python_coreml_stable_diffusion.torch2coreml --convert-unet -o <output-mlpackages-directory> && \
-python -m python_coreml_stable_diffusion.torch2coreml --convert-text-encoder -o <output-mlpackages-directory> && \
-python -m python_coreml_stable_diffusion.torch2coreml --convert-safety-checker -o <output-mlpackages-directory> &&
+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:
@@ -390,9 +593,10 @@ python -m python_coreml_stable_diffusion.torch2coreml --convert-unet --chunk-une
 <details>
 <summary> <b> Q5: </b> Every time I generate an image using the Python pipeline, loading all the Core ML models takes 2-3 minutes. Is this expected? </summary>
 
-<b> A5: </b> 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.
+<b> A5: </b> Both `.mlpackage` and `.mlmodelc` models are compiled (also known as "model preparation" in Core ML terms) upon first load when a specific compute unit is specified. `.mlpackage` does not cache this compiled asset so each model load retriggers this compilation which may take up to a few minutes. On the other hand, `.mlmodelc` files do cache this compiled asset and non-first load times are reduced to just a few seconds.
+
+In order to benefit from compilation caching, you may use the `.mlmodelc` assets instead of `.mlpackage` assets in both Swift (default) and Python (possible thanks to [@lopez-hector](https://github.com/lopez-hector)'s [contribution](https://github.com/apple/ml-stable-diffusion/commit/f3a212491cf531dd88493c89ad3d98d016db407f)) image generation pipelines.
 
-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.
 
 </details>
 
@@ -408,7 +612,7 @@ If your app crashes during image generation, consider adding the [Increased Memo
  
 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.
+We **strongly recommend** compressing your models following the recipes in [Advanced Weight Compression (Lower than 6-bits)](#compression-lower-than-6-bits) for iOS deployment. This reduces the peak RAM usage by up to 75% (from 16-bit to 4-bit) while preserving model output quality.
 
 </details>
 
@@ -426,7 +630,7 @@ We **strongly recommend** quantizing your models with `--quantize-nbits 6` for i
 
   <b> 1. Random Number Generator Behavior </b>
 
-  The main source of potentially different results across PyTorch and Core ML is the Random Number Generator ([RNG](https://en.wikipedia.org/wiki/Random_number_generation)) 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](https://github.com/apple/ml-stable-diffusion/pull/124), one can match the PyTorch (CPU) RNG behavior in Swift by specifying `--rng torch` which selects the `torchRNG` mode.
+  The main source of potentially different results across PyTorch and Core ML is the Random Number Generator ([RNG](https://en.wikipedia.org/wiki/Random_number_generation)) 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 [contributions](https://github.com/apple/ml-stable-diffusion/pull/124), one can match the PyTorch (CPU/GPU) RNG behavior in Swift by specifying `--rng torch/cuda` which selects the `torchRNG/cudaRNG` mode.
 
   <b> 2. PyTorch </b>
 
@@ -434,7 +638,7 @@ We **strongly recommend** quantizing your models with `--quantize-nbits 6` for i
 
   <b> 3. Model Function Drift During Conversion </b>
 
-  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](https://en.wikipedia.org/wiki/Peak_signal-to-noise_ratio) 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](#compression).
+  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](https://en.wikipedia.org/wiki/Peak_signal-to-noise_ratio) 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](#compression-6-bits-and-higher).
 
   <b> 4. Weights and Activations Data Type </b>
 

python_coreml_stable_diffusion/_version.py

@@ -1 +1 @@
-__version__ = "1.0.0"
+__version__ = "1.1.0"

python_coreml_stable_diffusion/chunk_mlprogram.py

@@ -31,8 +31,10 @@ import shutil
 import time
 
 
-def _verify_output_correctness_of_chunks(full_model, first_chunk_model,
-                                         second_chunk_model):
+def _verify_output_correctness_of_chunks(full_model,
+                                         first_chunk_model=None,
+                                         second_chunk_model=None,
+                                         pipeline_model=None,):
     """ Verifies the end-to-end output correctness of full (original) model versus chunked models
     """
     # Generate inputs for first chunk and full model
@@ -40,30 +42,40 @@ def _verify_output_correctness_of_chunks(full_model, first_chunk_model,
     for input_desc in full_model._spec.description.input:
         input_dict[input_desc.name] = random_gen_input_feature_type(input_desc)
 
-    # Generate outputs for first chunk and full model
+    # Generate outputs for full model
     outputs_from_full_model = full_model.predict(input_dict)
-    outputs_from_first_chunk_model = first_chunk_model.predict(input_dict)
 
-    # Prepare inputs for second chunk model from first chunk's outputs and regular inputs
-    second_chunk_input_dict = {}
-    for input_desc in second_chunk_model._spec.description.input:
-        if input_desc.name in outputs_from_first_chunk_model:
-            second_chunk_input_dict[
-                input_desc.name] = outputs_from_first_chunk_model[
+    if pipeline_model is not None:
+        outputs_from_pipeline_model = pipeline_model.predict(input_dict)
+        final_outputs = outputs_from_pipeline_model
+
+    elif first_chunk_model is not None and second_chunk_model is not None:
+        # Generate outputs for first chunk
+        outputs_from_first_chunk_model = first_chunk_model.predict(input_dict)
+
+        # Prepare inputs for second chunk model from first chunk's outputs and regular inputs
+        second_chunk_input_dict = {}
+        for input_desc in second_chunk_model._spec.description.input:
+            if input_desc.name in outputs_from_first_chunk_model:
+                second_chunk_input_dict[
+                    input_desc.name] = outputs_from_first_chunk_model[
+                        input_desc.name]
+            else:
+                second_chunk_input_dict[input_desc.name] = input_dict[
                     input_desc.name]
-        else:
-            second_chunk_input_dict[input_desc.name] = input_dict[
-                input_desc.name]
 
-    # Generate output for second chunk model
-    outputs_from_second_chunk_model = second_chunk_model.predict(
-        second_chunk_input_dict)
+        # Generate output for second chunk model
+        outputs_from_second_chunk_model = second_chunk_model.predict(
+            second_chunk_input_dict)
+        final_outputs = outputs_from_second_chunk_model
+    else:
+        raise ValueError
 
     # Verify correctness across all outputs from second chunk and full model
     for out_name in outputs_from_full_model.keys():
         torch2coreml.report_correctness(
             original_outputs=outputs_from_full_model[out_name],
-            final_outputs=outputs_from_second_chunk_model[out_name],
+            final_outputs=final_outputs[out_name],
             log_prefix=f"{out_name}")
 
 
@@ -105,7 +117,10 @@ def _get_op_idx_split_location(prog: Program):
             size_in_mb = op.val.val.size * op.val.val.itemsize / (1024 * 1024)
             cumulative_size_in_mb += size_in_mb
 
-        if (cumulative_size_in_mb > half_size and op.op_type != "const"
+        # Note: The condition "not op.op_type.startswith("const")" is to make sure that the
+        # incision op is neither of type "const" nor "constexpr_*" ops that
+        # are used to store compressed weights
+        if (cumulative_size_in_mb > half_size and not op.op_type.startswith("const")
                 and len(op.outputs) == 1
                 and len(op.outputs[0].child_ops) == 1):
             op_idx = main_block.operations.index(op)
@@ -119,12 +134,13 @@ def _get_first_chunk_outputs(block, op_idx):
     boundary_vars = set()
     for i in range(op_idx + 1):
         op = block.operations[i]
-        for var in op.outputs:
-            if var.val is None:  # only consider non const vars
-                for child_op in var.child_ops:
-                    child_op_idx = block.operations.index(child_op)
-                    if child_op_idx > op_idx:
-                        boundary_vars.add(var)
+        if not op.op_type.startswith("const"):
+            for var in op.outputs:
+                if var.val is None:  # only consider non const vars
+                    for child_op in var.child_ops:
+                        child_op_idx = block.operations.index(child_op)
+                        if child_op_idx > op_idx:
+                            boundary_vars.add(var)
     return list(boundary_vars)
 
 
@@ -192,6 +208,11 @@ def _make_second_chunk_prog(prog, op_idx):
                 anchor_op=boundary_op,
                 old_var=var,
                 new_var=new_var,
+                # This is needed if the program contains "constexpr_*" ops. In normal cases, there are stricter
+                # rules for removing them, and their presence may prevent replacing this var.
+                # However in this case, since we want to remove all the ops in chunk 1, we can safely
+                # set this to True.
+                force_replace=True,
             )
 
     PASS_REGISTRY["common::dead_code_elimination"](prog)
@@ -293,16 +314,32 @@ def main(args):
         shutil.rmtree(args.mlpackage_path)
         logger.info("Done.")
 
-    # Save the chunked models to disk
-    out_path_chunk1 = os.path.join(args.o, name + "_chunk1.mlpackage")
-    out_path_chunk2 = os.path.join(args.o, name + "_chunk2.mlpackage")
+    if args.merge_chunks_in_pipeline_model:
+        # Make a single pipeline model to manage the model chunks
+        pipeline_model = ct.utils.make_pipeline(model_chunk1, model_chunk2)
+        out_path_pipeline = os.path.join(args.o, name + "_chunked_pipeline.mlpackage")
 
-    logger.info(
-        f"Saved chunks in {args.o} with the suffix _chunk1.mlpackage and _chunk2.mlpackage"
-    )
-    model_chunk1.save(out_path_chunk1)
-    model_chunk2.save(out_path_chunk2)
-    logger.info("Done.")
+        # Save and reload to ensure CPU placement
+        pipeline_model.save(out_path_pipeline)
+        pipeline_model = ct.models.MLModel(out_path_pipeline, compute_units=ct.ComputeUnit.CPU_ONLY)
+
+        if args.check_output_correctness:
+            logger.info("Verifying output correctness of pipeline model")
+            _verify_output_correctness_of_chunks(
+                full_model=model,
+                pipeline_model=pipeline_model,
+            )
+    else:
+        # Save the chunked models to disk
+        out_path_chunk1 = os.path.join(args.o, name + "_chunk1.mlpackage")
+        out_path_chunk2 = os.path.join(args.o, name + "_chunk2.mlpackage")
+
+        logger.info(
+            f"Saved chunks in {args.o} with the suffix _chunk1.mlpackage and _chunk2.mlpackage"
+        )
+        model_chunk1.save(out_path_chunk1)
+        model_chunk2.save(out_path_chunk2)
+        logger.info("Done.")
 
 
 if __name__ == "__main__":
@@ -332,6 +369,12 @@ if __name__ == "__main__":
         ("If specified, compares the outputs of original Core ML model with that of pipelined CoreML model chunks and reports PSNR in dB. ",
          "Enabling this feature uses more memory. Disable it if your machine runs out of memory."
          ))
+    parser.add_argument(
+        "--merge-chunks-in-pipeline-model",
+        action="store_true",
+        help=
+        ("If specified, model chunks are managed inside a single pipeline model for easier asset maintenance"
+         ))
 
     args = parser.parse_args()
     main(args)

python_coreml_stable_diffusion/controlnet.py

@@ -69,6 +69,7 @@ class ControlNetModel(ModelMixin, ConfigMixin):
         norm_num_groups=32,
         norm_eps=1e-5,
         cross_attention_dim=1280,
+        transformer_layers_per_block=1,
         attention_head_dim=8,
         use_linear_projection=False,
         upcast_attention=False,
@@ -129,6 +130,9 @@ class ControlNetModel(ModelMixin, ConfigMixin):
         if isinstance(attention_head_dim, int):
             attention_head_dim = (attention_head_dim,) * len(down_block_types)
 
+        if isinstance(transformer_layers_per_block, int):
+            transformer_layers_per_block = [transformer_layers_per_block] * len(down_block_types)
+
         # down
         output_channel = block_out_channels[0]
 
@@ -142,6 +146,7 @@ class ControlNetModel(ModelMixin, ConfigMixin):
 
             down_block = get_down_block(
                 down_block_type,
+                transformer_layers_per_block=transformer_layers_per_block[i],
                 num_layers=layers_per_block,
                 in_channels=input_channel,
                 out_channels=output_channel,
@@ -151,6 +156,7 @@ class ControlNetModel(ModelMixin, ConfigMixin):
                 cross_attention_dim=cross_attention_dim,
                 attn_num_head_channels=attention_head_dim[i],
                 downsample_padding=downsample_padding,
+                add_downsample=not is_final_block,
             )
             self.down_blocks.append(down_block)
 

python_coreml_stable_diffusion/coreml_model.py

@@ -6,6 +6,7 @@
 import coremltools as ct
 
 import logging
+import json
 
 logging.basicConfig()
 logger = logging.getLogger(__name__)
@@ -21,14 +22,47 @@ class CoreMLModel:
     """ Wrapper for running CoreML models using coremltools
     """
 
-    def __init__(self, model_path, compute_unit):
-        assert os.path.exists(model_path) and model_path.endswith(".mlpackage")
+    def __init__(self, model_path, compute_unit, sources='packages'):
 
         logger.info(f"Loading {model_path}")
 
         start = time.time()
-        self.model = ct.models.MLModel(
-            model_path, compute_units=ct.ComputeUnit[compute_unit])
+        if sources == 'packages':
+            assert os.path.exists(model_path) and model_path.endswith(".mlpackage")
+
+            self.model = ct.models.MLModel(
+                model_path, compute_units=ct.ComputeUnit[compute_unit])
+            DTYPE_MAP = {
+                65552: np.float16,
+                65568: np.float32,
+                131104: np.int32,
+            }
+            self.expected_inputs = {
+                input_tensor.name: {
+                    "shape": tuple(input_tensor.type.multiArrayType.shape),
+                    "dtype": DTYPE_MAP[input_tensor.type.multiArrayType.dataType],
+                }
+                for input_tensor in self.model._spec.description.input
+            }
+        elif sources == 'compiled':
+            assert os.path.exists(model_path) and model_path.endswith(".mlmodelc")
+
+            self.model = ct.models.CompiledMLModel(model_path, ct.ComputeUnit[compute_unit])
+
+            # Grab expected inputs from metadata.json
+            with open(os.path.join(model_path, 'metadata.json'), 'r') as f:
+                config = json.load(f)[0]
+
+            self.expected_inputs = {
+                input_tensor['name']: {
+                    "shape": tuple(eval(input_tensor['shape'])),
+                    "dtype": np.dtype(input_tensor['dataType'].lower()),
+                }
+                for input_tensor in config['inputSchema']
+            }
+        else:
+            raise ValueError(f'Expected `packages` or `compiled` for sources, received {sources}')
+
         load_time = time.time() - start
         logger.info(f"Done. Took {load_time:.1f} seconds.")
 
@@ -38,21 +72,6 @@ class CoreMLModel:
                 "The Swift package we provide uses precompiled Core ML models (.mlmodelc) to avoid compile-on-load."
             )
 
-
-        DTYPE_MAP = {
-            65552: np.float16,
-            65568: np.float32,
-            131104: np.int32,
-        }
-
-        self.expected_inputs = {
-            input_tensor.name: {
-                "shape": tuple(input_tensor.type.multiArrayType.shape),
-                "dtype": DTYPE_MAP[input_tensor.type.multiArrayType.dataType],
-            }
-            for input_tensor in self.model._spec.description.input
-        }
-
     def _verify_inputs(self, **kwargs):
         for k, v in kwargs.items():
             if k in self.expected_inputs:
@@ -72,7 +91,7 @@ class CoreMLModel:
                         f"Expected shape {expected_shape}, got {v.shape} for input: {k}"
                     )
             else:
-                raise ValueError("Received unexpected input kwarg: {k}")
+                raise ValueError(f"Received unexpected input kwarg: {k}")
 
     def __call__(self, **kwargs):
         self._verify_inputs(**kwargs)
@@ -82,21 +101,77 @@ class CoreMLModel:
 LOAD_TIME_INFO_MSG_TRIGGER = 10  # seconds
 
 
-def _load_mlpackage(submodule_name, mlpackages_dir, model_version,
-                    compute_unit):
-    """ Load Core ML (mlpackage) models from disk (As exported by torch2coreml.py)
+def get_resource_type(resources_dir: str) -> str:
     """
-    logger.info(f"Loading {submodule_name} mlpackage")
+        Detect resource type based on filepath extensions.
+        returns:
+            `packages`: for .mlpackage resources
+            'compiled`: for .mlmodelc resources
+    """
+    directories = [f for f in os.listdir(resources_dir) if os.path.isdir(os.path.join(resources_dir, f))]
 
-    fname = f"Stable_Diffusion_version_{model_version}_{submodule_name}.mlpackage".replace(
-        "/", "_")
-    mlpackage_path = os.path.join(mlpackages_dir, fname)
+    # consider directories ending with extension
+    extensions = set([os.path.splitext(e)[1] for e in directories if os.path.splitext(e)[1]])
 
-    if not os.path.exists(mlpackage_path):
-        raise FileNotFoundError(
-            f"{submodule_name} CoreML model doesn't exist at {mlpackage_path}")
+    # if one extension present we may be able to infer sources type
+    if len(set(extensions)) == 1:
+        extension = extensions.pop()
+    else:
+        raise ValueError(f'Multiple file extensions found at {resources_dir}.'
+                         f'Cannot infer resource type from contents.')
+
+    if extension == '.mlpackage':
+        sources = 'packages'
+    elif extension == '.mlmodelc':
+        sources = 'compiled'
+    else:
+        raise ValueError(f'Did not find .mlpackage or .mlmodelc at {resources_dir}')
+
+    return sources
+
+
+def _load_mlpackage(submodule_name,
+                    mlpackages_dir,
+                    model_version,
+                    compute_unit,
+                    sources=None):
+    """
+        Load Core ML (mlpackage) models from disk (As exported by torch2coreml.py)
+
+    """
+
+    # if sources not provided, attempt to infer `packages` or `compiled` from the
+    # resources directory
+    if sources is None:
+        sources = get_resource_type(mlpackages_dir)
+
+    if sources == 'packages':
+        logger.info(f"Loading {submodule_name} mlpackage")
+        fname = f"Stable_Diffusion_version_{model_version}_{submodule_name}.mlpackage".replace(
+            "/", "_")
+        mlpackage_path = os.path.join(mlpackages_dir, fname)
+
+        if not os.path.exists(mlpackage_path):
+            raise FileNotFoundError(
+                f"{submodule_name} CoreML model doesn't exist at {mlpackage_path}")
+
+    elif sources == 'compiled':
+        logger.info(f"Loading {submodule_name} mlmodelc")
+
+        # FixMe: Submodule names and compiled resources names differ. Can change if names match in the future.
+        submodule_names = ["text_encoder", "text_encoder_2", "unet", "vae_decoder", "vae_encoder", "safety_checker"]
+        compiled_names = ['TextEncoder', 'TextEncoder2', 'Unet', 'VAEDecoder', 'VAEEncoder', 'SafetyChecker']
+        name_map = dict(zip(submodule_names, compiled_names))
+
+        cname = name_map[submodule_name] + '.mlmodelc'
+        mlpackage_path = os.path.join(mlpackages_dir, cname)
+
+        if not os.path.exists(mlpackage_path):
+            raise FileNotFoundError(
+                f"{submodule_name} CoreML model doesn't exist at {mlpackage_path}")
+
+    return CoreMLModel(mlpackage_path, compute_unit, sources=sources)
 
-    return CoreMLModel(mlpackage_path, compute_unit)
 
 def _load_mlpackage_controlnet(mlpackages_dir, model_version, compute_unit):
     """ Load Core ML (mlpackage) models from disk (As exported by torch2coreml.py)
@@ -115,5 +190,6 @@ def _load_mlpackage_controlnet(mlpackages_dir, model_version, compute_unit):
 
     return CoreMLModel(mlpackage_path, compute_unit)
 
+
 def get_available_compute_units():
     return tuple(cu for cu in ct.ComputeUnit._member_names_)

python_coreml_stable_diffusion/mixed_bit_compression_apply.py

@@ -0,0 +1,133 @@
+import argparse
+import gc
+import json
+import logging
+import os
+
+import coremltools as ct
+import coremltools.optimize.coreml as cto
+import numpy as np
+
+from python_coreml_stable_diffusion.torch2coreml import get_pipeline
+from python_coreml_stable_diffusion.mixed_bit_compression_pre_analysis import (
+    NBITS,
+    PALETTIZE_MIN_SIZE as MIN_SIZE
+)
+
+
+logging.basicConfig()
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+
+
+def main(args):
+    # Load Core ML model
+    coreml_model = ct.models.MLModel(args.mlpackage_path, compute_units=ct.ComputeUnit.CPU_ONLY)
+    logger.info(f"Loaded {args.mlpackage_path}")
+
+    # Load palettization recipe
+    with open(args.pre_analysis_json_path, 'r') as f:
+        pre_analysis = json.load(f)
+
+    if args.selected_recipe not in list(pre_analysis["recipes"]):
+        raise KeyError(
+            f"--selected-recipe ({args.selected_recipe}) not found in "
+            f"--pre-analysis-json-path ({args.pre_analysis_json_path}). "
+            f" Available recipes: {list(pre_analysis['recipes'])}"
+        )
+
+
+    recipe = pre_analysis["recipes"][args.selected_recipe]
+    assert all(nbits in NBITS + [16] for nbits in recipe.values()), \
+        f"Some nbits values in the recipe are illegal. Allowed values: {NBITS}"
+
+    # Hash tensors to be able to match torch tensor names to mil tensors
+    def get_tensor_hash(tensor):
+        assert tensor.dtype == np.float16
+        return tensor.ravel()[0] + np.prod(tensor.shape)
+
+    args.model_version = pre_analysis["model_version"]
+    pipe = get_pipeline(args)
+    torch_model = pipe.unet
+
+    hashed_recipe = {}
+    for torch_module_name, nbits in recipe.items():
+        tensor = [
+            tensor.cpu().numpy().astype(np.float16) for name,tensor in torch_model.named_parameters()
+            if name == torch_module_name + '.weight'
+        ][0]
+        hashed_recipe[get_tensor_hash(tensor)] = nbits
+
+    del pipe
+    gc.collect()
+
+    op_name_configs = {}
+    weight_metadata = cto.get_weights_metadata(coreml_model, weight_threshold=MIN_SIZE)
+    hashes = np.array(list(hashed_recipe))
+    for name, metadata in weight_metadata.items():
+        # Look up target bits for this weight
+        tensor_hash = get_tensor_hash(metadata.val)
+        pdist = np.abs(hashes - tensor_hash)
+        assert(pdist.min() < 0.01)
+        matched = pdist.argmin()
+        target_nbits = hashed_recipe[hashes[matched]]
+
+        if target_nbits == 16:
+            continue
+
+        op_name_configs[name] = cto.OpPalettizerConfig(
+            mode="kmeans",
+            nbits=target_nbits,
+            weight_threshold=int(MIN_SIZE)
+        )
+
+    config = ct.optimize.coreml.OptimizationConfig(op_name_configs=op_name_configs)
+    coreml_model = ct.optimize.coreml.palettize_weights(coreml_model, config)
+
+    coreml_model.save(args.o)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "-o",
+        required=True,
+        help="Output directory to save the custom palettized model"
+    )
+    parser.add_argument(
+        "--mlpackage-path",
+        required=True,
+        help="Path to .mlpackage model to be palettized"
+    )
+    parser.add_argument(
+        "--pre-analysis-json-path",
+        required=True,
+        type=str,
+        help=("The JSON file generated by mixed_bit_compression_pre_analysis.py"
+    ))
+    parser.add_argument(
+        "--selected-recipe",
+        required=True,
+        type=str,
+        help=("The string key into --pre-analysis-json-path's baselines dict"
+    ))
+    parser.add_argument(
+        "--custom-vae-version",
+        type=str,
+        default=None,
+        help=
+        ("Custom VAE checkpoint to override the pipeline's built-in VAE. "
+            "If specified, the specified VAE will be converted instead of the one associated to the `--model-version` checkpoint. "
+            "No precision override is applied when using a custom VAE."
+    ))
+
+    args = parser.parse_args()
+
+    if not os.path.exists(args.mlpackage_path):
+        raise FileNotFoundError
+    if not os.path.exists(args.pre_analysis_json_path):
+        raise FileNotFoundError
+    if not args.pre_analysis_json_path.endswith('.json'):
+        raise ValueError("--recipe-json-path should end with '.json'")
+
+    main(args)

python_coreml_stable_diffusion/mixed_bit_compression_pre_analysis.py

@@ -0,0 +1,583 @@
+from collections import OrderedDict
+from copy import deepcopy
+from functools import partial
+import argparse
+import gc
+import json
+
+import logging
+logging.basicConfig()
+logger = logging.getLogger()
+logger.setLevel('INFO')
+
+import numpy as np
+import os
+from PIL import Image
+from python_coreml_stable_diffusion.torch2coreml import compute_psnr, get_pipeline
+import time
+
+import torch
+import torch.nn as nn
+import requests
+torch.set_grad_enabled(False)
+
+from tqdm import tqdm
+
+# Bit-widths the Neural Engine is capable of accelerating
+NBITS = [1, 2, 4, 6, 8]
+
+# Minimum number of elements in a weight tensor to be considered for palettization
+# (saves pre-analysis time)
+PALETTIZE_MIN_SIZE = 1e5
+
+# Signal integrity is computed based on these 4 random prompts
+RANDOM_TEST_DATA = [
+    "a black and brown dog standing outside a door.",
+    "a person on a motorcycle makes a turn on the track.",
+    "inflatable boats sit on the arizona river, and on the bank",
+    "a white cat sitting under a white umbrella",
+    "black bear standing in a field of grass under a tree.",
+    "a train that is parked on tracks and has graffiti writing on it, with a mountain range in the background.",
+    "a cake inside of a pan sitting in an oven.",
+    "a table with paper plates and flowers in a home",
+]
+
+TEST_RESOLUTION = 768
+
+RANDOM_TEST_IMAGE_DATA = [
+    Image.open(
+        requests.get(path, stream=True).raw).convert("RGB").resize(
+            (TEST_RESOLUTION, TEST_RESOLUTION), Image.LANCZOS
+    ) for path in [
+        "http://farm1.staticflickr.com/106/298138827_19bb723252_z.jpg",
+        "http://farm4.staticflickr.com/3772/9666116202_648cd752d6_z.jpg",
+        "http://farm3.staticflickr.com/2238/2472574092_f5534bb2f7_z.jpg",
+        "http://farm1.staticflickr.com/220/475442674_47d81fdc2c_z.jpg",
+        "http://farm8.staticflickr.com/7231/7359341784_4c5358197f_z.jpg",
+        "http://farm8.staticflickr.com/7283/8737653089_d0c77b8597_z.jpg",
+        "http://farm3.staticflickr.com/2454/3989339438_2f32b76ebb_z.jpg",
+        "http://farm1.staticflickr.com/34/123005230_13051344b1_z.jpg",
+]]
+
+
+# Copied from https://github.com/apple/coremltools/blob/7.0b1/coremltools/optimize/coreml/_quantization_passes.py#L602
+from coremltools.converters.mil.mil import types
+def fake_linear_quantize(val, axis=-1, mode='LINEAR', dtype=types.int8):
+    from coremltools.optimize.coreml._quantization_passes import AffineQuantParams
+    from coremltools.converters.mil.mil.types.type_mapping import nptype_from_builtin
+
+    val_dtype = val.dtype
+    def _ensure_numerical_range_and_cast(val, low, high, np_dtype):
+        '''
+        For some cases, the computed quantized data might exceed the data range.
+        For instance, after rounding and addition, we might get `128` for the int8 quantization.
+        This utility function ensures the val in the data range before doing the cast.
+        '''
+        val = np.minimum(val, high)
+        val = np.maximum(val, low)
+        return val.astype(np_dtype)
+
+    mode_dtype_to_range = {
+        (types.int8, "LINEAR"): (-128, 127),
+        (types.int8, "LINEAR_SYMMETRIC"): (-127, 127),
+        (types.uint8, "LINEAR"): (0, 255),
+        (types.uint8, "LINEAR_SYMMETRIC"): (0, 254),
+    }
+
+    if not isinstance(val, (np.ndarray, np.generic)):
+        raise ValueError("Only numpy arrays are supported")
+
+    params = AffineQuantParams()
+    axes = tuple([i for i in range(len(val.shape)) if i != axis])
+    val_min = np.amin(val, axis=axes, keepdims=True)
+    val_max = np.amax(val, axis=axes, keepdims=True)
+
+    if mode == "LINEAR_SYMMETRIC":
+        # For the linear_symmetric mode, the range is symmetrical to 0
+        max_abs = np.maximum(np.abs(val_min), np.abs(val_max))
+        val_min = -max_abs
+        val_max = max_abs
+    else:
+        assert mode == "LINEAR"
+        # For the linear mode, we need to make sure the data range contains `0`
+        val_min = np.minimum(0.0, val_min)
+        val_max = np.maximum(0.0, val_max)
+
+    q_val_min, q_val_max = mode_dtype_to_range[(dtype, mode)]
+
+    # Set the zero point to symmetric mode
+    np_dtype = nptype_from_builtin(dtype)
+    if mode == "LINEAR_SYMMETRIC":
+        if dtype == types.int8:
+            params.zero_point = (0 * np.ones(val_min.shape)).astype(np.int8)
+        else:
+            assert dtype == types.uint8
+            params.zero_point = (127 * np.ones(val_min.shape)).astype(np.uint8)
+    else:
+        assert mode == "LINEAR"
+        params.zero_point = (q_val_min * val_max - q_val_max * val_min) / (val_max - val_min)
+        params.zero_point = np.round(params.zero_point)
+        params.zero_point = _ensure_numerical_range_and_cast(params.zero_point, q_val_min, q_val_max, np_dtype)
+
+    # compute the params
+    params.scale = (val_max - val_min) / (q_val_max - q_val_min)
+    params.scale = params.scale.astype(val.dtype).squeeze()
+
+    params.quantized_data = np.round(
+        val * (q_val_max - q_val_min) / (val_max - val_min)
+    )
+    params.quantized_data = (params.quantized_data + params.zero_point)
+    params.quantized_data = _ensure_numerical_range_and_cast(params.quantized_data, q_val_min, q_val_max, np_dtype)
+
+    params.zero_point = params.zero_point.squeeze()
+    params.axis = axis
+
+    return (params.quantized_data.astype(val_dtype) - params.zero_point.astype(val_dtype)) * params.scale
+
+
+# Copied from https://github.com/apple/coremltools/blob/7.0b1/coremltools/optimize/coreml/_quantization_passes.py#L423
+def fake_palettize(module, nbits, in_ngroups=1, out_ngroups=1):
+    """ Simulate weight palettization
+    """
+    from coremltools.models.neural_network.quantization_utils import _get_kmeans_lookup_table_and_weight
+
+    def compress_kmeans(val, nbits):
+        lut, indices = _get_kmeans_lookup_table_and_weight(nbits, val)
+        lut = lut.astype(val.dtype)
+        indices = indices.astype(np.uint8)
+        return lut, indices
+
+    dtype = module.weight.data.dtype
+    device = module.weight.data.device
+    val = module.weight.data.cpu().numpy().astype(np.float16)
+    
+    if out_ngroups == 1 and in_ngroups == 1:
+        lut, indices = compress_kmeans(val=val, nbits=nbits)
+        module.weight.data = torch.from_numpy(lut[indices]).reshape(val.shape).to(dtype)
+
+    elif out_ngroups > 1 and in_ngroups == 1:
+        assert val.shape[0] % out_ngroups == 0
+        rvals = [
+            compress_kmeans(val=chunked_val, nbits=nbits)
+            for chunked_val in np.split(val, out_ngroups, axis=0)
+        ]
+        shape = list(val.shape)
+        shape[0] = shape[0] // out_ngroups
+        module.weight.data = torch.cat([
+            torch.from_numpy(lut[indices]).reshape(shape)
+            for lut,indices in rvals
+        ], dim=0).to(dtype).to(device)
+ 
+    elif in_ngroups > 1 and out_ngroups == 1:
+        assert val.shape[1] % in_ngroups == 0
+        rvals = [
+            compress_kmeans(val=chunked_val, nbits=nbits)
+            for chunked_val in np.split(val, in_ngroups, axis=1)
+        ]
+        shape = list(val.shape)
+        shape[1] = shape[1] // in_ngroups
+        module.weight.data = torch.cat([
+            torch.from_numpy(lut[indices]).reshape(shape)
+            for lut,indices in rvals
+        ], dim=1).to(dtype).to(device)
+    else:
+        raise ValueError(f"in_ngroups={in_ngroups} & out_ngroups={out_ngroups} is illegal!!!")
+    
+    return torch.from_numpy(val).to(dtype)
+
+
+def restore_weight(module, value):
+    device = module.weight.data.device
+    module.weight.data = value.to(device)
+
+
+def get_palettizable_modules(unet, min_size=PALETTIZE_MIN_SIZE):
+    ret = [
+        (name, getattr(module, 'weight').data.numel()) for name, module in unet.named_modules()
+        if isinstance(module, (nn.Linear, nn.Conv2d))
+        if hasattr(module, 'weight') and getattr(module, 'weight').data.numel() > min_size
+    ]
+    candidates, sizes = [[a for a,b in ret], [b for a,b in ret]]
+    logger.info(f"{len(candidates)} candidate tensors with {sum(sizes)/1e6} M total params")
+    return candidates, sizes
+
+
+def fake_int8_quantize(module):
+    i = 0
+    for name, submodule in tqdm(module.named_modules()):
+        if hasattr(submodule, 'weight'):
+            i+=1
+            submodule.weight.data = torch.from_numpy(
+                fake_linear_quantize(submodule.weight.data.numpy()))
+    logger.info(f"{i} modules fake int8 quantized")
+    return module
+
+
+def fake_nbits_palette(module, nbits):
+    i = 0
+    for name, submodule in tqdm(module.named_modules()):
+        if hasattr(submodule, 'weight'):
+            i+=1
+            fake_palettize(submodule, nbits=nbits)
+    logger.info(f"{i} modules fake {nbits}-bits palettized")
+    return module
+
+
+def fake_palette_from_recipe(module, recipe):
+    tot_bits = 0
+    tot_numel = 0
+    for name, submodule in tqdm(module.named_modules()):
+        if hasattr(submodule, 'weight'):
+            tot_numel += submodule.weight.numel()
+            if name in recipe:
+                nbits = recipe[name]
+                assert nbits in NBITS + [16]
+                tot_bits += submodule.weight.numel() * nbits
+                if nbits == 16:
+                    continue
+                fake_palettize(submodule, nbits=nbits)
+            else:
+                tot_bits += submodule.weight.numel() * 16
+
+    logger.info(f"Palettized to {tot_bits/tot_numel:.2f}-bits mixed palette ({tot_bits/8e6} MB) ")
+
+# Globally synced RNG state
+rng = torch.Generator()
+rng_state = rng.get_state()
+
+def run_pipe(pipe):
+    if torch.backends.mps.is_available():
+        device = "mps"
+    elif torch.cuda.is_available():
+        device = "cuda"
+    else:
+        device = "cpu"
+    logger.debug(f"Placing pipe in {device}")
+
+    global rng, rng_state
+    rng.set_state(rng_state)
+    kwargs = dict(
+        prompt=RANDOM_TEST_DATA,
+        negative_prompt=[""] * len(RANDOM_TEST_DATA),
+        num_inference_steps=1,
+        height=TEST_RESOLUTION,
+        width=TEST_RESOLUTION,
+        output_type="latent",
+        generator=rng
+    )
+    if "Img2Img" in pipe.__class__.__name__:
+        kwargs["image"] = RANDOM_TEST_IMAGE_DATA
+        kwargs.pop("height")
+        kwargs.pop("width")
+
+        # Run a single denoising step
+        kwargs["num_inference_steps"] = 4
+        kwargs["strength"] = 0.25
+
+    return np.array([latent.cpu().numpy() for latent in pipe.to(device)(**kwargs).images])
+
+
+def benchmark_signal_integrity(pipe,
+                               candidates,
+                               nbits,
+                               cumulative,
+                               in_ngroups=1,
+                               out_ngroups=1,
+                               ref_out=None,
+                               ):
+    results = {}
+    results['metadata'] = {
+        'nbits': nbits,
+        'out_ngroups': out_ngroups,
+        'in_ngroups': in_ngroups,
+        'cumulative': cumulative,
+    }
+
+    # If reference outputs are not provided, treat current pipe as reference
+    if ref_out is None:
+        ref_out = run_pipe(pipe)
+
+    for candidate in tqdm(candidates):
+        palettized = False
+        for name, module in pipe.unet.named_modules():
+            if name == candidate:
+                orig_weight = fake_palettize(
+                    module,
+                    nbits,
+                    out_ngroups=out_ngroups,
+                    in_ngroups=in_ngroups,
+                )
+                palettized = True
+                break
+        if not palettized:
+            raise KeyError(name)
+
+        test_out = run_pipe(pipe)
+
+        if not cumulative:
+            restore_weight(module, orig_weight)
+
+        results[candidate] = [
+            float(f"{compute_psnr(r,t):.1f}")
+            for r,t in zip(ref_out, test_out)
+        ]
+        logger.info(f"{nbits}-bit: {candidate} = {results[candidate]}")
+
+    return results
+            
+
+def descending_psnr_order(results):
+    if 'metadata' in results:
+        results.pop('metadata')
+
+    return OrderedDict(sorted(results.items(), key=lambda items: -sum(items[1])))
+
+
+def simulate_quant_fn(ref_pipe, quantization_to_simulate):
+    simulated_pipe = deepcopy(ref_pipe.to('cpu'))
+    quantization_to_simulate(simulated_pipe.unet)
+    simulated_out = run_pipe(simulated_pipe)
+    del simulated_pipe
+    gc.collect()
+
+    ref_out = run_pipe(ref_pipe)
+    simulated_psnr = sum([
+        float(f"{compute_psnr(r, t):.1f}")
+        for r, t in zip(ref_out, simulated_out)
+    ]) / len(ref_out)
+
+    return simulated_out, simulated_psnr
+
+
+def build_recipe(results, sizes, psnr_threshold, default_nbits):
+    stats = {'nbits': 0}
+    recipe = {}
+
+    for key in results[str(NBITS[0])]:
+        if key == 'metadata':
+            continue
+
+        achieved_nbits = default_nbits
+        for nbits in NBITS:
+            avg_psnr = sum(results[str(nbits)][key])/len(RANDOM_TEST_DATA)
+            if avg_psnr > psnr_threshold:
+                achieved_nbits = nbits
+                break
+        recipe[key] = achieved_nbits
+        stats['nbits'] += achieved_nbits * sizes[key]
+
+    stats['size_mb'] = stats['nbits'] / (8*1e6)
+    tot_size = sum(list(sizes.values()))
+    stats['nbits'] /= tot_size
+
+    return recipe, stats
+
+
+def plot(results, args):
+    import matplotlib.pyplot as plt
+    max_model_size = sum(results['cumulative'][str(NBITS[0])]['metadata']['sizes'])
+    f, ax = plt.subplots(1, 1, figsize=(7, 5))
+
+    def compute_x_axis(sizes, nbits, default_nbits):
+        max_compression_percent = (default_nbits - nbits) / default_nbits
+        progress = np.cumsum(sizes)
+        normalized_progress = progress / progress.max()
+
+        return normalized_progress * max_compression_percent * 100
+
+    # Linear 8-bit baseline and the intercept points for mixed-bit recipes
+    linear8bit_baseline = results['baselines']['linear_8bit']
+
+    # Mark the linear 8-bit baseline
+    ax.plot(
+        8 / args.default_nbits * 100,
+        linear8bit_baseline,
+        'bx',
+        markersize=8,
+        label="8-bit (linear quant)")
+
+    # Plot the iso-dB line that matches the 8-bit baseline
+    ax.plot([0,100], [linear8bit_baseline]*2, '--b')
+
+    # Plot non-mixed-bit palettization curves
+    for idx, nbits in enumerate(NBITS):
+        size_keys = compute_x_axis(results['cumulative'][str(nbits)]['metadata']['sizes'], nbits, args.default_nbits)
+        psnr = [
+            sum(v) / len(RANDOM_TEST_DATA) # avg psnr
+            for k,v in results['cumulative'][str(nbits)].items() if k != 'metadata'
+        ]
+        ax.plot(
+            size_keys,
+            psnr,
+            label=f"{nbits}-bit")
+
+
+    # Plot mixed-bit results
+    mixed_palettes = [
+        (float(spec.rsplit('_')[1]), psnr)
+        for spec,psnr in results['baselines'].items()
+        if 'recipe' in spec
+    ]
+    mixedbit_sizes = [100. * (1. - a[0] / args.default_nbits) for a in mixed_palettes]
+    mixedbit_psnrs = [a[1] for a in mixed_palettes]
+    ax.plot(
+        mixedbit_sizes,
+        mixedbit_psnrs,
+        label="mixed-bit",
+    )
+
+    ax.set_xlabel("Model Size Reduction (%)")
+    ax.set_ylabel("Signal Integrity (PSNR in dB)")
+    ax.set_title(args.model_version)
+    ax.legend()
+
+    f.savefig(os.path.join(args.o, f"{args.model_version.replace('/','_')}_psnr_vs_size.png"))
+
+def main(args):
+
+    # Initialize pipe
+    pipe = get_pipeline(args)
+
+    # Preserve a pristine copy for reference outputs
+    ref_pipe = deepcopy(pipe)
+    if args.default_nbits != 16:
+        logger.info(f"Palettizing unet to default {args.default_nbits}-bit")
+        fake_nbits_palette(pipe.unet, args.default_nbits)
+        logger.info("Done.")
+
+    # Cache reference outputs
+    ref_out = run_pipe(pipe)
+
+    # Bookkeeping
+    os.makedirs(args.o, exist_ok=True)
+
+    results = {
+        'single_layer': {},
+        'cumulative': {},
+        'model_version': args.model_version,
+    }
+    json_name = f"{args.model_version.replace('/','-')}_palettization_recipe.json"
+    candidates, sizes = get_palettizable_modules(pipe.unet)
+
+    sizes_table = dict(zip(candidates, sizes))
+
+    if os.path.isfile(os.path.join(args.o, json_name)):
+        with open(os.path.join(args.o, json_name), "r") as f:
+            results = json.load(f)
+
+    # Analyze uniform-precision palettization impact on signal integrity
+    for nbits in NBITS:
+        if str(nbits) not in results['single_layer']:
+            # Measure the impact of palettization of each layer independently
+            results['single_layer'][str(nbits)] = benchmark_signal_integrity(
+                pipe,
+                candidates,
+                nbits,
+                cumulative=False,
+                ref_out=ref_out,
+            )
+            with open(os.path.join(args.o, json_name), 'w') as f:
+                json.dump(results, f, indent=2)
+
+        # Measure the cumulative impact of palettization based on ascending individual impact computed earlier
+        sorted_candidates = descending_psnr_order(results['single_layer'][str(nbits)])
+
+        if str(nbits) not in results['cumulative']:
+            results['cumulative'][str(nbits)] = benchmark_signal_integrity(
+                deepcopy(pipe),
+                sorted_candidates,
+                nbits,
+                cumulative=True,
+                ref_out=ref_out,
+            )
+            results['cumulative'][str(nbits)]['metadata'].update({
+                'candidates': list(sorted_candidates.keys()),
+                'sizes': [sizes_table[candidate] for candidate in sorted_candidates],
+            })
+
+            with open(os.path.join(args.o, json_name), 'w') as f:
+                json.dump(results, f, indent=2)
+
+    # Generate uniform-quantization baselines
+    results['baselines'] = {
+        "original": simulate_quant_fn(ref_pipe, lambda x: x)[1],
+        "linear_8bit": simulate_quant_fn(ref_pipe, fake_int8_quantize)[1],
+    }
+    with open(os.path.join(args.o, json_name), 'w') as f:
+        json.dump(results, f, indent=2)
+
+
+    # Generate mixed-bit recipes via decreasing PSNR thresholds
+    results['recipes'] = {}
+    recipe_psnr_thresholds = np.linspace(
+        results['baselines']['original'] - 1,
+        results['baselines']["linear_8bit"] + 5,
+        args.num_recipes,
+    )
+
+    for recipe_no, psnr_threshold in enumerate(recipe_psnr_thresholds):
+        logger.info(f"Building recipe #{recipe_no}")
+        recipe, stats = build_recipe(
+            results['cumulative'],
+            sizes_table,
+            psnr_threshold,
+            args.default_nbits,
+        )
+        achieved_psnr = simulate_quant_fn(ref_pipe, lambda x: partial(fake_palette_from_recipe, recipe=recipe)(x))[1]
+        logger.info(
+            f"Recipe #{recipe_no}: {stats['nbits']:.2f}-bits @ per-layer {psnr_threshold} dB, "
+            f"end-to-end {achieved_psnr} dB & "
+            f"{stats['size_mb']:.2f} MB"
+        )
+
+        # Save achieved PSNR and compressed size
+        recipe_key = f"recipe_{stats['nbits']:.2f}_bit_mixedpalette"
+        results['baselines'][recipe_key] = float(f"{achieved_psnr:.1f}")
+        results['recipes'][recipe_key] = recipe
+
+        with open(os.path.join(args.o, json_name), 'w') as f:
+            json.dump(results, f, indent=2)
+
+    # Plot model size vs signal integrity
+    plot(results, args)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "-o",
+        required=True,
+        help="Output directory to save the palettization artifacts (recipe json, PSNR plots etc.)"
+    )
+    parser.add_argument(
+        "--model-version",
+        required=True,
+        help=
+        ("The pre-trained model checkpoint and configuration to restore. "
+         "For available versions: https://huggingface.co/models?search=stable-diffusion"
+     ))
+    parser.add_argument(
+        "--default-nbits",
+        help="Default number of bits to use for palettization",
+        choices=tuple(NBITS + [16]),
+        default=16,
+        type=int,
+    )
+    parser.add_argument(
+        "--num-recipes",
+        help="Maximum number of recipes to generate (with decreasing model size and signal integrity)",
+        default=7,
+        type=int,
+    )
+    parser.add_argument(
+        "--custom-vae-version",
+        type=str,
+        default=None,
+        help=
+        ("Custom VAE checkpoint to override the pipeline's built-in VAE. "
+         "If specified, the specified VAE will be converted instead of the one associated to the `--model-version` checkpoint. "
+         "No precision override is applied when using a custom VAE."
+         ))
+
+    args = parser.parse_args()
+    main(args)

python_coreml_stable_diffusion/pipeline.py

@@ -5,7 +5,8 @@
 
 import argparse
 
-from diffusers.pipeline_utils import DiffusionPipeline
+from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
 from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
 from diffusers.schedulers import (
     DDIMScheduler,
@@ -39,7 +40,7 @@ from python_coreml_stable_diffusion.coreml_model import (
 import time
 import torch  # Only used for `torch.from_tensor` in `pipe.scheduler.step()`
 from transformers import CLIPFeatureExtractor, CLIPTokenizer
-from typing import List, Optional, Union
+from typing import List, Optional, Union, Tuple
 from PIL import Image
 
 
@@ -49,20 +50,27 @@ class CoreMLStableDiffusionPipeline(DiffusionPipeline):
     """
 
     def __init__(
-        self,
-        text_encoder: CoreMLModel,
-        unet: CoreMLModel,
-        vae_decoder: CoreMLModel,
-        feature_extractor: CLIPFeatureExtractor,
-        safety_checker: Optional[CoreMLModel],
-        scheduler: Union[DDIMScheduler,
-                         DPMSolverMultistepScheduler,
-                         EulerAncestralDiscreteScheduler,
-                         EulerDiscreteScheduler,
-                         LMSDiscreteScheduler,
-                         PNDMScheduler],
-        tokenizer: CLIPTokenizer,
-        controlnet: Optional[List[CoreMLModel]],
+            self,
+            text_encoder: CoreMLModel,
+            unet: CoreMLModel,
+            vae_decoder: CoreMLModel,
+            scheduler: Union[
+                DDIMScheduler,
+                DPMSolverMultistepScheduler,
+                EulerAncestralDiscreteScheduler,
+                EulerDiscreteScheduler,
+                LMSDiscreteScheduler,
+                PNDMScheduler
+            ],
+            tokenizer: CLIPTokenizer,
+            controlnet: Optional[List[CoreMLModel]],
+            xl: Optional[bool] = False,
+            force_zeros_for_empty_prompt: Optional[bool] = True,
+            feature_extractor: Optional[CLIPFeatureExtractor] = None,
+            safety_checker: Optional[CoreMLModel] = None,
+            text_encoder_2: Optional[CoreMLModel] = None,
+            tokenizer_2: Optional[CLIPTokenizer] = None
+
     ):
         super().__init__()
 
@@ -84,10 +92,14 @@ class CoreMLStableDiffusionPipeline(DiffusionPipeline):
                 " it only for use-cases that involve analyzing network behavior or auditing its results. For more"
                 " information, please have a look at https://github.com/huggingface/diffusers/pull/254 ."
             )
+        self.xl = xl
+        self.force_zeros_for_empty_prompt = force_zeros_for_empty_prompt
 
         # Register Core ML components of the pipeline
         self.safety_checker = safety_checker
         self.text_encoder = text_encoder
+        self.text_encoder_2 = text_encoder_2
+        self.tokenizer_2 = tokenizer_2
         self.unet = unet
         self.unet.in_channels = self.unet.expected_inputs["sample"]["shape"][1]
 
@@ -108,76 +120,147 @@ class CoreMLStableDiffusionPipeline(DiffusionPipeline):
             f"Stable Diffusion configured to generate {self.height}x{self.width} images"
         )
 
-    def _encode_prompt(self, prompt, num_images_per_prompt,
-                       do_classifier_free_guidance, negative_prompt):
+    def _encode_prompt(self,
+                       prompt,
+                       prompt_2: Optional[str] = None,
+                       do_classifier_free_guidance: bool = True,
+                       negative_prompt: Optional[str] = None,
+                       negative_prompt_2: Optional[str] = None,
+                       ):
+
         batch_size = len(prompt) if isinstance(prompt, list) else 1
 
-        text_inputs = self.tokenizer(
-            prompt,
-            padding="max_length",
-            max_length=self.tokenizer.model_max_length,
-            return_tensors="np",
-        )
-        text_input_ids = text_inputs.input_ids
+        if self.xl is True:
+            prompts = [prompt, prompt_2] if prompt_2 is not None else [prompt, prompt]
 
-        if text_input_ids.shape[-1] > self.tokenizer.model_max_length:
-            removed_text = self.tokenizer.batch_decode(
-                text_input_ids[:, self.tokenizer.model_max_length:])
-            logger.warning(
-                "The following part of your input was truncated because CLIP can only handle sequences up to"
-                f" {self.tokenizer.model_max_length} tokens: {removed_text}")
-            text_input_ids = text_input_ids[:, :self.tokenizer.
-                                            model_max_length]
+            # refiner uses only one tokenizer and text encoder (tokenizer_2 and text_encoder_2)
+            tokenizers = [self.tokenizer, self.tokenizer_2] if self.tokenizer is not None else [self.tokenizer_2]
 
-        text_embeddings = self.text_encoder(
-            input_ids=text_input_ids.astype(np.float32))["last_hidden_state"]
+            text_encoders = [self.text_encoder, self.text_encoder_2] if self.text_encoder is not None else [
+                self.text_encoder_2]
+            hidden_state_key = 'hidden_embeds'
+        else:
+            prompts = [prompt]
+            tokenizers = [self.tokenizer]
+            text_encoders = [self.text_encoder]
+            hidden_state_key = 'last_hidden_state'
+
+        prompt_embeds_list = []
+        for prompt, tokenizer, text_encoder in zip(prompts, tokenizers, text_encoders):
+
+            text_inputs = tokenizer(
+                prompt,
+                padding="max_length",
+                max_length=tokenizer.model_max_length,
+                truncation=True,
+                return_tensors="np",
+            )
+
+            text_input_ids = text_inputs.input_ids
+
+            # tokenize without max_length to catch any truncation
+            untruncated_ids = tokenizer(prompt, padding="longest", return_tensors="np").input_ids
+
+            if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not np.equal(
+                    text_input_ids, untruncated_ids
+            ):
+                removed_text = tokenizer.batch_decode(untruncated_ids[:, tokenizer.model_max_length - 1: -1])
+                logger.warning(
+                    "The following part of your input was truncated because CLIP can only handle sequences up to"
+                    f" {tokenizer.model_max_length} tokens: {removed_text}"
+                )
+
+            embeddings = text_encoder(input_ids=text_input_ids.astype(np.float32))
+
+            prompt_embeds_list.append(embeddings[hidden_state_key])
+
+            # We are only ALWAYS interested in the pooled output of the final text encoder
+            if self.xl:
+                pooled_prompt_embeds = embeddings['pooled_outputs']
+
+        prompt_embeds = np.concatenate(prompt_embeds_list, axis=-1)
+
+        if do_classifier_free_guidance and negative_prompt is None and self.force_zeros_for_empty_prompt:
+            negative_prompt_embeds = np.zeros_like(prompt_embeds)
+
+            if self.xl:
+                negative_pooled_prompt_embeds = np.zeros_like(pooled_prompt_embeds)
+        elif do_classifier_free_guidance:
+
+            negative_prompt = negative_prompt or ""
+            negative_prompt_2 = negative_prompt_2 or negative_prompt
+
+            # normalize str to list
+            negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
+            negative_prompt_2 = (
+                batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2
+            )
 
-        if do_classifier_free_guidance:
             uncond_tokens: List[str]
-            if negative_prompt is None:
-                uncond_tokens = [""] * batch_size
-            elif type(prompt) is not type(negative_prompt):
+            if prompts is not None and type(prompts) is not type(negative_prompt):
                 raise TypeError(
-                    "`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
-                    " {type(prompt)}.")
-            elif isinstance(negative_prompt, str):
-                uncond_tokens = [negative_prompt] * batch_size
+                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
+                    f" {type(prompt)}."
+                )
             elif batch_size != len(negative_prompt):
                 raise ValueError(
                     f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
                     f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
                     " the batch size of `prompt`.")
             else:
-                uncond_tokens = negative_prompt
+                uncond_tokens = [negative_prompt, negative_prompt_2]
 
-            max_length = text_input_ids.shape[-1]
-            uncond_input = self.tokenizer(
-                uncond_tokens,
-                padding="max_length",
-                max_length=max_length,
-                truncation=True,
-                return_tensors="np",
-            )
+            negative_prompt_embeds_list = []
+            for negative_prompt, tokenizer, text_encoder in zip(uncond_tokens, tokenizers, text_encoders):
 
-            uncond_embeddings = self.text_encoder(
-                input_ids=uncond_input.input_ids.astype(
-                    np.float32))["last_hidden_state"]
+                max_length = prompt_embeds.shape[1]
 
+                uncond_input = tokenizer(
+                    negative_prompt,
+                    padding="max_length",
+                    max_length=max_length,
+                    truncation=True,
+                    return_tensors="np",
+                )
+                uncond_input_ids = uncond_input.input_ids
+
+                negative_embeddings = text_encoder(
+                    input_ids=uncond_input_ids.astype(np.float32)
+                )
+
+                negative_text_embeddings = negative_embeddings[hidden_state_key]
+
+                negative_prompt_embeds_list.append(negative_text_embeddings)
+
+                # We are only ALWAYS interested in the pooled output of the final text encoder
+                if self.xl:
+                    negative_pooled_prompt_embeds = negative_embeddings['pooled_outputs']
+
+            negative_prompt_embeds = np.concatenate(negative_prompt_embeds_list, axis=-1)
+
+        if do_classifier_free_guidance:
             # For classifier free guidance, we need to do two forward passes.
             # Here we concatenate the unconditional and text embeddings into a single batch
             # to avoid doing two forward passes
-            text_embeddings = np.concatenate(
-                [uncond_embeddings, text_embeddings])
+            prompt_embeds = np.concatenate(
+                [negative_prompt_embeds, prompt_embeds])
 
-        text_embeddings = text_embeddings.transpose(0, 2, 1)[:, :, None, :]
+            if self.xl:
+                pooled_prompt_embeds = np.concatenate(
+                    [negative_pooled_prompt_embeds, pooled_prompt_embeds])
 
-        return text_embeddings
+        prompt_embeddings = prompt_embeds.transpose(0, 2, 1)[:, :, None, :]
 
-    def run_controlnet(self, 
-                       sample, 
-                       timestep, 
-                       encoder_hidden_states, 
-                       controlnet_cond, 
+        if self.xl:
+            return prompt_embeddings, pooled_prompt_embeds
+        else:
+            return prompt_embeddings, None
+
+    def run_controlnet(self,
+                       sample,
+                       timestep,
+                       encoder_hidden_states,
+                       controlnet_cond,
                        output_dtype=np.float16):
         if not self.controlnet:
             raise ValueError(
@@ -185,9 +268,9 @@ class CoreMLStableDiffusionPipeline(DiffusionPipeline):
 
         for i, (module, cond) in enumerate(zip(self.controlnet, controlnet_cond)):
             module_outputs = module(
-                sample=sample.astype(np.float16), 
-                timestep=timestep.astype(np.float16), 
-                encoder_hidden_states=encoder_hidden_states.astype(np.float16), 
+                sample=sample.astype(np.float16),
+                timestep=timestep.astype(np.float16),
+                encoder_hidden_states=encoder_hidden_states.astype(np.float16),
                 controlnet_cond=cond.astype(np.float16),
             )
             if i == 0:
@@ -195,9 +278,9 @@ class CoreMLStableDiffusionPipeline(DiffusionPipeline):
             else:
                 for key in outputs.keys():
                     outputs[key] += module_outputs[key]
-        
+
         outputs = {k: v.astype(output_dtype) for k, v in outputs.items()}
-        
+
         return outputs
 
     def run_safety_checker(self, image):
@@ -229,7 +312,8 @@ class CoreMLStableDiffusionPipeline(DiffusionPipeline):
 
     def decode_latents(self, latents):
         latents = 1 / 0.18215 * latents
-        image = self.vae_decoder(z=latents.astype(np.float16))["image"]
+        dtype = self.vae_decoder.expected_inputs['z']['dtype']
+        image = self.vae_decoder(z=latents.astype(dtype))["image"]
         image = np.clip(image / 2 + 0.5, 0, 1)
         image = image.transpose((0, 2, 3, 1))
 
@@ -250,14 +334,19 @@ class CoreMLStableDiffusionPipeline(DiffusionPipeline):
                 f"Unexpected latents shape, got {latents.shape}, expected {latents_shape}"
             )
 
-        latents = latents * self.scheduler.init_noise_sigma
+        init_noise = self.scheduler.init_noise_sigma
+
+        if isinstance(init_noise, torch.Tensor):
+            init_noise = init_noise.numpy()
+
+        latents = latents * init_noise
 
         return latents
 
-    def prepare_control_cond(self, 
-                             controlnet_cond, 
-                             do_classifier_free_guidance, 
-                             batch_size, 
+    def prepare_control_cond(self,
+                             controlnet_cond,
+                             do_classifier_free_guidance,
+                             batch_size,
                              num_images_per_prompt):
         processed_cond_list = []
         for cond in controlnet_cond:
@@ -306,26 +395,39 @@ class CoreMLStableDiffusionPipeline(DiffusionPipeline):
 
         return extra_step_kwargs
 
+    def _get_add_time_ids(self, original_size, crops_coords_top_left, target_size, dtype):
+        add_time_ids = list(original_size + crops_coords_top_left + target_size)
+        add_time_ids = np.array(add_time_ids).astype(dtype)
+        return add_time_ids
+
     def __call__(
-        self,
-        prompt,
-        height=512,
-        width=512,
-        num_inference_steps=50,
-        guidance_scale=7.5,
-        negative_prompt=None,
-        num_images_per_prompt=1,
-        eta=0.0,
-        latents=None,
-        output_type="pil",
-        return_dict=True,
-        callback=None,
-        callback_steps=1,
-        controlnet_cond=None,
-        **kwargs,
+            self,
+            prompt,
+            height=512,
+            width=512,
+            num_inference_steps=50,
+            guidance_scale=7.5,
+            negative_prompt=None,
+            num_images_per_prompt=1,
+            eta=0.0,
+            latents=None,
+            output_type="pil",
+            return_dict=True,
+            callback=None,
+            callback_steps=1,
+            controlnet_cond=None,
+            original_size: Optional[Tuple[int, int]] = None,
+            crops_coords_top_left: Tuple[int, int] = (0, 0),
+            target_size: Optional[Tuple[int, int]] = None,
+            **kwargs,
     ):
         # 1. Check inputs. Raise error if not correct
         self.check_inputs(prompt, height, width, callback_steps)
+        height = height or self.height
+        width = width or self.width
+
+        original_size = original_size or (height, width)
+        target_size = target_size or (height, width)
 
         # 2. Define call parameters
         batch_size = 1 if isinstance(prompt, str) else len(prompt)
@@ -340,18 +442,39 @@ class CoreMLStableDiffusionPipeline(DiffusionPipeline):
         do_classifier_free_guidance = guidance_scale > 1.0
 
         # 3. Encode input prompt
-        text_embeddings = self._encode_prompt(
-            prompt,
-            num_images_per_prompt,
-            do_classifier_free_guidance,
-            negative_prompt,
+        text_embeddings, pooled_prompt_embeds = self._encode_prompt(
+            prompt=prompt,
+            prompt_2=None,
+            do_classifier_free_guidance=do_classifier_free_guidance,
+            negative_prompt=negative_prompt,
+            negative_prompt_2=None
         )
 
-        # 4. Prepare timesteps
+        # 4. Prepare XL kwargs if needed
+        unet_additional_kwargs = {}
+
+        # we add pooled prompt embeds + time_ids to unet kwargs
+        if self.xl:
+            add_text_embeds = pooled_prompt_embeds
+            add_time_ids = self._get_add_time_ids(original_size, crops_coords_top_left, target_size,
+                                                  text_embeddings.dtype)
+            if do_classifier_free_guidance:
+
+                # TODO: This checks if the time_ids input is looking for time_ids.shape == (12,) or (2, 6)
+                # Remove once model input shapes are ubiquitous
+                if len(self.unet.expected_inputs['time_ids']['shape']) > 1:
+                    add_time_ids = [add_time_ids]
+
+                add_time_ids = np.concatenate([add_time_ids, add_time_ids])
+
+            unet_additional_kwargs.update({'text_embeds': add_text_embeds.astype(np.float16),
+                                           'time_ids': add_time_ids.astype(np.float16)})
+
+        # 5. Prepare timesteps
         self.scheduler.set_timesteps(num_inference_steps)
         timesteps = self.scheduler.timesteps
 
-        # 5. Prepare latent variables and controlnet cond
+        # 6. Prepare latent variables and controlnet cond
         num_channels_latents = self.unet.in_channels
         latents = self.prepare_latents(
             batch_size * num_images_per_prompt,
@@ -369,10 +492,10 @@ class CoreMLStableDiffusionPipeline(DiffusionPipeline):
                 num_images_per_prompt,
             )
 
-        # 6. Prepare extra step kwargs
+        # 7. Prepare extra step kwargs
         extra_step_kwargs = self.prepare_extra_step_kwargs(eta)
 
-        # 7. Denoising loop
+        # 8. Denoising loop
         for i, t in enumerate(self.progress_bar(timesteps)):
             # expand the latents if we are doing classifier free guidance
             latent_model_input = np.concatenate(
@@ -380,37 +503,42 @@ class CoreMLStableDiffusionPipeline(DiffusionPipeline):
             latent_model_input = self.scheduler.scale_model_input(
                 latent_model_input, t)
 
+            if isinstance(latent_model_input, torch.Tensor):
+                latent_model_input = latent_model_input.numpy()
+
             # controlnet
             if controlnet_cond:
-                additional_residuals = self.run_controlnet(
+                control_net_additional_residuals = self.run_controlnet(
                     sample=latent_model_input,
                     timestep=np.array([t, t]),
                     encoder_hidden_states=text_embeddings,
                     controlnet_cond=controlnet_cond,
                 )
             else:
-                additional_residuals = {}
+                control_net_additional_residuals = {}
 
             # predict the noise residual
+            unet_additional_kwargs.update(control_net_additional_residuals)
+
             noise_pred = self.unet(
                 sample=latent_model_input.astype(np.float16),
                 timestep=np.array([t, t], np.float16),
                 encoder_hidden_states=text_embeddings.astype(np.float16),
-                **additional_residuals,
+                **unet_additional_kwargs,
             )["noise_pred"]
 
             # perform guidance
             if do_classifier_free_guidance:
                 noise_pred_uncond, noise_pred_text = np.split(noise_pred, 2)
                 noise_pred = noise_pred_uncond + guidance_scale * (
-                    noise_pred_text - noise_pred_uncond)
+                        noise_pred_text - noise_pred_uncond)
 
             # compute the previous noisy sample x_t -> x_t-1
             latents = self.scheduler.step(torch.from_numpy(noise_pred),
                                           t,
                                           torch.from_numpy(latents),
                                           **extra_step_kwargs,
-            ).prev_sample.numpy()
+                                          ).prev_sample.numpy()
 
             # call the callback, if provided
             if callback is not None and i % callback_steps == 0:
@@ -444,18 +572,25 @@ def get_available_schedulers():
         schedulers[scheduler().__class__.__name__.replace("Scheduler", "")] = scheduler
     return schedulers
 
+
 SCHEDULER_MAP = get_available_schedulers()
 
+
 def get_coreml_pipe(pytorch_pipe,
                     mlpackages_dir,
                     model_version,
                     compute_unit,
                     delete_original_pipe=True,
                     scheduler_override=None,
-                    controlnet_models=None):
-    """ Initializes and returns a `CoreMLStableDiffusionPipeline` from an original
-    diffusers PyTorch pipeline
+                    controlnet_models=None,
+                    force_zeros_for_empty_prompt=True,
+                    sources=None):
     """
+    Initializes and returns a `CoreMLStableDiffusionPipeline` from an original
+    diffusers PyTorch pipeline
+        sources: 'packages' or 'compiled' forces creation of model from specified sources. sources must be in mlpackages_dir
+    """
+
     # Ensure `scheduler_override` object is of correct type if specified
     if scheduler_override is not None:
         assert isinstance(scheduler_override, SchedulerMixin)
@@ -464,15 +599,27 @@ def get_coreml_pipe(pytorch_pipe,
             f"Default={pytorch_pipe.scheduler}, Override={scheduler_override}")
 
     # Gather configured tokenizer and scheduler attributes from the original pipe
-    coreml_pipe_kwargs = {
-        "tokenizer": pytorch_pipe.tokenizer,
-        "scheduler": pytorch_pipe.scheduler if scheduler_override is None else scheduler_override,
-        "feature_extractor": pytorch_pipe.feature_extractor,
-    }
+    if 'xl' in model_version:
+        coreml_pipe_kwargs = {
+            "tokenizer": pytorch_pipe.tokenizer,
+            'tokenizer_2': pytorch_pipe.tokenizer_2,
+            "scheduler": pytorch_pipe.scheduler if scheduler_override is None else scheduler_override,
+            "force_zeros_for_empty_prompt": force_zeros_for_empty_prompt,
+            'xl': True
+        }
+
+        model_packages_to_load = ["text_encoder", "text_encoder_2", "unet", "vae_decoder"]
+
+    else:
+        coreml_pipe_kwargs = {
+            "tokenizer": pytorch_pipe.tokenizer,
+            "scheduler": pytorch_pipe.scheduler if scheduler_override is None else scheduler_override,
+            "feature_extractor": pytorch_pipe.feature_extractor,
+        }
+        model_packages_to_load = ["text_encoder", "unet", "vae_decoder"]
 
-    model_names_to_load = ["text_encoder", "unet", "vae_decoder"]
     if getattr(pytorch_pipe, "safety_checker", None) is not None:
-        model_names_to_load.append("safety_checker")
+        model_packages_to_load.append("safety_checker")
     else:
         logger.warning(
             f"Original diffusers pipeline for {model_version} does not have a safety_checker, "
@@ -485,16 +632,16 @@ def get_coreml_pipe(pytorch_pipe,
         logger.info("Removed PyTorch pipe to reduce peak memory consumption")
 
     if controlnet_models:
-        model_names_to_load.remove("unet")
+        model_packages_to_load.remove("unet")
         coreml_pipe_kwargs["unet"] = _load_mlpackage(
-            "control-unet",
-            mlpackages_dir,
-            model_version,
-            compute_unit,
+            submodule_name="control-unet",
+            mlpackages_dir=mlpackages_dir,
+            model_version=model_version,
+            compute_unit=compute_unit,
         )
         coreml_pipe_kwargs["controlnet"] = [_load_mlpackage_controlnet(
-            mlpackages_dir, 
-            model_version, 
+            mlpackages_dir,
+            model_version,
             compute_unit,
         ) for model_version in controlnet_models]
     else:
@@ -504,12 +651,13 @@ def get_coreml_pipe(pytorch_pipe,
     logger.info(f"Loading Core ML models in memory from {mlpackages_dir}")
     coreml_pipe_kwargs.update({
         model_name: _load_mlpackage(
-            model_name,
-            mlpackages_dir,
-            model_version,
-            compute_unit,
+            submodule_name=model_name,
+            mlpackages_dir=mlpackages_dir,
+            model_version=model_version,
+            compute_unit=compute_unit,
+            sources=sources,
         )
-        for model_name in model_names_to_load
+        for model_name in model_packages_to_load
     })
     logger.info("Done.")
 
@@ -536,32 +684,48 @@ def get_image_path(args, **override_kwargs):
 
     return os.path.join(out_folder, out_fname + ".png")
 
+
 def prepare_controlnet_cond(image_path, height, width):
     image = Image.open(image_path).convert("RGB")
     image = image.resize((height, width), resample=Image.LANCZOS)
     image = np.array(image).transpose(2, 0, 1) / 255.0
     return image
 
+
 def main(args):
     logger.info(f"Setting random seed to {args.seed}")
     np.random.seed(args.seed)
 
     logger.info("Initializing PyTorch pipe for reference configuration")
-    from diffusers import StableDiffusionPipeline
-    pytorch_pipe = StableDiffusionPipeline.from_pretrained(args.model_version,
-                                                           use_auth_token=True)
 
+    SDP = StableDiffusionXLPipeline if 'xl' in args.model_version else StableDiffusionPipeline
+
+    pytorch_pipe = SDP.from_pretrained(
+        args.model_version,
+        use_auth_token=True,
+    )
+
+    # Get Scheduler
     user_specified_scheduler = None
     if args.scheduler is not None:
         user_specified_scheduler = SCHEDULER_MAP[
             args.scheduler].from_config(pytorch_pipe.scheduler.config)
 
-    coreml_pipe = get_coreml_pipe(pytorch_pipe=pytorch_pipe,
-                                  mlpackages_dir=args.i,
-                                  model_version=args.model_version,
-                                  compute_unit=args.compute_unit,
-                                  scheduler_override=user_specified_scheduler,
-                                  controlnet_models=args.controlnet)
+    # Get Force Zeros Config if it exists
+    force_zeros_for_empty_prompt: bool = False
+    if 'force_zeros_for_empty_prompt' in pytorch_pipe.config:
+        force_zeros_for_empty_prompt = pytorch_pipe.config['force_zeros_for_empty_prompt']
+
+    coreml_pipe = get_coreml_pipe(
+        pytorch_pipe=pytorch_pipe,
+        mlpackages_dir=args.i,
+        model_version=args.model_version,
+        compute_unit=args.compute_unit,
+        scheduler_override=user_specified_scheduler,
+        controlnet_models=args.controlnet,
+        force_zeros_for_empty_prompt=force_zeros_for_empty_prompt,
+        sources=args.model_sources,
+    )
 
     if args.controlnet:
         controlnet_cond = []
@@ -624,7 +788,7 @@ if __name__ == "__main__":
         choices=tuple(SCHEDULER_MAP.keys()),
         default=None,
         help=("The scheduler to use for running the reverse diffusion process. "
-             "If not specified, the default scheduler from the diffusers pipeline is utilized"))
+              "If not specified, the default scheduler from the diffusers pipeline is utilized"))
     parser.add_argument(
         "--num-inference-steps",
         default=50,
@@ -637,20 +801,24 @@ if __name__ == "__main__":
         help="Controls the influence of the text prompt on sampling process (0=random images)")
     parser.add_argument(
         "--controlnet",
-        nargs="*", 
+        nargs="*",
         type=str,
         help=("Enables ControlNet and use control-unet instead of unet for additional inputs. "
-            "For Multi-Controlnet, provide the model names separated by spaces."))
+              "For Multi-Controlnet, provide the model names separated by spaces."))
     parser.add_argument(
         "--controlnet-inputs",
-        nargs="*", 
+        nargs="*",
         type=str,
         help=("Image paths for ControlNet inputs. "
-            "Please enter images corresponding to each controlnet provided at --controlnet option in same order."))
+              "Please enter images corresponding to each controlnet provided at --controlnet option in same order."))
     parser.add_argument(
         "--negative-prompt",
         default=None,
         help="The negative text prompt to be used for text-to-image generation.")
+    parser.add_argument('--model-sources',
+                        default=None,
+                        choices=['packages', 'compiled'],
+                        help='Force build from `packages` or `compiled`')
 
     args = parser.parse_args()
     main(args)

python_coreml_stable_diffusion/torch2coreml.py

@@ -11,7 +11,11 @@ import argparse
 from collections import OrderedDict, defaultdict
 from copy import deepcopy
 import coremltools as ct
-from diffusers import StableDiffusionPipeline, ControlNetModel
+from diffusers import (
+    StableDiffusionPipeline,
+    DiffusionPipeline,
+    ControlNetModel
+)
 import gc
 
 import logging
@@ -93,20 +97,14 @@ def _get_out_path(args, submodule_name):
     return os.path.join(args.o, fname)
 
 
-# https://github.com/apple/coremltools/issues/1680
-def _save_mlpackage(model, output_path):
-    # First recreate MLModel object using its in memory spec, then save
-    ct.models.MLModel(model._spec,
-                      weights_dir=model._weights_dir,
-                      is_temp_package=True).save(output_path)
-
-
 def _convert_to_coreml(submodule_name, torchscript_module, sample_inputs,
-                       output_names, args, out_path=None):
+                       output_names, args, out_path=None, precision=None, compute_unit=None):
 
     if out_path is None:
         out_path = _get_out_path(args, submodule_name)
 
+    compute_unit = compute_unit or ct.ComputeUnit[args.compute_unit]
+
     if os.path.exists(out_path):
         logger.info(f"Skipping export because {out_path} already exists")
         logger.info(f"Loading model from {out_path}")
@@ -116,11 +114,11 @@ def _convert_to_coreml(submodule_name, torchscript_module, sample_inputs,
         # The Swifty CLI we provide uses precompiled Core ML models (.mlmodelc) which incurs compilation only
         # upon first load and mitigates the load time in subsequent runs.
         coreml_model = ct.models.MLModel(
-            out_path, compute_units=ct.ComputeUnit[args.compute_unit])
+            out_path, compute_units=compute_unit)
         logger.info(
             f"Loading {out_path} took {time.time() - start:.1f} seconds")
 
-        coreml_model.compute_unit = ct.ComputeUnit[args.compute_unit]
+        coreml_model.compute_unit = compute_unit
     else:
         logger.info(f"Converting {submodule_name} to CoreML..")
         coreml_model = ct.convert(
@@ -128,24 +126,22 @@ def _convert_to_coreml(submodule_name, torchscript_module, sample_inputs,
             convert_to="mlprogram",
             minimum_deployment_target=ct.target.macOS13,
             inputs=_get_coreml_inputs(sample_inputs, args),
-            outputs=[ct.TensorType(name=name) for name in output_names],
-            compute_units=ct.ComputeUnit[args.compute_unit],
-            # skip_model_load=True,
+            outputs=[ct.TensorType(name=name, dtype=np.float32) for name in output_names],
+            compute_units=compute_unit,
+            compute_precision=precision,
+            skip_model_load=not args.check_output_correctness,
         )
 
         del torchscript_module
         gc.collect()
 
-        coreml_model.save(out_path)
-        logger.info(f"Saved {submodule_name} model to {out_path}")
-
     return coreml_model, out_path
 
 
 def quantize_weights(args):
     """ Quantize weights to args.quantize_nbits using a palette (look-up table)
     """
-    for model_name in ["text_encoder", "unet", "control-unet"]:
+    for model_name in ["text_encoder", "text_encoder_2", "unet", "refiner", "control-unet"]:
         logger.info(f"Quantizing {model_name} to {args.quantize_nbits}-bit precision")
         out_path = _get_out_path(args, model_name)
         _quantize_weights(
@@ -223,11 +219,15 @@ def bundle_resources_for_swift_cli(args):
 
     # Compile model using coremlcompiler (Significantly reduces the load time for unet)
     for source_name, target_name in [("text_encoder", "TextEncoder"),
+                                     ("text_encoder_2", "TextEncoder2"),
                                      ("vae_decoder", "VAEDecoder"),
                                      ("vae_encoder", "VAEEncoder"),
                                      ("unet", "Unet"),
                                      ("unet_chunk1", "UnetChunk1"),
                                      ("unet_chunk2", "UnetChunk2"),
+                                     ("refiner", "UnetRefiner"),
+                                     ("refiner_chunk1", "UnetRefinerChunk1"),
+                                     ("refiner_chunk2", "UnetRefinerChunk2"),
                                      ("control-unet", "ControlledUnet"),
                                      ("control-unet_chunk1", "ControlledUnetChunk1"),
                                      ("control-unet_chunk2", "ControlledUnetChunk2"),
@@ -274,24 +274,42 @@ def bundle_resources_for_swift_cli(args):
     return resources_dir
 
 
-def convert_text_encoder(pipe, args):
+from transformers.models.clip import modeling_clip
+
+# Copied from https://github.com/huggingface/transformers/blob/v4.30.0/src/transformers/models/clip/modeling_clip.py#L677C1-L692C1
+def patched_make_causal_mask(input_ids_shape, dtype, device, past_key_values_length: int = 0):
+    """ Patch to replace torch.finfo(dtype).min with -1e4
+    """
+    bsz, tgt_len = input_ids_shape
+    mask = torch.full((tgt_len, tgt_len), torch.tensor(-1e4, device=device), device=device)
+    mask_cond = torch.arange(mask.size(-1), device=device)
+    mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)
+    mask = mask.to(dtype)
+
+    if past_key_values_length > 0:
+        mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype, device=device), mask], dim=-1)
+    return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length)
+
+modeling_clip._make_causal_mask = patched_make_causal_mask
+
+def convert_text_encoder(text_encoder, tokenizer, submodule_name, args):
     """ Converts the text encoder component of Stable Diffusion
     """
-    out_path = _get_out_path(args, "text_encoder")
+    text_encoder = text_encoder.to(dtype=torch.float32)
+    out_path = _get_out_path(args, submodule_name)
     if os.path.exists(out_path):
         logger.info(
-            f"`text_encoder` already exists at {out_path}, skipping conversion."
+            f"`{submodule_name}` already exists at {out_path}, skipping conversion."
         )
         return
 
     # Create sample inputs for tracing, conversion and correctness verification
-    text_encoder_sequence_length = pipe.tokenizer.model_max_length
-    text_encoder_hidden_size = args.text_encoder_hidden_size or pipe.text_encoder.config.hidden_size
+    text_encoder_sequence_length = tokenizer.model_max_length
 
     sample_text_encoder_inputs = {
         "input_ids":
         torch.randint(
-            pipe.text_encoder.config.vocab_size,
+            text_encoder.config.vocab_size,
             (1, text_encoder_sequence_length),
             # https://github.com/apple/coremltools/issues/1423
             dtype=torch.float32,
@@ -303,41 +321,50 @@ def convert_text_encoder(pipe, args):
     }
     logger.info(f"Sample inputs spec: {sample_text_encoder_inputs_spec}")
 
-    def _build_causal_attention_mask(self, bsz, seq_len, dtype, device=None):
-        mask = torch.ones((bsz, seq_len, seq_len), dtype=dtype, device=device) * -1e4
-        mask.triu_(1)
-        mask = mask.unsqueeze(1)
-        return mask
-
     class TextEncoder(nn.Module):
 
-        def __init__(self):
+        def __init__(self, with_hidden_states_for_layer=None):
             super().__init__()
-            self.text_encoder = pipe.text_encoder
-            setattr(
-                self.text_encoder.text_model, "_build_causal_attention_mask",
-                MethodType(_build_causal_attention_mask,
-                           self.text_encoder.text_model))
+            self.text_encoder = text_encoder
+            self.with_hidden_states_for_layer = with_hidden_states_for_layer
 
         def forward(self, input_ids):
-            return self.text_encoder(input_ids, return_dict=False)
+            if self.with_hidden_states_for_layer is not None:
+                output = self.text_encoder(input_ids, output_hidden_states=True)
+                hidden_embeds = output.hidden_states[self.with_hidden_states_for_layer]
+                if "text_embeds" in output:
+                    return (hidden_embeds, output.text_embeds)
+                else:
+                    return (hidden_embeds, output.pooler_output)
+            else:
+                return self.text_encoder(input_ids, return_dict=False)
 
-    reference_text_encoder = TextEncoder().eval()
+    # SD XL uses the hidden states after the encoder layers from both encoders,
+    # and the pooled `text_embeds` output of the second encoder.
+    hidden_layer = -2 if args.xl_version else None
+    reference_text_encoder = TextEncoder(with_hidden_states_for_layer=hidden_layer).eval()
 
-    logger.info("JIT tracing text_encoder..")
+    logger.info(f"JIT tracing {submodule_name}..")
     reference_text_encoder = torch.jit.trace(
         reference_text_encoder,
         (sample_text_encoder_inputs["input_ids"].to(torch.int32), ),
     )
     logger.info("Done.")
 
+    if args.xl_version:
+        output_names = ["hidden_embeds", "pooled_outputs"]
+    else:
+        output_names = ["last_hidden_state", "pooled_outputs"]
     coreml_text_encoder, out_path = _convert_to_coreml(
-        "text_encoder", reference_text_encoder, sample_text_encoder_inputs,
-        ["last_hidden_state", "pooled_outputs"], args)
+        submodule_name, reference_text_encoder, sample_text_encoder_inputs,
+        output_names, args)
 
     # Set model metadata
     coreml_text_encoder.author = f"Please refer to the Model Card available at huggingface.co/{args.model_version}"
-    coreml_text_encoder.license = "OpenRAIL (https://huggingface.co/spaces/CompVis/stable-diffusion-license)"
+    if args.xl_version:
+        coreml_text_encoder.license = "OpenRAIL++-M (https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/LICENSE.md)"
+    else:
+        coreml_text_encoder.license = "OpenRAIL (https://huggingface.co/spaces/CompVis/stable-diffusion-license)"
     coreml_text_encoder.version = args.model_version
     coreml_text_encoder.short_description = \
         "Stable Diffusion generates images conditioned on text and/or other images as input through the diffusion process. " \
@@ -348,31 +375,41 @@ def convert_text_encoder(pipe, args):
         "input_ids"] = "The token ids that represent the input text"
 
     # Set the output descriptions
-    coreml_text_encoder.output_description[
-        "last_hidden_state"] = "The token embeddings as encoded by the Transformer model"
+    if args.xl_version:
+        coreml_text_encoder.output_description[
+            "hidden_embeds"] = "Hidden states after the encoder layers"
+    else:
+        coreml_text_encoder.output_description[
+            "last_hidden_state"] = "The token embeddings as encoded by the Transformer model"
     coreml_text_encoder.output_description[
         "pooled_outputs"] = "The version of the `last_hidden_state` output after pooling"
 
-    _save_mlpackage(coreml_text_encoder, out_path)
+    coreml_text_encoder.save(out_path)
 
     logger.info(f"Saved text_encoder into {out_path}")
 
     # Parity check PyTorch vs CoreML
     if args.check_output_correctness:
-        baseline_out = pipe.text_encoder(
+        baseline_out = text_encoder(
             sample_text_encoder_inputs["input_ids"].to(torch.int32),
-            return_dict=False,
-        )[1].numpy()
+            output_hidden_states=args.xl_version,
+            return_dict=True,
+        )
+        if args.xl_version:
+            # TODO: maybe check pooled_outputs too
+            baseline_out = baseline_out.hidden_states[hidden_layer].numpy()
+        else:
+            baseline_out = baseline_out.last_hidden_state.numpy()
 
-        coreml_out = list(
-            coreml_text_encoder.predict(
-                {k: v.numpy()
-                 for k, v in sample_text_encoder_inputs.items()}).values())[0]
+        coreml_out = coreml_text_encoder.predict(
+            {k: v.numpy() for k, v in sample_text_encoder_inputs.items()}
+        )
+        coreml_out = coreml_out["hidden_embeds" if args.xl_version else "last_hidden_state"]
         report_correctness(
             baseline_out, coreml_out,
             "text_encoder baseline PyTorch to reference CoreML")
 
-    del reference_text_encoder, coreml_text_encoder, pipe.text_encoder
+    del reference_text_encoder, coreml_text_encoder
     gc.collect()
 
 
@@ -443,11 +480,21 @@ def convert_vae_decoder(pipe, args):
         1,  # B
         pipe.vae.config.latent_channels,  # C
         args.latent_h or pipe.unet.config.sample_size,  # H
-        args.latent_w or pipe.unet.config.sample_size,  # w
+        args.latent_w or pipe.unet.config.sample_size,  # W
     )
 
+    if args.custom_vae_version is None and args.xl_version:
+        inputs_dtype = torch.float32
+        compute_precision = ct.precision.FLOAT32
+        # FIXME: Hardcoding to CPU_AND_GPU since ANE doesn't support FLOAT32
+        compute_unit = ct.ComputeUnit.CPU_AND_GPU
+    else:
+        inputs_dtype = torch.float16
+        compute_precision = None
+        compute_unit = None
+
     sample_vae_decoder_inputs = {
-        "z": torch.rand(*z_shape, dtype=torch.float16)
+        "z": torch.rand(*z_shape, dtype=inputs_dtype)
     }
 
     class VAEDecoder(nn.Module):
@@ -456,10 +503,8 @@ def convert_vae_decoder(pipe, args):
 
         def __init__(self):
             super().__init__()
-            self.post_quant_conv = pipe.vae.post_quant_conv
-            self.decoder = pipe.vae.decoder
-            # Disable torch 2.0 scaled dot-product attention: https://github.com/apple/coremltools/issues/1823
-            self.decoder.mid_block.attentions[0]._use_2_0_attn = False
+            self.post_quant_conv = pipe.vae.post_quant_conv.to(dtype=torch.float32)
+            self.decoder = pipe.vae.decoder.to(dtype=torch.float32)
 
         def forward(self, z):
             return self.decoder(self.post_quant_conv(z))
@@ -473,11 +518,14 @@ def convert_vae_decoder(pipe, args):
     modify_coremltools_torch_frontend_badbmm()
     coreml_vae_decoder, out_path = _convert_to_coreml(
         "vae_decoder", traced_vae_decoder, sample_vae_decoder_inputs,
-        ["image"], args)
+        ["image"], args, precision=compute_precision, compute_unit=compute_unit)
 
     # Set model metadata
     coreml_vae_decoder.author = f"Please refer to the Model Card available at huggingface.co/{args.model_version}"
-    coreml_vae_decoder.license = "OpenRAIL (https://huggingface.co/spaces/CompVis/stable-diffusion-license)"
+    if args.xl_version:
+        coreml_vae_decoder.license = "OpenRAIL++-M (https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/LICENSE.md)"
+    else:
+        coreml_vae_decoder.license = "OpenRAIL (https://huggingface.co/spaces/CompVis/stable-diffusion-license)"
     coreml_vae_decoder.version = args.model_version
     coreml_vae_decoder.short_description = \
         "Stable Diffusion generates images conditioned on text and/or other images as input through the diffusion process. " \
@@ -491,7 +539,7 @@ def convert_vae_decoder(pipe, args):
     coreml_vae_decoder.output_description[
         "image"] = "Generated image normalized to range [-1, 1]"
 
-    _save_mlpackage(coreml_vae_decoder, out_path)
+    coreml_vae_decoder.save(out_path)
 
     logger.info(f"Saved vae_decoder into {out_path}")
 
@@ -528,15 +576,25 @@ def convert_vae_encoder(pipe, args):
     height = (args.latent_h or pipe.unet.config.sample_size) * 8
     width = (args.latent_w or pipe.unet.config.sample_size) * 8
     
-    z_shape = (
+    x_shape = (
         1,  # B
         3,  # C (RGB range from -1 to 1)
         height,  # H
         width,  # w
     )
 
+    if args.xl_version:
+        inputs_dtype = torch.float32
+        compute_precision = ct.precision.FLOAT32
+        # FIXME: Hardcoding to CPU_AND_GPU since ANE doesn't support FLOAT32
+        compute_unit = ct.ComputeUnit.CPU_AND_GPU
+    else:
+        inputs_dtype = torch.float16
+        compute_precision = None
+        compute_unit = None
+
     sample_vae_encoder_inputs = {
-        "z": torch.rand(*z_shape, dtype=torch.float16)
+        "x": torch.rand(*x_shape, dtype=inputs_dtype)
     }
 
     class VAEEncoder(nn.Module):
@@ -545,48 +603,49 @@ def convert_vae_encoder(pipe, args):
 
         def __init__(self):
             super().__init__()
-            self.quant_conv = pipe.vae.quant_conv
-            self.encoder = pipe.vae.encoder
-            # Disable torch 2.0 scaled dot-product attention: https://github.com/apple/coremltools/issues/1823
-            self.encoder.mid_block.attentions[0]._use_2_0_attn = False
+            self.quant_conv = pipe.vae.quant_conv.to(dtype=torch.float32)
+            self.encoder = pipe.vae.encoder.to(dtype=torch.float32)
 
-        def forward(self, z):
-            return self.quant_conv(self.encoder(z))
+        def forward(self, x):
+            return self.quant_conv(self.encoder(x))
 
     baseline_encoder = VAEEncoder().eval()
 
     # No optimization needed for the VAE Encoder as it is a pure ConvNet
     traced_vae_encoder = torch.jit.trace(
-        baseline_encoder, (sample_vae_encoder_inputs["z"].to(torch.float32), ))
+        baseline_encoder, (sample_vae_encoder_inputs["x"].to(torch.float32), ))
 
     modify_coremltools_torch_frontend_badbmm()
     coreml_vae_encoder, out_path = _convert_to_coreml(
         "vae_encoder", traced_vae_encoder, sample_vae_encoder_inputs,
-        ["latent"], args)
+        ["latent"], args, precision=compute_precision, compute_unit=compute_unit)
 
     # Set model metadata
     coreml_vae_encoder.author = f"Please refer to the Model Card available at huggingface.co/{args.model_version}"
-    coreml_vae_encoder.license = "OpenRAIL (https://huggingface.co/spaces/CompVis/stable-diffusion-license)"
+    if args.xl_version:
+        coreml_vae_encoder.license = "OpenRAIL++-M (https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/LICENSE.md)"
+    else:
+        coreml_vae_encoder.license = "OpenRAIL (https://huggingface.co/spaces/CompVis/stable-diffusion-license)"
     coreml_vae_encoder.version = args.model_version
     coreml_vae_encoder.short_description = \
         "Stable Diffusion generates images conditioned on text and/or other images as input through the diffusion process. " \
         "Please refer to https://arxiv.org/abs/2112.10752 for details."
 
     # Set the input descriptions
-    coreml_vae_encoder.input_description["z"] = \
+    coreml_vae_encoder.input_description["x"] = \
         "The input image to base the initial latents on normalized to range [-1, 1]"
 
     # Set the output descriptions
     coreml_vae_encoder.output_description["latent"] = "The latent embeddings from the unet model from the input image."
 
-    _save_mlpackage(coreml_vae_encoder, out_path)
+    coreml_vae_encoder.save(out_path)
 
     logger.info(f"Saved vae_encoder into {out_path}")
 
     # Parity check PyTorch vs CoreML
     if args.check_output_correctness:
         baseline_out = baseline_encoder(
-            z=sample_vae_encoder_inputs["z"].to(torch.float32)).numpy()
+            x=sample_vae_encoder_inputs["x"].to(torch.float32)).numpy()
         coreml_out = list(
             coreml_vae_encoder.predict(
                 {k: v.numpy()
@@ -598,13 +657,13 @@ def convert_vae_encoder(pipe, args):
     gc.collect()
 
 
-def convert_unet(pipe, args):
+def convert_unet(pipe, args, model_name = None):
     """ Converts the UNet component of Stable Diffusion
     """
     if args.unet_support_controlnet:
         unet_name = "control-unet"
     else:
-        unet_name = "unet"
+        unet_name = model_name or "unet"
 
     out_path = _get_out_path(args, unet_name)
 
@@ -635,12 +694,19 @@ def convert_unet(pipe, args):
             raise RuntimeError(
                 "convert_text_encoder() deletes pipe.text_encoder to save RAM. "
                 "Please use convert_unet() before convert_text_encoder()")
+        
+        if hasattr(pipe, "text_encoder") and pipe.text_encoder is not None:
+            text_token_sequence_length = pipe.text_encoder.config.max_position_embeddings
+            hidden_size = pipe.text_encoder.config.hidden_size,
+        elif hasattr(pipe, "text_encoder_2") and pipe.text_encoder_2 is not None:
+            text_token_sequence_length = pipe.text_encoder_2.config.max_position_embeddings
+            hidden_size = pipe.text_encoder_2.config.hidden_size,
 
         encoder_hidden_states_shape = (
             batch_size,
-            args.text_encoder_hidden_size or pipe.text_encoder.config.hidden_size,
+            args.text_encoder_hidden_size or pipe.unet.config.cross_attention_dim or hidden_size,
             1,
-            args.text_token_sequence_length or pipe.text_encoder.config.max_position_embeddings,
+            args.text_token_sequence_length or text_token_sequence_length,
         )
 
         # Create the scheduled timesteps for downstream use
@@ -662,7 +728,51 @@ def convert_unet(pipe, args):
                 "encoder_hidden_states"].squeeze(2).transpose(1, 2)
 
         # Initialize reference unet
-        reference_unet = unet.UNet2DConditionModel(**pipe.unet.config).eval()
+        if args.xl_version:
+            unet_cls = unet.UNet2DConditionModelXL
+
+            # Sample time_ids
+            height = (args.latent_h or pipe.unet.config.sample_size) * 8
+            width = (args.latent_w or pipe.unet.config.sample_size) * 8
+
+            original_size = (height, width) # output_resolution
+            crops_coords_top_left = (0, 0) # topleft_crop_cond
+            target_size = (height, width) # resolution_cond
+            if hasattr(pipe.config, "requires_aesthetics_score") and pipe.config.requires_aesthetics_score:
+                # Part of SDXL's micro-conditioning as explained in section 2.2 of
+                # [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). Can be used to
+                # simulate an aesthetic score of the generated image by influencing the positive and negative text conditions.
+                aesthetic_score = 6.0 # default aesthetic_score
+                negative_aesthetic_score = 2.5 # default negative_aesthetic_score
+                add_time_ids = list(original_size + crops_coords_top_left + (aesthetic_score,))
+                add_neg_time_ids = list(original_size + crops_coords_top_left + (negative_aesthetic_score,))
+            else:
+                add_time_ids = list(original_size + crops_coords_top_left + target_size)
+                add_neg_time_ids = list(original_size + crops_coords_top_left + target_size)
+
+            time_ids = [
+                add_neg_time_ids,
+                add_time_ids
+            ]
+
+            # Pooled text embedding from text_encoder_2
+            text_embeds_shape = (
+                batch_size,
+                pipe.text_encoder_2.config.hidden_size
+            )
+
+            additional_xl_inputs = OrderedDict([
+                ("time_ids", torch.tensor(time_ids).to(torch.float32)),
+                ("text_embeds", torch.rand(*text_embeds_shape)),
+            ])
+
+            sample_unet_inputs.update(additional_xl_inputs)
+            baseline_sample_unet_inputs['added_cond_kwargs'] = additional_xl_inputs
+        else:
+            unet_cls = unet.UNet2DConditionModel
+
+        reference_unet = unet_cls(**pipe.unet.config).eval()
+
         load_state_dict_summary = reference_unet.load_state_dict(
             pipe.unet.state_dict())
 
@@ -717,7 +827,7 @@ def convert_unet(pipe, args):
         logger.info("Done.")
 
         if args.check_output_correctness:
-            baseline_out = pipe.unet(**baseline_sample_unet_inputs,
+            baseline_out = pipe.unet.to(torch.float32)(**baseline_sample_unet_inputs,
                                      return_dict=False)[0].numpy()
             reference_out = reference_unet(*sample_unet_inputs.values())[0].numpy()
             report_correctness(baseline_out, reference_out,
@@ -731,6 +841,7 @@ def convert_unet(pipe, args):
             for k, v in sample_unet_inputs.items()
         }
 
+
         coreml_unet, out_path = _convert_to_coreml(unet_name, reference_unet,
                                                    coreml_sample_unet_inputs,
                                                    ["noise_pred"], args)
@@ -739,8 +850,11 @@ def convert_unet(pipe, args):
 
         # Set model metadata
         coreml_unet.author = f"Please refer to the Model Card available at huggingface.co/{args.model_version}"
-        coreml_unet.license = "OpenRAIL (https://huggingface.co/spaces/CompVis/stable-diffusion-license)"
-        coreml_unet.version = args.model_version
+        if args.xl_version:
+            coreml_unet.license = "OpenRAIL++-M (https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/LICENSE.md)"
+        else:
+            coreml_unet.license = "OpenRAIL (https://huggingface.co/spaces/CompVis/stable-diffusion-license)"
+        coreml_unet.version = args.model_version if model_name != "refiner" or not hasattr(args, "refiner_version") else args.refiner_version
         coreml_unet.short_description = \
             "Stable Diffusion generates images conditioned on text or other images as input through the diffusion process. " \
             "Please refer to https://arxiv.org/abs/2112.10752 for details."
@@ -754,6 +868,11 @@ def convert_unet(pipe, args):
             "Output embeddings from the associated text_encoder model to condition to generated image on text. " \
             "A maximum of 77 tokens (~40 words) are allowed. Longer text is truncated. " \
             "Shorter text does not reduce computation."
+        if args.xl_version:
+            coreml_unet.input_description["time_ids"] = \
+                "Additional embeddings that if specified are added to the embeddings that are passed along to the UNet blocks."
+            coreml_unet.input_description["text_embeds"] = \
+                "Additional embeddings from text_encoder_2 that if specified are added to the embeddings that are passed along to the UNet blocks."
 
         # Set the output descriptions
         coreml_unet.output_description["noise_pred"] = \
@@ -764,7 +883,7 @@ def convert_unet(pipe, args):
         from python_coreml_stable_diffusion._version import __version__
         coreml_unet.user_defined_metadata["com.github.apple.ml-stable-diffusion.version"] = __version__
 
-        _save_mlpackage(coreml_unet, out_path)
+        coreml_unet.save(out_path)
         logger.info(f"Saved unet into {out_path}")
 
         # Parity check PyTorch vs CoreML
@@ -783,9 +902,10 @@ def convert_unet(pipe, args):
             f"`unet` already exists at {out_path}, skipping conversion.")
 
     if args.chunk_unet and not unet_chunks_exist:
-        logger.info("Chunking unet in two approximately equal MLModels")
+        logger.info(f"Chunking {model_name} in two approximately equal MLModels")
         args.mlpackage_path = out_path
         args.remove_original = False
+        args.merge_chunks_in_pipeline_model = False
         chunk_mlprogram.main(args)
 
 
@@ -806,6 +926,8 @@ def convert_safety_checker(pipe, args):
         )
         return
 
+    pipe.safety_checker = pipe.safety_checker.to(torch.float32)
+    
     im_h = pipe.vae.config.sample_size
     im_w = pipe.vae.config.sample_size
 
@@ -966,7 +1088,7 @@ def convert_safety_checker(pipe, args):
         "Concept scores are the scores before thresholding at zero yields the `has_nsfw_concepts` output. " \
         "These scores can be used to tune the `adjustment` input"
 
-    _save_mlpackage(coreml_safety_checker, out_path)
+    coreml_safety_checker.save(out_path)
 
     if args.check_output_correctness:
         baseline_out = pipe.safety_checker(
@@ -1138,7 +1260,7 @@ def convert_controlnet(pipe, args):
                 "One of the outputs of each downsampling block in ControlNet. " \
                 "The value added to the corresponding resnet output in UNet."
 
-        _save_mlpackage(coreml_controlnet, out_path)
+        coreml_controlnet.save(out_path)
         logger.info(f"Saved controlnet into {out_path}")
 
         # Parity check PyTorch vs CoreML
@@ -1153,16 +1275,36 @@ def convert_controlnet(pipe, args):
         del coreml_controlnet
         gc.collect()
 
+def get_pipeline(args):
+    model_version = args.model_version
+
+    logger.info(f"Initializing DiffusionPipeline with {model_version}..")
+    if args.custom_vae_version:
+        from diffusers import AutoencoderKL
+        vae = AutoencoderKL.from_pretrained(args.custom_vae_version, torch_dtype=torch.float16)
+        pipe = DiffusionPipeline.from_pretrained(model_version,
+                                            torch_dtype=torch.float16,
+                                            variant="fp16",
+                                            use_safetensors=True,
+                                            vae=vae,
+                                            use_auth_token=True)
+    else:
+        pipe = DiffusionPipeline.from_pretrained(model_version,
+                                            torch_dtype=torch.float16,
+                                            variant="fp16",
+                                            use_safetensors=True,
+                                            use_auth_token=True)
+
+    logger.info(f"Done. Pipeline in effect: {pipe.__class__.__name__}")
+
+    return pipe
+
 
 def main(args):
     os.makedirs(args.o, exist_ok=True)
 
     # Instantiate diffusers pipe as reference
-    logger.info(
-        f"Initializing StableDiffusionPipeline with {args.model_version}..")
-    pipe = StableDiffusionPipeline.from_pretrained(args.model_version,
-                                                   use_auth_token=True)
-    logger.info("Done.")
+    pipe = get_pipeline(args)
 
     # Register the selected attention implementation globally
     unet.ATTENTION_IMPLEMENTATION_IN_EFFECT = unet.AttentionImplementations[
@@ -1192,16 +1334,36 @@ def main(args):
         convert_unet(pipe, args)
         logger.info("Converted unet")
 
-    if args.convert_text_encoder:
+    if args.convert_text_encoder and hasattr(pipe, "text_encoder") and pipe.text_encoder is not None:
         logger.info("Converting text_encoder")
-        convert_text_encoder(pipe, args)
+        convert_text_encoder(pipe.text_encoder, pipe.tokenizer, "text_encoder", args)
+        del pipe.text_encoder
         logger.info("Converted text_encoder")
 
+    if args.convert_text_encoder and hasattr(pipe, "text_encoder_2") and pipe.text_encoder_2 is not None:
+        logger.info("Converting text_encoder_2")
+        convert_text_encoder(pipe.text_encoder_2, pipe.tokenizer_2, "text_encoder_2", args)
+        del pipe.text_encoder_2
+        logger.info("Converted text_encoder_2")
+
     if args.convert_safety_checker:
         logger.info("Converting safety_checker")
         convert_safety_checker(pipe, args)
         logger.info("Converted safety_checker")
 
+    if args.convert_unet and args.refiner_version is not None:
+        logger.info(f"Converting refiner")
+        del pipe
+        gc.collect()
+        original_model_version = args.model_version
+        args.model_version = args.refiner_version
+        pipe = get_pipeline(args)
+        args.model_version = original_model_version 
+        convert_unet(pipe, args, model_name="refiner")
+        del pipe
+        gc.collect()
+        logger.info(f"Converted refiner")
+
     if args.quantize_nbits is not None:
         logger.info(f"Quantizing weights to {args.quantize_nbits}-bit precision")
         quantize_weights(args)
@@ -1232,16 +1394,33 @@ def parser_spec():
     )
     parser.add_argument(
         "--model-version",
-        default="CompVis/stable-diffusion-v1-4",
+        required=True,
         help=
         ("The pre-trained model checkpoint and configuration to restore. "
          "For available versions: https://huggingface.co/models?search=stable-diffusion"
          ))
+    parser.add_argument(
+        "--refiner-version",
+        default=None,
+        help=
+        ("The pre-trained refiner model checkpoint and configuration to restore. "
+         "If specified, this argument will convert and bundle the refiner unet only alongside the model unet. "
+         "If you would like to convert a refiner model on it's own, use the --model-version argument instead."
+         "For available versions: https://huggingface.co/models?sort=trending&search=stable-diffusion+refiner"
+         ))
+    parser.add_argument(
+        "--custom-vae-version",
+        type=str,
+        default=None,
+        help=
+        ("Custom VAE checkpoint to override the pipeline's built-in VAE. "
+         "If specified, the specified VAE will be converted instead of the one associated to the `--model-version` checkpoint. "
+         "No precision override is applied when using a custom VAE."
+         ))
     parser.add_argument("--compute-unit",
                         choices=tuple(cu
                                       for cu in ct.ComputeUnit._member_names_),
                         default="ALL")
-
     parser.add_argument(
         "--latent-h",
         type=int,
@@ -1299,7 +1478,7 @@ def parser_spec():
     parser.add_argument(
         "--quantize-nbits",
         default=None,
-        choices=(2, 4, 6, 8),
+        choices=(1, 2, 4, 6, 8),
         type=int,
         help="If specified, quantized each model to nbits precision"
     )
@@ -1329,6 +1508,11 @@ def parser_spec():
         default=
         "https://huggingface.co/openai/clip-vit-base-patch32/resolve/main/merges.txt",
         help="The URL to the merged pairs used in by the text tokenizer.")
+    parser.add_argument(
+        "--xl-version",
+        action="store_true",
+        help=("If specified, the pre-trained model will be treated as an instantiation of "
+        "`diffusers.pipelines.StableDiffusionXLPipeline` instead of `diffusers.pipelines.StableDiffusionPipeline`"))
 
     return parser
 

python_coreml_stable_diffusion/unet.py

@@ -116,15 +116,8 @@ def linear_to_conv2d_map(state_dict, prefix, local_metadata, strict,
     """ Unsqueeze twice to map nn.Linear weights to nn.Conv2d weights
     """
     for k in state_dict:
-        is_internal_proj = all(substr in k for substr in ["to_", ".weight"])
-        is_ff_proj = all(substr in k for substr in ["ff.", ".weight"])
-        is_temb_proj = all(substr in k for substr in ["time_emb", ".weight"])
-        is_proj_in = "proj_in.weight" in k
-        is_proj_out = "proj_out.weight" in k
-
-        if is_internal_proj or is_ff_proj or is_temb_proj or is_proj_in or is_proj_out:
-            if len(state_dict[k].shape) == 2:
-                state_dict[k] = state_dict[k][:, :, None, None]
+        if 'weight' in k and len(state_dict[k].shape) == 2:
+            state_dict[k] = state_dict[k][:, :, None, None]
 
 # Note: torch.nn.LayerNorm and ane_transformers.reference.layer_norm.LayerNormANE
 # apply scale and bias terms in opposite orders. In order to accurately restore a
@@ -167,6 +160,7 @@ class CrossAttnUpBlock2D(nn.Module):
         output_scale_factor=1.0,
         downsample_padding=1,
         add_upsample=True,
+        transformer_layers_per_block=1,
     ):
         super().__init__()
         resnets = []
@@ -194,7 +188,7 @@ class CrossAttnUpBlock2D(nn.Module):
                     out_channels,
                     attn_num_head_channels,
                     out_channels // attn_num_head_channels,
-                    depth=1,
+                    depth=transformer_layers_per_block,
                     context_dim=cross_attention_dim,
                 ))
         self.attentions = nn.ModuleList(attentions)
@@ -285,6 +279,7 @@ class CrossAttnDownBlock2D(nn.Module):
         in_channels,
         out_channels,
         temb_channels,
+        transformer_layers_per_block=1,
         num_layers=1,
         resnet_eps=1e-6,
         resnet_time_scale_shift="default",
@@ -320,7 +315,7 @@ class CrossAttnDownBlock2D(nn.Module):
                     out_channels,
                     attn_num_head_channels,
                     out_channels // attn_num_head_channels,
-                    depth=1,
+                    depth=transformer_layers_per_block,
                     context_dim=cross_attention_dim,
                 ))
         self.attentions = nn.ModuleList(attentions)
@@ -360,6 +355,7 @@ class DownBlock2D(nn.Module):
         resnet_time_scale_shift="default",
         resnet_act_fn="swish",
         resnet_groups=32,
+        add_downsample=True,
     ):
         super().__init__()
         resnets = []
@@ -378,6 +374,11 @@ class DownBlock2D(nn.Module):
 
         self.resnets = nn.ModuleList(resnets)
 
+        if add_downsample:
+            self.downsamplers = nn.ModuleList([Downsample2D(out_channels)])
+        else:
+            self.downsamplers = None
+
     def forward(self, hidden_states, temb=None):
         output_states = ()
 
@@ -385,6 +386,13 @@ class DownBlock2D(nn.Module):
             hidden_states = resnet(hidden_states, temb)
             output_states += (hidden_states, )
 
+        if self.downsamplers is not None:
+            for downsampler in self.downsamplers:
+                hidden_states = downsampler(hidden_states)
+
+            output_states = output_states + (hidden_states,)
+
+
         return hidden_states, output_states
 
 
@@ -602,28 +610,68 @@ class GEGLU(nn.Module):
         return hidden_states * F.gelu(gate)
 
 
-class TimestepEmbedding(nn.Module):
+def get_activation(act_fn):
+    if act_fn in ["swish", "silu"]:
+        return nn.SiLU()
+    elif act_fn == "mish":
+        return nn.Mish()
+    elif act_fn == "gelu":
+        return nn.GELU()
+    else:
+        raise ValueError(f"Unsupported activation function: {act_fn}")
 
-    def __init__(self, channel, time_embed_dim, act_fn="silu"):
+class TimestepEmbedding(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        time_embed_dim,
+        act_fn = "silu",
+        out_dim = None,
+        post_act_fn = None,
+        cond_proj_dim=None,
+    ):
         super().__init__()
 
-        self.linear_1 = nn.Conv2d(channel, time_embed_dim, kernel_size=1)
-        self.act = None
-        if act_fn == "silu":
-            self.act = nn.SiLU()
-        self.linear_2 = nn.Conv2d(time_embed_dim,
-                                  time_embed_dim,
-                                  kernel_size=1)
+        self.linear_1 = nn.Conv2d(
+            in_channels, time_embed_dim, kernel_size=1)
 
-    def forward(self, sample):
+        if cond_proj_dim is not None:
+            self.cond_proj = nn.Conv2d(
+                cond_proj_dim, in_channels, kernel_size=1, bias=False)
+        else:
+            self.cond_proj = None
+
+        self.act = get_activation(act_fn)
+
+        if out_dim is not None:
+            time_embed_dim_out = out_dim
+        else:
+            time_embed_dim_out = time_embed_dim
+        self.linear_2 = nn.Conv2d(
+            time_embed_dim, time_embed_dim_out, kernel_size=1)
+
+        if post_act_fn is None:
+            self.post_act = None
+        else:
+            self.post_act = get_activation(post_act_fn)
+
+    def forward(self, sample, condition=None):
         if len(sample.shape) == 2:
             sample = sample.unsqueeze(-1).unsqueeze(-1)
+
+        if condition is not None:
+            if len(condition.shape) == 2:
+                condition = condition.unsqueeze(-1).unsqueeze(-1)
+            sample = sample + self.cond_proj(condition)
         sample = self.linear_1(sample)
 
         if self.act is not None:
             sample = self.act(sample)
 
         sample = self.linear_2(sample)
+
+        if self.post_act is not None:
+            sample = self.post_act(sample)
         return sample
 
 
@@ -687,6 +735,7 @@ class UNetMidBlock2DCrossAttn(nn.Module):
         attn_num_head_channels=1,
         attention_type="default",
         cross_attention_dim=768,
+        transformer_layers_per_block=1,
         **kwargs,
     ):
         super().__init__()
@@ -714,7 +763,7 @@ class UNetMidBlock2DCrossAttn(nn.Module):
                     in_channels,
                     attn_num_head_channels,
                     in_channels // attn_num_head_channels,
-                    depth=1,
+                    depth=transformer_layers_per_block,
                     context_dim=cross_attention_dim,
                 ))
             resnets.append(
@@ -756,6 +805,7 @@ class UNet2DConditionModel(ModelMixin, ConfigMixin):
             "CrossAttnDownBlock2D",
             "DownBlock2D",
         ),
+        mid_block_type="UNetMidBlock2DCrossAttn",
         up_block_types=("UpBlock2D", "CrossAttnUpBlock2D",
                         "CrossAttnUpBlock2D", "CrossAttnUpBlock2D"),
         only_cross_attention=False,
@@ -767,7 +817,11 @@ class UNet2DConditionModel(ModelMixin, ConfigMixin):
         norm_num_groups=32,
         norm_eps=1e-5,
         cross_attention_dim=768,
+        transformer_layers_per_block=1,
         attention_head_dim=8,
+        addition_embed_type=None,
+        addition_time_embed_dim=None,
+        projection_class_embeddings_input_dim=None,
         **kwargs,
     ):
         if kwargs.get("dual_cross_attention", None):
@@ -799,6 +853,21 @@ class UNet2DConditionModel(ModelMixin, ConfigMixin):
 
         self.time_proj = time_proj
         self.time_embedding = time_embedding
+        self.encoder_hid_proj = None
+
+        if addition_embed_type == "text":
+            raise NotImplementedError
+        elif addition_embed_type == "text_image":
+            raise NotImplementedError
+        elif addition_embed_type == "text_time":
+            self.add_time_proj = Timesteps(addition_time_embed_dim, flip_sin_to_cos, freq_shift)
+            self.add_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim)
+        elif addition_embed_type == "image":
+            raise NotImplementedError
+        elif addition_embed_type == "image_hint":
+            raise NotImplementedError
+        elif addition_embed_type is not None:
+            raise ValueError(f"addition_embed_type: {addition_embed_type} must be None, 'text' or 'text_image'.")
 
         self.down_blocks = nn.ModuleList([])
         self.mid_block = None
@@ -810,6 +879,9 @@ class UNet2DConditionModel(ModelMixin, ConfigMixin):
         if isinstance(attention_head_dim, int):
             attention_head_dim = (attention_head_dim,) * len(down_block_types)
 
+        if isinstance(transformer_layers_per_block, int):
+            transformer_layers_per_block = [transformer_layers_per_block] * len(down_block_types)
+
         # down
         output_channel = block_out_channels[0]
         for i, down_block_type in enumerate(down_block_types):
@@ -819,6 +891,7 @@ class UNet2DConditionModel(ModelMixin, ConfigMixin):
 
             down_block = get_down_block(
                 down_block_type,
+                transformer_layers_per_block=transformer_layers_per_block[i],
                 num_layers=layers_per_block,
                 in_channels=input_channel,
                 out_channels=output_channel,
@@ -828,12 +901,15 @@ class UNet2DConditionModel(ModelMixin, ConfigMixin):
                 cross_attention_dim=cross_attention_dim,
                 attn_num_head_channels=attention_head_dim[i],
                 downsample_padding=downsample_padding,
+                add_downsample=not is_final_block,
             )
             self.down_blocks.append(down_block)
 
         # mid
+        assert mid_block_type == "UNetMidBlock2DCrossAttn"
         self.mid_block = UNetMidBlock2DCrossAttn(
             in_channels=block_out_channels[-1],
+            transformer_layers_per_block=transformer_layers_per_block[-1],
             temb_channels=time_embed_dim,
             resnet_eps=norm_eps,
             resnet_act_fn=act_fn,
@@ -847,6 +923,8 @@ class UNet2DConditionModel(ModelMixin, ConfigMixin):
         # up
         reversed_block_out_channels = list(reversed(block_out_channels))
         reversed_attention_head_dim = list(reversed(attention_head_dim))
+        reversed_transformer_layers_per_block = list(reversed(transformer_layers_per_block))
+
         output_channel = reversed_block_out_channels[0]
         for i, up_block_type in enumerate(up_block_types):
             prev_output_channel = output_channel
@@ -860,6 +938,7 @@ class UNet2DConditionModel(ModelMixin, ConfigMixin):
             up_block = get_up_block(
                 up_block_type,
                 num_layers=layers_per_block + 1,
+                transformer_layers_per_block=reversed_transformer_layers_per_block[i],
                 in_channels=input_channel,
                 out_channels=output_channel,
                 prev_output_channel=prev_output_channel,
@@ -901,6 +980,110 @@ class UNet2DConditionModel(ModelMixin, ConfigMixin):
         # 2. pre-process
         sample = self.conv_in(sample)
 
+        # 3. down
+        down_block_res_samples = (sample, )
+        for downsample_block in self.down_blocks:
+            if hasattr(
+                    downsample_block,
+                    "attentions") and downsample_block.attentions is not None:
+                sample, res_samples = downsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    encoder_hidden_states=encoder_hidden_states)
+            else:
+                sample, res_samples = downsample_block(hidden_states=sample,
+                                                       temb=emb)
+
+            down_block_res_samples += res_samples
+
+        if additional_residuals:
+            new_down_block_res_samples = ()
+            for i, down_block_res_sample in enumerate(down_block_res_samples):
+                down_block_res_sample = down_block_res_sample + additional_residuals[i]
+                new_down_block_res_samples += (down_block_res_sample,)
+            down_block_res_samples = new_down_block_res_samples
+
+        # 4. mid
+        sample = self.mid_block(sample,
+                                emb,
+                                encoder_hidden_states=encoder_hidden_states)
+
+        if additional_residuals:
+            sample = sample + additional_residuals[-1]
+
+        # 5. up
+        for upsample_block in self.up_blocks:
+            res_samples = down_block_res_samples[-len(upsample_block.resnets):]
+            down_block_res_samples = down_block_res_samples[:-len(
+                upsample_block.resnets)]
+
+            if hasattr(upsample_block,
+                       "attentions") and upsample_block.attentions is not None:
+                sample = upsample_block(
+                    hidden_states=sample,
+                    temb=emb,
+                    res_hidden_states_tuple=res_samples,
+                    encoder_hidden_states=encoder_hidden_states,
+                )
+            else:
+                sample = upsample_block(hidden_states=sample,
+                                        temb=emb,
+                                        res_hidden_states_tuple=res_samples)
+
+        # 6. post-process
+        sample = self.conv_norm_out(sample)
+        sample = self.conv_act(sample)
+        sample = self.conv_out(sample)
+
+        return (sample, )
+
+
+class UNet2DConditionModelXL(UNet2DConditionModel):
+    """ UNet2DConditionModel variant for Stable Diffusion XL
+    with an extended forward() signature
+    """
+    def forward(
+        self,
+        sample,
+        timestep,
+        encoder_hidden_states,
+        time_ids,
+        text_embeds,
+        *additional_residuals,
+    ):
+        # 0. Project time embeddings
+        t_emb = self.time_proj(timestep)
+        emb = self.time_embedding(t_emb)
+
+        aug_emb = None
+
+        if self.config.addition_embed_type == "text":
+            raise NotImplementedError
+        elif self.config.addition_embed_type == "text_image":
+            raise NotImplementedError
+        elif self.config.addition_embed_type == "text_time":
+            assert time_ids is not None
+            assert text_embeds is not None
+
+            time_embeds = self.add_time_proj(time_ids.flatten())
+            time_embeds = time_embeds.reshape((text_embeds.shape[0], -1))
+
+            add_embeds = torch.concat([text_embeds, time_embeds], dim=-1)
+            aug_emb = self.add_embedding(add_embeds)
+        elif self.config.addition_embed_type == "image":
+            raise NotImplementedError
+        elif self.config.addition_embed_type == "image_hint":
+            raise NotImplementedError
+
+        emb = emb + aug_emb if aug_emb is not None else emb
+
+        # 1. center input if necessary
+        if self.config.center_input_sample:
+            sample = 2 * sample - 1.0
+
+        # 2. pre-process
+        sample = self.conv_in(sample)
+
         # 3. down
         down_block_res_samples = (sample, )
         for downsample_block in self.down_blocks:
@@ -968,8 +1151,10 @@ def get_down_block(
     resnet_eps,
     resnet_act_fn,
     attn_num_head_channels,
+    transformer_layers_per_block=1,
     cross_attention_dim=None,
     downsample_padding=None,
+    add_downsample=True,
 ):
     down_block_type = down_block_type[7:] if down_block_type.startswith(
         "UNetRes") else down_block_type
@@ -981,6 +1166,7 @@ def get_down_block(
             temb_channels=temb_channels,
             resnet_eps=resnet_eps,
             resnet_act_fn=resnet_act_fn,
+            add_downsample=add_downsample,
         )
     elif down_block_type == "CrossAttnDownBlock2D":
         if cross_attention_dim is None:
@@ -989,6 +1175,7 @@ def get_down_block(
             )
         return CrossAttnDownBlock2D(
             num_layers=num_layers,
+            transformer_layers_per_block=transformer_layers_per_block,
             in_channels=in_channels,
             out_channels=out_channels,
             temb_channels=temb_channels,
@@ -997,6 +1184,7 @@ def get_down_block(
             downsample_padding=downsample_padding,
             cross_attention_dim=cross_attention_dim,
             attn_num_head_channels=attn_num_head_channels,
+            add_downsample=add_downsample,
         )
 
 
@@ -1011,6 +1199,7 @@ def get_up_block(
     resnet_eps,
     resnet_act_fn,
     attn_num_head_channels,
+    transformer_layers_per_block=1,
     cross_attention_dim=None,
 ):
     up_block_type = up_block_type[7:] if up_block_type.startswith(
@@ -1041,6 +1230,7 @@ def get_up_block(
             resnet_act_fn=resnet_act_fn,
             cross_attention_dim=cross_attention_dim,
             attn_num_head_channels=attn_num_head_channels,
+            transformer_layers_per_block=transformer_layers_per_block,
         )
     raise ValueError(f"{up_block_type} does not exist.")
 

requirements.txt

@@ -1,7 +1,10 @@
-coremltools>=7.0b1
+coremltools>=7.0
 diffusers[torch]
 torch
 transformers==4.29.2
 scipy
 scikit-learn
-pytest
+pytest
+invisible-watermark
+safetensors
+matplotlib

setup.py

@@ -14,15 +14,18 @@ setup(
     long_description_content_type='text/markdown',
     author='Apple Inc.',
     install_requires=[
-        "coremltools>=7.0b1",
+        "coremltools>=7.0b2",
         "diffusers[torch]",
         "torch",
-        "transformers==4.29.2",
+        "transformers>=4.30.0",
         "huggingface-hub",
         "scipy",
         "numpy<1.24",
         "pytest",
         "scikit-learn",
+        "invisible-watermark",
+        "safetensors",
+        "matplotlib",
     ],
     packages=find_packages(),
     classifiers=[

swift/StableDiffusion/pipeline/CGImage+vImage.swift

@@ -4,6 +4,7 @@
 import Foundation
 import Accelerate
 import CoreML
+import CoreGraphics
 
 @available(iOS 16.0, macOS 13.0, *)
 extension CGImage {
@@ -57,13 +58,13 @@ extension CGImage {
         let cgImage = uint8Image.makeCGImage(cgImageFormat:
                 .init(bitsPerComponent: 8,
                       bitsPerPixel: 3*8,
-                      colorSpace: CGColorSpaceCreateDeviceRGB(),
+                      colorSpace: CGColorSpace(name: CGColorSpace.sRGB) ?? CGColorSpaceCreateDeviceRGB(),
                       bitmapInfo: bitmapInfo)!)!
 
         return cgImage
     }
     
-    public func plannerRGBShapedArray(minValue: Float, maxValue: Float)
+    public func planarRGBShapedArray(minValue: Float, maxValue: Float)
         throws -> MLShapedArray<Float32> {
             guard
                 var sourceFormat = vImage_CGImageFormat(cgImage: self),
@@ -77,10 +78,10 @@ extension CGImage {
             else {
                 throw ShapedArrayError.incorrectFormatsConvertingToShapedArray
             }
-            
+
             var sourceImageBuffer = try vImage_Buffer(cgImage: self)
             
-            var mediumDesination = try vImage_Buffer(width: Int(width), height: Int(height), bitsPerPixel: mediumFormat.bitsPerPixel)
+            var mediumDestination = try vImage_Buffer(width: Int(width), height: Int(height), bitsPerPixel: mediumFormat.bitsPerPixel)
             
             let converter = vImageConverter_CreateWithCGImageFormat(
                 &sourceFormat,
@@ -88,32 +89,109 @@ extension CGImage {
                 nil,
                 vImage_Flags(kvImagePrintDiagnosticsToConsole),
                 nil)
-            
+
             guard let converter = converter?.takeRetainedValue() else {
                 throw ShapedArrayError.vImageConverterNotInitialized
             }
             
-            vImageConvert_AnyToAny(converter, &sourceImageBuffer, &mediumDesination, nil, vImage_Flags(kvImagePrintDiagnosticsToConsole))
+            vImageConvert_AnyToAny(converter, &sourceImageBuffer, &mediumDestination, nil, vImage_Flags(kvImagePrintDiagnosticsToConsole))
             
             var destinationA = try vImage_Buffer(width: Int(width), height: Int(height), bitsPerPixel: 8 * UInt32(MemoryLayout<Float>.size))
             var destinationR = try vImage_Buffer(width: Int(width), height: Int(height), bitsPerPixel: 8 * UInt32(MemoryLayout<Float>.size))
             var destinationG = try vImage_Buffer(width: Int(width), height: Int(height), bitsPerPixel: 8 * UInt32(MemoryLayout<Float>.size))
             var destinationB = try vImage_Buffer(width: Int(width), height: Int(height), bitsPerPixel: 8 * UInt32(MemoryLayout<Float>.size))
-            
+
             var minFloat: [Float] = Array(repeating: minValue, count: 4)
             var maxFloat: [Float] = Array(repeating: maxValue, count: 4)
             
-            vImageConvert_ARGB8888toPlanarF(&mediumDesination, &destinationA, &destinationR, &destinationG, &destinationB, &maxFloat, &minFloat, .zero)
+            vImageConvert_ARGB8888toPlanarF(&mediumDestination, &destinationA, &destinationR, &destinationG, &destinationB, &maxFloat, &minFloat, .zero)
            
-            let redData = Data(bytes: destinationR.data, count: Int(width) * Int(height) * MemoryLayout<Float>.size)
-            let greenData = Data(bytes: destinationG.data, count: Int(width) * Int(height) * MemoryLayout<Float>.size)
-            let blueData = Data(bytes: destinationB.data, count: Int(width) * Int(height) * MemoryLayout<Float>.size)
+            let destAPtr = destinationA.data.assumingMemoryBound(to: Float.self)
+            let destRPtr = destinationR.data.assumingMemoryBound(to: Float.self)
+            let destGPtr = destinationG.data.assumingMemoryBound(to: Float.self)
+            let destBPtr = destinationB.data.assumingMemoryBound(to: Float.self)
+
+            for i in 0..<Int(width) * Int(height) {
+                if destAPtr.advanced(by: i).pointee == 0 {
+                    destRPtr.advanced(by: i).pointee = -1
+                    destGPtr.advanced(by: i).pointee = -1
+                    destBPtr.advanced(by: i).pointee = -1
+                }
+            }
             
+            let redData = destinationR.unpaddedData()
+            let greenData = destinationG.unpaddedData()
+            let blueData = destinationB.unpaddedData()
+
             let imageData = redData + greenData + blueData
 
             let shapedArray = MLShapedArray<Float32>(data: imageData, shape: [1, 3, self.height, self.width])
-            
+
+            return shapedArray
+    }
+
+    private func normalizePixelValues(pixel: UInt8) -> Float {
+        return (Float(pixel) / 127.5) - 1.0
+    }
+
+    public func toRGBShapedArray(minValue: Float, maxValue: Float)
+        throws -> MLShapedArray<Float32> {
+            let image = self
+            let width = image.width
+            let height = image.height
+            let alphaMaskValue: Float = minValue
+
+            guard let colorSpace = CGColorSpace(name: CGColorSpace.sRGB),
+                  let context = CGContext(data: nil, width: width, height: height, bitsPerComponent: 8, bytesPerRow: 4 * width, space: colorSpace, bitmapInfo: CGImageAlphaInfo.premultipliedLast.rawValue),
+                  let ptr = context.data?.bindMemory(to: UInt8.self, capacity: width * height * 4) else {
+                return []
+            }
+
+            context.draw(image, in: CGRect(x: 0, y: 0, width: width, height: height))
+
+            var redChannel = [Float](repeating: 0, count: width * height)
+            var greenChannel = [Float](repeating: 0, count: width * height)
+            var blueChannel = [Float](repeating: 0, count: width * height)
+
+            for y in 0..<height {
+                for x in 0..<width {
+                    let i = 4 * (y * width + x)
+                    if ptr[i+3] == 0 {
+                        // Alpha mask for controlnets
+                        redChannel[y * width + x] = alphaMaskValue
+                        greenChannel[y * width + x] = alphaMaskValue
+                        blueChannel[y * width + x] = alphaMaskValue
+                    } else {
+                        redChannel[y * width + x] = normalizePixelValues(pixel: ptr[i])
+                        greenChannel[y * width + x] = normalizePixelValues(pixel: ptr[i+1])
+                        blueChannel[y * width + x] = normalizePixelValues(pixel: ptr[i+2])
+                    }
+                }
+            }
+
+            let colorShape = [1, 1, height, width]
+            let redShapedArray = MLShapedArray<Float32>(scalars: redChannel, shape: colorShape)
+            let greenShapedArray = MLShapedArray<Float32>(scalars: greenChannel, shape: colorShape)
+            let blueShapedArray = MLShapedArray<Float32>(scalars: blueChannel, shape: colorShape)
+
+            let shapedArray = MLShapedArray<Float32>(concatenating: [redShapedArray, greenShapedArray, blueShapedArray], alongAxis: 1)
+
             return shapedArray
     }
 }
 
+extension vImage_Buffer {
+    func unpaddedData() -> Data {
+        let bytesPerPixel = self.rowBytes / Int(self.width)
+        let bytesPerRow = Int(self.width) * bytesPerPixel
+
+        var contiguousPixelData = Data(capacity: bytesPerRow * Int(self.height))
+        for row in 0..<Int(self.height) {
+            let rowStart = self.data!.advanced(by: row * self.rowBytes)
+            let rowData = Data(bytes: rowStart, count: bytesPerRow)
+            contiguousPixelData.append(rowData)
+        }
+
+        return contiguousPixelData
+    }
+}

swift/StableDiffusion/pipeline/ControlNet.swift

@@ -3,6 +3,7 @@
 
 import Foundation
 import CoreML
+import Accelerate
 
 @available(iOS 16.2, macOS 13.1, *)
 public struct ControlNet: ResourceManaging {
@@ -99,13 +100,15 @@ public struct ControlNet: ResourceManaging {
             for n in 0..<results.count {
                 let result = results.features(at: n)
                 for k in result.featureNames {
-                    let newValue = result.featureValue(for: k)!.shapedArrayValue(of: Float32.self)!
+                    let newValue = result.featureValue(for: k)!.multiArrayValue!
                     if modelIndex == 0 {
-                        outputs[n][k] = newValue
+                        outputs[n][k] = MLShapedArray<Float32>(newValue)
                     } else {
-                        outputs[n][k]!.withUnsafeMutableShapedBufferPointer { pt, _, _ in
-                            for (i, v) in newValue.scalars.enumerated() { pt[i] += v }
-                        }
+                        let outputArray = MLMultiArray(outputs[n][k]!)
+                        let count = newValue.count
+                        let inputPointer = newValue.dataPointer.assumingMemoryBound(to: Float.self)
+                        let outputPointer = outputArray.dataPointer.assumingMemoryBound(to: Float.self)
+                        vDSP_vadd(inputPointer, 1, outputPointer, 1, outputPointer, 1, vDSP_Length(count))
                     }
                 }
             }

swift/StableDiffusion/pipeline/DPMSolverMultistepScheduler.swift

@@ -4,6 +4,13 @@
 import Accelerate
 import CoreML
 
+/// How to space timesteps for inference
+public enum TimeStepSpacing {
+    case linspace
+    case leading
+    case karras
+}
+
 /// A scheduler used to compute a de-noised image
 ///
 ///  This implementation matches:
@@ -32,12 +39,14 @@ public final class DPMSolverMultistepScheduler: Scheduler {
     public let solverOrder = 2
     private(set) var lowerOrderStepped = 0
     
+    private var usingKarrasSigmas = false
+
     /// Whether to use lower-order solvers in the final steps. Only valid for less than 15 inference steps.
     /// We empirically find this trick can stabilize the sampling of DPM-Solver, especially with 10 or fewer steps.
     public let useLowerOrderFinal = true
     
     // Stores solverOrder (2) items
-    private(set) var modelOutputs: [MLShapedArray<Float32>] = []
+    public private(set) var modelOutputs: [MLShapedArray<Float32>] = []
 
     /// Create a scheduler that uses a second order DPM-Solver++ algorithm.
     ///
@@ -47,13 +56,15 @@ public final class DPMSolverMultistepScheduler: Scheduler {
     ///   - betaSchedule: Method to schedule betas from betaStart to betaEnd
     ///   - betaStart: The starting value of beta for inference
     ///   - betaEnd: The end value for beta for inference
+    ///   - timeStepSpacing: How to space time steps
     /// - Returns: A scheduler ready for its first step
     public init(
         stepCount: Int = 50,
         trainStepCount: Int = 1000,
         betaSchedule: BetaSchedule = .scaledLinear,
         betaStart: Float = 0.00085,
-        betaEnd: Float = 0.012
+        betaEnd: Float = 0.012,
+        timeStepSpacing: TimeStepSpacing = .linspace
     ) {
         self.trainStepCount = trainStepCount
         self.inferenceStepCount = stepCount
@@ -72,12 +83,51 @@ public final class DPMSolverMultistepScheduler: Scheduler {
         }
         self.alphasCumProd = alphasCumProd
 
-        // Currently we only support VP-type noise shedule
-        self.alpha_t = vForce.sqrt(self.alphasCumProd)
-        self.sigma_t = vForce.sqrt(vDSP.subtract([Float](repeating: 1, count: self.alphasCumProd.count), self.alphasCumProd))
-        self.lambda_t = zip(self.alpha_t, self.sigma_t).map { α, σ in log(α) - log(σ) }
+        switch timeStepSpacing {
+        case .linspace:
+            self.timeSteps = linspace(0, Float(self.trainStepCount-1), stepCount+1).dropFirst().reversed().map { Int(round($0)) }
+            self.alpha_t = vForce.sqrt(self.alphasCumProd)
+            self.sigma_t = vForce.sqrt(vDSP.subtract([Float](repeating: 1, count: self.alphasCumProd.count), self.alphasCumProd))
+        case .leading:
+            let lastTimeStep = trainStepCount - 1
+            let stepRatio = lastTimeStep / (stepCount + 1)
+            // Creates integer timesteps by multiplying by ratio
+            self.timeSteps = (0...stepCount).map { 1 + $0 * stepRatio }.dropFirst().reversed()
+            self.alpha_t = vForce.sqrt(self.alphasCumProd)
+            self.sigma_t = vForce.sqrt(vDSP.subtract([Float](repeating: 1, count: self.alphasCumProd.count), self.alphasCumProd))
+        case .karras:
+            // sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5)
+            let scaled = vDSP.multiply(
+                subtraction: ([Float](repeating: 1, count: self.alphasCumProd.count), self.alphasCumProd),
+                subtraction: (vDSP.divide(1, self.alphasCumProd), [Float](repeating: 0, count: self.alphasCumProd.count))
+            )
+            let sigmas = vForce.sqrt(scaled)
+            let logSigmas = sigmas.map { log($0) }
 
-        self.timeSteps = linspace(0, Float(self.trainStepCount-1), stepCount+1).dropFirst().reversed().map { Int(round($0)) }
+            let sigmaMin = sigmas.first!
+            let sigmaMax = sigmas.last!
+            let rho: Float = 7
+            let ramp = linspace(0, 1, stepCount)
+            let minInvRho = pow(sigmaMin, (1 / rho))
+            let maxInvRho = pow(sigmaMax, (1 / rho))
+
+            var karrasSigmas = ramp.map { pow(maxInvRho + $0 * (minInvRho - maxInvRho), rho) }
+            let karrasTimeSteps = karrasSigmas.map { sigmaToTimestep(sigma: $0, logSigmas: logSigmas) }
+            self.timeSteps = karrasTimeSteps
+
+            karrasSigmas.append(karrasSigmas.last!)
+
+            self.alpha_t = vDSP.divide(1, vForce.sqrt(vDSP.add(1, vDSP.square(karrasSigmas))))
+            self.sigma_t = vDSP.multiply(karrasSigmas, self.alpha_t)
+            usingKarrasSigmas = true
+        }
+
+        self.lambda_t = zip(self.alpha_t, self.sigma_t).map { α, σ in log(α) - log(σ) }
+    }
+    
+    func timestepToIndex(_ timestep: Int) -> Int {
+        guard usingKarrasSigmas else { return timestep }
+        return self.timeSteps.firstIndex(of: timestep) ?? 0
     }
     
     /// Convert the model output to the corresponding type the algorithm needs.
@@ -85,7 +135,8 @@ public final class DPMSolverMultistepScheduler: Scheduler {
     func convertModelOutput(modelOutput: MLShapedArray<Float32>, timestep: Int, sample: MLShapedArray<Float32>) -> MLShapedArray<Float32> {
         assert(modelOutput.scalarCount == sample.scalarCount)
         let scalarCount = modelOutput.scalarCount
-        let (alpha_t, sigma_t) = (self.alpha_t[timestep], self.sigma_t[timestep])
+        let sigmaIndex = timestepToIndex(timestep)
+        let (alpha_t, sigma_t) = (self.alpha_t[sigmaIndex], self.sigma_t[sigmaIndex])
 
         return MLShapedArray(unsafeUninitializedShape: modelOutput.shape) { scalars, _ in
             assert(scalars.count == scalarCount)
@@ -108,9 +159,11 @@ public final class DPMSolverMultistepScheduler: Scheduler {
         prevTimestep: Int,
         sample: MLShapedArray<Float32>
     ) -> MLShapedArray<Float32> {
-        let (p_lambda_t, lambda_s) = (Double(lambda_t[prevTimestep]), Double(lambda_t[timestep]))
-        let p_alpha_t = Double(alpha_t[prevTimestep])
-        let (p_sigma_t, sigma_s) = (Double(sigma_t[prevTimestep]), Double(sigma_t[timestep]))
+        let prevIndex = timestepToIndex(prevTimestep)
+        let currIndex = timestepToIndex(timestep)
+        let (p_lambda_t, lambda_s) = (Double(lambda_t[prevIndex]), Double(lambda_t[currIndex]))
+        let p_alpha_t = Double(alpha_t[prevIndex])
+        let (p_sigma_t, sigma_s) = (Double(sigma_t[prevIndex]), Double(sigma_t[currIndex]))
         let h = p_lambda_t - lambda_s
         // x_t = (sigma_t / sigma_s) * sample - (alpha_t * (torch.exp(-h) - 1.0)) * model_output
         let x_t = weightedSum(
@@ -130,9 +183,13 @@ public final class DPMSolverMultistepScheduler: Scheduler {
     ) -> MLShapedArray<Float32> {
         let (s0, s1) = (timesteps[back: 1], timesteps[back: 2])
         let (m0, m1) = (modelOutputs[back: 1], modelOutputs[back: 2])
-        let (p_lambda_t, lambda_s0, lambda_s1) = (Double(lambda_t[t]), Double(lambda_t[s0]), Double(lambda_t[s1]))
-        let p_alpha_t = Double(alpha_t[t])
-        let (p_sigma_t, sigma_s0) = (Double(sigma_t[t]), Double(sigma_t[s0]))
+        let (p_lambda_t, lambda_s0, lambda_s1) = (
+            Double(lambda_t[timestepToIndex(t)]),
+            Double(lambda_t[timestepToIndex(s0)]),
+            Double(lambda_t[timestepToIndex(s1)])
+        )
+        let p_alpha_t = Double(alpha_t[timestepToIndex(t)])
+        let (p_sigma_t, sigma_s0) = (Double(sigma_t[timestepToIndex(t)]), Double(sigma_t[timestepToIndex(s0)]))
         let (h, h_0) = (p_lambda_t - lambda_s0, lambda_s0 - lambda_s1)
         let r0 = h_0 / h
         let D0 = m0
@@ -186,3 +243,31 @@ public final class DPMSolverMultistepScheduler: Scheduler {
         return prevSample
     }
 }
+
+func sigmaToTimestep(sigma: Float, logSigmas: [Float]) -> Int {
+    let logSigma = log(sigma)
+    let dists = logSigmas.map { logSigma - $0 }
+
+    // last index that is not negative, clipped to last index - 1
+    var lowIndex = dists.reduce(-1) { partialResult, dist in
+        return dist >= 0 && partialResult < dists.endIndex-2 ? partialResult + 1 : partialResult
+    }
+    lowIndex = max(lowIndex, 0)
+    let highIndex = lowIndex + 1
+
+    let low = logSigmas[lowIndex]
+    let high = logSigmas[highIndex]
+
+    // Interpolate sigmas
+    let w = ((low - logSigma) / (low - high)).clipped(to: 0...1)
+
+    // transform interpolated value to time range
+    let t = (1 - w) * Float(lowIndex) + w * Float(highIndex)
+    return Int(round(t))
+}
+
+extension FloatingPoint {
+    func clipped(to range: ClosedRange<Self>) -> Self {
+        return min(max(self, range.lowerBound), range.upperBound)
+    }
+}

swift/StableDiffusion/pipeline/Decoder.swift

@@ -64,7 +64,7 @@ public struct Decoder: ResourceManaging {
             let result = results.features(at: i)
             let outputName = result.featureNames.first!
             let output = result.featureValue(for: outputName)!.multiArrayValue!
-            return try CGImage.fromShapedArray(MLShapedArray<Float32>(output))
+            return try CGImage.fromShapedArray(MLShapedArray<Float32>(converting: output))
         }
 
         return images

swift/StableDiffusion/pipeline/Encoder.swift

@@ -50,7 +50,7 @@ public struct Encoder: ResourceManaging {
         scaleFactor: Float32,
         random: inout RandomSource
     ) throws -> MLShapedArray<Float32> {
-        let imageData = try image.plannerRGBShapedArray(minValue: -1.0, maxValue: 1.0)
+        let imageData = try image.planarRGBShapedArray(minValue: -1.0, maxValue: 1.0)
         guard imageData.shape == inputShape else {
             // TODO: Consider auto resizing and croping similar to how Vision or CoreML auto-generated Swift code can accomplish with `MLFeatureValue`
             throw Error.sampleInputShapeNotCorrect
@@ -63,7 +63,7 @@ public struct Encoder: ResourceManaging {
         }
         let outputName = result.featureNames.first!
         let outputValue = result.featureValue(for: outputName)!.multiArrayValue!
-        let output = MLShapedArray<Float32>(outputValue)
+        let output = MLShapedArray<Float32>(converting: outputValue)
         
         // DiagonalGaussianDistribution
         let mean = output[0][0..<4]
@@ -93,7 +93,7 @@ public struct Encoder: ResourceManaging {
     
     var inputDescription: MLFeatureDescription {
         try! model.perform { model in
-            model.modelDescription.inputDescriptionsByName["z"]!
+            model.modelDescription.inputDescriptionsByName.first!.value
         }
     }
     

swift/StableDiffusion/pipeline/MultilingualTextEncoder.swift

@@ -5,6 +5,7 @@ import Foundation
 import NaturalLanguage
 import CoreML
 
+#if canImport(NaturalLanguage.NLContextualEmbedding)
 @available(iOS 17.0, macOS 14.0, *)
 public struct MultilingualTextEncoder: TextEncoderModel {
     let adapter: ManagedMLModel?
@@ -174,11 +175,13 @@ extension MultilingualTextEncoder {
         }
     }
 }
+#endif
 
 @available(iOS 16.2, macOS 13.1, *)
 public enum Script: String {
     case latin, cyrillic, cjk
 
+#if canImport(NaturalLanguage.NLScript)
     @available(iOS 17.0, macOS 14.0, *)
     var asNLScript: NLScript {
         switch self {
@@ -187,4 +190,5 @@ public enum Script: String {
         case .cjk: return .simplifiedChinese
         }
     }
+#endif
 }

swift/StableDiffusion/pipeline/NvRandomSource.swift

@@ -0,0 +1,91 @@
+import Foundation
+import CoreML
+
+/// A random source consistent with NVIDIA curandom
+///
+///  This implementation references to:
+///  https://github.com/dsnz/random/blob/master/philox.py for Philox_M4_32 configuration.
+///
+@available(iOS 16.2, macOS 13.1, *)
+struct NvRandomSource: RandomSource {
+  public let seed: UInt64
+  private var offset: UInt32
+
+  /// Initialize with a random seed
+  ///
+  /// - Parameters
+  ///     - seed: Seed for underlying Philox M4 32 generator
+  /// - Returns random source
+  public init(seed: UInt32) {
+    self.seed = UInt64(seed)
+    offset = 0
+  }
+
+  static private let PHILOX_M4_32: (UInt32, UInt32) = (0xD251_1F53, 0xCD9E_8D57)
+  static private let PHILOX_W_32: (UInt32, UInt32) = (0x9E37_79B9, 0xBB67_AE85)
+
+  static private func philox4Round(counter: inout [[UInt32]], key: [[UInt32]]) {
+    for i in 0..<counter[0].count {
+      let v1: UInt64 = UInt64(counter[0][i]) * UInt64(PHILOX_M4_32.0)
+      let v2: UInt64 = UInt64(counter[2][i]) * UInt64(PHILOX_M4_32.1)
+      counter[0][i] = UInt32(v2 >> 32) ^ counter[1][i] ^ key[0][i]
+      counter[1][i] = UInt32(v2 & 0xffff_ffff)
+      counter[2][i] = UInt32(v1 >> 32) ^ counter[3][i] ^ key[1][i]
+      counter[3][i] = UInt32(v1 & 0xffff_ffff)
+    }
+  }
+
+  static private func philox4Bumpkey(key: inout [[UInt32]]) {
+    for (i, element) in key[0].enumerated() {
+      key[0][i] = element &+ PHILOX_W_32.0
+    }
+    for (i, element) in key[1].enumerated() {
+      key[1][i] = element &+ PHILOX_W_32.1
+    }
+  }
+
+  static private func philox4_32(counter: inout [[UInt32]], key: inout [[UInt32]], rounds: Int = 10) {
+    for _ in 0..<(rounds - 1) {
+      philox4Round(counter: &counter, key: key)
+      philox4Bumpkey(key: &key)
+    }
+    philox4Round(counter: &counter, key: key)
+  }
+
+  private func boxMuller(_ counter1: [UInt32], _ counter2: [UInt32], mean: Double, stdev: Double) -> [Double] {
+    // Box-Muller transform
+    return zip(counter1, counter2).map {
+      let u: Double = Double($0) / 4294967296.0 + (1.0 / 8589934592.0)
+      let v: Double = Double($1) * (.pi / 2147483648.0) + (.pi / 4294967296.0)
+      let radius = stdev * sqrt(-2.0 * log(u))
+      return radius * sin(v) + mean
+    }
+  }
+
+  private mutating func normalArray(count: Int, mean: Double, stdev: Double) -> [Double] {
+    var counter: [[UInt32]] = [
+      Array(repeating: offset, count: count),
+      Array(repeating: 0, count: count),
+      Array(0..<UInt32(count)),
+      Array(repeating: 0, count: count),
+    ]
+    offset += 1
+    var key: [[UInt32]] = [
+      Array(repeating: UInt32(seed & 0xffff_ffff), count: count),
+      Array(repeating: UInt32(seed >> 32), count: count),
+    ]
+    Self.philox4_32(counter: &counter, key: &key)
+    return boxMuller(counter[0], counter[1], mean: mean, stdev: stdev)
+  }
+
+  /// Generates a random value from a normal distribution with given mean and standard deviation.
+  mutating func nextNormal(mean: Double = 0.0, stdev: Double = 1.0) -> Double {
+    return normalArray(count: 1, mean: mean, stdev: stdev)[0]
+  }
+
+  /// Generate a shaped array with scalars from a normal distribution with given mean and standard deviation.
+  mutating func normalShapedArray(_ shape: [Int], mean: Double = 0.0, stdev: Double = 1.0) -> MLShapedArray<Double> {
+    let count = shape.reduce(1, *)
+    return .init(scalars: normalArray(count: count, mean: mean, stdev: stdev), shape: shape)
+  }
+}

swift/StableDiffusion/pipeline/Scheduler.swift

@@ -30,6 +30,9 @@ public protocol Scheduler {
     /// Standard deviation of the initial noise distribution
     var initNoiseSigma: Float { get }
 
+    /// Denoised latents
+    var modelOutputs: [MLShapedArray<Float32>] { get }
+
     /// Compute a de-noised image sample and step scheduler state
     ///
     /// - Parameters:
@@ -139,6 +142,12 @@ public final class PNDMScheduler: Scheduler {
     public let alphasCumProd: [Float]
     public let timeSteps: [Int]
 
+    public let alpha_t: [Float]
+    public let sigma_t: [Float]
+    public let lambda_t: [Float]
+
+    public private(set) var modelOutputs: [MLShapedArray<Float32>] = []
+
     // Internal state
     var counter: Int
     var ets: [MLShapedArray<Float32>]
@@ -181,6 +190,10 @@ public final class PNDMScheduler: Scheduler {
             Int((Float($0) * stepRatio).rounded()) + stepsOffset
         }
 
+        self.alpha_t = vForce.sqrt(self.alphasCumProd)
+        self.sigma_t = vForce.sqrt(vDSP.subtract([Float](repeating: 1, count: self.alphasCumProd.count), self.alphasCumProd))
+        self.lambda_t = zip(self.alpha_t, self.sigma_t).map { α, σ in log(α) - log(σ) }
+
         var timeSteps: [Int] = []
         timeSteps.append(contentsOf: forwardSteps.dropLast(1))
         timeSteps.append(timeSteps.last!)
@@ -251,11 +264,33 @@ public final class PNDMScheduler: Scheduler {
             )
         }
 
+        let convertedOutput = convertModelOutput(modelOutput: modelOutput, timestep: timeStep, sample: sample)
+        modelOutputs.append(convertedOutput)
+
         let prevSample = previousSample(sample, timeStep, prevStep, modelOutput)
+
         counter += 1
         return prevSample
     }
 
+    /// Convert the model output to the corresponding type the algorithm needs.
+    func convertModelOutput(modelOutput: MLShapedArray<Float32>, timestep: Int, sample: MLShapedArray<Float32>) -> MLShapedArray<Float32> {
+        assert(modelOutput.scalarCount == sample.scalarCount)
+        let scalarCount = modelOutput.scalarCount
+        let (alpha_t, sigma_t) = (self.alpha_t[timestep], self.sigma_t[timestep])
+
+        return MLShapedArray(unsafeUninitializedShape: modelOutput.shape) { scalars, _ in
+            assert(scalars.count == scalarCount)
+            modelOutput.withUnsafeShapedBufferPointer { modelOutput, _, _ in
+                sample.withUnsafeShapedBufferPointer { sample, _, _ in
+                    for i in 0 ..< scalarCount {
+                        scalars.initializeElement(at: i, to: (sample[i] - modelOutput[i] * sigma_t) / alpha_t)
+                    }
+                }
+            }
+        }
+    }
+
     /// Compute  sample (denoised image) at previous step given a current time step
     ///
     /// - Parameters:

swift/StableDiffusion/pipeline/StableDiffusionPipeline+Resources.swift

@@ -69,8 +69,10 @@ public extension StableDiffusionPipeline {
         /// Expect URL of each resource
         let urls = ResourceURLs(resourcesAt: baseURL)
         let textEncoder: TextEncoderModel
+
+#if canImport(NaturalLanguage.NLScript)
         if useMultilingualTextEncoder {
-            guard #available(macOS 14.0, iOS 17.0, *) else { throw Error.unsupportedOSVersion }
+            guard #available(macOS 14.0, iOS 17.0, *) else { throw PipelineError.unsupportedOSVersion }
             textEncoder = MultilingualTextEncoder(
                 modelAt: urls.multilingualTextEncoderProjectionURL,
                 configuration: config,
@@ -80,6 +82,10 @@ public extension StableDiffusionPipeline {
             let tokenizer = try BPETokenizer(mergesAt: urls.mergesURL, vocabularyAt: urls.vocabURL)
             textEncoder = TextEncoder(tokenizer: tokenizer, modelAt: urls.textEncoderURL, configuration: config)
         }
+#else
+        let tokenizer = try BPETokenizer(mergesAt: urls.mergesURL, vocabularyAt: urls.vocabURL)
+        textEncoder = TextEncoder(tokenizer: tokenizer, modelAt: urls.textEncoderURL, configuration: config)
+#endif
 
         // ControlNet model
         var controlNet: ControlNet? = nil

swift/StableDiffusion/pipeline/StableDiffusionPipeline.Configuration.swift

@@ -4,65 +4,93 @@
 import Foundation
 import CoreGraphics
 
-@available(iOS 16.2, macOS 13.1, *)
-extension StableDiffusionPipeline {
+/// Type of processing that will be performed to generate an image
+public enum PipelineMode {
+    case textToImage
+    case imageToImage
+    // case inPainting
+}
+
+/// Image generation configuration
+public struct PipelineConfiguration: Hashable {
     
-    /// Tyoe of processing that will be performed to generate an image
-    public enum Mode {
-        case textToImage
-        case imageToImage
-        // case inPainting
+    /// Text prompt to guide sampling
+    public var prompt: String
+    /// Negative text prompt to guide sampling
+    public var negativePrompt: String = ""
+    /// Starting image for image2image or in-painting
+    public var startingImage: CGImage? = nil
+    /// Fraction of inference steps to be used in `.imageToImage` pipeline mode
+    /// Must be between 0 and 1
+    /// Higher values will result in greater transformation of the `startingImage`
+    public var strength: Float = 1.0
+    /// Fraction of inference steps to at which to start using the refiner unet if present in `textToImage` mode
+    /// Must be between 0 and 1
+    /// Higher values will result in fewer refiner steps
+    public var refinerStart: Float = 0.8
+    /// Number of images to generate
+    public var imageCount: Int = 1
+    /// Number of inference steps to perform
+    public var stepCount: Int = 50
+    /// Random seed which to start generation
+    public var seed: UInt32 = 0
+    /// Controls the influence of the text prompt on sampling process (0=random images)
+    public var guidanceScale: Float = 7.5
+    /// List of Images for available ControlNet Models
+    public var controlNetInputs: [CGImage] = []
+    /// Safety checks are only performed if `self.canSafetyCheck && !disableSafety`
+    public var disableSafety: Bool = false
+    /// Enables progress updates to decode `currentImages` from denoised latent images for better previews
+    public var useDenoisedIntermediates: Bool = false
+    /// The type of Scheduler to use.
+    public var schedulerType: StableDiffusionScheduler = .pndmScheduler
+    /// The spacing to use for scheduler sigmas and time steps. Only supported when using `.dpmppScheduler`.
+    public var schedulerTimestepSpacing: TimeStepSpacing = .linspace
+    /// The type of RNG to use
+    public var rngType: StableDiffusionRNG = .numpyRNG
+    /// Scale factor to use on the latent after encoding
+    public var encoderScaleFactor: Float32 = 0.18215
+    /// Scale factor to use on the latent before decoding
+    public var decoderScaleFactor: Float32 = 0.18215
+    /// If `originalSize` is not the same as `targetSize` the image will appear to be down- or upsampled.
+    /// Part of SDXL’s micro-conditioning as explained in section 2.2 of https://huggingface.co/papers/2307.01952.
+    public var originalSize: Float32 = 1024
+    /// `cropsCoordsTopLeft` can be used to generate an image that appears to be “cropped” from the position `cropsCoordsTopLeft` downwards.
+    /// Favorable, well-centered images are usually achieved by setting `cropsCoordsTopLeft` to (0, 0).
+    public var cropsCoordsTopLeft: Float32 = 0
+    /// For most cases, `target_size` should be set to the desired height and width of the generated image.
+    public var targetSize: Float32 = 1024
+    /// Used to simulate an aesthetic score of the generated image by influencing the positive text condition.
+    public var aestheticScore: Float32 = 6
+    /// Can be used to simulate an aesthetic score of the generated image by influencing the negative text condition.
+    public var negativeAestheticScore: Float32 = 2.5
+
+    /// Given the configuration, what mode will be used for generation
+    public var mode: PipelineMode {
+        guard startingImage != nil else {
+            return .textToImage
+        }
+        guard strength < 1.0 else {
+            return .textToImage
+        }
+        return .imageToImage
     }
-    
-    /// Image generation configuration
-    public struct Configuration: Hashable {
-        
-        /// Text prompt to guide sampling
-        public var prompt: String
-        /// Negative text prompt to guide sampling
-        public var negativePrompt: String = ""
-        /// Starting image for image2image or in-painting
-        public var startingImage: CGImage? = nil
-        //public var maskImage: CGImage? = nil
-        public var strength: Float = 1.0
-        /// Number of images to generate
-        public var imageCount: Int = 1
-        /// Number of inference steps to perform
-        public var stepCount: Int = 50
-        /// Random seed which to start generation
-        public var seed: UInt32 = 0
-        /// Controls the influence of the text prompt on sampling process (0=random images)
-        public var guidanceScale: Float = 7.5
-        /// List of Images for available ControlNet Models
-        public var controlNetInputs: [CGImage] = []
-        /// Safety checks are only performed if `self.canSafetyCheck && !disableSafety`
-        public var disableSafety: Bool = false
-        /// The type of Scheduler to use.
-        public var schedulerType: StableDiffusionScheduler = .pndmScheduler
-        /// The type of RNG to use
-        public var rngType: StableDiffusionRNG = .numpyRNG
-        /// Scale factor to use on the latent after encoding
-        public var encoderScaleFactor: Float32 = 0.18215
-        /// Scale factor to use on the latent before decoding
-        public var decoderScaleFactor: Float32 = 0.18215
-        
-        /// Given the configuration, what mode will be used for generation
-        public var mode: Mode {
-            guard startingImage != nil else {
-                return .textToImage
-            }
-            guard strength < 1.0 else {
-                return .textToImage
-            }
-            return .imageToImage
-        }
-        
-        public init(
-            prompt: String
-        ) {
-            self.prompt = prompt
-        }
-        
+
+    public init(
+        prompt: String
+    ) {
+        self.prompt = prompt
     }
 
 }
+
+
+@available(iOS 16.2, macOS 13.1, *)
+public extension StableDiffusionPipeline {
+
+    /// Type of processing that will be performed to generate an image
+    typealias Mode = PipelineMode
+
+    /// Image generation configuration
+    typealias Configuration = PipelineConfiguration
+}

swift/StableDiffusion/pipeline/StableDiffusionPipeline.swift

@@ -21,6 +21,35 @@ public enum StableDiffusionRNG {
     case numpyRNG
     /// RNG that matches PyTorch CPU implementation.
     case torchRNG
+    /// RNG that matches PyTorch CUDA implementation.
+    case nvidiaRNG
+}
+
+public enum PipelineError: String, Swift.Error {
+    case missingUnetInputs
+    case startingImageProvidedWithoutEncoder
+    case startingText2ImgWithoutTextEncoder
+    case unsupportedOSVersion
+}
+
+@available(iOS 16.2, macOS 13.1, *)
+public protocol StableDiffusionPipelineProtocol: ResourceManaging {
+    var canSafetyCheck: Bool { get }
+
+    func generateImages(
+        configuration config: PipelineConfiguration,
+        progressHandler: (PipelineProgress) -> Bool
+    ) throws -> [CGImage?]
+
+    func decodeToImages(
+        _ latents: [MLShapedArray<Float32>],
+        configuration config: PipelineConfiguration
+    ) throws -> [CGImage?]
+}
+
+@available(iOS 16.2, macOS 13.1, *)
+public extension StableDiffusionPipelineProtocol {
+    var canSafetyCheck: Bool { false }
 }
 
 /// A pipeline used to generate image samples from text input using stable diffusion
@@ -28,13 +57,8 @@ public enum StableDiffusionRNG {
 /// This implementation matches:
 /// [Hugging Face Diffusers Pipeline](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py)
 @available(iOS 16.2, macOS 13.1, *)
-public struct StableDiffusionPipeline: ResourceManaging {
-    
-    public enum Error: String, Swift.Error {
-        case startingImageProvidedWithoutEncoder
-        case unsupportedOSVersion
-    }
-    
+public struct StableDiffusionPipeline: StableDiffusionPipelineProtocol {
+
     /// Model to generate embeddings for tokenized input text
     var textEncoder: TextEncoderModel
 
@@ -204,12 +228,16 @@ public struct StableDiffusionPipeline: ResourceManaging {
         let scheduler: [Scheduler] = (0..<config.imageCount).map { _ in
             switch config.schedulerType {
             case .pndmScheduler: return PNDMScheduler(stepCount: config.stepCount)
-            case .dpmSolverMultistepScheduler: return DPMSolverMultistepScheduler(stepCount: config.stepCount)
+            case .dpmSolverMultistepScheduler: return DPMSolverMultistepScheduler(stepCount: config.stepCount, timeStepSpacing: config.schedulerTimestepSpacing)
             }
         }
 
         // Generate random latent samples from specified seed
         var latents: [MLShapedArray<Float32>] = try generateLatentSamples(configuration: config, scheduler: scheduler[0])
+
+        // Store denoised latents from scheduler to pass into decoder
+        var denoisedLatents: [MLShapedArray<Float32>] = latents.map { MLShapedArray(converting: $0) }
+
         if reduceMemory {
             encoder?.unloadResources()
         }
@@ -217,7 +245,7 @@ public struct StableDiffusionPipeline: ResourceManaging {
         
         // Convert cgImage for ControlNet into MLShapedArray
         let controlNetConds = try config.controlNetInputs.map { cgImage in
-            let shapedArray = try cgImage.plannerRGBShapedArray(minValue: 0.0, maxValue: 1.0)
+            let shapedArray = try cgImage.planarRGBShapedArray(minValue: 0.0, maxValue: 1.0)
             return MLShapedArray(
                 concatenating: [shapedArray, shapedArray],
                 alongAxis: 0
@@ -233,7 +261,7 @@ public struct StableDiffusionPipeline: ResourceManaging {
             let latentUnetInput = latents.map {
                 MLShapedArray<Float32>(concatenating: [$0, $0], alongAxis: 0)
             }
-            
+
             // Before Unet, execute controlNet and add the output into Unet inputs
             let additionalResiduals = try controlNet?.execute(
                 latents: latentUnetInput,
@@ -261,15 +289,19 @@ public struct StableDiffusionPipeline: ResourceManaging {
                     timeStep: t,
                     sample: latents[i]
                 )
+
+                denoisedLatents[i] = scheduler[i].modelOutputs.last ?? latents[i]
             }
 
+            let currentLatentSamples = config.useDenoisedIntermediates ? denoisedLatents : latents
+
             // Report progress
             let progress = Progress(
                 pipeline: self,
                 prompt: config.prompt,
                 step: step,
                 stepCount: timeSteps.count,
-                currentLatentSamples: latents,
+                currentLatentSamples: currentLatentSamples,
                 configuration: config
             )
             if !progressHandler(progress) {
@@ -284,16 +316,7 @@ public struct StableDiffusionPipeline: ResourceManaging {
         }
 
         // Decode the latent samples to images
-        return try decodeToImages(latents, configuration: config)
-    }
-
-    private func randomSource(from rng: StableDiffusionRNG, seed: UInt32) -> RandomSource {
-        switch rng {
-        case .numpyRNG:
-            return NumPyRandomSource(seed: seed)
-        case .torchRNG:
-            return TorchRandomSource(seed: seed)
-        }
+        return try decodeToImages(denoisedLatents, configuration: config)
     }
 
     func generateLatentSamples(configuration config: Configuration, scheduler: Scheduler) throws -> [MLShapedArray<Float32>] {
@@ -308,7 +331,7 @@ public struct StableDiffusionPipeline: ResourceManaging {
         }
         if let image = config.startingImage, config.mode == .imageToImage {
             guard let encoder else {
-                throw Error.startingImageProvidedWithoutEncoder
+                throw PipelineError.startingImageProvidedWithoutEncoder
             }
             let latent = try encoder.encode(image, scaleFactor: config.encoderScaleFactor, random: &random)
             return scheduler.addNoise(originalSample: latent, noise: samples, strength: config.strength)
@@ -316,6 +339,74 @@ public struct StableDiffusionPipeline: ResourceManaging {
         return samples
     }
 
+    public func decodeToImages(_ latents: [MLShapedArray<Float32>], configuration config: Configuration) throws -> [CGImage?] {
+        let images = try decoder.decode(latents, scaleFactor: config.decoderScaleFactor)
+        if reduceMemory {
+            decoder.unloadResources()
+        }
+
+        // If safety is disabled return what was decoded
+        if config.disableSafety {
+            return images
+        }
+
+        // If there is no safety checker return what was decoded
+        guard let safetyChecker = safetyChecker else {
+            return images
+        }
+
+        // Otherwise change images which are not safe to nil
+        let safeImages = try images.map { image in
+            try safetyChecker.isSafe(image) ? image : nil
+        }
+
+        if reduceMemory {
+            safetyChecker.unloadResources()
+        }
+
+        return safeImages
+    }
+
+}
+
+/// Sampling progress details
+@available(iOS 16.2, macOS 13.1, *)
+public struct PipelineProgress {
+    public let pipeline: StableDiffusionPipelineProtocol
+    public let prompt: String
+    public let step: Int
+    public let stepCount: Int
+    public let currentLatentSamples: [MLShapedArray<Float32>]
+    public let configuration: PipelineConfiguration
+    public var isSafetyEnabled: Bool {
+        pipeline.canSafetyCheck && !configuration.disableSafety
+    }
+    public var currentImages: [CGImage?] {
+        try! pipeline.decodeToImages(currentLatentSamples, configuration: configuration)
+    }
+}
+
+@available(iOS 16.2, macOS 13.1, *)
+public extension StableDiffusionPipeline {
+    /// Sampling progress details
+    typealias Progress = PipelineProgress
+}
+
+// Helper functions
+
+@available(iOS 16.2, macOS 13.1, *)
+extension StableDiffusionPipelineProtocol {
+    internal func randomSource(from rng: StableDiffusionRNG, seed: UInt32) -> RandomSource {
+        switch rng {
+        case .numpyRNG:
+            return NumPyRandomSource(seed: seed)
+        case .torchRNG:
+            return TorchRandomSource(seed: seed)
+        case .nvidiaRNG:
+            return NvRandomSource(seed: seed)
+        }
+    }
+
     func toHiddenStates(_ embedding: MLShapedArray<Float32>) -> MLShapedArray<Float32> {
         // Unoptimized manual transpose [0, 2, None, 1]
         // e.g. From [2, 77, 768] to [2, 768, 1, 77]
@@ -351,52 +442,4 @@ public struct StableDiffusionPipeline: ResourceManaging {
             }
         }
     }
-
-    func decodeToImages(_ latents: [MLShapedArray<Float32>], configuration config: Configuration) throws -> [CGImage?] {
-        let images = try decoder.decode(latents, scaleFactor: config.decoderScaleFactor)
-        if reduceMemory {
-            decoder.unloadResources()
-        }
-
-        // If safety is disabled return what was decoded
-        if config.disableSafety {
-            return images
-        }
-
-        // If there is no safety checker return what was decoded
-        guard let safetyChecker = safetyChecker else {
-            return images
-        }
-
-        // Otherwise change images which are not safe to nil
-        let safeImages = try images.map { image in
-            try safetyChecker.isSafe(image) ? image : nil
-        }
-
-        if reduceMemory {
-            safetyChecker.unloadResources()
-        }
-
-        return safeImages
-    }
-
-}
-
-@available(iOS 16.2, macOS 13.1, *)
-extension StableDiffusionPipeline {
-    /// Sampling progress details
-    public struct Progress {
-        public let pipeline: StableDiffusionPipeline
-        public let prompt: String
-        public let step: Int
-        public let stepCount: Int
-        public let currentLatentSamples: [MLShapedArray<Float32>]
-        public let configuration: Configuration
-        public var isSafetyEnabled: Bool {
-            pipeline.canSafetyCheck && !configuration.disableSafety
-        }
-        public var currentImages: [CGImage?] {
-            try! pipeline.decodeToImages(currentLatentSamples, configuration: configuration)
-        }
-    }
 }

swift/StableDiffusion/pipeline/StableDiffusionXL+Resources.swift

@@ -0,0 +1,119 @@
+// For licensing see accompanying LICENSE.md file.
+// Copyright (C) 2023 Apple Inc. All Rights Reserved.
+
+import Foundation
+import CoreML
+import NaturalLanguage
+
+@available(iOS 17.0, macOS 14.0, *)
+public extension StableDiffusionXLPipeline {
+
+    struct ResourceURLs {
+
+        public let textEncoderURL: URL
+        public let textEncoder2URL: URL
+        public let unetURL: URL
+        public let unetChunk1URL: URL
+        public let unetChunk2URL: URL
+        public let unetRefinerURL: URL
+        public let unetRefinerChunk1URL: URL
+        public let unetRefinerChunk2URL: URL
+        public let decoderURL: URL
+        public let encoderURL: URL
+        public let vocabURL: URL
+        public let mergesURL: URL
+
+        public init(resourcesAt baseURL: URL) {
+            textEncoderURL = baseURL.appending(path: "TextEncoder.mlmodelc")
+            textEncoder2URL = baseURL.appending(path: "TextEncoder2.mlmodelc")
+            unetURL = baseURL.appending(path: "Unet.mlmodelc")
+            unetChunk1URL = baseURL.appending(path: "UnetChunk1.mlmodelc")
+            unetChunk2URL = baseURL.appending(path: "UnetChunk2.mlmodelc")
+            unetRefinerURL = baseURL.appending(path: "UnetRefiner.mlmodelc")
+            unetRefinerChunk1URL = baseURL.appending(path: "UnetRefinerChunk1.mlmodelc")
+            unetRefinerChunk2URL = baseURL.appending(path: "UnetRefinerChunk2.mlmodelc")
+            decoderURL = baseURL.appending(path: "VAEDecoder.mlmodelc")
+            encoderURL = baseURL.appending(path: "VAEEncoder.mlmodelc")
+            vocabURL = baseURL.appending(path: "vocab.json")
+            mergesURL = baseURL.appending(path: "merges.txt")
+        }
+    }
+
+    /// Create stable diffusion pipeline using model resources at a
+    /// specified URL
+    ///
+    /// - Parameters:
+    ///   - baseURL: URL pointing to directory holding all model and tokenization resources
+    ///   - configuration: The configuration to load model resources with
+    ///   - reduceMemory: Setup pipeline in reduced memory mode
+    /// - Returns:
+    ///  Pipeline ready for image generation if all  necessary resources loaded
+    init(
+        resourcesAt baseURL: URL,
+        configuration config: MLModelConfiguration = .init(),
+        reduceMemory: Bool = false
+    ) throws {
+
+        /// Expect URL of each resource
+        let urls = ResourceURLs(resourcesAt: baseURL)
+        let tokenizer = try BPETokenizer(mergesAt: urls.mergesURL, vocabularyAt: urls.vocabURL)
+        let textEncoder: TextEncoderXL?
+        if FileManager.default.fileExists(atPath: urls.textEncoderURL.path) {
+            textEncoder = TextEncoderXL(tokenizer: tokenizer, modelAt: urls.textEncoderURL, configuration: config)
+        } else {
+            textEncoder = nil
+        }
+        
+        // padToken is different in the second XL text encoder
+        let tokenizer2 = try BPETokenizer(mergesAt: urls.mergesURL, vocabularyAt: urls.vocabURL, padToken: "!")
+        let textEncoder2 = TextEncoderXL(tokenizer: tokenizer2, modelAt: urls.textEncoder2URL, configuration: config)
+
+        // Unet model
+        let unet: Unet
+        if FileManager.default.fileExists(atPath: urls.unetChunk1URL.path) &&
+            FileManager.default.fileExists(atPath: urls.unetChunk2URL.path) {
+            unet = Unet(chunksAt: [urls.unetChunk1URL, urls.unetChunk2URL],
+                        configuration: config)
+        } else {
+            unet = Unet(modelAt: urls.unetURL, configuration: config)
+        }
+
+        // Refiner Unet model
+        let unetRefiner: Unet?
+        if FileManager.default.fileExists(atPath: urls.unetRefinerChunk1URL.path) &&
+            FileManager.default.fileExists(atPath: urls.unetRefinerChunk2URL.path) {
+            unetRefiner = Unet(chunksAt: [urls.unetRefinerChunk1URL, urls.unetRefinerChunk2URL],
+                               configuration: config)
+        } else if FileManager.default.fileExists(atPath: urls.unetRefinerURL.path) {
+            unetRefiner = Unet(modelAt: urls.unetRefinerURL, configuration: config)
+        } else {
+            unetRefiner = nil
+        }
+
+
+        // Image Decoder
+        // FIXME: Hardcoding to .cpuAndGPU since ANE doesn't support FLOAT32
+        let vaeConfig = config.copy() as! MLModelConfiguration
+        vaeConfig.computeUnits = .cpuAndGPU
+        let decoder = Decoder(modelAt: urls.decoderURL, configuration: vaeConfig)
+
+        // Optional Image Encoder
+        let encoder: Encoder?
+        if FileManager.default.fileExists(atPath: urls.encoderURL.path) {
+            encoder = Encoder(modelAt: urls.encoderURL, configuration: vaeConfig)
+        } else {
+            encoder = nil
+        }
+
+        // Construct pipeline
+        self.init(
+            textEncoder: textEncoder,
+            textEncoder2: textEncoder2,
+            unet: unet,
+            unetRefiner: unetRefiner,
+            decoder: decoder,
+            encoder: encoder,
+            reduceMemory: reduceMemory
+        )
+    }
+}

swift/StableDiffusion/pipeline/StableDiffusionXLPipeline.swift

@@ -0,0 +1,399 @@
+// For licensing see accompanying LICENSE.md file.
+// Copyright (C) 2023 Apple Inc. All Rights Reserved.
+
+import Accelerate
+import CoreGraphics
+import CoreML
+import Foundation
+import NaturalLanguage
+
+
+/// A pipeline used to generate image samples from text input using stable diffusion XL
+///
+/// This implementation matches:
+/// [Hugging Face Diffusers XL Pipeline](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py)
+@available(iOS 17.0, macOS 14.0, *)
+public struct StableDiffusionXLPipeline: StableDiffusionPipelineProtocol {
+    
+    public typealias Configuration = PipelineConfiguration
+    public typealias Progress = PipelineProgress
+    
+    /// Model to generate embeddings for tokenized input text
+    var textEncoder: TextEncoderXLModel?
+    var textEncoder2: TextEncoderXLModel
+
+    /// Model used to predict noise residuals given an input, diffusion time step, and conditional embedding
+    var unet: Unet
+    
+    /// Model used to refine the image, if present
+    var unetRefiner: Unet?
+
+    /// Model used to generate final image from latent diffusion process
+    var decoder: Decoder
+    
+    /// Model used to latent space for image2image, and soon, in-painting
+    var encoder: Encoder?
+    
+    /// Option to reduce memory during image generation
+    ///
+    /// If true, the pipeline will lazily load TextEncoder, Unet, Decoder, and SafetyChecker
+    /// when needed and aggressively unload their resources after
+    ///
+    /// This will increase latency in favor of reducing memory
+    var reduceMemory: Bool = false
+
+    /// Creates a pipeline using the specified models and tokenizer
+    ///
+    /// - Parameters:
+    ///   - textEncoder: Model for encoding tokenized text
+    ///   - textEncoder2: Second text encoding model
+    ///   - unet: Model for noise prediction on latent samples
+    ///   - decoder: Model for decoding latent sample to image
+    ///   - reduceMemory: Option to enable reduced memory mode
+    /// - Returns: Pipeline ready for image generation
+    public init(
+        textEncoder: TextEncoderXLModel?,
+        textEncoder2: TextEncoderXLModel,
+        unet: Unet,
+        unetRefiner: Unet?,
+        decoder: Decoder,
+        encoder: Encoder?,
+        reduceMemory: Bool = false
+    ) {
+        self.textEncoder = textEncoder
+        self.textEncoder2 = textEncoder2
+        self.unet = unet
+        self.unetRefiner = unetRefiner
+        self.decoder = decoder
+        self.encoder = encoder
+        self.reduceMemory = reduceMemory
+    }
+
+    /// Load required resources for this pipeline
+    ///
+    /// If reducedMemory is true this will instead call prewarmResources instead
+    /// and let the pipeline lazily load resources as needed
+    public func loadResources() throws {
+        if reduceMemory {
+            try prewarmResources()
+        } else {
+            try textEncoder2.loadResources()
+            try unet.loadResources()
+            try decoder.loadResources()
+
+            do {
+                try textEncoder?.loadResources()
+            } catch {
+                print("Error loading resources for textEncoder: \(error)")
+            }
+
+            // Only prewarm refiner unet on load so it's unloaded until needed
+            do {
+                try unetRefiner?.prewarmResources()
+            } catch {
+                print("Error loading resources for unetRefiner: \(error)")
+            }
+
+            do {
+                try encoder?.loadResources()
+            } catch {
+                print("Error loading resources for vae encoder: \(error)")
+            }
+        }
+    }
+
+    /// Unload the underlying resources to free up memory
+    public func unloadResources() {
+        textEncoder?.unloadResources()
+        textEncoder2.unloadResources()
+        unet.unloadResources()
+        unetRefiner?.unloadResources()
+        decoder.unloadResources()
+        encoder?.unloadResources()
+    }
+
+    /// Prewarm resources one at a time
+    public func prewarmResources() throws {
+        try textEncoder2.prewarmResources()
+        try unet.prewarmResources()
+        try decoder.prewarmResources()
+
+        do {
+            try textEncoder?.prewarmResources()
+        } catch {
+            print("Error prewarming resources for textEncoder: \(error)")
+        }
+
+        do {
+            try unetRefiner?.prewarmResources()
+        } catch {
+            print("Error prewarming resources for unetRefiner: \(error)")
+        }
+
+        do {
+            try encoder?.prewarmResources()
+        } catch {
+            print("Error prewarming resources for vae encoder: \(error)")
+        }
+    }
+    /// Image generation using stable diffusion
+    /// - Parameters:
+    ///   - configuration: Image generation configuration
+    ///   - progressHandler: Callback to perform after each step, stops on receiving false response
+    /// - Returns: An array of `imageCount` optional images.
+    ///            The images will be nil if safety checks were performed and found the result to be un-safe
+    public func generateImages(
+        configuration config: Configuration,
+        progressHandler: (Progress) -> Bool = { _ in true }
+    ) throws -> [CGImage?] {
+
+        // Determine input type of Unet
+        // SDXL Refiner has a latentTimeIdShape of [2, 5]
+        // SDXL Base has either [12] or [2, 6]
+        let isRefiner = unet.latentTimeIdShape.last == 5
+
+        // Setup geometry conditioning for base/refiner inputs
+        var baseInput: ModelInputs?
+        var refinerInput: ModelInputs?
+
+        // Check if the first textEncoder is available, which is required for base models
+        if textEncoder != nil {
+            baseInput = try generateConditioning(using: config, forRefiner: isRefiner)
+        }
+
+        // Check if the refiner unet exists, or if the current unet is a refiner
+        if unetRefiner != nil || isRefiner {
+            refinerInput = try generateConditioning(using: config, forRefiner: true)
+        }
+
+        if reduceMemory {
+            textEncoder?.unloadResources()
+            textEncoder2.unloadResources()
+        }
+
+        /// Setup schedulers
+        let scheduler: [Scheduler] = (0..<config.imageCount).map { _ in
+            switch config.schedulerType {
+            case .pndmScheduler: return PNDMScheduler(stepCount: config.stepCount)
+            case .dpmSolverMultistepScheduler: return DPMSolverMultistepScheduler(stepCount: config.stepCount, timeStepSpacing: config.schedulerTimestepSpacing)
+            }
+        }
+
+        // Generate random latent samples from specified seed
+        var latents: [MLShapedArray<Float32>] = try generateLatentSamples(configuration: config, scheduler: scheduler[0])
+
+        // Store denoised latents from scheduler to pass into decoder
+        var denoisedLatents: [MLShapedArray<Float32>] = latents.map { MLShapedArray(converting: $0) }
+
+        if reduceMemory {
+            encoder?.unloadResources()
+        }
+
+        let timestepStrength: Float? = config.mode == .imageToImage ? config.strength : nil
+
+        // Store current model
+        var unetModel = unet
+        var currentInput = baseInput ?? refinerInput
+
+        var unetHiddenStates = currentInput?.hiddenStates
+        var unetPooledStates = currentInput?.pooledStates
+        var unetGeometryConditioning = currentInput?.geometryConditioning
+
+        let timeSteps: [Int] = scheduler[0].calculateTimesteps(strength: timestepStrength)
+
+        // Calculate which step to swap to refiner
+        let refinerStartStep = Int(Float(timeSteps.count) * config.refinerStart)
+
+        // De-noising loop
+        for (step,t) in timeSteps.enumerated() {
+            // Expand the latents for classifier-free guidance
+            // and input to the Unet noise prediction model
+            let latentUnetInput = latents.map {
+                MLShapedArray<Float32>(concatenating: [$0, $0], alongAxis: 0)
+            }
+
+            // Switch to refiner if specified
+            if let refiner = unetRefiner, step == refinerStartStep {
+                unet.unloadResources()
+
+                unetModel = refiner
+                currentInput = refinerInput
+                unetHiddenStates = currentInput?.hiddenStates
+                unetPooledStates = currentInput?.pooledStates
+                unetGeometryConditioning = currentInput?.geometryConditioning
+            }
+
+            guard let hiddenStates = unetHiddenStates,
+                  let pooledStates = unetPooledStates,
+                  let geometryConditioning = unetGeometryConditioning else {
+                throw PipelineError.missingUnetInputs
+            }
+
+            // Predict noise residuals from latent samples
+            // and current time step conditioned on hidden states
+            var noise = try unetModel.predictNoise(
+                latents: latentUnetInput,
+                timeStep: t,
+                hiddenStates: hiddenStates,
+                pooledStates: pooledStates,
+                geometryConditioning: geometryConditioning
+            )
+
+            noise = performGuidance(noise, config.guidanceScale)
+
+            // Have the scheduler compute the previous (t-1) latent
+            // sample given the predicted noise and current sample
+            for i in 0..<config.imageCount {
+                latents[i] = scheduler[i].step(
+                    output: noise[i],
+                    timeStep: t,
+                    sample: latents[i]
+                )
+                
+                denoisedLatents[i] = scheduler[i].modelOutputs.last ?? latents[i]
+            }
+
+            let currentLatentSamples = config.useDenoisedIntermediates ? denoisedLatents : latents
+
+            // Report progress
+            let progress = Progress(
+                pipeline: self,
+                prompt: config.prompt,
+                step: step,
+                stepCount: timeSteps.count,
+                currentLatentSamples: currentLatentSamples,
+                configuration: config
+            )
+            if !progressHandler(progress) {
+                // Stop if requested by handler
+                return []
+            }
+        }
+
+        // Unload resources
+        if reduceMemory {
+            unet.unloadResources()
+        }
+        unetRefiner?.unloadResources()
+
+
+        // Decode the latent samples to images
+        return try decodeToImages(denoisedLatents, configuration: config)
+    }
+
+    func encodePrompt(_ prompt: String, forRefiner: Bool = false) throws -> (MLShapedArray<Float32>, MLShapedArray<Float32>) {
+        var embeds = MLShapedArray<Float32>()
+        var pooled = MLShapedArray<Float32>()
+
+        if forRefiner {
+            let (embeds2, pooledValue) = try textEncoder2.encode(prompt)
+
+            // Refiner only takes textEncoder2 embeddings
+            // [1, 77, 1280]
+            embeds = embeds2
+            pooled = pooledValue
+        } else {
+            guard let encoder = textEncoder else {
+                throw PipelineError.startingText2ImgWithoutTextEncoder
+            }
+            let (embeds1, _) = try encoder.encode(prompt)
+            let (embeds2, pooledValue) = try textEncoder2.encode(prompt)
+
+            // Base needs concatenated embeddings
+            // [1, 77, 768], [1, 77, 1280] -> [1, 77, 2048]
+            embeds = MLShapedArray<Float32>(
+                concatenating: [embeds1, embeds2],
+                alongAxis: 2
+            )
+            pooled = pooledValue
+        }
+
+        return (embeds, pooled)
+    }
+
+    func generateConditioning(using config: Configuration, forRefiner: Bool = false) throws -> ModelInputs {
+        // Encode the input prompt and negative prompt
+        let (promptEmbedding, pooled) = try encodePrompt(config.prompt, forRefiner: forRefiner)
+        let (negativePromptEmbedding, negativePooled) = try encodePrompt(config.negativePrompt, forRefiner: forRefiner)
+
+        // Convert to Unet hidden state representation
+        // Concatenate the prompt and negative prompt embeddings
+        let hiddenStates = toHiddenStates(MLShapedArray(concatenating: [negativePromptEmbedding, promptEmbedding], alongAxis: 0))
+        let pooledStates = MLShapedArray(concatenating: [negativePooled, pooled], alongAxis: 0)
+
+        // Inline helper functions for geometry creation
+        func refinerGeometry() -> MLShapedArray<Float32> {
+            let negativeGeometry = MLShapedArray<Float32>(
+                scalars: [
+                    config.originalSize, config.originalSize,
+                    config.cropsCoordsTopLeft, config.cropsCoordsTopLeft,
+                    config.negativeAestheticScore
+                ],
+                shape: [1, 5]
+            )
+            let positiveGeometry = MLShapedArray<Float32>(
+                scalars: [
+                    config.originalSize, config.originalSize,
+                    config.cropsCoordsTopLeft, config.cropsCoordsTopLeft,
+                    config.aestheticScore
+                ],
+                shape: [1, 5]
+            )
+            return MLShapedArray<Float32>(concatenating: [negativeGeometry, positiveGeometry], alongAxis: 0)
+        }
+
+        func baseGeometry() -> MLShapedArray<Float32> {
+            let geometry = MLShapedArray<Float32>(
+                scalars: [
+                    config.originalSize, config.originalSize,
+                    config.cropsCoordsTopLeft, config.cropsCoordsTopLeft,
+                    config.targetSize, config.targetSize
+                ],
+                // TODO: This checks if the time_ids input is looking for [12] or [2, 6]
+                // Remove once model input shapes are ubiquitous
+                shape: unet.latentTimeIdShape.count > 1 ? [1, 6] : [6]
+            )
+            return MLShapedArray<Float32>(concatenating: [geometry, geometry], alongAxis: 0)
+        }
+
+        let geometry = forRefiner ? refinerGeometry() : baseGeometry()
+
+        return ModelInputs(hiddenStates: hiddenStates, pooledStates: pooledStates, geometryConditioning: geometry)
+    }
+
+    func generateLatentSamples(configuration config: Configuration, scheduler: Scheduler) throws -> [MLShapedArray<Float32>] {
+        var sampleShape = unet.latentSampleShape
+        sampleShape[0] = 1
+        
+        let stdev = scheduler.initNoiseSigma
+        var random = randomSource(from: config.rngType, seed: config.seed)
+        let samples = (0..<config.imageCount).map { _ in
+            MLShapedArray<Float32>(
+                converting: random.normalShapedArray(sampleShape, mean: 0.0, stdev: Double(stdev)))
+        }
+        if let image = config.startingImage, config.mode == .imageToImage {
+            guard let encoder else {
+                throw PipelineError.startingImageProvidedWithoutEncoder
+            }
+            let latent = try encoder.encode(image, scaleFactor: config.encoderScaleFactor, random: &random)
+            return scheduler.addNoise(originalSample: latent, noise: samples, strength: config.strength)
+        }
+        return samples
+    }
+
+    public func decodeToImages(_ latents: [MLShapedArray<Float32>], configuration config: Configuration) throws -> [CGImage?] {
+        defer {
+            if reduceMemory {
+                decoder.unloadResources()
+            }
+        }
+        
+        return try decoder.decode(latents, scaleFactor: config.decoderScaleFactor)
+    }
+
+    struct ModelInputs {
+        var hiddenStates: MLShapedArray<Float32>
+        var pooledStates: MLShapedArray<Float32>
+        var geometryConditioning: MLShapedArray<Float32>
+    }
+}

swift/StableDiffusion/pipeline/TextEncoderXL.swift

@@ -0,0 +1,99 @@
+// For licensing see accompanying LICENSE.md file.
+// Copyright (C) 2023 Apple Inc. All Rights Reserved.
+
+import Foundation
+import CoreML
+
+@available(iOS 17.0, macOS 14.0, *)
+public protocol TextEncoderXLModel: ResourceManaging {
+    typealias TextEncoderXLOutput = (hiddenEmbeddings: MLShapedArray<Float32>, pooledOutputs: MLShapedArray<Float32>)
+    func encode(_ text: String) throws -> TextEncoderXLOutput
+}
+
+///  A model for encoding text, suitable for SDXL
+@available(iOS 17.0, macOS 14.0, *)
+public struct TextEncoderXL: TextEncoderXLModel {
+
+    /// Text tokenizer
+    var tokenizer: BPETokenizer
+
+    /// Embedding model
+    var model: ManagedMLModel
+
+    /// Creates text encoder which embeds a tokenized string
+    ///
+    /// - Parameters:
+    ///   - tokenizer: Tokenizer for input text
+    ///   - url: Location of compiled text encoding  Core ML model
+    ///   - configuration: configuration to be used when the model is loaded
+    /// - Returns: A text encoder that will lazily load its required resources when needed or requested
+    public init(tokenizer: BPETokenizer,
+                modelAt url: URL,
+                configuration: MLModelConfiguration) {
+        self.tokenizer = tokenizer
+        self.model = ManagedMLModel(modelAt: url, configuration: configuration)
+    }
+
+    /// Ensure the model has been loaded into memory
+    public func loadResources() throws {
+        try model.loadResources()
+    }
+
+    /// Unload the underlying model to free up memory
+    public func unloadResources() {
+       model.unloadResources()
+    }
+
+    /// Encode input text/string
+    ///
+    ///  - Parameters:
+    ///     - text: Input text to be tokenized and then embedded
+    ///  - Returns: Embedding representing the input text
+    public func encode(_ text: String) throws -> TextEncoderXLOutput {
+
+        // Get models expected input length
+        let inputLength = inputShape.last!
+
+        // Tokenize, padding to the expected length
+        var (tokens, ids) = tokenizer.tokenize(input: text, minCount: inputLength)
+
+        // Truncate if necessary
+        if ids.count > inputLength {
+            tokens = tokens.dropLast(tokens.count - inputLength)
+            ids = ids.dropLast(ids.count - inputLength)
+            let truncated = tokenizer.decode(tokens: tokens)
+            print("Needed to truncate input '\(text)' to '\(truncated)'")
+        }
+
+        // Use the model to generate the embedding
+        return try encode(ids: ids)
+    }
+
+    func encode(ids: [Int]) throws -> TextEncoderXLOutput {
+        let inputName = inputDescription.name
+        let inputShape = inputShape
+
+        let floatIds = ids.map { Float32($0) }
+        let inputArray = MLShapedArray<Float32>(scalars: floatIds, shape: inputShape)
+        let inputFeatures = try! MLDictionaryFeatureProvider(
+            dictionary: [inputName: MLMultiArray(inputArray)])
+
+        let result = try model.perform { model in
+            try model.prediction(from: inputFeatures)
+        }
+
+        let embeddingFeature = result.featureValue(for: "hidden_embeds")
+        let pooledFeature = result.featureValue(for: "pooled_outputs")
+        return (MLShapedArray<Float32>(converting: embeddingFeature!.multiArrayValue!), MLShapedArray<Float32>(converting: pooledFeature!.multiArrayValue!))
+    }
+
+    var inputDescription: MLFeatureDescription {
+        try! model.perform { model in
+            model.modelDescription.inputDescriptionsByName.first!.value
+        }
+    }
+
+    var inputShape: [Int] {
+        inputDescription.multiArrayConstraint!.shape.map { $0.intValue }
+    }
+}

swift/StableDiffusion/pipeline/Unet.swift

@@ -69,6 +69,17 @@ public struct Unet: ResourceManaging {
         latentSampleDescription.multiArrayConstraint!.shape.map { $0.intValue }
     }
 
+    var latentTimeIdDescription: MLFeatureDescription {
+        try! models.first!.perform { model in
+            model.modelDescription.inputDescriptionsByName["time_ids"]!
+        }
+    }
+
+    /// The expected shape of the geometry conditioning
+    public var latentTimeIdShape: [Int] {
+        latentTimeIdDescription.multiArrayConstraint!.shape.map { $0.intValue }
+    }
+
     /// Batch prediction noise from latent samples
     ///
     /// - Parameters:
@@ -127,6 +138,65 @@ public struct Unet: ResourceManaging {
         return noise
     }
 
+    /// Batch prediction noise from latent samples, for Stable Diffusion XL
+    ///
+    /// - Parameters:
+    ///   - latents: Batch of latent samples in an array
+    ///   - timeStep: Current diffusion timestep
+    ///   - hiddenStates: Hidden state to condition on
+    ///   - pooledStates: Additional text states to condition on
+    ///   - geometryConditioning: Condition on image geometry
+    /// - Returns: Array of predicted noise residuals
+    @available(iOS 17.0, macOS 14.0, *)
+    func predictNoise(
+        latents: [MLShapedArray<Float32>],
+        timeStep: Int,
+        hiddenStates: MLShapedArray<Float32>,
+        pooledStates: MLShapedArray<Float32>,
+        geometryConditioning: MLShapedArray<Float32>
+    ) throws -> [MLShapedArray<Float32>] {
+
+        // Match time step batch dimension to the model / latent samples
+        let t = MLShapedArray<Float32>(scalars:[Float(timeStep), Float(timeStep)],shape:[2])
+
+        // Form batch input to model
+        let inputs = try latents.enumerated().map {
+            let dict: [String: Any] = [
+                "sample" : MLMultiArray($0.element),
+                "timestep" : MLMultiArray(t),
+                "encoder_hidden_states": MLMultiArray(hiddenStates),
+                "text_embeds": MLMultiArray(pooledStates),
+                "time_ids": MLMultiArray(geometryConditioning)
+            ]
+            return try MLDictionaryFeatureProvider(dictionary: dict)
+        }
+        let batch = MLArrayBatchProvider(array: inputs)
+
+        // Make predictions
+        let results = try predictions(from: batch)
+
+        // Pull out the results in Float32 format
+        let noise = (0..<results.count).map { i in
+
+            let result = results.features(at: i)
+            let outputName = result.featureNames.first!
+
+            let outputNoise = result.featureValue(for: outputName)!.multiArrayValue!
+
+            // To conform to this func return type make sure we return float32
+            // Use the fact that the concatenating constructor for MLMultiArray
+            // can do type conversion:
+            let fp32Noise = MLMultiArray(
+                concatenating: [outputNoise],
+                axis: 0,
+                dataType: .float32
+            )
+            return MLShapedArray<Float32>(fp32Noise)
+        }
+
+        return noise
+    }
+
     func predictions(from batch: MLBatchProvider) throws -> MLBatchProvider {
 
         var results = try models.first!.perform { model in

swift/StableDiffusion/tokenizer/BPETokenizer.swift

@@ -12,15 +12,15 @@ public struct BPETokenizer {
     /// A dictionary from of tokens to identifiers.
     let vocabulary: [String: Int]
 
+    /// The token used for padding
+    let padToken: String
+
     /// The start token.
     let startToken: String = "<|startoftext|>"
 
     /// The end token.
     let endToken: String = "<|endoftext|>"
 
-    /// The token used for padding
-    let padToken: String = "<|endoftext|>"
-
     /// The unknown token.
     let unknownToken: String = "<|endoftext|>"
 
@@ -33,9 +33,10 @@ public struct BPETokenizer {
     /// - Parameters:
     ///   - merges: A dictionary that maps pairs of tokens to the rank/order of the merge.
     ///   - vocabulary: A dictionary from of tokens to identifiers.
-    public init(merges: [TokenPair: Int], vocabulary: [String: Int]) {
+    public init(merges: [TokenPair: Int], vocabulary: [String: Int], padToken: String = "<|endoftext|>") {
         self.merges = merges
         self.vocabulary = vocabulary
+        self.padToken = padToken
     }
 
     /// Creates a tokenizer by loading merges and vocabulary from URLs.
@@ -43,9 +44,10 @@ public struct BPETokenizer {
     /// - Parameters:
     ///   - mergesURL: The URL of a text file containing merges.
     ///   - vocabularyURL: The URL of a JSON file containing the vocabulary.
-    public init(mergesAt mergesURL: URL, vocabularyAt vocabularyURL: URL) throws {
+    public init(mergesAt mergesURL: URL, vocabularyAt vocabularyURL: URL, padToken: String = "<|endoftext|>") throws {
         self.merges = try Self.readMerges(url: mergesURL)
         self.vocabulary = try! Self.readVocabulary(url: vocabularyURL)
+        self.padToken = padToken
     }
 
     /// Tokenizes an input string.

swift/StableDiffusionCLI/main.swift

@@ -36,6 +36,9 @@ struct StableDiffusionSample: ParsableCommand {
     )
     var resourcePath: String = "./"
     
+    @Flag(name: .customLong("xl"), help: "The resources correspond to a Stable Diffusion XL model")
+    var isXL: Bool = false
+
     @Option(help: "Path to starting image.")
     var image: String? = nil
     
@@ -71,7 +74,7 @@ struct StableDiffusionSample: ParsableCommand {
     @Option(help: "Scheduler to use, one of {pndm, dpmpp}")
     var scheduler: SchedulerOption = .pndm
 
-    @Option(help: "Random number generator to use, one of {numpy, torch}")
+    @Option(help: "Random number generator to use, one of {numpy, torch, nvidia}")
     var rng: RNGOption = .numpy
     
     @Option(
@@ -109,17 +112,33 @@ struct StableDiffusionSample: ParsableCommand {
 
         log("Loading resources and creating pipeline\n")
         log("(Note: This can take a while the first time using these resources)\n")
-        let pipeline: StableDiffusionPipeline
+        let pipeline: StableDiffusionPipelineProtocol
+        var scaleFactor: Float32 = 0.18215
         if #available(macOS 14.0, iOS 17.0, *) {
-            pipeline = try StableDiffusionPipeline(
-                resourcesAt: resourceURL,
-                controlNet: controlnet,
-                configuration: config,
-                disableSafety: disableSafety,
-                reduceMemory: reduceMemory,
-                useMultilingualTextEncoder: useMultilingualTextEncoder,
-                script: script
-            )
+            if isXL {
+                scaleFactor = 0.13025
+                if !controlnet.isEmpty {
+                    throw RunError.unsupported("ControlNet is not supported for Stable Diffusion XL")
+                }
+                if useMultilingualTextEncoder {
+                    throw RunError.unsupported("Multilingual text encoder is not yet supported for Stable Diffusion XL")
+                }
+                pipeline = try StableDiffusionXLPipeline(
+                    resourcesAt: resourceURL,
+                    configuration: config,
+                    reduceMemory: reduceMemory
+                )
+            } else {
+                pipeline = try StableDiffusionPipeline(
+                    resourcesAt: resourceURL,
+                    controlNet: controlnet,
+                    configuration: config,
+                    disableSafety: disableSafety,
+                    reduceMemory: reduceMemory,
+                    useMultilingualTextEncoder: useMultilingualTextEncoder,
+                    script: script
+                )
+            }
         } else  {
             pipeline = try StableDiffusionPipeline(
                 resourcesAt: resourceURL,
@@ -176,17 +195,19 @@ struct StableDiffusionSample: ParsableCommand {
         pipelineConfig.guidanceScale = guidanceScale
         pipelineConfig.schedulerType = scheduler.stableDiffusionScheduler
         pipelineConfig.rngType = rng.stableDiffusionRNG
+        pipelineConfig.useDenoisedIntermediates = true
+        pipelineConfig.encoderScaleFactor = scaleFactor
+        pipelineConfig.decoderScaleFactor = scaleFactor
 
         let images = try pipeline.generateImages(
-            configuration: pipelineConfig,
-            progressHandler: { progress in
+            configuration: pipelineConfig) { progress in
                 sampleTimer.stop()
                 handleProgress(progress,sampleTimer)
                 if progress.stepCount != progress.step {
                     sampleTimer.start()
                 }
                 return true
-            })
+            }
 
         _ = try saveImages(images, logNames: true)
     }
@@ -285,6 +306,7 @@ struct StableDiffusionSample: ParsableCommand {
 enum RunError: Error {
     case resources(String)
     case saving(String)
+    case unsupported(String)
 }
 
 @available(iOS 16.2, macOS 13.1, *)
@@ -313,11 +335,12 @@ enum SchedulerOption: String, ExpressibleByArgument {
 
 @available(iOS 16.2, macOS 13.1, *)
 enum RNGOption: String, ExpressibleByArgument {
-    case numpy, torch
+    case numpy, torch, nvidia
     var stableDiffusionRNG: StableDiffusionRNG {
         switch self {
         case .numpy: return .numpyRNG
         case .torch: return .torchRNG
+        case .nvidia: return .nvidiaRNG
         }
     }
 }