Merge pull request #368 from Aintor/main

Fix for Mask Patch Failure and Quantization Issues in Latest `transformers` Versions

.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

@@ -1,4 +1,4 @@
-// swift-tools-version: 5.7
+// swift-tools-version: 5.8
 // The swift-tools-version declares the minimum version of Swift required to build this package.
 
 import PackageDescription
@@ -6,8 +6,8 @@ import PackageDescription
 let package = Package(
     name: "stable-diffusion",
     platforms: [
-        .macOS(.v11),
-        .iOS(.v14),
+        .macOS(.v13),
+        .iOS(.v16),
     ],
     products: [
         .library(
@@ -18,12 +18,15 @@ 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"),
+        .package(url: "https://github.com/huggingface/swift-transformers.git", exact: "0.1.8"),
     ],
     targets: [
         .target(
             name: "StableDiffusion",
-            dependencies: [],
+            dependencies:  [
+                .product(name: "Transformers", package: "swift-transformers"),
+            ],
             path: "swift/StableDiffusion"),
         .executableTarget(
             name: "StableDiffusionCLI",

README.md

@@ -2,7 +2,8 @@
 
 Run Stable Diffusion on Apple Silicon with Core ML
 
-<img src="assets/readme_reel.png">
+[\[Blog Post\]](https://machinelearning.apple.com/research/stable-diffusion-coreml-apple-silicon) [\[BibTeX\]](#bibtex)
+
 
 This repository comprises:
 
@@ -11,38 +12,466 @@ This repository comprises:
 
 If you run into issues during installation or runtime, please refer to the [FAQ](#faq) section. Please refer to the [System Requirements](#system-requirements) section before getting started.
 
-
-## <a name="example-results"></a> Example Results
-
-There are numerous versions of Stable Diffusion available on the [Hugging Face Hub](https://huggingface.co/models?search=stable-diffusion). Here are example results from three of those models:
-
-`--model-version` | [stabilityai/stable-diffusion-2-base](https://huggingface.co/stabilityai/stable-diffusion-2-base) |  [CompVis/stable-diffusion-v1-4](https://huggingface.co/CompVis/stable-diffusion-v1-4) |  [runwayml/stable-diffusion-v1-5](https://huggingface.co/runwayml/stable-diffusion-v1-5) |
-:------:|:------:|:------:|:------:
-Output | ![](assets/a_high_quality_photo_of_an_astronaut_riding_a_horse_in_space/randomSeed_11_computeUnit_CPU_AND_GPU_modelVersion_stabilityai_stable-diffusion-2-base.png) | ![](assets/a_high_quality_photo_of_an_astronaut_riding_a_horse_in_space/randomSeed_13_computeUnit_CPU_AND_NE_modelVersion_CompVis_stable-diffusion-v1-4.png) | ![](assets/a_high_quality_photo_of_an_astronaut_riding_a_horse_in_space/randomSeed_93_computeUnit_CPU_AND_NE_modelVersion_runwayml_stable-diffusion-v1-5.png)
-M1 iPad Pro 8GB Latency (s)     | 29 | 38 | 38 |
-M1 MacBook Pro 16GB Latency (s) | 24 | 35 | 35 |
-M2 MacBook Air 8GB Latency (s)  | 18 | 23 | 23 |
-
-Please see [Important Notes on Performance Benchmarks](#important-notes-on-performance-benchmarks) section for details.
+<img src="assets/readme_reel.png">
 
 ## <a name="system-requirements"></a> System Requirements
 
-The following is recommended to use all the functionality in this repository:
+<details>
+  <summary> Details (Click to expand) </summary>
 
-Python | macOS | Xcode | iPadOS, iOS |
-:------:|:------:|:------:|:------:|
-3.8 | 13.1 | 14.2 | 16.2 |
+Model Conversion:
+
+ macOS  | Python | coremltools |
+:------:|:------:|:-----------:|
+  13.1  | 3.8    |    7.0      |
+
+Project Build:
+
+  macOS | Xcode | Swift |
+:------:|:-----:|:-----:|
+  13.1  | 14.3  |  5.8  |
+
+Target Device Runtime:
+
+  macOS | iPadOS, iOS |
+:------:|:-----------:|
+  13.1  |     16.2    |
+
+Target Device Runtime ([With Memory Improvements](#compression-6-bits-and-higher)):
+
+  macOS | iPadOS, iOS |
+:------:|:-----------:|
+  14.0  |     17.0    |
+
+Target Device Hardware Generation:
+
+  Mac   |  iPad   | iPhone  |
+:------:|:-------:|:-------:|
+   M1   |   M1    |  A14    |
+
+</details>
+
+
+## <a name="performance-benchmark"></a> Performance Benchmarks
+
+<details>
+  <summary> Details (Click to expand) </summary>
+
+
+[`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 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-6-bits-and-higher) for details.
+- Activations are in float16 precision for both the GPU and the Neural Engine.
+- `*` 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-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.
+
+We demonstrate how data-free [post-training palettization](https://coremltools.readme.io/v7.0/docs/post-training-palettization) implemented in `coremltools.optimize.coreml.palettize_weights` enables us to achieve greatly improved performance for Stable Diffusion on mobile devices. This API implements the [Fast Exact k-Means](https://arxiv.org/abs/1701.07204) algorithm for optimal weight clustering which yields more accurate palettes. Using `--quantize-nbits {2,4,6,8}` during [conversion](#converting-models-to-coreml) is going to apply this compression to the unet and text_encoder models.
+
+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: 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"* |
+| :---------------:| :----------------: | ------------------------------------------------------  |
+| 6-bit            | cpuAndNeuralEngine | <img src="assets/palette6_cpuandne_readmereel.png"> |
+| 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 (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 progressively 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="activation-quant"></a> Activation Quantization
+
+<details>
+  <summary> Details (Click to expand) </summary>
+
+On newer hardware with A17 Pro or M4 chips, such as the iPhone 15 Pro, quantizing both activations and weight to int8 can leverage optimized compute on the Neural Engine which can be used to improve runtime latency in compute-bound models.
+
+In this section, we demonstrate how to apply [Post Training Activation Quantization](https://apple.github.io/coremltools/docs-guides/source/opt-quantization-algos.html#post-training-data-calibration-activation-quantization), using calibration data, on Stable Diffusion UNet model. 
+
+Similar to Mixed-Bit Palettization (MBP) described [above](#a-namecompression-lower-than-6-bitsa-advanced-weight-compression-lower-than-6-bits), first, a per-layer analysis is run to determine which intermediate activations are more sensitive to 8-bit compression.
+Less sensitive layers are weight and activation quantized (W8A8), whereas more sensitive layers are only weight quantized (W8A16).
+
+Here are the steps for applying this technique:
+
+**Step 1:** Generate calibration data 
+            
+```python
+python -m python_coreml_stable_diffusion.activation_quantization --model-version <model-version> --generate-calibration-data  -o <output-dir>
+```
+
+A set of calibration text prompts are run through StableDiffusionPipeline and UNet model inputs are recorded and stored as pickle files in `calibration_data_<model-version>` folder inside specified output directory.
+
+**Step 2:** Run layer-wise sensitivity analysis 
+
+```python
+python -m python_coreml_stable_diffusion.activation_quantization --model-version <model-version> --layerwise-sensitivity --calibration-nsamples <num-samples> -o <output-dir>
+```
+
+This will run the analysis on all Convolutional and Attention (Einsum) modules in the model.
+For each module, a compressed version is generated by quantizing only that layer’s weights and activations.
+Then the PSNR between the outputs of the compressed and original model is calculated, using the same random seed and text prompts.
+
+This analysis takes up to a few hours on a single GPU (cuda). The number of calibration samples used to quantize the model can be reduced to speed up the process. 
+ 
+The resulting JSON file looks like this:
+
+```json
+{
+  "conv": {
+    "conv_in": 30.74,
+    "down_blocks.0.attentions.0.proj_in": 38.93,
+    "down_blocks.0.attentions.0.transformer_blocks.0.attn1.to_q": 48.15,
+    "down_blocks.0.attentions.0.transformer_blocks.0.attn1.to_k": 50.13,
+    "down_blocks.0.attentions.0.transformer_blocks.0.attn1.to_v": 45.70,
+    "down_blocks.0.attentions.0.transformer_blocks.0.attn1.to_out.0": 39.56,
+    ...
+  },
+  "einsum": {
+    "down_blocks.0.attentions.0.transformer_blocks.0.attn1.einsum": 25.34,
+    "down_blocks.0.attentions.0.transformer_blocks.0.attn2.einsum": 31.76,
+    "down_blocks.0.attentions.1.transformer_blocks.0.attn1.einsum": 23.40,
+    "down_blocks.0.attentions.1.transformer_blocks.0.attn2.einsum": 31.56,
+    ...
+  },
+  "model_version": "stabilityai/stable-diffusion-2-1-base"
+}
+```
+
+**Step 3:** Generate quantized model 
+
+Using calibration data and layer-wise sensitivity the quantized CoreML model can be generated as follows:
+
+```python
+python -m python_coreml_stable_diffusion.activation_quantization --model-version <model-version> --quantize-pytorch --conv-psnr 38 --attn-psnr 26 -o <output-dir>
+```
+
+The PSNR thresholds determine which layers will be activation quantized. This number can be tuned to trade-off between output quality and inference latency.  
+
+</details>
+
+## <a name="using-stable-diffusion-3"></a> Using Stable Diffusion 3
+
+<details>
+  <summary> Details (Click to expand) </summary>
+
+### Model Conversion
+
+Stable Diffusion 3 uses some new and some old models to run. For the text encoders, the conversion can be done using a similar command as before with the `--sd3-version` flag.
+
+```bash
+python -m python_coreml_stable_diffusion.torch2coreml --model-version stabilityai/stable-diffusion-3-medium --bundle-resources-for-swift-cli --convert-text-encoder --sd3-version -o <output-dir>
+```
+
+For the new models (MMDiT, a new VAE with 16 channels, and the T5 text encoder), there are a number of new CLI flags that utilize the [DiffusionKit](https://www.github.com/argmaxinc/DiffusionKit) repo:
+
+- `--sd3-version`: Indicates to the converter to treat this as a Stable Diffusion 3 model
+- `--convert-mmdit`: Convert the MMDiT model
+- `--convert-vae-decoder`: Convert the new VAE model (this will use the 16 channel version if --sd3-version is set)
+- `--include-t5`: Downloads and includes a pre-converted T5 text encoder in the conversion
+
+e.g.:
+```bash
+python -m python_coreml_stable_diffusion.torch2coreml --model-version stabilityai/stable-diffusion-3-medium --bundle-resources-for-swift-cli --convert-vae-decoder --convert-mmdit  --include-t5 --sd3-version -o <output-dir>
+```
+
+To convert the full pipeline with at 1024x1024 resolution, the following command may be used:
+
+```bash
+python -m python_coreml_stable_diffusion.torch2coreml --model-version stabilityai/stable-diffusion-3-medium --bundle-resources-for-swift-cli --convert-text-encoder --convert-vae-decoder --convert-mmdit --include-t5 --sd3-version --latent-h 128 --latent-w 128 -o <output-dir>
+```
+
+Keep in mind that the MMDiT model is quite large and will require increasingly more memory and time to convert as the latent resolution increases.
+
+Also note that currently the MMDiT model requires fp32 and therefore only supports `CPU_AND_GPU` compute units and `ORIGINAL` attention implementation (the default for this pipeline).
+
+### Swift Inference
+
+Swift inference for Stable Diffusion 3 is similar to the previous versions. The only difference is that the `--sd3` flag should be used to indicate that the model is a Stable Diffusion 3 model.
+
+```bash
+swift run StableDiffusionSample <prompt> --resource-path <output-mlpackages-directory/Resources> --output-path <output-dir> --compute-units cpuAndGPU --sd3
+```
+
+</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 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:
+
+**Step 1:** Curate an image-text dataset with the desired languages.
+
+**Step 2:** Pre-compute the NLContextualEmbedding values and replace the text strings with these embedding vectors in your dataset.
+
+**Step 3:** Fine-tune a base model from Hugging Face Hub that is compatible with the [StableDiffusionPipeline](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/overview) by using your new dataset and replacing the default text_encoder with your pre-computed NLContextualEmbedding values.
+
+**Step 4:** In order to be able to swap the text_encoder of a base model without training new layers, the base model's `text_encoder.hidden_size` must match that of NLContextualEmbedding. If it doesn't, you will need to train a linear projection layer to map between the two dimensionalities. After fine-tuning, this linear layer should be converted to CoreML as follows:
+
+```shell
+python -m python_coreml_stable_diffusion.multilingual_projection --input-path <path-to-projection-torchscript> --output-dir <output-dir>
+```
+
+The command above will yield a `MultilingualTextEncoderProjection.mlmodelc` file under `--output-dir` and this should be colocated with the rest of the Core ML model assets that were generated through `--bundle-resources-for-swift-cli`.
+
+**Step 5:** The multilingual system text encoder can now be invoked by setting `useMultilingualTextEncoder` to true when initializing a pipeline or setting `--use-multilingual-text-encoder` in the CLI. Note that the model assets are distributed over-the-air so the first invocation will trigger asset downloads which is less than 100MB.
+
+
+Resources:
+- [WWDC23 Session Video: Explore Natural Language multilingual models](https://developer.apple.com/videos/play/wwdc2023/10042)
+- [NLContextualEmbedding API Documentation](https://developer.apple.com/documentation/naturallanguage/nlcontextualembedding)
+
+</details>
 
 ## <a name="using-converted-weights"></a> Using Ready-made Core ML Models from Hugging Face Hub
 
 <details>
   <summary> Click to expand </summary>
 
-🤗 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).
+🤗 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-coreml).
 
-* [`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)
+* 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)
+  - [`stabilityai/stable-diffusion-3-medium`](https://huggingface.co/stabilityai/stable-diffusion-3-medium)
 
 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).
 
@@ -66,33 +495,17 @@ git lfs install
 git clone https://huggingface.co/apple/coreml-stable-diffusion-v1-4
 ```
 
-If you prefer to download specific variants instead of cloning the repos, you can use the `huggingface_hub` Python library. For example, to do generation in Python using the `ORIGINAL` attention implementation (read [this section](#converting-models-to-core-ml) for details), you could use the following helper code:
+If you prefer to download specific variants instead of cloning the repos, you can use the `huggingface_hub` Python library. For example, to do generation in Python using the `ORIGINAL` attention implementation (read [this section](#converting-models-to-coreml) for details), you could use the following helper code:
 
 ```Python
 from huggingface_hub import snapshot_download
-from huggingface_hub.file_download import repo_folder_name
 from pathlib import Path
-import shutil
 
 repo_id = "apple/coreml-stable-diffusion-v1-4"
 variant = "original/packages"
 
-def download_model(repo_id, variant, output_dir):
-    destination = Path(output_dir) / (repo_id.split("/")[-1] + "_" + variant.replace("/", "_"))
-    if destination.exists():
-        raise Exception(f"Model already exists at {destination}")
-    
-    # Download and copy without symlinks
-    downloaded = snapshot_download(repo_id, allow_patterns=f"{variant}/*", cache_dir=output_dir)
-    downloaded_bundle = Path(downloaded) / variant
-    shutil.copytree(downloaded_bundle, destination)
-
-    # Remove all downloaded files
-    cache_folder = Path(output_dir) / repo_folder_name(repo_id=repo_id, repo_type="model")
-    shutil.rmtree(cache_folder)
-    return destination
-
-model_path = download_model(repo_id, variant, output_dir="./models")
+model_path = Path("./models") / (repo_id.split("/")[-1] + "_" + variant.replace("/", "_"))
+snapshot_download(repo_id, allow_patterns=f"{variant}/*", local_dir=model_path, local_dir_use_symlinks=False)
 print(f"Model downloaded at {model_path}")
 ```
 
@@ -121,21 +534,24 @@ 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).
 
-- `--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. 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. Chunking is for on-device deployment with Swift only.
+- `--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.
 
-- `--attention-implementation`: Defaults to `SPLIT_EINSUM` which is the implementation described in [Deploying Transformers on the Apple Neural Engine](https://machinelearning.apple.com/research/neural-engine-transformers). `--attention-implementation ORIGINAL` will switch to an alternative that should be used for CPU or GPU deployment. Please refer to the [Performance Benchmark](#performance-benchmark) section for further guidance.
+- `--chunk-unet`: Splits the Unet model in two approximately equal chunks (each with less than 1GB of weights) for mobile-friendly deployment. This is **required** for Neural Engine deployment on iOS and iPadOS if weights are not quantized to 6-bits or less (`--quantize-nbits {2,4,6}`). This is not required for macOS. Swift CLI is able to consume both the chunked and regular versions of the Unet model but prioritizes the former. Note that chunked unet is not compatible with the Python pipeline because Python pipeline is intended for macOS only.
+
+- `--attention-implementation`: Defaults to `SPLIT_EINSUM` which is the implementation described in [Deploying Transformers on the Apple Neural Engine](https://machinelearning.apple.com/research/neural-engine-transformers). `--attention-implementation SPLIT_EINSUM_V2` yields 10-30% improvement for mobile devices, still targeting the Neural Engine. `--attention-implementation ORIGINAL` will switch to an alternative implementation that should be used for CPU or GPU deployment on some Mac devices. Please refer to the [Performance Benchmark](#performance-benchmark) section for further guidance.
 
 - `--check-output-correctness`: Compares original PyTorch model's outputs to final Core ML model's outputs. This flag increases RAM consumption significantly so it is recommended only for debugging purposes.
 
@@ -143,6 +559,10 @@ 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.
 
+- `--unet-batch-one`: use a batch size of one for the unet, this is needed if you do not want to do classifier free guidance, i.e. using a `guidance-scale` of less than one.
+
+- `--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
@@ -153,16 +573,17 @@ 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.
 - `--controlnet`: ControlNet models specified with this option are used in image generation. Use this option in the format `--controlnet lllyasviel/sd-controlnet-mlsd lllyasviel/sd-controlnet-depth` and make sure to use `--controlnet-inputs` in conjunction.
 - `--controlnet-inputs`: Image inputs corresponding to each ControlNet model. Please provide image paths in same order as models in `--controlnet`, for example: `--controlnet-inputs image_mlsd image_depth`.
+- `--unet-batch-one`: Do not batch unet predictions for the prompt and negative prompt. This requires the unet has been converted with a batch size of one, see `--unet-batch-one` option in conversion script.
 
 </details>
 
@@ -171,26 +592,6 @@ Please refer to the help menu for all available arguments: `python -m python_cor
 <details>
   <summary> Click to expand </summary>
 
-### <a name="swift-requirements"></a> System Requirements
-
-**Building** (recommended):
-
-- Xcode 14.2
-- Command Line Tools for Xcode 14.2
-
-Check [developer.apple.com](https://developer.apple.com/download/all/?q=xcode) for the latest versions.
-
-**Running** (minimum):
-
-| Mac | iPad\* | iPhone\* |
-|:------:|:------:|:------:|
-| macOS 13.1 | iPadOS 16.2 | iOS 16.2 |
-| M1 |  M1  | iPhone 12 Pro |
-
-You will also need the resources generated by the `--bundle-resources-for-swift-cli` option described in [Converting Models to Core ML](#converting-models-to-coreml)
-
-\* Please see [FAQ](#faq) [Q6](#q-mobile-app) regarding deploying on iPad and iPhone.
-
 ### Example CLI Usage
 ```shell
 swift run StableDiffusionSample "a photo of an astronaut riding a horse on mars" --resource-path <output-mlpackages-directory>/Resources/ --seed 93 --output-path </path/to/output/image>
@@ -220,7 +621,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)
@@ -234,6 +635,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)
@@ -257,78 +662,6 @@ Hugging Face has made the app [available in the Mac App Store](https://apps.appl
 
 </details>
 
-## <a name="performance-benchmark"></a> Performance Benchmark
-
-<details>
-  <summary> Click to expand </summary>
-
-Standard [CompVis/stable-diffusion-v1-4](https://huggingface.co/CompVis/stable-diffusion-v1-4) Benchmark
-
-|        Device                      | `--compute-unit`| `--attention-implementation` | Latency (seconds) |
-| ---------------------------------- | --------------  | ---------------------------- | ----------------- |
-| Mac Studio (M1 Ultra, 64-core GPU) | `CPU_AND_GPU`   |     `ORIGINAL`               |      9            |
-| Mac Studio (M1 Ultra, 48-core GPU) | `CPU_AND_GPU`   |     `ORIGINAL`               |      13           |
-| MacBook Pro (M1 Max, 32-core GPU)  | `CPU_AND_GPU`   |     `ORIGINAL`               |      18           |
-| MacBook Pro (M1 Max, 24-core GPU)  | `CPU_AND_GPU`   |     `ORIGINAL`               |      20           |
-| MacBook Pro (M1 Pro, 16-core GPU)  |    `ALL`        |     `SPLIT_EINSUM (default)` |      26           |
-| MacBook Pro (M2)                   | `CPU_AND_NE`    |     `SPLIT_EINSUM (default)` |      23           |
-| MacBook Pro (M1)                   | `CPU_AND_NE`    |     `SPLIT_EINSUM (default)` |      35           |
-| iPad Pro (5th gen, M1)             | `CPU_AND_NE`    |     `SPLIT_EINSUM (default)` |      38           |
-
-
-Please see [Important Notes on Performance Benchmarks](#important-notes-on-performance-benchmarks) section for details.
-
-</details>
-
-## <a name="important-notes-on-performance-benchmarks"></a> Important Notes on Performance Benchmarks
-
-<details>
-  <summary> Click to expand </summary>
-
-- This benchmark was conducted by Apple using public beta versions of iOS 16.2, iPadOS 16.2 and macOS 13.1 in November 2022.
-- The executed program is `python_coreml_stable_diffusion.pipeline` for macOS devices and a minimal Swift test app built on the `StableDiffusion` Swift package for iOS and iPadOS devices.
-- The median value across 3 end-to-end executions is reported.
-- Performance may materially differ 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.
-- The image generation procedure follows the standard configuration: 50 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.
-- Pipelining across the 4 models is not optimized and these performance numbers are subject to variance under increased system load from other applications. Given these factors, we do not report sub-second variance in latency.
-- Weights and activations are in float16 precision for both the GPU and the Neural Engine.
-- The Swift CLI program consumes a peak memory of approximately 2.6GB (without the safety checker), 2.1GB of which is model weights in float16 precision. We applied [8-bit weight quantization](https://coremltools.readme.io/docs/compressing-ml-program-weights#use-affine-quantization) to reduce peak memory consumption by approximately 1GB. However, we observed that it had an adverse effect on generated image quality and we rolled it back. We encourage developers to experiment with other advanced weight compression techniques such as [palettization](https://coremltools.readme.io/docs/compressing-ml-program-weights#use-a-lookup-table) and/or [pruning](https://coremltools.readme.io/docs/compressing-ml-program-weights#use-sparse-representation) which may yield better results.
-- In the [benchmark table](performance-benchmark), 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.
-
-</details>
-
-
-## <a name="results-with-different-compute-units"></a> Results with Different Compute Units
-
-<details>
-  <summary> Click to expand </summary>
-
-It is highly probable that there will be slight differences across generated images using different compute units.
-
-The following images were generated on an M1 MacBook Pro and macOS 13.1 with the prompt *"a photo of an astronaut riding a horse on mars"* using the [runwayml/stable-diffusion-v1-5](https://huggingface.co/runwayml/stable-diffusion-v1-5) model version. The random seed was set to 93:
-
-  CPU_AND_NE  |  CPU_AND_GPU  |  ALL  |
-:------------:|:-------------:|:------:
-![](assets/a_high_quality_photo_of_an_astronaut_riding_a_horse_in_space/randomSeed_93_computeUnit_CPU_AND_NE_modelVersion_runwayml_stable-diffusion-v1-5.png)  |  ![](assets/a_high_quality_photo_of_an_astronaut_riding_a_horse_in_space/randomSeed_93_computeUnit_CPU_AND_GPU_modelVersion_runwayml_stable-diffusion-v1-5.png) | ![](assets/a_high_quality_photo_of_an_astronaut_riding_a_horse_in_space/randomSeed_93_computeUnit_ALL_modelVersion_runwayml_stable-diffusion-v1-5.png) |
-
-Differences may be less or more pronounced for different inputs. Please see the [FAQ](#faq) Q8 for a detailed explanation.
-
-</details>
-
-## <a name="results-with-controlnet"></a> Results with ControlNet
-
-<details>
-  <summary> Click to expand </summary>
-
-[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 CLI arguments in previous sections to exercise this new feature.
-
-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">
-
-</details>
-
 
 ##  <a name="faq"></a> FAQ
 
@@ -355,11 +688,11 @@ Example results using the prompt "a high quality photo of a surfing dog" conditi
 <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:
@@ -379,9 +712,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>
 
@@ -396,6 +730,9 @@ The image generation process in `StableDiffusion` can yield over 2 GB of peak me
 If your app crashes during image generation, consider adding the [Increased Memory Limit](https://developer.apple.com/documentation/bundleresources/entitlements/com_apple_developer_kernel_increased-memory-limit) capability to inform the system that some of your app’s core features may perform better by exceeding the default app memory limit on supported devices.
  
 On iOS, depending on the iPhone model, Stable Diffusion model versions, selected compute units, system load and design of your app, this may still not be sufficient to keep your apps peak memory under the limit. Please remember, because the device shares memory between apps and iOS processes, one app using too much memory can compromise the user experience across the whole device.
+
+We **strongly recommend** 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>
 
 <details>
@@ -412,7 +749,7 @@ On iOS, depending on the iPhone model, Stable Diffusion model versions, selected
 
   <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>
 
@@ -420,7 +757,7 @@ On iOS, depending on the iPhone model, Stable Diffusion model versions, selected
 
   <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](#results-with-different-compute-units).
+  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>
 
@@ -458,3 +795,14 @@ On iOS, depending on the iPhone model, Stable Diffusion model versions, selected
 </details>
 
 </details>
+
+## <a name="bibtex"></a> BibTeX Reference
+
+```latex
+@misc{stable-diffusion-coreml-apple-silicon,
+title = {Stable Diffusion with Core ML on Apple Silicon},
+author = {Atila Orhon and Michael Siracusa and Aseem Wadhwa},
+year = {2022},
+URL = {null}
+}
+```

python_coreml_stable_diffusion/_version.py

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

python_coreml_stable_diffusion/activation_quantization.py

@@ -0,0 +1,503 @@
+#
+# For licensing see accompanying LICENSE.md file.
+# Copyright (C) 2022 Apple Inc. All Rights Reserved.
+#
+
+import logging
+import operator
+
+import torch
+
+logging.basicConfig()
+logger = logging.getLogger()
+logger.setLevel('INFO')
+
+import argparse
+import gc
+import json
+import os
+import pickle
+from copy import deepcopy
+
+import coremltools as ct
+import numpy as np
+from coremltools.optimize.torch.quantization import (
+    LinearQuantizer, LinearQuantizerConfig, ModuleLinearQuantizerConfig)
+from diffusers import StableDiffusionPipeline
+from tqdm import tqdm
+
+from python_coreml_stable_diffusion import attention
+from python_coreml_stable_diffusion import unet
+from python_coreml_stable_diffusion.layer_norm import LayerNormANE
+from python_coreml_stable_diffusion.torch2coreml import compute_psnr
+from python_coreml_stable_diffusion.unet import Einsum
+
+attention.SPLIT_SOFTMAX = True
+
+CALIBRATION_DATA = [
+    "image of a transparent tall glass with ice, fruits and mint, photograph, commercial, food, warm background, beautiful image, detailed",
+    "picture of dimly lit living room, minimalist furniture, vaulted ceiling, huge room, floor to ceiling window with an ocean view, nighttime, 3D render, high quality, detailed",
+    "modern office building, 8 stories tall, glass and steel, 3D render style, wide angle view, very detailed, sharp photographic image, in an office park, bright sunny day, clear blue skies, trees and landscaping",
+    "cute small cat sitting in a movie theater eating popcorn, watching a movie, cozy indoor lighting, detailed, digital painting, character design",
+    "a highly detailed matte painting of a man on a hill watching a rocket launch in the distance by studio ghibli, volumetric lighting, octane render, 4K resolution, hyperrealism, highly detailed, insanely detailed, cinematic lighting, depth of field",
+    "an undersea world with several of fish, rocks, detailed, realistic, photograph, amazing, beautiful, high resolution",
+    "large ocean wave hitting a beach at sunset, photograph, detailed",
+    "pocket watch on a table, close up. macro, sharp, high gloss, brass, gears, sharp, detailed",
+    "pocket watch in the style of pablo picasso, painting",
+    "majestic royal tall ship on a calm sea, realistic painting, cloudy blue sky, in the style of edward hopper",
+    "german castle on a mountain, blue sky, realistic, photograph, dramatic, wide angle view",
+    "artificial intelligence, AI, concept art, blue line sketch",
+    "a humanoid robot, concept art, 3D render, high quality, detailed",
+    "donut with sprinkles and a cup of coffee on a wood table, detailed, photograph",
+    "orchard at sunset, beautiful, photograph, great composition, detailed, realistic, HDR",
+    "image of a map of a country, tattered, old, styled, illustration, for a video game style",
+    "blue and green woven fibers, nano fiber material, detailed, concept art, micro photography",
+]
+
+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",
+]
+
+def get_coreml_inputs(sample_inputs):
+    return [
+        ct.TensorType(
+            name=k,
+            shape=v.shape,
+            dtype=v.numpy().dtype if isinstance(v, torch.Tensor) else v.dtype,
+        ) for k, v in sample_inputs.items()
+    ]
+
+def convert_to_coreml(torchscript_module, sample_inputs):
+    logger.info("Converting model to CoreML..")
+    coreml_model = ct.convert(
+        torchscript_module,
+        convert_to="mlprogram",
+        minimum_deployment_target=ct.target.macOS14,
+        inputs=get_coreml_inputs(sample_inputs),
+        outputs=[ct.TensorType(name="noise_pred", dtype=np.float32)],
+        compute_units=ct.ComputeUnit.ALL,
+        skip_model_load=True,
+    )
+
+    return coreml_model
+
+
+def unet_data_loader(data_dir, device='cpu', calibration_nsamples=None):
+    """
+    Load calibration data from specified path.
+    Limit number of samples to calibration_nsamples, if specified.
+    """
+    dataloader = []
+    skip_load = False
+    for file in sorted(os.listdir(data_dir)):
+        if file.endswith('.pkl'):
+            filepath = os.path.join(data_dir, file)
+            with open(filepath, 'rb') as data:
+                try:
+                    while not skip_load:
+                        unet_data = pickle.load(data)
+                        for input in unet_data:
+                            dataloader.append([x.to(torch.float).to(device) for x in input])
+
+                            if calibration_nsamples:
+                                if len(dataloader) >= calibration_nsamples:
+                                    skip_load = True
+                                    break
+                except EOFError:
+                    pass
+        if skip_load: 
+            break
+
+    logger.info(f"Total calibration samples: {len(dataloader)}")
+    return dataloader
+
+def quantize_module_config(module_name):
+    """
+    Generate quantization config to apply W8A8 quantization for specified module.
+    Rest of the model is kept in FP32 precision.
+    """
+    config = LinearQuantizerConfig(
+        global_config=ModuleLinearQuantizerConfig(
+            milestones=[0, 1000, 1000, 0],
+            weight_dtype=torch.float32,
+            activation_dtype=torch.float32,
+        ),
+        module_name_configs={
+            module_name: ModuleLinearQuantizerConfig(
+                quantization_scheme="symmetric",
+                milestones=[0, 1000, 1000, 0],
+            ),
+        },
+    )
+    return config
+
+def quantize_cumulative_config(skip_conv_layers, skip_einsum_layers):
+    """
+    Generate quantization config to apply W8A8 quantization.
+    Skipped layers are kept in W8A32 precision.
+    """
+    logger.info(f"Skipping {len(skip_conv_layers)} conv layers and {len(skip_einsum_layers)} einsum layers")
+    w8config = ModuleLinearQuantizerConfig(
+            quantization_scheme="symmetric",
+            milestones=[0, 1000, 1000, 0],
+            activation_dtype=torch.float32)
+    conv_modules_config = {name: w8config for name in skip_conv_layers}
+    einsum_modules_config = {name: w8config for name in skip_einsum_layers}
+    module_name_config = {}
+    module_name_config.update(conv_modules_config)
+    module_name_config.update(einsum_modules_config)
+
+    config = LinearQuantizerConfig(
+        global_config=ModuleLinearQuantizerConfig(
+            quantization_scheme="symmetric",
+            milestones=[0, 1000, 1000, 0],
+        ),
+        module_name_configs=module_name_config,
+        module_type_configs={
+            torch.cat: None,
+            torch.nn.GroupNorm: None,
+            torch.nn.SiLU: None,
+            torch.nn.functional.gelu: None,
+            operator.add: None,
+        },
+    )
+    return config
+
+def quantize(model, config, calibration_data):
+    """
+    Apply post training activation quantization to specified model, using calibration data
+    """
+    submodules = dict(model.named_modules(remove_duplicate=True))
+    layer_norm_modules = [key for key, val in submodules.items() if isinstance(val, LayerNormANE)]
+    non_traceable_module_names = layer_norm_modules + [
+        "time_proj",
+        "time_embedding",
+    ]
+
+    # Mark certain modules as non-traceable to make the UNet model fx traceable
+    config.non_traceable_module_names = non_traceable_module_names
+    config.preserved_attributes = ['config', 'device']
+
+    sample_input = calibration_data[0]
+    quantizer = LinearQuantizer(model, config)
+    logger.info("Preparing model for quantization")
+    prepared_model = quantizer.prepare(example_inputs=(sample_input,))
+    prepared_model.eval()
+
+    quantizer.step()
+
+    logger.info("Calibrate")
+    for idx, data in enumerate(calibration_data):
+        logger.info(f"Calibration data sample: {idx}")
+        prepared_model(*data)
+
+    logger.info("Finalize model")
+    quantized_model = quantizer.finalize()
+    return quantized_model
+
+def get_quantizable_modules(unet):
+    quantizable_modules = []
+    for name, module in unet.named_modules():
+        if len(list(module.children())) > 0:
+            continue
+        if type(module) == torch.nn.modules.conv.Conv2d:
+            quantizable_modules.append(('conv', name))
+        if type(module) == Einsum:
+            quantizable_modules.append(('einsum', name))
+
+    return quantizable_modules
+
+def recipe_overrides_for_inference_speedup(conv_layers, skipped_conv):
+    """
+    Quantize the slowest conv layers, even if in skipped set based on PSNR, for good inference speedup
+    """
+    for layer in conv_layers:
+        if "up_blocks" in layer and "resnets" in layer and "conv1" in layer:
+            if layer in skipped_conv:
+                logger.info(f"removing {layer}")
+                skipped_conv.remove(layer)
+        if "upsamplers" in layer:
+            if layer in skipped_conv:
+                logger.info(f"removing {layer}")
+                skipped_conv.remove(layer)
+
+def recipe_overrides_for_quality(conv_layers, skipped_conv):
+    """
+    Do not quantize out projection layers to avoid quantizing outputs of preceding concat layers.
+    Quantizing output of concat layers can lead to quality degradation, due to sharing of scales
+    across concat inputs, which can have varied ranges. Since this is a constraint enforced during
+    model conversion, it may not be captured in layer-wise PSNR analysis of PyTorch model.
+    """
+    out_proj_layers = [layer for layer in conv_layers if "to_out" in layer]
+    for layer in out_proj_layers:
+        if layer not in skipped_conv:
+            logger.info(f"adding {layer}")
+            skipped_conv.add(layer)
+
+def register_input_log_hook(unet, inputs):
+    """
+    Register forward pre hook to save model inputs 
+    """
+    def hook(_, input):
+        input_copy = deepcopy(input)
+        input_copy = tuple(i.to('cpu') for i in input_copy)
+        inputs.append(input_copy)
+
+        # Return inputs unmodified
+        return input
+
+    return unet.register_forward_pre_hook(hook)
+
+def generate_calibration_data(pipe, args, calibration_dir):
+    # Register forward pre hook to record unet inputs
+    unet_inputs = []
+    handle = register_input_log_hook(pipe.unet, unet_inputs)
+
+    # If directory doesn't exist, create it
+    os.makedirs(calibration_dir, exist_ok=True)
+
+    # Run calibration prompts through the pipeline and
+    # serialize recorded UNet model inputs
+    for prompt in CALIBRATION_DATA:
+        gen = torch.manual_seed(args.seed)
+        # run forward pass
+        pipe(prompt=prompt, generator=gen)
+        # save unet inputs
+        filename = "_".join(prompt.split(" ")) + "_" + str(args.seed) + ".pkl"
+        filepath = os.path.join(calibration_dir, filename)
+        with open(filepath, 'wb') as f:
+            pickle.dump(unet_inputs, f)
+        # clear
+        unet_inputs.clear()
+
+    handle.remove()
+
+def register_input_preprocessing_hook(pipe):
+    """
+    Register forward pre hook to convert UNet inputs from HuggingFace StableDiffusionPipeline
+    to match expected model inputs in UNet2DConditionModel defined in unet.py
+    """
+    def hook(_, args, kwargs):
+        sample = args[0]
+        timestep = args[1]
+        if len(timestep.shape) == 0:
+            timestep = timestep[None]
+        timestep = timestep.expand(sample.shape[0])
+        encoder_hidden_states = kwargs["encoder_hidden_states"]
+        encoder_hidden_states = encoder_hidden_states.permute((0, 2, 1)).unsqueeze(2)
+        modified_args = (sample, timestep, encoder_hidden_states)
+        return (modified_args, {})
+
+    return pipe.unet.register_forward_pre_hook(hook, with_kwargs=True)
+
+def prepare_pipe(pipe, unet):
+    """
+    Create a new pipeline from `pipe` with `unet` as the noise predictor
+    """
+    new_pipe = deepcopy(pipe)
+    unet.to(new_pipe.unet.device)
+    new_pipe.unet = unet
+    pre_hook_handle = register_input_preprocessing_hook(new_pipe)
+    return new_pipe, pre_hook_handle
+
+def run_pipe(pipe):
+    gen = torch.manual_seed(args.seed)
+    kwargs = dict(
+        prompt=RANDOM_TEST_DATA,
+        output_type="latent",
+        generator=gen,
+    )
+    return np.array([latent.cpu().numpy() for latent in pipe(**kwargs).images])
+
+
+def get_reference_pipeline(model_version):
+    # Initialize pipe
+    pipe = StableDiffusionPipeline.from_pretrained(
+        model_version,
+        use_safetensors=True,
+        use_auth_token=True,
+    )
+    DEFAULT_NUM_INFERENCE_STEPS = 50
+    pipe.scheduler.set_timesteps(DEFAULT_NUM_INFERENCE_STEPS)
+
+    # Initialize reference unet
+    unet_cls = unet.UNet2DConditionModel
+    reference_unet = unet_cls(**pipe.unet.config).eval()
+    reference_unet.load_state_dict(pipe.unet.state_dict())
+
+    # Initialize reference pipeline
+    ref_pipe, _ = prepare_pipe(pipe, reference_unet)
+
+    del pipe
+    gc.collect()
+    return ref_pipe
+
+def main(args):
+    # Initialize reference pipeline
+    ref_pipe = get_reference_pipeline(args.model_version)
+    if torch.cuda.is_available():
+        device = "cuda"
+    else:
+        device = "cpu"
+    logger.debug(f"Placing pipe in {device}")
+    ref_pipe.to(device)
+    # Generate baseline outputs
+    ref_out = run_pipe(ref_pipe)
+
+    # Setup artifact file paths
+    os.makedirs(args.o, exist_ok=True)
+    recipe_json_path = os.path.join(args.o, f"{args.model_version.replace('/', '_')}_quantization_recipe.json")
+    calibration_dir = os.path.join(args.o, f"calibration_data_{args.model_version.replace('/', '_')}")
+
+    # Generate calibration data 
+    if args.generate_calibration_data:
+        generate_calibration_data(ref_pipe, args, calibration_dir)
+
+    # Compute layer-wise PSNR
+    if args.layerwise_sensitivity:
+        logger.info("Compute Layer-wise PSNR")
+        quantizable_modules = get_quantizable_modules(ref_pipe.unet)
+
+        results = {
+            'conv': {},
+            'einsum': {},
+            'model_version': args.model_version
+        }
+        dataloader = unet_data_loader(calibration_dir, device, args.calibration_nsamples)
+
+        for module_type, module_name in tqdm(quantizable_modules):
+            logger.info(f"Quantizing UNet Layer: {module_name}")
+            config = quantize_module_config(module_name)
+            quantized_unet = quantize(ref_pipe.unet, config, dataloader)
+
+            # Generate outputs from quantized model
+            q_pipe, _ = prepare_pipe(ref_pipe, quantized_unet)
+            test_out = run_pipe(q_pipe)
+
+            psnr = [float(f"{compute_psnr(r, t):.1f}") for r, t in zip(ref_out, test_out)]
+            logger.info(f"PSNR: {psnr}")
+            avg_psnr = sum(psnr) / len(psnr)
+            logger.info(f"AVG PSNR: {avg_psnr}")
+            results[module_type][module_name] = avg_psnr
+
+            del quantized_unet
+            del q_pipe
+            gc.collect()
+
+        with open(recipe_json_path, 'w') as f:
+            json.dump(results, f, indent=2)
+
+    if args.quantize_pytorch:
+        logger.info("Quantizing UNet PyTorch model")
+        dataloader = unet_data_loader(calibration_dir, device, args.calibration_nsamples)
+
+        with open(recipe_json_path, "r") as f:
+            results = json.load(f)
+
+        logger.info(f"Conv PSNR threshold: {args.conv_psnr}, Attn PSNR threshold: {args.attn_psnr}")
+        skipped_conv = set([layer for layer, psnr in results['conv'].items() if psnr < args.conv_psnr])
+        skipped_einsum = set([layer for layer, psnr in results['einsum'].items() if psnr < args.attn_psnr])
+
+        # Apply some overrides on PSNR based recipe for inference and quality improvements
+        # Users can disable these selectively based on specific targets
+        recipe_overrides_for_inference_speedup(results['conv'].keys(), skipped_conv)
+        recipe_overrides_for_quality(results['conv'].keys(), skipped_conv)
+
+        config = quantize_cumulative_config(skipped_conv, skipped_einsum)
+        quantized_unet = quantize(ref_pipe.unet, config, dataloader)
+
+        # Generate outputs from quantized model
+        q_pipe, handle = prepare_pipe(ref_pipe, quantized_unet)
+        test_out = run_pipe(q_pipe)
+
+        psnr = [float(f"{compute_psnr(r, t):.1f}") for r, t in zip(ref_out, test_out)]
+        logger.info(f"PSNR: {psnr}")
+        avg_psnr = sum(psnr) / len(psnr)
+        logger.info(f"AVG PSNR: {avg_psnr}")
+
+        handle.remove()
+        quantized_unet.to('cpu')
+        sample_unet_input = {
+            "sample": dataloader[0][0].to('cpu'),
+            "timestep": dataloader[0][1].to('cpu'),
+            "encoder_hidden_states": dataloader[0][2].to('cpu'),
+        }
+
+        logger.info("JIT tracing quantized model")
+        traced_model = torch.jit.trace(quantized_unet, example_inputs=list(sample_unet_input.values()))
+
+        logger.info("Converting to CoreML")
+        coreml_sample_unet_input = {
+            k: v.numpy().astype(np.float16)
+            for k, v in sample_unet_input.items()
+        }
+        coreml_model = convert_to_coreml(traced_model, coreml_sample_unet_input)
+        coreml_filename = f"Stable_Diffusion_version_{args.model_version.replace('/', '_')}_unet.mlpackage"
+        coreml_model.save(os.path.join(args.o, coreml_filename))
+
+        del q_pipe
+
+    del ref_pipe
+    gc.collect()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "-o",
+        required=True,
+        help="Output directory to save calibration data and quantization artifacts"
+    )
+    parser.add_argument(
+        "--model-version",
+        required=True,
+        choices=("runwayml/stable-diffusion-v1-5", "stabilityai/stable-diffusion-2-1-base"),
+        help=
+        ("The pre-trained model checkpoint and configuration to restore"
+    ))
+    parser.add_argument(
+        "--generate-calibration-data",
+        action="store_true",
+        help="Generate calibration data for UNet model"
+    )
+    parser.add_argument(
+        "--layerwise-sensitivity",
+        action="store_true",
+        help="Compute compression sensitivity per-layer, by quantizing one layer at a time"
+    )
+    parser.add_argument(
+        "--quantize-pytorch",
+        action="store_true",
+        help="Generate activation quantized UNet model by quantizing layers above specified PSNR threshold"
+    )
+    parser.add_argument(
+        "--calibration-nsamples",
+        type=int,
+        help="Number of samples to use for calibrating UNet model"
+    )
+    parser.add_argument("--seed",
+                        "-s",
+                        default=50,
+                        type=int,
+                        help="Random seed to be able to reproduce results"
+    )
+    parser.add_argument("--conv-psnr",
+                        default=40.0,
+                        type=float,
+                        help="PSNR threshold for convolutional layers (default for stabilityai/stable-diffusion-2-1-base)"
+    )
+    parser.add_argument("--attn-psnr",
+                        default=30.0,
+                        type=float,
+                        help="PSNR threshold for attention (Einsum) layers (default for stabilityai/stable-diffusion-2-1-base)"
+    )
+    args = parser.parse_args()
+
+    main(args)

python_coreml_stable_diffusion/attention.py

@@ -0,0 +1,168 @@
+import logging
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+
+import torch
+import math
+
+SPLIT_SOFTMAX = False
+
+def softmax(x, dim):
+    # Reduction max
+    max_x = x.max(dim=dim, keepdim=True).values
+    # EW sub
+    x -= max_x
+    # Scale for EXP to EXP2, Activation EXP2
+    scaled_x = x * (1 / math.log(2))
+    exp_act = torch.exp2(scaled_x)
+    # Reduction Sum + Inv
+    exp_sum_inv = 1 / exp_act.sum(dim=dim, keepdims=True)
+    # EW Mult
+    return exp_act * exp_sum_inv
+
+def split_einsum(q, k, v, mask, heads, dim_head):
+    """ Attention Implementation backing AttentionImplementations.SPLIT_EINSUM
+
+    - Implements https://machinelearning.apple.com/research/neural-engine-transformers
+    - Recommended for ANE
+    - Marginally slower on GPU
+    """
+    mh_q = [
+        q[:, head_idx * dim_head:(head_idx + 1) *
+          dim_head, :, :] for head_idx in range(heads)
+    ]  # (bs, dim_head, 1, max_seq_length) * heads
+
+    k = k.transpose(1, 3)
+    mh_k = [
+        k[:, :, :,
+          head_idx * dim_head:(head_idx + 1) * dim_head]
+        for head_idx in range(heads)
+    ]  # (bs, max_seq_length, 1, dim_head) * heads
+
+    mh_v = [
+        v[:, head_idx * dim_head:(head_idx + 1) *
+          dim_head, :, :] for head_idx in range(heads)
+    ]  # (bs, dim_head, 1, max_seq_length) * heads
+
+    attn_weights = [
+        torch.einsum("bchq,bkhc->bkhq", [qi, ki]) * (dim_head**-0.5)
+        for qi, ki in zip(mh_q, mh_k)
+    ]  # (bs, max_seq_length, 1, max_seq_length) * heads
+
+    if mask is not None:
+        for head_idx in range(heads):
+            attn_weights[head_idx] = attn_weights[head_idx] + mask
+
+    if SPLIT_SOFTMAX:
+        attn_weights = [
+            softmax(aw, dim=1) for aw in attn_weights
+        ]  # (bs, max_seq_length, 1, max_seq_length) * heads
+    else:
+        attn_weights = [
+            aw.softmax(dim=1) for aw in attn_weights
+        ]  # (bs, max_seq_length, 1, max_seq_length) * heads
+
+    attn = [
+        torch.einsum("bkhq,bchk->bchq", wi, vi)
+        for wi, vi in zip(attn_weights, mh_v)
+    ]  # (bs, dim_head, 1, max_seq_length) * heads
+
+    attn = torch.cat(attn, dim=1)  # (bs, dim, 1, max_seq_length)
+    return attn
+
+
+CHUNK_SIZE = 512
+
+def split_einsum_v2(q, k, v, mask, heads, dim_head):
+    """ Attention Implementation backing AttentionImplementations.SPLIT_EINSUM_V2
+
+    - Implements https://machinelearning.apple.com/research/neural-engine-transformers
+    - Recommended for ANE
+    - Marginally slower on GPU
+    - Chunks the query sequence to avoid large intermediate tensors and improves ANE performance
+    """
+    query_seq_length = q.size(3)
+    num_chunks = query_seq_length // CHUNK_SIZE
+    
+    if num_chunks == 0:
+        logger.info(
+            "AttentionImplementations.SPLIT_EINSUM_V2: query sequence too short to chunk "
+            f"({query_seq_length}<{CHUNK_SIZE}), fall back to AttentionImplementations.SPLIT_EINSUM (safe to ignore)")
+        return split_einsum(q, k, v, mask, heads, dim_head)
+    
+    logger.info(
+        "AttentionImplementations.SPLIT_EINSUM_V2: Splitting query sequence length of "
+        f"{query_seq_length} into {num_chunks} chunks")
+
+    mh_q = [
+        q[:, head_idx * dim_head:(head_idx + 1) *
+          dim_head, :, :] for head_idx in range(heads)
+    ]  # (bs, dim_head, 1, max_seq_length) * heads
+
+    # Chunk the query sequence for each head
+    mh_q_chunked = [
+        [h_q[..., chunk_idx * CHUNK_SIZE:(chunk_idx + 1) * CHUNK_SIZE] for chunk_idx in range(num_chunks)]
+        for h_q in mh_q
+    ]  # ((bs, dim_head, 1, QUERY_SEQ_CHUNK_SIZE) * num_chunks) * heads
+
+    k = k.transpose(1, 3)
+    mh_k = [
+        k[:, :, :,
+          head_idx * dim_head:(head_idx + 1) * dim_head]
+        for head_idx in range(heads)
+    ]  # (bs, max_seq_length, 1, dim_head) * heads
+
+    mh_v = [
+        v[:, head_idx * dim_head:(head_idx + 1) *
+          dim_head, :, :] for head_idx in range(heads)
+    ]  # (bs, dim_head, 1, max_seq_length) * heads
+
+    attn_weights = [
+        [
+            torch.einsum("bchq,bkhc->bkhq", [qi_chunk, ki]) * (dim_head**-0.5)
+            for qi_chunk in h_q_chunked
+        ] for h_q_chunked, ki in zip(mh_q_chunked, mh_k)
+    ]  # ((bs, max_seq_length, 1, chunk_size) * num_chunks) * heads
+
+    attn_weights = [
+        [aw_chunk.softmax(dim=1) for aw_chunk in aw_chunked]
+        for aw_chunked in attn_weights
+    ]  # ((bs, max_seq_length, 1, chunk_size) * num_chunks) * heads
+
+    attn = [
+        [
+            torch.einsum("bkhq,bchk->bchq", wi_chunk, vi)
+            for wi_chunk in wi_chunked
+        ] for wi_chunked, vi in zip(attn_weights, mh_v)
+    ]  # ((bs, dim_head, 1, chunk_size) * num_chunks) * heads
+
+    attn = torch.cat([
+        torch.cat(attn_chunked, dim=3) for attn_chunked in attn
+    ], dim=1)  # (bs, dim, 1, max_seq_length)
+
+    return attn
+
+
+def original(q, k, v, mask, heads, dim_head):
+    """ Attention Implementation backing AttentionImplementations.ORIGINAL
+
+    - Not recommended for ANE
+    - Recommended for GPU
+    """
+    bs = q.size(0)
+    mh_q = q.view(bs, heads, dim_head, -1)
+    mh_k = k.view(bs, heads, dim_head, -1)
+    mh_v = v.view(bs, heads, dim_head, -1)
+
+    attn_weights = torch.einsum("bhcq,bhck->bhqk", [mh_q, mh_k])
+    attn_weights.mul_(dim_head**-0.5)
+
+    if mask is not None:
+        attn_weights = attn_weights + mask
+
+    attn_weights = attn_weights.softmax(dim=3)
+
+    attn = torch.einsum("bhqk,bhck->bhcq", [attn_weights, mh_v])
+    attn = attn.contiguous().view(bs, heads * dim_head, 1, -1)
+    return attn

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}")
 
 
@@ -89,6 +101,7 @@ def _get_op_idx_split_location(prog: Program):
     """ Find the op that approximately bisects the graph as measure by weights size on each side
     """
     main_block = prog.functions["main"]
+    main_block.operations = list(main_block.operations)
     total_size_in_mb = 0
 
     for op in main_block.operations:
@@ -105,7 +118,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)
@@ -117,14 +133,16 @@ def _get_first_chunk_outputs(block, op_idx):
     # to the second program (all ops from op_idx+1 till the end). These all vars need to be made the output
     # of the first program and the input of the second program
     boundary_vars = set()
+    block.operations = list(block.operations)
     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)
 
 
@@ -165,6 +183,7 @@ def _make_second_chunk_prog(prog, op_idx):
     boundary_vars = _get_first_chunk_outputs(block, op_idx)
 
     # This op will not be included in this program. Its output var will be made into an input
+    block.operations = list(block.operations)
     boundary_op = block.operations[op_idx]
 
     # Add all boundary ops as inputs
@@ -192,6 +211,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)
@@ -207,7 +231,8 @@ def _make_second_chunk_prog(prog, op_idx):
     return prog
 
 
-def main(args):
+def _legacy_model_chunking(args):
+    # TODO: Remove this method after setting the coremltools dependency >= 8.0
     os.makedirs(args.o, exist_ok=True)
 
     # Check filename extension
@@ -286,6 +311,59 @@ def main(args):
             second_chunk_model=model_chunk2,
         )
 
+    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")
+
+        # 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.")
+
+
+def main(args):
+    ct_version = ct.__version__
+
+    if ct_version != "8.0b2" and ct_version < "8.0":
+        # With coremltools version <= 8.0b1,
+        # we use the legacy implementation.
+        # TODO: Remove the logic after setting the coremltools dependency >= 8.0.
+        logger.info(
+            f"coremltools version {ct_version} detected. Recommended upgrading the package version to "
+            f"'8.0b2' when you running chunk_mlprogram.py script for the latest supports and bug fixes."
+        )
+        _legacy_model_chunking(args)
+    else:
+        # Starting from coremltools==8.0b2, there is this `bisect_model` API that
+        # we can directly call into.
+        from coremltools.models.utils import bisect_model
+        logger.info(f"Start chunking model {args.mlpackage_path} into two pieces.")
+        ct.models.utils.bisect_model(
+            model=args.mlpackage_path,
+            output_dir=args.o,
+            merge_chunks_to_pipeline=args.merge_chunks_in_pipeline_model,
+            check_output_correctness=args.check_output_correctness,
+        )
+        logger.info(f"Model chunking is done.")
+
     # Remove original (non-chunked) model if requested
     if args.remove_original:
         logger.info(
@@ -293,17 +371,6 @@ 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")
-
-    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__":
     parser = argparse.ArgumentParser()
@@ -332,6 +399,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__)
@@ -15,20 +16,75 @@ import numpy as np
 
 import os
 import time
+import subprocess
+import sys
+
+
+def _macos_version():
+    """
+    Returns macOS version as a tuple of integers. On non-Macs, returns an empty tuple.
+    """
+    if sys.platform == "darwin":
+        try:
+            ver_str = subprocess.run(["sw_vers", "-productVersion"], stdout=subprocess.PIPE).stdout.decode('utf-8').strip('\n')
+            return tuple([int(v) for v in ver_str.split(".")])
+        except:
+            raise Exception("Unable to determine the macOS version")
+    return ()
 
 
 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', optimization_hints=None):
 
         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],
+                optimization_hints=optimization_hints,
+            )
+            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,
+                compute_units=ct.ComputeUnit[compute_unit],
+                optimization_hints=optimization_hints,
+            )
+
+            # 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 +94,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 +113,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 +123,85 @@ 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}")
+
+    # On macOS 15+, set fast prediction optimization hint for the unet.
+    optimization_hints = None
+    if submodule_name == "unet" and _macos_version() >= (15, 0):
+        optimization_hints = {"specializationStrategy": ct.SpecializationStrategy.FastPrediction}
+
+    return CoreMLModel(mlpackage_path,
+                       compute_unit,
+                       sources=sources,
+                       optimization_hints=optimization_hints)
 
-    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 +220,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/multilingual_projection.py

@@ -0,0 +1,60 @@
+from python_coreml_stable_diffusion.torch2coreml import _compile_coreml_model
+
+import argparse
+import coremltools as ct
+import numpy as np
+import os
+import torch
+import torch.nn as nn
+
+# TODO: Read these values off of the NLContextualEmbedding API to enforce dimensions and track API versioning
+MAX_SEQUENCE_LENGTH = 256
+EMBED_DIM = 512
+BATCH_SIZE = 1
+
+def main(args):
+    # Layer that was trained to map NLContextualEmbedding to your text_encoder.hidden_size dimensionality
+    text_encoder_projection = torch.jit.load(args.input_path)
+
+    # Prepare random inputs for tracing the network before conversion
+    random_input = torch.randn(BATCH_SIZE, MAX_SEQUENCE_LENGTH, EMBED_DIM)
+
+    # Create a class to bake in the reshape operations required to fit the existing model interface
+    class TextEncoderProjection(nn.Module):
+        def __init__(self, proj):
+            super().__init__()
+            self.proj = proj
+
+        def forward(self, x):
+            return self.proj(x).transpose(1, 2).unsqueeze(2) # BSC, BC1S
+
+    # Trace the torch model
+    text_encoder_projection = torch.jit.trace(TextEncoderProjection(text_encoder_projection), (random_input,))
+
+    # Convert the model to Core ML
+    mlpackage_path = os.path.join(args.output_dir, "MultilingualTextEncoderProjection.mlpackage")
+    ct.convert(
+        text_encoder_projection,
+        inputs=[ct.TensorType('nlcontextualembeddings_output', shape=(1, MAX_SEQUENCE_LENGTH, EMBED_DIM), dtype=np.float32)],
+        outputs=[ct.TensorType('encoder_hidden_states', dtype=np.float32)],
+        minimum_deployment_target=ct.target.macOS14,  # NLContextualEmbedding minimum availability build
+        convert_to='mlprogram',
+    ).save()
+
+    # Compile the model and save it under the specified directory
+    _compile_coreml_model(mlpackage_path, args.output_dir, final_name="MultilingualTextEncoderProjection")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--input-path",
+        help="Path to the torchscript file that contains the projection layer"
+    )
+    parser.add_argument(
+        "--output-dir",
+        help="Output directory in which the Core ML model should be saved",
+    )
+    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,40 @@ 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,
+            unet_batch_one=False,
+            **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 +443,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 +493,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 +504,69 @@ 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()
+
+            if do_classifier_free_guidance:
+                timestep = np.array([t, t], np.float16)
+            else:
+                timestep = np.array([t,], np.float16)
+
             # 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]),
+                    timestep=timestep,
                     encoder_hidden_states=text_embeddings,
                     controlnet_cond=controlnet_cond,
                 )
             else:
-                additional_residuals = {}
+                control_net_additional_residuals = {}
 
             # predict the noise residual
-            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,
-            )["noise_pred"]
+            unet_additional_kwargs.update(control_net_additional_residuals)
+
+            # get prediction from unet
+            if not (unet_batch_one and do_classifier_free_guidance):
+                noise_pred = self.unet(
+                    sample=latent_model_input.astype(np.float16),
+                    timestep=timestep,
+                    encoder_hidden_states=text_embeddings.astype(np.float16),
+                    **unet_additional_kwargs,
+                )["noise_pred"]
+
+                if do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = np.split(noise_pred, 2)
+            else:
+                # query unet sequentially
+                latent_model_input = latent_model_input.astype(np.float16)
+                text_embeddings = text_embeddings.astype(np.float16)
+                timestep = np.array([t,], np.float16)
+
+                noise_pred_uncond = self.unet(
+                    sample=np.expand_dims(latent_model_input[0], axis=0),
+                    timestep=timestep,
+                    encoder_hidden_states=np.expand_dims(text_embeddings[0], axis=0),
+                    **unet_additional_kwargs,
+                )["noise_pred"]
+                noise_pred_text = self.unet(
+                    sample=np.expand_dims(latent_model_input[1], axis=0),
+                    timestep=timestep,
+                    encoder_hidden_states=np.expand_dims(text_embeddings[1], axis=0),
+                    **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 +600,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 +627,28 @@ 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,
+            '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"]
+
+    coreml_pipe_kwargs["force_zeros_for_empty_prompt"] = force_zeros_for_empty_prompt
 
-    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 +661,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 +680,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 +713,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 'xl' in args.model_version and '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 = []
@@ -581,6 +774,7 @@ def main(args):
         guidance_scale=args.guidance_scale,
         controlnet_cond=controlnet_cond,
         negative_prompt=args.negative_prompt,
+        unet_batch_one=args.unet_batch_one,
     )
 
     out_path = get_image_path(args)
@@ -624,7 +818,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 +831,28 @@ 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(
+        "--unet-batch-one",
+        action="store_true",
+        help="Do not batch unet predictions for the prompt and negative prompt.")
+    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/unet.py

@@ -4,6 +4,7 @@
 #
 
 from python_coreml_stable_diffusion.layer_norm import LayerNormANE
+from python_coreml_stable_diffusion import attention
 
 from diffusers.configuration_utils import ConfigMixin, register_to_config
 from diffusers import ModelMixin
@@ -12,7 +13,6 @@ from enum import Enum
 
 import logging
 
-logger = logging.getLogger(__name__)
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
 
@@ -33,6 +33,7 @@ if not _macos_version() >= (13, 1):
 class AttentionImplementations(Enum):
     ORIGINAL = "ORIGINAL"
     SPLIT_EINSUM = "SPLIT_EINSUM"
+    SPLIT_EINSUM_V2 = "SPLIT_EINSUM_V2"
 
 
 ATTENTION_IMPLEMENTATION_IN_EFFECT = AttentionImplementations.SPLIT_EINSUM
@@ -41,6 +42,23 @@ WARN_MSG = \
     "This `nn.Module` is intended for Apple Silicon deployment only. " \
     "PyTorch-specific optimizations and training is disabled"
 
+class Einsum(nn.Module):
+    def __init__(self, heads, dim_head):
+        super().__init__()
+        self.heads = heads
+        self.dim_head = dim_head
+
+    def forward(self, q, k, v, mask):
+        if ATTENTION_IMPLEMENTATION_IN_EFFECT == AttentionImplementations.ORIGINAL:
+            return attention.original(q, k, v, mask, self.heads, self.dim_head)
+
+        elif ATTENTION_IMPLEMENTATION_IN_EFFECT == AttentionImplementations.SPLIT_EINSUM:
+            return attention.split_einsum(q, k, v, mask, self.heads, self.dim_head)
+
+        elif ATTENTION_IMPLEMENTATION_IN_EFFECT == AttentionImplementations.SPLIT_EINSUM_V2:
+            return attention.split_einsum_v2(q, k, v, mask, self.heads, self.dim_head)
+
+
 class CrossAttention(nn.Module):
     """ Apple Silicon friendly version of `diffusers.models.attention.CrossAttention`
     """
@@ -64,10 +82,11 @@ class CrossAttention(nn.Module):
                               bias=False)
         self.to_out = nn.Sequential(
             nn.Conv2d(inner_dim, query_dim, kernel_size=1, bias=True))
+        self.einsum = Einsum(self.heads, self.dim_head)
 
     def forward(self, hidden_states, context=None, mask=None):
-        if self.training:
-            raise NotImplementedError(WARN_MSG)
+        # if self.training:
+        #     raise NotImplementedError(WARN_MSG)
 
         batch_size, dim, _, sequence_length = hidden_states.shape
 
@@ -94,68 +113,7 @@ class CrossAttention(nn.Module):
                     f"Invalid shape for `mask` (Expected {expected_mask_shape}, got {list(mask.size())}"
                 )
 
-        if ATTENTION_IMPLEMENTATION_IN_EFFECT == AttentionImplementations.ORIGINAL:
-            # This version of the attention function is recommended for high GPU core count
-            # devices such as the M1 Max and M1 Ultra
-            bs = q.size(0)
-            mh_q = q.view(bs, self.heads, self.dim_head, -1)
-            mh_k = k.view(bs, self.heads, self.dim_head, -1)
-            mh_v = v.view(bs, self.heads, self.dim_head, -1)
-
-            attn_weights = torch.einsum("bhcq,bhck->bhqk", [mh_q, mh_k])
-            attn_weights.mul_(self.scale)
-
-            if mask is not None:
-                attn_weights = attn_weights + mask
-
-            attn_weights = attn_weights.softmax(dim=3)
-
-            attn = torch.einsum("bhqk,bhck->bhcq", [attn_weights, mh_v])
-            attn = attn.contiguous().view(bs, self.heads * self.dim_head, 1,
-                                          -1)
-
-        elif ATTENTION_IMPLEMENTATION_IN_EFFECT == AttentionImplementations.SPLIT_EINSUM:
-            # The split attention and einsum from https://machinelearning.apple.com/research/neural-engine-transformers
-            # are utilized to build an ANE implementation. This version is marginally slower on the GPU engine and is
-            # not recommended for Max and Ultra Mac variants
-            mh_q = [
-                q[:, head_idx * self.dim_head:(head_idx + 1) *
-                  self.dim_head, :, :] for head_idx in range(self.heads)
-            ]  # (bs, dim_head, 1, max_seq_length) * heads
-
-            k = k.transpose(1, 3)
-            mh_k = [
-                k[:, :, :,
-                  head_idx * self.dim_head:(head_idx + 1) * self.dim_head]
-                for head_idx in range(self.heads)
-            ]  # (bs, max_seq_length, 1, dim_head) * heads
-
-            mh_v = [
-                v[:, head_idx * self.dim_head:(head_idx + 1) *
-                  self.dim_head, :, :] for head_idx in range(self.heads)
-            ]  # (bs, dim_head, 1, max_seq_length) * heads
-
-            attn_weights = [
-                torch.einsum("bchq,bkhc->bkhq", [qi, ki]) * self.scale
-                for qi, ki in zip(mh_q, mh_k)
-            ]  # (bs, max_seq_length, 1, max_seq_length) * heads
-
-            if mask is not None:
-                for head_idx in range(self.heads):
-                    attn_weights[head_idx] = attn_weights[head_idx] + mask
-
-            attn_weights = [
-                aw.softmax(dim=1) for aw in attn_weights
-            ]  # (bs, max_seq_length, 1, max_seq_length) * heads
-            attn = [
-                torch.einsum("bkhq,bchk->bchq", wi, vi)
-                for wi, vi in zip(attn_weights, mh_v)
-            ]  # (bs, dim_head, 1, max_seq_length) * heads
-
-            attn = torch.cat(attn, dim=1)  # (bs, dim, 1, max_seq_length)
-
-        else:
-            raise ValueError(ATTENTION_IMPLEMENTATION_IN_EFFECT)
+        attn = self.einsum(q, k, v, mask)
 
         return self.to_out(attn)
 
@@ -165,15 +123,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
@@ -216,6 +167,7 @@ class CrossAttnUpBlock2D(nn.Module):
         output_scale_factor=1.0,
         downsample_padding=1,
         add_upsample=True,
+        transformer_layers_per_block=1,
     ):
         super().__init__()
         resnets = []
