chore: import upstream snapshot with attribution

swift/StableDiffusion/pipeline/CGImage+vImage.swift

@@ -0,0 +1,197 @@
+// For licensing see accompanying LICENSE.md file.
+// Copyright (C) 2022 Apple Inc. All Rights Reserved.
+
+import Foundation
+import Accelerate
+import CoreML
+import CoreGraphics
+
+@available(iOS 16.0, macOS 13.0, *)
+extension CGImage {
+    
+    typealias PixelBufferPFx1 = vImage.PixelBuffer<vImage.PlanarF>
+    typealias PixelBufferP8x3 = vImage.PixelBuffer<vImage.Planar8x3>
+    typealias PixelBufferIFx3 = vImage.PixelBuffer<vImage.InterleavedFx3>
+    typealias PixelBufferI8x3 = vImage.PixelBuffer<vImage.Interleaved8x3>
+    
+    public enum ShapedArrayError: String, Swift.Error {
+        case wrongNumberOfChannels
+        case incorrectFormatsConvertingToShapedArray
+        case vImageConverterNotInitialized
+    }
+    
+    public static func fromShapedArray(_ array: MLShapedArray<Float32>) throws -> CGImage {
+        
+        // array is [N,C,H,W], where C==3
+        let channelCount = array.shape[1]
+        guard channelCount == 3 else {
+            throw ShapedArrayError.wrongNumberOfChannels
+        }
+        
+        let height = array.shape[2]
+        let width = array.shape[3]
+
+        // Normalize each channel into a float between 0 and 1.0
+        let floatChannels = (0..<channelCount).map { i in
+
+            // Normalized channel output
+            let cOut = PixelBufferPFx1(width: width, height:height)
+
+            // Reference this channel in the array and normalize
+            array[0][i].withUnsafeShapedBufferPointer { ptr, _, strides in
+                let cIn = PixelBufferPFx1(data: .init(mutating: ptr.baseAddress!),
+                                          width: width, height: height,
+                                          byteCountPerRow: strides[0]*4)
+                // Map [-1.0 1.0] -> [0.0 1.0]
+                cIn.multiply(by: 0.5, preBias: 1.0, postBias: 0.0, destination: cOut)
+            }
+            return cOut
+        }
+
+        // Convert to interleaved and then to UInt8
+        let floatImage = PixelBufferIFx3(planarBuffers: floatChannels)
+        let uint8Image = PixelBufferI8x3(width: width, height: height)
+        floatImage.convert(to:uint8Image) // maps [0.0 1.0] -> [0 255] and clips
+
+        // Convert to uint8x3 to RGB CGImage (no alpha)
+        let bitmapInfo = CGBitmapInfo(rawValue: CGImageAlphaInfo.none.rawValue)
+        let cgImage = uint8Image.makeCGImage(cgImageFormat:
+                .init(bitsPerComponent: 8,
+                      bitsPerPixel: 3*8,
+                      colorSpace: CGColorSpace(name: CGColorSpace.sRGB) ?? CGColorSpaceCreateDeviceRGB(),
+                      bitmapInfo: bitmapInfo)!)!
+
+        return cgImage
+    }
+    
+    public func planarRGBShapedArray(minValue: Float, maxValue: Float)
+        throws -> MLShapedArray<Float32> {
+            guard
+                var sourceFormat = vImage_CGImageFormat(cgImage: self),
+                var mediumFormat = vImage_CGImageFormat(
+                    bitsPerComponent: 8 * MemoryLayout<UInt8>.size,
+                    bitsPerPixel: 8 * MemoryLayout<UInt8>.size * 4,
+                    colorSpace: CGColorSpaceCreateDeviceRGB(),
+                    bitmapInfo: CGBitmapInfo(rawValue: CGImageAlphaInfo.first.rawValue)),
+                let width = vImagePixelCount(exactly: self.width),
+                let height = vImagePixelCount(exactly: self.height)
+            else {
+                throw ShapedArrayError.incorrectFormatsConvertingToShapedArray
+            }
+
+            var sourceImageBuffer = try vImage_Buffer(cgImage: self)
+            
+            var mediumDestination = try vImage_Buffer(width: Int(width), height: Int(height), bitsPerPixel: mediumFormat.bitsPerPixel)
+            
+            let converter = vImageConverter_CreateWithCGImageFormat(
+                &sourceFormat,
+                &mediumFormat,
+                nil,
+                vImage_Flags(kvImagePrintDiagnosticsToConsole),
+                nil)
+
+            guard let converter = converter?.takeRetainedValue() else {
+                throw ShapedArrayError.vImageConverterNotInitialized
+            }
+            
+            vImageConvert_AnyToAny(converter, &sourceImageBuffer, &mediumDestination, nil, vImage_Flags(kvImagePrintDiagnosticsToConsole))
+            
+            var destinationA = try vImage_Buffer(width: Int(width), height: Int(height), bitsPerPixel: 8 * UInt32(MemoryLayout<Float>.size))
+            var destinationR = try vImage_Buffer(width: Int(width), height: Int(height), bitsPerPixel: 8 * UInt32(MemoryLayout<Float>.size))
+            var destinationG = try vImage_Buffer(width: Int(width), height: Int(height), bitsPerPixel: 8 * UInt32(MemoryLayout<Float>.size))
+            var destinationB = try vImage_Buffer(width: Int(width), height: Int(height), bitsPerPixel: 8 * UInt32(MemoryLayout<Float>.size))
+
+            var minFloat: [Float] = Array(repeating: minValue, count: 4)
+            var maxFloat: [Float] = Array(repeating: maxValue, count: 4)
+            
+            vImageConvert_ARGB8888toPlanarF(&mediumDestination, &destinationA, &destinationR, &destinationG, &destinationB, &maxFloat, &minFloat, .zero)
+           
+            let destAPtr = destinationA.data.assumingMemoryBound(to: Float.self)
+            let destRPtr = destinationR.data.assumingMemoryBound(to: Float.self)
+            let destGPtr = destinationG.data.assumingMemoryBound(to: Float.self)
+            let destBPtr = destinationB.data.assumingMemoryBound(to: Float.self)
+
+            for i in 0..<Int(width) * Int(height) {
+                if destAPtr.advanced(by: i).pointee == 0 {
+                    destRPtr.advanced(by: i).pointee = -1
+                    destGPtr.advanced(by: i).pointee = -1
+                    destBPtr.advanced(by: i).pointee = -1
+                }
+            }
+            
+            let redData = destinationR.unpaddedData()
+            let greenData = destinationG.unpaddedData()
+            let blueData = destinationB.unpaddedData()
+
+            let imageData = redData + greenData + blueData
+
+            let shapedArray = MLShapedArray<Float32>(data: imageData, shape: [1, 3, self.height, self.width])
+
+            return shapedArray
+    }
+
+    private func normalizePixelValues(pixel: UInt8) -> Float {
+        return (Float(pixel) / 127.5) - 1.0
+    }
+
+    public func toRGBShapedArray(minValue: Float, maxValue: Float)
+        throws -> MLShapedArray<Float32> {
+            let image = self
+            let width = image.width
+            let height = image.height
+            let alphaMaskValue: Float = minValue
+
+            guard let colorSpace = CGColorSpace(name: CGColorSpace.sRGB),
+                  let context = CGContext(data: nil, width: width, height: height, bitsPerComponent: 8, bytesPerRow: 4 * width, space: colorSpace, bitmapInfo: CGImageAlphaInfo.premultipliedLast.rawValue),
+                  let ptr = context.data?.bindMemory(to: UInt8.self, capacity: width * height * 4) else {
+                return []
+            }
+
+            context.draw(image, in: CGRect(x: 0, y: 0, width: width, height: height))
+
+            var redChannel = [Float](repeating: 0, count: width * height)
+            var greenChannel = [Float](repeating: 0, count: width * height)
+            var blueChannel = [Float](repeating: 0, count: width * height)
+
+            for y in 0..<height {
+                for x in 0..<width {
+                    let i = 4 * (y * width + x)
+                    if ptr[i+3] == 0 {
+                        // Alpha mask for controlnets
+                        redChannel[y * width + x] = alphaMaskValue
+                        greenChannel[y * width + x] = alphaMaskValue
+                        blueChannel[y * width + x] = alphaMaskValue
+                    } else {
+                        redChannel[y * width + x] = normalizePixelValues(pixel: ptr[i])
+                        greenChannel[y * width + x] = normalizePixelValues(pixel: ptr[i+1])
+                        blueChannel[y * width + x] = normalizePixelValues(pixel: ptr[i+2])
+                    }
+                }
+            }
+
+            let colorShape = [1, 1, height, width]
+            let redShapedArray = MLShapedArray<Float32>(scalars: redChannel, shape: colorShape)
+            let greenShapedArray = MLShapedArray<Float32>(scalars: greenChannel, shape: colorShape)
+            let blueShapedArray = MLShapedArray<Float32>(scalars: blueChannel, shape: colorShape)
+
+            let shapedArray = MLShapedArray<Float32>(concatenating: [redShapedArray, greenShapedArray, blueShapedArray], alongAxis: 1)
+
+            return shapedArray
+    }
+}
+
+extension vImage_Buffer {
+    func unpaddedData() -> Data {
+        let bytesPerPixel = self.rowBytes / Int(self.width)
+        let bytesPerRow = Int(self.width) * bytesPerPixel
+
+        var contiguousPixelData = Data(capacity: bytesPerRow * Int(self.height))
+        for row in 0..<Int(self.height) {
+            let rowStart = self.data!.advanced(by: row * self.rowBytes)
+            let rowData = Data(bytes: rowStart, count: bytesPerRow)
+            contiguousPixelData.append(rowData)
+        }
+
+        return contiguousPixelData
+    }
+}

swift/StableDiffusion/pipeline/ControlNet.swift

@@ -0,0 +1,130 @@
+// For licensing see accompanying LICENSE.md file.
+// Copyright (C) 2022 Apple Inc. All Rights Reserved.
+
+import Foundation
+import CoreML
+import Accelerate
+
+@available(iOS 16.2, macOS 13.1, *)
+public struct ControlNet: ResourceManaging {
+    
+    var models: [ManagedMLModel]
+    
+    public init(modelAt urls: [URL],
+                configuration: MLModelConfiguration) {
+        self.models = urls.map { ManagedMLModel(modelAt: $0, configuration: configuration) }
+    }
+    
+    /// Load resources.
+    public func loadResources() throws {
+        for model in models {
+            try model.loadResources()
+        }
+    }
+
+    /// Unload the underlying model to free up memory
+    public func unloadResources() {
+        for model in models {
+            model.unloadResources()
+        }
+    }
+    
+    /// Pre-warm resources
+    public func prewarmResources() throws {
+        // Override default to pre-warm each model
+        for model in models {
+            try model.loadResources()
+            model.unloadResources()
+        }
+    }
+    
+    var inputImageDescriptions: [MLFeatureDescription] {
+        models.map { model in
+            try! model.perform {
+                $0.modelDescription.inputDescriptionsByName["controlnet_cond"]!
+            }
+        }
+    }
+    
+    /// The expected shape of the models image input
+    public var inputImageShapes: [[Int]] {
+        inputImageDescriptions.map { desc in
+            desc.multiArrayConstraint!.shape.map { $0.intValue }
+        }
+    }
+    
+    /// Calculate additional inputs for Unet to generate intended image following provided images
+    ///
+    /// - Parameters:
+    ///   - latents: Batch of latent samples in an array
+    ///   - timeStep: Current diffusion timestep
+    ///   - hiddenStates: Hidden state to condition on
+    ///   - images: Images for each ControlNet
+    /// - Returns: Array of predicted noise residuals
+    func execute(
+        latents: [MLShapedArray<Float32>],
+        timeStep: Int,
+        hiddenStates: MLShapedArray<Float32>,
+        images: [MLShapedArray<Float32>]
+    ) throws -> [[String: MLShapedArray<Float32>]] {
+        // Match time step batch dimension to the model / latent samples
+        let t = MLShapedArray(scalars: [Float(timeStep), Float(timeStep)], shape: [2])
+        
+        var outputs: [[String: MLShapedArray<Float32>]] = []
+        
+        for (modelIndex, model) in models.enumerated() {
+            let inputs = try latents.map { latent in
+                let dict: [String: Any] = [
+                    "sample": MLMultiArray(latent),
+                    "timestep": MLMultiArray(t),
+                    "encoder_hidden_states": MLMultiArray(hiddenStates),
+                    "controlnet_cond": MLMultiArray(images[modelIndex])
+                ]
+                return try MLDictionaryFeatureProvider(dictionary: dict)
+            }
+            
+            let batch = MLArrayBatchProvider(array: inputs)
+            
+            let results = try model.perform {
+                try $0.predictions(fromBatch: batch)
+            }
+            
+            // pre-allocate MLShapedArray with a specific shape in outputs
+            if outputs.isEmpty {
+                outputs = initOutputs(
+                    batch: latents.count,
+                    shapes: results.features(at: 0).featureValueDictionary
+                )
+            }
+            
+            for n in 0..<results.count {
+                let result = results.features(at: n)
+                for k in result.featureNames {
+                    let newValue = result.featureValue(for: k)!.multiArrayValue!
+                    if modelIndex == 0 {
+                        outputs[n][k] = MLShapedArray<Float32>(newValue)
+                    } else {
+                        let outputArray = MLMultiArray(outputs[n][k]!)
+                        let count = newValue.count
+                        let inputPointer = newValue.dataPointer.assumingMemoryBound(to: Float.self)
+                        let outputPointer = outputArray.dataPointer.assumingMemoryBound(to: Float.self)
+                        vDSP_vadd(inputPointer, 1, outputPointer, 1, outputPointer, 1, vDSP_Length(count))
+                    }
+                }
+            }
+        }
+        
+        return outputs
+    }
+    
+    private func initOutputs(batch: Int, shapes: [String: MLFeatureValue]) -> [[String: MLShapedArray<Float32>]] {
+        var output: [String: MLShapedArray<Float32>] = [:]
+        for (outputName, featureValue) in shapes {
+            output[outputName] = MLShapedArray<Float32>(
+                repeating: 0.0,
+                shape: featureValue.multiArrayValue!.shape.map { $0.intValue }
+            )
+        }
+        return Array(repeating: output, count: batch)
+    }
+}

swift/StableDiffusion/pipeline/DPMSolverMultistepScheduler.swift

