feat(tts): expose Kokoro CoreML compute units through the binding

examples/kokoro.rs

@@ -1,23 +1,34 @@
 //! Example: synthesize TTS via FluidAudio Kokoro.
 //! Usage: cargo run --example kokoro --features tts -- "Hello world" af_heart en-us > out.wav
 //!        cargo run --example kokoro --features tts -- "你好" zf_001 zh > out.wav
+//!        cargo run --example kokoro --features tts -- "Hi" af_heart en-us cpu-and-gpu > out.wav
 //! `lang` selects the KokoroAne variant (`zh` → Mandarin, else English).
+//! The optional 4th arg picks CoreML compute units (`default`, `all-ane`,
+//! `cpu-and-gpu`, `cpu-only`). Hosts without a usable Neural Engine — a
+//! virtualised macOS guest, e.g. a GitHub-hosted `macos-14` runner — need
+//! `cpu-and-gpu`; the default mapping pins four stages to the ANE and fails
+//! there with "Failed to prepare the model for predictions".
 //! `af_heart` (English) and `zf_001` (Mandarin) are the built-in default voices.
 //! Note: English currently hosts only `af_heart`; Mandarin hosts ~100 voices (`zf_*`, `zm_*`, …)
 //! that download on demand from HuggingFace on first use. An id absent from the hosted bundle
 //! fails to load — or pre-stage it as `<voice>.bin` in the model cache.
