Compare commits
10 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 Atila Orhon
|
7f9c58a6c1 | ||
|
Zsolt Váradi
|
e3875a584b | ||
|
Oleg Ponomaryov
|
ea904e31d7 | ||
|
Alejandro Isaza
|
940dba02ee | ||
|
Atila Orhon
|
8fd725764f | ||
|
Atila Orhon
|
efea9d3812 | ||
|
Atila Orhon
|
ef5aa89c45 | ||
|
Pedro Cuenca
|
449a16d96e | ||
|
Pedro Cuenca
|
41174b937b | ||
|
Pedro Cuenca
|
4aafa737ba |
+1
-1
@@ -1,4 +1,4 @@
|
||||
// swift-tools-version: 5.7
|
||||
// swift-tools-version: 5.8
|
||||
// The swift-tools-version declares the minimum version of Swift required to build this package.
|
||||
|
||||
import PackageDescription
|
||||
|
||||
@@ -2,7 +2,8 @@
|
||||
|
||||
Run Stable Diffusion on Apple Silicon with Core ML
|
||||
|
||||
<img src="assets/readme_reel.png">
|
||||
[\[Blog Post\]](https://machinelearning.apple.com/research/stable-diffusion-coreml-apple-silicon) [\[BibTeX\]](#bibtex)
|
||||
|
||||
|
||||
This repository comprises:
|
||||
|
||||
@@ -11,27 +12,142 @@ This repository comprises:
|
||||
|
||||
If you run into issues during installation or runtime, please refer to the [FAQ](#faq) section. Please refer to the [System Requirements](#system-requirements) section before getting started.
|
||||
|
||||
|
||||
## <a name="example-results"></a> Example Results
|
||||
|
||||
There are numerous versions of Stable Diffusion available on the [Hugging Face Hub](https://huggingface.co/models?search=stable-diffusion). Here are example results from three of those models:
|
||||
|
||||
`--model-version` | [stabilityai/stable-diffusion-2-base](https://huggingface.co/stabilityai/stable-diffusion-2-base) | [CompVis/stable-diffusion-v1-4](https://huggingface.co/CompVis/stable-diffusion-v1-4) | [runwayml/stable-diffusion-v1-5](https://huggingface.co/runwayml/stable-diffusion-v1-5) |
|
||||
:------:|:------:|:------:|:------:
|
||||
Output |  |  | 
|
||||
M1 iPad Pro 8GB Latency (s) | 29 | 38 | 38 |
|
||||
M1 MacBook Pro 16GB Latency (s) | 24 | 35 | 35 |
|
||||
M2 MacBook Air 8GB Latency (s) | 18 | 23 | 23 |
|
||||
|
||||
Please see [Important Notes on Performance Benchmarks](#important-notes-on-performance-benchmarks) section for details.
|
||||
<img src="assets/readme_reel.png">
|
||||
|
||||
## <a name="system-requirements"></a> System Requirements
|
||||
|
||||
The following is recommended to use all the functionality in this repository:
|
||||
Model Conversion:
|
||||
|
||||
Python | macOS | Xcode | iPadOS, iOS |
|
||||
:------:|:------:|:------:|:------:|
|
||||
3.8 | 13.1 | 14.2 | 16.2 |
|
||||
macOS | Python | coremltools |
|
||||
:------:|:------:|:-----------:|
|
||||
13.1 | 3.8 | 7.0 |
|
||||
|
||||
Project Build:
|
||||
|
||||
macOS | Xcode | Swift |
|
||||
:------:|:-----:|:-----:|
|
||||
14.0 | 15.0 | 5.8 |
|
||||
|
||||
Target Device Runtime:
|
||||
|
||||
macOS | iPadOS, iOS |
|
||||
:------:|:-----------:|
|
||||
13.1 | 16.2 |
|
||||
|
||||
Target Device Runtime ([With Memory Improvements](#compression)):
|
||||
|
||||
macOS | iPadOS, iOS |
|
||||
:------:|:-----------:|
|
||||
14.0 | 17.0 |
|
||||
|
||||
Target Device Hardware Generation:
|
||||
|
||||
Mac | iPad | iPhone |
|
||||
:------:|:-------:|:-------:|
|
||||
M1 | M1 | A14 |
|
||||
|
||||
|
||||
## <a name="performance-benchmark"></a> Performance Benchmarks
|
||||
|
||||
|
||||
[`stabilityai/stable-diffusion-2-1-base`](https://huggingface.co/apple/coreml-stable-diffusion-2-1-base) Benchmark:
|
||||
|
||||
|
||||
| Device | `--compute-unit`| `--attention-implementation` | End-to-End Latency (s) | Diffusion Speed (iter/s) |
|
||||
| --------------------- | --------------- | ---------------------------- | ---------------------- | ------------------------ |
|
||||
| iPhone 12 Mini | `CPU_AND_NE` | `SPLIT_EINSUM_V2` | 20 | 1.3 |
|
||||
| iPhone 12 Pro Max | `CPU_AND_NE` | `SPLIT_EINSUM_V2` | 17 | 1.4 |
|
||||
| iPhone 13 | `CPU_AND_NE` | `SPLIT_EINSUM_V2` | 15 | 1.7 |
|
||||
| iPhone 13 Pro Max | `CPU_AND_NE` | `SPLIT_EINSUM_V2` | 12 | 1.8 |
|
||||
| iPhone 14 | `CPU_AND_NE` | `SPLIT_EINSUM_V2` | 13 | 1.8 |
|
||||
| iPhone 14 Pro Max | `CPU_AND_NE` | `SPLIT_EINSUM_V2` | 9 | 2.3 |
|
||||
| iPad Pro (M1) | `CPU_AND_NE` | `SPLIT_EINSUM_V2` | 11 | 2.1 |
|
||||
| iPad Pro (M2) | `CPU_AND_NE` | `SPLIT_EINSUM_V2` | 8 | 2.9 |
|
||||
| Mac Studio (M1 Ultra) | `CPU_AND_GPU` | `ORIGINAL` | 4 | 6.3 |
|
||||
| Mac Studio (M2 Ultra) | `CPU_AND_GPU` | `ORIGINAL` | 3 | 7.6 |
|
||||
|
||||
|
||||
<details>
|
||||
<summary> Details (Click to expand) </summary>
|
||||
|
||||
- This benchmark was conducted by Apple using public beta versions of iOS 17.0, iPadOS 17.0 and macOS 14.0 in June 2023.
|
||||
- The performance data was collected by running the `StableDiffusion` Swift pipeline.
|
||||
- Swift code is not fully optimized, introducing up to ~10% overhead unrelated to Core ML model execution.
|
||||
- The median latency value across 3 back-to-back end-to-end executions are reported
|
||||
- The image generation procedure follows the standard configuration: 20 inference steps, 512x512 output image resolution, 77 text token sequence length, classifier-free guidance (batch size of 2 for unet).
|
||||
- The actual prompt length does not impact performance because the Core ML model is converted with a static shape that computes the forward pass for all of the 77 elements (`tokenizer.model_max_length`) in the text token sequence regardless of the actual length of the input text.
|
||||
- Weights are compressed to 6 bit precision. Please refer to [this section](#compression) for details.
|
||||
- Activations are in float16 precision for both the GPU and the Neural Engine.
|
||||
- For iPhone 12 mini and iPhone 13, we enable the [reduceMemory](https://github.com/apple/ml-stable-diffusion/blob/main/swift/StableDiffusion/pipeline/StableDiffusionPipeline.swift#L65) option to load and unload models just-in-time to avoid memory issues. This added up to 2 seconds to the end-to-end latency.
|
||||
- In the benchmark table, we report the best performing `--compute-unit` and `--attention-implementation` values per device. The former does not modify the Core ML model and can be applied during runtime. The latter modifies the Core ML model. Note that the best performing compute unit is model version and hardware-specific.
|
||||
- Note that the performance optimizations in this repository (e.g. `--attention-implementation`) are generally applicable to Transformers and not customized to Stable Diffusion. Better performance may be observed upon custom kernel tuning. Therefore, these numbers do not represent **peak** HW capability.
|
||||
- Performance may vary across different versions of Stable Diffusion due to architecture changes in the model itself. Each reported number is specific to the model version mentioned in that context.
|
||||
- Performance may vary due to factors like increased system load from other applications or suboptimal device thermal state. Given these factors, we do not report sub-second variance in latency.
|
||||
|
||||
</details>
|
||||
|
||||
## <a name="compression"></a> Weight Compression
|
||||
|
||||
coremltools-7.0 supports advanced weight compression techniques for [pruning](https://coremltools.readme.io/v7.0/docs/pruning), [palettization](https://coremltools.readme.io/v7.0/docs/palettization-overview) and [linear 8-bit quantization](https://coremltools.readme.io/v7.0/docs/quantization-aware-training). For these techniques, `coremltools.optimize.torch.*` includes APIs that require fine-tuning to maintain accuracy at higher compression rates whereas `coremltools.optimize.coreml.*` includes APIs that are applied post-training and are data-free.
|
||||
|
||||
We demonstrate how data-free [post-training palettization](https://coremltools.readme.io/v7.0/docs/post-training-palettization) implemented in `coremltools.optimize.coreml.palettize_weights` enables us to achieve greatly improved performance for Stable Diffusion on mobile devices. This API implements the [Fast Exact k-Means](https://arxiv.org/abs/1701.07204) algorithm for optimal weight clustering which yields more accurate palettes. Using `--quantize-nbits {2,4,6,8}` during [conversion](#converting-models-to-coreml) is going to apply this compression to the unet and text_encoder models.
|
||||
|
||||
For best results, we recommend [training-time palettization](https://coremltools.readme.io/v7.0/docs/training-time-palettization): `coremltools.optimize.torch.palettization.DKMPalettizer` if fine-tuning your model is feasible. This API implements the [Differentiable k-Means (DKM)](https://machinelearning.apple.com/research/differentiable-k-means) learned palettization algorithm. In this exercise, we stick to post-training palettization for the sake of simplicity and ease of reproducibility.
|
||||
|
||||
The Neural Engine is capable of accelerating models with low-bit palettization: 2, 4, 6 or 8 bits. With iOS 17 and macOS 14, compressed weights for Core ML models can be just-in-time decompressed during runtime (as opposed to ahead-of-time decompression upon load) to match the precision of activation tensors. This yields significant memory savings and enables models to run on devices with smaller RAM (e.g. iPhone 12 Mini). In addition, compressed weights are faster to fetch from memory which reduces the latency of memory bandwidth-bound layers. The just-in-time decompression behavior depends on the compute unit, layer type and hardware generation.
|
||||
|
||||
| Weight Precision | `--compute-unit` | [`stabilityai/stable-diffusion-2-1-base`](https://huggingface.co/apple/coreml-stable-diffusion-2-1-base) generating *"a high quality photo of a surfing dog"* |
|
||||
| :---------------:| :----------------: | ------------------------------------------------------ |
|
||||
| 6-bit | cpuAndNeuralEngine | <img src="assets/palette6_cpuandne_readmereel.png"> |
|
||||
| 16-bit | cpuAndNeuralEngine | <img src="assets/float16_cpuandne_readmereel.png"> |
|
||||
| 16-bit | cpuAndGPU | <img src="assets/float16_gpu_readmereel.png"> |
|
||||
|
||||
Note that there are minor differences across 16-bit (float16) and 6-bit results. These differences are comparable to the differences across float16 and float32 or differences across compute units as exemplified above. We recommend a minimum of 6 bits for palettizing Stable Diffusion. Smaller number of bits (2 and 4) will require fine-tuning to recover image generation quality as previously mentioned.
|
||||
|
||||
Resources:
|
||||
- [Core ML Tools Docs: Optimizing Models](https://coremltools.readme.io/v7.0/docs/optimizing-models)
|
||||
- [WWDC23 Session Video: Use Core ML Tools for machine learning model compression](https://developer.apple.com/videos/play/wwdc2023/10047)
|
||||
|
||||
## <a name="using-controlnet"></a> Using ControlNet
|
||||
|
||||
Example results using the prompt *"a high quality photo of a surfing dog"* conditioned on the scribble (leftmost):
|
||||
|
||||
<img src="assets/controlnet_readme_reel.png">
|
||||
|
||||
[ControlNet](https://huggingface.co/lllyasviel/ControlNet) allows users to condition image generation with Stable Diffusion on signals such as edge maps, depth maps, segmentation maps, scribbles and pose. Thanks to [@ryu38's contribution](https://github.com/apple/ml-stable-diffusion/pull/153), both the Python CLI and the Swift package support ControlNet models. Please refer to [this section](#converting-models-to-coreml) for details on setting up Stable Diffusion with ControlNet.
|
||||
|
||||
|
||||
## <a name="system-multilingual-text-encoder"></a> Using the System Multilingual Text Encoder
|
||||
|
||||
<details>
|
||||
<summary> Details (Click to expand) </summary>
|
||||
|
||||
With iOS 17 and macOS 14, `NaturalLanguage` framework introduced the [NLContextualEmbedding](https://developer.apple.com/documentation/naturallanguage/nlcontextualembedding) which provides Transformer-based textual embeddings for Latin (20 languages), Cyrillic (4 langauges) and CJK (3 languages) scripts. The WWDC23 session titled [Explore Natural Language multilingual models](https://developer.apple.com/videos/play/wwdc2023/10042) demonstrated how this powerful new model can be used by developers to train downstream tasks such as multilingual image generation with Stable Diffusion.
|
||||
|
||||
The code to reproduce this demo workflow is made available in this repository. There are several ways in which this workflow can be implemented. Here is an example:
|
||||
|
||||
**Step 1:** Curate an image-text dataset with the desired languages.
|
||||
|
||||
**Step 2:** Pre-compute the NLContextualEmbedding values and replace the text strings with these embedding vectors in your dataset.
|
||||
|
||||
**Step 3:** Fine-tune a base model from Hugging Face Hub that is compatible with the [StableDiffusionPipeline](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/overview) by using your new dataset and replacing the default text_encoder with your pre-computed NLContextualEmbedding values.
|
||||
|
||||
**Step 4:** In order to be able to swap the text_encoder of a base model without training new layers, the base model's `text_encoder.hidden_size` must match that of NLContextualEmbedding. If it doesn't, you will need to train a linear projection layer to map between the two dimensionalities. After fine-tuning, this linear layer should be converted to CoreML as follows:
|
||||
|
||||
```shell
|
||||
python -m python_coreml_stable_diffusion.multilingual_projection --input-path <path-to-projection-torchscript> --output-dir <output-dir>
|
||||
```
|
||||
|
||||
The command above will yield a `MultilingualTextEncoderProjection.mlmodelc` file under `--output-dir` and this should be colocated with the rest of the Core ML model assets that were generated through `--bundle-resources-for-swift-cli`.
|
||||
|
||||
**Step 5:** The multilingual system text encoder can now be invoked by setting `useMultilingualTextEncoder` to true when initializing a pipeline or setting `--use-multilingual-text-encoder` in the CLI. Note that the model assets are distributed over-the-air so the first invocation will trigger asset downloads which is less than 100MB.
|
||||
|
||||
|
||||
Resources:
|
||||
- [WWDC23 Session Video: Explore Natural Language multilingual models](https://developer.apple.com/videos/play/wwdc2023/10042)
|
||||
- [NLContextualEmbedding API Documentation](https://developer.apple.com/documentation/naturallanguage/nlcontextualembedding)
|
||||
|
||||
</details>
|
||||
|
||||
## <a name="using-converted-weights"></a> Using Ready-made Core ML Models from Hugging Face Hub
|
||||
|
||||
@@ -43,6 +159,7 @@ Python | macOS | Xcode | iPadOS, iOS |
|
||||
* [`CompVis/stable-diffusion-v1-4`](https://huggingface.co/apple/coreml-stable-diffusion-v1-4)
|
||||
* [`runwayml/stable-diffusion-v1-5`](https://huggingface.co/apple/coreml-stable-diffusion-v1-5)
|
||||
* [`stabilityai/stable-diffusion-2-base`](https://huggingface.co/apple/coreml-stable-diffusion-2-base)
|
||||
* [`stabilityai/stable-diffusion-2-1-base`](https://huggingface.co/apple/coreml-stable-diffusion-2-1-base)
|
||||
|
||||
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).
|
||||
|
||||
@@ -70,29 +187,13 @@ If you prefer to download specific variants instead of cloning the repos, you ca
|
||||
|
||||
```Python
|
||||
from huggingface_hub import snapshot_download
|
||||
from huggingface_hub.file_download import repo_folder_name
|
||||
from pathlib import Path
|
||||
import shutil
|
||||
|
||||
repo_id = "apple/coreml-stable-diffusion-v1-4"
|
||||
variant = "original/packages"
|
||||
|
||||
def download_model(repo_id, variant, output_dir):
|
||||
destination = Path(output_dir) / (repo_id.split("/")[-1] + "_" + variant.replace("/", "_"))
|
||||
if destination.exists():
|
||||
raise Exception(f"Model already exists at {destination}")
|
||||
|
||||
# Download and copy without symlinks
|
||||
downloaded = snapshot_download(repo_id, allow_patterns=f"{variant}/*", cache_dir=output_dir)
|
||||
downloaded_bundle = Path(downloaded) / variant
|
||||
shutil.copytree(downloaded_bundle, destination)
|
||||
|
||||
# Remove all downloaded files
|
||||
cache_folder = Path(output_dir) / repo_folder_name(repo_id=repo_id, repo_type="model")
|
||||
shutil.rmtree(cache_folder)
|
||||
return destination
|
||||
|
||||
model_path = download_model(repo_id, variant, output_dir="./models")
|
||||
model_path = Path("./models") / (repo_id.split("/")[-1] + "_" + variant.replace("/", "_"))
|
||||
snapshot_download(repo_id, allow_patterns=f"{variant}/*", local_dir=model_path, local_dir_use_symlinks=False)
|
||||
print(f"Model downloaded at {model_path}")
|
||||
```
|
||||
|
||||
@@ -133,9 +234,11 @@ This generally takes 15-20 minutes on an M1 MacBook Pro. Upon successful executi
|
||||
|
||||
- `--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.
|
||||
|
||||
- `--chunk-unet`: Splits the Unet model in two approximately equal chunks (each with less than 1GB of weights) for mobile-friendly deployment. This is **required** for Neural Engine deployment on iOS and iPadOS. This is not required for macOS. Swift CLI is able to consume both the chunked and regular versions of the Unet model but prioritizes the former. Note that chunked unet is not compatible with the Python pipeline because Python pipeline is intended for macOS only. Chunking is for on-device deployment with Swift only.
|
||||
- `--quantize-nbits`: Quantizes the weights of unet and text_encoder models down to 2, 4, 6 or 8 bits using a globally optimal k-means clustering algorithm. By default all models are weight-quantized to 16 bits even if this argument is not specified. Please refer to [this section](#compression for details and further guidance on weight compression.
|
||||
|
||||
- `--attention-implementation`: Defaults to `SPLIT_EINSUM` which is the implementation described in [Deploying Transformers on the Apple Neural Engine](https://machinelearning.apple.com/research/neural-engine-transformers). `--attention-implementation ORIGINAL` will switch to an alternative that should be used for CPU or GPU deployment. Please refer to the [Performance Benchmark](#performance-benchmark) section for further guidance.
|
||||
- `--chunk-unet`: Splits the Unet model in two approximately equal chunks (each with less than 1GB of weights) for mobile-friendly deployment. This is **required** for Neural Engine deployment on iOS and iPadOS if weights are not quantized to 6-bits or less (`--quantize-nbits {2,4,6}`). This is not required for macOS. Swift CLI is able to consume both the chunked and regular versions of the Unet model but prioritizes the former. Note that chunked unet is not compatible with the Python pipeline because Python pipeline is intended for macOS only.
|
||||
|
||||
- `--attention-implementation`: Defaults to `SPLIT_EINSUM` which is the implementation described in [Deploying Transformers on the Apple Neural Engine](https://machinelearning.apple.com/research/neural-engine-transformers). `--attention-implementation SPLIT_EINSUM_V2` yields 10-30% improvement for mobile devices, still targeting the Neural Engine. `--attention-implementation ORIGINAL` will switch to an alternative implementation that should be used for CPU or GPU deployment on some Mac devices. Please refer to the [Performance Benchmark](#performance-benchmark) section for further guidance.
|
||||
|
||||
- `--check-output-correctness`: Compares original PyTorch model's outputs to final Core ML model's outputs. This flag increases RAM consumption significantly so it is recommended only for debugging purposes.
|
||||
|
||||
@@ -171,26 +274,6 @@ Please refer to the help menu for all available arguments: `python -m python_cor
|
||||
<details>
|
||||
<summary> Click to expand </summary>
|
||||
|
||||
### <a name="swift-requirements"></a> System Requirements
|
||||
|
||||
**Building** (recommended):
|
||||
|
||||
- Xcode 14.2
|
||||
- Command Line Tools for Xcode 14.2
|
||||
|
||||
Check [developer.apple.com](https://developer.apple.com/download/all/?q=xcode) for the latest versions.
|
||||
|
||||
**Running** (minimum):
|
||||
|
||||
| Mac | iPad\* | iPhone\* |
|
||||
|:------:|:------:|:------:|
|
||||
| macOS 13.1 | iPadOS 16.2 | iOS 16.2 |
|
||||
| M1 | M1 | iPhone 12 Pro |
|
||||
|
||||
You will also need the resources generated by the `--bundle-resources-for-swift-cli` option described in [Converting Models to Core ML](#converting-models-to-coreml)
|
||||
|
||||
\* Please see [FAQ](#faq) [Q6](#q-mobile-app) regarding deploying on iPad and iPhone.
|
||||
|
||||
### Example CLI Usage
|
||||
```shell
|
||||
swift run StableDiffusionSample "a photo of an astronaut riding a horse on mars" --resource-path <output-mlpackages-directory>/Resources/ --seed 93 --output-path </path/to/output/image>
|
||||
@@ -257,78 +340,6 @@ Hugging Face has made the app [available in the Mac App Store](https://apps.appl
|
||||
|
||||
</details>
|
||||
|
||||
## <a name="performance-benchmark"></a> Performance Benchmark
|
||||
|
||||
<details>
|
||||
<summary> Click to expand </summary>
|
||||
|
||||
Standard [CompVis/stable-diffusion-v1-4](https://huggingface.co/CompVis/stable-diffusion-v1-4) Benchmark
|
||||
|
||||
| Device | `--compute-unit`| `--attention-implementation` | Latency (seconds) |
|
||||
| ---------------------------------- | -------------- | ---------------------------- | ----------------- |
|
||||
| Mac Studio (M1 Ultra, 64-core GPU) | `CPU_AND_GPU` | `ORIGINAL` | 9 |
|
||||
| Mac Studio (M1 Ultra, 48-core GPU) | `CPU_AND_GPU` | `ORIGINAL` | 13 |
|
||||
| MacBook Pro (M1 Max, 32-core GPU) | `CPU_AND_GPU` | `ORIGINAL` | 18 |
|
||||
| MacBook Pro (M1 Max, 24-core GPU) | `CPU_AND_GPU` | `ORIGINAL` | 20 |
|
||||
| MacBook Pro (M1 Pro, 16-core GPU) | `ALL` | `SPLIT_EINSUM (default)` | 26 |
|
||||
| MacBook Pro (M2) | `CPU_AND_NE` | `SPLIT_EINSUM (default)` | 23 |
|
||||
| MacBook Pro (M1) | `CPU_AND_NE` | `SPLIT_EINSUM (default)` | 35 |
|
||||
| iPad Pro (5th gen, M1) | `CPU_AND_NE` | `SPLIT_EINSUM (default)` | 38 |
|
||||
|
||||
|
||||
Please see [Important Notes on Performance Benchmarks](#important-notes-on-performance-benchmarks) section for details.
|
||||
|
||||
</details>
|
||||
|
||||
## <a name="important-notes-on-performance-benchmarks"></a> Important Notes on Performance Benchmarks
|
||||
|
||||
<details>
|
||||
<summary> Click to expand </summary>
|
||||
|
||||
- This benchmark was conducted by Apple using public beta versions of iOS 16.2, iPadOS 16.2 and macOS 13.1 in November 2022.
|
||||
- The executed program is `python_coreml_stable_diffusion.pipeline` for macOS devices and a minimal Swift test app built on the `StableDiffusion` Swift package for iOS and iPadOS devices.
|
||||
- The median value across 3 end-to-end executions is reported.
|
||||
- Performance may materially differ across different versions of Stable Diffusion due to architecture changes in the model itself. Each reported number is specific to the model version mentioned in that context.
|
||||
- The image generation procedure follows the standard configuration: 50 inference steps, 512x512 output image resolution, 77 text token sequence length, classifier-free guidance (batch size of 2 for unet).
|
||||
- The actual prompt length does not impact performance because the Core ML model is converted with a static shape that computes the forward pass for all of the 77 elements (`tokenizer.model_max_length`) in the text token sequence regardless of the actual length of the input text.
|
||||
- Pipelining across the 4 models is not optimized and these performance numbers are subject to variance under increased system load from other applications. Given these factors, we do not report sub-second variance in latency.
|
||||
- Weights and activations are in float16 precision for both the GPU and the Neural Engine.
|
||||
- The Swift CLI program consumes a peak memory of approximately 2.6GB (without the safety checker), 2.1GB of which is model weights in float16 precision. We applied [8-bit weight quantization](https://coremltools.readme.io/docs/compressing-ml-program-weights#use-affine-quantization) to reduce peak memory consumption by approximately 1GB. However, we observed that it had an adverse effect on generated image quality and we rolled it back. We encourage developers to experiment with other advanced weight compression techniques such as [palettization](https://coremltools.readme.io/docs/compressing-ml-program-weights#use-a-lookup-table) and/or [pruning](https://coremltools.readme.io/docs/compressing-ml-program-weights#use-sparse-representation) which may yield better results.
|
||||
- In the [benchmark table](performance-benchmark), we report the best performing `--compute-unit` and `--attention-implementation` values per device. The former does not modify the Core ML model and can be applied during runtime. The latter modifies the Core ML model. Note that the best performing compute unit is model version and hardware-specific.
|
||||
|
||||
</details>
|
||||
|
||||
|
||||
## <a name="results-with-different-compute-units"></a> Results with Different Compute Units
|
||||
|
||||
<details>
|
||||
<summary> Click to expand </summary>
|
||||
|
||||
It is highly probable that there will be slight differences across generated images using different compute units.
|
||||
|
||||
The following images were generated on an M1 MacBook Pro and macOS 13.1 with the prompt *"a photo of an astronaut riding a horse on mars"* using the [runwayml/stable-diffusion-v1-5](https://huggingface.co/runwayml/stable-diffusion-v1-5) model version. The random seed was set to 93:
|
||||
|
||||
CPU_AND_NE | CPU_AND_GPU | ALL |
|
||||
:------------:|:-------------:|:------:
|
||||
 |  |  |
|
||||
|
||||
Differences may be less or more pronounced for different inputs. Please see the [FAQ](#faq) Q8 for a detailed explanation.
|
||||
|
||||
</details>
|
||||
|
||||
## <a name="results-with-controlnet"></a> Results with ControlNet
|
||||
|
||||
<details>
|
||||
<summary> Click to expand </summary>
|
||||
|
||||
[ControlNet](https://huggingface.co/lllyasviel/ControlNet) allows users to condition image generation with Stable Diffusion on signals such as edge maps, depth maps, segmentation maps, scribbles and pose. Thanks to [@ryu38's contribution](https://github.com/apple/ml-stable-diffusion/pull/153), both the Python CLI and the Swift package support ControlNet models. Please refer to CLI arguments in previous sections to exercise this new feature.
|
||||
|
||||
Example results using the prompt "a high quality photo of a surfing dog" conditioned on the scribble (leftmost):
|
||||
|
||||
<img src="assets/controlnet_readme_reel.png">
|
||||
|
||||
</details>
|
||||
|
||||
|
||||
## <a name="faq"></a> FAQ
|
||||
|
||||
@@ -396,6 +407,9 @@ The image generation process in `StableDiffusion` can yield over 2 GB of peak me
|
||||
If your app crashes during image generation, consider adding the [Increased Memory Limit](https://developer.apple.com/documentation/bundleresources/entitlements/com_apple_developer_kernel_increased-memory-limit) capability to inform the system that some of your app’s core features may perform better by exceeding the default app memory limit on supported devices.
|
||||
|
||||
On iOS, depending on the iPhone model, Stable Diffusion model versions, selected compute units, system load and design of your app, this may still not be sufficient to keep your apps peak memory under the limit. Please remember, because the device shares memory between apps and iOS processes, one app using too much memory can compromise the user experience across the whole device.
|
||||
|
||||
We **strongly recommend** quantizing your models with `--quantize-nbits 6` for iOS deployment. This reduces the peak RAM usage by 1GB or more depending on the model version and robustly enables inference even on iPhone 12 Mini.
|
||||
|
||||
</details>
|
||||
|
||||
<details>
|
||||
@@ -420,7 +434,7 @@ On iOS, depending on the iPhone model, Stable Diffusion model versions, selected
|
||||
|
||||
<b> 3. Model Function Drift During Conversion </b>
|
||||
|
||||
The difference in outputs across corresponding PyTorch and Core ML models is a potential cause. The signal integrity is tested during the conversion process (enabled via `--check-output-correctness` argument to `python_coreml_stable_diffusion.torch2coreml`) and it is verified to be above a minimum [PSNR](https://en.wikipedia.org/wiki/Peak_signal-to-noise_ratio) value as tested on random inputs. Note that this is simply a sanity check and does not guarantee this minimum PSNR across all possible inputs. Furthermore, the results are not guaranteed to be identical when executing the same Core ML models across different compute units. This is not expected to be a major source of difference as the sample visual results indicate in [this section](#results-with-different-compute-units).
|
||||
The difference in outputs across corresponding PyTorch and Core ML models is a potential cause. The signal integrity is tested during the conversion process (enabled via `--check-output-correctness` argument to `python_coreml_stable_diffusion.torch2coreml`) and it is verified to be above a minimum [PSNR](https://en.wikipedia.org/wiki/Peak_signal-to-noise_ratio) value as tested on random inputs. Note that this is simply a sanity check and does not guarantee this minimum PSNR across all possible inputs. Furthermore, the results are not guaranteed to be identical when executing the same Core ML models across different compute units. This is not expected to be a major source of difference as the sample visual results indicate in [this section](#compression).
|
||||
|
||||
<b> 4. Weights and Activations Data Type </b>
|
||||
|
||||
@@ -458,3 +472,14 @@ On iOS, depending on the iPhone model, Stable Diffusion model versions, selected
|
||||
</details>
|
||||
|
||||
</details>
|
||||
|
||||
## <a name="bibtex"></a> BibTeX Reference
|
||||
|
||||
```latex
|
||||
@misc{stable-diffusion-coreml-apple-silicon,
|
||||
title = {Stable Diffusion with Core ML on Apple Silicon},
|
||||
author = {Atila Orhon and Michael Siracusa and Aseem Wadhwa},
|
||||
year = {2022},
|
||||
URL = {null}
|
||||
}
|
||||
```
|
||||
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 2.2 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.2 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.2 MiB |
@@ -1 +1 @@
|
||||
__version__ = "0.1.0"
|
||||
__version__ = "1.0.0"
|
||||
|
||||
@@ -0,0 +1,146 @@
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
logger.setLevel(logging.INFO)
|
||||
|
||||
import torch
|
||||
|
||||
def split_einsum(q, k, v, mask, heads, dim_head):
|
||||
""" Attention Implementation backing AttentionImplementations.SPLIT_EINSUM
|
||||
|
||||
- Implements https://machinelearning.apple.com/research/neural-engine-transformers
|
||||
- Recommended for ANE
|
||||
- Marginally slower on GPU
|
||||
"""
|
||||
mh_q = [
|
||||
q[:, head_idx * dim_head:(head_idx + 1) *
|
||||
dim_head, :, :] for head_idx in range(heads)
|
||||
] # (bs, dim_head, 1, max_seq_length) * heads
|
||||
|
||||
k = k.transpose(1, 3)
|
||||
mh_k = [
|
||||
k[:, :, :,
|
||||
head_idx * dim_head:(head_idx + 1) * dim_head]
|
||||
for head_idx in range(heads)
|
||||
] # (bs, max_seq_length, 1, dim_head) * heads
|
||||
|
||||
mh_v = [
|
||||
v[:, head_idx * dim_head:(head_idx + 1) *
|
||||
dim_head, :, :] for head_idx in range(heads)
|
||||
] # (bs, dim_head, 1, max_seq_length) * heads
|
||||
|
||||
attn_weights = [
|
||||
torch.einsum("bchq,bkhc->bkhq", [qi, ki]) * (dim_head**-0.5)
|
||||
for qi, ki in zip(mh_q, mh_k)
|
||||
] # (bs, max_seq_length, 1, max_seq_length) * heads
|
||||
|
||||
if mask is not None:
|
||||
for head_idx in range(heads):
|
||||
attn_weights[head_idx] = attn_weights[head_idx] + mask
|
||||
|
||||
attn_weights = [
|
||||
aw.softmax(dim=1) for aw in attn_weights
|
||||
] # (bs, max_seq_length, 1, max_seq_length) * heads
|
||||
attn = [
|
||||
torch.einsum("bkhq,bchk->bchq", wi, vi)
|
||||
for wi, vi in zip(attn_weights, mh_v)
|
||||
] # (bs, dim_head, 1, max_seq_length) * heads
|
||||
|
||||
attn = torch.cat(attn, dim=1) # (bs, dim, 1, max_seq_length)
|
||||
return attn
|
||||
|
||||
|
||||
CHUNK_SIZE = 512
|
||||
|
||||
def split_einsum_v2(q, k, v, mask, heads, dim_head):
|
||||
""" Attention Implementation backing AttentionImplementations.SPLIT_EINSUM_V2
|
||||
|
||||
- Implements https://machinelearning.apple.com/research/neural-engine-transformers
|
||||
- Recommended for ANE
|
||||
- Marginally slower on GPU
|
||||
- Chunks the query sequence to avoid large intermediate tensors and improves ANE performance
|
||||
"""
|
||||
query_seq_length = q.size(3)
|
||||
num_chunks = query_seq_length // CHUNK_SIZE
|
||||
|
||||
if num_chunks == 0:
|
||||
logger.info(
|
||||
"AttentionImplementations.SPLIT_EINSUM_V2: query sequence too short to chunk "
|
||||
f"({query_seq_length}<{CHUNK_SIZE}), fall back to AttentionImplementations.SPLIT_EINSUM (safe to ignore)")
|
||||
return split_einsum(q, k, v, mask, heads, dim_head)
|
||||
|
||||
logger.info(
|
||||
"AttentionImplementations.SPLIT_EINSUM_V2: Splitting query sequence length of "
|
||||
f"{query_seq_length} into {num_chunks} chunks")
|
||||
|
||||
mh_q = [
|
||||
q[:, head_idx * dim_head:(head_idx + 1) *
|
||||
dim_head, :, :] for head_idx in range(heads)
|
||||
] # (bs, dim_head, 1, max_seq_length) * heads
|
||||
|
||||
# Chunk the query sequence for each head
|
||||
mh_q_chunked = [
|
||||
[h_q[..., chunk_idx * CHUNK_SIZE:(chunk_idx + 1) * CHUNK_SIZE] for chunk_idx in range(num_chunks)]
|
||||
for h_q in mh_q
|
||||
] # ((bs, dim_head, 1, QUERY_SEQ_CHUNK_SIZE) * num_chunks) * heads
|
||||
|
||||
k = k.transpose(1, 3)
|
||||
mh_k = [
|
||||
k[:, :, :,
|
||||
head_idx * dim_head:(head_idx + 1) * dim_head]
|
||||
for head_idx in range(heads)
|
||||
] # (bs, max_seq_length, 1, dim_head) * heads
|
||||
|
||||
mh_v = [
|
||||
v[:, head_idx * dim_head:(head_idx + 1) *
|
||||
dim_head, :, :] for head_idx in range(heads)
|
||||
] # (bs, dim_head, 1, max_seq_length) * heads
|
||||
|
||||
attn_weights = [
|
||||
[
|
||||
torch.einsum("bchq,bkhc->bkhq", [qi_chunk, ki]) * (dim_head**-0.5)
|
||||
for qi_chunk in h_q_chunked
|
||||
] for h_q_chunked, ki in zip(mh_q_chunked, mh_k)
|
||||
] # ((bs, max_seq_length, 1, chunk_size) * num_chunks) * heads
|
||||
|
||||
attn_weights = [
|
||||
[aw_chunk.softmax(dim=1) for aw_chunk in aw_chunked]
|
||||
for aw_chunked in attn_weights
|
||||
] # ((bs, max_seq_length, 1, chunk_size) * num_chunks) * heads
|
||||
|
||||
attn = [
|
||||
[
|
||||
torch.einsum("bkhq,bchk->bchq", wi_chunk, vi)
|
||||
for wi_chunk in wi_chunked
|
||||
] for wi_chunked, vi in zip(attn_weights, mh_v)
|
||||
] # ((bs, dim_head, 1, chunk_size) * num_chunks) * heads
|
||||
|
||||
attn = torch.cat([
|
||||
torch.cat(attn_chunked, dim=3) for attn_chunked in attn
|
||||
], dim=1) # (bs, dim, 1, max_seq_length)
|
||||
|
||||
return attn
|
||||
|
||||
|
||||
def original(q, k, v, mask, heads, dim_head):
|
||||
""" Attention Implementation backing AttentionImplementations.ORIGINAL
|
||||
|
||||
- Not recommended for ANE
|
||||
- Recommended for GPU
|
||||
"""
|
||||
bs = q.size(0)
|
||||
mh_q = q.view(bs, heads, dim_head, -1)
|
||||
mh_k = k.view(bs, heads, dim_head, -1)
|
||||
mh_v = v.view(bs, heads, dim_head, -1)
|
||||
|
||||
attn_weights = torch.einsum("bhcq,bhck->bhqk", [mh_q, mh_k])
|
||||
attn_weights.mul_(dim_head**-0.5)
|
||||
|
||||
if mask is not None:
|
||||
attn_weights = attn_weights + mask
|
||||
|
||||
attn_weights = attn_weights.softmax(dim=3)
|
||||
|
||||
attn = torch.einsum("bhqk,bhck->bhcq", [attn_weights, mh_v])
|
||||
attn = attn.contiguous().view(bs, heads * dim_head, 1, -1)
|
||||
return attn
|
||||
@@ -0,0 +1,60 @@
|
||||
from python_coreml_stable_diffusion.torch2coreml import _compile_coreml_model
|
||||
|
||||
import argparse
|
||||
import coremltools as ct
|
||||
import numpy as np
|
||||
import os
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
# TODO: Read these values off of the NLContextualEmbedding API to enforce dimensions and track API versioning
|
||||
MAX_SEQUENCE_LENGTH = 256
|
||||
EMBED_DIM = 512
|
||||
BATCH_SIZE = 1
|
||||
|
||||
def main(args):
|
||||
# Layer that was trained to map NLContextualEmbedding to your text_encoder.hidden_size dimensionality
|
||||
text_encoder_projection = torch.jit.load(args.input_path)
|
||||
|
||||
# Prepare random inputs for tracing the network before conversion
|
||||
random_input = torch.randn(BATCH_SIZE, MAX_SEQUENCE_LENGTH, EMBED_DIM)
|
||||
|
||||
# Create a class to bake in the reshape operations required to fit the existing model interface
|
||||
class TextEncoderProjection(nn.Module):
|
||||
def __init__(self, proj):
|
||||
super().__init__()
|
||||
self.proj = proj
|
||||
|
||||
def forward(self, x):
|
||||
return self.proj(x).transpose(1, 2).unsqueeze(2) # BSC, BC1S
|
||||
|
||||
# Trace the torch model
|
||||
text_encoder_projection = torch.jit.trace(TextEncoderProjection(text_encoder_projection), (random_input,))
|
||||
|
||||
# Convert the model to Core ML
|
||||
mlpackage_path = os.path.join(args.output_dir, "MultilingualTextEncoderProjection.mlpackage")
|
||||
ct.convert(
|
||||
text_encoder_projection,
|
||||
inputs=[ct.TensorType('nlcontextualembeddings_output', shape=(1, MAX_SEQUENCE_LENGTH, EMBED_DIM), dtype=np.float32)],
|
||||
outputs=[ct.TensorType('encoder_hidden_states', dtype=np.float32)],
|
||||
minimum_deployment_target=ct.target.macOS14, # NLContextualEmbedding minimum availability build
|
||||
convert_to='mlprogram',
|
||||
).save()
|
||||
|
||||
# Compile the model and save it under the specified directory
|
||||
_compile_coreml_model(mlpackage_path, args.output_dir, final_name="MultilingualTextEncoderProjection")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
"--input-path",
|
||||
help="Path to the torchscript file that contains the projection layer"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-dir",
|
||||
help="Output directory in which the Core ML model should be saved",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
main(args)
|
||||
@@ -3,7 +3,9 @@
|
||||
# Copyright (C) 2022 Apple Inc. All Rights Reserved.
|
||||
#
|
||||
|
||||
from python_coreml_stable_diffusion import unet, controlnet
|
||||
from python_coreml_stable_diffusion import (
|
||||
unet, controlnet, chunk_mlprogram
|
||||
)
|
||||
|
||||
import argparse
|
||||
from collections import OrderedDict, defaultdict
|
||||
@@ -20,7 +22,6 @@ logger.setLevel(logging.INFO)
|
||||
|
||||
import numpy as np
|
||||
import os
|
||||
from python_coreml_stable_diffusion import chunk_mlprogram
|
||||
import requests
|
||||
import shutil
|
||||
import time
|
||||
@@ -141,30 +142,49 @@ def _convert_to_coreml(submodule_name, torchscript_module, sample_inputs,
|
||||
return coreml_model, out_path
|
||||
|
||||
|
||||
def quantize_weights_to_8bits(args):
|
||||
for model_name in [
|
||||
"text_encoder", "vae_decoder", "vae_encoder", "unet", "unet_chunk1", "unet_chunk2",
|
||||
"control-unet", "control-unet_chunk1", "control-unet_chunk2", "safety_checker"
|
||||
]:
|
||||
def quantize_weights(args):
|
||||
""" Quantize weights to args.quantize_nbits using a palette (look-up table)
|
||||
"""
|
||||
for model_name in ["text_encoder", "unet", "control-unet"]:
|
||||
logger.info(f"Quantizing {model_name} to {args.quantize_nbits}-bit precision")
|
||||
out_path = _get_out_path(args, model_name)
|
||||
_quantize_and_save_8bits_model(out_path, model_name)
|
||||
_quantize_weights(
|
||||
out_path,
|
||||
model_name,
|
||||
args.quantize_nbits
|
||||
)
|
||||
|
||||
if args.convert_controlnet:
|
||||
for controlnet_model_version in args.convert_controlnet:
|
||||
controlnet_model_name = controlnet_model_version.replace("/", "_")
|
||||
logger.info(f"Quantizing {controlnet_model_name} to {args.quantize_nbits}-bit precision")
|
||||
fname = f"ControlNet_{controlnet_model_name}.mlpackage"
|
||||
out_path = os.path.join(args.o, fname)
|
||||
_quantize_and_save_8bits_model(out_path, controlnet_model_name)
|
||||
|
||||
_quantize_weights(
|
||||
out_path,
|
||||
controlnet_model_name,
|
||||
args.quantize_nbits
|
||||
)
|
||||
|
||||
def _quantize_and_save_8bits_model(out_path, model_name):
|
||||
def _quantize_weights(out_path, model_name, nbits):
|
||||
if os.path.exists(out_path):
|
||||
logger.info(f"Quantizing {model_name}")
|
||||
mlmodel = ct.models.MLModel(out_path,
|
||||
compute_units=ct.ComputeUnit.CPU_ONLY)
|
||||
mlmodel = ct.compression_utils.affine_quantize_weights(
|
||||
mlmodel, mode="linear")
|
||||
mlmodel.save(out_path)
|
||||
|
||||
op_config = ct.optimize.coreml.OpPalettizerConfig(
|
||||
mode="kmeans",
|
||||
nbits=nbits,
|
||||
)
|
||||
|
||||
config = ct.optimize.coreml.OptimizationConfig(
|
||||
global_config=op_config,
|
||||
op_type_configs={
|
||||
"gather": None # avoid quantizing the embedding table
|
||||
}
|
||||
)
|
||||
|
||||
model = ct.optimize.coreml.palettize_weights(mlmodel, config=config).save(out_path)
|
||||
logger.info("Done")
|
||||
else:
|
||||
logger.info(
|
||||
@@ -266,7 +286,7 @@ def convert_text_encoder(pipe, args):
|
||||
|
||||
# Create sample inputs for tracing, conversion and correctness verification
|
||||
text_encoder_sequence_length = pipe.tokenizer.model_max_length
|
||||
text_encoder_hidden_size = pipe.text_encoder.config.hidden_size
|
||||
text_encoder_hidden_size = args.text_encoder_hidden_size or pipe.text_encoder.config.hidden_size
|
||||
|
||||
sample_text_encoder_inputs = {
|
||||
"input_ids":
|
||||
@@ -283,8 +303,8 @@ def convert_text_encoder(pipe, args):
|
||||
}
|
||||
logger.info(f"Sample inputs spec: {sample_text_encoder_inputs_spec}")
|
||||
|
||||
def _build_causal_attention_mask(self, bsz, seq_len, dtype):
|
||||
mask = torch.ones((bsz, seq_len, seq_len), dtype=dtype) * -1e4
|
||||
def _build_causal_attention_mask(self, bsz, seq_len, dtype, device=None):
|
||||
mask = torch.ones((bsz, seq_len, seq_len), dtype=dtype, device=device) * -1e4
|
||||
mask.triu_(1)
|
||||
mask = mask.unsqueeze(1)
|
||||
return mask
|
||||
@@ -438,6 +458,8 @@ def convert_vae_decoder(pipe, args):
|
||||
super().__init__()
|
||||
self.post_quant_conv = pipe.vae.post_quant_conv
|
||||
self.decoder = pipe.vae.decoder
|
||||
# Disable torch 2.0 scaled dot-product attention: https://github.com/apple/coremltools/issues/1823
|
||||
self.decoder.mid_block.attentions[0]._use_2_0_attn = False
|
||||
|
||||
def forward(self, z):
|
||||
return self.decoder(self.post_quant_conv(z))
|
||||
@@ -525,6 +547,8 @@ def convert_vae_encoder(pipe, args):
|
||||
super().__init__()
|
||||
self.quant_conv = pipe.vae.quant_conv
|
||||
self.encoder = pipe.vae.encoder
|
||||
# Disable torch 2.0 scaled dot-product attention: https://github.com/apple/coremltools/issues/1823
|
||||
self.encoder.mid_block.attentions[0]._use_2_0_attn = False
|
||||
|
||||
def forward(self, z):
|
||||
return self.quant_conv(self.encoder(z))
|
||||
@@ -614,9 +638,9 @@ def convert_unet(pipe, args):
|
||||
|
||||
encoder_hidden_states_shape = (
|
||||
batch_size,
|
||||
pipe.text_encoder.config.hidden_size,
|
||||
args.text_encoder_hidden_size or pipe.text_encoder.config.hidden_size,
|
||||
1,
|
||||
pipe.text_encoder.config.max_position_embeddings,
|
||||
args.text_token_sequence_length or pipe.text_encoder.config.max_position_embeddings,
|
||||
)
|
||||
|
||||
# Create the scheduled timesteps for downstream use
|
||||
@@ -736,6 +760,10 @@ def convert_unet(pipe, args):
|
||||
"Same shape and dtype as the `sample` input. " \
|
||||
"The predicted noise to facilitate the reverse diffusion (denoising) process"
|
||||
|
||||
# Set package version metadata
|
||||
from python_coreml_stable_diffusion._version import __version__
|
||||
coreml_unet.user_defined_metadata["com.github.apple.ml-stable-diffusion.version"] = __version__
|
||||
|
||||
_save_mlpackage(coreml_unet, out_path)
|
||||
logger.info(f"Saved unet into {out_path}")
|
||||
|
||||
@@ -952,6 +980,11 @@ def convert_safety_checker(pipe, args):
|
||||
del traced_safety_checker, coreml_safety_checker, pipe.safety_checker
|
||||
gc.collect()
|
||||
|
||||
def _get_controlnet_base_model(controlnet_model_version):
|
||||
from huggingface_hub import model_info
|
||||
info = model_info(controlnet_model_version)
|
||||
return info.cardData.get("base_model", None)
|
||||
|
||||
def convert_controlnet(pipe, args):
|
||||
""" Converts each ControlNet for Stable Diffusion
|
||||
"""
|
||||
@@ -965,12 +998,20 @@ def convert_controlnet(pipe, args):
|
||||
"convert_text_encoder() deletes pipe.text_encoder to save RAM. "
|
||||
"Please use convert_unet() before convert_text_encoder()")
|
||||
|
||||
if args.model_version != "runwayml/stable-diffusion-v1-5":
|
||||
logger.warning(
|
||||
"The original ControlNet models were trained using Stable Diffusion v1.5. "
|
||||
"It is possible that the converted model may not be compatible with controlnet.")
|
||||
|
||||
for i, controlnet_model_version in enumerate(args.convert_controlnet):
|
||||
base_model = _get_controlnet_base_model(controlnet_model_version)
|
||||
|
||||
if base_model is None and args.model_version != "runwayml/stable-diffusion-v1-5":
|
||||
logger.warning(
|
||||
f"The original ControlNet models were trained using Stable Diffusion v1.5. "
|
||||
f"It is possible that model {args.model_version} is not compatible with controlnet.")
|
||||
if base_model is not None and base_model != args.model_version:
|
||||
raise RuntimeError(
|
||||
f"ControlNet model {controlnet_model_version} was trained using "
|
||||
f"Stable Diffusion model {base_model}.\n However, you specified "
|
||||
f"version {args.model_version} in the command line. Please, use "
|
||||
f"--model-version {base_model} to convert this model.")
|
||||
|
||||
controlnet_model_name = controlnet_model_version.replace("/", "_")
|
||||
fname = f"ControlNet_{controlnet_model_name}.mlpackage"
|
||||
out_path = os.path.join(args.o, fname)
|
||||
@@ -992,9 +1033,9 @@ def convert_controlnet(pipe, args):
|
||||
|
||||
encoder_hidden_states_shape = (
|
||||
batch_size,
|
||||
pipe.text_encoder.config.hidden_size,
|
||||
args.text_encoder_hidden_size or pipe.text_encoder.config.hidden_size,
|
||||
1,
|
||||
pipe.text_encoder.config.max_position_embeddings,
|
||||
args.text_token_sequence_length or pipe.text_encoder.config.max_position_embeddings,
|
||||
)
|
||||
|
||||
controlnet_cond_shape = (
|
||||
@@ -1051,18 +1092,10 @@ def convert_controlnet(pipe, args):
|
||||
baseline_out = original_controlnet(**baseline_sample_controlnet_inputs,
|
||||
return_dict=False)
|
||||
reference_out = reference_controlnet(*sample_controlnet_inputs.values())
|
||||
|
||||
baseline_down_residuals, baseline_mid_residuals = baseline_out
|
||||
baseline_out = baseline_down_residuals + (baseline_mid_residuals,)
|
||||
reference_down_residuals, reference_mid_residuals = reference_out
|
||||
reference_out = reference_down_residuals +(reference_mid_residuals,)
|
||||
|
||||
for key, b_out, r_out in zip(output_keys, baseline_out, reference_out):
|
||||
b_out = b_out.numpy()
|
||||
r_out = r_out.numpy()
|
||||
logger.info(f"Check {key} correctness")
|
||||
report_correctness(b_out, r_out,
|
||||
f"controlnet({controlnet_model_name}) baseline to reference PyTorch")
|
||||
report_correctness(
|
||||
baseline_out[-1].numpy(),
|
||||
reference_out[-1].numpy(),
|
||||
f"{controlnet_model_name} baseline to reference PyTorch")
|
||||
|
||||
del original_controlnet
|
||||
gc.collect()
|
||||
@@ -1111,12 +1144,12 @@ def convert_controlnet(pipe, args):
|
||||
# Parity check PyTorch vs CoreML
|
||||
if args.check_output_correctness:
|
||||
coreml_out = coreml_controlnet.predict(coreml_sample_controlnet_inputs)
|
||||
for key, b_out in zip(output_keys, baseline_out):
|
||||
b_out = b_out.numpy()
|
||||
logger.info(f"Check {key} correctness")
|
||||
report_correctness(b_out, coreml_out[key],
|
||||
"controlnet baseline PyTorch to reference CoreML")
|
||||
|
||||
report_correctness(
|
||||
baseline_out[-1].numpy(),
|
||||
coreml_out[output_keys[-1]],
|
||||
"controlnet baseline PyTorch to reference CoreML"
|
||||
)
|
||||
|
||||
del coreml_controlnet
|
||||
gc.collect()
|
||||
|
||||
@@ -1169,17 +1202,16 @@ def main(args):
|
||||
convert_safety_checker(pipe, args)
|
||||
logger.info("Converted safety_checker")
|
||||
|
||||
if args.quantize_nbits is not None:
|
||||
logger.info(f"Quantizing weights to {args.quantize_nbits}-bit precision")
|
||||
quantize_weights(args)
|
||||
logger.info(f"Quantized weights to {args.quantize_nbits}-bit precision")
|
||||
|
||||
if args.bundle_resources_for_swift_cli:
|
||||
logger.info("Bundling resources for the Swift CLI")
|
||||
bundle_resources_for_swift_cli(args)
|
||||
logger.info("Bundled resources for the Swift CLI")
|
||||
|
||||
if args.quantize_weights_to_8bits:
|
||||
# Note: Not recommended, significantly degrades generated image quality
|
||||
logger.info("Quantizing weights to 8-bit precision")
|
||||
quantize_weights_to_8bits(args)
|
||||
logger.info("Quantized weights to 8-bit precision")
|
||||
|
||||
|
||||
def parser_spec():
|
||||
parser = argparse.ArgumentParser()
|
||||
@@ -1224,6 +1256,20 @@ def parser_spec():
|
||||
help=
|
||||
"The spatial resolution (number of cols) of the latent space. `Defaults to pipe.unet.config.sample_size`",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--text-token-sequence-length",
|
||||
type=int,
|
||||
default=None,
|
||||
help=
|
||||
"The token sequence length for the text encoder. `Defaults to pipe.text_encoder.config.max_position_embeddings`",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--text-encoder-hidden-size",
|
||||
type=int,
|
||||
default=None,
|
||||
help=
|
||||
"The hidden size for the text encoder. `Defaults to pipe.text_encoder.config.hidden_size`",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--attention-implementation",
|
||||
choices=tuple(ai
|
||||
@@ -1240,39 +1286,39 @@ def parser_spec():
|
||||
"--check-output-correctness",
|
||||
action="store_true",
|
||||
help=
|
||||
("If specified, compares the outputs of original PyTorch and final CoreML models and reports PSNR in dB. ",
|
||||
"Enabling this feature uses more memory. Disable it if your machine runs out of memory."
|
||||
))
|
||||
"If specified, compares the outputs of original PyTorch and final CoreML models and reports PSNR in dB. "
|
||||
"Enabling this feature uses more memory. Disable it if your machine runs out of memory."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--chunk-unet",
|
||||
action="store_true",
|
||||
help=
|
||||
("If specified, generates two mlpackages out of the unet model which approximately equal weights sizes. "
|
||||
"This is required for ANE deployment on iOS and iPadOS. Not required for macOS."
|
||||
))
|
||||
"If specified, generates two mlpackages out of the unet model which approximately equal weights sizes. "
|
||||
"This is required for ANE deployment on iOS and iPadOS. Not required for macOS."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--quantize-weights-to-8bits",
|
||||
action="store_true",
|
||||
help=
|
||||
("If specified, quantize 16-bits weights to 8-bits weights in-place for all models. "
|
||||
"Not recommended as the generated image quality degraded significantly after 8-bit weight quantization"
|
||||
))
|
||||
"--quantize-nbits",
|
||||
default=None,
|
||||
choices=(2, 4, 6, 8),
|
||||
type=int,
|
||||
help="If specified, quantized each model to nbits precision"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--unet-support-controlnet",
|
||||
action="store_true",
|
||||
help=
|
||||
("If specified, enable unet to receive additional inputs from controlnet. "
|
||||
"Each input added to corresponding resnet output."
|
||||
))
|
||||
"If specified, enable unet to receive additional inputs from controlnet. "
|
||||
"Each input added to corresponding resnet output."
|
||||
)
|
||||
|
||||
# Swift CLI Resource Bundling
|
||||
parser.add_argument(
|
||||
"--bundle-resources-for-swift-cli",
|
||||
action="store_true",
|
||||
help=
|
||||
("If specified, creates a resources directory compatible with the sample Swift CLI. "
|
||||
"It compiles all four models and adds them to a StableDiffusionResources directory "
|
||||
"along with a `vocab.json` and `merges.txt` for the text tokenizer"))
|
||||
"If specified, creates a resources directory compatible with the sample Swift CLI. "
|
||||
"It compiles all four models and adds them to a StableDiffusionResources directory "
|
||||
"along with a `vocab.json` and `merges.txt` for the text tokenizer")
|
||||
parser.add_argument(
|
||||
"--text-encoder-vocabulary-url",
|
||||
default=
|
||||
|
||||
@@ -4,6 +4,7 @@
|
||||
#
|
||||
|
||||
from python_coreml_stable_diffusion.layer_norm import LayerNormANE
|
||||
from python_coreml_stable_diffusion import attention
|
||||
|
||||
from diffusers.configuration_utils import ConfigMixin, register_to_config
|
||||
from diffusers import ModelMixin
|
||||
@@ -33,6 +34,7 @@ if not _macos_version() >= (13, 1):
|
||||
class AttentionImplementations(Enum):
|
||||
ORIGINAL = "ORIGINAL"
|
||||
SPLIT_EINSUM = "SPLIT_EINSUM"
|
||||
SPLIT_EINSUM_V2 = "SPLIT_EINSUM_V2"
|
||||
|
||||
|
||||
ATTENTION_IMPLEMENTATION_IN_EFFECT = AttentionImplementations.SPLIT_EINSUM
|
||||
@@ -95,64 +97,13 @@ class CrossAttention(nn.Module):
|
||||
)
|
||||
|
||||
if ATTENTION_IMPLEMENTATION_IN_EFFECT == AttentionImplementations.ORIGINAL:
|
||||
# This version of the attention function is recommended for high GPU core count
|
||||
# devices such as the M1 Max and M1 Ultra
|
||||
bs = q.size(0)
|
||||
mh_q = q.view(bs, self.heads, self.dim_head, -1)
|
||||
mh_k = k.view(bs, self.heads, self.dim_head, -1)
|
||||
mh_v = v.view(bs, self.heads, self.dim_head, -1)
|
||||
|
||||
attn_weights = torch.einsum("bhcq,bhck->bhqk", [mh_q, mh_k])
|
||||
attn_weights.mul_(self.scale)
|
||||
|
||||
if mask is not None:
|
||||
attn_weights = attn_weights + mask
|
||||
|
||||
attn_weights = attn_weights.softmax(dim=3)
|
||||
|
||||
attn = torch.einsum("bhqk,bhck->bhcq", [attn_weights, mh_v])
|
||||
attn = attn.contiguous().view(bs, self.heads * self.dim_head, 1,
|
||||
-1)
|
||||
attn = attention.original(q, k, v, mask, self.heads, self.dim_head)
|
||||
|
||||
elif ATTENTION_IMPLEMENTATION_IN_EFFECT == AttentionImplementations.SPLIT_EINSUM:
|
||||
# The split attention and einsum from https://machinelearning.apple.com/research/neural-engine-transformers
|
||||
# are utilized to build an ANE implementation. This version is marginally slower on the GPU engine and is
|
||||
# not recommended for Max and Ultra Mac variants
|
||||
mh_q = [
|
||||
q[:, head_idx * self.dim_head:(head_idx + 1) *
|
||||
self.dim_head, :, :] for head_idx in range(self.heads)
|
||||
] # (bs, dim_head, 1, max_seq_length) * heads
|
||||
attn = attention.split_einsum(q, k, v, mask, self.heads, self.dim_head)
|
||||
|
||||
k = k.transpose(1, 3)
|
||||
mh_k = [
|
||||
k[:, :, :,
|
||||
head_idx * self.dim_head:(head_idx + 1) * self.dim_head]
|
||||
for head_idx in range(self.heads)
|
||||
] # (bs, max_seq_length, 1, dim_head) * heads
|
||||
|
||||
mh_v = [
|
||||
v[:, head_idx * self.dim_head:(head_idx + 1) *
|
||||
self.dim_head, :, :] for head_idx in range(self.heads)
|
||||
] # (bs, dim_head, 1, max_seq_length) * heads
|
||||
|
||||
attn_weights = [
|
||||
torch.einsum("bchq,bkhc->bkhq", [qi, ki]) * self.scale
|
||||
for qi, ki in zip(mh_q, mh_k)
|
||||
] # (bs, max_seq_length, 1, max_seq_length) * heads
|
||||
|
||||
if mask is not None:
|
||||
for head_idx in range(self.heads):
|
||||
attn_weights[head_idx] = attn_weights[head_idx] + mask
|
||||
|
||||
attn_weights = [
|
||||
aw.softmax(dim=1) for aw in attn_weights
|
||||
] # (bs, max_seq_length, 1, max_seq_length) * heads
|
||||
attn = [
|
||||
torch.einsum("bkhq,bchk->bchq", wi, vi)
|
||||
for wi, vi in zip(attn_weights, mh_v)
|
||||
] # (bs, dim_head, 1, max_seq_length) * heads
|
||||
|
||||
attn = torch.cat(attn, dim=1) # (bs, dim, 1, max_seq_length)
|
||||
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)
|
||||
|
||||
+5
-3
@@ -1,5 +1,7 @@
|
||||
coremltools
|
||||
coremltools>=7.0b1
|
||||
diffusers[torch]
|
||||
torch
|
||||
transformers
|
||||
scipy
|
||||
transformers==4.29.2
|
||||
scipy
|
||||
scikit-learn
|
||||
pytest
|
||||
@@ -14,12 +14,15 @@ setup(
|
||||
long_description_content_type='text/markdown',
|
||||
author='Apple Inc.',
|
||||
install_requires=[
|
||||
"coremltools>=6.1",
|
||||
"coremltools>=7.0b1",
|
||||
"diffusers[torch]",
|
||||
"torch",
|
||||
"transformers",
|
||||
"transformers==4.29.2",
|
||||
"huggingface-hub",
|
||||
"scipy",
|
||||
"numpy<1.24",
|
||||
"pytest",
|
||||
"scikit-learn",
|
||||
],
|
||||
packages=find_packages(),
|
||||
classifiers=[
|
||||
|
||||
@@ -0,0 +1,190 @@
|
||||
// For licensing see accompanying LICENSE.md file.
|
||||
// Copyright (C) 2023 Apple Inc. All Rights Reserved.
|
||||
|
||||
import Foundation
|
||||
import NaturalLanguage
|
||||
import CoreML
|
||||
|
||||
@available(iOS 17.0, macOS 14.0, *)
|
||||
public struct MultilingualTextEncoder: TextEncoderModel {
|
||||
let adapter: ManagedMLModel?
|
||||
|
||||
let embeddingModel: NLContextualEmbedding
|
||||
|
||||
// TODO: use maximum sequence length from embedding.
|
||||
let maximumEmbeddingSequenceLength = 256
|
||||
|
||||
/// Creates a multilingual text encoder.
|
||||
///
|
||||
/// - Parameters:
|
||||
/// - url: The location of the compiled Core ML adapter model. The model is a linear projection layer that
|
||||
/// transforms the contextual embedding size of 512 to the default text encoder CLIP size of 768.
|
||||
/// - configuration: The configuration to be used when the model is loaded.
|
||||
/// - script: The scipt of the contextual embedding.
|
||||
public init(
|
||||
modelAt url: URL? = nil,
|
||||
configuration: MLModelConfiguration = .init(),
|
||||
script: Script = .latin
|
||||
) {
|
||||
if let url {
|
||||
self.adapter = ManagedMLModel(modelAt: url, configuration: configuration)
|
||||
} else {
|
||||
self.adapter = nil
|
||||
}
|
||||
self.embeddingModel = NLContextualEmbedding(script: script.asNLScript)!
|
||||
self.embeddingModel.requestAssets { _, _ in }
|
||||
}
|
||||
|
||||
/// Loads model resources into memory.
|
||||
public func loadResources() throws {
|
||||
try adapter?.loadResources()
|
||||
try embeddingModel.load()
|
||||
}
|
||||
|
||||
/// Unloads the model resources to free up memory.
|
||||
public func unloadResources() {
|
||||
adapter?.unloadResources()
|
||||
embeddingModel.unload()
|
||||
}
|
||||
|
||||
/// Encodes the input text.
|
||||
///
|
||||
/// - Parameter text: The input text.
|
||||
/// - Returns: An embedding shaped array.
|
||||
public func encode(_ text: String) throws -> MLShapedArray<Float> {
|
||||
guard embeddingModel.hasAvailableAssets else {
|
||||
throw Error.missingEmbeddingResource
|
||||
}
|
||||
|
||||
// Create the text embedding result.
|
||||
let embedding = try embeddingModel.embeddingResult(for: text, language: nil)
|
||||
|
||||
// Create embedding array from token vectors.
|
||||
var shapedEmbeddings = MLShapedArray<Double>(
|
||||
repeating: 0.0,
|
||||
shape: [1, maximumEmbeddingSequenceLength, embeddingModel.dimension]
|
||||
)
|
||||
shapedEmbeddings.withUnsafeMutableShapedBufferPointer { pointer, _, _ in
|
||||
var tokenIndex = 0
|
||||
embedding.enumerateTokenVectors(in: text.startIndex ..< text.endIndex) { (tokenEmbeddings, _) -> Bool in
|
||||
for tokenEmbeddingIndex in 0 ..< tokenEmbeddings.count {
|
||||
pointer[tokenIndex * embeddingModel.dimension + tokenEmbeddingIndex] = tokenEmbeddings[tokenEmbeddingIndex]
|
||||
}
|
||||
tokenIndex += 1
|
||||
return true
|
||||
}
|
||||
}
|
||||
|
||||
if adapter == nil {
|
||||
// Return embeddings with shape [1, 256, 512].
|
||||
return MLShapedArray(converting: shapedEmbeddings)
|
||||
} else {
|
||||
// Project the embeddings to the correct CLIP model input shape of [1, 768, 1, 256].
|
||||
return try projectEmbeddings(shapedEmbeddings)
|
||||
}
|
||||
}
|
||||
|
||||
/// Creates the adapter model input feature provider.
|
||||
private func prepareProjectionInput(_ input: MLShapedArray<Double>) throws -> MLDictionaryFeatureProvider {
|
||||
guard let adapter else {
|
||||
fatalError("Cannot prepare projection input without an adapter.")
|
||||
}
|
||||
return try adapter.perform { model in
|
||||
guard let inputDescription = model.modelDescription.inputDescriptionsByName.first?.value else {
|
||||
throw Error.missingAdapterInput
|
||||
}
|
||||
return try MLDictionaryFeatureProvider(dictionary: [inputDescription.name: MLMultiArray(input)])
|
||||
}
|
||||
}
|
||||
|
||||
/// Processes the adapter model output feature provider.
|
||||
private func processProjectionOutput(_ output: MLFeatureProvider) throws -> MLShapedArray<Float> {
|
||||
guard let adapter else {
|
||||
fatalError("Cannot process projection output without an adapter.")
|
||||
}
|
||||
return try adapter.perform { model in
|
||||
guard let outputDescription = model.modelDescription.outputDescriptionsByName.first?.value else {
|
||||
throw Error.missingAdapterOutput
|
||||
}
|
||||
guard let result = output
|
||||
.featureValue(for: outputDescription.name)?
|
||||
.multiArrayValue else {
|
||||
|
||||
throw Error.incompatibleAdapterOutputDataFormat(
|
||||
expected: .multiArray,
|
||||
actual: outputDescription.type
|
||||
)
|
||||
}
|
||||
|
||||
return MLShapedArray(converting: result)
|
||||
}
|
||||
}
|
||||
|
||||
/// Projects the embeddings.
|
||||
private func projectEmbeddings(_ embeddings: MLShapedArray<Double>) throws -> MLShapedArray<Float> {
|
||||
guard let adapter else {
|
||||
fatalError("Cannot project embeddings without an adapter.")
|
||||
}
|
||||
let inputFeatureProvider = try prepareProjectionInput(embeddings)
|
||||
let projection = try adapter.perform { model in
|
||||
return try model.prediction(from: inputFeatureProvider)
|
||||
}
|
||||
return try processProjectionOutput(projection)
|
||||
}
|
||||
}
|
||||
|
||||
@available(iOS 17.0, macOS 14.0, *)
|
||||
extension MultilingualTextEncoder {
|
||||
/// A multilingual text encoder error.
|
||||
public enum Error: Swift.Error, LocalizedError, Equatable, CustomDebugStringConvertible {
|
||||
/// An error that indicates that the resource for the embedding is missing.
|
||||
case missingEmbeddingResource
|
||||
|
||||
/// An error that indicates that the adapter model input data has the wrong format.
|
||||
case incompatibleAdapterInputDataFormat(expected: MLFeatureType, actual: MLFeatureType)
|
||||
|
||||
/// An error that indicates that the adapter model output data has the wrong format.
|
||||
case incompatibleAdapterOutputDataFormat(expected: MLFeatureType, actual: MLFeatureType)
|
||||
|
||||
/// An error that indicates that the adapter model is missing an input.
|
||||
case missingAdapterInput
|
||||
|
||||
/// An error that indicates that the adapter model is missing an output.
|
||||
case missingAdapterOutput
|
||||
|
||||
/// A debug description of the error.
|
||||
public var errorDescription: String? {
|
||||
debugDescription
|
||||
}
|
||||
|
||||
/// A text representation of the error.
|
||||
public var debugDescription: String {
|
||||
switch self {
|
||||
case .missingEmbeddingResource:
|
||||
return "Resources required for generating embeddings are missing. Make sure that your device is connected to the internet and try again."
|
||||
case .incompatibleAdapterInputDataFormat(expected: let expected, actual: let actual):
|
||||
return "The adapter model input expected to be \(expected) but is \(actual)."
|
||||
case .incompatibleAdapterOutputDataFormat(expected: let expected, actual: let actual):
|
||||
return "The adapter model output expected to be \(expected) but is \(actual)."
|
||||
case .missingAdapterInput:
|
||||
return "The adapter model is missing an input."
|
||||
case .missingAdapterOutput:
|
||||
return "The adapter model is missing an output."
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@available(iOS 16.2, macOS 13.1, *)
|
||||
public enum Script: String {
|
||||
case latin, cyrillic, cjk
|
||||
|
||||
@available(iOS 17.0, macOS 14.0, *)
|
||||
var asNLScript: NLScript {
|
||||
switch self {
|
||||
case .latin: return .latin
|
||||
case .cyrillic: return .cyrillic
|
||||
case .cjk: return .simplifiedChinese
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -3,6 +3,7 @@
|
||||
|
||||
import Foundation
|
||||
import CoreML
|
||||
import NaturalLanguage
|
||||
|
||||
@available(iOS 16.2, macOS 13.1, *)
|
||||
public extension StableDiffusionPipeline {
|
||||
@@ -22,6 +23,7 @@ public extension StableDiffusionPipeline {
|
||||
public let controlledUnetURL: URL
|
||||
public let controlledUnetChunk1URL: URL
|
||||
public let controlledUnetChunk2URL: URL
|
||||
public let multilingualTextEncoderProjectionURL: URL
|
||||
|
||||
public init(resourcesAt baseURL: URL) {
|
||||
textEncoderURL = baseURL.appending(path: "TextEncoder.mlmodelc")
|
||||
@@ -37,6 +39,7 @@ public extension StableDiffusionPipeline {
|
||||
controlledUnetURL = baseURL.appending(path: "ControlledUnet.mlmodelc")
|
||||
controlledUnetChunk1URL = baseURL.appending(path: "ControlledUnetChunk1.mlmodelc")
|
||||
controlledUnetChunk2URL = baseURL.appending(path: "ControlledUnetChunk2.mlmodelc")
|
||||
multilingualTextEncoderProjectionURL = baseURL.appending(path: "MultilingualTextEncoderProjection.mlmodelc")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -44,29 +47,40 @@ public extension StableDiffusionPipeline {
|
||||
/// specified URL
|
||||
///
|
||||
/// - Parameters:
|
||||
/// - baseURL: URL pointing to directory holding all model
|
||||
/// and tokenization resources
|
||||
/// - baseURL: URL pointing to directory holding all model and tokenization resources
|
||||
/// - controlNetModelNames: Specify ControlNet models to use in generation
|
||||
/// - configuration: The configuration to load model resources with
|
||||
/// - disableSafety: Load time disable of safety to save memory
|
||||
/// - reduceMemory: Setup pipeline in reduced memory mode
|
||||
/// - useMultilingualTextEncoder: Option to use system multilingual NLContextualEmbedding as encoder
|
||||
/// - script: Optional natural language script to use for the text encoder.
|
||||
/// - Returns:
|
||||
/// Pipeline ready for image generation if all necessary resources loaded
|
||||
init(resourcesAt baseURL: URL,
|
||||
controlNet controlNetModelNames: [String],
|
||||
configuration config: MLModelConfiguration = .init(),
|
||||
disableSafety: Bool = false,
|
||||
reduceMemory: Bool = false) throws {
|
||||
init(
|
||||
resourcesAt baseURL: URL,
|
||||
controlNet controlNetModelNames: [String],
|
||||
configuration config: MLModelConfiguration = .init(),
|
||||
disableSafety: Bool = false,
|
||||
reduceMemory: Bool = false,
|
||||
useMultilingualTextEncoder: Bool = false,
|
||||
script: Script? = nil
|
||||
) throws {
|
||||
|
||||
/// Expect URL of each resource
|
||||
let urls = ResourceURLs(resourcesAt: baseURL)
|
||||
let textEncoder: TextEncoderModel
|
||||
if useMultilingualTextEncoder {
|
||||
guard #available(macOS 14.0, iOS 17.0, *) else { throw Error.unsupportedOSVersion }
|
||||
textEncoder = MultilingualTextEncoder(
|
||||
modelAt: urls.multilingualTextEncoderProjectionURL,
|
||||
configuration: config,
|
||||
script: script ?? .latin
|
||||
)
|
||||
} else {
|
||||
let tokenizer = try BPETokenizer(mergesAt: urls.mergesURL, vocabularyAt: urls.vocabURL)
|
||||
textEncoder = TextEncoder(tokenizer: tokenizer, modelAt: urls.textEncoderURL, configuration: config)
|
||||
}
|
||||
|
||||
// Text tokenizer and encoder
|
||||
let tokenizer = try BPETokenizer(mergesAt: urls.mergesURL, vocabularyAt: urls.vocabURL)
|
||||
let textEncoder = TextEncoder(tokenizer: tokenizer,
|
||||
modelAt: urls.textEncoderURL,
|
||||
configuration: config)
|
||||
|
||||
// ControlNet model
|
||||
var controlNet: ControlNet? = nil
|
||||
let controlNetURLs = controlNetModelNames.map { model in
|
||||
@@ -119,12 +133,28 @@ public extension StableDiffusionPipeline {
|
||||
}
|
||||
|
||||
// Construct pipeline
|
||||
self.init(textEncoder: textEncoder,
|
||||
unet: unet,
|
||||
decoder: decoder,
|
||||
encoder: encoder,
|
||||
controlNet: controlNet,
|
||||
safetyChecker: safetyChecker,
|
||||
reduceMemory: reduceMemory)
|
||||
if #available(macOS 14.0, iOS 17.0, *) {
|
||||
self.init(
|
||||
textEncoder: textEncoder,
|
||||
unet: unet,
|
||||
decoder: decoder,
|
||||
encoder: encoder,
|
||||
controlNet: controlNet,
|
||||
safetyChecker: safetyChecker,
|
||||
reduceMemory: reduceMemory,
|
||||
useMultilingualTextEncoder: useMultilingualTextEncoder,
|
||||
script: script
|
||||
)
|
||||
} else {
|
||||
self.init(
|
||||
textEncoder: textEncoder,
|
||||
unet: unet,
|
||||
decoder: decoder,
|
||||
encoder: encoder,
|
||||
controlNet: controlNet,
|
||||
safetyChecker: safetyChecker,
|
||||
reduceMemory: reduceMemory
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,10 +1,11 @@
|
||||
// For licensing see accompanying LICENSE.md file.
|
||||
// Copyright (C) 2022 Apple Inc. All Rights Reserved.
|
||||
|
||||
import Foundation
|
||||
import CoreML
|
||||
import Accelerate
|
||||
import CoreGraphics
|
||||
import CoreML
|
||||
import Foundation
|
||||
import NaturalLanguage
|
||||
|
||||
/// Schedulers compatible with StableDiffusionPipeline
|
||||
public enum StableDiffusionScheduler {
|
||||
@@ -31,10 +32,11 @@ public struct StableDiffusionPipeline: ResourceManaging {
|
||||
|
||||
public enum Error: String, Swift.Error {
|
||||
case startingImageProvidedWithoutEncoder
|
||||
case unsupportedOSVersion
|
||||
}
|
||||
|
||||
/// Model to generate embeddings for tokenized input text
|
||||
var textEncoder: TextEncoder
|
||||
var textEncoder: TextEncoderModel
|
||||
|
||||
/// Model used to predict noise residuals given an input, diffusion time step, and conditional embedding
|
||||
var unet: Unet
|
||||
@@ -64,6 +66,12 @@ public struct StableDiffusionPipeline: ResourceManaging {
|
||||
/// This will increase latency in favor of reducing memory
|
||||
var reduceMemory: Bool = false
|
||||
|
||||
/// Option to use system multilingual NLContextualEmbedding as encoder
|
||||
var useMultilingualTextEncoder: Bool = false
|
||||
|
||||
/// Optional natural language script to use for the text encoder.
|
||||
var script: Script? = nil
|
||||
|
||||
/// Creates a pipeline using the specified models and tokenizer
|
||||
///
|
||||
/// - Parameters:
|
||||
@@ -74,13 +82,15 @@ public struct StableDiffusionPipeline: ResourceManaging {
|
||||
/// - safetyChecker: Optional model for checking safety of generated images
|
||||
/// - reduceMemory: Option to enable reduced memory mode
|
||||
/// - Returns: Pipeline ready for image generation
|
||||
public init(textEncoder: TextEncoder,
|
||||
unet: Unet,
|
||||
decoder: Decoder,
|
||||
encoder: Encoder?,
|
||||
controlNet: ControlNet? = nil,
|
||||
safetyChecker: SafetyChecker? = nil,
|
||||
reduceMemory: Bool = false) {
|
||||
public init(
|
||||
textEncoder: TextEncoderModel,
|
||||
unet: Unet,
|
||||
decoder: Decoder,
|
||||
encoder: Encoder?,
|
||||
controlNet: ControlNet? = nil,
|
||||
safetyChecker: SafetyChecker? = nil,
|
||||
reduceMemory: Bool = false
|
||||
) {
|
||||
self.textEncoder = textEncoder
|
||||
self.unet = unet
|
||||
self.decoder = decoder
|
||||
@@ -90,6 +100,41 @@ public struct StableDiffusionPipeline: ResourceManaging {
|
||||
self.reduceMemory = reduceMemory
|
||||
}
|
||||
|
||||
/// Creates a pipeline using the specified models and tokenizer
|
||||
///
|
||||
/// - Parameters:
|
||||
/// - textEncoder: Model for encoding tokenized text
|
||||
/// - unet: Model for noise prediction on latent samples
|
||||
/// - decoder: Model for decoding latent sample to image
|
||||
/// - controlNet: Optional model to control generated images by additonal inputs
|
||||
/// - safetyChecker: Optional model for checking safety of generated images
|
||||
/// - reduceMemory: Option to enable reduced memory mode
|
||||
/// - useMultilingualTextEncoder: Option to use system multilingual NLContextualEmbedding as encoder
|
||||
/// - script: Optional natural language script to use for the text encoder.
|
||||
/// - Returns: Pipeline ready for image generation
|
||||
@available(iOS 17.0, macOS 14.0, *)
|
||||
public init(
|
||||
textEncoder: TextEncoderModel,
|
||||
unet: Unet,
|
||||
decoder: Decoder,
|
||||
encoder: Encoder?,
|
||||
controlNet: ControlNet? = nil,
|
||||
safetyChecker: SafetyChecker? = nil,
|
||||
reduceMemory: Bool = false,
|
||||
useMultilingualTextEncoder: Bool = false,
|
||||
script: Script? = nil
|
||||
) {
|
||||
self.textEncoder = textEncoder
|
||||
self.unet = unet
|
||||
self.decoder = decoder
|
||||
self.encoder = encoder
|
||||
self.controlNet = controlNet
|
||||
self.safetyChecker = safetyChecker
|
||||
self.reduceMemory = reduceMemory
|
||||
self.useMultilingualTextEncoder = useMultilingualTextEncoder
|
||||
self.script = script
|
||||
}
|
||||
|
||||
/// Load required resources for this pipeline
|
||||
///
|
||||
/// If reducedMemory is true this will instead call prewarmResources instead
|
||||
@@ -98,9 +143,10 @@ public struct StableDiffusionPipeline: ResourceManaging {
|
||||
if reduceMemory {
|
||||
try prewarmResources()
|
||||
} else {
|
||||
try textEncoder.loadResources()
|
||||
try unet.loadResources()
|
||||
try textEncoder.loadResources()
|
||||
try decoder.loadResources()
|
||||
try encoder?.loadResources()
|
||||
try controlNet?.loadResources()
|
||||
try safetyChecker?.loadResources()
|
||||
}
|
||||
@@ -111,6 +157,7 @@ public struct StableDiffusionPipeline: ResourceManaging {
|
||||
textEncoder.unloadResources()
|
||||
unet.unloadResources()
|
||||
decoder.unloadResources()
|
||||
encoder?.unloadResources()
|
||||
controlNet?.unloadResources()
|
||||
safetyChecker?.unloadResources()
|
||||
}
|
||||
@@ -120,6 +167,7 @@ public struct StableDiffusionPipeline: ResourceManaging {
|
||||
try textEncoder.prewarmResources()
|
||||
try unet.prewarmResources()
|
||||
try decoder.prewarmResources()
|
||||
try encoder?.prewarmResources()
|
||||
try controlNet?.prewarmResources()
|
||||
try safetyChecker?.prewarmResources()
|
||||
}
|
||||
@@ -150,7 +198,7 @@ public struct StableDiffusionPipeline: ResourceManaging {
|
||||
alongAxis: 0
|
||||
)
|
||||
|
||||
let hiddenStates = toHiddenStates(concatEmbedding)
|
||||
let hiddenStates = useMultilingualTextEncoder ? concatEmbedding : toHiddenStates(concatEmbedding)
|
||||
|
||||
/// Setup schedulers
|
||||
let scheduler: [Scheduler] = (0..<config.imageCount).map { _ in
|
||||
@@ -162,6 +210,9 @@ public struct StableDiffusionPipeline: ResourceManaging {
|
||||
|
||||
// Generate random latent samples from specified seed
|
||||
var latents: [MLShapedArray<Float32>] = try generateLatentSamples(configuration: config, scheduler: scheduler[0])
|
||||
if reduceMemory {
|
||||
encoder?.unloadResources()
|
||||
}
|
||||
let timestepStrength: Float? = config.mode == .imageToImage ? config.strength : nil
|
||||
|
||||
// Convert cgImage for ControlNet into MLShapedArray
|
||||
|
||||
@@ -4,9 +4,15 @@
|
||||
import Foundation
|
||||
import CoreML
|
||||
|
||||
@available(iOS 16.2, macOS 13.1, *)
|
||||
public protocol TextEncoderModel: ResourceManaging {
|
||||
|
||||
func encode(_ text: String) throws -> MLShapedArray<Float32>
|
||||
}
|
||||
|
||||
/// A model for encoding text
|
||||
@available(iOS 16.2, macOS 13.1, *)
|
||||
public struct TextEncoder: ResourceManaging {
|
||||
public struct TextEncoder: TextEncoderModel {
|
||||
|
||||
/// Text tokenizer
|
||||
var tokenizer: BPETokenizer
|
||||
|
||||
@@ -9,6 +9,7 @@ import StableDiffusion
|
||||
import UniformTypeIdentifiers
|
||||
import Cocoa
|
||||
import CoreImage
|
||||
import NaturalLanguage
|
||||
|
||||
@available(iOS 16.2, macOS 13.1, *)
|
||||
struct StableDiffusionSample: ParsableCommand {
|
||||
@@ -91,6 +92,12 @@ struct StableDiffusionSample: ParsableCommand {
|
||||
@Flag(help: "Reduce memory usage")
|
||||
var reduceMemory: Bool = false
|
||||
|
||||
@Flag(help: "Use system multilingual NLContextualEmbedding as encoder model")
|
||||
var useMultilingualTextEncoder: Bool = false
|
||||
|
||||
@Option(help: "The natural language script for the multilingual contextual embedding")
|
||||
var script: Script = .latin
|
||||
|
||||
mutating func run() throws {
|
||||
guard FileManager.default.fileExists(atPath: resourcePath) else {
|
||||
throw RunError.resources("Resource path does not exist \(resourcePath)")
|
||||
@@ -102,11 +109,27 @@ struct StableDiffusionSample: ParsableCommand {
|
||||
|
||||
log("Loading resources and creating pipeline\n")
|
||||
log("(Note: This can take a while the first time using these resources)\n")
|
||||
let pipeline = try StableDiffusionPipeline(resourcesAt: resourceURL,
|
||||
controlNet: controlnet,
|
||||
configuration: config,
|
||||
disableSafety: disableSafety,
|
||||
reduceMemory: reduceMemory)
|
||||
let pipeline: StableDiffusionPipeline
|
||||
if #available(macOS 14.0, iOS 17.0, *) {
|
||||
pipeline = try StableDiffusionPipeline(
|
||||
resourcesAt: resourceURL,
|
||||
controlNet: controlnet,
|
||||
configuration: config,
|
||||
disableSafety: disableSafety,
|
||||
reduceMemory: reduceMemory,
|
||||
useMultilingualTextEncoder: useMultilingualTextEncoder,
|
||||
script: script
|
||||
)
|
||||
} else {
|
||||
pipeline = try StableDiffusionPipeline(
|
||||
resourcesAt: resourceURL,
|
||||
controlNet: controlnet,
|
||||
configuration: config,
|
||||
disableSafety: disableSafety,
|
||||
reduceMemory: reduceMemory
|
||||
)
|
||||
}
|
||||
|
||||
try pipeline.loadResources()
|
||||
|
||||
let startingImage: CGImage?
|
||||
@@ -153,7 +176,7 @@ struct StableDiffusionSample: ParsableCommand {
|
||||
pipelineConfig.guidanceScale = guidanceScale
|
||||
pipelineConfig.schedulerType = scheduler.stableDiffusionScheduler
|
||||
pipelineConfig.rngType = rng.stableDiffusionRNG
|
||||
|
||||
|
||||
let images = try pipeline.generateImages(
|
||||
configuration: pipelineConfig,
|
||||
progressHandler: { progress in
|
||||
@@ -299,6 +322,9 @@ enum RNGOption: String, ExpressibleByArgument {
|
||||
}
|
||||
}
|
||||
|
||||
@available(iOS 16.2, macOS 13.1, *)
|
||||
extension Script: ExpressibleByArgument {}
|
||||
|
||||
if #available(iOS 16.2, macOS 13.1, *) {
|
||||
StableDiffusionSample.main()
|
||||
} else {
|
||||
|
||||
Reference in New Issue
Block a user