@@ -0,0 +1,273 @@
+// For licensing see accompanying LICENSE.md file.
+// Copyright (C) 2022 Apple Inc. and The HuggingFace Team. All Rights Reserved.
+
+import Accelerate
+import CoreML
+
+/// How to space timesteps for inference
+public enum TimeStepSpacing {
+    case linspace
+    case leading
+    case karras
+}
+
+/// A scheduler used to compute a de-noised image
+///
+///  This implementation matches:
+///  [Hugging Face Diffusers DPMSolverMultistepScheduler](https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py)
+///
+/// It uses the DPM-Solver++ algorithm: [code](https://github.com/LuChengTHU/dpm-solver) [paper](https://arxiv.org/abs/2211.01095).
+/// Limitations:
+///  - Only implemented for DPM-Solver++ algorithm (not DPM-Solver).
+///  - Second order only.
+///  - Assumes the model predicts epsilon.
+///  - No dynamic thresholding.
+///  - `midpoint` solver algorithm.
+@available(iOS 16.2, macOS 13.1, *)
+public final class DPMSolverMultistepScheduler: Scheduler {
+    public let trainStepCount: Int
+    public let inferenceStepCount: Int
+    public let betas: [Float]
+    public let alphas: [Float]
+    public let alphasCumProd: [Float]
+    public let timeSteps: [Int]
+
+    public let alpha_t: [Float]
+    public let sigma_t: [Float]
+    public let lambda_t: [Float]
+    
+    public let solverOrder = 2
+    private(set) var lowerOrderStepped = 0
+    
+    private var usingKarrasSigmas = false
+
+    /// Whether to use lower-order solvers in the final steps. Only valid for less than 15 inference steps.
+    /// We empirically find this trick can stabilize the sampling of DPM-Solver, especially with 10 or fewer steps.
+    public let useLowerOrderFinal = true
+    
+    // Stores solverOrder (2) items
+    public private(set) var modelOutputs: [MLShapedArray<Float32>] = []
+
+    /// Create a scheduler that uses a second order DPM-Solver++ algorithm.
+    ///
+    /// - Parameters:
+    ///   - stepCount: Number of inference steps to schedule
+    ///   - trainStepCount: Number of training diffusion steps
+    ///   - betaSchedule: Method to schedule betas from betaStart to betaEnd
+    ///   - betaStart: The starting value of beta for inference
+    ///   - betaEnd: The end value for beta for inference
+    ///   - timeStepSpacing: How to space time steps
+    /// - Returns: A scheduler ready for its first step
+    public init(
+        stepCount: Int = 50,
+        trainStepCount: Int = 1000,
+        betaSchedule: BetaSchedule = .scaledLinear,
+        betaStart: Float = 0.00085,
+        betaEnd: Float = 0.012,
+        timeStepSpacing: TimeStepSpacing = .linspace
+    ) {
+        self.trainStepCount = trainStepCount
+        self.inferenceStepCount = stepCount
+        
+        switch betaSchedule {
+        case .linear:
+            self.betas = linspace(betaStart, betaEnd, trainStepCount)
+        case .scaledLinear:
+            self.betas = linspace(pow(betaStart, 0.5), pow(betaEnd, 0.5), trainStepCount).map({ $0 * $0 })
+        }
+        
+        self.alphas = betas.map({ 1.0 - $0 })
+        var alphasCumProd = self.alphas
+        for i in 1..<alphasCumProd.count {
+            alphasCumProd[i] *= alphasCumProd[i -  1]
+        }
+        self.alphasCumProd = alphasCumProd
+
+        switch timeStepSpacing {
+        case .linspace:
+            self.timeSteps = linspace(0, Float(self.trainStepCount-1), stepCount+1).dropFirst().reversed().map { Int(round($0)) }
+            self.alpha_t = vForce.sqrt(self.alphasCumProd)
+            self.sigma_t = vForce.sqrt(vDSP.subtract([Float](repeating: 1, count: self.alphasCumProd.count), self.alphasCumProd))
+        case .leading:
+            let lastTimeStep = trainStepCount - 1
+            let stepRatio = lastTimeStep / (stepCount + 1)
+            // Creates integer timesteps by multiplying by ratio
+            self.timeSteps = (0...stepCount).map { 1 + $0 * stepRatio }.dropFirst().reversed()
+            self.alpha_t = vForce.sqrt(self.alphasCumProd)
+            self.sigma_t = vForce.sqrt(vDSP.subtract([Float](repeating: 1, count: self.alphasCumProd.count), self.alphasCumProd))
+        case .karras:
+            // sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5)
+            let scaled = vDSP.multiply(
+                subtraction: ([Float](repeating: 1, count: self.alphasCumProd.count), self.alphasCumProd),
+                subtraction: (vDSP.divide(1, self.alphasCumProd), [Float](repeating: 0, count: self.alphasCumProd.count))
+            )
+            let sigmas = vForce.sqrt(scaled)
+            let logSigmas = sigmas.map { log($0) }
+
+            let sigmaMin = sigmas.first!
+            let sigmaMax = sigmas.last!
+            let rho: Float = 7
+            let ramp = linspace(0, 1, stepCount)
+            let minInvRho = pow(sigmaMin, (1 / rho))
+            let maxInvRho = pow(sigmaMax, (1 / rho))
+
+            var karrasSigmas = ramp.map { pow(maxInvRho + $0 * (minInvRho - maxInvRho), rho) }
+            let karrasTimeSteps = karrasSigmas.map { sigmaToTimestep(sigma: $0, logSigmas: logSigmas) }
+            self.timeSteps = karrasTimeSteps
+
+            karrasSigmas.append(karrasSigmas.last!)
+
+            self.alpha_t = vDSP.divide(1, vForce.sqrt(vDSP.add(1, vDSP.square(karrasSigmas))))
+            self.sigma_t = vDSP.multiply(karrasSigmas, self.alpha_t)
+            usingKarrasSigmas = true
+        }
+
+        self.lambda_t = zip(self.alpha_t, self.sigma_t).map { α, σ in log(α) - log(σ) }
+    }
+    
+    func timestepToIndex(_ timestep: Int) -> Int {
+        guard usingKarrasSigmas else { return timestep }
+        return self.timeSteps.firstIndex(of: timestep) ?? 0
+    }
+    
+    /// Convert the model output to the corresponding type the algorithm needs.
+    /// This implementation is for second-order DPM-Solver++ assuming epsilon prediction.
+    func convertModelOutput(modelOutput: MLShapedArray<Float32>, timestep: Int, sample: MLShapedArray<Float32>) -> MLShapedArray<Float32> {
+        assert(modelOutput.scalarCount == sample.scalarCount)
+        let scalarCount = modelOutput.scalarCount
+        let sigmaIndex = timestepToIndex(timestep)
+        let (alpha_t, sigma_t) = (self.alpha_t[sigmaIndex], self.sigma_t[sigmaIndex])
+
+        return MLShapedArray(unsafeUninitializedShape: modelOutput.shape) { scalars, _ in
+            assert(scalars.count == scalarCount)
+            modelOutput.withUnsafeShapedBufferPointer { modelOutput, _, _ in
+                sample.withUnsafeShapedBufferPointer { sample, _, _ in
+                    for i in 0 ..< scalarCount {
+                        scalars.initializeElement(at: i, to: (sample[i] - modelOutput[i] * sigma_t) / alpha_t)
+                    }
+                }
+            }
+        }
+    }
+
+    /// One step for the first-order DPM-Solver (equivalent to DDIM).
+    /// See https://arxiv.org/abs/2206.00927 for the detailed derivation.
+    /// var names and code structure mostly follow https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py
+    func firstOrderUpdate(
+        modelOutput: MLShapedArray<Float32>,
+        timestep: Int,
+        prevTimestep: Int,
+        sample: MLShapedArray<Float32>
+    ) -> MLShapedArray<Float32> {
+        let prevIndex = timestepToIndex(prevTimestep)
+        let currIndex = timestepToIndex(timestep)
+        let (p_lambda_t, lambda_s) = (Double(lambda_t[prevIndex]), Double(lambda_t[currIndex]))
+        let p_alpha_t = Double(alpha_t[prevIndex])
+        let (p_sigma_t, sigma_s) = (Double(sigma_t[prevIndex]), Double(sigma_t[currIndex]))
+        let h = p_lambda_t - lambda_s
+        // x_t = (sigma_t / sigma_s) * sample - (alpha_t * (torch.exp(-h) - 1.0)) * model_output
+        let x_t = weightedSum(
+            [p_sigma_t / sigma_s, -p_alpha_t * (exp(-h) - 1)],
+            [sample, modelOutput]
+        )
+        return x_t
+    }
+
+    /// One step for the second-order multistep DPM-Solver++ algorithm, using the midpoint method.
+    /// var names and code structure mostly follow https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py
+    func secondOrderUpdate(
+        modelOutputs: [MLShapedArray<Float32>],
+        timesteps: [Int],
+        prevTimestep t: Int,
+        sample: MLShapedArray<Float32>
+    ) -> MLShapedArray<Float32> {
+        let (s0, s1) = (timesteps[back: 1], timesteps[back: 2])
+        let (m0, m1) = (modelOutputs[back: 1], modelOutputs[back: 2])
+        let (p_lambda_t, lambda_s0, lambda_s1) = (
+            Double(lambda_t[timestepToIndex(t)]),
+            Double(lambda_t[timestepToIndex(s0)]),
+            Double(lambda_t[timestepToIndex(s1)])
+        )
+        let p_alpha_t = Double(alpha_t[timestepToIndex(t)])
+        let (p_sigma_t, sigma_s0) = (Double(sigma_t[timestepToIndex(t)]), Double(sigma_t[timestepToIndex(s0)]))
+        let (h, h_0) = (p_lambda_t - lambda_s0, lambda_s0 - lambda_s1)
+        let r0 = h_0 / h
+        let D0 = m0
+        
+        // D1 = (1.0 / r0) * (m0 - m1)
+        let D1 = weightedSum(
+            [1/r0, -1/r0],
+            [m0, m1]
+        )
+        
+        // See https://arxiv.org/abs/2211.01095 for detailed derivations
+        // x_t = (
+        //     (sigma_t / sigma_s0) * sample
+        //     - (alpha_t * (torch.exp(-h) - 1.0)) * D0
+        //     - 0.5 * (alpha_t * (torch.exp(-h) - 1.0)) * D1
+        // )
+        let x_t = weightedSum(
+            [p_sigma_t/sigma_s0, -p_alpha_t * (exp(-h) - 1), -0.5 * p_alpha_t * (exp(-h) - 1)],
+            [sample, D0, D1]
+        )
+        return x_t
+    }
+
+    public func step(output: MLShapedArray<Float32>, timeStep t: Int, sample: MLShapedArray<Float32>) -> MLShapedArray<Float32> {
+        let stepIndex = timeSteps.firstIndex(of: t) ?? timeSteps.count - 1
+        let prevTimestep = stepIndex == timeSteps.count - 1 ? 0 : timeSteps[stepIndex + 1]
+
+        let lowerOrderFinal = useLowerOrderFinal && stepIndex == timeSteps.count - 1 && timeSteps.count < 15
+        let lowerOrderSecond = useLowerOrderFinal && stepIndex == timeSteps.count - 2 && timeSteps.count < 15
+        let lowerOrder = lowerOrderStepped < 1 || lowerOrderFinal || lowerOrderSecond
+        
+        let modelOutput = convertModelOutput(modelOutput: output, timestep: t, sample: sample)
+        if modelOutputs.count == solverOrder { modelOutputs.removeFirst() }
+        modelOutputs.append(modelOutput)
+        
+        let prevSample: MLShapedArray<Float32>
+        if lowerOrder {
+            prevSample = firstOrderUpdate(modelOutput: modelOutput, timestep: t, prevTimestep: prevTimestep, sample: sample)
+        } else {
+            prevSample = secondOrderUpdate(
+                modelOutputs: modelOutputs,
+                timesteps: [timeSteps[stepIndex - 1], t],
+                prevTimestep: prevTimestep,
+                sample: sample
+            )
+        }
+        if lowerOrderStepped < solverOrder {
+            lowerOrderStepped += 1
+        }
+        
+        return prevSample
+    }
+}
+
+func sigmaToTimestep(sigma: Float, logSigmas: [Float]) -> Int {
+    let logSigma = log(sigma)
+    let dists = logSigmas.map { logSigma - $0 }
+
+    // last index that is not negative, clipped to last index - 1
+    var lowIndex = dists.reduce(-1) { partialResult, dist in
+        return dist >= 0 && partialResult < dists.endIndex-2 ? partialResult + 1 : partialResult
+    }
+    lowIndex = max(lowIndex, 0)
+    let highIndex = lowIndex + 1
+
+    let low = logSigmas[lowIndex]
+    let high = logSigmas[highIndex]
+
+    // Interpolate sigmas
+    let w = ((low - logSigma) / (low - high)).clipped(to: 0...1)
+
+    // transform interpolated value to time range
+    let t = (1 - w) * Float(lowIndex) + w * Float(highIndex)
+    return Int(round(t))
+}
+
+extension FloatingPoint {
+    func clipped(to range: ClosedRange<Self>) -> Self {
+        return min(max(self, range.lowerBound), range.upperBound)
+    }
+}

swift/StableDiffusion/pipeline/Decoder.swift

@@ -0,0 +1,80 @@
+// For licensing see accompanying LICENSE.md file.
+// Copyright (C) 2024 Apple Inc. All Rights Reserved.
+
+import Foundation
+import CoreML
+
+/// A decoder model which produces RGB images from latent samples
+@available(iOS 16.2, macOS 13.1, *)
+public struct Decoder: ResourceManaging {
+
+    /// VAE decoder model
+    var model: ManagedMLModel
+
+    /// Create decoder from Core ML 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)
+    }
+
+    /// 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
+    ///
+    ///  - Parameters:
+    ///    - latents: Batch of latent samples to decode
+    ///    - scaleFactor: scalar divisor on latents before decoding
+    ///  - Returns: decoded images
+    public func decode(
+        _ latents: [MLShapedArray<Float32>],
+        scaleFactor: Float32,
+        shiftFactor: Float32 = 0.0
+    ) throws -> [CGImage] {
+
+        // Form batch inputs for model
+        let inputs: [MLFeatureProvider] = try latents.map { sample in
+            // Reference pipeline scales the latent samples before decoding
+            let sampleScaled = MLShapedArray<Float32>(
+                scalars: sample.scalars.map { $0 / scaleFactor + shiftFactor },
+                shape: sample.shape)
+
+            let dict = [inputName: MLMultiArray(sampleScaled)]
+            return try MLDictionaryFeatureProvider(dictionary: dict)
+        }
+        let batch = MLArrayBatchProvider(array: inputs)
+
+        // Batch predict with model
+        let results = try model.perform { model in
+            try model.predictions(fromBatch: batch)
+        }
+
+        // Transform the outputs to CGImages
+        let images: [CGImage] = try (0..<results.count).map { i in
+            let result = results.features(at: i)
+            let outputName = result.featureNames.first!
+            let output = result.featureValue(for: outputName)!.multiArrayValue!
+            return try CGImage.fromShapedArray(MLShapedArray<Float32>(converting: output))
+        }
+
+        return images
+    }
+
+    var inputName: String {
+        try! model.perform { model in
+            model.modelDescription.inputDescriptionsByName.first!.key
+        }
+    }
+
+}

