feat: add CosyVoice3 MLX implementation - #861
Conversation
| component (v2 uses a UNet ConditionalDecoder). | ||
| 3. HiFT : CausalHiFTGenerator vocoder (mel -> waveform, NSF + ISTFT). | ||
|
|
||
| Reference (PyTorch): /Users/admin/CosyVoice/cosyvoice |
There was a problem hiding this comment.
Can you remove the implementation-specific comments here and elsewhere? We only want comments that improve the understanding of the code, not porting notes.
| import mlx.core as mx | ||
|
|
||
|
|
||
| def nucleus_sampling(logits: mx.array, top_p: float = 0.8, top_k: int = 25) -> int: |
There was a problem hiding this comment.
Can we just use make_sampler here or is there something specific needed for these?
| streaming_interval: float = 2.0, | ||
| save: bool = False, | ||
| use_zero_spk_emb: bool = False, | ||
| cross_lingual: bool = False, |
There was a problem hiding this comment.
These should be passed through to the model as extra kwargs, we don't want these in the standard generate interface.
| claude.md | ||
| .claude/site/ | ||
| .generated/legacy-docs-generated/ | ||
|
|
|
@walkoncross Thanks for the submission! This needs some cleanup, I left a few comments inline. Also, the examples should just live in the README.md, we don't want model-specific artifacts in there. If you can make those changes we can consider this. |
lucasnewman
left a comment
There was a problem hiding this comment.
See comments above.
Remove CosyVoice3 porting notes from runtime code, move model-specific arguments out of the shared generate interface, and align tests and examples with the reviewer feedback. Co-Authored-By: Claude <noreply@anthropic.com>
Summary
Adds a full MLX port of CosyVoice3 (
mlx_audio/tts/models/cosyvoice3/), following the existing model contract used by other TTS models in this repo (e.g. chatterbox, spark).instruct2) generation, matching the referenceCosyVoice3Python implementation's behavior.spk2infosystem for caching/reusing zero-shot speaker prompts (--add_spk,--save_spkinfo,--load_spkinfo).CosyVoice3Tokenizerspecial-token registration for fine-grained control tokens.convert.pyso official CosyVoice3 checkpoints can be converted to MLX weights./CosyVoice/cosyvoice): causal HiFT rewrite, DiT RoPE pairing/application order, CFM noise seeding to remove pitch bias, sample-rate handling for timbre cloning, prompt-feature/silence-token alignment with the reference, and a defensive post-hoc fix for a rare stochastic LLM-sampling tail-hiss artifact (acoustic-signature detection + fade, since the exact triggering token sequence could not be reliably reproduced after 100+ attempts — seeREADME.mdunder this model's## Statussection for the full investigation).## Statussection tracking known issues/investigations.45 commits, squash-mergeable if preferred. See individual commit messages for a detailed history of the implementation and each fix.
Test plan
pytest mlx_audio/tts/tests/test_cosyvoice3.py— 31 passed (2 pre-existing failures unrelated to this branch, confirmed viagit stashagainstmain:TestLLM::test_autoregressive_decode,TestModelIntegration::test_generate_end_to_end_with_mock_frontend, bothAttributeErroron an incompleteMockFrontEndtest double)generate()validated against real CosyVoice3-0.5B checkpoint weights for zero-shot, cross-lingual, and instruct2 modes