Use GPU for audio encoder on macOS 13 (#83)

* Use gpu on macos 13 for the audio encoder * Use gpu on macos 13 for the audio encoder * Update Sources/WhisperKitCLI/CLIArguments.swift
argmaxinc · Mar 19, 2024 · cf75348 · cf75348
1 parent bd75392
commit cf75348
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 6 deletions.
diff --git a/Sources/WhisperKit/Core/Models.swift b/Sources/WhisperKit/Core/Models.swift
@@ -135,7 +135,7 @@ public struct ModelComputeOptions {
 
     public init(
         melCompute: MLComputeUnits = .cpuAndGPU,
-        audioEncoderCompute: MLComputeUnits = .cpuAndNeuralEngine,
+        audioEncoderCompute: MLComputeUnits? = nil,
         textDecoderCompute: MLComputeUnits = .cpuAndNeuralEngine,
         prefillCompute: MLComputeUnits = .cpuOnly
     ) {
@@ -146,10 +146,16 @@ public struct ModelComputeOptions {
             self.prefillCompute = .cpuOnly
             return
         }
+
         self.melCompute = melCompute
-        self.audioEncoderCompute = audioEncoderCompute
-        self.textDecoderCompute = textDecoderCompute
         self.prefillCompute = prefillCompute
+        self.textDecoderCompute = textDecoderCompute
+
+        if #available(macOS 14.0, iOS 17.0, watchOS 10, visionOS 1, *) {
+            self.audioEncoderCompute = audioEncoderCompute ?? .cpuAndNeuralEngine
+        } else {
+            self.audioEncoderCompute = audioEncoderCompute ?? .cpuAndGPU
+        }
     }
 }
 

diff --git a/Sources/WhisperKitCLI/CLIArguments.swift b/Sources/WhisperKitCLI/CLIArguments.swift
@@ -21,7 +21,6 @@ struct CLIArguments: ParsableArguments {
 
     @Option(help: "Compute units for audio encoder model with {all,cpuOnly,cpuAndGPU,cpuAndNeuralEngine,random}")
     var audioEncoderComputeUnits: ComputeUnits = .cpuAndNeuralEngine
-
     @Option(help: "Compute units for text decoder model with {all,cpuOnly,cpuAndGPU,cpuAndNeuralEngine,random}")
     var textDecoderComputeUnits: ComputeUnits = .cpuAndNeuralEngine
 

diff --git a/Sources/WhisperKitCLI/Transcribe.swift b/Sources/WhisperKitCLI/Transcribe.swift
@@ -32,9 +32,19 @@ struct Transcribe: AsyncParsableCommand {
             print("Transcribing audio at \(cliArguments.audioPath)")
         }
 
+        var audioEncoderComputeUnits = cliArguments.audioEncoderComputeUnits.asMLComputeUnits
+        let textDecoderComputeUnits = cliArguments.textDecoderComputeUnits.asMLComputeUnits
+
+        // Use gpu for audio encoder on macOS below 14
+        if audioEncoderComputeUnits == .cpuAndNeuralEngine {
+            if #unavailable(macOS 14.0) {
+                audioEncoderComputeUnits = .cpuAndGPU
+            }
+        }
+
         let computeOptions = ModelComputeOptions(
-            audioEncoderCompute: cliArguments.audioEncoderComputeUnits.asMLComputeUnits,
-            textDecoderCompute: cliArguments.textDecoderComputeUnits.asMLComputeUnits
+            audioEncoderCompute: audioEncoderComputeUnits,
+            textDecoderCompute: textDecoderComputeUnits
         )
 
         let downloadTokenizerFolder: URL? =