From 7bc8360e039236c99cbd569d98f1101bad7714dc Mon Sep 17 00:00:00 2001 From: Alex-Wengg Date: Wed, 19 Aug 2026 10:44:44 -0400 Subject: [PATCH] fix(tts/kokoro-ane): adopt COLA-corrected KokoroTail_v2 + native output level for all variants MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Point all three variants (ANE/, ANE-zh/, ANE-ja/) at KokoroTail_v2.mlmodelc and drop peak normalization for English/Mandarin, closing the follow-up scoped in #699 and reported in #852. The v1 tail omits the overlap-add/COLA normalization torch.istft applies (interior envelope sum(w^2) = 1.5 for periodic Hann, win 20 / hop 5), leaving raw output exactly 1.5x the PyTorch reference. #699 fixed the conversion and moved Japanese to native output level, but the corrected tail was never published — .japanese has been shipping unmasked 1.5x-hot audio. The corrected tails are now live as KokoroTail_v2.mlmodelc alongside the originals (kokoro-82m-coreml commit acac8811; conversion fix recommitted in mobius #83, reconstruction matches the #699-measured artifact to 1 ulp). The _v2 rename (same pattern as KokoroNoise_v2) doubles as cache invalidation: DownloadUtils skips files that already exist and the weight.bin size is unchanged, so a same-path replacement would never reach existing caches. With the rename, existing caches fetch just the missing file on next ensureModels. Verified e2e on M5 Pro: existing cache self-healed (only KokoroTail_v2 downloaded per variant); native render levels en peak 0.343 / zh 0.400 / ja 0.354 — reference territory (#699 measured PyTorch jf_alpha at 0.299), vs 1.0 forced by the old 0 dBFS normalization. CoreML A/B of v1 vs v2 tails: interior per-sample ratio 1.500000, corr 1.0 (pure scalar, no spectral change), for both the en/ja and zh tails. Fixes #852. --- Sources/FluidAudio/ModelNames.swift | 5 ++++- .../FluidAudio/TTS/KokoroAne/KokoroAneConstants.swift | 2 +- .../FluidAudio/TTS/KokoroAne/KokoroAneManager.swift | 10 +++++----- .../Pipeline/KokoroAneSynthesizer+Types.swift | 2 +- Sources/FluidAudioCLI/Commands/TTSCommand.swift | 4 +++- 5 files changed, 14 insertions(+), 9 deletions(-) diff --git a/Sources/FluidAudio/ModelNames.swift b/Sources/FluidAudio/ModelNames.swift index 6a015d237..edfa91bdb 100644 --- a/Sources/FluidAudio/ModelNames.swift +++ b/Sources/FluidAudio/ModelNames.swift @@ -1428,7 +1428,10 @@ public enum ModelNames { // re-download. See mobius laishere-coreml docs/trials-and-errors.md. public static let noise = "KokoroNoise_v2.mlmodelc" public static let vocoder = "KokoroVocoder.mlmodelc" - public static let tail = "KokoroTail.mlmodelc" + // v2: COLA-normalized iSTFT deconv weights (raw output was exactly 1.5x + // the PyTorch reference). Renamed (not overwritten) so cached clients + // re-download. See issue #852. + public static let tail = "KokoroTail_v2.mlmodelc" /// Auxiliary (non-CoreML) files that must accompany the mlmodelc bundles. public static let vocab = "vocab.json" diff --git a/Sources/FluidAudio/TTS/KokoroAne/KokoroAneConstants.swift b/Sources/FluidAudio/TTS/KokoroAne/KokoroAneConstants.swift index c42807a96..2ba2623af 100644 --- a/Sources/FluidAudio/TTS/KokoroAne/KokoroAneConstants.swift +++ b/Sources/FluidAudio/TTS/KokoroAne/KokoroAneConstants.swift @@ -16,7 +16,7 @@ public enum KokoroAneConstants { /// Default voice id for the Japanese (`ANE-ja/`) variant. public static let defaultVoiceJapanese = "jf_alpha" - /// Output sample rate of the iSTFT in `KokoroTail.mlpackage`. + /// Output sample rate of the iSTFT in `KokoroTail_v2.mlpackage`. public static let sampleRate = 24_000 /// BOS / EOS token id used by both `convert-coreml.py` and the iOS demo. diff --git a/Sources/FluidAudio/TTS/KokoroAne/KokoroAneManager.swift b/Sources/FluidAudio/TTS/KokoroAne/KokoroAneManager.swift index 80f02285c..e604ef525 100644 --- a/Sources/FluidAudio/TTS/KokoroAne/KokoroAneManager.swift +++ b/Sources/FluidAudio/TTS/KokoroAne/KokoroAneManager.swift @@ -358,14 +358,14 @@ public actor KokoroAneManager { private func wavData(from result: KokoroAneSynthesisResult) throws -> Data { do { - // Japanese writes at the model's native level (no peak-normalization) - // so the output matches the PyTorch reference instead of being - // slammed to 0 dBFS. English/Mandarin keep peak-normalization until - // their tails get the same COLA-corrected iSTFT (#698 follow-up). + // All variants write at the model's native level (no + // peak-normalization) so the output matches the PyTorch reference + // instead of being slammed to 0 dBFS. Requires the COLA-corrected + // KokoroTail_v2 (#852). return try AudioWAV.data( from: result.samples, sampleRate: Double(result.sampleRate), - normalize: variant != .japanese) + normalize: false) } catch { throw KokoroAneError.audioConversionFailed(error.localizedDescription) } diff --git a/Sources/FluidAudio/TTS/KokoroAne/Pipeline/KokoroAneSynthesizer+Types.swift b/Sources/FluidAudio/TTS/KokoroAne/Pipeline/KokoroAneSynthesizer+Types.swift index a7fddc7b5..c00de1884 100644 --- a/Sources/FluidAudio/TTS/KokoroAne/Pipeline/KokoroAneSynthesizer+Types.swift +++ b/Sources/FluidAudio/TTS/KokoroAne/Pipeline/KokoroAneSynthesizer+Types.swift @@ -96,7 +96,7 @@ public enum KokoroAneStage: String, CaseIterable, Sendable { case .prosody: return "KokoroProsody.mlmodelc" case .noise: return "KokoroNoise_v2.mlmodelc" // v2: atan2 phase-correction (HF-noise fix) case .vocoder: return "KokoroVocoder.mlmodelc" - case .tail: return "KokoroTail.mlmodelc" + case .tail: return "KokoroTail_v2.mlmodelc" // v2: COLA-normalized iSTFT (level fix, #852) } } } diff --git a/Sources/FluidAudioCLI/Commands/TTSCommand.swift b/Sources/FluidAudioCLI/Commands/TTSCommand.swift index 43575a473..0306d3088 100644 --- a/Sources/FluidAudioCLI/Commands/TTSCommand.swift +++ b/Sources/FluidAudioCLI/Commands/TTSCommand.swift @@ -879,10 +879,12 @@ public struct TTS { detailed = try await manager.synthesizeDetailed( text: text, voice: resolvedVoice, speed: 1.0) } + // Native level for all variants — matches the PyTorch reference + // now that KokoroTail_v2 carries the COLA-corrected iSTFT (#852). let wav = try AudioWAV.data( from: detailed.samples, sampleRate: Double(detailed.sampleRate), - normalize: variant != .japanese) + normalize: false) let tSynth1 = Date() let outURL = resolveInputURL(output)