-use fluidaudio_rs::FluidAudio;
+use fluidaudio_rs::{FluidAudio, KokoroComputeUnits};
 use std::io::Write;
 
 fn main() -> Result<(), Box<dyn std::error::Error>> {
     let args: Vec<String> = std::env::args().collect();
     let text = args.get(1).map(String::as_str).unwrap_or("Hello world");
     let voice = args.get(2).map(String::as_str).unwrap_or("af_heart");
     let lang = args.get(3).map(String::as_str).unwrap_or("en-us");
+    let compute_units: KokoroComputeUnits =
+        args.get(4).map_or(Ok(Default::default()), |s| s.parse())?;
 
     let audio = FluidAudio::new()?;
-    eprintln!("Initializing Kokoro (downloads model on first run)...");
-    audio.init_kokoro(voice, lang)?;
+    eprintln!(
+        "Initializing Kokoro on {} compute units (downloads model on first run)...",
+        compute_units.as_str()
+    );
+    audio.init_kokoro_with_compute_units(voice, lang, compute_units)?;
     eprintln!("Kokoro available: {}", audio.is_kokoro_available());
 
     let wav = audio.synthesize_kokoro(text, voice, 1.0)?;

src/ffi/bridge.rs

@@ -188,6 +188,12 @@ extern "C" {
         default_voice: *const i8,
         lang: *const i8,
     ) -> i32;
+    fn fluidaudio_initialize_kokoro_with_compute_units(
+        bridge: *mut std::ffi::c_void,
+        default_voice: *const i8,
+        lang: *const i8,
+        compute_units: *const i8,
+    ) -> i32;
     fn fluidaudio_kokoro_synthesize(
         bridge: *mut std::ffi::c_void,
         text: *const i8,
@@ -261,6 +267,47 @@ impl FluidAudioBridge {
         }
     }
 
+    /// The C entry point reports only success or `-1`, so the returned error
+    /// names the preset and nothing else. The actionable text — the CoreML
+    /// domain and code, or the download status for a missing voice pack — is
+    /// written to **stderr** by the Swift side, not carried back across the
+    /// boundary. A caller diagnosing a failure has to read it there. Two real
+    /// examples of what only stderr shows:
+    ///
+    /// ```text
+    /// Kokoro init error: invalidResponse(description: "am_michael voice pack", statusCode: 404)
+    /// Kokoro synthesize error: predictionFailed(stage: "vocoder", ... "Failed to prepare
+    ///   the model for predictions. ML program was KokoroVocoder ...")
+    /// ```
+    ///
+    /// Both are indistinguishable from here — a missing voice and an
+    /// unusable Neural Engine arrive as the same `-1`.
+    pub fn initialize_kokoro_with_compute_units(
+        &self,
+        default_voice: &str,
+        lang: &str,
+        compute_units: &str,
+    ) -> Result<(), String> {
+        let c_voice = CString::new(default_voice).map_err(|_| "Invalid voice")?;
+        let c_lang = CString::new(lang).map_err(|_| "Invalid lang")?;
+        let c_units = CString::new(compute_units).map_err(|_| "Invalid compute units")?;
+        let result = unsafe {
+            fluidaudio_initialize_kokoro_with_compute_units(
+                self.ptr,
+                c_voice.as_ptr(),
+                c_lang.as_ptr(),
+                c_units.as_ptr(),
+            )
+        };
+        if result == 0 {
+            Ok(())
+        } else {
+            Err(format!(
+                "Failed to initialize Kokoro (compute units: {compute_units})"
+            ))
+        }
+    }
+
     /// Synthesize `text` with `voice` at `speed`; returns the complete WAV bytes
     /// produced by FluidAudio's KokoroAneManager (24 kHz mono 16-bit PCM (i16), peak-normalized).
     pub fn kokoro_synthesize(&self, text: &str, voice: &str, speed: f32) -> Result<Vec<u8>, String> {

src/lib.rs

@@ -65,6 +65,55 @@ impl From<String> for FluidAudioError {
     }
 }
 
+/// CoreML compute-unit preset for the Kokoro TTS pipeline.
+///
+/// Mirrors FluidAudio's `TtsComputeUnitPreset`; [`Self::as_str`] emits the same
+/// kebab-case spellings its `init?(cliValue:)` parser accepts, so the value
+/// round-trips across the FFI boundary.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
+pub enum KokoroComputeUnits {
+    /// FluidAudio's empirical per-stage mapping (Albert / PostAlbert /
+    /// Alignment / Vocoder on the Neural Engine).
+    #[default]
+    Default,
+    /// Every stage on `.cpuAndNeuralEngine`.
+    AllAne,
+    /// Every stage on `.cpuAndGPU` — skips the ANE entirely.
+    CpuAndGpu,
+    /// Every stage on `.cpuOnly`.
+    CpuOnly,
+}
+
+impl KokoroComputeUnits {
+    /// Canonical kebab-case name, as accepted by FluidAudio's
+    /// `TtsComputeUnitPreset(cliValue:)`.
+    pub fn as_str(self) -> &'static str {
+        match self {
+            Self::Default => "default",
+            Self::AllAne => "all-ane",
+            Self::CpuAndGpu => "cpu-and-gpu",
+            Self::CpuOnly => "cpu-only",
+        }
+    }
+}
+
+impl std::str::FromStr for KokoroComputeUnits {
+    type Err = FluidAudioError;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s.to_ascii_lowercase().as_str() {
+            "default" => Ok(Self::Default),
+            "all-ane" | "ane" | "neural-engine" => Ok(Self::AllAne),
+            "cpu-and-gpu" | "cpuandgpu" | "gpu" => Ok(Self::CpuAndGpu),
+            "cpu-only" | "cpu" | "cpuonly" => Ok(Self::CpuOnly),
+            other => Err(FluidAudioError::BridgeError(format!(
+                "unknown Kokoro compute-units preset '{other}' \
+                 (expected: default, all-ane, cpu-and-gpu, cpu-only)"
+            ))),
+        }
+    }
+}
+
 /// Main FluidAudio interface for Rust
 ///
 /// Provides access to ASR and VAD functionality.
