Repository navigation
feat(models): Kokoro-82M TTS with iSTFTNet vocoder for /v1/audio/speech - #374
Merged
Merged
Conversation
Kokoro and other phoneme-input TTS models consume IPA phonemes, not raw text, so this adds a self-contained American-English g2p front-end (the host-side analogue of an STT mel front-end). Pipeline: normalize text (lower-case, speak integers, keep common punctuation), look each word up in a bundled lexicon, and fall back to deterministic letter-to-sound rules for out-of-vocabulary words. Output uses only the IPA symbols in Kokoro's vocab, so no symbol is silently dropped at tokenization. No external binary or download is required. Modules: normalize (tokenizer + number expansion), lexicon (compiled-in word to IPA table with a symbol-audit test), rules (digraph and single-letter letter-to-sound with stress placement), and the orchestrator text_to_phonemes. Language scope is American English; per-language packs are future work.
Port the Kokoro-82M text-to-speech model to MLX/Rust. Kokoro is a StyleTTS2 phoneme-to-mel acoustic model with a built-in iSTFTNet vocoder, so it produces a 24 kHz mono waveform directly with no separate neural codec.
Forward pass: phoneme ids to a PLBert text encoder (a custom ALBERT with one weight-shared layer reused 12 times, gelu_new activation), a duration predictor (DurationEncoder bi-LSTMs + AdaLayerNorm, a duration head, and an alignment matrix that expands per-token features to per-frame), F0 and N prosody (a shared bi-LSTM plus AdainResBlk1d stacks), an acoustic text encoder, and the iSTFTNet decoder (encode/decode AdainResBlk1d blocks, transposed-conv upsampling, an NSF harmonic source, Snake-activated MRF residual blocks, conv_post, and an inverse STFT with magnitude = exp and phase = sin).
Implementation notes: weight-norm conv kernels are reconstructed from weight_g/weight_v and transposed from PyTorch layout to MLX channels-last; grouped (depthwise) transposed convs use a distinct weight transpose. MLX has no recurrent primitive, so bidirectional LSTMs are scanned by hand. The inverse STFT uses an explicit real inverse-DFT with Hermitian bin doubling and the standard 1/N normalization, matching the PyTorch reference (mag * exp(j*phase), periodic Hann, no phase unwrap), which avoids depending on MLX complex64 byte layout. The NSF source is deterministic (zero initial phase, no additive noise) so output is reproducible.
Registry: add ModelType::Kokoro with metadata ("Kokoro (StyleTTS2 + iSTFTNet)", "Text-to-speech"), register it in ALL_MODEL_TYPES and the static descriptor and tensor-parallel fallback tables (served via the audio endpoints, never text generation), and add the Text-to-speech family to the `mlxcel arch` ordering.
Tests: STFT round-trip at unit scale, alignment-matrix construction, voice-pack loading and indexing, detection by the istftnet block, and a checkpoint-gated end-to-end synthesis that checks 24 kHz output, non-trivial duration, finite bounded samples, and that speed scales duration. The end-to-end tests load the real prince-canuma/Kokoro-82M checkpoint when present (KOKORO_MODEL_DIR or models/kokoro-82m) and skip otherwise.
Also rewrite mean_prompt_len with checked_div to satisfy the workspace clippy gate.
Connect the Kokoro model to the existing audio seam so `POST /v1/audio/speech` synthesizes real audio. KokoroTtsProvider mirrors the Whisper STT provider: it owns an AudioWorker that loads the checkpoint and runs every synthesis on one dedicated, stream-initialized thread (MLX work is thread-affine), and the provider itself holds no MLX handles. The engine runs the g2p front-end (text to phonemes) then the acoustic model and built-in vocoder, returning f32 PCM that the route encodes to WAV via the existing #364 writer. The voice request field selects a voices/ pack (validated, default af_heart), speed scales duration, and response_format honors wav. Detection: Kokoro's config.json has no top-level model_type, so get_model_type checks is_kokoro_checkpoint (the istftnet config block or the kokoro-v1_0.safetensors filename) before the model_type dispatch that would otherwise error. Startup loads KokoroTtsProvider when the checkpoint resolves to ModelType::Kokoro. No existing model's detection is affected. Docs: add a Text-to-speech section to supported-models.md describing the flow, detection, request fields, the g2p approach and its English-only scope, and the single-model serving constraint. Closes #366
inureyes
added
type:enhancement
The decoder handed the NSF harmonic source the conv-downsampled F0 curve (f0_a, length T) instead of the raw predictor F0 (f0_pred, length 2T). The generator upsamples its F0 input by 300, so the downsampled curve produced a harmonic source at T*300 samples, half the intended 2T*300, which after the align_time trim left the synthesized waveform at roughly half its correct duration (and pitched off the smoothed, strided curve rather than the raw pitch). Both the upstream MLX reference (istftnet.py Decoder.__call__ passes the raw F0_pred to generator, computing the downsampled F0 only for the encode/decode concat) and the NumPy oracle pass the full-resolution F0_pred. The conv-downsampled f0_a stays correct for the feature concatenation; only the generator's source curve changes to f0_pred. Verified against the oracle: for the test phonemes the fixed path emits exactly 31800 samples (1.325 s, = 2T*300 with 2T=106), matching ref_audio.npy and doubling the prior 15900. cargo check --lib --tests, clippy --lib --tests -D warnings, fmt, and test --lib models::kokoro (7/7, real-checkpoint e2e included) are clean.
Member
Author
Implementation Review SummaryIntent
Findings Addressed
Remaining Items
Verification
Robustness (spot-checked, all sound)
|
Two hardening fixes on the attacker-reachable /v1/audio/speech path. DoS: the TTS route inherits the 25 MiB audio body limit (sized for STT file uploads), so request.input can be ~25 MB of free text. KokoroEngine::synthesize ran the full g2p front-end (normalize plus per-word lexicon/rules) over the entire string before the model truncates to 510 tokens, so the 510-token cap protected only the MLX stage. Because synthesis runs on the single AudioWorker thread, one oversized request monopolized it and blocked all audio requests. Cap the text to MAX_INPUT_CHARS (4096) on a char boundary before any g2p work; truncate rather than reject so well-behaved long inputs still synthesize. The cap is unit-tested independently of the model (short passthrough, exact-cap, overlong, and multibyte boundary cases). Availability: the worker loop called engine.synthesize/transcribe with no panic boundary. The kokoro forward path carries many expect/unwrap calls reachable from request input; a panic there unwound out of the loop, closed the channel, and made every later audio request fail with WORKER_GONE for the process lifetime. Wrap each engine call in catch_unwind (AssertUnwindSafe over the single-thread-owned engine), still synchronize the stream on both paths, and convert a caught panic into a per-request Inference error so the thread keeps serving. A regression test proves the worker survives a panicking request and serves the next one. Note: this backstop only takes effect when the binary unwinds; release builds set panic = "abort" in Cargo.toml, so a production build still aborts on panic. Removing that global profile setting is out of scope here and is called out in the PR review for the maintainers to decide.
Member
Author
Security & performance review (PR #374, Kokoro TTS)Reviewed the new code on the trust boundary FixedHIGH, DoS: unbounded g2p work before the 510-token cap. The TTS route inherits the 25 MiB audio body limit ( HIGH, availability: a synthesis panic kills the audio worker. The worker loop called Needs a maintainer decision (not changed here)
Verified as already-safe (no change needed)
VerificationToolchain 1.93.1, narrow selectors: |
3 tasks done
Add narrow, deterministic tests that were absent:
- g2p/mod.rs: possessive clitic strip-and-retry path ("cat's" -> kætz), whitespace-only input
- kokoro/mod.rs: weight-filename detection branch (no istftnet block, file present)
- kokoro/voices.rs: resolve_voice fallback to DEFAULT_VOICE (empty dir, unknown name, unsafe name, None); resolve_voice picks known voice from tmpdir
Apply cargo fmt to audio_worker.rs (two run_guarded call sites reformatted by the formatter; no logic change).
Update docs/audio-api.md: remove the stale "TTS always returns 501" statement in the Current Status paragraph, add kokoro_tts.rs to the source map table, and add a Kokoro TTS setup section covering checkpoint detection, voices, language scope, input length cap, and current limitations.
Member
Author
PR Finalization CompleteTests addedThree new test groups covering previously untested code paths (all deterministic, no checkpoint required):
Running total: 19 g2p tests, 7 kokoro tests (non-checkpoint), 3 voice tests, 4 kokoro_tts tests, 4 audio_worker tests. Documentation updateddocs/audio-api.md:
docs/supported-models.md: already accurate from prior commits; no changes needed. Lint / format
Verification
|
inureyes
added a commit
that referenced
this pull request
Jun 20, 2026
Audio I/O on the OpenAI /v1/audio/* surface (Whisper speech-to-text #371, Kokoro-82M text-to-speech #374, endpoint plumbing #368), reasoning_content on non-streaming chat (#359), pre-M5 decode +8-12% via a hardware-gated op cap (#360), Turbo4Asym decode rerouted through dequant-SDPA (#369), batched sampling fixes (#339, #340, #344, #356), and the Turbo KV advisor/docs aligned to the measured sweep with ADR 0002 (#343, #376). README: consolidate the v0.3.0 and v0.3.1 highlights under a single v0.3 entry and add v0.3.2.
2 of 3 tasks
7 tasks done
inureyes
added a commit
that referenced
this pull request
Jun 21, 2026
…ind with explicit core-thread abort (#383) * fix(server): make audio synthesis panic-safe in release via unwind The release profile set `panic = "abort"`, which defeats the `catch_unwind` request-isolation backstop PR #374 added to the audio worker (`run_guarded` in src/server/audio_worker.rs). Under abort a synthesis panic does not unwind, so the boundary never runs and the whole server process aborts instead of the offending request degrading to a per-request error. The same setting silenced the pipeline stage-isolation `catch_unwind` in production. Flip `[profile.release]` to `panic = "unwind"` so those deliberate containment boundaries work in the shipped binary as they already do under `cargo test` (the test harness always unwinds). The unwinding-table size and runtime cost are negligible for an MLX-linked inference server whose decode hot path does not panic. Unwind alone would let a panic on a core generation thread silently kill that thread and leave the server alive but unable to generate, so re-impose fail-fast explicitly where it is wanted: add `run_core_thread_or_abort`, a narrow `catch_unwind` + `process::abort()` wrapper that logs without request content, and wrap the two core generation worker threads (batched and legacy sequential) in src/server/model_worker.rs. The batch scheduler, the DiffusionGemma batch-1 loop, and the disaggregated serving-role loop all run inside those two thread bodies, so the two wrappers cover every core generation path. The audio worker `run_guarded` and the pipeline stage boundaries deliberately contain panics and are not wrapped. There is no global abort panic hook on purpose: a hook runs before unwinding and would defeat those backstops. The default panic hook (prints panic + backtrace) is kept. An MLX C++ FFI exception still becomes `std::terminate` rather than a Rust panic and terminates the process; that residual vector is unchanged here and tracked as issue #382. Adds a happy-path unit test for `run_core_thread_or_abort` (the abort path terminates the test runner and is verified manually), documents that the existing `worker_survives_engine_panic_and_keeps_serving` test now represents release behavior, and records the decision and rejected options in docs/adr/0003. * fix(server): wrap remote pipeline stage thread for fail-fast parity Release `panic = "unwind"` (issue #375) made a panic on the remote pipeline stage service thread (`RemoteStageService::run` in src/distributed/pipeline/remote_service.rs) unwind and silently kill the thread while `serve_remote_pipeline_stage` stays parked on `ctrl_c`, leaving a zombie stage process a supervisor will not restart. Under the previous `panic = "abort"` such a panic aborted the stage cleanly, so the unwind flip was a fail-fast regression for a core inference thread that this PR had not covered. Relocate `run_core_thread_or_abort` from src/server/model_worker.rs to a shared crate-level module (src/worker_failfast.rs) reachable from both the server workers and the distributed pipeline, generalizing it to forward the body's return value so a recoverable `Err` propagates unchanged and only a panic aborts. Wrap the remote stage service thread with it; the `JoinHandle<Result<()>>` type and the `shutdown()` join handling are unchanged on the non-panic path. The local batched and legacy `--no-batch` workers keep their existing wrapping, no global panic hook is added, and the audio worker `run_guarded` containment is untouched. Correct the ADR 0003, the Cargo.toml profile comment, and architecture.md, which claimed the pipeline stage "stays contained" with its own `catch_unwind`. There is no production `catch_unwind` under src/distributed/ (the only one is `#[cfg(test)]`). The audio worker is now described as the only deliberately contained boundary, all core inference threads (batched, legacy, and remote stage) fail-fast explicitly through `run_core_thread_or_abort`, and pipeline faults are handled at the coordinator by `Result` propagation plus stage timeout and health probing. The MLX FFI `std::terminate` residual (#382) and the always-unwind `cargo test` notes are kept. * test(worker): add Err-forwarding test for run_core_thread_or_abort The existing happy-path test exercises the void return case. This adds a test with a Result-returning body that returns Err, confirming the wrapper forwards a recoverable error unchanged and does not abort.
be-free-snowboy
pushed a commit
to freestylejs/snowbrain-mlxcel
that referenced
this pull request
Oct 27, 2026
…ch (lablup#374) * feat(models): add English grapheme-to-phoneme front-end for TTS Kokoro and other phoneme-input TTS models consume IPA phonemes, not raw text, so this adds a self-contained American-English g2p front-end (the host-side analogue of an STT mel front-end). Pipeline: normalize text (lower-case, speak integers, keep common punctuation), look each word up in a bundled lexicon, and fall back to deterministic letter-to-sound rules for out-of-vocabulary words. Output uses only the IPA symbols in Kokoro's vocab, so no symbol is silently dropped at tokenization. No external binary or download is required. Modules: normalize (tokenizer + number expansion), lexicon (compiled-in word to IPA table with a symbol-audit test), rules (digraph and single-letter letter-to-sound with stress placement), and the orchestrator text_to_phonemes. Language scope is American English; per-language packs are future work. * feat(models): port Kokoro-82M StyleTTS2 + iSTFTNet TTS to MLX Port the Kokoro-82M text-to-speech model to MLX/Rust. Kokoro is a StyleTTS2 phoneme-to-mel acoustic model with a built-in iSTFTNet vocoder, so it produces a 24 kHz mono waveform directly with no separate neural codec. Forward pass: phoneme ids to a PLBert text encoder (a custom ALBERT with one weight-shared layer reused 12 times, gelu_new activation), a duration predictor (DurationEncoder bi-LSTMs + AdaLayerNorm, a duration head, and an alignment matrix that expands per-token features to per-frame), F0 and N prosody (a shared bi-LSTM plus AdainResBlk1d stacks), an acoustic text encoder, and the iSTFTNet decoder (encode/decode AdainResBlk1d blocks, transposed-conv upsampling, an NSF harmonic source, Snake-activated MRF residual blocks, conv_post, and an inverse STFT with magnitude = exp and phase = sin). Implementation notes: weight-norm conv kernels are reconstructed from weight_g/weight_v and transposed from PyTorch layout to MLX channels-last; grouped (depthwise) transposed convs use a distinct weight transpose. MLX has no recurrent primitive, so bidirectional LSTMs are scanned by hand. The inverse STFT uses an explicit real inverse-DFT with Hermitian bin doubling and the standard 1/N normalization, matching the PyTorch reference (mag * exp(j*phase), periodic Hann, no phase unwrap), which avoids depending on MLX complex64 byte layout. The NSF source is deterministic (zero initial phase, no additive noise) so output is reproducible. Registry: add ModelType::Kokoro with metadata ("Kokoro (StyleTTS2 + iSTFTNet)", "Text-to-speech"), register it in ALL_MODEL_TYPES and the static descriptor and tensor-parallel fallback tables (served via the audio endpoints, never text generation), and add the Text-to-speech family to the `mlxcel arch` ordering. Tests: STFT round-trip at unit scale, alignment-matrix construction, voice-pack loading and indexing, detection by the istftnet block, and a checkpoint-gated end-to-end synthesis that checks 24 kHz output, non-trivial duration, finite bounded samples, and that speed scales duration. The end-to-end tests load the real prince-canuma/Kokoro-82M checkpoint when present (KOKORO_MODEL_DIR or models/kokoro-82m) and skip otherwise. Also rewrite mean_prompt_len with checked_div to satisfy the workspace clippy gate. * feat(server): wire Kokoro TTS provider to /v1/audio/speech Connect the Kokoro model to the existing audio seam so `POST /v1/audio/speech` synthesizes real audio. KokoroTtsProvider mirrors the Whisper STT provider: it owns an AudioWorker that loads the checkpoint and runs every synthesis on one dedicated, stream-initialized thread (MLX work is thread-affine), and the provider itself holds no MLX handles. The engine runs the g2p front-end (text to phonemes) then the acoustic model and built-in vocoder, returning f32 PCM that the route encodes to WAV via the existing lablup#364 writer. The voice request field selects a voices/ pack (validated, default af_heart), speed scales duration, and response_format honors wav. Detection: Kokoro's config.json has no top-level model_type, so get_model_type checks is_kokoro_checkpoint (the istftnet config block or the kokoro-v1_0.safetensors filename) before the model_type dispatch that would otherwise error. Startup loads KokoroTtsProvider when the checkpoint resolves to ModelType::Kokoro. No existing model's detection is affected. Docs: add a Text-to-speech section to supported-models.md describing the flow, detection, request fields, the g2p approach and its English-only scope, and the single-model serving constraint. Closes lablup#366 * fix(kokoro): feed raw 2T predictor F0 to the iSTFTNet generator The decoder handed the NSF harmonic source the conv-downsampled F0 curve (f0_a, length T) instead of the raw predictor F0 (f0_pred, length 2T). The generator upsamples its F0 input by 300, so the downsampled curve produced a harmonic source at T*300 samples, half the intended 2T*300, which after the align_time trim left the synthesized waveform at roughly half its correct duration (and pitched off the smoothed, strided curve rather than the raw pitch). Both the upstream MLX reference (istftnet.py Decoder.__call__ passes the raw F0_pred to generator, computing the downsampled F0 only for the encode/decode concat) and the NumPy oracle pass the full-resolution F0_pred. The conv-downsampled f0_a stays correct for the feature concatenation; only the generator's source curve changes to f0_pred. Verified against the oracle: for the test phonemes the fixed path emits exactly 31800 samples (1.325 s, = 2T*300 with 2T=106), matching ref_audio.npy and doubling the prior 15900. cargo check --lib --tests, clippy --lib --tests -D warnings, fmt, and test --lib models::kokoro (7/7, real-checkpoint e2e included) are clean. * fix(kokoro): bound TTS input and survive synthesis panics Two hardening fixes on the attacker-reachable /v1/audio/speech path. DoS: the TTS route inherits the 25 MiB audio body limit (sized for STT file uploads), so request.input can be ~25 MB of free text. KokoroEngine::synthesize ran the full g2p front-end (normalize plus per-word lexicon/rules) over the entire string before the model truncates to 510 tokens, so the 510-token cap protected only the MLX stage. Because synthesis runs on the single AudioWorker thread, one oversized request monopolized it and blocked all audio requests. Cap the text to MAX_INPUT_CHARS (4096) on a char boundary before any g2p work; truncate rather than reject so well-behaved long inputs still synthesize. The cap is unit-tested independently of the model (short passthrough, exact-cap, overlong, and multibyte boundary cases). Availability: the worker loop called engine.synthesize/transcribe with no panic boundary. The kokoro forward path carries many expect/unwrap calls reachable from request input; a panic there unwound out of the loop, closed the channel, and made every later audio request fail with WORKER_GONE for the process lifetime. Wrap each engine call in catch_unwind (AssertUnwindSafe over the single-thread-owned engine), still synchronize the stream on both paths, and convert a caught panic into a per-request Inference error so the thread keeps serving. A regression test proves the worker survives a panicking request and serves the next one. Note: this backstop only takes effect when the binary unwinds; release builds set panic = "abort" in Cargo.toml, so a production build still aborts on panic. Removing that global profile setting is out of scope here and is called out in the PR review for the maintainers to decide. * test(kokoro): add missing unit tests and fix fmt; update audio-api.md Add narrow, deterministic tests that were absent: - g2p/mod.rs: possessive clitic strip-and-retry path ("cat's" -> kætz), whitespace-only input - kokoro/mod.rs: weight-filename detection branch (no istftnet block, file present) - kokoro/voices.rs: resolve_voice fallback to DEFAULT_VOICE (empty dir, unknown name, unsafe name, None); resolve_voice picks known voice from tmpdir Apply cargo fmt to audio_worker.rs (two run_guarded call sites reformatted by the formatter; no logic change). Update docs/audio-api.md: remove the stale "TTS always returns 501" statement in the Current Status paragraph, add kokoro_tts.rs to the source map table, and add a Kokoro TTS setup section covering checkpoint detection, voices, language scope, input length cap, and current limitations.
be-free-snowboy
pushed a commit
to freestylejs/snowbrain-mlxcel
that referenced
this pull request
Oct 27, 2026
Audio I/O on the OpenAI /v1/audio/* surface (Whisper speech-to-text lablup#371, Kokoro-82M text-to-speech lablup#374, endpoint plumbing lablup#368), reasoning_content on non-streaming chat (lablup#359), pre-M5 decode +8-12% via a hardware-gated op cap (lablup#360), Turbo4Asym decode rerouted through dequant-SDPA (lablup#369), batched sampling fixes (lablup#339, lablup#340, lablup#344, lablup#356), and the Turbo KV advisor/docs aligned to the measured sweep with ADR 0002 (lablup#343, lablup#376). README: consolidate the v0.3.0 and v0.3.1 highlights under a single v0.3 entry and add v0.3.2.
be-free-snowboy
pushed a commit
to freestylejs/snowbrain-mlxcel
that referenced
this pull request
Oct 27, 2026
…ind with explicit core-thread abort (lablup#383) * fix(server): make audio synthesis panic-safe in release via unwind The release profile set `panic = "abort"`, which defeats the `catch_unwind` request-isolation backstop PR lablup#374 added to the audio worker (`run_guarded` in src/server/audio_worker.rs). Under abort a synthesis panic does not unwind, so the boundary never runs and the whole server process aborts instead of the offending request degrading to a per-request error. The same setting silenced the pipeline stage-isolation `catch_unwind` in production. Flip `[profile.release]` to `panic = "unwind"` so those deliberate containment boundaries work in the shipped binary as they already do under `cargo test` (the test harness always unwinds). The unwinding-table size and runtime cost are negligible for an MLX-linked inference server whose decode hot path does not panic. Unwind alone would let a panic on a core generation thread silently kill that thread and leave the server alive but unable to generate, so re-impose fail-fast explicitly where it is wanted: add `run_core_thread_or_abort`, a narrow `catch_unwind` + `process::abort()` wrapper that logs without request content, and wrap the two core generation worker threads (batched and legacy sequential) in src/server/model_worker.rs. The batch scheduler, the DiffusionGemma batch-1 loop, and the disaggregated serving-role loop all run inside those two thread bodies, so the two wrappers cover every core generation path. The audio worker `run_guarded` and the pipeline stage boundaries deliberately contain panics and are not wrapped. There is no global abort panic hook on purpose: a hook runs before unwinding and would defeat those backstops. The default panic hook (prints panic + backtrace) is kept. An MLX C++ FFI exception still becomes `std::terminate` rather than a Rust panic and terminates the process; that residual vector is unchanged here and tracked as issue lablup#382. Adds a happy-path unit test for `run_core_thread_or_abort` (the abort path terminates the test runner and is verified manually), documents that the existing `worker_survives_engine_panic_and_keeps_serving` test now represents release behavior, and records the decision and rejected options in docs/adr/0003. * fix(server): wrap remote pipeline stage thread for fail-fast parity Release `panic = "unwind"` (issue lablup#375) made a panic on the remote pipeline stage service thread (`RemoteStageService::run` in src/distributed/pipeline/remote_service.rs) unwind and silently kill the thread while `serve_remote_pipeline_stage` stays parked on `ctrl_c`, leaving a zombie stage process a supervisor will not restart. Under the previous `panic = "abort"` such a panic aborted the stage cleanly, so the unwind flip was a fail-fast regression for a core inference thread that this PR had not covered. Relocate `run_core_thread_or_abort` from src/server/model_worker.rs to a shared crate-level module (src/worker_failfast.rs) reachable from both the server workers and the distributed pipeline, generalizing it to forward the body's return value so a recoverable `Err` propagates unchanged and only a panic aborts. Wrap the remote stage service thread with it; the `JoinHandle<Result<()>>` type and the `shutdown()` join handling are unchanged on the non-panic path. The local batched and legacy `--no-batch` workers keep their existing wrapping, no global panic hook is added, and the audio worker `run_guarded` containment is untouched. Correct the ADR 0003, the Cargo.toml profile comment, and architecture.md, which claimed the pipeline stage "stays contained" with its own `catch_unwind`. There is no production `catch_unwind` under src/distributed/ (the only one is `#[cfg(test)]`). The audio worker is now described as the only deliberately contained boundary, all core inference threads (batched, legacy, and remote stage) fail-fast explicitly through `run_core_thread_or_abort`, and pipeline faults are handled at the coordinator by `Result` propagation plus stage timeout and health probing. The MLX FFI `std::terminate` residual (lablup#382) and the always-unwind `cargo test` notes are kept. * test(worker): add Err-forwarding test for run_core_thread_or_abort The existing happy-path test exercises the void return case. This adds a test with a Result-returning body that returns Err, confirming the wrapper forwards a recoverable error unchanged and does not abort.
7 of 8 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Ports the Kokoro-82M text-to-speech model (a StyleTTS2 phoneme-to-mel acoustic model with a built-in iSTFTNet vocoder) to MLX/Rust and wires it to
POST /v1/audio/speech, plus a self-contained English grapheme-to-phoneme front-end since Kokoro consumes IPA phonemes rather than raw text. The audio scaffolding from #364 (route, WAV writer, binary response) and #365 (the dedicated-threadAudioWorkerandAudioModelProviderseam) is reused unchanged; this PR adds thesynthesizedirection.A real synthesis through the full pipeline produces playable audio: for the phonemes of "hello" the model emits a 24 kHz mono waveform of 15900 samples (0.663 s), finite and bounded (rms 0.012, peak 0.083), written through the existing WAV writer with a valid RIFF header.
What changed
g2p front-end (
src/models/g2p/): an American-English phonemizer.normalizetokenizes text and speaks integers;lexiconis a compiled-in word-to-IPA table;rulesis a deterministic letter-to-sound fallback for out-of-vocabulary words;modorchestratestext_to_phonemes. Output uses only symbols in Kokoro's vocab (audited by a test). No external binary or download.Kokoro model (
src/models/kokoro/): the StyleTTS2 + iSTFTNet forward pass split into focused modules:bert(PLBert / weight-shared ALBERT),predictor(duration + F0/N prosody with AdaLayerNorm and AdaIN),text_encoder(acoustic),decoder(iSTFTNet generator with NSF source and Snake MRF blocks),stft(STFT/iSTFT + harmonic source),lstm(hand-rolled bi-LSTM, since MLX has no recurrent op),blocks(shared AdaIN residual blocks),ops/weights(MLX helpers, weight-norm reconstruction, channels-last conv).modis theKokoroModel::load/synthesizeentry plus theis_kokoro_checkpointdetector.Provider (
src/server/kokoro_tts.rs):KokoroTtsProvidermirrors the Whisper STT provider, loading and running all MLX on oneAudioWorkerthread and forwarding requests over its channel.Registry and wiring:
ModelType::Kokoroadded to the enum,ALL_MODEL_TYPES, metadata, the static-descriptor and tensor-parallel fallback tables, and themlxcel archfamily ordering; detection inget_model_typekeys off theistftnetconfig block or thekokoro-v1_0.safetensorsfilename (Kokoro has nomodel_type) before themodel_typedispatch; startup loads the provider for a Kokoro checkpoint.docs/supported-models.mdgains a Text-to-speech section.Design decisions and scope
g2p approach: a bundled lexicon plus letter-to-sound rules, scoped to American English (the default voice
af_heartand the othera*voices). This is the lowest-risk self-contained choice: no mandatory external binary (the repo has no espeak/phonemizer dependency) and it yields vocab-valid IPA for the default voice. Limit: non-English voices still load and synthesize, but their phonemes come from the English front-end, so pronunciation quality is reduced; per-language g2p (the analogue of upstreammisaki[xx]) is future work.iSTFT: the inverse follows the PyTorch reference (
magnitude * exp(j*phase), periodic Hann, no phase unwrap). It is computed with an explicit real inverse-DFT (Hermitian bin doubling,1/Nnormalization) rather than MLX complex64 array construction, which a round-trip test pinned as the correct unit-scale path. The NSF harmonic source is deterministic (zero initial phase, no additive noise) so synthesis is reproducible.response_format: WAV is implemented (via the #364 writer). mp3 is not added; no audio codec is currently a dependency, so it stays a documented follow-up and the route returns 400 for non-wav as before.Validation
cargo check --lib --tests: clean.cargo fmt --check: clean.cargo clippy --lib --tests -- -D warnings: clean (also rewrote a pre-existingmean_prompt_lendivision that the gate flagged).Unit tests (narrow selectors): 17 g2p tests pass (tokenizer, number expansion, lexicon lookup, symbol audit, rules, OOV, punctuation). 7 Kokoro tests pass including the STFT round-trip, alignment-matrix construction, voice-pack loading, detection, a real-checkpoint end-to-end synthesis (24 kHz, non-trivial duration, finite bounded samples), and a speed-scaling check (speed 1.5 yields fewer samples than 0.7). The end-to-end tests load
prince-canuma/Kokoro-82Mwhen present and skip otherwise.Acceptance criteria
All six issue checkboxes are satisfied: end-to-end synthesis through
/v1/audio/speechreturns playable 24 kHz audio;voice,speed, andresponse_format(wav) are honored; detection resolves a Kokoro checkpoint without amodel_type;mlxcel archlists the TTS architecture; unit tests are added and fmt/clippy/check are clean;docs/supported-models.mdis updated.Known gaps and risks
mp3 output is deferred. g2p is English-only. The bundled lexicon is curated for high-frequency words; rarer words use the rule-based fallback, which is approximate. The deterministic NSF makes the waveform slightly quieter than a randomized run but perceptually equivalent; exact perceptual parity with the upstream MLX/PyTorch implementations was not bit-compared (a NumPy reference oracle was used during development to validate the forward-pass structure and the unit-scale iSTFT).
Closes #366