TTS-0.75B support on moshi_mlx

Due diligence

• I have done my due diligence in trying to find the answer myself.

Topic

The MLX implementation

Question

Hi Kyutai team,

I've been running moshi_mlx for TTS inference on Apple Silicon. kyutai/tts-1.6b-en_fr works out of the box, but kyutai/tts-0.75b-en-public fails immediately — it looks like moshi_mlx's TTS loading/inference path doesn't currently support this checkpoint's config (vocab sizes, codebook count, and conditioner setup all differ from the 1.6B).

Does moshi_mlx officially support tts-0.75b-en-public? If not, is support planned, or is the PyTorch backend the only supported path for this checkpoint right now?

I attempted to patch moshi_mlx myself to get it running, and wanted to share what I found in case it's useful, and to ask whether these patches are actually correct/intended or just happen to avoid crashing.

Environment

• moshi_mlx (installed via pip, .venv, Python 3.12)
• macOS, Apple Silicon
• Command: python -m moshi_mlx.run_tts --quantize 8 --nq 16 --hf-repo kyutai/tts-0.75b-en-public input.jsonl

Patches attempted

1. lm.py — text_out_vocab_size reads the wrong config field

# was:
text_out_vocab_size=data["text_card"],
# changed to:
text_out_vocab_size=data["text_card_out"],

Error before patch: Expected shape (8000, 1024) but received shape (5, 1024) for parameter text_linear.weight

2. run_tts.py — Mimi audio tokenizer loaded with the wrong codebook count

# was:
audio_tokenizer = models.mimi.Mimi(models.mimi_202407(generated_codebooks))
# changed to:
audio_tokenizer = models.mimi.Mimi(models.mimi_202407(32))

The 0.75B checkpoint only uses 16 generated codebooks downstream, but Mimi itself needs to be initialized with the full 32 to load its weights without missing-key errors; codebooks get truncated later in get_prefix.

3. tts.py, get_prefix() — dtype overflow + wrong codebook count

prefix = self.mimi.encode(mx.array(wav)[None])[0, :, :-2].astype(mx.int64)  # was uint32
prefix = prefix[: self.lm.n_q]  # truncate 32 -> 16 codebooks

Error before patch: ValueError: Converting -1 to uint32 would result in overflow — token_ids.zero is -1 for this checkpoint (no CFG distillation), and MLX won't implicitly cast -1 into the uint32 dtype the encoded prefix defaults to.

4. lm.py, Lm.__init__ — condition_provider set to None when conditioners config is empty

# was:
if len(cfg.conditioners) > 0:
    self.condition_provider = ConditionProvider(cfg.transformer.d_model, cfg.conditioners)
else:
    self.condition_provider = None
# changed to:
self.condition_provider = ConditionProvider(cfg.transformer.d_model, cfg.conditioners)

plus tts.py, generate() — iterate over self.lm.condition_provider.conditioners (matching how PyTorch's ConditionProvider.prepare()/_collate_text() drive iteration off the registered conditioners) rather than over attributes[i].text.keys(), since tts-0.75b-en-public's config has "conditioners": {} and the original code crashed looking up an unsupported "control" key.

Error before patch: assert self.lm.condition_provider is not None → AssertionError, then (after patch 4a alone) ValueError: unsupported conditioner control.

5. tts.py, generate() — stale num_codebooks attribute

# was:
assert K == self.lm.num_codebooks
# changed to:
assert K == self.lm.audio_offset + self.lm.n_q

Lm doesn't define num_codebooks anywhere in moshi_mlx (confirmed via LmConfig/Lm — only n_q, dep_q, audio_offset, delays exist as properties). PyTorch's Lm does define num_codebooks, so this looks like a naming drift between the two ports that never got reconciled in this one assert.

6. tts.py, _on_audio_hook() — broadcasting shape mismatch

# was:
audio_codes = audio_prefix.popleft()
# changed to:
audio_codes = audio_prefix.popleft()[:, None]

Error before patch: ValueError: Cannot broadcast array of shape (16,16) into shape (1,16,1) — the popped (16,) row needs an explicit trailing axis to align with audio_tokens[b]'s (16,1) shape for mx.where.

Result

With all six patches applied, the script runs to completion and produces a .wav file — but the generated audio is garbled/unintelligible babble, not coherent speech. So while these patches get past every crash, I don't think they're numerically correct — likely one or more of them (my best guesses: the Mimi codebook truncation in patch 2/3, or a subtlety in the delay/prefix handling in _on_audio_hook) is producing valid-shaped but wrong data rather than reproducing the actual intended computation.

A few things stood out that might help narrow down where the bug actually is:

• Voice cloning appears to work. The generated babble is clearly in the reference speaker's voice/timbre, not a generic/default voice — suggesting the audio-prefix conditioning path (get_prefix) is at least partially correct, even though the resulting speech content isn't.
• Duration scales correctly with input text length. A short sentence produces ~3s of audio; a longer sentence produces ~8s; both feel like plausible durations for the corresponding text. This suggests the text-to-audio-frame timing/delay logic is roughly right, and the bug is more likely in codebook/token content than in sequence length or delay-step handling.

Together these make me lean toward the bug being in something that affects per-token values rather than shape/timing — which is why I suspect the Mimi codebook truncation (patch 2/3) most: if the wrong 16 of 32 codebooks are being kept, or codebooks are being truncated in the wrong order/axis, you'd expect exactly this kind of result — right voice, right length, wrong content.

Attached:

• input.jsonl: {"turns": ["Hey there, how are you. Do you think this is working?"], "voices": ["ref.wav"], "id": "test1"}
• reference voice .wav ref.mp3
• generated output .wav output.mp3

Questions

1. Is tts-0.75b-en-public intended to be supported by moshi_mlx, or is PyTorch currently the only supported backend for this checkpoint?
2. If MLX support is intended/planned, is there a reference implementation or WIP branch I should be testing against instead of patching blind?
3. Of the patches above, which (if any) are actually correct, and which are likely masking a deeper issue (silently producing wrong values rather than raising an error)? In particular I suspect the Mimi-codebook-truncation step (patch 2/3) and the audio-prefix delay logic (patch 6) as the most likely sources of the garbled output.

Thanks for any guidance — happy to share the full diff or a minimal repro if useful.