@@ -243,7 +195,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)
@@ -334,6 +286,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",
@@ -369,7 +322,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)
@@ -409,6 +362,7 @@ class DownBlock2D(nn.Module):
         resnet_time_scale_shift="default",
         resnet_act_fn="swish",
         resnet_groups=32,
+        add_downsample=True,
     ):
         super().__init__()
         resnets = []
@@ -427,6 +381,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 = ()
 
@@ -434,6 +393,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
 
 
@@ -651,28 +617,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
 
 
@@ -736,6 +742,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__()
@@ -763,7 +770,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(
@@ -805,6 +812,7 @@ class UNet2DConditionModel(ModelMixin, ConfigMixin):
             "CrossAttnDownBlock2D",
             "DownBlock2D",
         ),
+        mid_block_type="UNetMidBlock2DCrossAttn",
         up_block_types=("UpBlock2D", "CrossAttnUpBlock2D",
                         "CrossAttnUpBlock2D", "CrossAttnUpBlock2D"),
         only_cross_attention=False,
@@ -816,7 +824,12 @@ 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,
+        support_controlnet=False,
         **kwargs,
     ):
         if kwargs.get("dual_cross_attention", None):
@@ -831,6 +844,8 @@ class UNet2DConditionModel(ModelMixin, ConfigMixin):
         super().__init__()
         self._register_load_state_dict_pre_hook(linear_to_conv2d_map)
 
