Merge pull request #368 from Aintor/main

Fix for Mask Patch Failure and Quantization Issues in Latest `transformers` Versions
Add --min-deployment-target argument for configurable iOS/macOS deployment targets
2025-07-03 11:01:02 -07:00 · 2025-06-14 08:27:06 -07:00 · 2025-01-23 08:28:43 -08:00 · 2024-11-15 16:05:22 -08:00 · 2024-11-12 11:19:16 +08:00 · 2024-11-12 10:36:17 +08:00
31 changed files with 2280 additions and 256 deletions
@@ -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!
@@ -1,3 +1,5 @@
+*~
+
 # Swift Package
 .DS_Store
 /.build
@@ -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

@@ -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.
@@ -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",
@@ -97,6 +97,7 @@ Target Device Hardware Generation:

 |        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              |
@@ -182,7 +183,7 @@ Resources:
 <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 from coremltools-7.0](https://apple.github.io/coremltools/docs-guides/source/post-training-palettization.html).
+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.

@@ -193,7 +194,7 @@ MBP builds a per-layer "palettization recipe" by picking a suitable number of bi

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

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

@@ -245,6 +246,121 @@ An example `<selected-recipe-string-key>` would be `"recipe_4.50_bit_mixedpalett

 </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>
@@ -336,7 +452,7 @@ Resources:
 <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).

 * 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)
@@ -355,6 +471,7 @@ Resources:
  - [`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).

@@ -378,7 +495,7 @@ 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
@@ -428,7 +545,7 @@ This generally takes 15-20 minutes on an M1 MacBook Pro. Upon successful executi

 - `--refiner-version`: The refiner version name as published on the [Hugging Face Hub](https://huggingface.co/models?search=stable-diffusion). This is optional and if specified, this argument will convert and bundle the refiner unet alongside the model unet.

- `--bundle-resources-for-swift-cli`: Compiles all 4 models and bundles them along with necessary resources for text tokenization into `<output-mlpackages-directory>/Resources` which should provided as input to the Swift package. This flag is not necessary for the diffusers-based Python pipeline.
+- `--bundle-resources-for-swift-cli`: Compiles all 4 models and bundles them along with necessary resources for text tokenization into `<output-mlpackages-directory>/Resources` which should provided as input to the Swift package. This flag is not necessary for the diffusers-based Python pipeline. [However using these compiled models in Python will significantly speed up inference](https://apple.github.io/coremltools/docs-guides/source/model-prediction.html#why-use-a-compiled-model).

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

@@ -442,6 +559,8 @@ This generally takes 15-20 minutes on an M1 MacBook Pro. Upon successful executi

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

+- `--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>
@@ -454,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. If you had specified `--bundle-resources-for-swift-cli` during conversion, then `-i` should point to the resulting `Resources` folder which holds the compiled `.mlmodelc` files. The compiled models load much faster after first use.
+- `-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>

@@ -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)
@@ -4,6 +4,22 @@ 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
@@ -38,9 +54,15 @@ def split_einsum(q, k, v, mask, heads, dim_head):
        for head_idx in range(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
+    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)
@@ -101,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:
@@ -132,6 +133,7 @@ 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]
        if not op.op_type.startswith("const"):
@@ -181,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
@@ -228,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
@@ -307,13 +311,6 @@ def main(args):
            second_chunk_model=model_chunk2,
        )

-    # Remove original (non-chunked) model if requested
-    if args.remove_original:
-        logger.info(
-            "Removing original (non-chunked) model at {args.mlpackage_path}")
-        shutil.rmtree(args.mlpackage_path)
-        logger.info("Done.")
-
    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)
@@ -342,6 +339,39 @@ def main(args):
        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(
+            "Removing original (non-chunked) model at {args.mlpackage_path}")
+        shutil.rmtree(args.mlpackage_path)
+        logger.info("Done.")
+
+
 if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
@@ -16,13 +16,28 @@ 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, sources='packages'):
+    def __init__(self, model_path, compute_unit, sources='packages', optimization_hints=None):

        logger.info(f"Loading {model_path}")