swift/StableDiffusion/pipeline/DiscreteFlowScheduler.swift

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

swift/StableDiffusion/pipeline/Encoder.swift

@@ -0,0 +1,108 @@
+// For licensing see accompanying LICENSE.md file.
+// Copyright (C) 2022 Apple Inc. All Rights Reserved.
+
+import Foundation
+import CoreML
+
+/// A encoder model which produces latent samples from RGB images
+@available(iOS 16.2, macOS 13.1, *)
+public struct Encoder: ResourceManaging {
+    
+    public enum Error: String, Swift.Error {
+        case sampleInputShapeNotCorrect
+    }
+    
+    /// VAE encoder model + post math and adding noise from schedular
+    var model: ManagedMLModel
+    
+    /// Create encoder from Core ML model
+    ///
+    /// - Parameters:
+    ///     - url: Location of compiled VAE encoder Core ML model
+    ///     - configuration: configuration to be used when the model is loaded
+    /// - Returns: An encoder 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)
+    }
+    
+    /// 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()
+    }
+    
+    /// Prediction queue
+    let queue = DispatchQueue(label: "encoder.predict")
+
+    /// Encode image into latent sample
+    ///
+    ///  - Parameters:
+    ///    - image: Input image
+    ///    - scaleFactor: scalar multiplier on latents before encoding image
+    ///    - random
+    ///  - Returns: The encoded latent space as MLShapedArray
+    public func encode(
+        _ image: CGImage,
+        scaleFactor: Float32,
+        random: inout RandomSource
+    ) throws -> MLShapedArray<Float32> {
+        let imageData = try image.planarRGBShapedArray(minValue: -1.0, maxValue: 1.0)
+        guard imageData.shape == inputShape else {
+            // TODO: Consider auto resizing and croping similar to how Vision or CoreML auto-generated Swift code can accomplish with `MLFeatureValue`
+            throw Error.sampleInputShapeNotCorrect
+        }
+        let dict = [inputName: MLMultiArray(imageData)]
+        let input = try MLDictionaryFeatureProvider(dictionary: dict)
+        
+        let result = try model.perform { model in
+            try model.prediction(from: input)
+        }
+        let outputName = result.featureNames.first!
+        let outputValue = result.featureValue(for: outputName)!.multiArrayValue!
+        let output = MLShapedArray<Float32>(converting: outputValue)
+        
+        // DiagonalGaussianDistribution
+        let mean = output[0][0..<4]
+        let logvar = MLShapedArray<Float32>(
+            scalars: output[0][4..<8].scalars.map { min(max($0, -30), 20) },
+            shape: mean.shape
+        )
+        let std = MLShapedArray<Float32>(
+            scalars: logvar.scalars.map { exp(0.5 * $0) },
+            shape: logvar.shape
+        )
+        let latent = MLShapedArray<Float32>(
+            scalars: zip(mean.scalars, std.scalars).map {
+                Float32(random.nextNormal(mean: Double($0), stdev: Double($1)))
+            },
+            shape: logvar.shape
+        )
+        
+        // Reference pipeline scales the latent after encoding
+        let latentScaled = MLShapedArray<Float32>(
+            scalars: latent.scalars.map { $0 * scaleFactor },
+            shape: [1] + latent.shape
+        )
+
+        return latentScaled
+    }
+    
+    var inputDescription: MLFeatureDescription {
+        try! model.perform { model in
+            model.modelDescription.inputDescriptionsByName.first!.value
+        }
+    }
+    
+    var inputName: String {
+        inputDescription.name
+    }
+    
+    /// The expected shape of the models latent sample input
+    var inputShape: [Int] {
+        inputDescription.multiArrayConstraint!.shape.map { $0.intValue }
+    }
+}

swift/StableDiffusion/pipeline/ManagedMLModel.swift

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

swift/StableDiffusion/pipeline/MultiModalDiffusionTransformer.swift

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

swift/StableDiffusion/pipeline/MultilingualTextEncoder.swift

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

swift/StableDiffusion/pipeline/NumPyRandomSource.swift

@@ -0,0 +1,119 @@
+// For licensing see accompanying LICENSE.md file.
+// Copyright (C) 2022 Apple Inc. All Rights Reserved.
+
+import Foundation
+import CoreML
+
+/// A random source consistent with NumPy
+///
+///  This implementation matches:
+///  [NumPy's older randomkit.c](https://github.com/numpy/numpy/blob/v1.0/numpy/random/mtrand/randomkit.c)
+///
+@available(iOS 16.2, macOS 13.1, *)
+struct NumPyRandomSource: RandomNumberGenerator, RandomSource {
+
+    struct State {
+        var key = [UInt32](repeating: 0, count: 624)
+        var pos: Int = 0
+        var nextGauss: Double? = nil
+    }
+
+    var state: State
+
+    /// Initialize with a random seed
+    ///
+    /// - Parameters
+    ///     - seed: Seed for underlying Mersenne Twister 19937 generator
+    /// - Returns random source
+    init(seed: UInt32) {
+        state = .init()
+        var s = seed & 0xffffffff
+        for i in 0 ..< state.key.count {
+            state.key[i] = s
+            s = UInt32((UInt64(1812433253) * UInt64(s ^ (s >> 30)) + UInt64(i) + 1) & 0xffffffff)
+        }
+        state.pos = state.key.count
+        state.nextGauss = nil
+    }
+
+    /// Generate next UInt32 using fast 32bit Mersenne Twister
+    mutating func nextUInt32() -> UInt32 {
+        let n = 624
+        let m = 397
+        let matrixA: UInt64    = 0x9908b0df
+        let upperMask: UInt32  = 0x80000000
+        let lowerMask: UInt32  = 0x7fffffff
+
+        var y: UInt32
+        if state.pos == state.key.count {
+            for i in 0 ..< (n - m) {
+                y = (state.key[i] & upperMask) | (state.key[i + 1] & lowerMask)
+                state.key[i] = state.key[i + m] ^ (y >> 1) ^ UInt32((UInt64(~(y & 1)) + 1) & matrixA)
+            }
+            for i in (n - m) ..< (n - 1) {
+                y = (state.key[i] & upperMask) | (state.key[i + 1] & lowerMask)
+                state.key[i] = state.key[i + (m - n)] ^ (y >> 1) ^ UInt32((UInt64(~(y & 1)) + 1) & matrixA)
+            }
+            y = (state.key[n - 1] & upperMask) | (state.key[0] & lowerMask)
+            state.key[n - 1] = state.key[m - 1] ^ (y >> 1) ^ UInt32((UInt64(~(y & 1)) + 1) & matrixA)
+            state.pos = 0
+        }
+        y = state.key[state.pos]
+        state.pos += 1
+
+        y ^= (y >> 11)
+        y ^= (y << 7) & 0x9d2c5680
+        y ^= (y << 15) & 0xefc60000
+        y ^= (y >> 18)
+
+        return y
+    }
+
+    mutating func next() -> UInt64 {
+        let low = nextUInt32()
+        let high = nextUInt32()
+        return (UInt64(high) << 32) | UInt64(low)
+    }
+
+    /// Generate next random double value
+    mutating func nextDouble() -> Double {
+        let a = Double(nextUInt32() >> 5)
+        let b = Double(nextUInt32() >> 6)
+        return (a * 67108864.0 + b) / 9007199254740992.0
+    }
+
+    /// Generate next random value from a standard normal
+    mutating func nextGauss() -> Double {
+        if let nextGauss = state.nextGauss {
+            state.nextGauss = nil
+            return nextGauss
+        }
+        var x1, x2, r2: Double
+        repeat {
+            x1 = 2.0 * nextDouble() - 1.0
+            x2 = 2.0 * nextDouble() - 1.0
+            r2 = x1 * x1 + x2 * x2
+        } while r2 >= 1.0 || r2 == 0.0
+
+        // Box-Muller transform
+        let f = sqrt(-2.0 * log(r2) / r2)
+        state.nextGauss = f * x1
+        return f * x2
+    }
+
+    /// Generates a random value from a normal distribution with given mean and standard deviation.
+    mutating func nextNormal(mean: Double = 0.0, stdev: Double = 1.0) -> Double {
+        nextGauss() * stdev + mean
+    }
+
+    /// Generates an array of random values from a normal distribution with given mean and standard deviation.
+    mutating func normalArray(count: Int, mean: Double = 0.0, stdev: Double = 1.0) -> [Double] {
+        (0 ..< count).map { _ in nextNormal(mean: mean, stdev: stdev) }
+    }
+
+    /// Generate a shaped array with scalars from a normal distribution with given mean and standard deviation.
+    mutating func normalShapedArray(_ shape: [Int], mean: Double = 0.0, stdev: Double = 1.0) -> MLShapedArray<Double> {
+        let count = shape.reduce(1, *)
+        return .init(scalars: normalArray(count: count, mean: mean, stdev: stdev), shape: shape)
+    }
+}

swift/StableDiffusion/pipeline/NvRandomSource.swift

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

swift/StableDiffusion/pipeline/RandomSource.swift

@@ -0,0 +1,8 @@
+import CoreML
+
+@available(iOS 16.2, macOS 13.1, *)
+public protocol RandomSource {
+    mutating func nextNormal(mean: Double, stdev: Double) -> Double
+    
+    mutating func normalShapedArray(_ shape: [Int], mean: Double, stdev: Double) -> MLShapedArray<Double>
+}

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()
+    }
+}

swift/StableDiffusion/pipeline/SafetyChecker.swift

@@ -0,0 +1,167 @@
+// For licensing see accompanying LICENSE.md file.
+// Copyright (C) 2022 Apple Inc. All Rights Reserved.
+
+import Foundation
+import CoreML
+import Accelerate
+
+/// Image safety checking model
+@available(iOS 16.2, macOS 13.1, *)
+public struct SafetyChecker: ResourceManaging {
+
+    /// Safety checking Core ML model
+    var model: ManagedMLModel
+
+    /// Creates safety checker
+    ///
+    /// - Parameters:
+    ///     - url: Location of compiled safety checking  Core ML model
+    ///     - configuration: configuration to be used when the model is loaded
+    /// - Returns: A safety cherker 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)
+    }
+
+    /// 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()
+    }
+
+    typealias PixelBufferPFx1 = vImage.PixelBuffer<vImage.PlanarF>
+    typealias PixelBufferP8x1 = vImage.PixelBuffer<vImage.Planar8>
+    typealias PixelBufferPFx3 = vImage.PixelBuffer<vImage.PlanarFx3>
+    typealias PixelBufferP8x3 = vImage.PixelBuffer<vImage.Planar8x3>
+    typealias PixelBufferIFx3 = vImage.PixelBuffer<vImage.InterleavedFx3>
+    typealias PixelBufferI8x3 = vImage.PixelBuffer<vImage.Interleaved8x3>
+    typealias PixelBufferI8x4 = vImage.PixelBuffer<vImage.Interleaved8x4>
+
+    enum SafetyCheckError: Error {
+        case imageResizeFailure
+        case imageToFloatFailure
+        case modelInputFailure
+        case unexpectedModelOutput
+    }
+
+    /// Check if image is safe
+    ///
+    /// - Parameters:
+    ///     - image: Image to check
+    /// - Returns: Whether the model considers the image to be safe
+    public func isSafe(_ image: CGImage) throws -> Bool {
+
+        let inputName = "clip_input"
+        let adjustmentName = "adjustment"
+        let imagesNames = "images"
+
+        let inputInfo = try model.perform { model in
+            model.modelDescription.inputDescriptionsByName
+        }
+        let inputShape = inputInfo[inputName]!.multiArrayConstraint!.shape
+
+        let width = inputShape[2].intValue
+        let height = inputShape[3].intValue
+
+        let resizedImage = try resizeToRGBA(image, width: width, height: height)
+
+        let bufferP8x3 = try getRGBPlanes(of: resizedImage)
+
+        let arrayPFx3 = normalizeToFloatShapedArray(bufferP8x3)
+
+        guard let input = try? MLDictionaryFeatureProvider(
+            dictionary:[
+                // Input that is analyzed for safety
+                inputName      : MLMultiArray(arrayPFx3),
+                // No adjustment, use default threshold
+                adjustmentName : MLMultiArray(MLShapedArray<Float32>(scalars: [0], shape: [1])),
+                // Supplying dummy images to be filtered (will be ignored)
+                imagesNames    : MLMultiArray(shape:[1, 512, 512, 3], dataType: .float16)
+            ]
+        ) else {
+            throw SafetyCheckError.modelInputFailure
+        }
+
+        let result = try model.perform { model in
+            try model.prediction(from: input)
+        }
+
+        let output = result.featureValue(for: "has_nsfw_concepts")
+
+        guard let unsafe = output?.multiArrayValue?[0].boolValue else {
+            throw SafetyCheckError.unexpectedModelOutput
+        }
+
+        return !unsafe
+    }
+
+    func resizeToRGBA(_ image: CGImage,
+                      width: Int, height: Int) throws -> CGImage {
+
+        guard let context = CGContext(
+            data: nil,
+            width: width,
+            height: height,
+            bitsPerComponent: 8,
+            bytesPerRow: width*4,
+            space: CGColorSpaceCreateDeviceRGB(),
+            bitmapInfo: CGImageAlphaInfo.noneSkipLast.rawValue) else {
+            throw SafetyCheckError.imageResizeFailure
+        }
+
+        context.interpolationQuality = .high
+        context.draw(image, in: CGRect(x: 0, y: 0, width: width, height: height))
+        guard let resizedImage = context.makeImage() else {
+            throw SafetyCheckError.imageResizeFailure
+        }
+
+        return resizedImage
+    }
+
+    func getRGBPlanes(of rgbaImage: CGImage) throws -> PixelBufferP8x3 {
+        // Reference as interleaved 8 bit vImage PixelBuffer
+        var emptyFormat = vImage_CGImageFormat()
+        guard let bufferI8x4 = try? PixelBufferI8x4(
+            cgImage: rgbaImage,
+            cgImageFormat:&emptyFormat) else {
+            throw SafetyCheckError.imageToFloatFailure
+        }
+
+        // Drop the alpha channel, keeping RGB
+        let bufferI8x3 = PixelBufferI8x3(width: rgbaImage.width, height:rgbaImage.height)
+        bufferI8x4.convert(to: bufferI8x3, channelOrdering: .RGBA)
+
+        // De-interleave into 8-bit planes
+        return PixelBufferP8x3(interleavedBuffer: bufferI8x3)
+    }
+
+    func normalizeToFloatShapedArray(_ bufferP8x3: PixelBufferP8x3) -> MLShapedArray<Float32> {
+        let width = bufferP8x3.width
+        let height = bufferP8x3.height
+
+        let means = [0.485, 0.456, 0.406] as [Float]
+        let stds  = [0.229, 0.224, 0.225] as [Float]
+
+        // Convert to normalized float 1x3xWxH input (plannar)
+        let arrayPFx3 = MLShapedArray<Float32>(repeating: 0.0, shape: [1, 3, width, height])
+        for c in 0..<3 {
+            arrayPFx3[0][c].withUnsafeShapedBufferPointer { ptr, _, strides in
+                let floatChannel = PixelBufferPFx1(data: .init(mutating: ptr.baseAddress!),
+                                                   width: width, height: height,
+                                                   byteCountPerRow: strides[0]*4)
+
+                bufferP8x3.withUnsafePixelBuffer(at: c) { uint8Channel in
+                    uint8Channel.convert(to: floatChannel) // maps [0 255] -> [0 1]
+                    floatChannel.multiply(by: 1.0/stds[c],
+                                          preBias: -means[c],
+                                          postBias: 0.0,
+                                          destination: floatChannel)
+                }
+            }
+        }
+        return arrayPFx3
+    }
+}

