diff --git a/examples/kokoro.rs b/examples/kokoro.rs index 7de730d..3d01ae1 100644 --- a/examples/kokoro.rs +++ b/examples/kokoro.rs @@ -1,12 +1,18 @@ //! 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 `.bin` in the model cache. -use fluidaudio_rs::FluidAudio; +use fluidaudio_rs::{FluidAudio, KokoroComputeUnits}; use std::io::Write; fn main() -> Result<(), Box> { @@ -14,10 +20,15 @@ fn main() -> Result<(), Box> { 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)?; diff --git a/src/ffi/bridge.rs b/src/ffi/bridge.rs index c73da21..628017d 100644 --- a/src/ffi/bridge.rs +++ b/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, String> { diff --git a/src/lib.rs b/src/lib.rs index 4cde3b5..5cdb1c3 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -65,6 +65,55 @@ impl From 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 { + 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::().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::().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::().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::().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}" + ); + } } diff --git a/swift/FluidAudioBridge.swift b/swift/FluidAudioBridge.swift index a88e23d..d7ba533 100644 --- a/swift/FluidAudioBridge.swift +++ b/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 { diff --git a/swift/Kokoro_ffi.swift b/swift/Kokoro_ffi.swift index 41b2e01..3dc26c7 100644 --- a/swift/Kokoro_ffi.swift +++ b/swift/Kokoro_ffi.swift @@ -1,3 +1,4 @@ +import FluidAudio import Foundation // MARK: - Kokoro TTS C FFI @@ -13,18 +14,37 @@ 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?, - _ lang: UnsafePointer? + _ lang: UnsafePointer?, + _ computeUnits: UnsafePointer? ) -> Int32 { guard let ptr = ptr else { return -1 } let bridge = Unmanaged.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)") @@ -32,6 +52,31 @@ public func fluidaudio_initialize_kokoro( } } +/// 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?, + _ lang: UnsafePointer? +) -> 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?, + _ lang: UnsafePointer?, + _ computeUnits: UnsafePointer? +) -> 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`.