@@ -31,7 +46,10 @@ class CoreMLModel:
            assert os.path.exists(model_path) and model_path.endswith(".mlpackage")

            self.model = ct.models.MLModel(
-                model_path, compute_units=ct.ComputeUnit[compute_unit])
+                model_path,
+                compute_units=ct.ComputeUnit[compute_unit],
+                optimization_hints=optimization_hints,
+            )
            DTYPE_MAP = {
                65552: np.float16,
                65568: np.float32,
@@ -47,7 +65,11 @@ class CoreMLModel:
        elif sources == 'compiled':
            assert os.path.exists(model_path) and model_path.endswith(".mlmodelc")

-            self.model = ct.models.CompiledMLModel(model_path, ct.ComputeUnit[compute_unit])
+            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:
@@ -159,8 +181,8 @@ def _load_mlpackage(submodule_name,
        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"]
-        compiled_names = ['TextEncoder', 'TextEncoder2', 'Unet', 'VAEDecoder', 'VAEEncoder']
+        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'
@@ -170,7 +192,15 @@ def _load_mlpackage(submodule_name,
            raise FileNotFoundError(
                f"{submodule_name} CoreML model doesn't exist at {mlpackage_path}")

-    return CoreMLModel(mlpackage_path, compute_unit, sources=sources)
+    # 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)


 def _load_mlpackage_controlnet(mlpackages_dir, model_version, compute_unit):
@@ -1,18 +1,20 @@
-from pprint import pprint
 import argparse
-import coremltools as ct
 import gc
 import json
 import logging
-import numpy as np
 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)
@@ -23,9 +25,6 @@ def main(args):
    coreml_model = ct.models.MLModel(args.mlpackage_path, compute_units=ct.ComputeUnit.CPU_ONLY)
    logger.info(f"Loaded {args.mlpackage_path}")

-    # Keep track of precision stats
-    precision_stats = {nbits:{'num_tensors': 0, 'numel': 0} for nbits in NBITS}
-    
    # Load palettization recipe
    with open(args.pre_analysis_json_path, 'r') as f:
        pre_analysis = json.load(f)
@@ -62,53 +61,29 @@ def main(args):
    del pipe
    gc.collect()

-    current_nbits: int
-
-    def op_selector(const):
-        parameter_tensor = const.val.val
-        if parameter_tensor.size < MIN_SIZE:
-            return False
-
-        if parameter_tensor.dtype != np.float16:
-            # These are the tensors that were compressed to look-up indices in previous passes
-            return False
-
-        tensor_hash = get_tensor_hash(parameter_tensor)
-        tensor_spec = f"{tensor_hash} with shape {parameter_tensor.shape}"
-
-
-        hashes = list(hashed_recipe)
-        pdist = np.abs(np.array(hashes) - tensor_hash)
+    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()
-        logger.debug(f"{tensor_spec}: {tensor_hash} matched with {hashes[matched]} (hash error={pdist.min()})")
-
        target_nbits = hashed_recipe[hashes[matched]]
-        
-        do_palettize = current_nbits == target_nbits
-        if do_palettize:
-            logger.debug(f"{tensor_spec}: Palettizing to {target_nbits}-bit palette")
-            precision_stats[current_nbits]['num_tensors'] += 1
-            precision_stats[current_nbits]['numel'] +=  np.prod(parameter_tensor.shape)
-            return True
-        return False

-    for nbits in NBITS:
-        logger.info(f"Processing tensors targeting {nbits}-bit palettes")
-        current_nbits = nbits
+        if target_nbits == 16:
+            continue

-        config = ct.optimize.coreml.OptimizationConfig(
-           global_config=ct.optimize.coreml.OpPalettizerConfig(mode="kmeans", nbits=nbits, weight_threshold=None,),
-           is_deprecated=True,
-           op_selector=op_selector,
+        op_name_configs[name] = cto.OpPalettizerConfig(
+            mode="kmeans",
+            nbits=target_nbits,
+            weight_threshold=int(MIN_SIZE)
        )
-        coreml_model = ct.optimize.coreml.palettize_weights(coreml_model, config=config)
-        logger.info(f"{precision_stats[nbits]['num_tensors']} tensors are palettized with {nbits} bits")

+    config = ct.optimize.coreml.OptimizationConfig(op_name_configs=op_name_configs)
+    coreml_model = ct.optimize.coreml.palettize_weights(coreml_model, config)

-    tot_numel = sum([precision_stats[nbits]['numel'] for nbits in NBITS])
-    final_size = sum([precision_stats[nbits]['numel'] * nbits for nbits in NBITS])
-    logger.info(f"Palettization result: {final_size / tot_numel:.2f}-bits resulting in {final_size / (8*1e6)} MB")
-    pprint(precision_stats)
    coreml_model.save(args.o)


@@ -136,6 +111,15 @@ if __name__ == "__main__":
        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()

@@ -18,9 +18,10 @@ import time

 import torch
 import torch.nn as nn
+import requests
 torch.set_grad_enabled(False)

-from tqdm import tqdm, trange
+from tqdm import tqdm

 # Bit-widths the Neural Engine is capable of accelerating
 NBITS = [1, 2, 4, 6, 8]
@@ -341,8 +342,8 @@ def simulate_quant_fn(ref_pipe, quantization_to_simulate):

    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)
+        float(f"{compute_psnr(r, t):.1f}")
+        for r, t in zip(ref_out, simulated_out)
    ]) / len(ref_out)

    return simulated_out, simulated_psnr
@@ -458,9 +459,7 @@ def main(args):
    json_name = f"{args.model_version.replace('/','-')}_palettization_recipe.json"
    candidates, sizes = get_palettizable_modules(pipe.unet)

-    sizes_table = {
-        candidate:size for candidate, size in zip(candidates, sizes)
-    }
+    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:
@@ -570,6 +569,15 @@ if __name__ == "__main__":
        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)
@@ -5,6 +5,7 @@

 import argparse

+from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline
 from diffusers.pipelines.pipeline_utils import DiffusionPipeline
 from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
 from diffusers.schedulers import (
@@ -415,9 +416,10 @@ class CoreMLStableDiffusionPipeline(DiffusionPipeline):
            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,
+            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
@@ -505,11 +507,16 @@ class CoreMLStableDiffusionPipeline(DiffusionPipeline):
            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:
                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,
                )
@@ -517,18 +524,40 @@ class CoreMLStableDiffusionPipeline(DiffusionPipeline):
                control_net_additional_residuals = {}

            # predict the noise residual
-            unet_additional_kwargs = unet_additional_kwargs | control_net_additional_residuals
+            unet_additional_kwargs.update(control_net_additional_residuals)

-            noise_pred = self.unet(
-                sample=latent_model_input.astype(np.float16),
-                timestep=np.array([t, t], np.float16),
-                encoder_hidden_states=text_embeddings.astype(np.float16),
-                **unet_additional_kwargs,
-            )["noise_pred"]
+            # 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)

@@ -603,8 +632,7 @@ def get_coreml_pipe(pytorch_pipe,
            "tokenizer": pytorch_pipe.tokenizer,
            'tokenizer_2': pytorch_pipe.tokenizer_2,
            "scheduler": pytorch_pipe.scheduler if scheduler_override is None else scheduler_override,
-            "force_zeros_for_empty_prompt": force_zeros_for_empty_prompt,
-            'xl': True
+            'xl': True,
        }

        model_packages_to_load = ["text_encoder", "text_encoder_2", "unet", "vae_decoder"]