+        self.config.time_cond_proj_dim = None
+        self.support_controlnet = support_controlnet
         self.sample_size = sample_size
         time_embed_dim = block_out_channels[0] * 4
 
@@ -848,6 +863,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
@@ -859,6 +889,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):
@@ -868,6 +901,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,
@@ -877,12 +911,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,
@@ -896,6 +933,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
@@ -909,6 +948,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,
@@ -966,7 +1006,111 @@ class UNet2DConditionModel(ModelMixin, ConfigMixin):
 
             down_block_res_samples += res_samples
 
-        if additional_residuals:
+        if self.support_controlnet:
+            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 self.support_controlnet:
+            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:
+            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 self.support_controlnet:
             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]
@@ -978,7 +1122,7 @@ class UNet2DConditionModel(ModelMixin, ConfigMixin):
                                 emb,
                                 encoder_hidden_states=encoder_hidden_states)
         
-        if additional_residuals:
+        if self.support_controlnet:
             sample = sample + additional_residuals[-1]
 
         # 5. up
@@ -1017,8 +1161,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
@@ -1030,6 +1176,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:
@@ -1038,6 +1185,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,
@@ -1046,6 +1194,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,
         )
 
 
@@ -1060,6 +1209,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(
@@ -1090,6 +1240,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,5 +1,11 @@
-coremltools
-diffusers[torch]
+coremltools>=8.0
+diffusers[torch]==0.30.2
+diffusionkit==0.4.0
 torch
-transformers
-scipy
+transformers==4.44.2
+scipy
+scikit-learn
+pytest
+invisible-watermark
+safetensors
+matplotlib

setup.py

@@ -14,12 +14,19 @@ setup(
     long_description_content_type='text/markdown',
     author='Apple Inc.',
     install_requires=[
-        "coremltools>=6.1",
-        "diffusers[torch]",
+        "coremltools>=8.0",
+        "diffusers[torch]==0.30.2",
         "torch",
-        "transformers",
+        "transformers==4.44.2",
+        "huggingface-hub==0.24.6",
         "scipy",
         "numpy<1.24",
+        "pytest",
+        "scikit-learn",
+        "invisible-watermark",
+        "safetensors",
+        "matplotlib",
+        "diffusionkit==0.4.0",
     ],
     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

@@ -1,5 +1,5 @@
 // For licensing see accompanying LICENSE.md file.
-// Copyright (C) 2022 Apple Inc. All Rights Reserved.
+// Copyright (C) 2024 Apple Inc. All Rights Reserved.
 
 import Foundation
 import CoreML
@@ -28,7 +28,7 @@ public struct Decoder: ResourceManaging {
 
     /// Unload the underlying model to free up memory
     public func unloadResources() {
-       model.unloadResources()
+        model.unloadResources()
     }
 
     /// Batch decode latent samples into images
@@ -39,14 +39,15 @@ public struct Decoder: ResourceManaging {
     ///  - Returns: decoded images
     public func decode(
         _ latents: [MLShapedArray<Float32>],
-        scaleFactor: Float32
+        scaleFactor: Float32,
+        shiftFactor: Float32 = 0.0
     ) throws -> [CGImage] {
 
         // Form batch inputs for model
         let inputs: [MLFeatureProvider] = try latents.map { sample in
             // Reference pipeline scales the latent samples before decoding
             let sampleScaled = MLShapedArray<Float32>(
-                scalars: sample.scalars.map { $0 / scaleFactor },
+                scalars: sample.scalars.map { $0 / scaleFactor + shiftFactor },
                 shape: sample.shape)
 
             let dict = [inputName: MLMultiArray(sampleScaled)]
@@ -64,7 +65,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/DiscreteFlowScheduler.swift

@@ -0,0 +1,123 @@
+// For licensing see accompanying LICENSE.md file.
+// Copyright (C) 2024 Apple Inc. All Rights Reserved.
+
+import CoreML
+
+/// A scheduler used to compute a de-noised image
+@available(iOS 16.2, macOS 13.1, *)
+public final class DiscreteFlowScheduler: Scheduler {
+    public let trainStepCount: Int
+    public let inferenceStepCount: Int
+    public var timeSteps = [Int]()
+    public var betas = [Float]()
+    public var alphas = [Float]()
+    public var alphasCumProd = [Float]()
+
+    public private(set) var modelOutputs: [MLShapedArray<Float32>] = []
+
+    var trainSteps: Float
+    var shift: Float
+    var counter: Int
+    var sigmas = [Float]()
+
+    /// Create a scheduler that uses a second order DPM-Solver++ algorithm.
+    ///
+    /// - Parameters:
+    ///   - stepCount: Number of inference steps to schedule
+    ///   - trainStepCount: Number of training diffusion steps
+    ///   - timeStepShift: Amount to shift the timestep schedule
+    /// - Returns: A scheduler ready for its first step
+    public init(
+        stepCount: Int = 50,
+        trainStepCount: Int = 1000,
+        timeStepShift: Float = 3.0
+    ) {
+        self.trainStepCount = trainStepCount
+        self.inferenceStepCount = stepCount
+        self.trainSteps = Float(trainStepCount)
+        self.shift = timeStepShift
+        self.counter = 0
+
+        let sigmaDistribution = linspace(1, trainSteps, Int(trainSteps)).map { sigmaFromTimestep($0) }
+        let timeStepDistribution = linspace(sigmaDistribution.first!, sigmaDistribution.last!, stepCount).reversed()
+        self.timeSteps = timeStepDistribution.map { Int($0 * trainSteps) }
+        self.sigmas = timeStepDistribution.map { sigmaFromTimestep($0 * trainSteps) }
+    }
+
+    func sigmaFromTimestep(_ timestep: Float) -> Float {
+        if shift == 1.0 {
+            return timestep / trainSteps
+        } else {
+            // shift * timestep / (1 + (shift - 1) * timestep)
+            let t = timestep / trainSteps
+            return shift * t / (1 + (shift - 1) * t)
+        }
+    }
+
+    func timestepsFromSigmas() -> [Float] {
+        return sigmas.map { $0 * trainSteps }
+    }
+
+    /// 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 stepIndex = timeSteps.firstIndex(of: timestep) ?? counter
+        let sigma = sigmas[stepIndex]
+
+        return MLShapedArray<Float>(unsafeUninitializedShape: modelOutput.shape) { result, _ in
+            modelOutput.withUnsafeShapedBufferPointer { noiseScalars, _, _ in
+                sample.withUnsafeShapedBufferPointer { latentScalars, _, _ in
+                    for i in 0..<result.count {
+                        let denoised = latentScalars[i] - noiseScalars[i] * sigma
+                        result.initializeElement(
+                            at: i,
+                            to: denoised
+                        )
+                    }
+                }
+            }
+        }
+    }
+
+    public func calculateTimestepsFromSigmas(strength: Float?) -> [Float] {
+        guard let strength else { return timestepsFromSigmas() }
+        let startStep = max(inferenceStepCount - Int(Float(inferenceStepCount) * strength), 0)
+        let actualTimesteps = Array(timestepsFromSigmas()[startStep...])
+        return actualTimesteps
+    }
+
+    public func step(output: MLShapedArray<Float32>, timeStep t: Int, sample: MLShapedArray<Float32>) -> MLShapedArray<Float32> {
+        let stepIndex = timeSteps.firstIndex(of: t) ?? counter // TODO: allow float timesteps in scheduler step protocol
+        let modelOutput = convertModelOutput(modelOutput: output, timestep: t, sample: sample)
+        modelOutputs.append(modelOutput)
+
+        let sigma = sigmas[stepIndex]
+        var dt = sigma
+        var prevSigma: Float = 0
+        if stepIndex < sigmas.count - 1 {
+            prevSigma = sigmas[stepIndex + 1]
+            dt = prevSigma - sigma
+        }
+
+        let prevSample: MLShapedArray<Float32> = MLShapedArray<Float>(unsafeUninitializedShape: modelOutput.shape) { result, _ in
+            modelOutput.withUnsafeShapedBufferPointer { noiseScalars, _, _ in
+                sample.withUnsafeShapedBufferPointer { latentScalars, _, _ in
+                    for i in 0..<result.count {
+                        let denoised = noiseScalars[i]
+                        let x = latentScalars[i]
+
+                        let d = (x - denoised) / sigma
+                        let prev_x = x + d * dt
+                        result.initializeElement(
+                            at: i,
+                            to: prev_x
+                        )
+                    }
+                }
+            }
+        }
+
+        counter += 1
+        return prevSample
+    }
+}

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/ManagedMLModel.swift

@@ -72,6 +72,56 @@ public final class ManagedMLModel: ResourceManaging {
                                       configuration: configuration)
         }
     }
-
-
+}
+
+@available(iOS 16.2, macOS 13.1, *)
+public extension Array where Element == ManagedMLModel {
+    /// Performs batch predictions using an array of `[ManagedMLModel]` instances in a pipeline.
+    /// - Parameter batch: Inputs for btached predictions.
+    /// - Returns: Final prediction results after processing through all models.
+    /// - Throws: Errors if the array is empty, predictions fail, or results can't be combined.
+    func predictions(from batch: MLBatchProvider) throws -> MLBatchProvider {
+        var results = try self.first!.perform { model in
+            try model.predictions(fromBatch: batch)
+        }
+
+        if self.count == 1 {
+            return results
+        }
+
+        // Manual pipeline batch prediction
+        let inputs = batch.arrayOfFeatureValueDictionaries
+        for stage in self.dropFirst() {
+            // Combine the original inputs with the outputs of the last stage
+            let next = try results.arrayOfFeatureValueDictionaries
+                .enumerated().map { index, dict in
+                    let nextDict = dict.merging(inputs[index]) { out, _ in out }
+                    return try MLDictionaryFeatureProvider(dictionary: nextDict)
+                }
+            let nextBatch = MLArrayBatchProvider(array: next)
+
+            // Predict
+            results = try stage.perform { model in
+                try model.predictions(fromBatch: nextBatch)
+            }
+        }
+
+        return results
+    }
+}
+
+extension MLFeatureProvider {
+    var featureValueDictionary: [String : MLFeatureValue] {
+        self.featureNames.reduce(into: [String : MLFeatureValue]()) { result, name in
+            result[name] = self.featureValue(for: name)
+        }
+    }
+}
+
+extension MLBatchProvider {
+    var arrayOfFeatureValueDictionaries: [[String : MLFeatureValue]] {
+        (0..<self.count).map {
+            self.features(at: $0).featureValueDictionary
+        }
+    }
 }

swift/StableDiffusion/pipeline/MultiModalDiffusionTransformer.swift

@@ -0,0 +1,125 @@
+// For licensing see accompanying LICENSE.md file.
+// Copyright (C) 2022 Apple Inc. All Rights Reserved.
+
+import Foundation
+
+import CoreML
+
+/// MMDiT noise prediction model for stable diffusion
+@available(iOS 16.2, macOS 13.1, *)
+public struct MultiModalDiffusionTransformer: ResourceManaging {
+    /// Model used to predict noise residuals given an input, diffusion time step, and conditional embedding
+    ///
+    /// It can be in the form of a single model or multiple stages
+    var models: [ManagedMLModel]
+
+    /// Creates a MMDiT noise prediction model
+    ///
+    /// - Parameters:
+    ///   - url: Location of single MMDiT compiled Core ML model
+    ///   - configuration: Configuration to be used when the model is loaded
+    /// - Returns: MMDiT model that will lazily load its required resources when needed or requested
+    public init(modelAt url: URL,
+                configuration: MLModelConfiguration)
+    {
+        self.models = [ManagedMLModel(modelAt: url, configuration: configuration)]
+    }
+
+    /// Load resources.
+    public func loadResources() throws {
+        for model in models {
+            try model.loadResources()
+        }
+    }
+
+    /// Unload the underlying model to free up memory
+    public func unloadResources() {
+        for model in models {
+            model.unloadResources()
+        }
+    }
+
+    /// Pre-warm resources
+    public func prewarmResources() throws {
+        // Override default to pre-warm each model
+        for model in models {
+            try model.loadResources()
+            model.unloadResources()
+        }
+    }
+
+    var latentImageEmbeddingsDescription: MLFeatureDescription {
+        try! models.first!.perform { model in
+            model.modelDescription.inputDescriptionsByName["latent_image_embeddings"]!
+        }
+    }
+
+    /// The expected shape of the models latent sample input
+    public var latentImageEmbeddingsShape: [Int] {
+        latentImageEmbeddingsDescription.multiArrayConstraint!.shape.map { $0.intValue }
+    }
+
+    var tokenLevelTextEmbeddingsDescription: MLFeatureDescription {
+        try! models.first!.perform { model in
+            model.modelDescription.inputDescriptionsByName["token_level_text_embeddings"]!
+        }
+    }
+
+    /// The expected shape of the geometry conditioning
+    public var tokenLevelTextEmbeddingsShape: [Int] {
+        tokenLevelTextEmbeddingsDescription.multiArrayConstraint!.shape.map { $0.intValue }
+    }
+
+    /// Batch prediction noise from latent samples
+    ///
+    /// - Parameters:
+    ///   - latents: Batch of latent samples in an array
+    ///   - timeStep: Current diffusion timestep
+    ///   - hiddenStates: Hidden state to condition on
+    /// - Returns: Array of predicted noise residuals
+    func predictNoise(
+        latents: [MLShapedArray<Float32>],
+        timeStep: Float,
+        tokenLevelTextEmbeddings: MLShapedArray<Float32>,
+        pooledTextEmbeddings: MLShapedArray<Float32>
+    ) throws -> [MLShapedArray<Float32>] {
+        // Match time step batch dimension to the model / latent samples
+        let t = MLShapedArray<Float32>(scalars: [timeStep, timeStep], shape: [2])
+
+        // Form batch input to model
+        let inputs = try latents.enumerated().map {
+            let dict: [String: Any] = [
+                "latent_image_embeddings": MLMultiArray($0.element),
+                "timestep": MLMultiArray(t),
+                "token_level_text_embeddings": MLMultiArray(tokenLevelTextEmbeddings),
+                "pooled_text_embeddings": MLMultiArray(pooledTextEmbeddings),
+            ]
+            return try MLDictionaryFeatureProvider(dictionary: dict)
+        }
+        let batch = MLArrayBatchProvider(array: inputs)
+
+        // Make predictions
+        let results = try models.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
+    }
+}

swift/StableDiffusion/pipeline/MultilingualTextEncoder.swift

@@ -0,0 +1,194 @@
+// For licensing see accompanying LICENSE.md file.
+// Copyright (C) 2023 Apple Inc. All Rights Reserved.
+
+import Foundation
+import NaturalLanguage
+import CoreML
+
+#if canImport(NaturalLanguage.NLContextualEmbedding)
+@available(iOS 17.0, macOS 14.0, *)
+public struct MultilingualTextEncoder: TextEncoderModel {
+    let adapter: ManagedMLModel?
+
+    let embeddingModel: NLContextualEmbedding
+
+    // TODO: use maximum sequence length from embedding.
+    let maximumEmbeddingSequenceLength = 256
+
+    /// Creates a multilingual text encoder.
+    ///
+    /// - Parameters:
+    ///   - url: The location of the compiled Core ML adapter model. The model is a linear projection layer that
+    ///   transforms the contextual embedding size of 512 to the default text encoder CLIP size of 768.
+    ///   - configuration: The configuration to be used when the model is loaded.
+    ///   - script: The scipt of the contextual embedding.
+    public init(
+        modelAt url: URL? = nil,
+        configuration: MLModelConfiguration = .init(),
+        script: Script = .latin
+    ) {
+        if let url {
+            self.adapter = ManagedMLModel(modelAt: url, configuration: configuration)
+        } else {
+            self.adapter = nil
+        }
+        self.embeddingModel = NLContextualEmbedding(script: script.asNLScript)!
+        self.embeddingModel.requestAssets { _, _ in }
+    }
+
+    /// Loads model resources into memory.
+    public func loadResources() throws {
+        try adapter?.loadResources()
+        try embeddingModel.load()
+    }
+
+    /// Unloads the model resources to free up memory.
+    public func unloadResources() {
+        adapter?.unloadResources()
+        embeddingModel.unload()
+    }
+
+    /// Encodes the input text.
+    ///
+    ///  - Parameter text: The input text.
+    ///  - Returns: An embedding shaped array.
+    public func encode(_ text: String) throws -> MLShapedArray<Float> {
+        guard embeddingModel.hasAvailableAssets else {
+            throw Error.missingEmbeddingResource
+        }
+
+        // Create the text embedding result.
+        let embedding = try embeddingModel.embeddingResult(for: text, language: nil)
+
+        // Create embedding array from token vectors.
+        var shapedEmbeddings = MLShapedArray<Double>(
+            repeating: 0.0,
+            shape: [1, maximumEmbeddingSequenceLength, embeddingModel.dimension]
+        )
+        shapedEmbeddings.withUnsafeMutableShapedBufferPointer { pointer, _, _ in
+            var tokenIndex = 0
+            embedding.enumerateTokenVectors(in: text.startIndex ..< text.endIndex) { (tokenEmbeddings, _) -> Bool in
+                for tokenEmbeddingIndex in 0 ..< tokenEmbeddings.count {
+                    pointer[tokenIndex * embeddingModel.dimension + tokenEmbeddingIndex] = tokenEmbeddings[tokenEmbeddingIndex]
+                }
+                tokenIndex += 1
+                return true
+            }
+        }
+
+        if adapter == nil {
+            // Return embeddings with shape [1, 256, 512].
+            return MLShapedArray(converting: shapedEmbeddings)
+        } else {
+            // Project the embeddings to the correct CLIP model input shape of [1, 768, 1, 256].
+            return try projectEmbeddings(shapedEmbeddings)
+        }
+    }
+
+    /// Creates the adapter model input feature provider.
+    private func prepareProjectionInput(_ input: MLShapedArray<Double>) throws -> MLDictionaryFeatureProvider {
+        guard let adapter else {
+            fatalError("Cannot prepare projection input without an adapter.")
+        }
+        return try adapter.perform { model in
+            guard let inputDescription = model.modelDescription.inputDescriptionsByName.first?.value else {
+                throw Error.missingAdapterInput
+            }
+            return try MLDictionaryFeatureProvider(dictionary: [inputDescription.name: MLMultiArray(input)])
+        }
+    }
+
+    /// Processes the adapter model output feature provider.
+    private func processProjectionOutput(_ output: MLFeatureProvider) throws -> MLShapedArray<Float> {
+        guard let adapter else {
+            fatalError("Cannot process projection output without an adapter.")
+        }
+        return try adapter.perform { model in
+            guard let outputDescription = model.modelDescription.outputDescriptionsByName.first?.value else {
+                throw Error.missingAdapterOutput
+            }
+            guard let result = output
+                .featureValue(for: outputDescription.name)?
+                .multiArrayValue else {
+
+                throw Error.incompatibleAdapterOutputDataFormat(
+                    expected: .multiArray,
+                    actual: outputDescription.type
+                )
+            }
+
+            return MLShapedArray(converting: result)
+        }
+    }
+
+    /// Projects the embeddings.
+    private func projectEmbeddings(_ embeddings: MLShapedArray<Double>) throws -> MLShapedArray<Float> {
+        guard let adapter else {
+            fatalError("Cannot project embeddings without an adapter.")
+        }
+        let inputFeatureProvider = try prepareProjectionInput(embeddings)
+        let projection = try adapter.perform { model in
+            return try model.prediction(from: inputFeatureProvider)
+        }
+        return try processProjectionOutput(projection)
+    }
+}
+
+@available(iOS 17.0, macOS 14.0, *)
+extension MultilingualTextEncoder {
+    /// A multilingual text encoder error.
+    public enum Error: Swift.Error, LocalizedError, Equatable, CustomDebugStringConvertible {
+        /// An error that indicates that the resource for the embedding is missing.
+        case missingEmbeddingResource
+
+        /// An error that indicates that the adapter model input data has the wrong format.
+        case incompatibleAdapterInputDataFormat(expected: MLFeatureType, actual: MLFeatureType)
+
+        /// An error that indicates that the adapter model output data has the wrong format.
+        case incompatibleAdapterOutputDataFormat(expected: MLFeatureType, actual: MLFeatureType)
+
+        /// An error that indicates that the adapter model is missing an input.
+        case missingAdapterInput
+
+        /// An error that indicates that the adapter model is missing an output.
+        case missingAdapterOutput
+
+        /// A debug description of the error.
+        public var errorDescription: String? {
+            debugDescription
+        }
+
+        /// A text representation of the error.
+        public var debugDescription: String {
+            switch self {
+            case .missingEmbeddingResource:
+                return "Resources required for generating embeddings are missing. Make sure that your device is connected to the internet and try again."
+            case .incompatibleAdapterInputDataFormat(expected: let expected, actual: let actual):
+                return "The adapter model input expected to be \(expected) but is \(actual)."
+            case .incompatibleAdapterOutputDataFormat(expected: let expected, actual: let actual):
+                return "The adapter model output expected to be \(expected) but is \(actual)."
+            case .missingAdapterInput:
+                return "The adapter model is missing an input."
+            case .missingAdapterOutput:
+                return "The adapter model is missing an output."
+            }
+        }
+    }
+}
+#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 {
+        case .latin: return .latin
+        case .cyrillic: return .cyrillic
+        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/StableDiffusion3Pipeline+Resources.swift

@@ -0,0 +1,97 @@
+// For licensing see accompanying LICENSE.md file.
+// Copyright (C) 2024 Apple Inc. All Rights Reserved.
+
+import CoreML
+import Foundation
+import Tokenizers
+import Hub
+
+@available(iOS 17.0, macOS 14.0, *)
+public extension StableDiffusion3Pipeline {
+    struct ResourceURLs {
+        public let textEncoderURL: URL
+        public let textEncoder2URL: URL
+        public let textEncoderT5URL: URL
+        public let mmditURL: URL
+        public let decoderURL: URL
+        public let encoderURL: URL
+        public let vocabURL: URL
+        public let mergesURL: URL
+        public let configT5URL: URL
+        public let dataT5URL: URL
+
+        public init(resourcesAt baseURL: URL) {
+            textEncoderURL = baseURL.appending(path: "TextEncoder.mlmodelc")
+            textEncoder2URL = baseURL.appending(path: "TextEncoder2.mlmodelc")
+            textEncoderT5URL = baseURL.appending(path: "TextEncoderT5.mlmodelc")
+            mmditURL = baseURL.appending(path: "MultiModalDiffusionTransformer.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")
+            configT5URL = baseURL.appending(path: "tokenizer_config.json")
+            dataT5URL = baseURL.appending(path: "tokenizer.json")
+        }
+    }
+
+    /// 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(tokenizer: tokenizer, modelAt: urls.textEncoderURL, configuration: config)
+
+        // 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)
+
+        // Optional T5 encoder
+        var textEncoderT5: TextEncoderT5?
+        if FileManager.default.fileExists(atPath: urls.configT5URL.path),
+           FileManager.default.fileExists(atPath: urls.dataT5URL.path),
+           FileManager.default.fileExists(atPath: urls.textEncoderT5URL.path)
+        {
+            let tokenizerT5 = try PreTrainedTokenizer(tokenizerConfig: Config(fileURL: urls.configT5URL), tokenizerData: Config(fileURL: urls.dataT5URL))
+            textEncoderT5 = TextEncoderT5(tokenizer: tokenizerT5, modelAt: urls.textEncoderT5URL, configuration: config)
+        } else {
+            textEncoderT5 = nil
+        }
+
+        // Denoiser model
+        let mmdit = MultiModalDiffusionTransformer(modelAt: urls.mmditURL, configuration: config)
+
+        // Image Decoder
+        let decoder = Decoder(modelAt: urls.decoderURL, configuration: config)
+
+        // Optional Image Encoder
+        let encoder: Encoder?
+        if FileManager.default.fileExists(atPath: urls.encoderURL.path) {
+            encoder = Encoder(modelAt: urls.encoderURL, configuration: config)
+        } else {
+            encoder = nil
+        }
+
+        // Construct pipeline
+        self.init(
+            textEncoder: textEncoder,
+            textEncoder2: textEncoder2,
+            textEncoderT5: textEncoderT5,
+            mmdit: mmdit,
+            decoder: decoder,
+            encoder: encoder,
+            reduceMemory: reduceMemory
+        )
+    }
+}

swift/StableDiffusion/pipeline/StableDiffusion3Pipeline.swift

@@ -0,0 +1,486 @@
+// For licensing see accompanying LICENSE.md file.
+// Copyright (C) 2024 Apple Inc. All Rights Reserved.
+
+import Accelerate
+import CoreGraphics
+import CoreImage
+import CoreML
+import Foundation
+
+@available(iOS 17.0, macOS 14.0, *)
+public struct StableDiffusion3Pipeline: StableDiffusionPipelineProtocol {
+    public typealias Configuration = PipelineConfiguration
+    public typealias Progress = PipelineProgress
+
+    /// Model to generate embeddings for tokenized input text
+    var textEncoder: TextEncoderXLModel
+    var textEncoder2: TextEncoderXLModel
+    var textEncoderT5: TextEncoderT5Model?
+
+    /// Model used to predict noise residuals given an input, diffusion time step, and conditional embedding
+    var mmdit: MultiModalDiffusionTransformer
+
+    /// 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
+    ///   - mmdit: 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,
+        textEncoderT5: TextEncoderT5?,
+        mmdit: MultiModalDiffusionTransformer,
+        decoder: Decoder,
+        encoder: Encoder?,
+        reduceMemory: Bool = false
+    ) {
+        self.textEncoder = textEncoder
+        self.textEncoder2 = textEncoder2
+        self.textEncoderT5 = textEncoderT5
+        self.mmdit = mmdit
+        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 textEncoder.loadResources()
+            try textEncoder2.loadResources()
+            try textEncoderT5?.loadResources()
+            try mmdit.loadResources()
+            try decoder.loadResources()
+
+            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()
+        textEncoderT5?.unloadResources()
+        mmdit.unloadResources()
+        decoder.unloadResources()
+        encoder?.unloadResources()
+    }
+
+    /// Prewarm resources one at a time
+    public func prewarmResources() throws {
+        try textEncoder.prewarmResources()
+        try textEncoder2.prewarmResources()
+        try textEncoderT5?.prewarmResources()
+        try mmdit.prewarmResources()
+        try decoder.prewarmResources()
+
+        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?] {
+        // Setup geometry conditioning for base/refiner inputs
+        let sd3Input: ModelInputs = try generateConditioning(using: config)
+
+        if reduceMemory {
+            textEncoder.unloadResources()
+            textEncoder2.unloadResources()
+            textEncoderT5?.unloadResources()
+        }
+
+        // Setup schedulers
+        let scheduler: [DiscreteFlowScheduler] = (0..<config.imageCount).map { _ in
+            DiscreteFlowScheduler(stepCount: config.stepCount, timeStepShift: config.schedulerTimestepShift)
+        }
+
+        // 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
+        let mmditModel = mmdit
+
+        let mmditHiddenStates = sd3Input.hiddenStates
+        let mmditPooledStates = sd3Input.pooledStates
+
+        let timeSteps: [Float] = scheduler[0].calculateTimestepsFromSigmas(strength: timestepStrength)
+
+        // De-noising loop
+        for (step, t) in timeSteps.enumerated() {
+            // Expand the latents for classifier-free guidance
+            // and input to the MMDiT noise prediction model
+            let latentUnetInput = latents.map {
+                MLShapedArray<Float32>(concatenating: [$0, $0], alongAxis: 0)
+            }
+
+            // Predict noise residuals from latent samples
+            // and current time step conditioned on hidden states
+            var noise = try mmditModel.predictNoise(
+                latents: latentUnetInput,
+                timeStep: t,
+                tokenLevelTextEmbeddings: mmditHiddenStates,
+                pooledTextEmbeddings: mmditPooledStates
+            )
+
+            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: scheduler[i].timeSteps[step], // TODO: allow float timesteps in scheduler step protocol
+                    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 {
+            mmdit.unloadResources()
+        }
+
+        // Decode the latent samples to images
+        return try decodeToImages(denoisedLatents, configuration: config)
+    }
+
+    func encodePrompt(_ prompt: String) throws -> (MLShapedArray<Float32>, MLShapedArray<Float32>) {
+        var embeds = MLShapedArray<Float32>()
+        var pooled = MLShapedArray<Float32>()
+
+        let (embeds1, pooledValue1) = try textEncoder.encode(prompt)
+        let (embeds2, pooledValue2) = try textEncoder2.encode(prompt)
+        var embedsT5 = try textEncoderT5?.encode(prompt).encoderHiddenStates ?? MLShapedArray<Float32>(repeating: 0, shape: [1, 4096, 1, 77])
+
+        // Truncate T5
+        embedsT5 = truncatedT5Embeds(embedsT5)
+
+        let padding1 = MLShapedArray<Float32>(repeating: 0, shape: [1, 77, 2048])
+
+        // Base needs concatenated embeddings
+        // [1, 77, 768], [1, 77, 1280], [1, 77, 2048] -> [1, 77, 4096]
+        embeds = MLShapedArray<Float32>(
+            concatenating: [embeds1, embeds2, padding1],
+            alongAxis: 2
+        )
+
+        // [1, 77, 4096] -> [1, 4096, 1 77]
+        embeds = toHiddenStates(embeds)
+
+        // [1, 4096, 1 77], [1, 4096, 1, 77] -> [1, 4096, 1, 154]
+        embeds = MLShapedArray<Float32>(
+            concatenating: [embeds, embedsT5],
+            alongAxis: 3
+        )
+
+        // [1, 768], [1, 1280] -> [1, 2048]
+        pooled = MLShapedArray<Float32>(
+            concatenating: [pooledValue1, pooledValue2],
+            alongAxis: 1
+        )
+
+        return (embeds, pooled)
+    }
+
+    func generateConditioning(using config: Configuration) throws -> ModelInputs {
+        // Encode the input prompt and negative prompt
+        let (promptEmbedding, pooled) = try encodePrompt(config.prompt)
+        let (negativePromptEmbedding, negativePooled) = try encodePrompt(config.negativePrompt)
+
+        // Convert to Unet hidden state representation
+        // Concatenate the prompt and negative prompt embeddings
+        let hiddenStates = MLShapedArray(concatenating: [promptEmbedding, negativePromptEmbedding], alongAxis: 0)
+        let pooledScalars = MLShapedArray(concatenating: [pooled, negativePooled], alongAxis: 0)
+
+        let pooledStates = MLShapedArray<Float32>(
+            scalars: pooledScalars.scalars,
+            shape: [2, 2048, 1, 1]
+        )
+
+        return ModelInputs(hiddenStates: hiddenStates, pooledStates: pooledStates)
+    }
+
+    func generateLatentSamples(configuration config: Configuration, scheduler: Scheduler) throws -> [MLShapedArray<Float32>] {
+        var sampleShape = mmdit.latentImageEmbeddingsShape
+        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
+    }
+
+    func performGuidance(_ noise: [MLShapedArray<Float32>], _ guidanceScale: Float) -> [MLShapedArray<Float32>] {
+        noise.map { performGuidance($0, guidanceScale) }
+    }
+
+    func performGuidance(_ noise: MLShapedArray<Float32>, _ guidanceScale: Float) -> MLShapedArray<Float32> {
+        var shape = noise.shape
+        shape[0] = 1
+        return MLShapedArray<Float>(unsafeUninitializedShape: shape) { result, _ in
+            noise.withUnsafeShapedBufferPointer { scalars, _, strides in
+                for i in 0..<result.count {
+                    // unconditioned + guidance*(text - unconditioned)
+                    let text = scalars[i]
+                    let negText = scalars[strides[0] + i]
+                    let guidance = negText + guidanceScale * (text - negText)
+                    result.initializeElement(
+                        at: i,
+                        to: guidance
+                    )
+                }
+            }
+        }
+    }
+
+    public func decodeToImages(_ latents: [MLShapedArray<Float32>], configuration config: Configuration) throws -> [CGImage?] {
+        defer {
+            if reduceMemory {
+                decoder.unloadResources()
+            }
+        }
+
+        return try decoder.decode(latents, scaleFactor: config.decoderScaleFactor, shiftFactor: config.decoderShiftFactor)
+
+        // TODO: use latent rgb factors with blur for preview images
+        // This will require a method to decode with either the vae or the rgb factors depending on config
+        // return try decodePreviewImage(latents, scaleFactor: config.decoderScaleFactor)
+    }
+
+    /// Shape 16 x 3
+    let rgbFactors: [[Float]] = [
+        [-0.0645,  0.0177,  0.1052], [ 0.0028,  0.0312,  0.0650],
+        [ 0.1848,  0.0762,  0.0360], [ 0.0944,  0.0360,  0.0889],
+        [ 0.0897,  0.0506, -0.0364], [-0.0020,  0.1203,  0.0284],
+        [ 0.0855,  0.0118,  0.0283], [-0.0539,  0.0658,  0.1047],
+        [-0.0057,  0.0116,  0.0700], [-0.0412,  0.0281, -0.0039],
+        [ 0.1106,  0.1171,  0.1220], [-0.0248,  0.0682, -0.0481],
+        [ 0.0815,  0.0846,  0.1207], [-0.0120, -0.0055, -0.0867],
+        [-0.0749, -0.0634, -0.0456], [-0.1418, -0.1457, -0.1259]
+    ]
+
+    public func decodePreviewImage(
+        _ latents: [MLShapedArray<Float32>],
+        scaleFactor: Float32
+    ) throws -> [CGImage] {
+        let height = 64
+        let width = 64
+        let channels = 16
+        let outputChannels = 3
+
+        // Ensure there is a first element in latents and extract its scalars
+        guard let latentScalars = latents.first?.scalars else {
+            throw NSError(domain: "DecodeError", code: 0, userInfo: [NSLocalizedDescriptionKey: "Invalid latent array"])
+        }
+
+        // The latentScalars is a flat array, we need to reshape and multiply
+        var reshapedLatent = [Float32](repeating: 0, count: height * width * channels)
+
+        // We reorder the indices manually to switch from [channels, height, width] to [height, width, channels]
+        for h in 0..<height {
+            for w in 0..<width {
+                for c in 0..<channels {
+                    let oldIndex = c * height * width + h * width + w
+                    let newIndex = h * width * channels + w * channels + c
+                    reshapedLatent[newIndex] = latentScalars[oldIndex] // 1.5305 + 0.0609
+                }
+            }
+        }
+
+        // Prepare to hold the result of the multiplication
+        var imageArray = [Float32](repeating: 0, count: height * width * outputChannels)
+
+        // Perform matrix multiplication using Accelerate
+        vDSP_mmul(reshapedLatent, 1,
+                  rgbFactors.flatMap { $0 }, 1,
+                  &imageArray, 1,
+                  vDSP_Length(height * width),  // number of rows in output
+                  vDSP_Length(outputChannels),  // number of columns in output
+                  vDSP_Length(channels))        // common dimension
+
+        // Convert imageArray into a CGImage
+        let latentImage = imageArray.toCGImage(width: width, height: height)
+
+        // Apply a Gaussian blur to the preview image to reduce pixeled look
+        let ciImage = CIImage(cgImage: latentImage!)
+        let blurFilter = CIFilter(name: "CIGaussianBlur")!
+        blurFilter.setValue(ciImage, forKey: kCIInputImageKey)
+        blurFilter.setValue(4.0, forKey: kCIInputRadiusKey)
+
+        let context = CIContext()
+        guard let outputImage = blurFilter.outputImage,
+              let cgBlurredPreview = context.createCGImage(outputImage, from: ciImage.extent)
+        else {
+            throw PipelineError.errorCreatingPreview
+        }
+
+        return [cgBlurredPreview]
+    }
+
+    struct ModelInputs {
+        var hiddenStates: MLShapedArray<Float32>
+        var pooledStates: MLShapedArray<Float32>
+    }
+
+    /// Helper function to truncate the T5 embeddings
+    func truncatedT5Embeds(_ embedding: MLShapedArray<Float32>) -> MLShapedArray<Float32> {
+        // Unoptimized manual truncation
+        // e.g. From [1, 4096, 1, 128] to [1, 4096, 1, 77]
+        let fromShape = embedding.shape
+        let stateShape = [fromShape[0], fromShape[1], fromShape[2], 77]
+        var states = MLShapedArray<Float32>(repeating: 0.0, shape: stateShape)
+        for i0 in 0..<fromShape[0] {
+            for i1 in 0..<fromShape[1] {
+                for i2 in 0..<fromShape[2] {
+                    for i3 in 0..<stateShape[3] {
+                        states[scalarAt: i0, i1, i2, i3] = embedding[scalarAt: i0, i1, i2, i3]
+                    }
+                }
+            }
+        }
+        return states
+    }
+}
+
+extension Array where Element == Float32 {
+    func toCGImage(width: Int, height: Int) -> CGImage? {
+        // Define color space and bitmap info
+        let colorSpace = CGColorSpaceCreateDeviceRGB()
+        let bitmapInfo = CGBitmapInfo.byteOrder32Big.rawValue | CGImageAlphaInfo.premultipliedLast.rawValue
+
+        // Calculate bytes per pixel and bytes per row
+        let bytesPerPixel = 4
+        let bytesPerRow = width * bytesPerPixel
+
+        // Allocate memory for the pixel data
+        var data = [UInt8](repeating: 0, count: height * bytesPerRow)
+
+        // Fill the data array with pixel data
+        for h in 0..<height {
+            for w in 0..<width {
+                let pixelIndex = h * width + w
+                let dataIndex = h * bytesPerRow + w * bytesPerPixel
+                let pixelBase = pixelIndex * 3 // Base index for R, G, B values in the source array
+
+                // Ensure your source array has enough data
+                if (pixelBase + 3) < self.count {
+                    let redValue = (self[pixelBase] + 1) / 2 * 255
+                    let bluValue = (self[pixelBase + 1] + 1) / 2 * 255
+                    let grnValue = (self[pixelBase + 2] + 1) / 2 * 255
+                    data[dataIndex] = UInt8(clamp(value: redValue, lower: 0, upper: 255))     // Red
+                    data[dataIndex + 1] = UInt8(clamp(value: bluValue, lower: 0, upper: 255)) // Green
+                    data[dataIndex + 2] = UInt8(clamp(value: grnValue, lower: 0, upper: 255)) // Blue
+                    data[dataIndex + 3] = 255 // Alpha
+                }
+            }
+        }
+
+        // Create the context
+        guard let context = CGContext(data: &data, width: width, height: height, bitsPerComponent: 8, bytesPerRow: bytesPerRow, space: colorSpace, bitmapInfo: bitmapInfo) else {
+            print("Failed to create CGContext.")
+            return nil
+        }
+
+        // Create a CGImage from context
+        guard let smallImage = context.makeImage() else {
+            return nil
+        }
+
+        // Define the upscaled dimensions
+        let scaledWidth = width * 8
+        let scaledHeight = height * 8
+
+        // Create a new context with scaled dimensions
+        guard let largeContext = CGContext(data: nil, width: scaledWidth, height: scaledHeight, bitsPerComponent: 8, bytesPerRow: scaledWidth * 4, space: colorSpace, bitmapInfo: bitmapInfo) else {
+            return nil
+        }
+
+        // Draw the small image into the large context
+        largeContext.interpolationQuality = .high
+        largeContext.draw(smallImage, in: CGRect(x: 0, y: 0, width: scaledWidth, height: scaledHeight))
+
+        // Convert the upscaled context to a CGImage
+        return largeContext.makeImage()
+    }
+
+    /// Helper function to clamp the values within the specified range
+    private func clamp(value: Float32, lower: UInt8, upper: UInt8) -> UInt8 {
+        return UInt8(Swift.max(Float32(lower), Swift.min(value, Float32(upper))))
+    }
+}

swift/StableDiffusion/pipeline/StableDiffusionPipeline+Resources.swift

@@ -3,6 +3,7 @@
 
 import Foundation
 import CoreML
+import NaturalLanguage
 
 @available(iOS 16.2, macOS 13.1, *)
 public extension StableDiffusionPipeline {
@@ -22,6 +23,7 @@ public extension StableDiffusionPipeline {
         public let controlledUnetURL: URL
         public let controlledUnetChunk1URL: URL
         public let controlledUnetChunk2URL: URL
+        public let multilingualTextEncoderProjectionURL: URL
 
         public init(resourcesAt baseURL: URL) {
             textEncoderURL = baseURL.appending(path: "TextEncoder.mlmodelc")
@@ -37,6 +39,7 @@ public extension StableDiffusionPipeline {
             controlledUnetURL = baseURL.appending(path: "ControlledUnet.mlmodelc")
             controlledUnetChunk1URL = baseURL.appending(path: "ControlledUnetChunk1.mlmodelc")
             controlledUnetChunk2URL = baseURL.appending(path: "ControlledUnetChunk2.mlmodelc")
+            multilingualTextEncoderProjectionURL = baseURL.appending(path: "MultilingualTextEncoderProjection.mlmodelc")
         }
     }
 
@@ -44,29 +47,46 @@ public extension StableDiffusionPipeline {
     /// specified URL
     ///
     /// - Parameters:
-    ///    - baseURL: URL pointing to directory holding all model
-    ///               and tokenization resources
+    ///   - baseURL: URL pointing to directory holding all model and tokenization resources
     ///   - controlNetModelNames: Specify ControlNet models to use in generation
     ///   - configuration: The configuration to load model resources with
     ///   - disableSafety: Load time disable of safety to save memory
     ///   - reduceMemory: Setup pipeline in reduced memory mode
+    ///   - useMultilingualTextEncoder: Option to use system multilingual NLContextualEmbedding as encoder
+    ///   - script: Optional natural language script to use for the text encoder.
     /// - Returns:
     ///  Pipeline ready for image generation if all  necessary resources loaded
-    init(resourcesAt baseURL: URL,
-         controlNet controlNetModelNames: [String],
-         configuration config: MLModelConfiguration = .init(),
-         disableSafety: Bool = false,
-         reduceMemory: Bool = false) throws {
+    init(
+        resourcesAt baseURL: URL,
+        controlNet controlNetModelNames: [String],
+        configuration config: MLModelConfiguration = .init(),
+        disableSafety: Bool = false,
+        reduceMemory: Bool = false,
+        useMultilingualTextEncoder: Bool = false,
+        script: Script? = nil
+    ) throws {
 
         /// Expect URL of each resource
         let urls = ResourceURLs(resourcesAt: baseURL)
+        let textEncoder: TextEncoderModel
 
-        // Text tokenizer and encoder
+#if canImport(NaturalLanguage.NLScript)
+        if useMultilingualTextEncoder {
+            guard #available(macOS 14.0, iOS 17.0, *) else { throw PipelineError.unsupportedOSVersion }
+            textEncoder = MultilingualTextEncoder(
+                modelAt: urls.multilingualTextEncoderProjectionURL,
+                configuration: config,
+                script: script ?? .latin
+            )
+        } else {
+            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)
-        let textEncoder = TextEncoder(tokenizer: tokenizer,
-                                      modelAt: urls.textEncoderURL,
-                                      configuration: config)
-        
+        textEncoder = TextEncoder(tokenizer: tokenizer, modelAt: urls.textEncoderURL, configuration: config)
+#endif
+
         // ControlNet model
         var controlNet: ControlNet? = nil
         let controlNetURLs = controlNetModelNames.map { model in
@@ -119,12 +139,28 @@ public extension StableDiffusionPipeline {
         }
 
         // Construct pipeline
-        self.init(textEncoder: textEncoder,
-                  unet: unet,
-                  decoder: decoder,
-                  encoder: encoder,
-                  controlNet: controlNet,
-                  safetyChecker: safetyChecker,
-                  reduceMemory: reduceMemory)
+        if #available(macOS 14.0, iOS 17.0, *) {
+            self.init(
+                textEncoder: textEncoder,
+                unet: unet,
+                decoder: decoder,
+                encoder: encoder,
+                controlNet: controlNet,
+                safetyChecker: safetyChecker,
+                reduceMemory: reduceMemory,
+                useMultilingualTextEncoder: useMultilingualTextEncoder,
+                script: script
+            )
+        } else {
+            self.init(
+                textEncoder: textEncoder,
+                unet: unet,
+                decoder: decoder,
+                encoder: encoder,
+                controlNet: controlNet,
+                safetyChecker: safetyChecker,
+                reduceMemory: reduceMemory
+            )
+        }
     }
 }

swift/StableDiffusion/pipeline/StableDiffusionPipeline.Configuration.swift

@@ -4,65 +4,97 @@
 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
+    /// Resolution dependent shifting of timestep schedules
+    public var schedulerTimestepShift: Float = 3.0
+    /// 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
+    /// Shift factor to use on the latent before decoding
+    public var decoderShiftFactor: Float32 = 0.0
+    /// 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

@@ -1,10 +1,11 @@
 // For licensing see accompanying LICENSE.md file.
 // Copyright (C) 2022 Apple Inc. All Rights Reserved.
 
-import Foundation
-import CoreML
 import Accelerate
 import CoreGraphics
+import CoreML
+import Foundation
+import NaturalLanguage
 
 /// Schedulers compatible with StableDiffusionPipeline
 public enum StableDiffusionScheduler {
@@ -12,6 +13,8 @@ public enum StableDiffusionScheduler {
     case pndmScheduler
     /// Scheduler that uses a second order DPM-Solver++ algorithm
     case dpmSolverMultistepScheduler
+    /// Scheduler for rectified flow based multimodal diffusion transformer models
+    case discreteFlowScheduler
 }
 
 /// RNG compatible with StableDiffusionPipeline
@@ -20,6 +23,36 @@ 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
+    case errorCreatingPreview
+}
+
+@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
@@ -27,14 +60,10 @@ 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
-    }
-    
+public struct StableDiffusionPipeline: StableDiffusionPipelineProtocol {
+
     /// Model to generate embeddings for tokenized input text
-    var textEncoder: TextEncoder
+    var textEncoder: TextEncoderModel
 
     /// Model used to predict noise residuals given an input, diffusion time step, and conditional embedding
     var unet: Unet
@@ -64,6 +93,12 @@ public struct StableDiffusionPipeline: ResourceManaging {
     /// This will increase latency in favor of reducing memory
     var reduceMemory: Bool = false
 
+    /// Option to use system multilingual NLContextualEmbedding as encoder
+    var useMultilingualTextEncoder: Bool = false
+
+    /// Optional natural language script to use for the text encoder.
+    var script: Script? = nil
+
     /// Creates a pipeline using the specified models and tokenizer
     ///
     /// - Parameters:
@@ -74,13 +109,15 @@ public struct StableDiffusionPipeline: ResourceManaging {
     ///   - safetyChecker: Optional model for checking safety of generated images
     ///   - reduceMemory: Option to enable reduced memory mode
     /// - Returns: Pipeline ready for image generation
-    public init(textEncoder: TextEncoder,
-                unet: Unet,
-                decoder: Decoder,
-                encoder: Encoder?,
-                controlNet: ControlNet? = nil,
-                safetyChecker: SafetyChecker? = nil,
-                reduceMemory: Bool = false) {
+    public init(
+        textEncoder: TextEncoderModel,
+        unet: Unet,
+        decoder: Decoder,
+        encoder: Encoder?,
+        controlNet: ControlNet? = nil,
+        safetyChecker: SafetyChecker? = nil,
+        reduceMemory: Bool = false
+    ) {
         self.textEncoder = textEncoder
         self.unet = unet
         self.decoder = decoder
@@ -90,6 +127,41 @@ public struct StableDiffusionPipeline: ResourceManaging {
         self.reduceMemory = reduceMemory
     }
 
+    /// Creates a pipeline using the specified models and tokenizer
+    ///
+    /// - Parameters:
+    ///   - textEncoder: Model for encoding tokenized text
+    ///   - unet: Model for noise prediction on latent samples
+    ///   - decoder: Model for decoding latent sample to image
+    ///   - controlNet: Optional model to control generated images by additonal inputs
+    ///   - safetyChecker: Optional model for checking safety of generated images
+    ///   - reduceMemory: Option to enable reduced memory mode
+    ///   - useMultilingualTextEncoder: Option to use system multilingual NLContextualEmbedding as encoder
+    ///   - script: Optional natural language script to use for the text encoder.
+    /// - Returns: Pipeline ready for image generation
+    @available(iOS 17.0, macOS 14.0, *)
+    public init(
+        textEncoder: TextEncoderModel,
+        unet: Unet,
+        decoder: Decoder,
+        encoder: Encoder?,
+        controlNet: ControlNet? = nil,
+        safetyChecker: SafetyChecker? = nil,
+        reduceMemory: Bool = false,
+        useMultilingualTextEncoder: Bool = false,
+        script: Script? = nil
+    ) {
+        self.textEncoder = textEncoder
+        self.unet = unet
+        self.decoder = decoder
+        self.encoder = encoder
+        self.controlNet = controlNet
+        self.safetyChecker = safetyChecker
+        self.reduceMemory = reduceMemory
+        self.useMultilingualTextEncoder = useMultilingualTextEncoder
+        self.script = script
+    }
+
     /// Load required resources for this pipeline
     ///
     /// If reducedMemory is true this will instead call prewarmResources instead
@@ -98,9 +170,10 @@ public struct StableDiffusionPipeline: ResourceManaging {
         if reduceMemory {
             try prewarmResources()
         } else {
-            try textEncoder.loadResources()
             try unet.loadResources()
+            try textEncoder.loadResources()
             try decoder.loadResources()
+            try encoder?.loadResources()
             try controlNet?.loadResources()
             try safetyChecker?.loadResources()
         }
@@ -111,6 +184,7 @@ public struct StableDiffusionPipeline: ResourceManaging {
         textEncoder.unloadResources()
         unet.unloadResources()
         decoder.unloadResources()
+        encoder?.unloadResources()
         controlNet?.unloadResources()
         safetyChecker?.unloadResources()
     }
@@ -120,6 +194,7 @@ public struct StableDiffusionPipeline: ResourceManaging {
         try textEncoder.prewarmResources()
         try unet.prewarmResources()
         try decoder.prewarmResources()
+        try encoder?.prewarmResources()
         try controlNet?.prewarmResources()
         try safetyChecker?.prewarmResources()
     }
@@ -135,38 +210,48 @@ public struct StableDiffusionPipeline: ResourceManaging {
         progressHandler: (Progress) -> Bool = { _ in true }
     ) throws -> [CGImage?] {
 
-        // Encode the input prompt and negative prompt
-        let promptEmbedding = try textEncoder.encode(config.prompt)
-        let negativePromptEmbedding = try textEncoder.encode(config.negativePrompt)
+        // Encode the input prompt
+        var promptEmbedding = try textEncoder.encode(config.prompt)
+
+        if config.guidanceScale >= 1.0 {
+            // Convert to Unet hidden state representation
+            // Concatenate the prompt and negative prompt embeddings
+            let negativePromptEmbedding = try textEncoder.encode(config.negativePrompt)
+            promptEmbedding = MLShapedArray<Float32>(
+                concatenating: [negativePromptEmbedding, promptEmbedding],
+                alongAxis: 0
+            )
+        }
 
         if reduceMemory {
             textEncoder.unloadResources()
         }
 
-        // Convert to Unet hidden state representation
-        // Concatenate the prompt and negative prompt embeddings
-        let concatEmbedding = MLShapedArray<Float32>(
-            concatenating: [negativePromptEmbedding, promptEmbedding],
-            alongAxis: 0
-        )
-
-        let hiddenStates = toHiddenStates(concatEmbedding)
+        let hiddenStates = useMultilingualTextEncoder ? promptEmbedding : toHiddenStates(promptEmbedding)
 
         /// 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)
+            case .dpmSolverMultistepScheduler: return DPMSolverMultistepScheduler(stepCount: config.stepCount, timeStepSpacing: config.schedulerTimestepSpacing)
+            case .discreteFlowScheduler: return DiscreteFlowScheduler(stepCount: config.stepCount, timeStepShift: config.schedulerTimestepShift)
             }
         }
 
         // 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
         
         // 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
@@ -179,10 +264,15 @@ public struct StableDiffusionPipeline: ResourceManaging {
 
             // 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)
+            let latentUnetInput: [MLShapedArray<Float32>]
+            if config.guidanceScale >= 1.0 {
+                latentUnetInput = latents.map {
+                    MLShapedArray<Float32>(concatenating: [$0, $0], alongAxis: 0)
+                }
+            } else {
+                latentUnetInput = latents
             }
-            
+
             // Before Unet, execute controlNet and add the output into Unet inputs
             let additionalResiduals = try controlNet?.execute(
                 latents: latentUnetInput,
@@ -193,14 +283,42 @@ public struct StableDiffusionPipeline: ResourceManaging {
             
             // Predict noise residuals from latent samples
             // and current time step conditioned on hidden states
-            var noise = try unet.predictNoise(
-                latents: latentUnetInput,
-                timeStep: t,
-                hiddenStates: hiddenStates,
-                additionalResiduals: additionalResiduals
-            )
+            var noise : [MLShapedArray<Float32>]
+            if unet.latentSampleShape[0] >= 2 || config.guidanceScale < 1.0 {
+                // One predict call from the uNet, using batching if needed
+                noise = try unet.predictNoise(
+                  latents: latentUnetInput,
+                  timeStep: t,
+                  hiddenStates: hiddenStates,
+                  additionalResiduals: additionalResiduals
+                )
+            } else {
+                // Serial predictions from uNet
+                var hidden0 = MLShapedArray<Float32>(converting: hiddenStates[0])
+                hidden0 = MLShapedArray(scalars: hidden0.scalars, shape: [1]+hidden0.shape)
+                let noise_pred_uncond = try unet.predictNoise(
+                  latents: latents,
+                  timeStep: t,
+                  hiddenStates: hidden0,
+                  additionalResiduals: additionalResiduals
+                )
 
-            noise = performGuidance(noise, config.guidanceScale)
+                var hidden1 = MLShapedArray<Float32>(converting: hiddenStates[1])
+                hidden1 = MLShapedArray(scalars: hidden1.scalars, shape: [1]+hidden1.shape)
+                let noise_pred_text = try unet.predictNoise(
+                  latents: latents,
+                  timeStep: t,
+                  hiddenStates: hidden1,
+                  additionalResiduals: additionalResiduals
+                )
+
+                noise = [MLShapedArray<Float32>(concatenating: [noise_pred_uncond[0], noise_pred_text[0]],
+                                                alongAxis: 0)]
+            }
+
+            if config.guidanceScale >= 1.0 {
+                noise = performGuidance(noise, config.guidanceScale)
+            }
 
             // Have the scheduler compute the previous (t-1) latent
             // sample given the predicted noise and current sample
@@ -210,15 +328,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) {
@@ -233,16 +355,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>] {
@@ -257,7 +370,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)
@@ -265,6 +378,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]
@@ -300,52 +481,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,400 @@
+// 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)
+            case .discreteFlowScheduler: return DiscreteFlowScheduler(stepCount: config.stepCount, timeStepShift: config.schedulerTimestepShift)
+            }
+        }
+
+        // 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/TextEncoder.swift

@@ -4,9 +4,15 @@
 import Foundation
 import CoreML
 
+@available(iOS 16.2, macOS 13.1, *)
+public protocol TextEncoderModel: ResourceManaging {
+
+    func encode(_ text: String) throws -> MLShapedArray<Float32>
+}
+
 ///  A model for encoding text
 @available(iOS 16.2, macOS 13.1, *)
-public struct TextEncoder: ResourceManaging {
+public struct TextEncoder: TextEncoderModel {
 
     /// Text tokenizer
     var tokenizer: BPETokenizer

swift/StableDiffusion/pipeline/TextEncoderT5.swift

@@ -0,0 +1,124 @@
+// For licensing see accompanying LICENSE.md file.
+// Copyright (C) 2023 Apple Inc. All Rights Reserved.
+
+import Foundation
+import CoreML
+import Tokenizers
+
+@available(iOS 17.0, macOS 14.0, *)
+public protocol TextEncoderT5Model: ResourceManaging {
+    func encode(_ text: String) throws -> TextEncoderT5Output
+}
+
+@available(iOS 17.0, macOS 14.0, *)
+public struct TextEncoderT5Output {
+    public let encoderHiddenStates: MLShapedArray<Float32>
+}
+
+///  A model for encoding text, suitable for SD3
+@available(iOS 17.0, macOS 14.0, *)
+public struct TextEncoderT5: TextEncoderT5Model {
+
+    /// Text tokenizer
+    var tokenizer: Tokenizer
+
+    /// 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: Tokenizer,
+                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 -> TextEncoderT5Output {
+
+        // Get models expected input length
+        let inputLength = inputShape.last!
+
+        // Tokenize, padding to the expected length
+        var tokens = tokenizer.tokenize(text: text)
+        var ids = tokens.map { tokenizer.convertTokenToId($0) ?? 0 }
+
+        // Truncate if necessary
+        if ids.count > inputLength {
+            tokens = tokens.dropLast(tokens.count - inputLength)
+            ids = ids.dropLast(ids.count - inputLength)
+            print("Needed to truncate input for TextEncoderT5")
+        }
+
+        // Use the model to generate the embedding
+        let encodedText = try encode(ids: ids)
+        return encodedText
+    }
+
+    func encode(ids: [Int]) throws -> TextEncoderT5Output {
+        let inputName = "input_ids"
+        let inputShape = inputShape
+        let inputLength = inputShape[1]
+                
+        let bosToken = tokenizer.bosTokenId ?? 0
+        let eosToken = tokenizer.eosTokenId ?? 1
+        let padToken = bosToken
+        let maskToken = eosToken
+
+        // Truncate and pad input to the expected length
+        let truncatedIds = ids.prefix(inputLength - 1) + [eosToken]
+        let inputIds = truncatedIds + Array(repeating: padToken, count: inputLength - truncatedIds.count)
+
+        let attentionMaskName = "attention_mask"
+        var attentionMask: [Int] = inputIds.map { token in
+            token == padToken ? maskToken : padToken
+        }
+        attentionMask[0] = bosToken
+
+        let floatIds = inputIds.map { Float32($0) }
+        let floatMask = attentionMask.map { Float32($0) }
+
+        let inputArray = MLShapedArray<Float32>(scalars: floatIds, shape: inputShape)
+        let maskArray = MLShapedArray<Float32>(scalars: floatMask, shape: inputShape)
+        let inputFeatures = try! MLDictionaryFeatureProvider(
+            dictionary: [inputName: MLMultiArray(inputArray),
+                         attentionMaskName: MLMultiArray(maskArray)])
+
+        let result = try model.perform { model in
+            try model.prediction(from: inputFeatures)
+        }
+
+        let embeddingFeature = result.featureValue(for: "encoder_hidden_states")
+        return TextEncoderT5Output(encoderHiddenStates: MLShapedArray<Float32>(converting: embeddingFeature!.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/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:
@@ -84,7 +95,12 @@ public struct Unet: ResourceManaging {
     ) throws -> [MLShapedArray<Float32>] {
 
         // Match time step batch dimension to the model / latent samples
-        let t = MLShapedArray<Float32>(scalars:[Float(timeStep), Float(timeStep)],shape:[2])
+        let t: MLShapedArray<Float32>
+        if hiddenStates.shape[0] == 2 {
+            t = MLShapedArray(scalars: [Float(timeStep), Float(timeStep)], shape: [2])
+        } else {
+            t = MLShapedArray(scalars: [Float(timeStep)], shape: [1])
+        }
 
         // Form batch input to model
         let inputs = try latents.enumerated().map {
@@ -103,7 +119,7 @@ public struct Unet: ResourceManaging {
         let batch = MLArrayBatchProvider(array: inputs)
 
         // Make predictions
-        let results = try predictions(from: batch)
+        let results = try models.predictions(from: batch)
 
         // Pull out the results in Float32 format
         let noise = (0..<results.count).map { i in
@@ -127,50 +143,62 @@ public struct Unet: ResourceManaging {
         return noise
     }
 
-    func predictions(from batch: MLBatchProvider) throws -> MLBatchProvider {
+    /// 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>] {
 
-        var results = try models.first!.perform { model in
-            try model.predictions(fromBatch: batch)
+        // 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 models.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)
         }
 
-        if models.count == 1 {
-            return results
-        }
-
-        // Manual pipeline batch prediction
-        let inputs = batch.arrayOfFeatureValueDictionaries
-        for stage in models.dropFirst() {
-
-            // Combine the original inputs with the outputs of the last stage
-            let next = try results.arrayOfFeatureValueDictionaries
-                .enumerated().map { (index, dict) in
-                    let nextDict =  dict.merging(inputs[index]) { (out, _) in out }
-                    return try MLDictionaryFeatureProvider(dictionary: nextDict)
-            }
-            let nextBatch = MLArrayBatchProvider(array: next)
-
-            // Predict
-            results = try stage.perform { model in
-                try model.predictions(fromBatch: nextBatch)
-            }
-        }
-
-        return results
-    }
-}
-
-extension MLFeatureProvider {
-    var featureValueDictionary: [String : MLFeatureValue] {
-        self.featureNames.reduce(into: [String : MLFeatureValue]()) { result, name in
-            result[name] = self.featureValue(for: name)
-        }
-    }
-}
-
-extension MLBatchProvider {
-    var arrayOfFeatureValueDictionaries: [[String : MLFeatureValue]] {
-        (0..<self.count).map {
-            self.features(at: $0).featureValueDictionary
-        }
+        return noise
     }
 }

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/StableDiffusion/tokenizer/T5Tokenizer.swift

@@ -0,0 +1,21 @@
+// For licensing see accompanying LICENSE.md file.
+// Copyright (C) 2024 Apple Inc. All Rights Reserved.
+
+import Foundation
+import Hub
+import Tokenizers
+
+/// Extension to swift-transfomers Hub.swift to load local Config files
+public extension Config {
+    /// Assumes the file is already present at local url.
+    /// `fileURL` is a complete local file path for the given model
+    public init(fileURL: URL) throws  {
+        let data = try Data(contentsOf: fileURL)
+        let parsed = try JSONSerialization.jsonObject(with: data, options: [])
+        guard var dictionary = parsed as? [String: Any] else { throw Hub.HubClientError.parse }
+        
+        // Necessary override for loading local tokenizer configs
+        dictionary["tokenizer_class"] = "T5Tokenizer"
+        self.init(dictionary)
+    }
+}

swift/StableDiffusionCLI/main.swift

@@ -9,6 +9,7 @@ import StableDiffusion
 import UniformTypeIdentifiers
 import Cocoa
 import CoreImage
+import NaturalLanguage
 
 @available(iOS 16.2, macOS 13.1, *)
 struct StableDiffusionSample: ParsableCommand {
@@ -35,6 +36,12 @@ struct StableDiffusionSample: ParsableCommand {
     )
     var resourcePath: String = "./"
     
+    @Flag(name: .customLong("xl"), help: "The resources correspond to a Stable Diffusion XL model")
+    var isXL: Bool = false
+
+    @Flag(name: .customLong("sd3"), help: "The resources correspond to a Stable Diffusion 3 model")
+    var isSD3: Bool = false
+
     @Option(help: "Path to starting image.")
     var image: String? = nil
     
@@ -70,7 +77,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(
@@ -91,6 +98,12 @@ struct StableDiffusionSample: ParsableCommand {
     @Flag(help: "Reduce memory usage")
     var reduceMemory: Bool = false
 
+    @Flag(help: "Use system multilingual NLContextualEmbedding as encoder model")
+    var useMultilingualTextEncoder: Bool = false
+
+    @Option(help: "The natural language script for the multilingual contextual embedding")
+    var script: Script = .latin
+
     mutating func run() throws {
         guard FileManager.default.fileExists(atPath: resourcePath) else {
             throw RunError.resources("Resource path does not exist \(resourcePath)")
@@ -102,11 +115,60 @@ 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 = try StableDiffusionPipeline(resourcesAt: resourceURL,
-                                                   controlNet: controlnet,
-                                                   configuration: config,
-                                                   disableSafety: disableSafety,
-                                                   reduceMemory: reduceMemory)
+        let pipeline: StableDiffusionPipelineProtocol
+        var scaleFactor: Float32 = 0.18215
+        var shiftFactor: Float32 = 0.0
+        var timestepShift: Float32 = 1.0
+        if #available(macOS 14.0, iOS 17.0, *) {
+            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 if isSD3 {
+                scaleFactor = 1.5305
+                shiftFactor = 0.0609
+                timestepShift = 3.0
+                if !controlnet.isEmpty {
+                    throw RunError.unsupported("ControlNet is not supported for Stable Diffusion 3")
+                }
+                if useMultilingualTextEncoder {
+                    throw RunError.unsupported("Multilingual text encoder is not yet supported for Stable Diffusion 3")
+                }
+                pipeline = try StableDiffusion3Pipeline(
+                    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,
+                controlNet: controlnet,
+                configuration: config,
+                disableSafety: disableSafety,
+                reduceMemory: reduceMemory
+            )
+        }
+
         try pipeline.loadResources()
         
         let startingImage: CGImage?
@@ -153,17 +215,21 @@ struct StableDiffusionSample: ParsableCommand {
         pipelineConfig.guidanceScale = guidanceScale
         pipelineConfig.schedulerType = scheduler.stableDiffusionScheduler
         pipelineConfig.rngType = rng.stableDiffusionRNG
-        
+        pipelineConfig.useDenoisedIntermediates = true
+        pipelineConfig.encoderScaleFactor = scaleFactor
+        pipelineConfig.decoderScaleFactor = scaleFactor
+        pipelineConfig.decoderShiftFactor = shiftFactor
+        pipelineConfig.schedulerTimestepShift = timestepShift
+
         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)
     }
@@ -262,6 +328,7 @@ struct StableDiffusionSample: ParsableCommand {
 enum RunError: Error {
     case resources(String)
     case saving(String)
+    case unsupported(String)
 }
 
 @available(iOS 16.2, macOS 13.1, *)
@@ -290,15 +357,19 @@ 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
         }
     }
 }
 
+@available(iOS 16.2, macOS 13.1, *)
+extension Script: ExpressibleByArgument {}
+
 if #available(iOS 16.2, macOS 13.1, *) {
     StableDiffusionSample.main()
 } else {