rust/delayed-streams-modeling
1
0
Fork 0
mirror of https://github.com/kyutai-labs/delayed-streams-modeling.git synced 2025-12-22 19:09:57 +00:00
Code Issues 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
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)