@@ -617,6 +645,8 @@ def get_coreml_pipe(pytorch_pipe,
        }
        model_packages_to_load = ["text_encoder", "unet", "vae_decoder"]

+    coreml_pipe_kwargs["force_zeros_for_empty_prompt"] = force_zeros_for_empty_prompt
+
    if getattr(pytorch_pipe, "safety_checker", None) is not None:
        model_packages_to_load.append("safety_checker")
    else:
@@ -697,7 +727,6 @@ def main(args):

    logger.info("Initializing PyTorch pipe for reference configuration")

-    from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline
    SDP = StableDiffusionXLPipeline if 'xl' in args.model_version else StableDiffusionPipeline

    pytorch_pipe = SDP.from_pretrained(
@@ -713,7 +742,7 @@ def main(args):

    # Get Force Zeros Config if it exists
    force_zeros_for_empty_prompt: bool = False
-    if 'force_zeros_for_empty_prompt' in pytorch_pipe.config:
+    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(
@@ -745,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)
@@ -815,6 +845,10 @@ if __name__ == "__main__":
        "--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'],
@@ -16,7 +16,12 @@ from diffusers import (
    DiffusionPipeline,
    ControlNetModel
 )
+from diffusionkit.tests.torch2coreml import (
+    convert_mmdit_to_mlpackage,
+    convert_vae_to_mlpackage
+)
 import gc
+from huggingface_hub import snapshot_download

 import logging

@@ -121,10 +126,11 @@ def _convert_to_coreml(submodule_name, torchscript_module, sample_inputs,
        coreml_model.compute_unit = compute_unit
    else:
        logger.info(f"Converting {submodule_name} to CoreML..")
+        deployment_target = _get_deployment_target(args.min_deployment_target)
        coreml_model = ct.convert(
            torchscript_module,
            convert_to="mlprogram",
-            minimum_deployment_target=ct.target.macOS13,
+            minimum_deployment_target=deployment_target,
            inputs=_get_coreml_inputs(sample_inputs, args),
            outputs=[ct.TensorType(name=name, dtype=np.float32) for name in output_names],
            compute_units=compute_unit,
@@ -138,6 +144,41 @@ def _convert_to_coreml(submodule_name, torchscript_module, sample_inputs,
    return coreml_model, out_path


+def _get_deployment_target(target_string):
+    """
+    Convert deployment target string to coremltools target object.
+    
+    Args:
+        target_string (str): Target deployment string (e.g., "macOS13", "iOS18")
+        
+    Returns:
+        coremltools target object
+    """
+    target_map = {
+        "macOS13": ct.target.macOS13,
+        "macOS14": ct.target.macOS14,
+        "iOS16": ct.target.iOS16,
+        "iOS17": ct.target.iOS17,
+    }
+    
+    # Handle newer targets that might not be available in older coremltools versions
+    try:
+        if target_string == "macOS15":
+            return ct.target.macOS15
+        elif target_string == "iOS18":
+            return ct.target.iOS18
+    except AttributeError:
+        logger.warning(f"Deployment target {target_string} not available in this coremltools version. "
+                      f"Using macOS14 as fallback.")
+        return ct.target.macOS14
+    
+    if target_string in target_map:
+        return target_map[target_string]
+    else:
+        logger.warning(f"Unknown deployment target {target_string}. Using macOS13 as fallback.")
+        return ct.target.macOS13
+
+
 def quantize_weights(args):
    """ Quantize weights to args.quantize_nbits using a palette (look-up table)
    """
@@ -207,6 +248,26 @@ def _compile_coreml_model(source_model_path, output_dir, final_name):
    return target_path


+def _download_t5_model(args, t5_save_path):
+    t5_url = args.text_encoder_t5_url
+    match = re.match(r'https://huggingface.co/(.+)/resolve/main/(.+)', t5_url)
+    if not match:
+        raise ValueError(f"Invalid Hugging Face URL: {t5_url}")
+    repo_id, model_subpath = match.groups()
+
+    download_path = snapshot_download(
+        repo_id=repo_id,
+        revision="main",
+        allow_patterns=[f"{model_subpath}/*"]
+    )
+    logger.info(f"Downloaded T5 model to {download_path}")
+
+    # Move the downloaded model to the top level of the Resources directory
+    logger.info(f"Copying T5 model from {download_path} to {t5_save_path}")
+    cache_path = os.path.join(download_path, model_subpath)
+    shutil.copytree(cache_path, t5_save_path)
+
+
 def bundle_resources_for_swift_cli(args):
    """
    - Compiles Core ML models from mlpackage into mlmodelc format
@@ -228,6 +289,7 @@ def bundle_resources_for_swift_cli(args):
                                     ("refiner", "UnetRefiner"),
                                     ("refiner_chunk1", "UnetRefinerChunk1"),
                                     ("refiner_chunk2", "UnetRefinerChunk2"),
+                                     ("mmdit", "MultiModalDiffusionTransformer"),
                                     ("control-unet", "ControlledUnet"),
                                     ("control-unet_chunk1", "ControlledUnetChunk1"),
                                     ("control-unet_chunk2", "ControlledUnetChunk2"),
@@ -241,7 +303,7 @@ def bundle_resources_for_swift_cli(args):
            logger.warning(
                f"{source_path} not found, skipping compilation to {target_name}.mlmodelc"
            )
-            
+
    if args.convert_controlnet:
        for controlnet_model_version in args.convert_controlnet:
            controlnet_model_name = controlnet_model_version.replace("/", "_")
@@ -271,12 +333,33 @@ def bundle_resources_for_swift_cli(args):
        f.write(requests.get(args.text_encoder_merges_url).content)
    logger.info("Done")

+    # Fetch and save pre-converted T5 text encoder model
+    t5_model_name = "TextEncoderT5.mlmodelc"
+    t5_save_path = os.path.join(resources_dir, t5_model_name)
+    if args.include_t5:
+        if not os.path.exists(t5_save_path):
+            logger.info("Downloading pre-converted T5 encoder model TextEncoderT5.mlmodelc")
+            _download_t5_model(args, t5_save_path)
+            logger.info("Done")
+        else:
+            logger.info(f"Skipping T5 download as {t5_save_path} already exists")
+            
+        # Fetch and save T5 text tokenizer JSON files
+        logger.info("Downloading and saving T5 tokenizer files tokenizer_config.json and tokenizer.json")
+        with open(os.path.join(resources_dir, "tokenizer_config.json"), "wb") as f:
+            f.write(requests.get(args.text_encoder_t5_config_url).content)
+        with open(os.path.join(resources_dir, "tokenizer.json"), "wb") as f:
+            f.write(requests.get(args.text_encoder_t5_data_url).content)
+        logger.info("Done")
+
    return resources_dir


 from transformers.models.clip import modeling_clip

 # Copied from https://github.com/huggingface/transformers/blob/v4.30.0/src/transformers/models/clip/modeling_clip.py#L677C1-L692C1
+# Starting from transformers >= 4.35.0, the _make_causal_mask function is replaced by _create_4d_causal_attention_mask in modeling_clip.
+# For backward compatibility with versions < 4.35.0, both functions are patched here.
 def patched_make_causal_mask(input_ids_shape, dtype, device, past_key_values_length: int = 0):
    """ Patch to replace torch.finfo(dtype).min with -1e4
    """
@@ -289,8 +372,9 @@ def patched_make_causal_mask(input_ids_shape, dtype, device, past_key_values_len
    if past_key_values_length > 0:
        mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype, device=device), mask], dim=-1)
    return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length)
-
-modeling_clip._make_causal_mask = patched_make_causal_mask
+    
+modeling_clip._make_causal_mask = patched_make_causal_mask # For transformers >= 4.30.0 and transformers < 4.35.0
+modeling_clip._create_4d_causal_attention_mask = patched_make_causal_mask # For transformers >= 4.35.0

 def convert_text_encoder(text_encoder, tokenizer, submodule_name, args):
    """ Converts the text encoder component of Stable Diffusion
@@ -480,7 +564,7 @@ def convert_vae_decoder(pipe, args):
        1,  # B
        pipe.vae.config.latent_channels,  # C
        args.latent_h or pipe.unet.config.sample_size,  # H
-        args.latent_w or pipe.unet.config.sample_size,  # w
+        args.latent_w or pipe.unet.config.sample_size,  # W
    )

    if args.custom_vae_version is None and args.xl_version:
@@ -557,6 +641,61 @@ def convert_vae_decoder(pipe, args):
    del traced_vae_decoder, pipe.vae.decoder, coreml_vae_decoder
    gc.collect()

+def convert_vae_decoder_sd3(args):
+    """ Converts the VAE component of Stable Diffusion 3
+    """
+    out_path = _get_out_path(args, "vae_decoder")
+    if os.path.exists(out_path):
+        logger.info(
+            f"`vae_decoder` already exists at {out_path}, skipping conversion."
+        )
+        return
+    
+    # Convert the VAE Decoder model via DiffusionKit
+    converted_vae_path = convert_vae_to_mlpackage(
+        model_version=args.model_version, 
+        latent_h=args.latent_h, 
+        latent_w=args.latent_w, 
+        output_dir=args.o,
+    )
+
+    # Load converted model
+    coreml_vae_decoder = ct.models.MLModel(converted_vae_path)
+
+    # Set model metadata
+    coreml_vae_decoder.author = f"Please refer to the Model Card available at huggingface.co/{args.model_version}"
+    coreml_vae_decoder.license = "Stability AI Community License (https://huggingface.co/stabilityai/stable-diffusion-3-medium/blob/main/LICENSE.md)"
+    coreml_vae_decoder.version = args.model_version
+    coreml_vae_decodershort_description = \
+        "Stable Diffusion 3 generates images conditioned on text or other images as input through the diffusion process. " \
+        "Please refer to https://arxiv.org/pdf/2403.03206 for details."
+
+    # Set the input descriptions
+    coreml_vae_decoder.input_description["z"] = \
+        "The denoised latent embeddings from the unet model after the last step of reverse diffusion"
+
+    # Set the output descriptions
+    coreml_vae_decoder.output_description[
+        "image"] = "Generated image normalized to range [-1, 1]"
+
+    # Set package version metadata
+    from python_coreml_stable_diffusion._version import __version__
+    coreml_vae_decoder.user_defined_metadata["com.github.apple.ml-stable-diffusion.version"] = __version__
+    from diffusionkit.version import __version__
+    coreml_vae_decoder.user_defined_metadata["com.github.argmax.diffusionkit.version"] = __version__
+
+    # Save the updated model
+    coreml_vae_decoder.save(out_path)
+
+    logger.info(f"Saved vae_decoder into {out_path}")
+
+    # Delete the original file
+    if os.path.exists(converted_vae_path):
+        shutil.rmtree(converted_vae_path)
+
+    del coreml_vae_decoder
+    gc.collect()
+

 def convert_vae_encoder(pipe, args):
    """ Converts the VAE Encoder component of Stable Diffusion
@@ -657,7 +796,7 @@ def convert_vae_encoder(pipe, args):
    gc.collect()


-def convert_unet(pipe, args, model_name = None):
+def convert_unet(pipe, args, model_name=None):
    """ Converts the UNet component of Stable Diffusion
    """
    if args.unet_support_controlnet:
@@ -683,6 +822,8 @@ def convert_unet(pipe, args, model_name = None):
    elif not os.path.exists(out_path):
        # Prepare sample input shapes and values
        batch_size = 2  # for classifier-free guidance
+        if args.unet_batch_one:
+            batch_size = 1  # for not using classifier-free guidance
        sample_shape = (
            batch_size,                    # B
            pipe.unet.config.in_channels,  # C
@@ -771,7 +912,7 @@ def convert_unet(pipe, args, model_name = None):
        else:
            unet_cls = unet.UNet2DConditionModel

-        reference_unet = unet_cls(**pipe.unet.config).eval()
+        reference_unet = unet_cls(support_controlnet=args.unet_support_controlnet, **pipe.unet.config).eval()

        load_state_dict_summary = reference_unet.load_state_dict(
            pipe.unet.state_dict())
@@ -827,7 +968,7 @@ def convert_unet(pipe, args, model_name = None):
        logger.info("Done.")

        if args.check_output_correctness:
-            baseline_out = pipe.unet(**baseline_sample_unet_inputs,
+            baseline_out = pipe.unet.to(torch.float32)(**baseline_sample_unet_inputs,
                                     return_dict=False)[0].numpy()
            reference_out = reference_unet(*sample_unet_inputs.values())[0].numpy()
            report_correctness(baseline_out, reference_out,
@@ -909,6 +1050,72 @@ def convert_unet(pipe, args, model_name = None):
        chunk_mlprogram.main(args)


+def convert_mmdit(args):
+    """ Converts the MMDiT component of Stable Diffusion 3
+    """
+    out_path = _get_out_path(args, "mmdit")
+    if os.path.exists(out_path):
+        logger.info(
+            f"`mmdit` already exists at {out_path}, skipping conversion."
+        )
+        return
+    
+    # Convert the MMDiT model via DiffusionKit
+    converted_mmdit_path = convert_mmdit_to_mlpackage(
+        model_version=args.model_version, 
+        latent_h=args.latent_h, 
+        latent_w=args.latent_w, 
+        output_dir=args.o,
+        # FIXME: Hardcoding to CPU_AND_GPU since ANE doesn't support FLOAT32
+        compute_precision=ct.precision.FLOAT32,
+        compute_unit=ct.ComputeUnit.CPU_AND_GPU,
+    )
+
+    # Load converted model
+    coreml_mmdit = ct.models.MLModel(converted_mmdit_path)
+
+    # Set model metadata
+    coreml_mmdit.author = f"Please refer to the Model Card available at huggingface.co/{args.model_version}"
+    coreml_mmdit.license = "Stability AI Community License (https://huggingface.co/stabilityai/stable-diffusion-3-medium/blob/main/LICENSE.md)"
+    coreml_mmdit.version = args.model_version
+    coreml_mmdit.short_description = \
+    "Stable Diffusion 3 generates images conditioned on text or other images as input through the diffusion process. " \
+    "Please refer to https://arxiv.org/pdf/2403.03206 for details."
+
+    # Set the input descriptions
+    coreml_mmdit.input_description["latent_image_embeddings"] = \
+        "The low resolution latent feature maps being denoised through reverse diffusion"
+    coreml_mmdit.input_description["token_level_text_embeddings"] = \
+        "Output embeddings from the associated text_encoder model to condition to generated image on text. " \
+        "A maximum of 77 tokens (~40 words) are allowed. Longer text is truncated. "
+    coreml_mmdit.input_description["pooled_text_embeddings"] = \
+        "Additional embeddings that if specified are added to the embeddings that are passed along to the MMDiT model."
+    coreml_mmdit.input_description["timestep"] = \
+        "A value emitted by the associated scheduler object to condition the model on a given noise schedule"
+    
+    # Set the output descriptions
+    coreml_mmdit.output_description["denoiser_output"] = \
+        "Same shape and dtype as the `latent_image_embeddings` input. " \
+        "The predicted noise to facilitate the reverse diffusion (denoising) process"
+
+    # Set package version metadata
+    from python_coreml_stable_diffusion._version import __version__
+    coreml_mmdit.user_defined_metadata["com.github.apple.ml-stable-diffusion.version"] = __version__
+    from diffusionkit.version import __version__
+    coreml_mmdit.user_defined_metadata["com.github.argmax.diffusionkit.version"] = __version__
+
+    # Save the updated model
+    coreml_mmdit.save(out_path)
+
+    logger.info(f"Saved vae_decoder into {out_path}")
+
+    # Delete the original file
+    if os.path.exists(converted_mmdit_path):
+        shutil.rmtree(converted_mmdit_path)
+
+    del coreml_mmdit
+    gc.collect()
+
 def convert_safety_checker(pipe, args):
    """ Converts the Safety Checker component of Stable Diffusion
    """
@@ -926,6 +1133,8 @@ def convert_safety_checker(pipe, args):
        )
        return

+    pipe.safety_checker = pipe.safety_checker.to(torch.float32)
+    
    im_h = pipe.vae.config.sample_size
    im_w = pipe.vae.config.sample_size

@@ -1286,6 +1495,16 @@ def get_pipeline(args):
                                            use_safetensors=True,
                                            vae=vae,
                                            use_auth_token=True)
+    elif args.sd3_version:
+        # SD3 uses standard SDXL diffusers pipeline besides the vae, denoiser, and T5 text encoder
+        sdxl_base_version = "stabilityai/stable-diffusion-xl-base-1.0"
+        args.xl_version = True
+        logger.info(f"SD3 version specified, initializing DiffusionPipeline with {sdxl_base_version} for non-SD3 components..")
+        pipe = DiffusionPipeline.from_pretrained(sdxl_base_version,
+                                            torch_dtype=torch.float16,
+                                            variant="fp16",
+                                            use_safetensors=True,
+                                            use_auth_token=True)
    else:
        pipe = DiffusionPipeline.from_pretrained(model_version,
                                            torch_dtype=torch.float16,
@@ -1314,7 +1533,10 @@ def main(args):
    # Convert models
    if args.convert_vae_decoder:
        logger.info("Converting vae_decoder")
-        convert_vae_decoder(pipe, args)
+        if args.sd3_version:
+            convert_vae_decoder_sd3(args)
+        else:
+            convert_vae_decoder(pipe, args)
        logger.info("Converted vae_decoder")

    if args.convert_vae_encoder:
@@ -1361,6 +1583,11 @@ def main(args):
        del pipe
        gc.collect()
        logger.info(f"Converted refiner")
+    
+    if args.convert_mmdit:
+        logger.info("Converting mmdit")
+        convert_mmdit(args)
+        logger.info("Converted mmdit")

    if args.quantize_nbits is not None:
        logger.info(f"Quantizing weights to {args.quantize_nbits}-bit precision")
@@ -1381,6 +1608,7 @@ def parser_spec():
    parser.add_argument("--convert-vae-decoder", action="store_true")
    parser.add_argument("--convert-vae-encoder", action="store_true")
    parser.add_argument("--convert-unet", action="store_true")
+    parser.add_argument("--convert-mmdit", action="store_true")
    parser.add_argument("--convert-safety-checker", action="store_true")
    parser.add_argument(
        "--convert-controlnet", 
@@ -1487,6 +1715,14 @@ def parser_spec():
        "If specified, enable unet to receive additional inputs from controlnet. "
        "Each input added to corresponding resnet output."
        )
+    parser.add_argument(
+        "--unet-batch-one",
+        action="store_true",
+        help=
+        "If specified, a batch size of one will be used for the unet, this is needed if you do not want to do "
+        "classifier free guidance. Default unet batch size is two, which is needed for classifier free guidance."
+        )
+    parser.add_argument("--include-t5", action="store_true")

    # Swift CLI Resource Bundling
    parser.add_argument(
@@ -1506,11 +1742,40 @@ def parser_spec():
        default=
        "https://huggingface.co/openai/clip-vit-base-patch32/resolve/main/merges.txt",
        help="The URL to the merged pairs used in by the text tokenizer.")
+    parser.add_argument(
+        "--text-encoder-t5-url",
+        default=
+        "https://huggingface.co/argmaxinc/coreml-stable-diffusion-3-medium/resolve/main/TextEncoderT5.mlmodelc",
+        help="The URL to the pre-converted T5 encoder model.")
+    parser.add_argument(
+        "--text-encoder-t5-config-url",
+        default=
+        "https://huggingface.co/google-t5/t5-small/resolve/main/tokenizer_config.json",
+        help="The URL to the merged pairs used in by the text tokenizer.")
+    parser.add_argument(
+        "--text-encoder-t5-data-url",
+        default=
+        "https://huggingface.co/google-t5/t5-small/resolve/main/tokenizer.json",
+        help="The URL to the merged pairs used in by the text tokenizer.")
+    parser.add_argument(
+        "--min-deployment-target",
+        default="macOS13",
+        help=(
+            "Minimum deployment target for Core ML models. "
+            "Valid options include macOS13, macOS14, macOS15, iOS16, iOS17, iOS18. "
+            "For iOS 18 compatibility with advanced quantization features, use iOS18. "
+            "Default is macOS13 for backwards compatibility."
+        )
+    )
    parser.add_argument(
        "--xl-version",
        action="store_true",
        help=("If specified, the pre-trained model will be treated as an instantiation of "
        "`diffusers.pipelines.StableDiffusionXLPipeline` instead of `diffusers.pipelines.StableDiffusionPipeline`"))
+    parser.add_argument(
+        "--sd3-version",
+        action="store_true",
+        help=("If specified, the pre-trained model will be treated as an SD3 model."))

    return parser

@@ -13,7 +13,6 @@ from enum import Enum

 import logging

-logger = logging.getLogger(__name__)
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)

@@ -43,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`
    """
@@ -66,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

@@ -96,17 +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:
-            attn = attention.original(q, k, v, mask, self.heads, self.dim_head)
-
-        elif ATTENTION_IMPLEMENTATION_IN_EFFECT == AttentionImplementations.SPLIT_EINSUM:
-            attn = attention.split_einsum(q, k, v, mask, self.heads, self.dim_head)
-
-        elif ATTENTION_IMPLEMENTATION_IN_EFFECT == AttentionImplementations.SPLIT_EINSUM_V2:
-            attn = attention.split_einsum_v2(q, k, v, mask, self.heads, self.dim_head)
-
-        else:
-            raise ValueError(ATTENTION_IMPLEMENTATION_IN_EFFECT)
+        attn = self.einsum(q, k, v, mask)

        return self.to_out(attn)

@@ -822,6 +829,7 @@ class UNet2DConditionModel(ModelMixin, ConfigMixin):
        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):
@@ -836,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

@@ -996,7 +1006,7 @@ 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]
@@ -1008,7 +1018,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
@@ -1100,7 +1110,7 @@ class UNet2DConditionModelXL(UNet2DConditionModel):

            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]