@@ -445,12 +494,41 @@ impl FluidAudio {
     /// `lang` selects the KokoroAne variant in the Swift bridge (`zh` → Mandarin,
     /// everything else → English). Downloads the variant's model on first run
     /// (FluidAudio-managed cache).
+    ///
+    /// Uses [`KokoroComputeUnits::Default`]; see
+    /// [`init_kokoro_with_compute_units`](Self::init_kokoro_with_compute_units)
+    /// when the host has no usable Neural Engine.
     pub fn init_kokoro(&self, default_voice: &str, lang: &str) -> Result<(), FluidAudioError> {
         self.bridge
             .initialize_kokoro(default_voice, lang)
             .map_err(FluidAudioError::from)
     }
 
+    /// Initialize Kokoro TTS on an explicit CoreML compute-unit preset.
+    ///
+    /// FluidAudio's default per-stage mapping pins the Albert, PostAlbert,
+    /// Alignment and Vocoder stages to the Neural Engine. Where no ANE is
+    /// exposed — notably a virtualised macOS guest, such as a GitHub-hosted
+    /// `macos-14` runner — CoreML defers the failure past model load: this call
+    /// succeeds, and the *first* [`synthesize_kokoro`](Self::synthesize_kokoro)
+    /// then fails with `predictionFailed(stage: "vocoder", ...)` wrapping
+    /// "Failed to prepare the model for predictions". `initialize` only
+    /// downloads and loads the mlmodelcs, so there is no prediction at init to
+    /// surface the problem earlier. Passing [`KokoroComputeUnits::CpuAndGpu`]
+    /// (or `CpuOnly`) keeps synthesis working there, and doubles as the
+    /// debugging baseline FluidAudio's `KokoroAne.md` recommends for artefact
+    /// investigations.
+    pub fn init_kokoro_with_compute_units(
+        &self,
+        default_voice: &str,
+        lang: &str,
+        compute_units: KokoroComputeUnits,
+    ) -> Result<(), FluidAudioError> {
+        self.bridge
+            .initialize_kokoro_with_compute_units(default_voice, lang, compute_units.as_str())
+            .map_err(FluidAudioError::from)
+    }
+
     /// Synthesize text via Kokoro TTS in the language the engine was initialized
     /// with (`init_kokoro`'s `lang`): English by default, or Mandarin when
     /// `lang` was `zh` (the `.mandarin` KokoroAne variant).
@@ -804,4 +882,72 @@ mod tests {
         // For now, just test the types exist
         let _ = FluidAudioError::NotInitialized("test".to_string());
     }
+
+    #[test]
+    fn compute_units_emit_the_spellings_fluidaudio_parses() {
+        // The FFI passes the string, not the enum, so `as_str` has to land on a
+        // case of FluidAudio's `TtsComputeUnitPreset.init?(cliValue:)`. Nothing
+        // in this crate can execute that Swift parser, so the expectations below
+        // are transcribed from it and must be re-checked when the pinned
+        // FluidAudio version moves (Package.swift, currently 0.14.8 —
+        // Sources/FluidAudio/TTS/Shared/TtsComputeUnitPreset.swift). Asserting
+        // the literals — not just a `FromStr` round-trip — is what makes a
+        // rename of `as_str`'s output fail here instead of at runtime on the
+        // Swift side.
+        for (units, cli_value) in [
+            (KokoroComputeUnits::Default, "default"),
+            (KokoroComputeUnits::AllAne, "all-ane"),
+            (KokoroComputeUnits::CpuAndGpu, "cpu-and-gpu"),
+            (KokoroComputeUnits::CpuOnly, "cpu-only"),
+        ] {
+            assert_eq!(units.as_str(), cli_value);
+            assert_eq!(cli_value.parse::<KokoroComputeUnits>().unwrap(), units);
+        }
+        assert_eq!(KokoroComputeUnits::default(), KokoroComputeUnits::Default);
+    }
+
+    #[test]
+    fn compute_units_accepts_aliases_and_rejects_junk() {
+        // Every alias FluidAudio 0.14.8's `init?(cliValue:)` accepts, so a
+        // caller that learned a spelling from FluidAudio's own `--compute-units`
+        // flag is not rejected here before the string ever reaches Swift. Same
+        // manual-sync caveat as the test above.
+        for (spelling, expected) in [
+            ("default", KokoroComputeUnits::Default),
+            ("all-ane", KokoroComputeUnits::AllAne),
+            ("ane", KokoroComputeUnits::AllAne),
+            ("neural-engine", KokoroComputeUnits::AllAne),
+            ("cpu-and-gpu", KokoroComputeUnits::CpuAndGpu),
+            ("cpuandgpu", KokoroComputeUnits::CpuAndGpu),
+            ("gpu", KokoroComputeUnits::CpuAndGpu),
+            ("cpu-only", KokoroComputeUnits::CpuOnly),
+            ("cpu", KokoroComputeUnits::CpuOnly),
+            ("cpuonly", KokoroComputeUnits::CpuOnly),
+        ] {
+            assert_eq!(
+                spelling.parse::<KokoroComputeUnits>().unwrap(),
+                expected,
+                "alias {spelling:?} should parse"
+            );
+        }
+
+        // Swift lowercases before matching; so must we, or a mixed-case value
+        // would fail on this side and never reach the parser that accepts it.
+        assert_eq!(
+            "CPU-Only".parse::<KokoroComputeUnits>().unwrap(),
+            KokoroComputeUnits::CpuOnly
+        );
+
+        // Unknown presets fail in Rust, before the FFI call — the error names
+        // the canonical spellings rather than surfacing a bare `-1` from Swift.
+        let err = "tpu".parse::<KokoroComputeUnits>().unwrap_err().to_string();
+        assert!(
+            err.contains("tpu"),
+            "error should quote the bad value: {err}"
+        );
+        assert!(
+            err.contains("cpu-and-gpu"),
+            "error should list the accepted spellings: {err}"
+        );
+    }
 }

