README improvements and reduceMemory option in Swift

Co-authored-by: msiracusa <[email protected]>

Author: atiorh

README.md

@@ -9,7 +9,7 @@ This repository comprises:
 - `python_coreml_stable_diffusion`, a Python package for converting PyTorch models to Core ML format and performing image generation with Hugging Face [diffusers](https://github.com/huggingface/diffusers) in Python
 - `StableDiffusion`, a Swift package that developers can add to their Xcode projects as a dependency to deploy image generation capabilities in their apps. The Swift package relies on the Core ML model files generated by `python_coreml_stable_diffusion`
 
-If you run into issues during installation or runtime, please refer to the [FAQ](#FAQ) section.
+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
@@ -25,6 +25,13 @@ M2 MacBook Air 8GB Latency (s)  | 18 | 23 | 23 |
 
 Please see [Important Notes on Performance Benchmarks](#important-notes-on-performance-benchmarks) section for details.
 
+## <a name="system-requirements"></a> System Requirements
+
+The following is recommended to use all the functionality in this repository:
+
+Python | macOS | Xcode | iPadOS, iOS |
+:------:|:------:|:------:|:------:|
+3.8 | 13.1 | 14.2 | 16.2 |
 
 ## <a name="converting-models-to-coreml"></a> Converting Models to Core ML
 
@@ -50,7 +57,7 @@ pip install -e .
 python -m python_coreml_stable_diffusion.torch2coreml --convert-unet --convert-text-encoder --convert-vae-decoder --convert-safety-checker -o <output-mlpackages-directory>
 ```
 
-**WARNING:** This command will download several GB worth of PyTorch checkpoints from Hugging Face.
+**WARNING:** This command will download several GB worth of PyTorch checkpoints from Hugging Face. Please ensure that you are on Wi-Fi and have enough disk space.
 
 This generally takes 15-20 minutes on an M1 MacBook Pro. Upon successful execution, the 4 neural network models that comprise Stable Diffusion will have been converted from PyTorch to Core ML (`.mlpackage`) and saved into the specified `<output-mlpackages-directory>`. Some additional notable arguments:
 
@@ -59,9 +66,9 @@ 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 ANE 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.
+- `--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.
 
-- `--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 non-ANE deployment. Please refer to the [Performance Benchmark](#performance-benchmark) section for further guidance.
+- `--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.
 
 - `--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.
 
@@ -86,27 +93,30 @@ Please refer to the help menu for all available arguments: `python -m python_cor
 
 </details>
 
-## Image Generation with Swift
+## <a name="image-gen-swift"></a> Image Generation with Swift
 
 <details>
   <summary> Click to expand </summary>
 
 ### <a name="swift-requirements"></a> System Requirements
-Building the Swift projects require:
-- macOS 13 or newer
-- Xcode 14.1 or newer with command line tools installed. Please check [developer.apple.com](https://developer.apple.com/download/all/?q=xcode) for the latest version.
-- Core ML models and tokenization resources. Please see `--bundle-resources-for-swift-cli` from the [Converting Models to Core ML](#converting-models-to-coreml) section above
-
-If deploying this model to:
-- iPhone
-  - iOS 16.2 or newer
-  - iPhone 12 or newer
-- iPad
-  - iPadOS 16.2 or newer
-  - M1 or newer
-- Mac
-  - macOS 13.1 or newer
-  - M1 or newer
+
+**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
@@ -123,8 +133,10 @@ Please use the `--help` flag to learn about batched generation and more.
 import StableDiffusion
 ...
 let pipeline = try StableDiffusionPipeline(resourcesAt: resourceURL)
+pipeline.loadResources()
 let image = try pipeline.generateImages(prompt: prompt, seed: seed).first
 ```
+On iOS, the `reduceMemory` option should bet set to `true` when constructing `StableDiffusionPipeline`
 
 ### Swift Package Details
 
@@ -184,7 +196,7 @@ Please see [Important Notes on Performance Benchmarks](#important-notes-on-perfo
 - 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 ANE.
+- 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.
 
@@ -208,7 +220,7 @@ Differences may be less or more pronounced for different inputs. Please see the
 
 </details>
 
-## FAQ
+##  <a name="faq"></a> FAQ
 
 <details>
   <summary> Click to expand </summary>
@@ -228,7 +240,7 @@ Differences may be less or more pronounced for different inputs. Please see the
 </details>
 
 <details>
-<summary> <b> Q3: </b> My Mac has 8GB RAM and I am converting models to Core ML using the example command. The process is geting killed because of memory issues. How do I fix this issue? </summary>
+<summary> <b> <a name="low-mem-conversion"></a> Q3: </b> My Mac has 8GB RAM and I am converting models to Core ML using the example command. The process is geting killed because of memory issues. How do I fix this issue? </summary>
 
 <b> A3: </b>  In order to minimize the memory impact of the model conversion process, please execute the following command instead:
 
@@ -257,12 +269,22 @@ python -m python_coreml_stable_diffusion.torch2coreml --convert-unet --chunk-une
 <summary> <b> Q5: </b> Every time I generate an image using the Python pipeline, loading all the Core ML models takes 2-3 minutes. Is this expected? </summary>
 
 <b> A5: </b> Yes and using the Swift library reduces this to just a few seconds. The reason is that `coremltools` loads Core ML models (`.mlpackage`) and each model is compiled to be run on the requested compute unit during load time. Because of the size and number of operations of the unet model, it takes around 2-3 minutes to compile it for Neural Engine execution. Other models should take at most a few seconds. Note that `coremltools` does not cache the compiled model for later loads so each load takes equally long. In order to benefit from compilation caching, `StableDiffusion` Swift package by default relies on compiled Core ML models (`.mlmodelc`) which will be compiled down for the requested compute unit upon first load but then the cache will be reused on subsequent loads until it is purged due to lack of use.
+
+If you intend to use the Python pipeline in an application, we recommend initializing the pipeline once so that the load time is only incurred once. Afterwards, generating images using different prompts and random seeds will not incur the load time for the current session of your application.
+
 </details>
 
+
 <details>
-<summary> <b> Q6: </b> I want to deploy <code>StableDiffusion</code>, the Swift package, in my mobile app. What should I be aware of?" </summary>
+<summary> <b> <a name="q-mobile-app"></a> Q6: </b> I want to deploy <code>StableDiffusion</code>, the Swift package, in my mobile app. What should I be aware of? </summary>
 
-<b> A6: </b> [This section](#swift-requirements) describes the minimum SDK and OS versions as well as the device models supported by this package. In addition to these requirements, for best practice, we recommend testing the package on the device with the least amount of RAM available among your deployment targets. This is due to the fact that `StableDiffusion` consumes approximately 2.6GB of peak memory during runtime while using `.cpuAndNeuralEngine` (the Swift equivalent of `coremltools.ComputeUnit.CPU_AND_NE`). Other compute units may have a higher peak memory consumption so `.cpuAndNeuralEngine` is recommended for iOS and iPadOS deployment (Please refer to this [section](#swift-requirements) for minimum device model requirements). If your app crashes during image generation, please try adding the [Increased Memory Limit](https://developer.apple.com/documentation/bundleresources/entitlements/com_apple_developer_kernel_increased-memory-limit) capability to your Xcode project which should significantly increase your app's memory limit.
+<b> A6: </b>The [Image Generation with Swift](#image-gen-swift) section describes the minimum SDK and OS versions as well as the device models supported by this package. We recommend carefully testing the package on the device with the least amount of RAM available among your deployment targets.
+
+The image generation process in `StableDiffusion` can yield over 2 GB of peak memory during runtime depending on the compute units selected.  On iPadOS, we recommend using `.cpuAndNeuralEngine` in your configuration and the `reduceMemory` option when constructing a `StableDiffusionPipeline` to minimize memory pressure.
+
+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.
 </details>
 
 <details>
@@ -302,4 +324,26 @@ python -m python_coreml_stable_diffusion.torch2coreml --convert-unet --chunk-une
 
 </details>
 
+<details>
+<summary> <b> Q10: </b>  <code> `Could not initialize NNPACK! Reason: Unsupported hardware`  </code> </summary>
+
+<b> A10: </b> This warning is safe to ignore in the context of this repository.
+
+</details>
+
+<details>
+<summary> <b> Q11: </b>  <code> TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect </code> </summary>
+
+<b> A11: </b> This warning is safe to ignore in the context of this repository.
+</details>
+
+<details>
+<summary> <b> Q12: </b>  <code> UserWarning: resource_tracker: There appear to be 1 leaked semaphore objects to clean up at shutdown </code> </summary>
+
+<b> A12: </b> If this warning is printed right after <code> zsh: killed     python -m python_coreml_stable_diffusion.torch2coreml ... </code>, then it is highly likely that your Mac has run out of memory while converting models to Core ML. Please see [Q3](#low-mem-conversion) from above for the solution.
+
+</details>
+
+</details>
+
 </details>

swift/StableDiffusion/pipeline/Decoder.swift

@@ -6,21 +6,30 @@ import CoreML
 import Accelerate
 
 /// A decoder model which produces RGB images from latent samples
-public struct Decoder {
+public struct Decoder: ResourceManaging {
 
     /// VAE decoder model
-    var model: MLModel
+    var model: ManagedMLModel
 
     /// Create decoder from Core ML model
     ///
-    /// - Parameters
-    ///     - model: Core ML model for VAE decoder
-    public init(model: MLModel) {
-        self.model = model
+    /// - Parameters:
+    ///     - url: Location of compiled VAE decoder Core ML model
+    ///     - configuration: configuration to be used when the model is loaded
+    /// - Returns: A decoder that will lazily load its required resources when needed or requested
+    public init(modelAt url: URL, configuration: MLModelConfiguration) {
+        self.model = ManagedMLModel(modelAt: url, configuration: configuration)
     }
 
-    /// Prediction queue
-    let queue = DispatchQueue(label: "decoder.predict")
+    /// 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()
+    }
 
     /// Batch decode latent samples into images
     ///
@@ -42,7 +51,9 @@ public struct Decoder {
         let batch = MLArrayBatchProvider(array: inputs)
 
         // Batch predict with model
-        let results = try queue.sync { try model.predictions(fromBatch: batch) }
+        let results = try model.perform { model in
+            try model.predictions(fromBatch: batch)
+        }
 
         // Transform the outputs to CGImages
         let images: [CGImage] = (0..<results.count).map { i in
@@ -57,7 +68,9 @@ public struct Decoder {
     }
 
     var inputName: String {
-        model.modelDescription.inputDescriptionsByName.first!.key
+        try! model.perform { model in
+            model.modelDescription.inputDescriptionsByName.first!.key
+        }
     }
 
     typealias PixelBufferPFx1 = vImage.PixelBuffer<vImage.PlanarF>

swift/StableDiffusion/pipeline/ManagedMLModel.swift

@@ -0,0 +1,76 @@
+// For licensing see accompanying LICENSE.md file.
+// Copyright (C) 2022 Apple Inc. All Rights Reserved.
+
+import CoreML
+
+/// A class to manage and gate access to a Core ML model
+///
+/// It will automatically load a model into memory when needed or requested
+/// It allows one to request to unload the model from memory
+public final class ManagedMLModel: ResourceManaging {
+
+    /// The location of the model
+    var modelURL: URL
+
+    /// The configuration to be used when the model is loaded
+    var configuration: MLModelConfiguration
+
+    /// The loaded model (when loaded)
+    var loadedModel: MLModel?
+
+    /// Queue to protect access to loaded model
+    var queue: DispatchQueue
+
+    /// Create a managed model given its location and desired loaded configuration
+    ///
+    /// - Parameters:
+    ///     - url: The location of the model
+    ///     - configuration: The configuration to be used when the model is loaded/used
+    /// - Returns: A managed model that has not been loaded
+    public init(modelAt url: URL, configuration: MLModelConfiguration) {
+        self.modelURL = url
+        self.configuration = configuration
+        self.loadedModel = nil
+        self.queue = DispatchQueue(label: "managed.\(url.lastPathComponent)")
+    }
+
+    /// Instantiation and load model into memory
+    public func loadResources() throws {
+        try queue.sync {
+            try loadModel()
+        }
+    }
+
+    /// Unload the model if it was loaded
+    public func unloadResources() {
+        queue.sync {
+            loadedModel = nil
+        }
+    }
+
+    /// Perform an operation with the managed model via a supplied closure.
+    ///  The model will be loaded and supplied to the closure and should only be
+    ///  used within the closure to ensure all resource management is synchronized
+    ///
+    /// - Parameters:
+    ///     - body: Closure which performs and action on a loaded model
+    /// - Returns: The result of the closure
+    /// - Throws: An error if the model cannot be loaded or if the closure throws
+    public func perform<R>(_ body: (MLModel) throws -> R) throws -> R {
+        return try queue.sync {
+            try autoreleasepool {
+                try loadModel()
+                return try body(loadedModel!)
+            }
+        }
+    }
+
+    private func loadModel() throws {
+        if loadedModel == nil {
+            loadedModel = try MLModel(contentsOf: modelURL,
+                                      configuration: configuration)
+        }
+    }
+
+
+}

swift/StableDiffusion/pipeline/ResourceManaging.swift

@@ -0,0 +1,20 @@
+// For licensing see accompanying LICENSE.md file.
+// Copyright (C) 2022 Apple Inc. All Rights Reserved.
+
+/// Protocol for managing internal resources
+public protocol ResourceManaging {
+
+    /// Request resources to be loaded and ready if possible
+    func loadResources() throws
+
+    /// Request resources are unloaded / remove from memory if possible
+    func unloadResources()
+}
+
+extension ResourceManaging {
+    /// Request resources are pre-warmed by loading and unloading
+    func prewarmResources() throws {
+        try loadResources()
+        unloadResources()
+    }
+}