@@ -1112,7 +1122,7 @@ class UNet2DConditionModelXL(UNet2DConditionModel):
                                emb,
                                encoder_hidden_states=encoder_hidden_states)
        
-        if additional_residuals:
+        if self.support_controlnet:
            sample = sample + additional_residuals[-1]

        # 5. up
@@ -1,7 +1,8 @@
-coremltools>=7.0
-diffusers[torch]
+coremltools>=8.0
+diffusers[torch]==0.30.2
+diffusionkit==0.4.0
 torch
-transformers==4.29.2
+transformers==4.44.2
 scipy
 scikit-learn
 pytest
@@ -14,11 +14,11 @@ setup(
    long_description_content_type='text/markdown',
    author='Apple Inc.',
    install_requires=[
-        "coremltools>=7.0b2",
-        "diffusers[torch]",
+        "coremltools>=8.0",
+        "diffusers[torch]==0.30.2",
        "torch",
-        "transformers>=4.30.0",
-        "huggingface-hub",
+        "transformers==4.44.2",
+        "huggingface-hub==0.24.6",
        "scipy",
        "numpy<1.24",
        "pytest",
@@ -26,6 +26,7 @@ setup(
        "invisible-watermark",
        "safetensors",
        "matplotlib",
+        "diffusionkit==0.4.0",
    ],
    packages=find_packages(),
    classifiers=[
@@ -58,7 +58,7 @@ 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
@@ -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)]
@@ -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
+    }
+}
@@ -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
+        }
+    }
 }