swift/StableDiffusion/pipeline/SampleTimer.swift

@@ -0,0 +1,78 @@
+// For licensing see accompanying LICENSE.md file.
+// Copyright (C) 2022 Apple Inc. All Rights Reserved.
+
+import Foundation
+
+/// A utility for timing events and tracking time statistics
+///
+/// Typical usage
+/// ```
+/// let timer: SampleTimer
+///
+/// for i in 0...<iterationCount {
+///     timer.start()
+///     doStuff()
+///     timer.stop()
+/// }
+///
+/// print(String(format: "mean: %.2f, var: %.2f",
+///              timer.mean, timer.variance))
+/// ```
+@available(iOS 16.2, macOS 13.1, *)
+public final class SampleTimer: Codable {
+    var startTime: CFAbsoluteTime?
+    var sum: Double = 0.0
+    var sumOfSquares: Double = 0.0
+    var count = 0
+    var samples: [Double] = []
+
+    public init() {}
+
+    /// Start a sample, noting the current time
+    public func start() {
+        startTime = CFAbsoluteTimeGetCurrent()
+    }
+
+    // Stop a sample and record the elapsed time
+    @discardableResult public func stop() -> Double {
+        guard let startTime = startTime else {
+            return 0
+        }
+
+        let elapsed = CFAbsoluteTimeGetCurrent() - startTime
+        sum += elapsed
+        sumOfSquares += elapsed * elapsed
+        count += 1
+        samples.append(elapsed)
+        return elapsed
+    }
+
+    /// Mean of all sampled times
+    public var mean: Double { sum / Double(count) }
+
+    /// Variance of all sampled times
+    public var variance: Double {
+        guard count > 1 else {
+            return 0.0
+        }
+        return sumOfSquares / Double(count - 1) - mean * mean
+    }
+
+    /// Standard deviation of all sampled times
+    public var stdev: Double { variance.squareRoot() }
+
+    /// Median of all sampled times
+    public var median: Double {
+        let sorted = samples.sorted()
+        let (q, r) = sorted.count.quotientAndRemainder(dividingBy: 2)
+        if r == 0 {
+            return (sorted[q] + sorted[q - 1]) / 2.0
+        } else {
+            return Double(sorted[q])
+        }
+    }
+
+    public var allSamples: [Double] {
+        samples
+    }
+}

swift/StableDiffusion/pipeline/Scheduler.swift

@@ -0,0 +1,363 @@
+// For licensing see accompanying LICENSE.md file.
+// Copyright (C) 2022 Apple Inc. All Rights Reserved.
+
+import Accelerate
+import CoreML
+
+@available(iOS 16.2, macOS 13.1, *)
+public protocol Scheduler {
+    /// Number of diffusion steps performed during training
+    var trainStepCount: Int { get }
+
+    /// Number of inference steps to be performed
+    var inferenceStepCount: Int { get }
+
+    /// Training diffusion time steps index by inference time step
+    var timeSteps: [Int] { get }
+
+    /// Training diffusion time steps index by inference time step
+    func calculateTimesteps(strength: Float?) -> [Int]
+
+    /// Schedule of betas which controls the amount of noise added at each timestep
+    var betas: [Float] { get }
+
+    /// 1 - betas
+    var alphas: [Float] { get }
+
+    /// Cached cumulative product of alphas
+    var alphasCumProd: [Float] { get }
+
+    /// Standard deviation of the initial noise distribution
+    var initNoiseSigma: Float { get }
+
+    /// Denoised latents
+    var modelOutputs: [MLShapedArray<Float32>] { get }
+
+    /// Compute a de-noised image sample and step scheduler state
+    ///
+    /// - Parameters:
+    ///   - output: The predicted residual noise output of learned diffusion model
+    ///   - timeStep: The current time step in the diffusion chain
+    ///   - sample: The current input sample to the diffusion model
+    /// - Returns: Predicted de-noised sample at the previous time step
+    /// - Postcondition: The scheduler state is updated.
+    ///   The state holds the current sample and history of model output noise residuals
+    func step(
+        output: MLShapedArray<Float32>,
+        timeStep t: Int,
+        sample s: MLShapedArray<Float32>
+    ) -> MLShapedArray<Float32>
+}
+
+@available(iOS 16.2, macOS 13.1, *)
+public extension Scheduler {
+    var initNoiseSigma: Float { 1 }
+}
+
+@available(iOS 16.2, macOS 13.1, *)
+public extension Scheduler {
+    /// Compute weighted sum of shaped arrays of equal shapes
+    ///
+    /// - Parameters:
+    ///   - weights: The weights each array is multiplied by
+    ///   - values: The arrays to be weighted and summed
+    /// - Returns: sum_i weights[i]*values[i]
+    func weightedSum(_ weights: [Double], _ values: [MLShapedArray<Float32>]) -> MLShapedArray<Float32> {
+        let scalarCount = values.first!.scalarCount
+        assert(weights.count > 1 && values.count == weights.count)
+        assert(values.allSatisfy({ $0.scalarCount == scalarCount }))
+
+        return MLShapedArray(unsafeUninitializedShape: values.first!.shape) { scalars, _ in
+            scalars.initialize(repeating: 0.0)
+            for i in 0 ..< values.count {
+                let w = Float(weights[i])
+                values[i].withUnsafeShapedBufferPointer { buffer, _, _ in
+                    assert(buffer.count == scalarCount)
+                    // scalars[j] = w * values[i].scalars[j]
+                    cblas_saxpy(Int32(scalarCount), w, buffer.baseAddress, 1, scalars.baseAddress, 1)
+                }
+            }
+        }
+    }
+    
+    func addNoise(
+        originalSample: MLShapedArray<Float32>,
+        noise: [MLShapedArray<Float32>],
+        strength: Float
+    ) -> [MLShapedArray<Float32>] {
+        let startStep = max(inferenceStepCount - Int(Float(inferenceStepCount) * strength), 0)
+        let alphaProdt = alphasCumProd[timeSteps[startStep]]
+        let betaProdt = 1 - alphaProdt
+        let sqrtAlphaProdt = sqrt(alphaProdt)
+        let sqrtBetaProdt = sqrt(betaProdt)
+        
+        let noisySamples = noise.map {
+            weightedSum(
+                [Double(sqrtAlphaProdt), Double(sqrtBetaProdt)],
+                [originalSample, $0]
+            )
+        }
+
+        return noisySamples
+    }
+}
+
+// MARK: - Timesteps
+
+@available(iOS 16.2, macOS 13.1, *)
+public extension Scheduler {
+    func calculateTimesteps(strength: Float?) -> [Int] {
+        guard let strength else { return timeSteps }
+        let startStep = max(inferenceStepCount - Int(Float(inferenceStepCount) * strength), 0)
+        let actualTimesteps = Array(timeSteps[startStep...])
+        return actualTimesteps
+    }
+}
+
+// MARK: - BetaSchedule
+
+/// How to map a beta range to a sequence of betas to step over
+@available(iOS 16.2, macOS 13.1, *)
+public enum BetaSchedule {
+    /// Linear stepping between start and end
+    case linear
+    /// Steps using linspace(sqrt(start),sqrt(end))^2
+    case scaledLinear
+}
+
+// MARK: - PNDMScheduler
+
+/// A scheduler used to compute a de-noised image
+///
+///  This implementation matches:
+///  [Hugging Face Diffusers PNDMScheduler](https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_pndm.py)
+///
+/// This scheduler uses the pseudo linear multi-step (PLMS) method only, skipping pseudo Runge-Kutta (PRK) steps
+@available(iOS 16.2, macOS 13.1, *)
+public final class PNDMScheduler: Scheduler {
+    public let trainStepCount: Int
+    public let inferenceStepCount: Int
+    public let betas: [Float]
+    public let alphas: [Float]
+    public let alphasCumProd: [Float]
+    public let timeSteps: [Int]
+
+    public let alpha_t: [Float]
+    public let sigma_t: [Float]
+    public let lambda_t: [Float]
+
+    public private(set) var modelOutputs: [MLShapedArray<Float32>] = []
+
+    // Internal state
+    var counter: Int
+    var ets: [MLShapedArray<Float32>]
+    var currentSample: MLShapedArray<Float32>?
+
+    /// Create a scheduler that uses a pseudo linear multi-step (PLMS)  method
+    ///
+    /// - Parameters:
+    ///   - stepCount: Number of inference steps to schedule
+    ///   - trainStepCount: Number of training diffusion steps
+    ///   - betaSchedule: Method to schedule betas from betaStart to betaEnd
+    ///   - betaStart: The starting value of beta for inference
+    ///   - betaEnd: The end value for beta for inference
+    /// - Returns: A scheduler ready for its first step
+    public init(
+        stepCount: Int = 50,
+        trainStepCount: Int = 1000,
+        betaSchedule: BetaSchedule = .scaledLinear,
+        betaStart: Float = 0.00085,
+        betaEnd: Float = 0.012
+    ) {
+        self.trainStepCount = trainStepCount
+        self.inferenceStepCount = stepCount
+
+        switch betaSchedule {
+        case .linear:
+            self.betas = linspace(betaStart, betaEnd, trainStepCount)
+        case .scaledLinear:
+            self.betas = linspace(pow(betaStart, 0.5), pow(betaEnd, 0.5), trainStepCount).map({ $0 * $0 })
+        }
+        self.alphas = betas.map({ 1.0 - $0 })
+        var alphasCumProd = self.alphas
+        for i in 1..<alphasCumProd.count {
+            alphasCumProd[i] *= alphasCumProd[i -  1]
+        }
+        self.alphasCumProd = alphasCumProd
+        let stepsOffset = 1 // For stable diffusion
+        let stepRatio = Float(trainStepCount / stepCount )
+        let forwardSteps = (0..<stepCount).map {
+            Int((Float($0) * stepRatio).rounded()) + stepsOffset
+        }
+
+        self.alpha_t = vForce.sqrt(self.alphasCumProd)
+        self.sigma_t = vForce.sqrt(vDSP.subtract([Float](repeating: 1, count: self.alphasCumProd.count), self.alphasCumProd))
+        self.lambda_t = zip(self.alpha_t, self.sigma_t).map { α, σ in log(α) - log(σ) }
+
+        var timeSteps: [Int] = []
+        timeSteps.append(contentsOf: forwardSteps.dropLast(1))
+        timeSteps.append(timeSteps.last!)
+        timeSteps.append(forwardSteps.last!)
+        timeSteps.reverse()
+
+        self.timeSteps = timeSteps
+        self.counter = 0
+        self.ets = []
+        self.currentSample = nil
+    }
+
+    /// Compute a de-noised image sample and step scheduler state
+    ///
+    /// - Parameters:
+    ///   - output: The predicted residual noise output of learned diffusion model
+    ///   - timeStep: The current time step in the diffusion chain
+    ///   - sample: The current input sample to the diffusion model
+    /// - Returns: Predicted de-noised sample at the previous time step
+    /// - Postcondition: The scheduler state is updated.
+    ///   The state holds the current sample and history of model output noise residuals
+    public func step(
+        output: MLShapedArray<Float32>,
+        timeStep t: Int,
+        sample s: MLShapedArray<Float32>
+    ) -> MLShapedArray<Float32> {
+        
+        var timeStep = t
+        let stepInc = (trainStepCount / inferenceStepCount)
+        var prevStep = timeStep - stepInc
+        var modelOutput = output
+        var sample = s
+
+        if counter != 1 {
+            if ets.count > 3 {
+                ets = Array(ets[(ets.count - 3)..<ets.count])
+            }
+            ets.append(output)
+        } else {
+            prevStep = timeStep
+            timeStep = timeStep + stepInc
+        }
+
+        if ets.count == 1 && counter == 0 {
+            modelOutput = output
+            currentSample = sample
+        } else if ets.count == 1 && counter == 1 {
+            modelOutput = weightedSum(
+                [1.0/2.0, 1.0/2.0],
+                [output,  ets[back: 1]]
+            )
+            sample = currentSample!
+            currentSample = nil
+        } else if ets.count == 2 {
+            modelOutput = weightedSum(
+                [3.0/2.0,      -1.0/2.0],
+                [ets[back: 1], ets[back: 2]]
+            )
+        } else if ets.count == 3 {
+            modelOutput = weightedSum(
+                [23.0/12.0,    -16.0/12.0,   5.0/12.0],
+                [ets[back: 1], ets[back: 2], ets[back: 3]]
+            )
+        } else {
+            modelOutput = weightedSum(
+                [55.0/24.0,    -59.0/24.0,   37/24.0,      -9/24.0],
+                [ets[back: 1], ets[back: 2], ets[back: 3], ets[back: 4]]
+            )
+        }
+
+        let convertedOutput = convertModelOutput(modelOutput: modelOutput, timestep: timeStep, sample: sample)
+        modelOutputs.append(convertedOutput)
+
+        let prevSample = previousSample(sample, timeStep, prevStep, modelOutput)
+
+        counter += 1
+        return prevSample
+    }
+
+    /// Convert the model output to the corresponding type the algorithm needs.
+    func convertModelOutput(modelOutput: MLShapedArray<Float32>, timestep: Int, sample: MLShapedArray<Float32>) -> MLShapedArray<Float32> {
+        assert(modelOutput.scalarCount == sample.scalarCount)
+        let scalarCount = modelOutput.scalarCount
+        let (alpha_t, sigma_t) = (self.alpha_t[timestep], self.sigma_t[timestep])
+
+        return MLShapedArray(unsafeUninitializedShape: modelOutput.shape) { scalars, _ in
+            assert(scalars.count == scalarCount)
+            modelOutput.withUnsafeShapedBufferPointer { modelOutput, _, _ in
+                sample.withUnsafeShapedBufferPointer { sample, _, _ in
+                    for i in 0 ..< scalarCount {
+                        scalars.initializeElement(at: i, to: (sample[i] - modelOutput[i] * sigma_t) / alpha_t)
+                    }
+                }
+            }
+        }
+    }
+
+    /// Compute  sample (denoised image) at previous step given a current time step
+    ///
+    /// - Parameters:
+    ///   - sample: The current input to the model x_t
+    ///   - timeStep: The current time step t
+    ///   - prevStep: The previous time step t−δ
+    ///   - modelOutput: Predicted noise residual the current time step e_θ(x_t, t)
+    /// - Returns: Computes previous sample x_(t−δ)
+    func previousSample(
+        _ sample: MLShapedArray<Float32>,
+        _ timeStep: Int,
+        _ prevStep: Int,
+        _ modelOutput: MLShapedArray<Float32>
+    ) ->  MLShapedArray<Float32> {
+
+        // Compute x_(t−δ) using formula (9) from
+        // "Pseudo Numerical Methods for Diffusion Models on Manifolds",
+        // Luping Liu, Yi Ren, Zhijie Lin & Zhou Zhao.
+        // ICLR 2022
+        //
+        // Notation:
+        //
+        // alphaProdt       α_t
+        // alphaProdtPrev   α_(t−δ)
+        // betaProdt        (1 - α_t)
+        // betaProdtPrev    (1 - α_(t−δ))
+        let alphaProdt = alphasCumProd[timeStep]
+        let alphaProdtPrev = alphasCumProd[max(0,prevStep)]
+        let betaProdt = 1 - alphaProdt
+        let betaProdtPrev = 1 - alphaProdtPrev
+
+        // sampleCoeff = (α_(t−δ) - α_t) divided by
+        // denominator of x_t in formula (9) and plus 1
+        // Note: (α_(t−δ) - α_t) / (sqrt(α_t) * (sqrt(α_(t−δ)) + sqr(α_t))) =
+        // sqrt(α_(t−δ)) / sqrt(α_t))
+        let sampleCoeff = sqrt(alphaProdtPrev / alphaProdt)
+
+        // Denominator of e_θ(x_t, t) in formula (9)
+        let modelOutputDenomCoeff = alphaProdt * sqrt(betaProdtPrev)
+        + sqrt(alphaProdt * betaProdt * alphaProdtPrev)
+
+        // full formula (9)
+        let modelCoeff = -(alphaProdtPrev - alphaProdt)/modelOutputDenomCoeff
+        let prevSample = weightedSum(
+            [Double(sampleCoeff), Double(modelCoeff)],
+            [sample, modelOutput]
+        )
+
+        return prevSample
+    }
+}
+
+/// Evenly spaced floats between specified interval
+///
+/// - Parameters:
+///   - start: Start of the interval
+///   - end: End of the interval
+///   - count: The number of floats to return between [*start*, *end*]
+/// - Returns: Float array with *count* elements evenly spaced between at *start* and *end*
+func linspace(_ start: Float, _ end: Float, _ count: Int) -> [Float] {
+    let scale = (end - start) / Float(count - 1)
+    return (0..<count).map { Float($0)*scale + start }
+}
+
+extension Collection {
+    /// Collection element index from the back. *self[back: 1]* yields the last element
+    public subscript(back i: Int) -> Element {
+        return self[index(endIndex, offsetBy: -i)]
+    }
+}

