tipi/kyutai
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Workaround for the mlx kv-cache bug.

Browse Source
This commit is contained in:
Laurent 2025-08-04 16:33:43 +02:00
parent 09468c239a
commit 9de66be9ff
2 changed files with 6 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
scripts/tts_mlx.py
View File

@ -76,6 +76,9 @@ def main():
moshi_weights = hf_get(moshi_name, args.hf_repo)
tokenizer = hf_get(raw_config["tokenizer_name"], args.hf_repo)
lm_config = models.LmConfig.from_config_dict(raw_config)
# There is a bug in moshi_mlx <= 0.3.0 handling of the ring kv cache.
# The following line gets around it for now.
lm_config.transformer.max_seq_len = lm_config.transformer.context
model = models.Lm(lm_config)
model.set_dtype(mx.bfloat16)

3
scripts/tts_mlx_streaming.py
View File

@ -205,6 +205,9 @@ def main():
moshi_weights = hf_get(moshi_name, args.hf_repo)
tokenizer = hf_get(raw_config["tokenizer_name"], args.hf_repo)
lm_config = models.LmConfig.from_config_dict(raw_config)
# There is a bug in moshi_mlx <= 0.3.0 handling of the ring kv cache.
# The following line gets around it for now.
lm_config.transformer.max_seq_len = lm_config.transformer.context
model = models.Lm(lm_config)
model.set_dtype(mx.bfloat16)