@@ -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
+    }
+}
@@ -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
+        )
+    }
+}
@@ -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))))
+    }
+}
@@ -46,12 +46,16 @@ public struct PipelineConfiguration: Hashable {
    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
@@ -13,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
@@ -30,6 +32,7 @@ public enum PipelineError: String, Swift.Error {
    case startingImageProvidedWithoutEncoder
    case startingText2ImgWithoutTextEncoder
    case unsupportedOSVersion
+    case errorCreatingPreview
 }

@available(iOS 16.2, macOS 13.1, *)
@@ -207,28 +210,31 @@ public struct StableDiffusionPipeline: StableDiffusionPipelineProtocol {
        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 = useMultilingualTextEncoder ? concatEmbedding : 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, timeStepSpacing: config.schedulerTimestepSpacing)
+            case .discreteFlowScheduler: return DiscreteFlowScheduler(stepCount: config.stepCount, timeStepShift: config.schedulerTimestepShift)
            }
        }

@@ -258,8 +264,13 @@ public struct StableDiffusionPipeline: StableDiffusionPipelineProtocol {

            // 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
@@ -272,14 +283,42 @@ public struct StableDiffusionPipeline: StableDiffusionPipelineProtocol {
            
            // 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
@@ -176,6 +176,7 @@ public struct StableDiffusionXLPipeline: StableDiffusionPipelineProtocol {
            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)
            }
        }

