diff --git a/Sources/FluidAudio/TTS/KokoroAne/KokoroAneManager.swift b/Sources/FluidAudio/TTS/KokoroAne/KokoroAneManager.swift index cd5f643b9..38fdb4b1c 100644 --- a/Sources/FluidAudio/TTS/KokoroAne/KokoroAneManager.swift +++ b/Sources/FluidAudio/TTS/KokoroAne/KokoroAneManager.swift @@ -15,7 +15,11 @@ import Foundation /// * One default voice per variant (`af_heart` for English, `zf_001` for /// Mandarin); additional voices download on demand via ``setDefaultVoice`` /// / `voice:` / `initialize(preloadVoices:)`. -/// * IPA input capped at 512 tokens — chunk longer prompts upstream. +/// * IPA input capped at 512 tokens. The high-level text API +/// (``synthesize(text:voice:speed:)`` / ``synthesizeDetailed(text:voice:speed:)``) +/// auto-chunks longer prompts at whitespace / pause punctuation (#712, #940); +/// the low-level ``synthesizeFromPhonemes(_:voice:speed:)`` stays strict +/// and throws ``KokoroAneError/phonemeSequenceTooLong(_:)`` past the cap. /// * Loads from HF path `kokoro-82m-coreml/ANE/` (English, Spanish, /// French), `ANE-zh/` (Mandarin) or `ANE-ja/` (Japanese). /// @@ -236,8 +240,23 @@ public actor KokoroAneManager { speed: Float = KokoroAneConstants.defaultSpeed ) async throws -> KokoroAneSynthesisResult { let frontend = try await resolveFrontend(for: text) - return try await runChain( - phonemes: frontend.phonemes, normalizedText: frontend.normalizedText, voice: voice, speed: speed) + // Chunk after normalization / G2P so written forms are never split; the + // cap is in Unicode scalars, as the vocab encoder counts it (#712, #940). + let chunks = PhonemeChunker.chunk( + frontend.phonemes, maxLength: KokoroAneConstants.maxPhonemeLength, countsUnicodeScalars: true) + guard chunks.count > 1 else { + return try await runChain( + phonemes: frontend.phonemes, normalizedText: frontend.normalizedText, voice: voice, speed: speed) + } + var parts: [KokoroAneSynthesisResult] = [] + for chunk in chunks { + try Task.checkCancellation() + parts.append(try await runChain(phonemes: chunk, normalizedText: nil, voice: voice, speed: speed)) + } + var result = KokoroAneSynthesisResult.concatenating(parts) + result.normalizedText = frontend.normalizedText + result.phonemes = frontend.phonemes + return result } /// Resolve the exact phoneme string ``synthesize(text:voice:speed:)`` diff --git a/Sources/FluidAudio/TTS/KokoroAne/Pipeline/KokoroAneSynthesizer+Types.swift b/Sources/FluidAudio/TTS/KokoroAne/Pipeline/KokoroAneSynthesizer+Types.swift index 971c62b65..01e2e5d93 100644 --- a/Sources/FluidAudio/TTS/KokoroAne/Pipeline/KokoroAneSynthesizer+Types.swift +++ b/Sources/FluidAudio/TTS/KokoroAne/Pipeline/KokoroAneSynthesizer+Types.swift @@ -42,7 +42,9 @@ public struct KokoroAneSynthesisResult: Sendable { public let acousticFrames: Int /// Token ids passed to the Kokoro chain, including BOS/EOS. /// - /// Indices align one-to-one with ``predictedDurations``. + /// Indices align one-to-one with ``predictedDurations``. When the text API + /// chunks long input, this is the per-chunk ids concatenated in order, so + /// each chunk contributes its own BOS/EOS pair. public let inputIds: [Int32] /// PostAlbert `pred_dur`: acoustic-frame counts for each input token. /// @@ -64,6 +66,8 @@ public struct KokoroAneSynthesisResult: Sendable { /// Phoneme string handed to the vocab encoder. ``inputIds`` is this string /// with characters missing from `vocab.json` dropped and BOS/EOS added, so /// the two lengths differ when the string carries out-of-vocab scalars. + /// For chunked text input this is the full resolved string; the chunks + /// drop the whitespace at each split and add one BOS/EOS pair per chunk. public internal(set) var phonemes: String /// Per-stage timings. public let timings: KokoroAneStageTimings @@ -94,6 +98,24 @@ public struct KokoroAneSynthesisResult: Sendable { self.phonemes = phonemes self.timings = timings } + + /// Join per-chunk results in order: samples, ids and durations are + /// concatenated; token/frame counts and stage timings are summed. Level is + /// left untouched (no per-chunk normalization). Text fields are left empty + /// for the caller to set. + static func concatenating(_ parts: [KokoroAneSynthesisResult]) -> KokoroAneSynthesisResult { + var timings = KokoroAneStageTimings() + for part in parts { timings.add(part.timings) } + return KokoroAneSynthesisResult( + samples: parts.flatMap(\.samples), + sampleRate: parts.first?.sampleRate ?? KokoroAneConstants.sampleRate, + encoderTokens: parts.reduce(0) { $0 + $1.encoderTokens }, + acousticFrames: parts.reduce(0) { $0 + $1.acousticFrames }, + timings: timings, + inputIds: parts.flatMap(\.inputIds), + predictedDurations: parts.flatMap(\.predictedDurations) + ) + } } /// One of the 7 stages in the laishere chain. diff --git a/Sources/FluidAudio/TTS/Shared/PhonemeChunker.swift b/Sources/FluidAudio/TTS/Shared/PhonemeChunker.swift index f5efa6f27..31ab66a44 100644 --- a/Sources/FluidAudio/TTS/Shared/PhonemeChunker.swift +++ b/Sources/FluidAudio/TTS/Shared/PhonemeChunker.swift @@ -31,17 +31,23 @@ enum PhonemeChunker { /// is hard-split at the cap (rare for real phoneme strings). Leading and /// trailing whitespace is trimmed from every chunk. /// + /// - Parameter countsUnicodeScalars: measure length in Unicode scalars + /// instead of `Character`s, for vocabs that encode combining marks as + /// their own token (KokoroAne: `ɑ̃` is two symbols). /// - Returns: `[phonemes]` (trimmed) when the input already fits, /// `[]` for blank input, and otherwise the ordered chunks. Length is - /// counted in `Character`s to match the cap TTS vocabularies enforce. + /// counted in `Character`s by default to match the cap TTS vocabularies + /// enforce. static func chunk( _ phonemes: String, maxLength: Int, - boundaryPunctuation: Set = defaultBoundaryPunctuation + boundaryPunctuation: Set = defaultBoundaryPunctuation, + countsUnicodeScalars: Bool = false ) -> [String] { precondition(maxLength > 0, "maxLength must be positive") - let characters = Array(phonemes) + let characters: [Character] = + countsUnicodeScalars ? phonemes.unicodeScalars.map(Character.init) : Array(phonemes) let count = characters.count if count == 0 { return [] } diff --git a/Tests/FluidAudioTests/TTS/KokoroAne/KokoroAneSynthesizerTests.swift b/Tests/FluidAudioTests/TTS/KokoroAne/KokoroAneSynthesizerTests.swift index 5f4efb5cd..f33677ab1 100644 --- a/Tests/FluidAudioTests/TTS/KokoroAne/KokoroAneSynthesizerTests.swift +++ b/Tests/FluidAudioTests/TTS/KokoroAne/KokoroAneSynthesizerTests.swift @@ -18,6 +18,40 @@ final class KokoroAneStageBundleNameTests: XCTestCase { } } +/// Joining per-chunk results for long text input (#712, #940). +final class KokoroAneResultConcatenationTests: XCTestCase { + + private func part( + samples: [Float], ids: [Int32], durations: [Int32], albertMs: Double + ) + -> KokoroAneSynthesisResult + { + var timings = KokoroAneStageTimings() + timings.albert = albertMs + return KokoroAneSynthesisResult( + samples: samples, sampleRate: 24_000, encoderTokens: ids.count, + acousticFrames: Int(durations.reduce(0, +)), timings: timings, + inputIds: ids, predictedDurations: durations, normalizedText: "chunk", phonemes: "x") + } + + func testConcatenatesInOrderAndSumsCounts() { + let a = part(samples: [0.1, 0.2], ids: [0, 5, 0], durations: [1, 2, 1], albertMs: 3) + let b = part(samples: [0.3], ids: [0, 7, 8, 0], durations: [1, 1, 1, 1], albertMs: 4) + let joined = KokoroAneSynthesisResult.concatenating([a, b]) + + XCTAssertEqual(joined.samples, [0.1, 0.2, 0.3]) + XCTAssertEqual(joined.sampleRate, 24_000) + XCTAssertEqual(joined.inputIds, [0, 5, 0, 0, 7, 8, 0]) + XCTAssertEqual(joined.predictedDurations, [1, 2, 1, 1, 1, 1, 1]) + XCTAssertEqual(joined.inputIds.count, joined.predictedDurations.count) + XCTAssertEqual(joined.encoderTokens, 7) + XCTAssertEqual(joined.acousticFrames, 8) + XCTAssertEqual(joined.timings.albert, 7) + XCTAssertNil(joined.normalizedText) + XCTAssertEqual(joined.phonemes, "") + } +} + /// Lightweight tests for the pure duration-rounding helper (no models needed). final class KokoroAnePredictedDurationTests: XCTestCase { diff --git a/Tests/FluidAudioTests/TTS/Shared/PhonemeChunkerTests.swift b/Tests/FluidAudioTests/TTS/Shared/PhonemeChunkerTests.swift index 10acf15e1..b9904ec5e 100644 --- a/Tests/FluidAudioTests/TTS/Shared/PhonemeChunkerTests.swift +++ b/Tests/FluidAudioTests/TTS/Shared/PhonemeChunkerTests.swift @@ -80,4 +80,20 @@ final class PhonemeChunkerTests: XCTestCase { let roundTrip = chunks.joined().replacingOccurrences(of: " ", with: "") XCTAssertEqual(roundTrip, original) } + + // MARK: - Unicode-scalar counting + + func testScalarCountingKeepsCombiningMarksUnderTheCap() { + // "ɑ̃" is one Character but two scalars (U+0251 U+0303): 4 words of + // 6 scalars + spaces = 27 scalars, 15 Characters. + let word = String(repeating: "ɑ̃", count: 3) + let text = [word, word, word, word].joined(separator: " ") + XCTAssertEqual(PhonemeChunker.chunk(text, maxLength: 15), [text]) + + let chunks = PhonemeChunker.chunk(text, maxLength: 14, countsUnicodeScalars: true) + XCTAssertEqual(chunks, ["\(word) \(word)", "\(word) \(word)"]) + for piece in chunks { + XCTAssertLessThanOrEqual(piece.unicodeScalars.count, 14) + } + } }