Fix memory leak from non-async MLModel prediction (#56)

argmaxinc · Mar 7, 2024 · df15f89 · df15f89
1 parent 20c9001
commit df15f89
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 4 deletions.
diff --git a/Sources/WhisperKit/Core/AudioEncoder.swift b/Sources/WhisperKit/Core/AudioEncoder.swift
@@ -46,7 +46,7 @@ public class AudioEncoder: AudioEncoding, WhisperMLModel {
 
         try Task.checkCancellation()
 
-        let outputFeatures = try await model.prediction(from: modelInputs, options: MLPredictionOptions())
+        let outputFeatures = try await model.asyncPrediction(from: modelInputs, options: MLPredictionOptions())
 
         let output = AudioEncoderOutput(features: outputFeatures)
 

diff --git a/Sources/WhisperKit/Core/FeatureExtractor.swift b/Sources/WhisperKit/Core/FeatureExtractor.swift
@@ -35,7 +35,7 @@ public class FeatureExtractor: FeatureExtracting, WhisperMLModel {
 
         try Task.checkCancellation()
 
-        let outputFeatures = try await model.prediction(from: modelInputs, options: MLPredictionOptions())
+        let outputFeatures = try await model.asyncPrediction(from: modelInputs, options: MLPredictionOptions())
 
         let output = MelSpectrogramOutput(features: outputFeatures)
 

diff --git a/Sources/WhisperKit/Core/TextDecoder.swift b/Sources/WhisperKit/Core/TextDecoder.swift
@@ -176,7 +176,7 @@ public extension TextDecoding {
 
         try Task.checkCancellation()
 
-        let outputFeatures = try await prefillModel.prediction(from: modelInputs, options: MLPredictionOptions())
+        let outputFeatures = try await prefillModel.asyncPrediction(from: modelInputs, options: MLPredictionOptions())
 
         let output = TextDecoderCachePrefillOutput(features: outputFeatures)
 
@@ -291,7 +291,7 @@ public class TextDecoder: TextDecoding, WhisperMLModel {
 
         try Task.checkCancellation()
 
-        let outputFeatures = try await model.prediction(from: modelInputs, options: MLPredictionOptions())
+        let outputFeatures = try await model.asyncPrediction(from: modelInputs, options: MLPredictionOptions())
 
         let output = TextDecoderOutput(features: outputFeatures)
 

diff --git a/Sources/WhisperKit/Core/Utils.swift b/Sources/WhisperKit/Core/Utils.swift
@@ -39,6 +39,21 @@ extension MLMultiArray {
     }
 }
 
+extension MLModel {
+    func asyncPrediction(
+        from input: MLFeatureProvider,
+        options: MLPredictionOptions
+    ) async throws -> MLFeatureProvider {
+        if #available(macOS 14, iOS 17, watchOS 10, visionOS 1, *) {
+            return try await prediction(from: input, options: options)
+        } else {
+            return try await Task {
+                try prediction(from: input, options: options)
+            }.value
+        }
+    }
+}
+
 @available(macOS 13, iOS 16, watchOS 10, visionOS 1, *)
 func initMLMultiArray(shape: [NSNumber], dataType: MLMultiArrayDataType, initialValue: Any) -> MLMultiArray {
     let multiArray = try! MLMultiArray(shape: shape, dataType: dataType)