@@ -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 }
+    }
+}
@@ -95,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 {
@@ -114,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
@@ -173,7 +178,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
@@ -196,51 +201,4 @@ public struct Unet: ResourceManaging {

        return noise
    }
-
-    func predictions(from batch: MLBatchProvider) throws -> MLBatchProvider {
-
-        var results = try models.first!.perform { model in
-            try model.predictions(fromBatch: batch)
-        }
-
-        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
-        }
-    }
 }
@@ -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)
+    }
+}
@@ -39,6 +39,9 @@ struct StableDiffusionSample: ParsableCommand {
    @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
    
@@ -114,6 +117,8 @@ struct StableDiffusionSample: ParsableCommand {
        log("(Note: This can take a while the first time using these resources)\n")
        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
@@ -128,6 +133,21 @@ struct StableDiffusionSample: ParsableCommand {
                    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,
@@ -198,6 +218,8 @@ struct StableDiffusionSample: ParsableCommand {
        pipelineConfig.useDenoisedIntermediates = true
        pipelineConfig.encoderScaleFactor = scaleFactor
        pipelineConfig.decoderScaleFactor = scaleFactor
+        pipelineConfig.decoderShiftFactor = shiftFactor
+        pipelineConfig.schedulerTimestepShift = timestepShift

        let images = try pipeline.generateImages(
            configuration: pipelineConfig) { progress in
Author	SHA1	Message	Date
Aseem Wadhwa	e12202c1f6	Merge pull request #368 from Aintor/main Fix for Mask Patch Failure and Quantization Issues in Latest `transformers` Versions	2025-07-03 11:01:02 -07:00
Kushvinth-Madhavan	e17915a9cb	Add --min-deployment-target argument for configurable iOS/macOS deployment targets - Added --min-deployment-target argument with default value macOS13 for backward compatibility - Supports macOS13, macOS14, macOS15, iOS16, iOS17, iOS18 targets - Graceful fallback for newer targets not available in older coremltools versions - Enables targeting iOS 18 for advanced quantization features while maintaining compatibility - Includes comprehensive help text explaining usage and compatibility requirements	2025-06-14 08:27:06 -07:00
Adrian (Mircea) Nenu	e5d960c41a	Update README.md for converting-models-to-coreml reference Fix typo that does not anchor the links correctly	2025-01-23 08:28:43 -08:00
NehalBhandari	f8e36141a9	Add support for mixed precision activation quantization for UNet model (#365 )	2024-11-15 16:05:22 -08:00
Jason Wong	50bb4c43d0	Update torch2coreml.py	2024-11-12 11:19:16 +08:00
Jason Wong	7e3df29361	Update torch2coreml.py	2024-11-12 10:36:17 +08:00
Jason Wong	4ef83f13d0	Update torch2coreml.py	2024-11-12 10:31:26 +08:00
Toby Roseman	200f255cf4	Merge pull request #367 from TobyRoseman/swift-serial-unet-predictions In Swift, allow serial uNet predictions	2024-11-11 09:01:34 -08:00
Jason Wong	244703046f	Update torch2coreml.py	2024-11-11 00:13:05 +08:00
Toby Roseman	bdbe459eaf	In Swift, allow serial unet predictions	2024-11-08 09:18:39 -08:00
Toby Roseman	bef26ae4ff	Merge pull request #364 from TobyRoseman/add-unet-sequential-option Add Option for Sequential Unet Predictions	2024-10-23 16:40:45 -07:00
Toby Roseman	7505f91208	Add option for sequential unet predictions	2024-10-23 14:23:19 -07:00
Alex Kennedy	6365319010	fix: Remove duplicate logger instantiation	2024-10-15 19:41:54 -07:00
Alejandro Isaza	6775c2c01a	Allow not using classifier free guidance for inference in Swift code	2024-10-08 12:06:39 -07:00
Toby Roseman	1c194dafb4	Merge pull request #360 from TobyRoseman/unet-batch-1-or-2 Allow not using classifier free guidance	2024-10-04 13:57:20 -07:00
Toby Roseman	e497107793	Allow not using classifier free guidance	2024-10-04 11:10:18 -07:00
Toby Roseman	cab24e9a51	Merge pull request #357 from TobyRoseman/dev Fix Cartoonish Output for non-XL model	2024-09-18 15:23:47 -07:00
Toby Roseman	3e04bcecf3	On macOS 15+, set fast prediction optimization hint for the unet.	2024-09-18 09:26:45 -07:00
Toby Roseman	32be4af031	Pin versions for important depedencies	2024-09-17 13:47:02 -07:00
Toby Roseman	aa0ec6e569	Do not force zeros embeddings for empty prompt, unless using XL	2024-09-17 13:21:33 -07:00
Toby Roseman	73a4fca864	Merge pull request #355 from junpeiz/patch-1 Add diffusionkit to installation dependency	2024-08-28 10:25:37 -07:00
junpeiz	5d5f155512	Add diffusionkit into setup.py Adds diffusionkit to package installation dependency to avoid erroring out when converting UNet.	2024-08-26 15:29:13 -07:00
Toby Roseman	70cfcce955	Merge pull request #354 from jakesabathia2/eng/PR-henry-refactor-chunk-script Refactor the model chunking script to work with all version of coremltools	2024-08-23 11:39:52 -07:00
HenryT	c6884bf3b5	add some logging	2024-08-23 11:31:46 -07:00
HenryT	3fc577cc41	fix typo	2024-08-22 16:20:24 -07:00
HenryT	f7b2618c0d	refactor the model chunking script to work with all version of coremltools	2024-08-22 14:24:46 -07:00
Aseem Wadhwa	cf16df8207	Merge pull request #348 from lamtrinhdev/Update_README.md Correct wording in README.md	2024-08-11 09:47:48 -07:00
LamTrinh.Dev	cf4e1e8d6e	Update README.md Correct wording.	2024-08-11 20:07:36 +07:00
Zach Nagengast	c891f4357b	Add SD3 Pipeline (#329 ) * Add SD3 Pipeline Co-authored-by: atiorh <atiorh@users.noreply.github.com> Co-authored-by: arda-argmax <arda-argmax@users.noreply.github.com> * Use swift-transformers for tokenization * Use diffusionkit converters in torch2coreml * Documentation and cleanup * Add model link * Consolidate batch prediction logic * Remove DecoderSD3.swift and consolidate logic into Decoder.swift * Remove DiffusionKit MLX inference reference from README --------- Co-authored-by: atiorh <atiorh@users.noreply.github.com> Co-authored-by: arda-argmax <arda-argmax@users.noreply.github.com> Co-authored-by: atila <atiorh@icloud.com>	2024-07-22 21:16:50 -07:00
msiracusa	5a170d29cf	Update pull_request_template.md after changing to MIT license	2024-05-03 18:03:21 -07:00
msiracusa	c5be2739d0	Update CONTRIBUTING.md to reflect the MIT license	2024-05-03 17:59:39 -07:00
msiracusa	5d8cdbb92d	Update LICENSE.md to MIT Changing repository license to MIT	2024-05-03 17:58:28 -07:00
Joshua Park	f7d2f5e9fb	Update swift-argument-parser to latest minor version This will also make ml-stable-diffusion more tolerant with regard to which swift-argument-parser version it needs, resulting in fewer version conflicts in Xcode projects.	2024-03-12 12:38:44 -07:00
Toby Roseman	d456a972cd	Merge pull request #311 from TobyRoseman/mbp-new-api Applying mixed bit compression using new optimize API	2024-01-23 16:42:12 -08:00
Toby Roseman	ec09fe064f	code review feedback	2024-01-23 16:41:35 -08:00
Toby Roseman	a9b84c658e	Update Documentation	2024-01-23 08:32:22 -08:00
Toby Roseman	9a3d463397	Applying mixed bit compression using new optimize API	2024-01-17 15:43:08 -08:00
Stuart Moore	7449ce46a4	Update CGImage+vImage.swift	2023-11-14 10:35:27 -08:00
Stuart Moore	5c62fab073	Use sRGB colorspace	2023-11-14 10:35:27 -08:00
Toby Roseman	3a0aeba5e5	Merge pull request #296 from TobyRoseman/compiled_models Compiled Models in Python	2023-11-07 13:18:01 -08:00
Toby Roseman	46c94e8a86	Link to Core ML documentation about why compiled models are faster	2023-11-07 12:46:13 -08:00
Toby Roseman	4c54204bb0	Compiled Models in Python	2023-11-07 12:13:31 -08:00
Toby Roseman	bea04420b5	Merge pull request #287 from TobyRoseman/minor-fix Fix pipeline.py for Python 3.8	2023-10-24 08:26:53 -07:00
Toby Roseman	d159b9254c	Fix pipeline.py for Python 3.8	2023-10-23 10:27:56 -07:00
Atila Orhon	567dbf2757	fixes #281	2023-10-20 22:45:43 -07:00
Jiangdi	f4466f0b73	fix error (#279 )	2023-10-10 12:15:08 -07:00
Zach Nagengast	7fe1bf23b9	Cast safety_checker to float32 (#276 )	2023-10-06 20:54:39 -07:00
Atila Orhon	11ba3a0965	Add custom_vae_version in MBP pre-analysis	2023-09-30 19:44:01 -07:00
Atila Orhon	96c044e07b	Add iPhone 12 Pro Benchmark	2023-09-28 14:34:09 -07:00