swift/StableDiffusion/pipeline/StableDiffusion3Pipeline+Resources.swift

@@ -0,0 +1,97 @@
+// For licensing see accompanying LICENSE.md file.
+// Copyright (C) 2024 Apple Inc. All Rights Reserved.
+
+import CoreML
+import Foundation
+import Tokenizers
+import Hub
+
+@available(iOS 17.0, macOS 14.0, *)
+public extension StableDiffusion3Pipeline {
+    struct ResourceURLs {
+        public let textEncoderURL: URL
+        public let textEncoder2URL: URL
+        public let textEncoderT5URL: URL
+        public let mmditURL: URL
+        public let decoderURL: URL
+        public let encoderURL: URL
+        public let vocabURL: URL
+        public let mergesURL: URL
+        public let configT5URL: URL
+        public let dataT5URL: URL
+
+        public init(resourcesAt baseURL: URL) {
+            textEncoderURL = baseURL.appending(path: "TextEncoder.mlmodelc")
+            textEncoder2URL = baseURL.appending(path: "TextEncoder2.mlmodelc")
+            textEncoderT5URL = baseURL.appending(path: "TextEncoderT5.mlmodelc")
+            mmditURL = baseURL.appending(path: "MultiModalDiffusionTransformer.mlmodelc")
+            decoderURL = baseURL.appending(path: "VAEDecoder.mlmodelc")
+            encoderURL = baseURL.appending(path: "VAEEncoder.mlmodelc")
+            vocabURL = baseURL.appending(path: "vocab.json")
+            mergesURL = baseURL.appending(path: "merges.txt")
+            configT5URL = baseURL.appending(path: "tokenizer_config.json")
+            dataT5URL = baseURL.appending(path: "tokenizer.json")
+        }
+    }
+
+    /// Create stable diffusion pipeline using model resources at a
+    /// specified URL
+    ///
+    /// - Parameters:
+    ///   - baseURL: URL pointing to directory holding all model and tokenization resources
+    ///   - configuration: The configuration to load model resources with
+    ///   - reduceMemory: Setup pipeline in reduced memory mode
+    /// - Returns:
+    ///  Pipeline ready for image generation if all  necessary resources loaded
+    init(
+        resourcesAt baseURL: URL,
+        configuration config: MLModelConfiguration = .init(),
+        reduceMemory: Bool = false
+    ) throws {
+        // Expect URL of each resource
+        let urls = ResourceURLs(resourcesAt: baseURL)
+        let tokenizer = try BPETokenizer(mergesAt: urls.mergesURL, vocabularyAt: urls.vocabURL)
+        let textEncoder = TextEncoderXL(tokenizer: tokenizer, modelAt: urls.textEncoderURL, configuration: config)
+
+        // padToken is different in the second XL text encoder
+        let tokenizer2 = try BPETokenizer(mergesAt: urls.mergesURL, vocabularyAt: urls.vocabURL, padToken: "!")
+        let textEncoder2 = TextEncoderXL(tokenizer: tokenizer2, modelAt: urls.textEncoder2URL, configuration: config)
+
+        // Optional T5 encoder
+        var textEncoderT5: TextEncoderT5?
+        if FileManager.default.fileExists(atPath: urls.configT5URL.path),
+           FileManager.default.fileExists(atPath: urls.dataT5URL.path),
+           FileManager.default.fileExists(atPath: urls.textEncoderT5URL.path)
+        {
+            let tokenizerT5 = try PreTrainedTokenizer(tokenizerConfig: Config(fileURL: urls.configT5URL), tokenizerData: Config(fileURL: urls.dataT5URL))
+            textEncoderT5 = TextEncoderT5(tokenizer: tokenizerT5, modelAt: urls.textEncoderT5URL, configuration: config)
+        } else {
+            textEncoderT5 = nil
+        }
+
+        // Denoiser model
+        let mmdit = MultiModalDiffusionTransformer(modelAt: urls.mmditURL, configuration: config)
+
+        // Image Decoder
+        let decoder = Decoder(modelAt: urls.decoderURL, configuration: config)
+
+        // Optional Image Encoder
+        let encoder: Encoder?
+        if FileManager.default.fileExists(atPath: urls.encoderURL.path) {
+            encoder = Encoder(modelAt: urls.encoderURL, configuration: config)
+        } else {
+            encoder = nil
+        }
+
+        // Construct pipeline
+        self.init(
+            textEncoder: textEncoder,
+            textEncoder2: textEncoder2,
+            textEncoderT5: textEncoderT5,
+            mmdit: mmdit,
+            decoder: decoder,
+            encoder: encoder,
+            reduceMemory: reduceMemory
+        )
+    }
+}

swift/StableDiffusion/pipeline/StableDiffusion3Pipeline.swift

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

swift/StableDiffusion/pipeline/StableDiffusionPipeline+Resources.swift

@@ -0,0 +1,166 @@
+// For licensing see accompanying LICENSE.md file.
+// Copyright (C) 2022 Apple Inc. All Rights Reserved.
+
+import Foundation
+import CoreML
+import NaturalLanguage
+
+@available(iOS 16.2, macOS 13.1, *)
+public extension StableDiffusionPipeline {
+
+    struct ResourceURLs {
+
+        public let textEncoderURL: URL
+        public let unetURL: URL
+        public let unetChunk1URL: URL
+        public let unetChunk2URL: URL
+        public let decoderURL: URL
+        public let encoderURL: URL
+        public let safetyCheckerURL: URL
+        public let vocabURL: URL
+        public let mergesURL: URL
+        public let controlNetDirURL: URL
+        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")
+            unetURL = baseURL.appending(path: "Unet.mlmodelc")
+            unetChunk1URL = baseURL.appending(path: "UnetChunk1.mlmodelc")
+            unetChunk2URL = baseURL.appending(path: "UnetChunk2.mlmodelc")
+            decoderURL = baseURL.appending(path: "VAEDecoder.mlmodelc")
+            encoderURL = baseURL.appending(path: "VAEEncoder.mlmodelc")
+            safetyCheckerURL = baseURL.appending(path: "SafetyChecker.mlmodelc")
+            vocabURL = baseURL.appending(path: "vocab.json")
+            mergesURL = baseURL.appending(path: "merges.txt")
+            controlNetDirURL = baseURL.appending(path: "controlnet")
+            controlledUnetURL = baseURL.appending(path: "ControlledUnet.mlmodelc")
+            controlledUnetChunk1URL = baseURL.appending(path: "ControlledUnetChunk1.mlmodelc")
+            controlledUnetChunk2URL = baseURL.appending(path: "ControlledUnetChunk2.mlmodelc")
+            multilingualTextEncoderProjectionURL = baseURL.appending(path: "MultilingualTextEncoderProjection.mlmodelc")
+        }
+    }
+
+    /// Create stable diffusion pipeline using model resources at a
+    /// specified URL
+    ///
+    /// - Parameters:
+    ///   - 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,
+        useMultilingualTextEncoder: Bool = false,
+        script: Script? = nil
+    ) throws {
+
+        /// Expect URL of each resource
+        let urls = ResourceURLs(resourcesAt: baseURL)
+        let textEncoder: TextEncoderModel
+
+#if canImport(NaturalLanguage.NLScript)
+        if useMultilingualTextEncoder {
+            guard #available(macOS 14.0, iOS 17.0, *) else { throw PipelineError.unsupportedOSVersion }
+            textEncoder = MultilingualTextEncoder(
+                modelAt: urls.multilingualTextEncoderProjectionURL,
+                configuration: config,
+                script: script ?? .latin
+            )
+        } else {
+            let tokenizer = try BPETokenizer(mergesAt: urls.mergesURL, vocabularyAt: urls.vocabURL)
+            textEncoder = TextEncoder(tokenizer: tokenizer, modelAt: urls.textEncoderURL, configuration: config)
+        }
+#else
+        let tokenizer = try BPETokenizer(mergesAt: urls.mergesURL, vocabularyAt: urls.vocabURL)
+        textEncoder = TextEncoder(tokenizer: tokenizer, modelAt: urls.textEncoderURL, configuration: config)
+#endif
+
+        // ControlNet model
+        var controlNet: ControlNet? = nil
+        let controlNetURLs = controlNetModelNames.map { model in
+            let fileName = model + ".mlmodelc"
+            return urls.controlNetDirURL.appending(path: fileName)
+        }
+        if !controlNetURLs.isEmpty {
+            controlNet = ControlNet(modelAt: controlNetURLs, configuration: config)
+        }
+
+        // Unet model
+        let unet: Unet
+        let unetURL: URL, unetChunk1URL: URL, unetChunk2URL: URL
+        
+        // if ControlNet available, Unet supports additional inputs from ControlNet
+        if controlNet == nil {
+            unetURL = urls.unetURL
+            unetChunk1URL = urls.unetChunk1URL
+            unetChunk2URL = urls.unetChunk2URL
+        } else {
+            unetURL = urls.controlledUnetURL
+            unetChunk1URL = urls.controlledUnetChunk1URL
+            unetChunk2URL = urls.controlledUnetChunk2URL
+        }
+        
+        if FileManager.default.fileExists(atPath: unetChunk1URL.path) &&
+            FileManager.default.fileExists(atPath: unetChunk2URL.path) {
+            unet = Unet(chunksAt: [unetChunk1URL, unetChunk2URL],
+                        configuration: config)
+        } else {
+            unet = Unet(modelAt: unetURL, configuration: config)
+        }
+
+        // Image Decoder
+        let decoder = Decoder(modelAt: urls.decoderURL, configuration: config)
+
+        // Optional safety checker
+        var safetyChecker: SafetyChecker? = nil
+        if !disableSafety &&
+            FileManager.default.fileExists(atPath: urls.safetyCheckerURL.path) {
+            safetyChecker = SafetyChecker(modelAt: urls.safetyCheckerURL, configuration: config)
+        }
+        
+        // Optional Image Encoder
+        let encoder: Encoder?
+        if FileManager.default.fileExists(atPath: urls.encoderURL.path) {
+            encoder = Encoder(modelAt: urls.encoderURL, configuration: config)
+        } else {
+            encoder = nil
+        }
+
+        // Construct pipeline
+        if #available(macOS 14.0, iOS 17.0, *) {
+            self.init(
+                textEncoder: textEncoder,
+                unet: unet,
+                decoder: decoder,
+                encoder: encoder,
+                controlNet: controlNet,
+                safetyChecker: safetyChecker,
+                reduceMemory: reduceMemory,
+                useMultilingualTextEncoder: useMultilingualTextEncoder,
+                script: script
+            )
+        } else {
+            self.init(
+                textEncoder: textEncoder,
+                unet: unet,
+                decoder: decoder,
+                encoder: encoder,
+                controlNet: controlNet,
+                safetyChecker: safetyChecker,
+                reduceMemory: reduceMemory
+            )
+        }
+    }
+}

