chore: import upstream snapshot with attribution

2026-07-13 13:28:46 +08:00
commit 8199cf3c39
90 changed files with 62775 additions and 0 deletions
@@ -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)
+    }
+}