diff --git a/README.md b/README.md index e571542..d2f325a 100644 --- a/README.md +++ b/README.md @@ -11,6 +11,9 @@ a flexible formulation for streaming, multimodal sequence-to-sequence learning. Hugging Face + + Open In Colab + **More details can be found on the [project page](https://kyutai.org/next/stt).** @@ -44,7 +47,8 @@ Here is how to choose which one to use: MLX is Apple's ML framework that allows you to use hardware acceleration on Apple silicon. If you want to run the model on a Mac or an iPhone, choose the MLX implementation. -### PyTorch implementation +
+PyTorch implementation Hugging Face @@ -99,8 +103,10 @@ In the heart of an ancient forest, where the trees whispered secrets of the past Apart from nudging the model for a specific spelling of a word, other potential use-cases include speaker adaptation and steering the model towards a specific formatting style or even a language. However, please bear in mind that is an experimental feature and its behavior is very sensitive to the prompt provided. +
-### Rust server +
+Rust server Hugging Face @@ -143,8 +149,10 @@ The script limits the decoding speed to simulates real-time processing of the au Faster processing can be triggered by setting the real-time factor, e.g. `--rtf 1000` will process the data as fast as possible. +
-### Rust standalone +
+Rust standalone Hugging Face @@ -157,8 +165,10 @@ cargo run --features cuda -r -- audio/bria.mp3 ``` You can get the timestamps by adding the `--timestamps` flag, and see the output of the semantic VAD by adding the `--vad` flag. +
-### MLX implementation +
+MLX implementation Hugging Face @@ -187,9 +197,13 @@ python scripts/stt_from_mic_mlx.py The MLX models can also be used in swift using the [moshi-swift codebase](https://github.com/kyutai-labs/moshi-swift), the 1b model has been tested to work fine on an iPhone 16 Pro. +
## Kyutai Text-to-Speech + + Hugging Face + Open In Colab @@ -200,7 +214,8 @@ We provide different implementations of Kyutai TTS for different use cases. Here - Rust: for production. If you want to serve Kyutai TTS in a production setting, use our Rust server. Our robust Rust server provides streaming access to the model over websockets. We use this server to run Unmute. - MLX: for on-device inference on iPhone and Mac. MLX is Apple's ML framework that allows you to use hardware acceleration on Apple silicon. If you want to run the model on a Mac or an iPhone, choose the MLX implementation. -### PyTorch implementation +
+PyTorch implementation Open In Colab @@ -219,8 +234,10 @@ python scripts/tts_pytorch.py text_to_say.txt audio_output.wav This requires the [moshi package](https://pypi.org/project/moshi/), which can be installed via pip. If you have [uv](https://docs.astral.sh/uv/) installed, you can skip the installation step and just prefix the command above with `uvx --with moshi`. +
-### Rust server +
+Rust server The Rust implementation provides a server that can process multiple streaming @@ -251,18 +268,23 @@ echo "Hey, how are you?" | python scripts/tts_rust_server.py - - # From text file to audio file python scripts/tts_rust_server.py text_to_say.txt audio_output.wav ``` +
-### MLX implementation +
+MLX implementation [MLX](https://ml-explore.github.io/mlx/build/html/index.html) is Apple's ML framework that allows you to use hardware acceleration on Apple silicon. Use our example script to run Kyutai TTS on MLX. The script takes text from stdin or a file and can output to a file or stream the resulting audio. +When streaming the output, if the model is not fast enough to keep with +real-time, you can use the `--quantize 8` or `--quantize 4` flags to quantize +the model resulting in faster inference. ```bash # From stdin, plays audio immediately -echo "Hey, how are you?" | python scripts/tts_mlx.py - - +echo "Hey, how are you?" | python scripts/tts_mlx.py - - --quantize 8 # From text file to audio file python scripts/tts_mlx.py text_to_say.txt audio_output.wav @@ -271,6 +293,7 @@ python scripts/tts_mlx.py text_to_say.txt audio_output.wav This requires the [moshi-mlx package](https://pypi.org/project/moshi-mlx/), which can be installed via pip. If you have [uv](https://docs.astral.sh/uv/) installed, you can skip the installation step and just prefix the command above with `uvx --with moshi-mlx`. +
## License @@ -290,4 +313,4 @@ pip install pre-commit pre-commit install ``` -If you're using `uv`, you can replace the two commands with `uvx pre-commit install`. \ No newline at end of file +If you're using `uv`, you can replace the two commands with `uvx pre-commit install`. diff --git a/scripts/tts_mlx.py b/scripts/tts_mlx.py index 8e08677..e59f962 100644 --- a/scripts/tts_mlx.py +++ b/scripts/tts_mlx.py @@ -2,7 +2,7 @@ # requires-python = ">=3.12" # dependencies = [ # "huggingface_hub", -# "moshi_mlx>=0.2.8", +# "moshi_mlx @ git+https://git@github.com/kyutai-labs/moshi#egg=moshi_mlx&subdirectory=moshi_mlx", # "numpy", # "sounddevice", # ] @@ -142,10 +142,10 @@ def main(): wav_frames = queue.Queue() - def _on_audio_hook(audio_tokens): - if (audio_tokens == -1).any(): + def _on_frame(frame): + if (frame == -1).any(): return - _pcm = tts_model.mimi.decode_step(audio_tokens[None, :, None]) + _pcm = tts_model.mimi.decode_step(frame[:, :, None]) _pcm = np.array(mx.clip(_pcm[0, 0], -1, 1)) wav_frames.put_nowait(_pcm) @@ -157,7 +157,7 @@ def main(): all_attributes, cfg_is_no_prefix=cfg_is_no_prefix, cfg_is_no_text=cfg_is_no_text, - on_audio_hook=_on_audio_hook, + on_frame=_on_frame, ) frames = mx.concat(result.frames, axis=-1) total_duration = frames.shape[0] * frames.shape[-1] / mimi.frame_rate @@ -188,6 +188,7 @@ def main(): break time.sleep(1) else: + run() frames = [] while True: try: