Skip to content

Commit b68f484

Browse files
Alex-WenggAlex-Wengg
authored
feat(tts/chatterbox-nano): extended ~30 s output capacity + honest budget errors (#925)
Closes #924 ## Problem As #924 documents (and this PR's investigation confirmed against the shipped safetensors), the raw caps hide the voice's own footprint: | | Raw shape | Voice footprint | Usable budget | |---|---|---|---| | Input text | 512-token prefill | 376 cond rows + BOS | **≤135 BPE tokens** (~500–550 chars) | | Generated audio | 500-token flow bucket | 250 prompt + 3 silence | **≤247 tokens ≈ 9.9 s** | The output cap binds first in practice. The buckets are conversion-time static-shape choices, not model limits — the checkpoint's `wpe` is `[8196, 768]` and the decode model's `M1536` KV cache already supports ~1020 generated tokens, so only the S3Gen pair needed a larger export. ## Changes - **`ChatterboxNanoOutputCapacity`** — `.standard` (N500/T1000, ≈9.9 s) or `.extended` (N1000/T2000, ≈29.9 s), chosen at `ChatterboxNanoManager(outputCapacity:)`. The extended pair is a separate ~280 MB download from [FluidInference/chatterbox-nano-coreml](https://huggingface.co/FluidInference/chatterbox-nano-coreml) (exported in FluidInference/mobius#93, fp16 parity in the N500 class) and roughly doubles flow/vocoder latency per call, so it stays opt-in. T3 prefill/decode are shared. - **Fail fast + honest errors** — decode throws the moment the generation budget is exhausted instead of decoding to EOS first; `textTooLong`/`generationTooLong` now report the *usable* budget (text tokens vs prefill-minus-conditioning; generated tokens vs bucket-minus-prompt, with seconds) for both Nano and Multilingual. - **Docs** — effective-budget table in `Documentation/TTS/Chatterbox.md` + doc comments on the constants; CLI `--extended-output`. - **Tests** — budget math, per-capacity model sets, repo variant mapping. ## Verification (M5 Pro, release build) - 95-word text, `--extended-output`: **25.6 s WAV**, verbatim Parakeet round-trip - Same text, `.standard`: fails fast with `generationTooLong(tokens: 248, max: 247)` (previously decoded to EOS before throwing) - Short text, `.standard`: unchanged (3.0 s, same seed)
1 parent 69e42da commit b68f484

10 files changed

Lines changed: 240 additions & 33 deletions

File tree

‎Documentation/TTS/Chatterbox.md‎

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -66,9 +66,19 @@ takes (the T3 stage samples stochastically).
6666

6767
- One-shot synthesis (no streaming): the AR decode + flow + vocoder finish
6868
before audio is available.
69-
- Generation is capped by the 500-token flow bucket ≈ **10 s of audio per
70-
call** after the built-in voice's prompt tokens; split long text into
71-
sentences.
69+
- Per-call budgets are the static model shapes *minus the voice's own
70+
footprint* (#924). With the built-in voice:
71+
72+
| | Raw shape | Voice footprint | Usable budget |
73+
|---|---|---|---|
74+
| Input text | 512-token prefill | 376 cond rows + BOS | **≤135 BPE tokens** (~500–550 chars) |
75+
| Generated audio | 500-token flow bucket | 250 prompt + 3 silence | **≤247 tokens ≈ 9.9 s** |
76+
77+
The output cap binds first in practice. Nano can trade download size for
78+
headroom: `ChatterboxNanoManager(outputCapacity: .extended)` (CLI
79+
`--extended-output`) loads an `N1000`/`T2000` S3Gen pair — **≈29.9 s per
80+
call**, extra ~280 MB download, roughly double the flow/vocoder latency.
81+
Otherwise split long text into sentences.
7282
- Never force `.cpuOnly` — the Multilingual T3 packages hard-crash there
7383
(Nano untested; both load `.cpuAndGPU`).
7484
- **Benchmark with `-c release`.** Debug builds spend ~12 ms/token in

‎Sources/FluidAudio/ModelNames.swift‎

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1578,6 +1578,10 @@ public enum ModelNames {
15781578
public static let decodeFile = "T3Nano-Decode-M1536-fp16-stateful.mlmodelc"
15791579
public static let flowFile = "FlowMean-N500-fp16.mlmodelc"
15801580
public static let vocoderFile = "HiFT-T1000-fp16.mlmodelc"
1581+
/// Larger S3Gen bucket pair (`ChatterboxNanoOutputCapacity.extended`,
1582+
/// ~30 s of generated audio) — downloaded only when requested.
1583+
public static let flowFileExtended = "FlowMean-N1000-fp16.mlmodelc"
1584+
public static let vocoderFileExtended = "HiFT-T2000-fp16.mlmodelc"
15811585
public static let tablesFile = "tables/tables.safetensors"
15821586
public static let defaultVoiceFile = "tables/voice-default.safetensors"
15831587
public static let vocabFile = "tokenizer/vocab.json"
@@ -1590,6 +1594,16 @@ public enum ModelNames {
15901594
flowFile,
15911595
vocoderFile,
15921596
]
1597+
/// Required model set for an output capacity ("extended" swaps in
1598+
/// the N1000/T2000 S3Gen pair).
1599+
public static func requiredModels(capacity: ChatterboxNanoOutputCapacity) -> Set<String> {
1600+
switch capacity {
1601+
case .standard:
1602+
return requiredModels
1603+
case .extended:
1604+
return [prefillFile, decodeFile, flowFileExtended, vocoderFileExtended]
1605+
}
1606+
}
15931607
/// Non-model assets fetched individually (nested under `tables/` and
15941608
/// `tokenizer/`, which the repo-root model walk does not descend into).
15951609
public static let auxFiles: [String] = [
@@ -1703,7 +1717,9 @@ public enum ModelNames {
17031717
case .chatterbox:
17041718
return ModelNames.Chatterbox.requiredModels
17051719
case .chatterboxNano:
1706-
return ModelNames.ChatterboxNano.requiredModels
1720+
// Variant: "extended" → N1000/T2000 S3Gen bucket pair (~30 s).
1721+
let capacity = ChatterboxNanoOutputCapacity(rawValue: variant ?? "") ?? .standard
1722+
return ModelNames.ChatterboxNano.requiredModels(capacity: capacity)
17071723
case .luxtts:
17081724
// Variants: "gpu" (macOS) / "ane" (iOS); nil → platform default.
17091725
return ModelNames.LuxTts.requiredFiles(variant: variant)

‎Sources/FluidAudio/TTS/Chatterbox/ChatterboxError.swift‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,9 +24,15 @@ public enum ChatterboxError: Error, LocalizedError {
2424
+ "(needs language-specific text preprocessing). Supported: "
2525
+ ChatterboxConstants.supportedLanguages.sorted().joined(separator: ", ")
2626
case .textTooLong(let tokens, let max):
27-
return "Text tokenizes to \(tokens) tokens; the prefill window holds \(max)"
27+
return
28+
"Text tokenizes to \(tokens) BPE tokens but the usable budget is \(max) "
29+
+ "(the prefill window minus the voice's conditioning); split the text"
2830
case .generationTooLong(let tokens, let max):
29-
return "Generated \(tokens) speech tokens; the flow bucket holds \(max)"
31+
return
32+
"Generated \(tokens) speech tokens but the usable budget is \(max) "
33+
+ "≈ \(String(format: "%.1f", Double(max) / 25.0)) s of audio "
34+
+ "(the flow bucket minus the voice's prompt tokens); split the text "
35+
+ "or load ChatterboxNanoOutputCapacity.extended (Nano only)"
3036
case .processingFailed(let detail):
3137
return "Chatterbox synthesis failed: \(detail)"
3238
}

‎Sources/FluidAudio/TTS/Chatterbox/Nano/ChatterboxNanoConstants.swift‎

Lines changed: 49 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,11 @@ public enum ChatterboxNanoConstants {
1818
public static let samplesPerMelFrame = 480
1919

2020
// ---- T3 (token generator) ----
21-
/// Static prefill window baked into the prefill model.
21+
/// Static prefill window baked into the prefill model. The window holds
22+
/// `[voice conditioning, text BPE tokens, BOS]`, and the built-in voice's
23+
/// conditioning is 376 rows — so the *usable text budget* is
24+
/// 512 − 376 − 1 = **135 BPE tokens** (roughly 500–550 characters), not
25+
/// 512 (#924).
2226
public static let prefillLength = 512
2327
/// KV-cache capacity baked into the decode model.
2428
public static let maxContext = 1536
@@ -41,9 +45,15 @@ public enum ChatterboxNanoConstants {
4145
/// Upstream appends three silence tokens before vocoding (`S3GEN_SIL`).
4246
public static let silenceToken = 4299
4347
public static let silenceTokenCount = 3
44-
/// Flow token bucket (prompt + generated) baked into `FlowMean-N500`.
48+
/// Flow token bucket (prompt + generated) of the `.standard` capacity
49+
/// (`FlowMean-N500`). The built-in voice's 250 prompt tokens and the 3
50+
/// appended silence tokens live inside this bucket, so the *usable
51+
/// generation budget* is 500 − 250 − 3 = **247 speech tokens ≈ 9.9 s of
52+
/// audio** (#924). Use `ChatterboxNanoOutputCapacity.extended` for ~3×
53+
/// that. Prefer `ChatterboxNanoOutputCapacity.flowTokenBucket`.
4554
public static let flowTokenBucket = 500
46-
/// Mel frames produced by the flow bucket (2 per token) = HiFT bucket.
55+
/// Mel frames produced by the `.standard` flow bucket (2 per token) =
56+
/// HiFT bucket. Prefer `ChatterboxNanoOutputCapacity.melFrameBucket`.
4757
public static let melFrameBucket = 1000
4858
/// Harmonic channels in the HiFT source module (harmonics + fundamental).
4959
public static let hiftHarmonics = 9
@@ -57,3 +67,39 @@ public enum ChatterboxNanoConstants {
5767

5868
public static let defaultVoice = "default"
5969
}
70+
71+
/// Which S3Gen flow/vocoder bucket pair to download and load. The flow
72+
/// bucket holds `voice prompt tokens + generated speech tokens + 3 silence
73+
/// tokens`, so the audio each capacity can generate depends on the voice:
74+
/// with the built-in voice (250 prompt tokens) `.standard` yields ≤247
75+
/// generated tokens ≈ 9.9 s per call and `.extended` ≤747 ≈ 29.9 s.
76+
///
77+
/// `.extended` is a separate ~270 MB download (`FlowMean-N1000` +
78+
/// `HiFT-T2000`) and roughly doubles the flow/vocoder latency per call —
79+
/// the buckets are static shapes, so short outputs pay the full bucket.
80+
public enum ChatterboxNanoOutputCapacity: String, CaseIterable, Sendable {
81+
/// `FlowMean-N500` + `HiFT-T1000` — ≈9.9 s of generated audio with the
82+
/// built-in voice.
83+
case standard
84+
/// `FlowMean-N1000` + `HiFT-T2000` — ≈29.9 s of generated audio with
85+
/// the built-in voice.
86+
case extended
87+
88+
/// Flow token bucket (prompt + generated + silence) baked into the
89+
/// capacity's `FlowMean` model.
90+
public var flowTokenBucket: Int {
91+
switch self {
92+
case .standard: return 500
93+
case .extended: return 1000
94+
}
95+
}
96+
97+
/// Mel frames produced by the flow bucket (2 per token) = HiFT bucket.
98+
public var melFrameBucket: Int { 2 * flowTokenBucket }
99+
100+
/// Speech tokens available for generation once `promptTokens` (the
101+
/// loaded voice's prompt) and the appended silence tokens are counted.
102+
public func generationBudget(promptTokens: Int) -> Int {
103+
max(0, flowTokenBucket - promptTokens - ChatterboxNanoConstants.silenceTokenCount)
104+
}
105+
}

‎Sources/FluidAudio/TTS/Chatterbox/Nano/ChatterboxNanoManager.swift‎

Lines changed: 17 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,10 +6,15 @@ import Foundation
66
/// Requires macOS 15 / iOS 18: the T3 decode step keeps its KV cache in
77
/// CoreML `MLState` buffers.
88
///
9+
/// Per-call budgets are set by the static model shapes *minus the voice's
10+
/// own footprint* (#924): with the built-in voice, text is capped at 135
11+
/// BPE tokens (~500–550 characters) and generated audio at ≈9.9 s
12+
/// (`.standard`) or ≈29.9 s (`.extended`, an extra ~280 MB download).
13+
///
914
/// - Note: Beta — this is a beta model conversion; API, model artifacts, and accuracy may change.
1015
///
1116
/// ```swift
12-
/// let manager = ChatterboxNanoManager()
17+
/// let manager = ChatterboxNanoManager() // or (outputCapacity: .extended)
1318
/// try await manager.initialize()
1419
/// let audio = try await manager.synthesize(
1520
/// text: "Well that went better than expected [chuckle], see you tomorrow.")
@@ -25,14 +30,22 @@ public actor ChatterboxNanoManager {
2530
}
2631

2732
private var models: ChatterboxNanoModels?
33+
private let outputCapacity: ChatterboxNanoOutputCapacity
2834

29-
public init() {}
35+
/// - Parameter outputCapacity: which S3Gen bucket pair to load —
36+
/// `.standard` (≈9.9 s of generated audio per call with the built-in
37+
/// voice) or `.extended` (≈29.9 s; separate ~280 MB download, roughly
38+
/// double the flow/vocoder latency per call).
39+
public init(outputCapacity: ChatterboxNanoOutputCapacity = .standard) {
40+
self.outputCapacity = outputCapacity
41+
}
3042

3143
/// Download (if needed) and load the four CoreML models + tables + tokenizer.
3244
public func initialize(progressHandler: ProgressHandler? = nil) async throws {
3345
guard models == nil else { return }
34-
models = try await ChatterboxNanoModels.load(progressHandler: progressHandler)
35-
Self.logger.info("Chatterbox Nano models ready")
46+
models = try await ChatterboxNanoModels.load(
47+
capacity: outputCapacity, progressHandler: progressHandler)
48+
Self.logger.info("Chatterbox Nano models ready (\(self.outputCapacity.rawValue) capacity)")
3649
}
3750

3851
/// Synthesize English `text` with the built-in voice. Paralinguistic

‎Sources/FluidAudio/TTS/Chatterbox/Nano/ChatterboxNanoModels.swift‎

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,20 +15,22 @@ struct ChatterboxNanoModels: Sendable {
1515
let decode: MLModel
1616
let flow: MLModel
1717
let vocoder: MLModel
18+
let capacity: ChatterboxNanoOutputCapacity
1819
let tokenizer: ChatterboxNanoTokenizer
1920
let tables: ChatterboxTables.Nano
2021
let voice: ChatterboxTables.Voice
2122
let repoDir: URL
2223

2324
static func load(
2425
directory: URL? = nil,
26+
capacity: ChatterboxNanoOutputCapacity = .standard,
2527
progressHandler: ProgressHandler? = nil
2628
) async throws -> ChatterboxNanoModels {
2729
let modelsRoot = try directory ?? defaultCacheRoot()
2830
let repoDir = modelsRoot.appendingPathComponent(Repo.chatterboxNano.folderName)
2931

3032
let requiredPaths =
31-
ModelNames.ChatterboxNano.requiredModels.map { $0 }
33+
ModelNames.ChatterboxNano.requiredModels(capacity: capacity).map { $0 }
3234
+ ModelNames.ChatterboxNano.auxFiles
3335
let allPresent = requiredPaths.allSatisfy {
3436
FileManager.default.fileExists(atPath: repoDir.appendingPathComponent($0).path)
@@ -37,6 +39,7 @@ struct ChatterboxNanoModels: Sendable {
3739
logger.info("Downloading Chatterbox Nano CoreML assets from HuggingFace…")
3840
try await ModelHub.download(
3941
.chatterboxNano, to: modelsRoot,
42+
variant: capacity == .standard ? nil : capacity.rawValue,
4043
progressHandler: progressHandler)
4144
// The repo walk only descends into the required .mlmodelc bundles;
4245
// the tables + tokenizer assets live in subdirectories and are
@@ -61,11 +64,17 @@ struct ChatterboxNanoModels: Sendable {
6164
let decode = try await MLModel.load(
6265
contentsOf: repoDir.appendingPathComponent(ModelNames.ChatterboxNano.decodeFile),
6366
configuration: makeConfig(.cpuAndGPU))
67+
let flowFile =
68+
capacity == .standard
69+
? ModelNames.ChatterboxNano.flowFile : ModelNames.ChatterboxNano.flowFileExtended
70+
let vocoderFile =
71+
capacity == .standard
72+
? ModelNames.ChatterboxNano.vocoderFile : ModelNames.ChatterboxNano.vocoderFileExtended
6473
let flow = try await MLModel.load(
65-
contentsOf: repoDir.appendingPathComponent(ModelNames.ChatterboxNano.flowFile),
74+
contentsOf: repoDir.appendingPathComponent(flowFile),
6675
configuration: makeConfig(.cpuAndGPU))
6776
let vocoder = try await MLModel.load(
68-
contentsOf: repoDir.appendingPathComponent(ModelNames.ChatterboxNano.vocoderFile),
77+
contentsOf: repoDir.appendingPathComponent(vocoderFile),
6978
configuration: makeConfig(.cpuAndGPU))
7079

7180
func loadAux() throws -> (ChatterboxNanoTokenizer, ChatterboxTables.Nano, ChatterboxTables.Voice) {
@@ -99,6 +108,7 @@ struct ChatterboxNanoModels: Sendable {
99108
decode: decode,
100109
flow: flow,
101110
vocoder: vocoder,
111+
capacity: capacity,
102112
tokenizer: tokenizer,
103113
tables: tables,
104114
voice: voice,

‎Sources/FluidAudio/TTS/Chatterbox/Nano/ChatterboxNanoSynthesizer.swift‎

Lines changed: 19 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -44,9 +44,12 @@ struct ChatterboxNanoSynthesizer {
4444
let textIds = models.tokenizer.encode(normalized)
4545
let condLen = models.voice.condEmb.rows
4646
let contextLen = condLen + textIds.count + 1 // single BOS embed
47+
// Report the budget the caller can actually influence: text tokens
48+
// vs. what remains of the prefill window after the voice
49+
// conditioning and BOS (#924).
50+
let textBudget = ChatterboxNanoConstants.prefillLength - condLen - 1
4751
guard contextLen <= ChatterboxNanoConstants.prefillLength else {
48-
throw ChatterboxError.textTooLong(
49-
tokens: contextLen, max: ChatterboxNanoConstants.prefillLength)
52+
throw ChatterboxError.textTooLong(tokens: textIds.count, max: textBudget)
5053
}
5154

5255
let prefillEmbeds = try buildPrefillEmbeds(textIds: textIds)
@@ -91,6 +94,11 @@ struct ChatterboxNanoSynthesizer {
9194
let maxSteps = min(
9295
ChatterboxNanoConstants.maxNewTokens,
9396
ChatterboxNanoConstants.maxContext - contextLen - 1)
97+
// Speech tokens the loaded flow bucket can hold beyond the voice's
98+
// prompt and the appended silence — fail as soon as it's exhausted
99+
// rather than decoding to EOS first (#924).
100+
let generationBudget = models.capacity.generationBudget(
101+
promptTokens: models.voice.promptTokens.count)
94102

95103
for step in 0..<maxSteps {
96104
// Upstream's first sample penalizes the BOS id (its input_ids
@@ -106,6 +114,10 @@ struct ChatterboxNanoSynthesizer {
106114
generatedIds.append(token)
107115
if token == ChatterboxNanoConstants.stopSpeechToken { break }
108116
if token < ChatterboxNanoConstants.speechVocabSize { speechTokens.append(token) }
117+
if speechTokens.count > generationBudget {
118+
throw ChatterboxError.generationTooLong(
119+
tokens: speechTokens.count, max: generationBudget)
120+
}
109121

110122
try fillStepEmbeds(stepEmbeds, token: token)
111123
curLenArr[0] = NSNumber(value: contextLen + step)
@@ -126,17 +138,12 @@ struct ChatterboxNanoSynthesizer {
126138
guard !speechTokens.isEmpty else {
127139
throw ChatterboxError.processingFailed("no speech tokens generated")
128140
}
129-
// Upstream appends three silence tokens before vocoding.
141+
// Upstream appends three silence tokens before vocoding. The
142+
// in-loop budget check already reserved room for them.
130143
speechTokens.append(
131144
contentsOf: [Int](
132145
repeating: ChatterboxNanoConstants.silenceToken,
133146
count: ChatterboxNanoConstants.silenceTokenCount))
134-
let promptLen = models.voice.promptTokens.count
135-
let totalTokens = promptLen + speechTokens.count
136-
guard totalTokens <= ChatterboxNanoConstants.flowTokenBucket else {
137-
throw ChatterboxError.generationTooLong(
138-
tokens: totalTokens, max: ChatterboxNanoConstants.flowTokenBucket)
139-
}
140147

141148
// ---- S3Gen: meanflow (mel) + HiFT (waveform) ----
142149
let flowStart = Date()
@@ -285,8 +292,8 @@ struct ChatterboxNanoSynthesizer {
285292
speechTokens: [Int], rng: inout SplitMix64,
286293
isolation: isolated (any Actor)? = #isolation
287294
) async throws -> MLMultiArray {
288-
let bucket = ChatterboxNanoConstants.flowTokenBucket
289-
let melBucket = ChatterboxNanoConstants.melFrameBucket
295+
let bucket = models.capacity.flowTokenBucket
296+
let melBucket = models.capacity.melFrameBucket
290297
let voice = models.voice
291298
let promptLen = voice.promptTokens.count
292299
let totalLen = promptLen + speechTokens.count
@@ -347,7 +354,7 @@ struct ChatterboxNanoSynthesizer {
347354
mel: MLMultiArray, melFrames: Int, rng: inout SplitMix64,
348355
isolation: isolated (any Actor)? = #isolation
349356
) async throws -> [Float] {
350-
let melBucket = ChatterboxNanoConstants.melFrameBucket
357+
let melBucket = models.capacity.melFrameBucket
351358
let promptFrames = 2 * models.voice.promptTokens.count
352359
let melValues = try ChatterboxMLSupport.floatBuffer(mel) // [80 * melBucket]
353360

‎Sources/FluidAudio/TTS/Chatterbox/Pipeline/ChatterboxSynthesizer.swift‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -52,9 +52,12 @@ struct ChatterboxSynthesizer {
5252
textIds.append(ChatterboxConstants.stopTextToken)
5353
let condLen = models.voice.condEmb.rows
5454
let contextLen = condLen + textIds.count + 2 // two BOS embeds
55+
// Report text tokens vs. what remains of the prefill window after
56+
// the voice conditioning and BOS embeds (#924).
5557
guard contextLen <= ChatterboxConstants.prefillLength else {
5658
throw ChatterboxError.textTooLong(
57-
tokens: contextLen, max: ChatterboxConstants.prefillLength)
59+
tokens: textIds.count,
60+
max: ChatterboxConstants.prefillLength - condLen - 2)
5861
}
5962

6063
let prefillEmbeds = try buildPrefillEmbeds(textIds: textIds)
@@ -144,9 +147,12 @@ struct ChatterboxSynthesizer {
144147
}
145148
let promptLen = models.voice.promptTokens.count
146149
let totalTokens = promptLen + speechTokens.count
150+
// Report generated tokens vs. what remains of the flow bucket after
151+
// the voice's prompt tokens (#924).
147152
guard totalTokens <= ChatterboxConstants.flowTokenBucket else {
148153
throw ChatterboxError.generationTooLong(
149-
tokens: totalTokens, max: ChatterboxConstants.flowTokenBucket)
154+
tokens: speechTokens.count,
155+
max: ChatterboxConstants.flowTokenBucket - promptLen)
150156
}
151157

152158
// ---- S3Gen: flow (mel) + HiFT (waveform) ----

0 commit comments

Comments
 (0)