swift/StableDiffusion/pipeline/StableDiffusionPipeline.Configuration.swift

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

swift/StableDiffusion/pipeline/StableDiffusionPipeline.swift

@@ -0,0 +1,484 @@
+// For licensing see accompanying LICENSE.md file.
+// Copyright (C) 2022 Apple Inc. All Rights Reserved.
+
+import Accelerate
+import CoreGraphics
+import CoreML
+import Foundation
+import NaturalLanguage
+
+/// Schedulers compatible with StableDiffusionPipeline
+public enum StableDiffusionScheduler {
+    /// Scheduler that uses a pseudo-linear multi-step (PLMS) method
+    case pndmScheduler
+    /// Scheduler that uses a second order DPM-Solver++ algorithm
+    case dpmSolverMultistepScheduler
+    /// Scheduler for rectified flow based multimodal diffusion transformer models
+    case discreteFlowScheduler
+}
+
+/// RNG compatible with StableDiffusionPipeline
+public enum StableDiffusionRNG {
+    /// RNG that matches numpy implementation
+    case numpyRNG
+    /// RNG that matches PyTorch CPU implementation.
+    case torchRNG
+    /// RNG that matches PyTorch CUDA implementation.
+    case nvidiaRNG
+}
+
+public enum PipelineError: String, Swift.Error {
+    case missingUnetInputs
+    case startingImageProvidedWithoutEncoder
+    case startingText2ImgWithoutTextEncoder
+    case unsupportedOSVersion
+    case errorCreatingPreview
+}
+
+@available(iOS 16.2, macOS 13.1, *)
+public protocol StableDiffusionPipelineProtocol: ResourceManaging {
+    var canSafetyCheck: Bool { get }
+
+    func generateImages(
+        configuration config: PipelineConfiguration,
+        progressHandler: (PipelineProgress) -> Bool
+    ) throws -> [CGImage?]
+
+    func decodeToImages(
+        _ latents: [MLShapedArray<Float32>],
+        configuration config: PipelineConfiguration
+    ) throws -> [CGImage?]
+}
+
+@available(iOS 16.2, macOS 13.1, *)
+public extension StableDiffusionPipelineProtocol {
+    var canSafetyCheck: Bool { false }
+}
+
+/// A pipeline used to generate image samples from text input using stable diffusion
+///
+/// This implementation matches:
+/// [Hugging Face Diffusers Pipeline](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py)
+@available(iOS 16.2, macOS 13.1, *)
+public struct StableDiffusionPipeline: StableDiffusionPipelineProtocol {
+
+    /// Model to generate embeddings for tokenized input text
+    var textEncoder: TextEncoderModel
+
+    /// Model used to predict noise residuals given an input, diffusion time step, and conditional embedding
+    var unet: Unet
+
+    /// Model used to generate final image from latent diffusion process
+    var decoder: Decoder
+    
+    /// Model used to latent space for image2image, and soon, in-painting
+    var encoder: Encoder?
+
+    /// Optional model for checking safety of generated image
+    var safetyChecker: SafetyChecker? = nil
+    
+    /// Optional model used before Unet to control generated images by additonal inputs
+    var controlNet: ControlNet? = nil
+
+    /// Reports whether this pipeline can perform safety checks
+    public var canSafetyCheck: Bool {
+        safetyChecker != nil
+    }
+
+    /// Option to reduce memory during image generation
+    ///
+    /// If true, the pipeline will lazily load TextEncoder, Unet, Decoder, and SafetyChecker
+    /// when needed and aggressively unload their resources after
+    ///
+    /// This will increase latency in favor of reducing memory
+    var reduceMemory: Bool = false
+
+    /// 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:
+    ///   - 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
+    /// - Returns: Pipeline ready for image generation
+    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
+        self.encoder = encoder
+        self.controlNet = controlNet
+        self.safetyChecker = safetyChecker
+        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
+    /// and let the pipeline lazily load resources as needed
+    public func loadResources() throws {
+        if reduceMemory {
+            try prewarmResources()
+        } else {
+            try unet.loadResources()
+            try textEncoder.loadResources()
+            try decoder.loadResources()
+            try encoder?.loadResources()
+            try controlNet?.loadResources()
+            try safetyChecker?.loadResources()
+        }
+    }
+
+    /// Unload the underlying resources to free up memory
+    public func unloadResources() {
+        textEncoder.unloadResources()
+        unet.unloadResources()
+        decoder.unloadResources()
+        encoder?.unloadResources()
+        controlNet?.unloadResources()
+        safetyChecker?.unloadResources()
+    }
+
+    // Prewarm resources one at a time
+    public func prewarmResources() throws {
+        try textEncoder.prewarmResources()
+        try unet.prewarmResources()
+        try decoder.prewarmResources()
+        try encoder?.prewarmResources()
+        try controlNet?.prewarmResources()
+        try safetyChecker?.prewarmResources()
+    }
+
+    /// Image generation using stable diffusion
+    /// - Parameters:
+    ///   - configuration: Image generation configuration
+    ///   - progressHandler: Callback to perform after each step, stops on receiving false response
+    /// - Returns: An array of `imageCount` optional images.
+    ///            The images will be nil if safety checks were performed and found the result to be un-safe
+    public func generateImages(
+        configuration config: Configuration,
+        progressHandler: (Progress) -> Bool = { _ in true }
+    ) throws -> [CGImage?] {
+
+        // Encode the input prompt
+        var promptEmbedding = try textEncoder.encode(config.prompt)
+
+        if config.guidanceScale >= 1.0 {
+            // Convert to Unet hidden state representation
+            // Concatenate the prompt and negative prompt embeddings
+            let negativePromptEmbedding = try textEncoder.encode(config.negativePrompt)
+            promptEmbedding = MLShapedArray<Float32>(
+                concatenating: [negativePromptEmbedding, promptEmbedding],
+                alongAxis: 0
+            )
+        }
+
+        if reduceMemory {
+            textEncoder.unloadResources()
+        }
+
+        let hiddenStates = useMultilingualTextEncoder ? promptEmbedding : toHiddenStates(promptEmbedding)
+
+        /// Setup schedulers
+        let scheduler: [Scheduler] = (0..<config.imageCount).map { _ in
+            switch config.schedulerType {
+            case .pndmScheduler: return PNDMScheduler(stepCount: config.stepCount)
+            case .dpmSolverMultistepScheduler: return DPMSolverMultistepScheduler(stepCount: config.stepCount, timeStepSpacing: config.schedulerTimestepSpacing)
+            case .discreteFlowScheduler: return DiscreteFlowScheduler(stepCount: config.stepCount, timeStepShift: config.schedulerTimestepShift)
+            }
+        }
+
+        // Generate random latent samples from specified seed
+        var latents: [MLShapedArray<Float32>] = try generateLatentSamples(configuration: config, scheduler: scheduler[0])
+
+        // Store denoised latents from scheduler to pass into decoder
+        var denoisedLatents: [MLShapedArray<Float32>] = latents.map { MLShapedArray(converting: $0) }
+
+        if reduceMemory {
+            encoder?.unloadResources()
+        }
+        let timestepStrength: Float? = config.mode == .imageToImage ? config.strength : nil
+        
+        // Convert cgImage for ControlNet into MLShapedArray
+        let controlNetConds = try config.controlNetInputs.map { cgImage in
+            let shapedArray = try cgImage.planarRGBShapedArray(minValue: 0.0, maxValue: 1.0)
+            return MLShapedArray(
+                concatenating: [shapedArray, shapedArray],
+                alongAxis: 0
+            )
+        }
+
+        // De-noising loop
+        let timeSteps: [Int] = scheduler[0].calculateTimesteps(strength: timestepStrength)
+        for (step,t) in timeSteps.enumerated() {
+
+            // Expand the latents for classifier-free guidance
+            // and input to the Unet noise prediction model
+            let latentUnetInput: [MLShapedArray<Float32>]
+            if config.guidanceScale >= 1.0 {
+                latentUnetInput = latents.map {
+                    MLShapedArray<Float32>(concatenating: [$0, $0], alongAxis: 0)
+                }
+            } else {
+                latentUnetInput = latents
+            }
+
+            // Before Unet, execute controlNet and add the output into Unet inputs
+            let additionalResiduals = try controlNet?.execute(
+                latents: latentUnetInput,
+                timeStep: t,
+                hiddenStates: hiddenStates,
+                images: controlNetConds
+            )
+            
+            // Predict noise residuals from latent samples
+            // and current time step conditioned on hidden states
+            var noise : [MLShapedArray<Float32>]
+            if unet.latentSampleShape[0] >= 2 || config.guidanceScale < 1.0 {
+                // One predict call from the uNet, using batching if needed
+                noise = try unet.predictNoise(
+                  latents: latentUnetInput,
+                  timeStep: t,
+                  hiddenStates: hiddenStates,
+                  additionalResiduals: additionalResiduals
+                )
+            } else {
+                // Serial predictions from uNet
+                var hidden0 = MLShapedArray<Float32>(converting: hiddenStates[0])
+                hidden0 = MLShapedArray(scalars: hidden0.scalars, shape: [1]+hidden0.shape)
+                let noise_pred_uncond = try unet.predictNoise(
+                  latents: latents,
+                  timeStep: t,
+                  hiddenStates: hidden0,
+                  additionalResiduals: additionalResiduals
+                )
+
+                var hidden1 = MLShapedArray<Float32>(converting: hiddenStates[1])
+                hidden1 = MLShapedArray(scalars: hidden1.scalars, shape: [1]+hidden1.shape)
+                let noise_pred_text = try unet.predictNoise(
+                  latents: latents,
+                  timeStep: t,
+                  hiddenStates: hidden1,
+                  additionalResiduals: additionalResiduals
+                )
+
+                noise = [MLShapedArray<Float32>(concatenating: [noise_pred_uncond[0], noise_pred_text[0]],
+                                                alongAxis: 0)]
+            }
+
+            if config.guidanceScale >= 1.0 {
+                noise = performGuidance(noise, config.guidanceScale)
+            }
+
+            // Have the scheduler compute the previous (t-1) latent
+            // sample given the predicted noise and current sample
+            for i in 0..<config.imageCount {
+                latents[i] = scheduler[i].step(
+                    output: noise[i],
+                    timeStep: t,
+                    sample: latents[i]
+                )
+
+                denoisedLatents[i] = scheduler[i].modelOutputs.last ?? latents[i]
+            }
+
+            let currentLatentSamples = config.useDenoisedIntermediates ? denoisedLatents : latents
+
+            // Report progress
+            let progress = Progress(
+                pipeline: self,
+                prompt: config.prompt,
+                step: step,
+                stepCount: timeSteps.count,
+                currentLatentSamples: currentLatentSamples,
+                configuration: config
+            )
+            if !progressHandler(progress) {
+                // Stop if requested by handler
+                return []
+            }
+        }
+
+        if reduceMemory {
+            controlNet?.unloadResources()
+            unet.unloadResources()
+        }
+
+        // Decode the latent samples to images
+        return try decodeToImages(denoisedLatents, configuration: config)
+    }
+
+    func generateLatentSamples(configuration config: Configuration, scheduler: Scheduler) throws -> [MLShapedArray<Float32>] {
+        var sampleShape = unet.latentSampleShape
+        sampleShape[0] = 1
+        
+        let stdev = scheduler.initNoiseSigma
+        var random = randomSource(from: config.rngType, seed: config.seed)
+        let samples = (0..<config.imageCount).map { _ in
+            MLShapedArray<Float32>(
+                converting: random.normalShapedArray(sampleShape, mean: 0.0, stdev: Double(stdev)))
+        }
+        if let image = config.startingImage, config.mode == .imageToImage {
+            guard let encoder else {
+                throw PipelineError.startingImageProvidedWithoutEncoder
+            }
+            let latent = try encoder.encode(image, scaleFactor: config.encoderScaleFactor, random: &random)
+            return scheduler.addNoise(originalSample: latent, noise: samples, strength: config.strength)
+        }
+        return samples
+    }
+
+    public func decodeToImages(_ latents: [MLShapedArray<Float32>], configuration config: Configuration) throws -> [CGImage?] {
+        let images = try decoder.decode(latents, scaleFactor: config.decoderScaleFactor)
+        if reduceMemory {
+            decoder.unloadResources()
+        }
+
+        // If safety is disabled return what was decoded
+        if config.disableSafety {
+            return images
+        }
+
+        // If there is no safety checker return what was decoded
+        guard let safetyChecker = safetyChecker else {
+            return images
+        }
+
+        // Otherwise change images which are not safe to nil
+        let safeImages = try images.map { image in
+            try safetyChecker.isSafe(image) ? image : nil
+        }
+
+        if reduceMemory {
+            safetyChecker.unloadResources()
+        }
+
+        return safeImages
+    }
+
+}
+
+/// Sampling progress details
+@available(iOS 16.2, macOS 13.1, *)
+public struct PipelineProgress {
+    public let pipeline: StableDiffusionPipelineProtocol
+    public let prompt: String
+    public let step: Int
+    public let stepCount: Int
+    public let currentLatentSamples: [MLShapedArray<Float32>]
+    public let configuration: PipelineConfiguration
+    public var isSafetyEnabled: Bool {
+        pipeline.canSafetyCheck && !configuration.disableSafety
+    }
+    public var currentImages: [CGImage?] {
+        try! pipeline.decodeToImages(currentLatentSamples, configuration: configuration)
+    }
+}
+
+@available(iOS 16.2, macOS 13.1, *)
+public extension StableDiffusionPipeline {
+    /// Sampling progress details
+    typealias Progress = PipelineProgress
+}
+
+// Helper functions
+
+@available(iOS 16.2, macOS 13.1, *)
+extension StableDiffusionPipelineProtocol {
+    internal func randomSource(from rng: StableDiffusionRNG, seed: UInt32) -> RandomSource {
+        switch rng {
+        case .numpyRNG:
+            return NumPyRandomSource(seed: seed)
+        case .torchRNG:
+            return TorchRandomSource(seed: seed)
+        case .nvidiaRNG:
+            return NvRandomSource(seed: seed)
+        }
+    }
+
+    func toHiddenStates(_ embedding: MLShapedArray<Float32>) -> MLShapedArray<Float32> {
+        // Unoptimized manual transpose [0, 2, None, 1]
+        // e.g. From [2, 77, 768] to [2, 768, 1, 77]
+        let fromShape = embedding.shape
+        let stateShape = [fromShape[0],fromShape[2], 1, fromShape[1]]
+        var states = MLShapedArray<Float32>(repeating: 0.0, shape: stateShape)
+        for i0 in 0..<fromShape[0] {
+            for i1 in 0..<fromShape[1] {
+                for i2 in 0..<fromShape[2] {
+                    states[scalarAt:i0,i2,0,i1] = embedding[scalarAt:i0, i1, i2]
+                }
+            }
+        }
+        return states
+    }
+
+    func performGuidance(_ noise: [MLShapedArray<Float32>], _ guidanceScale: Float) -> [MLShapedArray<Float32>] {
+        noise.map { performGuidance($0, guidanceScale) }
+    }
+
+    func performGuidance(_ noise: MLShapedArray<Float32>, _ guidanceScale: Float) -> MLShapedArray<Float32> {
+        var shape = noise.shape
+        shape[0] = 1
+        return MLShapedArray<Float>(unsafeUninitializedShape: shape) { result, _ in
+            noise.withUnsafeShapedBufferPointer { scalars, _, strides in
+                for i in 0 ..< result.count {
+                    // unconditioned + guidance*(text - unconditioned)
+                    result.initializeElement(
+                        at: i,
+                        to: scalars[i] + guidanceScale * (scalars[strides[0] + i] - scalars[i])
+                    )
+                }
+            }
+        }
+    }
+}

