MLX talker backend + SpeechDecoder overflow chunking - #3
Open
dbrkn wants to merge 3 commits into
Open
Conversation
dbrkn
force-pushed
the
berkin/voice-clone-coreml-talker
branch
2 times, most recently
from
July 21, 2026 12:56
ef6208f to
a95a911
Compare
dbrkn
force-pushed
the
berkin/voice-clone-coreml
branch
from
July 21, 2026 12:56
f75887d to
99611a3
Compare
dbrkn
force-pushed
the
berkin/voice-clone-coreml-talker
branch
from
July 21, 2026 13:40
a95a911 to
730e3ef
Compare
Generations longer than the SpeechDecoder cache window (~21s for kv_len_256) previously clamped KV writes with only a log line, silently degrading the audio tail. Port the Python reference's windowed-decode semantics to the streaming design: when the cache fills, reset it and re-decode the last 24 frames as context (audio discarded — already emitted), then continue. Extra cost only on long generations (~10% of decode calls per window). Unit-tested with a cache-advancing mock: 300 frames through a 256-slot cache, byte-exact emitted audio, >=1 re-prime. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Add an MLX-Swift port of the Qwen3-TTS talker as an opt-in CodeDecoder backend for voice cloning, integrated into the existing argmax-cli (no second CLI). `argmax-cli tts --code-decoder-backend mlx` swaps the CoreML talker for `MlxCodeDecoder` via the `TTSKitConfig.codeDecoder` override; every other component (voice-clone reference encoders, embedders, MultiCodeDecoder, SpeechDecoder) stays on CoreML. The MLX talker prefills the whole ICL prefix in one batched forward (~17x faster time-to-first-audio on 150-300-token voice-clone prompts) and has no compile-time KV cap. Default remains coreml; --mlx-model-dir and --mlx-max-sequence-length configure the MLX path. The new TTSKitMLX library target lives in the root package: mlx-swift floors at macOS 14, so the root macOS platform floor is bumped 13 -> 14. This is justified because TTS generation already requires macOS 15 at runtime (multifunction SpeechDecoder guard); iOS 16 / watchOS 10 / visionOS 1 floors are unchanged. The MLX products are macOS-conditional dependencies and all TTSKitMLX sources compile behind #if canImport(MLX), so iOS/watchOS/visionOS builds compile the target as empty. A `graft-mlx-metallib` Makefile target encapsulates the known workaround for running the mlx backend from a command-line SwiftPM build (xcodebuild compiles mlx-swift's Metal shaders into mlx-swift_Cmlx.bundle; graft it next to the release binary). TalkerParityTests assert parity against goldens exported from the Python mlx-audio reference: batched-prefill last-position logits and hidden cosine 1.000000, 20/20 greedy decode tokens exact, batched == sequential prefill, and KV bookkeeping through the TTSKit KVCache interface. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
dbrkn
force-pushed
the
berkin/voice-clone-coreml-talker
branch
from
July 21, 2026 16:29
730e3ef to
3c30e2d
Compare
The MLX talker keeps a single private KV cache per decoder instance; concurrent chunk workers interleave batched prefills into it, crashing with broadcast_shapes mask/score mismatches or MLXArray retain-count faults on multi-chunk texts. Force --concurrent-worker-count 1 for --code-decoder-backend mlx (matches the sequential Python prototype). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
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
1 participant
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.
This adds two things on top of the voice-clone base:
Extensions/TTSKitMLX — an opt-in MLX-Swift port of the Qwen3-TTS talker (CodeDecoder). Voice-clone ICL prefixes are 150–300+ tokens, and the CoreML talker prefills one position per call (~50 tok/s ⇒ seconds of time-to-first-audio per clone) with a compile-time KV cap on prompt + generation. The MLX talker prefills the whole prefix in one batched forward (~900 tok/s measured on M-series, ~17×) and has no KV cap; decode-step latency is at parity (steps are dominated by the CoreML MultiCodeDecoder/SpeechDecoder). Parity vs the Python reference: prefill logits cosine 1.000000, 20/20 greedy-decode tokens exact. Lives in a nested package because mlx-swift floors at macOS 14 (no watchOS) and SwiftPM platform floors are package-wide — TTSKit itself gains no dependency and keeps its platforms. Injected via the existing TTSKitConfig.codeDecoder override; ttskit-mlx-cli tts drives the full hybrid (CoreML encoders + MLX talker + CoreML decode) end-to-end.
SpeechStreamWriter: KV-cache re-prime on overflow. Generations longer than the SpeechDecoder's cache window (~21s on the kv_len_256 asset) previously clamped KV writes with only a log line, silently degrading the audio tail. When the cache fills, the writer now resets it and re-decodes the last 24 frames as context (audio discarded — already emitted), then continues. Unit-tested: 300 frames through a 256-slot cache, byte-exact emitted audio.
Requires: the Base-family mlx-community talker checkpoint (opt-in path only; a pruned talker-only export is a follow-up).