swift/FluidAudioBridge.swift

@@ -156,7 +156,18 @@ class FluidAudioBridgeInternal {
         }
     }
 
-    func initializeKokoro(defaultVoice: String, lang: String) throws {
+    /// `computeUnits` maps onto `KokoroAneComputeUnits`. `.default` keeps
+    /// FluidAudio's empirical per-stage mapping, which pins Albert / PostAlbert /
+    /// Alignment / Vocoder to the Neural Engine. Callers running where no ANE is
+    /// exposed — a virtualised macOS guest, e.g. a GitHub-hosted `macos-14`
+    /// runner — must pass `.cpuAndGpu` or `.cpuOnly`. Note the failure does not
+    /// land here: `KokoroAneManager.initialize` only downloads and loads the
+    /// mlmodelcs, so this call succeeds and `synthesizeKokoro` then throws
+    /// `predictionFailed(stage: "vocoder", ...)` wrapping "Failed to prepare the
+    /// model for predictions" on the first prediction.
+    func initializeKokoro(
+        defaultVoice: String, lang: String, computeUnits: TtsComputeUnitPreset = .default
+    ) throws {
         let semaphore = DispatchSemaphore(value: 0)
         var initError: Error?
 
@@ -165,7 +176,10 @@ class FluidAudioBridgeInternal {
                 // `KokoroAneManager` feeds `speed` as a real model input tensor,
                 // so `--rate` applies correctly (unlike the prior `voiceSpeed:` path).
                 let variant = Self.kokoroVariant(for: lang)
-                let manager = KokoroAneManager(variant: variant, defaultVoice: defaultVoice)
+                let manager = KokoroAneManager(
+                    variant: variant,
+                    defaultVoice: defaultVoice,
+                    computeUnits: KokoroAneComputeUnits(preset: computeUnits))
                 try await manager.initialize(preloadVoices: [defaultVoice])
                 self.kokoroManager = manager
             } catch {

swift/Kokoro_ffi.swift

@@ -1,3 +1,4 @@
+import FluidAudio
 import Foundation
 
 // MARK: - Kokoro TTS C FFI
@@ -13,25 +14,69 @@ private func kokoroLog(_ message: String) {
     FileHandle.standardError.write(Data((message + "\n").utf8))
 }
 
-@_cdecl("fluidaudio_initialize_kokoro")
-public func fluidaudio_initialize_kokoro(
+/// Shared body of both init entry points. `computeUnits` is the kebab-case preset
+/// name `TtsComputeUnitPreset(cliValue:)` accepts (`default`, `all-ane`,
+/// `cpu-and-gpu`, `cpu-only`). NULL or empty keeps the backend's empirical
+/// per-stage mapping. An unrecognised value is a caller bug, so it fails rather
+/// than silently synthesising on units the caller did not ask for.
+private func initializeKokoro(
     _ ptr: UnsafeMutableRawPointer?,
     _ defaultVoice: UnsafePointer<CChar>?,
-    _ lang: UnsafePointer<CChar>?
+    _ lang: UnsafePointer<CChar>?,
+    _ computeUnits: UnsafePointer<CChar>?
 ) -> Int32 {
     guard let ptr = ptr else { return -1 }
     let bridge = Unmanaged<FluidAudioBridgeInternal>.fromOpaque(ptr).takeUnretainedValue()
     let voice = defaultVoice.map { String(cString: $0) } ?? "af_heart"
     let langString = lang.map { String(cString: $0) } ?? ""
+    let unitsString = computeUnits.map { String(cString: $0) } ?? ""
+
+    let preset: TtsComputeUnitPreset
+    if unitsString.isEmpty {
+        preset = .default
+    } else if let parsed = TtsComputeUnitPreset(cliValue: unitsString) {
+        preset = parsed
+    } else {
+        kokoroLog(
+            "Kokoro init error: unknown compute-units preset '\(unitsString)' "
+                + "(expected one of: \(TtsComputeUnitPreset.allCases.map(\.cliValue).joined(separator: ", ")))")
+        return -1
+    }
+
     do {
-        try bridge.initializeKokoro(defaultVoice: voice, lang: langString)
+        try bridge.initializeKokoro(defaultVoice: voice, lang: langString, computeUnits: preset)
         return 0
     } catch {
         kokoroLog("Kokoro init error: \(error)")
         return -1
     }
 }
 
+/// Initialize on FluidAudio's empirical per-stage compute-unit mapping.
+///
+/// Kept at three parameters so the existing C symbol's arity is unchanged —
+/// anything linking these entry points from an older build keeps working.
+@_cdecl("fluidaudio_initialize_kokoro")
+public func fluidaudio_initialize_kokoro(
+    _ ptr: UnsafeMutableRawPointer?,
+    _ defaultVoice: UnsafePointer<CChar>?,
+    _ lang: UnsafePointer<CChar>?
+) -> Int32 {
+    initializeKokoro(ptr, defaultVoice, lang, nil)
+}
+
+/// Same, but pins every pipeline stage to an explicit preset. See
+/// `initializeKokoro` for the accepted spellings and the failure contract.
+@_cdecl("fluidaudio_initialize_kokoro_with_compute_units")
+public func fluidaudio_initialize_kokoro_with_compute_units(
+    _ ptr: UnsafeMutableRawPointer?,
+    _ defaultVoice: UnsafePointer<CChar>?,
+    _ lang: UnsafePointer<CChar>?,
+    _ computeUnits: UnsafePointer<CChar>?
+) -> Int32 {
+    initializeKokoro(ptr, defaultVoice, lang, computeUnits)
+}
+
 /// Synthesize `text` with `voice` at `speed`; returns a complete WAV byte buffer
 /// (24 kHz mono 16-bit PCM (i16), peak-normalized) via `outBytes`/`outLen`. The caller owns the buffer and must
 /// free it with `fluidaudio_kokoro_free_bytes`.