swift/StableDiffusion/pipeline/StableDiffusionXL+Resources.swift

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

swift/StableDiffusion/pipeline/StableDiffusionXLPipeline.swift

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

swift/StableDiffusion/pipeline/TextEncoder.swift

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

swift/StableDiffusion/pipeline/TextEncoderT5.swift

@@ -0,0 +1,124 @@
+// For licensing see accompanying LICENSE.md file.
+// Copyright (C) 2023 Apple Inc. All Rights Reserved.
+
+import Foundation
+import CoreML
+import Tokenizers
+
+@available(iOS 17.0, macOS 14.0, *)
+public protocol TextEncoderT5Model: ResourceManaging {
+    func encode(_ text: String) throws -> TextEncoderT5Output
+}
+
+@available(iOS 17.0, macOS 14.0, *)
+public struct TextEncoderT5Output {
+    public let encoderHiddenStates: MLShapedArray<Float32>
+}
+
+///  A model for encoding text, suitable for SD3
+@available(iOS 17.0, macOS 14.0, *)
+public struct TextEncoderT5: TextEncoderT5Model {
+
+    /// Text tokenizer
+    var tokenizer: Tokenizer
+
+    /// Embedding model
+    var model: ManagedMLModel
+
+    /// Creates text encoder which embeds a tokenized string
+    ///
+    /// - Parameters:
+    ///   - tokenizer: Tokenizer for input text
+    ///   - url: Location of compiled text encoding  Core ML model
+    ///   - configuration: configuration to be used when the model is loaded
+    /// - Returns: A text encoder that will lazily load its required resources when needed or requested
+    public init(tokenizer: Tokenizer,
+                modelAt url: URL,
+                configuration: MLModelConfiguration) {
+        self.tokenizer = tokenizer
+        self.model = ManagedMLModel(modelAt: url, configuration: configuration)
+    }
+
+    /// Ensure the model has been loaded into memory
+    public func loadResources() throws {
+        try model.loadResources()
+    }
+
+    /// Unload the underlying model to free up memory
+    public func unloadResources() {
+       model.unloadResources()
+    }
+
+    /// Encode input text/string
+    ///
+    ///  - Parameters:
+    ///     - text: Input text to be tokenized and then embedded
+    ///  - Returns: Embedding representing the input text
+    public func encode(_ text: String) throws -> TextEncoderT5Output {
+
+        // Get models expected input length
+        let inputLength = inputShape.last!
+
+        // Tokenize, padding to the expected length
+        var tokens = tokenizer.tokenize(text: text)
+        var ids = tokens.map { tokenizer.convertTokenToId($0) ?? 0 }
+
+        // Truncate if necessary
+        if ids.count > inputLength {
+            tokens = tokens.dropLast(tokens.count - inputLength)
+            ids = ids.dropLast(ids.count - inputLength)
+            print("Needed to truncate input for TextEncoderT5")
+        }
+
+        // Use the model to generate the embedding
+        let encodedText = try encode(ids: ids)
+        return encodedText
+    }
+
+    func encode(ids: [Int]) throws -> TextEncoderT5Output {
+        let inputName = "input_ids"
+        let inputShape = inputShape
+        let inputLength = inputShape[1]
+                
+        let bosToken = tokenizer.bosTokenId ?? 0
+        let eosToken = tokenizer.eosTokenId ?? 1
+        let padToken = bosToken
+        let maskToken = eosToken
+
+        // Truncate and pad input to the expected length
+        let truncatedIds = ids.prefix(inputLength - 1) + [eosToken]
+        let inputIds = truncatedIds + Array(repeating: padToken, count: inputLength - truncatedIds.count)
+
+        let attentionMaskName = "attention_mask"
+        var attentionMask: [Int] = inputIds.map { token in
+            token == padToken ? maskToken : padToken
+        }
+        attentionMask[0] = bosToken
+
+        let floatIds = inputIds.map { Float32($0) }
+        let floatMask = attentionMask.map { Float32($0) }
+
+        let inputArray = MLShapedArray<Float32>(scalars: floatIds, shape: inputShape)
+        let maskArray = MLShapedArray<Float32>(scalars: floatMask, shape: inputShape)
+        let inputFeatures = try! MLDictionaryFeatureProvider(
+            dictionary: [inputName: MLMultiArray(inputArray),
+                         attentionMaskName: MLMultiArray(maskArray)])
+
+        let result = try model.perform { model in
+            try model.prediction(from: inputFeatures)
+        }
+
+        let embeddingFeature = result.featureValue(for: "encoder_hidden_states")
+        return TextEncoderT5Output(encoderHiddenStates: MLShapedArray<Float32>(converting: embeddingFeature!.multiArrayValue!))
+    }
+
+    var inputDescription: MLFeatureDescription {
+        try! model.perform { model in
+            model.modelDescription.inputDescriptionsByName.first!.value
+        }
+    }
+    
+    var inputShape: [Int] {
+        inputDescription.multiArrayConstraint!.shape.map { $0.intValue }
+    }
+}

swift/StableDiffusion/pipeline/TextEncoderXL.swift

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

swift/StableDiffusion/pipeline/TorchRandomSource.swift

@@ -0,0 +1,157 @@
+// For licensing see accompanying LICENSE.md file.
+// Copyright (C) 2022 Apple Inc. All Rights Reserved.
+
+import Foundation
+import CoreML
+
+/// A random source consistent with PyTorch
+///
+///  This implementation matches:
+///  https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/core/DistributionsHelper.h
+///  https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cpu/DistributionTemplates.h
+///  https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cpu/DistributionKernels.cpp
+///  https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/core/TransformationHelper.h
+///
+@available(iOS 16.2, macOS 13.1, *)
+struct TorchRandomSource: RandomNumberGenerator, RandomSource {
+
+  struct State {
+    var key = [UInt32](repeating: 0, count: 624)
+    var pos: Int = 0
+    var nextGauss: Double? = nil
+  }
+
+  var state: State
+
+  /// Initialize with a random seed
+  ///
+  /// - Parameters
+  ///     - seed: Seed for underlying Mersenne Twister 19937 generator
+  /// - Returns random source
+  init(seed: UInt32) {
+    state = .init()
+    var s = seed & 0xffff_ffff
+    for i in 0..<state.key.count {
+      state.key[i] = s
+      s = UInt32((UInt64(1_812_433_253) * UInt64(s ^ (s >> 30)) + UInt64(i) + 1) & 0xffff_ffff)
+    }
+    state.pos = state.key.count
+    state.nextGauss = nil
+  }
+
+  /// Generate next UInt32 using fast 32bit Mersenne Twister
+  mutating func nextUInt32() -> UInt32 {
+    let n = 624
+    let m = 397
+    let matrixA: UInt64 = 0x9908_b0df
+    let upperMask: UInt32 = 0x8000_0000
+    let lowerMask: UInt32 = 0x7fff_ffff
+
+    var y: UInt32
+    if state.pos == state.key.count {
+      for i in 0..<(n - m) {
+        y = (state.key[i] & upperMask) | (state.key[i + 1] & lowerMask)
+        state.key[i] = state.key[i + m] ^ (y >> 1) ^ UInt32((UInt64(~(y & 1)) + 1) & matrixA)
+      }
+      for i in (n - m)..<(n - 1) {
+        y = (state.key[i] & upperMask) | (state.key[i + 1] & lowerMask)
+        state.key[i] = state.key[i + (m - n)] ^ (y >> 1) ^ UInt32((UInt64(~(y & 1)) + 1) & matrixA)
+      }
+      y = (state.key[n - 1] & upperMask) | (state.key[0] & lowerMask)
+      state.key[n - 1] = state.key[m - 1] ^ (y >> 1) ^ UInt32((UInt64(~(y & 1)) + 1) & matrixA)
+      state.pos = 0
+    }
+    y = state.key[state.pos]
+    state.pos += 1
+
+    y ^= (y >> 11)
+    y ^= (y << 7) & 0x9d2c_5680
+    y ^= (y << 15) & 0xefc6_0000
+    y ^= (y >> 18)
+
+    return y
+  }
+
+  mutating func next() -> UInt64 {
+    let high = nextUInt32()
+    let low = nextUInt32()
+    return (UInt64(high) << 32) | UInt64(low)
+  }
+
+  /// Generate next random double value
+  mutating func nextDouble() -> Double {
+    let a = next()
+    return Double(a & 9_007_199_254_740_991) * (1.0 / 9007199254740992.0)
+  }
+
+  /// Generate next random float value
+  mutating func nextFloat() -> Float {
+    let a = nextUInt32()
+    return Float(a & 16_777_215) * (1.0 / 16777216.0)
+  }
+
+  /// Generate next random value from a standard normal
+  mutating func nextGauss() -> Double {
+    if let nextGauss = state.nextGauss {
+      state.nextGauss = nil
+      return nextGauss
+    }
+    // Box-Muller transform
+    let u1: Double = nextDouble()
+    let u2: Double = 1 - nextDouble()
+    let radius = sqrt(-2.0 * log(u2))
+    let theta = 2.0 * .pi * u1
+    state.nextGauss = radius * sin(theta)
+    return radius * cos(theta)
+  }
+    
+  /// Generates a random value from a normal distribution with given mean and standard deviation.
+  mutating func nextNormal(mean: Double = 0.0, stdev: Double = 1.0) -> Double {
+    nextGauss() * stdev + mean
+  }
+
+  /// Generates an array of random values from a normal distribution with given mean and standard deviation.
+  /// This simulates torch.randn([1, 4, 64, 64], dtype=torch.float), note that for dtype=torch.double, it
+  /// will be slightly different.
+  mutating func normalArray(count: Int, mean: Double = 0.0, stdev: Double = 1.0) -> [Double] {
+    // If it is smaller than 16 elements, Torch generates from Box-Muller transform directly.
+    // Note that even if this is used to generate Float, it will use Double underneath.
+    guard count >= 16 else {
+      return (0..<count).map { _ in nextNormal(mean: mean, stdev: stdev) }
+    }
+    // Otherwise, Torch first fill a uniform distribution array, then do Box-Muller
+    // transformation over this array.
+    var data = (0..<count).map { _ in Double(nextFloat()) }
+    for i in stride(from: 0, to: count - 15, by: 16) {
+      for j in 0..<8 {
+        let u1 = 1 - data[i + j]
+        let u2 = data[i + j + 8]
+        let radius = sqrt(-2.0 * log(u1))
+        let theta = 2.0 * .pi * u2
+        data[i + j] = radius * cos(theta) * stdev + mean
+        data[i + j + 8] = radius * sin(theta) * stdev + mean
+      }
+    }
+    if count % 16 != 0 {
+      for i in (count - 16)..<count {
+        data[i] = nextDouble()
+      }
+      let i = count - 16
+      for j in 0..<8 {
+        let u1 = 1 - data[i + j]
+        let u2 = data[i + j + 8]
+        let radius = sqrt(-2.0 * log(u1))
+        let theta = 2.0 * .pi * u2
+        data[i + j] = radius * cos(theta) * stdev + mean
+        data[i + j + 8] = radius * sin(theta) * stdev + mean
+      }
+    }
+    return data
+  }
+
+  /// Generate a shaped array with scalars from a normal distribution with given mean and standard deviation.
+  mutating func normalShapedArray(_ shape: [Int], mean: Double = 0.0, stdev: Double = 1.0) -> MLShapedArray<Double> {
+    let count = shape.reduce(1, *)
+    return .init(scalars: normalArray(count: count, mean: mean, stdev: stdev), shape: shape)
+  }
+}

swift/StableDiffusion/pipeline/Unet.swift

@@ -0,0 +1,204 @@
+// For licensing see accompanying LICENSE.md file.
+// Copyright (C) 2022 Apple Inc. All Rights Reserved.
+
+import Foundation
+import CoreML
+
+/// U-Net noise prediction model for stable diffusion
+@available(iOS 16.2, macOS 13.1, *)
+public struct Unet: ResourceManaging {
+
+    /// Model used to predict noise residuals given an input, diffusion time step, and conditional embedding
+    ///
+    /// It can be in the form of a single model or multiple stages
+    var models: [ManagedMLModel]
+
+    /// Creates a U-Net noise prediction model
+    ///
+    /// - Parameters:
+    ///   - url: Location of single U-Net  compiled Core ML model
+    ///   - configuration: Configuration to be used when the model is loaded
+    /// - Returns: U-net model that will lazily load its required resources when needed or requested
+    public init(modelAt url: URL,
+                configuration: MLModelConfiguration) {
+        self.models = [ManagedMLModel(modelAt: url, configuration: configuration)]
+    }
+
+    /// Creates a U-Net noise prediction model
+    ///
+    /// - Parameters:
+    ///   - urls: Location of chunked U-Net via urls to each compiled chunk
+    ///   - configuration: Configuration to be used when the model is loaded
+    /// - Returns: U-net model that will lazily load its required resources when needed or requested
+    public init(chunksAt urls: [URL],
+                configuration: MLModelConfiguration) {
+        self.models = urls.map { ManagedMLModel(modelAt: $0, configuration: configuration) }
+    }
+
+    /// Load resources.
+    public func loadResources() throws {
+        for model in models {
+            try model.loadResources()
+        }
+    }
+
+    /// Unload the underlying model to free up memory
+    public func unloadResources() {
+        for model in models {
+            model.unloadResources()
+        }
+    }
+
+    /// Pre-warm resources
+    public func prewarmResources() throws {
+        // Override default to pre-warm each model
+        for model in models {
+            try model.loadResources()
+            model.unloadResources()
+        }
+    }
+
+    var latentSampleDescription: MLFeatureDescription {
+        try! models.first!.perform { model in
+            model.modelDescription.inputDescriptionsByName["sample"]!
+        }
+    }
+
+    /// The expected shape of the models latent sample input
+    public var latentSampleShape: [Int] {
+        latentSampleDescription.multiArrayConstraint!.shape.map { $0.intValue }
+    }
+
+    var latentTimeIdDescription: MLFeatureDescription {
+        try! models.first!.perform { model in
+            model.modelDescription.inputDescriptionsByName["time_ids"]!
+        }
+    }
+
+    /// The expected shape of the geometry conditioning
+    public var latentTimeIdShape: [Int] {
+        latentTimeIdDescription.multiArrayConstraint!.shape.map { $0.intValue }
+    }
+
+    /// Batch prediction noise from latent samples
+    ///
+    /// - Parameters:
+    ///   - latents: Batch of latent samples in an array
+    ///   - timeStep: Current diffusion timestep
+    ///   - hiddenStates: Hidden state to condition on
+    /// - Returns: Array of predicted noise residuals
+    func predictNoise(
+        latents: [MLShapedArray<Float32>],
+        timeStep: Int,
+        hiddenStates: MLShapedArray<Float32>,
+        additionalResiduals: [[String: MLShapedArray<Float32>]]? = nil
+    ) throws -> [MLShapedArray<Float32>] {
+
+        // Match time step batch dimension to the model / latent samples
+        let t: MLShapedArray<Float32>
+        if hiddenStates.shape[0] == 2 {
+            t = MLShapedArray(scalars: [Float(timeStep), Float(timeStep)], shape: [2])
+        } else {
+            t = MLShapedArray(scalars: [Float(timeStep)], shape: [1])
+        }
+
+        // Form batch input to model
+        let inputs = try latents.enumerated().map {
+            var dict: [String: Any] = [
+                "sample" : MLMultiArray($0.element),
+                "timestep" : MLMultiArray(t),
+                "encoder_hidden_states": MLMultiArray(hiddenStates)
+            ]
+            if let residuals = additionalResiduals?[$0.offset] {
+                for (k, v) in residuals {
+                    dict[k] = MLMultiArray(v)
+                }
+            }
+            return try MLDictionaryFeatureProvider(dictionary: dict)
+        }
+        let batch = MLArrayBatchProvider(array: inputs)
+
+        // Make predictions
+        let results = try models.predictions(from: batch)
+
+        // Pull out the results in Float32 format
+        let noise = (0..<results.count).map { i in
+
+            let result = results.features(at: i)
+            let outputName = result.featureNames.first!
+
+            let outputNoise = result.featureValue(for: outputName)!.multiArrayValue!
+
+            // To conform to this func return type make sure we return float32
+            // Use the fact that the concatenating constructor for MLMultiArray
+            // can do type conversion:
+            let fp32Noise = MLMultiArray(
+                concatenating: [outputNoise],
+                axis: 0,
+                dataType: .float32
+            )
+            return MLShapedArray<Float32>(fp32Noise)
+        }
+
+        return noise
+    }
+
+    /// Batch prediction noise from latent samples, for Stable Diffusion XL
+    ///
+    /// - Parameters:
+    ///   - latents: Batch of latent samples in an array
+    ///   - timeStep: Current diffusion timestep
+    ///   - hiddenStates: Hidden state to condition on
+    ///   - pooledStates: Additional text states to condition on
+    ///   - geometryConditioning: Condition on image geometry
+    /// - Returns: Array of predicted noise residuals
+    @available(iOS 17.0, macOS 14.0, *)
+    func predictNoise(
+        latents: [MLShapedArray<Float32>],
+        timeStep: Int,
+        hiddenStates: MLShapedArray<Float32>,
+        pooledStates: MLShapedArray<Float32>,
+        geometryConditioning: MLShapedArray<Float32>
+    ) throws -> [MLShapedArray<Float32>] {
+
+        // Match time step batch dimension to the model / latent samples
+        let t = MLShapedArray<Float32>(scalars:[Float(timeStep), Float(timeStep)],shape:[2])
+
+        // Form batch input to model
+        let inputs = try latents.enumerated().map {
+            let dict: [String: Any] = [
+                "sample" : MLMultiArray($0.element),
+                "timestep" : MLMultiArray(t),
+                "encoder_hidden_states": MLMultiArray(hiddenStates),
+                "text_embeds": MLMultiArray(pooledStates),
+                "time_ids": MLMultiArray(geometryConditioning)
+            ]
+            return try MLDictionaryFeatureProvider(dictionary: dict)
+        }
+        let batch = MLArrayBatchProvider(array: inputs)
+
+        // Make predictions
+        let results = try models.predictions(from: batch)
+
+        // Pull out the results in Float32 format
+        let noise = (0..<results.count).map { i in
+
+            let result = results.features(at: i)
+            let outputName = result.featureNames.first!
+
+            let outputNoise = result.featureValue(for: outputName)!.multiArrayValue!
+
+            // To conform to this func return type make sure we return float32
+            // Use the fact that the concatenating constructor for MLMultiArray
+            // can do type conversion:
+            let fp32Noise = MLMultiArray(
+                concatenating: [outputNoise],
+                axis: 0,
+                dataType: .float32
+            )
+            return MLShapedArray<Float32>(fp32Noise)
+        }
+
+        return noise
+    }
+}

swift/StableDiffusion/tokenizer/BPETokenizer+Reading.swift

@@ -0,0 +1,58 @@
+// For licensing see accompanying LICENSE.md file.
+// Copyright (C) 2022 Apple Inc. All Rights Reserved.
+
+import Foundation
+
+@available(iOS 16.2, macOS 13.1, *)
+extension BPETokenizer {
+    enum FileReadError: Error {
+        case invalidMergeFileLine(Int)
+    }
+
+    /// Read vocab.json file at URL into a dictionary mapping a String to its Int token id
+    static func readVocabulary(url: URL) throws -> [String: Int] {
+        let content = try Data(contentsOf: url)
+        return try JSONDecoder().decode([String: Int].self, from: content)
+    }
+
+    /// Read merges.txt file at URL into a dictionary mapping bigrams to the line number/rank/priority
+    static func readMerges(url: URL) throws -> [TokenPair: Int] {
+        let data = try Data(contentsOf: url)
+        var merges = [TokenPair: Int]()
+        var index = 0
+        var line = [UInt8]()
+        for byte in data {
+            if byte == UInt8(ascii: "\n") {
+                if let pair = try parseMergesLine(line, index: index) {
+                    merges[pair] = index
+                }
+                line.removeAll(keepingCapacity: true)
+                index += 1
+            } else {
+                line.append(byte)
+            }
+        }
+
+        return merges
+    }
+
+    static func parseMergesLine(_ line: [UInt8], index: Int) throws -> TokenPair? {
+        if line.isEmpty || line.first == UInt8(ascii: "#") {
+            return nil
+        }
+        let pair = line.split(separator: UInt8(ascii: " "))
+        if pair.count != 2 {
+            throw FileReadError.invalidMergeFileLine(index + 1)
+        }
+        return TokenPair( String(bytes: pair[0]), String(bytes: pair[1]))
+    }
+}
+
+extension String {
+    init(bytes: some Collection<UInt8>) {
+        self.init(unsafeUninitializedCapacity: bytes.count) { pointer in
+            _ = pointer.initialize(fromContentsOf: bytes)
+            return bytes.count
+        }
+    }
+}

swift/StableDiffusion/tokenizer/BPETokenizer.swift

@@ -0,0 +1,185 @@
+// For licensing see accompanying LICENSE.md file.
+// Copyright (C) 2022 Apple Inc. All Rights Reserved.
+
+import Foundation
+
+/// A tokenizer based on byte pair encoding.
+@available(iOS 16.2, macOS 13.1, *)
+public struct BPETokenizer {
+    /// A dictionary that maps pairs of tokens to the rank/order of the merge.
+    let merges: [TokenPair : Int]
+
+    /// A dictionary from of tokens to identifiers.
+    let vocabulary: [String: Int]
+
+    /// The token used for padding
+    let padToken: String
+
+    /// The start token.
+    let startToken: String = "<|startoftext|>"
+
+    /// The end token.
+    let endToken: String = "<|endoftext|>"
+
+    /// The unknown token.
+    let unknownToken: String = "<|endoftext|>"
+
+    var unknownTokenID: Int {
+        vocabulary[unknownToken, default: 0]
+    }
+
+    /// Creates a tokenizer.
+    ///
+    /// - Parameters:
+    ///   - merges: A dictionary that maps pairs of tokens to the rank/order of the merge.
+    ///   - vocabulary: A dictionary from of tokens to identifiers.
+    public init(merges: [TokenPair: Int], vocabulary: [String: Int], padToken: String = "<|endoftext|>") {
+        self.merges = merges
+        self.vocabulary = vocabulary
+        self.padToken = padToken
+    }
+
+    /// Creates a tokenizer by loading merges and vocabulary from URLs.
+    ///
+    /// - Parameters:
+    ///   - mergesURL: The URL of a text file containing merges.
+    ///   - vocabularyURL: The URL of a JSON file containing the vocabulary.
+    public init(mergesAt mergesURL: URL, vocabularyAt vocabularyURL: URL, padToken: String = "<|endoftext|>") throws {
+        self.merges = try Self.readMerges(url: mergesURL)
+        self.vocabulary = try! Self.readVocabulary(url: vocabularyURL)
+        self.padToken = padToken
+    }
+
+    /// Tokenizes an input string.
+    ///
+    /// - Parameters:
+    ///   - input: A string.
+    ///   - minCount: The minimum number of tokens to return.
+    /// - Returns: An array of tokens and an array of token identifiers.
+    public func tokenize(input: String, minCount: Int? = nil) -> (tokens: [String], tokenIDs: [Int]) {
+        var tokens: [String] = []
+
+        tokens.append(startToken)
+        tokens.append(contentsOf: encode(input: input))
+        tokens.append(endToken)
+
+        // Pad if there was a min length specified
+        if let minLen = minCount, minLen > tokens.count {
+            tokens.append(contentsOf: repeatElement(padToken, count: minLen - tokens.count))
+        }
+
+        let ids = tokens.map({ vocabulary[$0, default: unknownTokenID] })
+        return (tokens: tokens, tokenIDs: ids)
+    }
+
+    /// Returns the token identifier for a token.
+    public func tokenID(for token: String) -> Int? {
+        vocabulary[token]
+    }
+
+    /// Returns the token for a token identifier.
+    public func token(id: Int) -> String? {
+        vocabulary.first(where: { $0.value == id })?.key
+    }
+
+    /// Decodes a sequence of tokens into a fully formed string
+    public func decode(tokens: [String]) -> String {
+        String(tokens.joined())
+            .replacingOccurrences(of: "</w>", with: " ")
+            .replacingOccurrences(of: startToken, with: "")
+            .replacingOccurrences(of: endToken, with: "")
+    }
+
+    /// Encode an input string to a sequence of tokens
+    func encode(input: String) -> [String] {
+        let normalized = input.trimmingCharacters(in: .whitespacesAndNewlines).lowercased()
+        let words = normalized.split(separator: " ")
+        return words.flatMap({ encode(word: $0) })
+    }
+
+    /// Encode a single word into a sequence of tokens
+    func encode(word: Substring) -> [String] {
+        var tokens = word.map { String($0) }
+        if let last = tokens.indices.last {
+            tokens[last] = tokens[last] + "</w>"
+        }
+
+        while true {
+            let pairs = pairs(for: tokens)
+            let canMerge = pairs.filter { merges[$0] != nil }
+
+            if canMerge.isEmpty {
+                break
+            }
+
+            // If multiple merges are found, use the one with the lowest rank
+            let shouldMerge = canMerge.min { merges[$0]! < merges[$1]! }!
+            tokens = update(tokens, merging: shouldMerge)
+        }
+        return tokens
+    }
+
+    /// Get  the set of adjacent pairs / bigrams from a sequence of tokens
+    func pairs(for tokens: [String]) -> Set<TokenPair> {
+        guard tokens.count > 1 else {
+            return Set()
+        }
+
+        var pairs = Set<TokenPair>(minimumCapacity: tokens.count - 1)
+        var prev = tokens.first!
+        for current in tokens.dropFirst() {
+            pairs.insert(TokenPair(prev, current))
+            prev = current
+        }
+        return pairs
+    }
+
+    /// Update the sequence of tokens by greedily merging instance of a specific bigram
+    func update(_ tokens: [String], merging bigram: TokenPair) -> [String] {
+        guard tokens.count > 1 else {
+            return []
+        }
+
+        var newTokens = [String]()
+        newTokens.reserveCapacity(tokens.count - 1)
+
+        var index = 0
+        while index < tokens.count {
+            let remainingTokens = tokens[index...]
+            if let startMatchIndex = remainingTokens.firstIndex(of: bigram.first) {
+                // Found a possible match, append everything before it
+                newTokens.append(contentsOf: tokens[index..<startMatchIndex])
+
+                if index < tokens.count - 1 && tokens[startMatchIndex + 1] == bigram.second {
+                    // Full match, merge
+                    newTokens.append(bigram.first + bigram.second)
+                    index = startMatchIndex + 2
+                } else {
+                    // Only matched the first, no merge
+                    newTokens.append(bigram.first)
+                    index = startMatchIndex + 1
+                }
+            } else {
+                // Didn't find any more matches, append the rest unmerged
+                newTokens.append(contentsOf: remainingTokens)
+                break
+            }
+        }
+        return newTokens
+    }
+}
+
+@available(iOS 16.2, macOS 13.1, *)
+extension BPETokenizer {
+
+    /// A hashable tuple of strings
+    public struct TokenPair: Hashable {
+        let first: String
+        let second: String
+
+        init(_ first: String, _ second: String) {
+            self.first = first
+            self.second = second
+        }
+    }
+}

swift/StableDiffusion/tokenizer/T5Tokenizer.swift

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