diff --git a/README.md b/README.md index 28c723f..ba69232 100644 --- a/README.md +++ b/README.md @@ -24,11 +24,11 @@ Supports Python 3.10, 3.11, 3.12, 3.13 and 3.14. Requires PyTorch 2.5+. Does not * Uses only 2 CPU cores * Python API and CLI * Voice cloning -* English only at the moment +* Multi-language support: english, french, german, portuguese, italian, spanish * Can handle infinitely long text inputs * [Can run on client-side in the browser](#in-browser-implementations) -More languages are planned: See our [official announcement](https://github.com/kyutai-labs/pocket-tts/issues/118) +Additional languages may be added in the future. ## Trying it from the website, without installing anything @@ -48,18 +48,33 @@ uvx pocket-tts generate pocket-tts generate ``` Modify the voice with `--voice` and the text with `--text`. We provide a small catalog of voices. +Choose a pretrained language model with `--language` when running `generate`, `export-voice`, or `serve` (default: `english`). +The `--config` option accepts only a local YAML path for custom weights. You can take a look at [this page](https://huggingface.co/kyutai/tts-voices) which details the licenses for each voice. * [alba](https://huggingface.co/kyutai/tts-voices/blob/main/alba-mackenna/casual.wav) -* [marius](https://huggingface.co/kyutai/tts-voices/blob/main/voice-donations/Selfie.wav) -* [javert](https://huggingface.co/kyutai/tts-voices/blob/main/voice-donations/Butter.wav) -* [jean](https://huggingface.co/kyutai/tts-voices/blob/main/ears/p010/freeform_speech_01.wav) -* [fantine](https://huggingface.co/kyutai/tts-voices/blob/main/vctk/p244_023.wav) +* [anna](https://huggingface.co/kyutai/tts-voices/blob/main/vctk/p228_023_enhanced.wav) +* [azelma](https://huggingface.co/kyutai/tts-voices/blob/main/vctk/p303_023_enhanced.wav) +* [bill_boerst](https://huggingface.co/kyutai/tts-voices/blob/main/voice-zero/bill_boerst.wav) +* [caro_davy](https://huggingface.co/kyutai/tts-voices/blob/main/voice-zero/caro_davy.wav) +* [charles](https://huggingface.co/kyutai/tts-voices/blob/main/vctk/p254_023_enhanced.wav) * [cosette](https://huggingface.co/kyutai/tts-voices/blob/main/expresso/ex04-ex02_confused_001_channel1_499s.wav) -* [eponine](https://huggingface.co/kyutai/tts-voices/blob/main/vctk/p262_023.wav) -* [azelma](https://huggingface.co/kyutai/tts-voices/blob/main/vctk/p303_023.wav) +* [eponine](https://huggingface.co/kyutai/tts-voices/blob/main/vctk/p262_023_enhanced.wav) +* [eve](https://huggingface.co/kyutai/tts-voices/blob/main/vctk/p361_023_enhanced.wav) +* [fantine](https://huggingface.co/kyutai/tts-voices/blob/main/vctk/p244_023_enhanced.wav) +* [george](https://huggingface.co/kyutai/tts-voices/blob/main/vctk/p315_023_enhanced.wav) +* [jane](https://huggingface.co/kyutai/tts-voices/blob/main/vctk/p339_023_enhanced.wav) +* [jean](https://huggingface.co/kyutai/tts-voices/blob/main/ears/p010/freeform_speech_01_enhanced.wav) +* [javert](https://huggingface.co/kyutai/tts-voices/blob/main/voice-donations/Butter.wav) +* [marius](https://huggingface.co/kyutai/tts-voices/blob/main/voice-donations/Selfie.wav) +* [mary](https://huggingface.co/kyutai/tts-voices/blob/main/vctk/p333_023_enhanced.wav) +* [michael](https://huggingface.co/kyutai/tts-voices/blob/main/vctk/p360_023_enhanced.wav) +* [paul](https://huggingface.co/kyutai/tts-voices/blob/main/vctk/p259_023_enhanced.wav) +* [peter_yearsley](https://huggingface.co/kyutai/tts-voices/blob/main/voice-zero/peter_yearsley.wav) +* [stuart_bell](https://huggingface.co/kyutai/tts-voices/blob/main/voice-zero/stuart_bell.wav) +* [vera](https://huggingface.co/kyutai/tts-voices/blob/main/vctk/p229_023_enhanced.wav) The `--voice` argument can also take a plain wav file as input for voice cloning. You can use your own or check out our [voice repository](https://huggingface.co/kyutai/tts-voices). diff --git a/docs/API Reference/python-api.md b/docs/API Reference/python-api.md index 866e291..23fb45c 100644 --- a/docs/API Reference/python-api.md +++ b/docs/API Reference/python-api.md @@ -37,16 +37,18 @@ The main class for text-to-speech generation. #### Class Methods -##### `load_model(config="b6369a24", temp=0.7, lsd_decode_steps=1, noise_clamp=None, eos_threshold=-4.0)` +##### `load_model(language=None, config=None, temp=0.7, lsd_decode_steps=1, noise_clamp=None, eos_threshold=-4.0, quantize=False)` Load and return a TTSModel instance with pre-trained weights. **Parameters:** -- `config` (str): Path to model config YAML file or a variant identifier (default: "b6369a24") +- `language` (str | None): Name of built-in language config to load. Supported values: `"english_2026-01"`, `"english_2026-04"`, `"english"`, `"french_24l"`, `"german_24l"`, `"portuguese_24l"`, `"italian_24l"`, `"spanish_24l"`. If both `language` and `config` are omitted, defaults to `"english"`, which is the same model as `"english_2026-04"`. The "24l" variants are larger models that are not distilled yet and are here only as a preview. +- `config` (str | None): Path to model config YAML file. Incompatible with `language`. - `temp` (float): Sampling temperature for generation (default: 0.7) - `lsd_decode_steps` (int): Number of generation steps (default: 1) - `noise_clamp` (float | None): Maximum value for noise sampling (default: None) - `eos_threshold` (float): Threshold for end-of-sequence detection (default: -4.0) +- `quantize` (bool): Enable int8 quantization when loading the model (default: `False`) **Returns:** - `TTSModel`: Loaded model instance on CPU @@ -59,7 +61,9 @@ from pocket_tts import TTSModel model = TTSModel.load_model() # Load with custom parameters -model = TTSModel.load_model(variant="b6369a24", temp=0.5, lsd_decode_steps=5, eos_threshold=-3.0) +model = TTSModel.load_model( + language="english_2026-01", temp=0.5, lsd_decode_steps=5, eos_threshold=-3.0 +) ``` #### Properties @@ -255,4 +259,4 @@ scipy.io.wavfile.write("batch_output.wav", model.sample_rate, full_audio.numpy() ### Streaming to File You can refer to our CLI implementation which can stream audio to a wav file. -For more information about the command-line interface, see the [Generate Documentation](generate.md) or [Serve Documentation](serve.md). \ No newline at end of file +For more information about the command-line interface, see the [Generate Documentation](generate.md) or [Serve Documentation](serve.md). diff --git a/docs/CLI Commands/export_voice.md b/docs/CLI Commands/export_voice.md index 442acdd..3575359 100644 --- a/docs/CLI Commands/export_voice.md +++ b/docs/CLI Commands/export_voice.md @@ -24,7 +24,8 @@ Only the first 30 seconds of the audio file will be processed. ### Options - `--quiet`: Do not print any output except errors. -- `--config`: Model config path or signature +- `--language`: Language for the TTS model, one of `'english_2026-01'`, `'english_2026-04'`, `'english'`, `'french_24l'`, `'german_24l'`, `'portuguese_24l'`, `'italian_24l'`, `'spanish_24l'` (default: `english`, which is the same model as `'english_2026-04'`). Incompatible with `--config`. The "24l" variants are bigger models, not distilled yet and here only as preview. +- `--config`: Model config local yaml path. Incompatible with `--language`. ## Examples diff --git a/docs/CLI Commands/generate.md b/docs/CLI Commands/generate.md index 6f21674..c60d155 100644 --- a/docs/CLI Commands/generate.md +++ b/docs/CLI Commands/generate.md @@ -19,10 +19,11 @@ This will generate a WAV file `./tts_output.wav` with the default text and voice - `--text TEXT`: Text to generate (default: "Hello world! I am Kyutai Pocket TTS. I'm fast enough to run on small CPUs. I hope you'll like me.") - `--voice VOICE`: Path to audio conditioning file (voice to clone) (default: "hf://kyutai/tts-voices/alba-mackenna/casual.wav"). Urls and local paths are supported. - `--output-path OUTPUT_PATH`: Output path for generated audio (default: "./tts_output.wav") +- `--language LANGUAGE`: Language for the TTS model, one of `'english_2026-01'`, `'english_2026-04'`, `'english'`, `'french_24l'`, `'german_24l'`, `'portuguese_24l'`, `'italian_24l'`, `'spanish_24l'` (default: `english`, which is the same model as `'english_2026-04'`). Incompatible with `--config`. The "24l" variants are bigger models, not distilled yet and here only as preview. ### Generation Parameters -- `--config CONFIG_PATH`: Path to custom config.yaml (for loading local model files) or model signature (default: "b6369a24") +- `--config CONFIG_PATH`: Path to custom config.yaml (for loading local model files). Incompatible with `--language`. - `--lsd-decode-steps LSD_DECODE_STEPS`: Number of generation steps (default: 1) - `--temperature TEMPERATURE`: Temperature for generation (default: 0.7) - `--noise-clamp NOISE_CLAMP`: Noise clamp value (default: None) @@ -32,6 +33,7 @@ This will generate a WAV file `./tts_output.wav` with the default text and voice ### Performance Options - `--device DEVICE`: Device to use (default: "cpu", you may not get a speedup by using a gpu since it's a small model) +- `--quantize`: Use int8 quantization for the model (default: False). This can reduce memory usage and increase speed, with minimal impact on audio quality. - `--quiet`, `-q`: Disable logging output ## Examples @@ -79,7 +81,7 @@ pocket-tts generate --eos-threshold -3.0 If you'd like to override the paths from which the models are loaded, you can provide a custom YAML configuration. -Copy pocket_tts/config/b6369a24.yaml and change weights_path:, weights_path_without_voice_cloning: and tokenizer_path: to the paths of the models you want to load. +Copy one of the files in `pocket_tts/config` (for example `pocket_tts/config/english.yaml`) and change `weights_path`, `weights_path_without_voice_cloning`, and `tokenizer_path` to the paths of the models you want to load. Then, use the --config option to point to your newly created config. diff --git a/docs/CLI Commands/serve.md b/docs/CLI Commands/serve.md index b18609f..7fc0486 100644 --- a/docs/CLI Commands/serve.md +++ b/docs/CLI Commands/serve.md @@ -14,12 +14,12 @@ This starts a server on `http://localhost:8000` with the default voice model. ## Command Options -- `--voice VOICE`: Path to voice prompt audio file (voice to clone) (default: "hf://kyutai/tts-voices/alba-mackenna/casual.wav") - `--host HOST`: Host to bind to (default: "localhost") - `--port PORT`: Port to bind to (default: 8000) - `--reload`: Enable auto-reload for development -- `--config`: Path to a custom config .yaml - +- `--language`: Language for the TTS model, one of `'english_2026-01'`, `'english_2026-04'`, `'english'`, `'french_24l'`, `'german_24l'`, `'portuguese_24l'`, `'italian_24l'`, `'spanish_24l'` (default: `english`, which is the same model as `'english_2026-04'`). Incompatible with `--config`. The "24l" variants are bigger models, not distilled yet and here only as preview. +- `--config`: Path to a custom config .yaml. Incompatible with `--language`. +- `--quantize`: Use int8 quantization for the model (default: False). This can reduce memory usage and increase speed, with minimal impact on audio quality. ## Examples ### Basic Server @@ -32,21 +32,17 @@ pocket-tts serve pocket-tts serve --host "localhost" --port 8080 ``` -### Custom Voice - +### Custom Language +To select the default language model, pass `--language`: ```bash -# Use different voice -pocket-tts serve --voice "hf://kyutai/tts-voices/jessica-jian/casual.wav" - -# Use local voice file -pocket-tts serve --voice "./my_voice.wav" +pocket-tts serve --language french_24l ``` ### Custom Model Config If you'd like to override the paths from which the models are loaded, you can provide a custom YAML configuration. -Copy pocket_tts/config/b6369a24.yaml and change weights_path:, weights_path_without_voice_cloning: and tokenizer_path: to the paths of the models you want to load. +Copy one of the files in `pocket_tts/config` (for example `pocket_tts/config/english.yaml`) and change `weights_path`, `weights_path_without_voice_cloning:`, and `tokenizer_path:` to the paths of the models you want to load. Then, use the --config option to point to your newly created config. @@ -60,4 +56,3 @@ pocket-tts serve --config "C://pocket-tts/my_config.yaml" Once the server is running, navigate to `http://localhost:8000` to access the web interface. For more advanced usage, see the [Python API documentation](python-api.md) for direct integration with the TTS model. - diff --git a/pocket_tts/config/english.yaml b/pocket_tts/config/english.yaml new file mode 100644 index 0000000..6877c37 --- /dev/null +++ b/pocket_tts/config/english.yaml @@ -0,0 +1,61 @@ +# sig: a0ac5076.yaml + +weights_path: hf://kyutai/pocket-tts/languages/english/model.safetensors@39592ff23c9ef80098bb74895d104c26275fe2c9 +weights_path_without_voice_cloning: hf://kyutai/pocket-tts-without-voice-cloning/languages/english/model.safetensors@d29db7978e464fb90cb3359ee0c69a273b9142cc + + +flow_lm: + insert_bos_before_voice: true + dtype: float32 + flow: + depth: 6 + dim: 512 + transformer: + d_model: 1024 + hidden_scale: 4 + max_period: 10000 + num_heads: 16 + num_layers: 6 + lookup_table: + dim: 1024 + n_bins: 4000 + tokenizer: sentencepiece + tokenizer_path: hf://kyutai/pocket-tts-without-voice-cloning/languages/english/tokenizer.model@d29db7978e464fb90cb3359ee0c69a273b9142cc + #weights_path: final.safetensors + +mimi: + dtype: float32 + sample_rate: 24000 + inner_dim: 32 + outer_dim: 512 + channels: 1 + frame_rate: 12.5 + seanet: + dimension: 512 + channels: 1 + n_filters: 64 + n_residual_layers: 1 + ratios: + - 6 + - 5 + - 4 + kernel_size: 7 + residual_kernel_size: 3 + last_kernel_size: 3 + dilation_base: 2 + pad_mode: constant + compress: 2 + transformer: + d_model: 512 + num_heads: 8 + num_layers: 2 + layer_scale: 0.01 + context: 250 + dim_feedforward: 2048 + input_dimension: 512 + output_dimensions: + - 512 + quantizer: + dimension: 32 + output_dimension: 512 + #weights_path: codec.safetensors diff --git a/pocket_tts/config/b6369a24.yaml b/pocket_tts/config/english_2026-01.yaml similarity index 73% rename from pocket_tts/config/b6369a24.yaml rename to pocket_tts/config/english_2026-01.yaml index 68c01ce..a730f60 100644 --- a/pocket_tts/config/b6369a24.yaml +++ b/pocket_tts/config/english_2026-01.yaml @@ -1,9 +1,12 @@ # sig: b6369a24 -weights_path: hf://kyutai/pocket-tts/tts_b6369a24.safetensors@427e3d61b276ed69fdd03de0d185fa8a8d97fc5b -weights_path_without_voice_cloning: hf://kyutai/pocket-tts-without-voice-cloning/tts_b6369a24.safetensors@d4fdd22ae8c8e1cb3634e150ebeff1dab2d16df3 +weights_path: hf://kyutai/pocket-tts/languages/english_2026-01/model.safetensors@39592ff23c9ef80098bb74895d104c26275fe2c9 +weights_path_without_voice_cloning: hf://kyutai/pocket-tts-without-voice-cloning/languages/english_2026-01/model.safetensors@d29db7978e464fb90cb3359ee0c69a273b9142cc + +pad_with_spaces_for_short_inputs: true flow_lm: + insert_bos_before_voice: false dtype: float32 flow: depth: 6 @@ -18,12 +21,14 @@ flow_lm: dim: 1024 n_bins: 4000 tokenizer: sentencepiece - tokenizer_path: hf://kyutai/pocket-tts-without-voice-cloning/tokenizer.model@d4fdd22ae8c8e1cb3634e150ebeff1dab2d16df3 + tokenizer_path: hf://kyutai/pocket-tts-without-voice-cloning/tokenizer.model@d29db7978e464fb90cb3359ee0c69a273b9142cc #weights_path: flow_lm_b6369a24.safetensors mimi: dtype: float32 sample_rate: 24000 + inner_dim: 512 + outer_dim: 512 channels: 1 frame_rate: 12.5 seanet: diff --git a/pocket_tts/config/english_2026-04.yaml b/pocket_tts/config/english_2026-04.yaml new file mode 100644 index 0000000..a610346 --- /dev/null +++ b/pocket_tts/config/english_2026-04.yaml @@ -0,0 +1,61 @@ +# sig: a0ac5076.yaml + +weights_path: hf://kyutai/pocket-tts/languages/english_2026-04/model.safetensors@19f95fe2df36e79fbd9f10008595cc4c977a0fcc +weights_path_without_voice_cloning: hf://kyutai/pocket-tts-without-voice-cloning/languages/english_2026-04/model.safetensors@d29db7978e464fb90cb3359ee0c69a273b9142cc + + +flow_lm: + insert_bos_before_voice: true + dtype: float32 + flow: + depth: 6 + dim: 512 + transformer: + d_model: 1024 + hidden_scale: 4 + max_period: 10000 + num_heads: 16 + num_layers: 6 + lookup_table: + dim: 1024 + n_bins: 4000 + tokenizer: sentencepiece + tokenizer_path: hf://kyutai/pocket-tts-without-voice-cloning/languages/english_2026-04/tokenizer.model@d29db7978e464fb90cb3359ee0c69a273b9142cc + #weights_path: final.safetensors + +mimi: + dtype: float32 + sample_rate: 24000 + inner_dim: 32 + outer_dim: 512 + channels: 1 + frame_rate: 12.5 + seanet: + dimension: 512 + channels: 1 + n_filters: 64 + n_residual_layers: 1 + ratios: + - 6 + - 5 + - 4 + kernel_size: 7 + residual_kernel_size: 3 + last_kernel_size: 3 + dilation_base: 2 + pad_mode: constant + compress: 2 + transformer: + d_model: 512 + num_heads: 8 + num_layers: 2 + layer_scale: 0.01 + context: 250 + dim_feedforward: 2048 + input_dimension: 512 + output_dimensions: + - 512 + quantizer: + dimension: 32 + output_dimension: 512 + #weights_path: codec.safetensors diff --git a/pocket_tts/config/french_24l.yaml b/pocket_tts/config/french_24l.yaml new file mode 100644 index 0000000..311448a --- /dev/null +++ b/pocket_tts/config/french_24l.yaml @@ -0,0 +1,63 @@ +# sig: 709d9f84.yaml + +weights_path: hf://kyutai/pocket-tts/languages/french_24l/model.safetensors@39592ff23c9ef80098bb74895d104c26275fe2c9 +weights_path_without_voice_cloning: hf://kyutai/pocket-tts-without-voice-cloning/languages/french_24l/model.safetensors@d29db7978e464fb90cb3359ee0c69a273b9142cc + +remove_semicolons: true +model_recommended_frames_after_eos: 12 + +flow_lm: + insert_bos_before_voice: true + dtype: float32 + flow: + depth: 6 + dim: 512 + transformer: + d_model: 1024 + hidden_scale: 4 + max_period: 10000 + num_heads: 16 + num_layers: 24 + lookup_table: + dim: 1024 + n_bins: 4000 + tokenizer: sentencepiece + tokenizer_path: hf://kyutai/pocket-tts-without-voice-cloning/languages/french_24l/tokenizer.model@d29db7978e464fb90cb3359ee0c69a273b9142cc + #weights_path: partial.safetensors + +mimi: + dtype: float32 + sample_rate: 24000 + inner_dim: 32 + outer_dim: 512 + channels: 1 + frame_rate: 12.5 + seanet: + dimension: 512 + channels: 1 + n_filters: 64 + n_residual_layers: 1 + ratios: + - 6 + - 5 + - 4 + kernel_size: 7 + residual_kernel_size: 3 + last_kernel_size: 3 + dilation_base: 2 + pad_mode: constant + compress: 2 + transformer: + d_model: 512 + num_heads: 8 + num_layers: 2 + layer_scale: 0.01 + context: 250 + dim_feedforward: 2048 + input_dimension: 512 + output_dimensions: + - 512 + quantizer: + dimension: 32 + output_dimension: 512 + #weights_path: codec.safetensors diff --git a/pocket_tts/config/german.yaml b/pocket_tts/config/german.yaml new file mode 100644 index 0000000..dbda4de --- /dev/null +++ b/pocket_tts/config/german.yaml @@ -0,0 +1,61 @@ +# sig: ba9816ab.yaml + +weights_path: hf://kyutai/pocket-tts/languages/german/model.safetensors@39592ff23c9ef80098bb74895d104c26275fe2c9 +weights_path_without_voice_cloning: hf://kyutai/pocket-tts-without-voice-cloning/languages/german/model.safetensors@d29db7978e464fb90cb3359ee0c69a273b9142cc + + +flow_lm: + insert_bos_before_voice: true + dtype: float32 + flow: + depth: 6 + dim: 512 + transformer: + d_model: 1024 + hidden_scale: 4 + max_period: 10000 + num_heads: 16 + num_layers: 6 + lookup_table: + dim: 1024 + n_bins: 4000 + tokenizer: sentencepiece + tokenizer_path: hf://kyutai/pocket-tts-without-voice-cloning/languages/german/tokenizer.model@d29db7978e464fb90cb3359ee0c69a273b9142cc + #weights_path: partial.safetensors + +mimi: + dtype: float32 + sample_rate: 24000 + inner_dim: 32 + outer_dim: 512 + channels: 1 + frame_rate: 12.5 + seanet: + dimension: 512 + channels: 1 + n_filters: 64 + n_residual_layers: 1 + ratios: + - 6 + - 5 + - 4 + kernel_size: 7 + residual_kernel_size: 3 + last_kernel_size: 3 + dilation_base: 2 + pad_mode: constant + compress: 2 + transformer: + d_model: 512 + num_heads: 8 + num_layers: 2 + layer_scale: 0.01 + context: 250 + dim_feedforward: 2048 + input_dimension: 512 + output_dimensions: + - 512 + quantizer: + dimension: 32 + output_dimension: 512 + #weights_path: codec.safetensors diff --git a/pocket_tts/config/german_24l.yaml b/pocket_tts/config/german_24l.yaml new file mode 100644 index 0000000..8d0cd25 --- /dev/null +++ b/pocket_tts/config/german_24l.yaml @@ -0,0 +1,61 @@ +# sig: 085f1405.yaml + +weights_path: hf://kyutai/pocket-tts/languages/german_24l/model.safetensors@39592ff23c9ef80098bb74895d104c26275fe2c9 +weights_path_without_voice_cloning: hf://kyutai/pocket-tts-without-voice-cloning/languages/german_24l/model.safetensors@d29db7978e464fb90cb3359ee0c69a273b9142cc + + +flow_lm: + insert_bos_before_voice: true + dtype: float32 + flow: + depth: 6 + dim: 512 + transformer: + d_model: 1024 + hidden_scale: 4 + max_period: 10000 + num_heads: 16 + num_layers: 24 + lookup_table: + dim: 1024 + n_bins: 4000 + tokenizer: sentencepiece + tokenizer_path: hf://kyutai/pocket-tts-without-voice-cloning/languages/german_24l/tokenizer.model@d29db7978e464fb90cb3359ee0c69a273b9142cc + #weights_path: partial.safetensors + +mimi: + dtype: float32 + sample_rate: 24000 + inner_dim: 32 + outer_dim: 512 + channels: 1 + frame_rate: 12.5 + seanet: + dimension: 512 + channels: 1 + n_filters: 64 + n_residual_layers: 1 + ratios: + - 6 + - 5 + - 4 + kernel_size: 7 + residual_kernel_size: 3 + last_kernel_size: 3 + dilation_base: 2 + pad_mode: constant + compress: 2 + transformer: + d_model: 512 + num_heads: 8 + num_layers: 2 + layer_scale: 0.01 + context: 250 + dim_feedforward: 2048 + input_dimension: 512 + output_dimensions: + - 512 + quantizer: + dimension: 32 + output_dimension: 512 + #weights_path: codec.safetensors diff --git a/pocket_tts/config/italian.yaml b/pocket_tts/config/italian.yaml new file mode 100644 index 0000000..5cf7392 --- /dev/null +++ b/pocket_tts/config/italian.yaml @@ -0,0 +1,61 @@ +# sig: 78c8126d.yaml + +weights_path: hf://kyutai/pocket-tts/languages/italian/model.safetensors@39592ff23c9ef80098bb74895d104c26275fe2c9 +weights_path_without_voice_cloning: hf://kyutai/pocket-tts-without-voice-cloning/languages/italian/model.safetensors@d29db7978e464fb90cb3359ee0c69a273b9142cc + + +flow_lm: + insert_bos_before_voice: true + dtype: float32 + flow: + depth: 6 + dim: 512 + transformer: + d_model: 1024 + hidden_scale: 4 + max_period: 10000 + num_heads: 16 + num_layers: 6 + lookup_table: + dim: 1024 + n_bins: 4000 + tokenizer: sentencepiece + tokenizer_path: hf://kyutai/pocket-tts-without-voice-cloning/languages/italian/tokenizer.model@d29db7978e464fb90cb3359ee0c69a273b9142cc + #weights_path: partial.safetensors + +mimi: + dtype: float32 + sample_rate: 24000 + inner_dim: 32 + outer_dim: 512 + channels: 1 + frame_rate: 12.5 + seanet: + dimension: 512 + channels: 1 + n_filters: 64 + n_residual_layers: 1 + ratios: + - 6 + - 5 + - 4 + kernel_size: 7 + residual_kernel_size: 3 + last_kernel_size: 3 + dilation_base: 2 + pad_mode: constant + compress: 2 + transformer: + d_model: 512 + num_heads: 8 + num_layers: 2 + layer_scale: 0.01 + context: 250 + dim_feedforward: 2048 + input_dimension: 512 + output_dimensions: + - 512 + quantizer: + dimension: 32 + output_dimension: 512 + #weights_path: codec.safetensors diff --git a/pocket_tts/config/italian_24l.yaml b/pocket_tts/config/italian_24l.yaml new file mode 100644 index 0000000..a1ed4f6 --- /dev/null +++ b/pocket_tts/config/italian_24l.yaml @@ -0,0 +1,61 @@ +# sig: 6ec79d0f + +weights_path: hf://kyutai/pocket-tts/languages/italian_24l/model.safetensors@39592ff23c9ef80098bb74895d104c26275fe2c9 +weights_path_without_voice_cloning: hf://kyutai/pocket-tts-without-voice-cloning/languages/italian_24l/model.safetensors@d29db7978e464fb90cb3359ee0c69a273b9142cc + + +flow_lm: + insert_bos_before_voice: true + dtype: float32 + flow: + depth: 6 + dim: 512 + transformer: + d_model: 1024 + hidden_scale: 4 + max_period: 10000 + num_heads: 16 + num_layers: 24 + lookup_table: + dim: 1024 + n_bins: 4000 + tokenizer: sentencepiece + tokenizer_path: hf://kyutai/pocket-tts-without-voice-cloning/languages/italian_24l/tokenizer.model@d29db7978e464fb90cb3359ee0c69a273b9142cc + #weights_path: partial.safetensors + +mimi: + dtype: float32 + sample_rate: 24000 + inner_dim: 32 + outer_dim: 512 + channels: 1 + frame_rate: 12.5 + seanet: + dimension: 512 + channels: 1 + n_filters: 64 + n_residual_layers: 1 + ratios: + - 6 + - 5 + - 4 + kernel_size: 7 + residual_kernel_size: 3 + last_kernel_size: 3 + dilation_base: 2 + pad_mode: constant + compress: 2 + transformer: + d_model: 512 + num_heads: 8 + num_layers: 2 + layer_scale: 0.01 + context: 250 + dim_feedforward: 2048 + input_dimension: 512 + output_dimensions: + - 512 + quantizer: + dimension: 32 + output_dimension: 512 + #weights_path: codec.safetensors diff --git a/pocket_tts/config/portuguese.yaml b/pocket_tts/config/portuguese.yaml new file mode 100644 index 0000000..ea57c23 --- /dev/null +++ b/pocket_tts/config/portuguese.yaml @@ -0,0 +1,61 @@ +# sig: 4a4b1e3a.yaml + +weights_path: hf://kyutai/pocket-tts/languages/portuguese/model.safetensors@39592ff23c9ef80098bb74895d104c26275fe2c9 +weights_path_without_voice_cloning: hf://kyutai/pocket-tts-without-voice-cloning/languages/portuguese/model.safetensors@d29db7978e464fb90cb3359ee0c69a273b9142cc + + +flow_lm: + insert_bos_before_voice: true + dtype: float32 + flow: + depth: 6 + dim: 512 + transformer: + d_model: 1024 + hidden_scale: 4 + max_period: 10000 + num_heads: 16 + num_layers: 6 + lookup_table: + dim: 1024 + n_bins: 4000 + tokenizer: sentencepiece + tokenizer_path: hf://kyutai/pocket-tts-without-voice-cloning/languages/portuguese/tokenizer.model@d29db7978e464fb90cb3359ee0c69a273b9142cc + #weights_path: partial.safetensors + +mimi: + dtype: float32 + sample_rate: 24000 + inner_dim: 32 + outer_dim: 512 + channels: 1 + frame_rate: 12.5 + seanet: + dimension: 512 + channels: 1 + n_filters: 64 + n_residual_layers: 1 + ratios: + - 6 + - 5 + - 4 + kernel_size: 7 + residual_kernel_size: 3 + last_kernel_size: 3 + dilation_base: 2 + pad_mode: constant + compress: 2 + transformer: + d_model: 512 + num_heads: 8 + num_layers: 2 + layer_scale: 0.01 + context: 250 + dim_feedforward: 2048 + input_dimension: 512 + output_dimensions: + - 512 + quantizer: + dimension: 32 + output_dimension: 512 + #weights_path: codec.safetensors diff --git a/pocket_tts/config/portuguese_24l.yaml b/pocket_tts/config/portuguese_24l.yaml new file mode 100644 index 0000000..d84bde1 --- /dev/null +++ b/pocket_tts/config/portuguese_24l.yaml @@ -0,0 +1,61 @@ +# sig: 56579f81 + +weights_path: hf://kyutai/pocket-tts/languages/portuguese_24l/model.safetensors@39592ff23c9ef80098bb74895d104c26275fe2c9 +weights_path_without_voice_cloning: hf://kyutai/pocket-tts-without-voice-cloning/languages/portuguese_24l/model.safetensors@d29db7978e464fb90cb3359ee0c69a273b9142cc + + +flow_lm: + insert_bos_before_voice: true + dtype: float32 + flow: + depth: 6 + dim: 512 + transformer: + d_model: 1024 + hidden_scale: 4 + max_period: 10000 + num_heads: 16 + num_layers: 24 + lookup_table: + dim: 1024 + n_bins: 4000 + tokenizer: sentencepiece + tokenizer_path: hf://kyutai/pocket-tts-without-voice-cloning/languages/portuguese_24l/tokenizer.model@d29db7978e464fb90cb3359ee0c69a273b9142cc + #weights_path: partial.safetensors + +mimi: + dtype: float32 + sample_rate: 24000 + inner_dim: 32 + outer_dim: 512 + channels: 1 + frame_rate: 12.5 + seanet: + dimension: 512 + channels: 1 + n_filters: 64 + n_residual_layers: 1 + ratios: + - 6 + - 5 + - 4 + kernel_size: 7 + residual_kernel_size: 3 + last_kernel_size: 3 + dilation_base: 2 + pad_mode: constant + compress: 2 + transformer: + d_model: 512 + num_heads: 8 + num_layers: 2 + layer_scale: 0.01 + context: 250 + dim_feedforward: 2048 + input_dimension: 512 + output_dimensions: + - 512 + quantizer: + dimension: 32 + output_dimension: 512 + # weights_path: codec.safetensors diff --git a/pocket_tts/config/spanish.yaml b/pocket_tts/config/spanish.yaml new file mode 100644 index 0000000..f974781 --- /dev/null +++ b/pocket_tts/config/spanish.yaml @@ -0,0 +1,61 @@ +# sig: 8bcc8b3b.yaml + +weights_path: hf://kyutai/pocket-tts/languages/spanish/model.safetensors@39592ff23c9ef80098bb74895d104c26275fe2c9 +weights_path_without_voice_cloning: hf://kyutai/pocket-tts-without-voice-cloning/languages/spanish/model.safetensors@d29db7978e464fb90cb3359ee0c69a273b9142cc + + +flow_lm: + insert_bos_before_voice: true + dtype: float32 + flow: + depth: 6 + dim: 512 + transformer: + d_model: 1024 + hidden_scale: 4 + max_period: 10000 + num_heads: 16 + num_layers: 6 + lookup_table: + dim: 1024 + n_bins: 4000 + tokenizer: sentencepiece + tokenizer_path: hf://kyutai/pocket-tts-without-voice-cloning/languages/spanish/tokenizer.model@d29db7978e464fb90cb3359ee0c69a273b9142cc + #weights_path: partial.safetensors + +mimi: + dtype: float32 + sample_rate: 24000 + inner_dim: 32 + outer_dim: 512 + channels: 1 + frame_rate: 12.5 + seanet: + dimension: 512 + channels: 1 + n_filters: 64 + n_residual_layers: 1 + ratios: + - 6 + - 5 + - 4 + kernel_size: 7 + residual_kernel_size: 3 + last_kernel_size: 3 + dilation_base: 2 + pad_mode: constant + compress: 2 + transformer: + d_model: 512 + num_heads: 8 + num_layers: 2 + layer_scale: 0.01 + context: 250 + dim_feedforward: 2048 + input_dimension: 512 + output_dimensions: + - 512 + quantizer: + dimension: 32 + output_dimension: 512 + #weights_path: codec.safetensors diff --git a/pocket_tts/config/spanish_24l.yaml b/pocket_tts/config/spanish_24l.yaml new file mode 100644 index 0000000..1651623 --- /dev/null +++ b/pocket_tts/config/spanish_24l.yaml @@ -0,0 +1,61 @@ +# sig: afd6403b + +weights_path: hf://kyutai/pocket-tts/languages/spanish_24l/model.safetensors@39592ff23c9ef80098bb74895d104c26275fe2c9 +weights_path_without_voice_cloning: hf://kyutai/pocket-tts-without-voice-cloning/languages/spanish_24l/model.safetensors@d29db7978e464fb90cb3359ee0c69a273b9142cc + + +flow_lm: + insert_bos_before_voice: true + dtype: float32 + flow: + depth: 6 + dim: 512 + transformer: + d_model: 1024 + hidden_scale: 4 + max_period: 10000 + num_heads: 16 + num_layers: 24 + lookup_table: + dim: 1024 + n_bins: 4000 + tokenizer: sentencepiece + tokenizer_path: hf://kyutai/pocket-tts-without-voice-cloning/languages/spanish_24l/tokenizer.model@d29db7978e464fb90cb3359ee0c69a273b9142cc + #weights_path: partial.safetensors + +mimi: + dtype: float32 + sample_rate: 24000 + inner_dim: 32 + outer_dim: 512 + channels: 1 + frame_rate: 12.5 + seanet: + dimension: 512 + channels: 1 + n_filters: 64 + n_residual_layers: 1 + ratios: + - 6 + - 5 + - 4 + kernel_size: 7 + residual_kernel_size: 3 + last_kernel_size: 3 + dilation_base: 2 + pad_mode: constant + compress: 2 + transformer: + d_model: 512 + num_heads: 8 + num_layers: 2 + layer_scale: 0.01 + context: 250 + dim_feedforward: 2048 + input_dimension: 512 + output_dimensions: + - 512 + quantizer: + dimension: 32 + output_dimension: 512 + #weights_path: codec.safetensors diff --git a/pocket_tts/default_parameters.py b/pocket_tts/default_parameters.py index 871e0aa..5027016 100644 --- a/pocket_tts/default_parameters.py +++ b/pocket_tts/default_parameters.py @@ -1,8 +1,9 @@ DEFAULT_AUDIO_PROMPT = "alba" -DEFAULT_VARIANT = "b6369a24" +DEFAULT_LANGUAGE = "english" DEFAULT_TEMPERATURE = 0.7 DEFAULT_LSD_DECODE_STEPS = 1 DEFAULT_NOISE_CLAMP = None DEFAULT_EOS_THRESHOLD = -4.0 DEFAULT_FRAMES_AFTER_EOS = None +# TODO: make this dynamic since english_2026-04 supports bigger chunks MAX_TOKEN_PER_CHUNK = 50 diff --git a/pocket_tts/main.py b/pocket_tts/main.py index 7c4a74d..f63694e 100644 --- a/pocket_tts/main.py +++ b/pocket_tts/main.py @@ -22,12 +22,11 @@ from pocket_tts.default_parameters import ( DEFAULT_LSD_DECODE_STEPS, DEFAULT_NOISE_CLAMP, DEFAULT_TEMPERATURE, - DEFAULT_VARIANT, MAX_TOKEN_PER_CHUNK, ) from pocket_tts.models.tts_model import TTSModel, export_model_state from pocket_tts.utils.logging_utils import enable_logging -from pocket_tts.utils.utils import PREDEFINED_VOICES, size_of_dict +from pocket_tts.utils.utils import _ORIGINS_OF_PREDEFINED_VOICES logger = logging.getLogger(__name__) @@ -42,7 +41,6 @@ cli_app = typer.Typer( # Global model instance tts_model: TTSModel | None = None -global_model_state = None web_app = FastAPI( title="Kyutai Pocket TTS API", description="Text-to-Speech generation API", version="1.0.0" @@ -124,12 +122,15 @@ def text_to_speech( Args: text: Text to convert to speech - voice_url: Optional voice URL (http://, https://, or hf://) + voice_url: Optional built-in voice name (e.g., "alba"), or voice URL (http://, https://, or hf://) voice_wav: Optional uploaded voice file (mutually exclusive with voice_url) """ if not text.strip(): raise HTTPException(status_code=400, detail="Text cannot be empty") + if voice_url is None and voice_wav is None: + voice_url = DEFAULT_AUDIO_PROMPT + if voice_url is not None and voice_wav is not None: raise HTTPException(status_code=400, detail="Cannot provide both voice_url and voice_wav") @@ -139,7 +140,7 @@ def text_to_speech( voice_url.startswith("http://") or voice_url.startswith("https://") or voice_url.startswith("hf://") - or voice_url in PREDEFINED_VOICES + or voice_url in _ORIGINS_OF_PREDEFINED_VOICES ): raise HTTPException( status_code=400, detail="voice_url must start with http://, https://, or hf://" @@ -161,8 +162,7 @@ def text_to_speech( finally: os.unlink(temp_file_path) else: - # Use default global model state - model_state = global_model_state + raise HTTPException(status_code=500, detail="This should never happen.") return StreamingResponse( generate_data_with_state(text, model_state), @@ -176,30 +176,33 @@ def text_to_speech( @cli_app.command() def serve( - voice: Annotated[ - str, typer.Option(help="Path to voice prompt audio file (voice to clone)") - ] = DEFAULT_AUDIO_PROMPT, host: Annotated[str, typer.Option(help="Host to bind to")] = "localhost", port: Annotated[int, typer.Option(help="Port to bind to")] = 8000, reload: Annotated[bool, typer.Option(help="Enable auto-reload")] = False, - config: Annotated[ - str, + language: Annotated[ + str | None, typer.Option( - help="Path to locally-saved model config .yaml file or model variant signature" + help="Language for the TTS model. " + "'english_2026-01', 'english_2026-04', 'english', 'french_24l', 'german_24l', 'portuguese', 'italian', 'spanish'." + " Incompatible with the config argument. Default is 'english', which is the same model as 'english_2026-04'.", + show_default=False, ), - ] = DEFAULT_VARIANT, + ] = None, + config: Annotated[ + str | None, + typer.Option( + help="Path to locally-saved model config .yaml file. " + "Incompatible with the language argument. If not provided, will use the default English model." + ), + ] = None, quantize: Annotated[ bool, typer.Option(help="Apply int8 quantization to reduce memory usage") ] = False, ): """Start the FastAPI server.""" - global tts_model, global_model_state - tts_model = TTSModel.load_model(config, quantize=quantize) - - # Pre-load the voice prompt - global_model_state = tts_model.get_state_for_audio_prompt(voice) - logger.info(f"The size of the model state is {size_of_dict(global_model_state) // 1e6} MB") + global tts_model + tts_model = TTSModel.load_model(language=language, config=config, quantize=quantize) uvicorn.run("pocket_tts.main:web_app", host=host, port=port, reload=reload) @@ -218,9 +221,28 @@ def generate( str, typer.Option(help="Path to audio conditioning file (voice to clone)") ] = DEFAULT_AUDIO_PROMPT, quiet: Annotated[bool, typer.Option("-q", "--quiet", help="Disable logging output")] = False, + language: Annotated[ + str | None, + typer.Option( + help=( + "Language for the TTS model. " + "'english_2026-01', 'english_2026-04', 'english', 'french_24l', 'spanish_24l'," + "'german_24l', 'portuguese_24l', 'italian_24l'." + " Incompatible with the config argument. Default is 'english', which is the same model as 'english_2026-04'. " + "The '24l' variants are bigger models, " + "not distilled yet and here only as preview. They're not the final " + "models for those languages." + ), + show_default=False, + ), + ] = None, config: Annotated[ - str, typer.Option(help="Model signature or path to config .yaml file") - ] = DEFAULT_VARIANT, + str | None, + typer.Option( + help="Path to locally-saved model config .yaml file. " + "Incompatible with the language argument. If not provided, will use the default English model." + ), + ] = None, lsd_decode_steps: Annotated[ int, typer.Option(help="Number of generation steps") ] = DEFAULT_LSD_DECODE_STEPS, @@ -254,7 +276,13 @@ def generate( logger.error("No input received from stdin.") raise typer.Exit(code=1) tts_model = TTSModel.load_model( - config, temperature, lsd_decode_steps, noise_clamp, eos_threshold, quantize=quantize + language=language, + config=config, + temp=temperature, + lsd_decode_steps=lsd_decode_steps, + noise_clamp=noise_clamp, + eos_threshold=eos_threshold, + quantize=quantize, ) tts_model.to(device) @@ -293,13 +321,33 @@ def export_voice( ], export_path: Annotated[str, typer.Argument(help="Output file or directory")], quiet: Annotated[bool, typer.Option("-q", "--quiet", help="Disable logging output")] = False, - config: Annotated[str, typer.Option(help="Model config path or signature")] = DEFAULT_VARIANT, + language: Annotated[ + str | None, + typer.Option( + help=( + "Language for the TTS model. " + "'english_2026-01', 'english_2026-04', 'english', 'french_24l', 'german_24l','spanish_24l'," + " 'portuguese_24l', 'italian_24l'." + " Incompatible with the config argument. Default is 'english', which is the same model as 'english_2026-04'. " + "The '24l' variants are bigger models, " + "not distilled yet and here only as preview." + ), + show_default=False, + ), + ] = None, + config: Annotated[ + str | None, + typer.Option( + help="Path to locally-saved model config .yaml file. " + "Incompatible with the language argument. If not provided, will use the default English model." + ), + ] = None, ): """Convert and save audio to .safetensors file""" log_level = logging.ERROR if quiet else logging.INFO with enable_logging("pocket_tts", log_level): - tts_model = TTSModel.load_model(config) + tts_model = TTSModel.load_model(language=language, config=config) model_state = tts_model.get_state_for_audio_prompt( audio_conditioning=audio_path, truncate=True ) diff --git a/pocket_tts/models/flow_lm.py b/pocket_tts/models/flow_lm.py index 6e0189c..07877de 100644 --- a/pocket_tts/models/flow_lm.py +++ b/pocket_tts/models/flow_lm.py @@ -65,6 +65,7 @@ class FlowLMModel(nn.Module): stats_ema_decay: float = 0.999, text_padding_weight: float = 1.0, dtype=None, + insert_bos_before_voice: bool = False, ): super().__init__() self.conditioner = conditioner @@ -78,6 +79,10 @@ class FlowLMModel(nn.Module): self.register_buffer("emb_std", torch.ones(ldim, dtype=dtype)) self.register_buffer("emb_mean", torch.zeros(ldim, dtype=dtype)) self.bos_emb = torch.nn.Parameter(torch.randn(ldim, dtype=dtype)) + self.insert_bos_before_voice = insert_bos_before_voice + if self.insert_bos_before_voice: + # Add BOS value that's to be inserted before the voice condition + self.bos_before_voice = torch.nn.Parameter(torch.randn((1, 1, self.dim), dtype=dtype)) self.input_linear = nn.Linear(self.ldim, dim, bias=False, dtype=dtype) self.transformer = transformer @@ -185,7 +190,9 @@ class FlowLMModel(nn.Module): return result @classmethod - def from_pydantic_config(cls, config: FlowLMConfig, latent_dim: int) -> Self: + def from_pydantic_config( + cls, config: FlowLMConfig, latent_dim: int, insert_bos_before_voice: bool + ) -> Self: d_model = config.transformer.d_model flow_mlp = SimpleMLPAdaLN.from_pydantic_config(config, latent_dim, d_model) @@ -205,4 +212,5 @@ class FlowLMModel(nn.Module): conditioner=conditioner, ldim=latent_dim, dtype=getattr(torch, config.dtype), + insert_bos_before_voice=insert_bos_before_voice, ) diff --git a/pocket_tts/models/mimi.py b/pocket_tts/models/mimi.py index 2251631..c1ef706 100644 --- a/pocket_tts/models/mimi.py +++ b/pocket_tts/models/mimi.py @@ -22,6 +22,8 @@ class MimiModel(nn.Module): encoder_frame_rate: float, sample_rate: int, channels: int, + inner_dim: int | None, + outer_dim: int | None, encoder_transformer: ProjectedTransformer, decoder_transformer: ProjectedTransformer, ): @@ -50,8 +52,14 @@ class MimiModel(nn.Module): assert downsample_stride == int(downsample_stride), ( f"Only integer strides are supported, got {downsample_stride}" ) - self.downsample = ConvDownsample1d(int(downsample_stride), dimension=dimension) - self.upsample = ConvTrUpsample1d(int(downsample_stride), dimension=dimension) + + # v2 of models + self.downsample = ConvDownsample1d( + int(downsample_stride), dimension=dimension, out_dimension=inner_dim + ) + self.upsample = ConvTrUpsample1d( + int(downsample_stride), dimension=dimension, in_dimension=outer_dim + ) @property def frame_size(self) -> int: diff --git a/pocket_tts/models/tts_model.py b/pocket_tts/models/tts_model.py index 1ce9237..fd46a71 100644 --- a/pocket_tts/models/tts_model.py +++ b/pocket_tts/models/tts_model.py @@ -11,6 +11,7 @@ from pathlib import Path import safetensors import safetensors.torch +import scipy.io.wavfile import torch from torch import nn from torch.nn import functional as F @@ -21,10 +22,10 @@ from pocket_tts.data.audio import audio_read from pocket_tts.data.audio_utils import convert_audio from pocket_tts.default_parameters import ( DEFAULT_EOS_THRESHOLD, + DEFAULT_LANGUAGE, DEFAULT_LSD_DECODE_STEPS, DEFAULT_NOISE_CLAMP, DEFAULT_TEMPERATURE, - DEFAULT_VARIANT, MAX_TOKEN_PER_CHUNK, ) from pocket_tts.models.flow_lm import FlowLMModel @@ -34,11 +35,13 @@ from pocket_tts.modules.dummy_quantizer import DummyQuantizer from pocket_tts.modules.seanet import SEANetDecoder, SEANetEncoder from pocket_tts.modules.stateful_module import StatefulModule, increment_steps, init_states from pocket_tts.quantization import RECOMMENDED_CONFIG, apply_dynamic_int8 -from pocket_tts.utils.config import Config, load_config +from pocket_tts.utils.config import CONFIGS_DIR, Config, load_config from pocket_tts.utils.utils import ( - PREDEFINED_VOICES, + _ORIGINS_OF_PREDEFINED_VOICES, + DEBUG_MIMI, display_execution_time, download_if_necessary, + get_predefined_voice, size_of_dict, ) from pocket_tts.utils.weights_loading import get_flow_lm_state_dict, get_mimi_state_dict @@ -49,7 +52,7 @@ logger = logging.getLogger(__name__) VOICE_CLONING_UNSUPPORTED = ( f"We could not download the weights for the model with voice cloning, " f"but you're trying to use voice cloning. " - f"Without voice cloning, you can use our catalog of voices {list(PREDEFINED_VOICES)}. " + f"Without voice cloning, you can use our catalog of voices {list(_ORIGINS_OF_PREDEFINED_VOICES.keys())}. " f"If you want access to the model with voice cloning, go to " f"https://huggingface.co/kyutai/pocket-tts and accept the terms, " f"then make sure you're logged in locally with `uvx hf auth login`." @@ -68,6 +71,10 @@ class TTSModel(nn.Module): noise_clamp: float | None, eos_threshold, config: Config, + origin: Path | None = None, + pad_with_spaces_for_short_inputs: bool = False, + model_recommended_frames_after_eos: int | None = None, + remove_semicolons: bool = False, ): super().__init__() self.flow_lm = flow_lm @@ -77,6 +84,10 @@ class TTSModel(nn.Module): self.eos_threshold = eos_threshold self.config = config self.has_voice_cloning = True + self.origin = origin + self.pad_with_spaces_for_short_inputs: bool = pad_with_spaces_for_short_inputs + self.model_recommended_frames_after_eos = model_recommended_frames_after_eos + self.remove_semicolons = remove_semicolons @property def device(self) -> str: @@ -88,23 +99,54 @@ class TTSModel(nn.Module): @classmethod def _from_pydantic_config( - cls, config: Config, temp, lsd_decode_steps, noise_clamp: float | None, eos_threshold + cls, + config: Config, + temp, + lsd_decode_steps, + noise_clamp: float | None, + eos_threshold, + origin: Path | None, ) -> Self: flow_lm = FlowLMModel.from_pydantic_config( - config.flow_lm, latent_dim=config.mimi.quantizer.dimension + config.flow_lm, + latent_dim=config.mimi.quantizer.dimension, + insert_bos_before_voice=config.flow_lm.insert_bos_before_voice, + ) + tts_model = cls( + flow_lm, + temp, + lsd_decode_steps, + noise_clamp, + eos_threshold, + config, + origin=origin, + pad_with_spaces_for_short_inputs=config.pad_with_spaces_for_short_inputs, + model_recommended_frames_after_eos=config.model_recommended_frames_after_eos, + remove_semicolons=config.remove_semicolons, ) - tts_model = cls(flow_lm, temp, lsd_decode_steps, noise_clamp, eos_threshold, config) return tts_model @classmethod def _from_pydantic_config_with_weights( - cls, config: Config, temp, lsd_decode_steps, noise_clamp: float | None, eos_threshold + cls, + config: Config, + temp, + lsd_decode_steps, + noise_clamp: float | None, + eos_threshold, + origin: Path | None = None, ) -> Self: tts_model = cls._from_pydantic_config( - config, temp, lsd_decode_steps, noise_clamp, eos_threshold + config, temp, lsd_decode_steps, noise_clamp, eos_threshold, origin=origin ) tts_model.flow_lm.speaker_proj_weight = torch.nn.Parameter( - torch.zeros((1024, 512), dtype=torch.float32) + torch.zeros( + ( + config.flow_lm.transformer.d_model, + config.mimi.inner_dim or config.mimi.seanet.dimension, + ), + dtype=torch.float32, + ) ) if config.flow_lm.weights_path is not None: if config.mimi.weights_path is None: @@ -137,6 +179,8 @@ class TTSModel(nn.Module): sample_rate=mimi_config["sample_rate"], frame_rate=mimi_config["frame_rate"], encoder_frame_rate=mimi_config["sample_rate"] / encoder.hop_length, + inner_dim=mimi_config["inner_dim"], + outer_dim=mimi_config["outer_dim"], encoder_transformer=encoder_transformer, decoder_transformer=decoder_transformer, ).to(device="cpu") @@ -153,11 +197,7 @@ class TTSModel(nn.Module): tts_model.mimi.load_state_dict(mimi_state, strict=True) tts_model.mimi.eval() - # tts_model.to(dtype=torch.float32) - # uncomment to save the weights - # tts_model = tts_model.to(dtype=torch.bfloat16) - # safetensors.torch.save_file(tts_model.state_dict(), "tts_b6369a24.safetensors") if config.weights_path is not None: logger.info(f"Loading TTSModel weights from {config.weights_path}") try: @@ -174,6 +214,10 @@ class TTSModel(nn.Module): "No weights_path specified for FlowLM or TTSModel, model is uninitialized!" ) size_in_mb = size_of_dict(tts_model.state_dict()) // 1e6 + if os.environ.get("POCKET_TTS_SAVE_WEIGHTS", "0") == "1": + save_path = "./model.safetensors" + safetensors.torch.save_file(tts_model.state_dict(), save_path) + logger.info(f"Saved TTSModel weights to {save_path}") logging.info(f"TTS Model loaded successfully. Its size is {size_in_mb} MB") # TODO: move this in the __init__ and make self.mimi in __init__ @@ -188,7 +232,8 @@ class TTSModel(nn.Module): @classmethod def load_model( cls, - config: str | Path = DEFAULT_VARIANT, + language: str | None = None, + config: str | Path | None = None, temp: float | int = DEFAULT_TEMPERATURE, lsd_decode_steps: int = DEFAULT_LSD_DECODE_STEPS, noise_clamp: float | int | None = DEFAULT_NOISE_CLAMP, @@ -202,8 +247,11 @@ class TTSModel(nn.Module): with the specified generation parameters and ready for inference. Args: - config: a path to a custom YAML config file saved locally (e.g., C://pocket_tts/pocket_tts_config.yaml) - or a model variant identifier (e.g., '610b0b2c'; must match a YAML file in the config directory). + language: Optional language identifier to select a predefined config. Incompatible with + the `config` argument. Available options + are `"english_2026-01"`, `"english_2026-04"`, `"english"`, `"french_24l"`, `"german_24l"`, `"portuguese"`, `"italian"`, `"spanish_24l"`. + If neither `config` nor `language` is provided, defaults to `"english", which is the same model as 'english_2026-04'`. + config: A path to a custom YAML config file saved locally (e.g., `"C://pocket_tts/pocket_tts_config.yaml"`). temp: Sampling temperature for generation. Higher values produce more diverse but potentially lower quality output. lsd_decode_steps: Number of steps for Lagrangian Self Distillation @@ -238,15 +286,23 @@ class TTSModel(nn.Module): model = TTSModel.load_model(quantize=True) ``` """ - if str(config).endswith(".yaml"): - config_path = Path(config) - config = load_config(config_path) - logger.info(f"Loading model from config at {config_path}...") - else: - config = load_config(Path(__file__).parents[1] / f"config/{config}.yaml") + if config is not None and language is not None: + raise ValueError( + "Cannot specify both config and language, please choose one or the other." + ) + if config is None and language is None: + language = DEFAULT_LANGUAGE + if language is not None: + config = CONFIGS_DIR / f"{language}.yaml" + config = Path(config) + if config.suffix not in (".yaml", ".yml"): + raise ValueError("Config should be a path to a YAML file ending with .yaml") + config_path = Path(config) + config = load_config(config_path) + logger.info(f"Loading model from config at {config_path}...") tts_model = TTSModel._from_pydantic_config_with_weights( - config, temp, lsd_decode_steps, noise_clamp, eos_threshold + config, temp, lsd_decode_steps, noise_clamp, eos_threshold, origin=config_path ) if quantize: @@ -306,8 +362,23 @@ class TTSModel(nn.Module): ) return output_embeddings[:, None, :], is_eos + def _decode_and_dump(self, encoded: torch.Tensor, filename: str): + mimi_state = init_states(self.mimi, batch_size=1, sequence_length=10000) + if encoded.shape[1] == self.mimi.quantizer.dimension: + latent_to_decode = self.mimi.quantizer(encoded) + else: + latent_to_decode = encoded + resored_audio = self.mimi.decode_from_latent(latent_to_decode, mimi_state) + scipy.io.wavfile.write(filename, self.sample_rate, resored_audio.numpy()) + logger.info("Saved restored audio from Mimi encoding to %s for debugging", filename) + def _encode_audio(self, audio: torch.Tensor) -> torch.Tensor: encoded = self.mimi.encode_to_latent(audio) + + if DEBUG_MIMI: + # sanity check + self._decode_and_dump(encoded, "debug_encoded_latent_decoded.wav") + latents = encoded.transpose(-1, -2).to(torch.float32) conditioning = F.linear(latents, self.flow_lm.speaker_proj_weight) return conditioning @@ -525,17 +596,25 @@ class TTSModel(nn.Module): and audio decoding. Generation performance is logged including real-time factor (RTF) metrics. """ + if frames_after_eos is None: + frames_after_eos = self.model_recommended_frames_after_eos # This is a very simplistic way of handling long texts. We could do much better # by using teacher forcing, but it would be a bit slower. # TODO: add the teacher forcing method for long texts where we use the audio of one chunk # as conditioning for the next chunk. chunks = split_into_best_sentences( - self.flow_lm.conditioner.tokenizer, text_to_generate, max_tokens + self.flow_lm.conditioner.tokenizer, + text_to_generate, + max_tokens, + self.pad_with_spaces_for_short_inputs, + remove_semicolons=self.remove_semicolons, ) for chunk in chunks: - text_to_generate, frames_after_eos_guess = prepare_text_prompt(chunk) + text_to_generate, frames_after_eos_guess = prepare_text_prompt( + chunk, self.pad_with_spaces_for_short_inputs, self.remove_semicolons + ) frames_after_eos_guess += 2 effective_frames = ( frames_after_eos if frames_after_eos is not None else frames_after_eos_guess @@ -767,9 +846,22 @@ class TTSModel(nn.Module): return _import_model_state(audio_conditioning) - elif isinstance(audio_conditioning, str) and audio_conditioning in PREDEFINED_VOICES: + elif ( + isinstance(audio_conditioning, str) + and audio_conditioning in _ORIGINS_OF_PREDEFINED_VOICES + ): # We get the audio conditioning directly from the safetensors file. - return _import_model_state(download_if_necessary(PREDEFINED_VOICES[audio_conditioning])) + if self.origin is None or not self.origin.is_relative_to(CONFIGS_DIR): + raise ValueError( + f"Cannot use predefined voices when the model " + f"is not loaded from a config associated with a language." + f"Here the origin is {self.origin}" + ) + return _import_model_state( + download_if_necessary( + get_predefined_voice(language=self.origin.stem, name=audio_conditioning) + ) + ) if not self.has_voice_cloning and isinstance(audio_conditioning, (str, Path)): raise ValueError(VOICE_CLONING_UNSUPPORTED) @@ -793,6 +885,9 @@ class TTSModel(nn.Module): with display_execution_time("Encoding audio prompt"): prompt = self._encode_audio(audio_conditioning.unsqueeze(0).to(self.device)) + if self.flow_lm.insert_bos_before_voice: + prompt = torch.cat([self.flow_lm.bos_before_voice, prompt], dim=1) + model_state = init_states(self.flow_lm, batch_size=1, sequence_length=prompt.shape[1]) with display_execution_time("Prompting audio"): @@ -810,11 +905,15 @@ class TTSModel(nn.Module): return math.ceil(gen_len_sec * frame_rate) -def prepare_text_prompt(text: str) -> tuple[str, int]: +def prepare_text_prompt( + text: str, pad_with_spaces_for_short_inputs: bool, remove_semicolons: bool +) -> tuple[str, int]: text = text.strip() if text == "": raise ValueError("Text prompt cannot be empty") text = text.replace("\n", " ").replace("\r", " ").replace(" ", " ") + if remove_semicolons: + text = text.replace(";", ",") number_of_words = len(text.split()) if number_of_words <= 4: frames_after_eos_guess = 3 @@ -832,7 +931,7 @@ def prepare_text_prompt(text: str) -> tuple[str, int]: # The model does not perform well when there are very few tokens, so # we can add empty spaces at the beginning to increase the token count. - if len(text.split()) < 5: + if pad_with_spaces_for_short_inputs and len(text.split()) < 5: text = " " * 8 + text return text, frames_after_eos_guess @@ -871,8 +970,16 @@ def _segments_from_boundaries( return segments -def split_into_best_sentences(tokenizer, text_to_generate: str, max_tokens: int) -> list[str]: - text_to_generate, _ = prepare_text_prompt(text_to_generate) +def split_into_best_sentences( + tokenizer, + text_to_generate: str, + max_tokens: int, + pad_with_spaces_for_short_inputs: bool, + remove_semicolons: bool, +) -> list[str]: + text_to_generate, _ = prepare_text_prompt( + text_to_generate, pad_with_spaces_for_short_inputs, remove_semicolons + ) text_to_generate = text_to_generate.strip() tokens = tokenizer(text_to_generate) list_of_tokens = tokens.tokens[0].tolist() diff --git a/pocket_tts/modules/resample.py b/pocket_tts/modules/resample.py index da1ae08..f0dea7c 100644 --- a/pocket_tts/modules/resample.py +++ b/pocket_tts/modules/resample.py @@ -10,11 +10,14 @@ class ConvDownsample1d(nn.Module): with a kernel size of twice the stride. """ - def __init__(self, stride: int, dimension: int): + def __init__(self, stride: int, dimension: int, out_dimension: int | None = None): super().__init__() + if out_dimension is None: + out_dimension = dimension + self.conv = StreamingConv1d( dimension, - dimension, + out_dimension, kernel_size=2 * stride, stride=stride, groups=1, @@ -31,10 +34,12 @@ class ConvTrUpsample1d(nn.Module): Upsample by some integer amount `stride` using transposed convolutions. """ - def __init__(self, stride: int, dimension: int): + def __init__(self, stride: int, dimension: int, in_dimension: int | None = None): super().__init__() + if in_dimension is None: + in_dimension = dimension self.convtr = StreamingConvTranspose1d( - dimension, + in_dimension, dimension, kernel_size=2 * stride, stride=stride, diff --git a/pocket_tts/utils/config.py b/pocket_tts/utils/config.py index 6ce1c6d..0a9625d 100644 --- a/pocket_tts/utils/config.py +++ b/pocket_tts/utils/config.py @@ -5,6 +5,8 @@ from pathlib import Path import yaml from pydantic import BaseModel, ConfigDict +CONFIGS_DIR = Path(__file__).parent.parent / "config" + class StrictModel(BaseModel): model_config = ConfigDict(extra="forbid") @@ -45,6 +47,7 @@ class FlowLMConfig(StrictModel): # conditioning lookup_table: LookupTable weights_path: str | None = None + insert_bos_before_voice: bool = False # SEANet configuration @@ -101,6 +104,8 @@ class MimiConfig(StrictModel): # Quantizer quantizer: QuantizerConfig weights_path: str | None = None + inner_dim: int | None = None + outer_dim: int | None = None class Config(StrictModel): @@ -108,13 +113,21 @@ class Config(StrictModel): mimi: MimiConfig weights_path: str | None = None weights_path_without_voice_cloning: str | None = None + pad_with_spaces_for_short_inputs: bool = False + remove_semicolons: bool = False + model_recommended_frames_after_eos: int | None = None def load_config(yaml_path: str | Path) -> Config: yaml_path = Path(yaml_path) if not yaml_path.exists(): - raise FileNotFoundError(f"Config file not found: {yaml_path}") + if yaml_path.is_relative_to(CONFIGS_DIR): + raise FileNotFoundError( + f"Config file not found: {yaml_path}. " + f"Did you make a typo? Available languages: {[p.stem for p in CONFIGS_DIR.glob('*.yaml')]}" + ) + raise FileNotFoundError(f"Config file not found: {yaml_path}. Did you make a typo?") with open(yaml_path, "r") as f: config_dict = yaml.safe_load(f) diff --git a/pocket_tts/utils/utils.py b/pocket_tts/utils/utils.py index cfcb80d..39636ee 100644 --- a/pocket_tts/utils/utils.py +++ b/pocket_tts/utils/utils.py @@ -1,15 +1,16 @@ import hashlib import logging +import os import time from pathlib import Path import requests -import safetensors.torch import torch from huggingface_hub import hf_hub_download from torch import nn PROJECT_ROOT = Path(__file__).parent.parent.parent +DEBUG_MIMI = os.environ.get("DEBUG_MIMI", "0") == "1" _ORIGINS_OF_PREDEFINED_VOICES = { "cosette": "hf://kyutai/tts-voices/expresso/ex04-ex02_confused_001_channel1_499s.wav", @@ -35,11 +36,9 @@ _ORIGINS_OF_PREDEFINED_VOICES = { "caro_davy": "hf://kyutai/tts-voices/voice-zero/caro_davy.wav", } -PREDEFINED_VOICES = { - # don't forget to change this - x: f"hf://kyutai/pocket-tts-without-voice-cloning/embeddings_v3/{x}.safetensors@075c0abfe7e41450521b0200b5168cfbc16bc77b" - for x in _ORIGINS_OF_PREDEFINED_VOICES -} + +def get_predefined_voice(language: str, name: str) -> str: + return f"hf://kyutai/pocket-tts-without-voice-cloning/languages/{language}/embeddings/{name}.safetensors@d29db7978e464fb90cb3359ee0c69a273b9142cc" def make_cache_directory() -> Path: @@ -113,14 +112,3 @@ def download_if_necessary(file_path: str) -> Path: return Path(cached_file) else: return Path(file_path) - - -def load_predefined_voice(voice_name: str) -> torch.Tensor: - if voice_name not in PREDEFINED_VOICES: - raise ValueError( - f"Predefined voice '{voice_name}' not found" - f", available voices are {list(PREDEFINED_VOICES)}." - ) - voice_file = download_if_necessary(PREDEFINED_VOICES[voice_name]) - # There is only one tensor in the file. - return safetensors.torch.load_file(voice_file)["audio_prompt"] diff --git a/pocket_tts/utils/weights_loading.py b/pocket_tts/utils/weights_loading.py index 44eb228..3f41bda 100644 --- a/pocket_tts/utils/weights_loading.py +++ b/pocket_tts/utils/weights_loading.py @@ -1,6 +1,7 @@ from pathlib import Path import safetensors +import torch def get_flow_lm_state_dict(path: Path) -> dict: @@ -11,6 +12,7 @@ def get_flow_lm_state_dict(path: Path) -> dict: key.startswith("flow.w_s_t.") or key == "condition_provider.conditioners.transcript_in_segment.learnt_padding" or key == "condition_provider.conditioners.speaker_wavs.learnt_padding" + or key == "num_ema_updates" ): # skip lookup table weights continue @@ -19,6 +21,11 @@ def get_flow_lm_state_dict(path: Path) -> dict: new_name = "conditioner.embed.weight" if key == "condition_provider.conditioners.speaker_wavs.output_proj.weight": new_name = "speaker_proj_weight" + if key == "fuser.padding_value": + new_name = "bos_before_voice" + + new_name = new_name.replace(".self_attn.in_proj_weight", ".self_attn.in_proj.weight") + state_dict[new_name] = f.get_tensor(key) return state_dict @@ -27,9 +34,46 @@ def get_mimi_state_dict(path: Path) -> dict: state_dict = {} with safetensors.safe_open(path, framework="pt", device="cpu") as f: for key in f.keys(): - if key.startswith("model.quantizer.vq.") or key == "model.quantizer.logvar_proj.weight": + if ( + key.startswith("model.quantizer.vq.") + or key == "model.quantizer.logvar_proj.weight" + or "_codebook" in key + or key.endswith(".weight_v") + or key == "quantizer.logvar_proj.weight" # this is new + ): # skip vq weights continue - state_dict[key.removeprefix("model.")] = f.get_tensor(key) + # weight = weight_g * torch.nn.functional.normalize(weight_v, dim=0) + if key.endswith(".weight_g"): + key_g = key + + new_key = key.removesuffix("_g") + key_v = new_key + "_v" + weight_v = f.get_tensor(key_v) + weight_g = f.get_tensor(key_g) + + new_key = new_key.replace(".conv.conv.", ".conv.").replace( + ".convtr.convtr.", ".convtr." + ) + state_dict[new_key] = torch._weight_norm(weight_v, weight_g, dim=0) + continue + + if key in [ + "wavlm_emb_downsample.conv.conv.weight", + "wavlm_input_resample.kernel", + "wavlm_proj.weight", + "quantizer.logvar_param", + ]: + continue + + if "wavlm_emb_downsample" in key: + continue + + state_dict[ + key.removeprefix("model.") + .replace(".conv.conv.", ".conv.") + .replace(".convtr.convtr.", ".convtr.") + .replace("in_proj_weight", "in_proj.weight") + ] = f.get_tensor(key) return state_dict diff --git a/scripts/generate_default_voices.py b/scripts/generate_default_voices.py index 2dbaac3..b145c67 100644 --- a/scripts/generate_default_voices.py +++ b/scripts/generate_default_voices.py @@ -1,9 +1,11 @@ +import sys + import scipy.io.wavfile from pocket_tts import TTSModel, export_model_state from pocket_tts.utils.utils import _ORIGINS_OF_PREDEFINED_VOICES -model = TTSModel.load_model() +model = TTSModel.load_model(language=sys.argv[1]) for voice_name, voice_origin in _ORIGINS_OF_PREDEFINED_VOICES.items(): print(f"Processing voice: {voice_name} from origin: {voice_origin}") @@ -16,7 +18,8 @@ for voice_name, voice_origin in _ORIGINS_OF_PREDEFINED_VOICES.items(): ) audio = model.generate_audio( - model_state_copy, "Hello, it's a good day, isn't it? I hope you are doing well." + model_state_copy, + "An Sommernachmittagen spielten die Kinder fröhlich auf dem Platz der kleinen Stadt.", ) scipy.io.wavfile.write( f"./built-in-voices-generated/{voice_name}.wav", model.sample_rate, audio.numpy() diff --git a/swarm-config.yaml b/swarm-config.yaml new file mode 100644 index 0000000..4cf7300 --- /dev/null +++ b/swarm-config.yaml @@ -0,0 +1,89 @@ +services: + traefik: + image: traefik:v3.6.7 + command: + # Swarm provider configuration + - "--providers.swarm.endpoint=unix:///var/run/docker.sock" + - "--providers.swarm.exposedByDefault=false" + + # EntryPoints for HTTP and HTTPS + - "--entrypoints.web.address=:80" + - "--entrypoints.websecure.address=:443" + + # Redirect HTTP to HTTPS + - "--entrypoints.web.http.redirections.entryPoint.to=websecure" + - "--entrypoints.web.http.redirections.entryPoint.scheme=https" + + # Let's Encrypt configuration + - "--certificatesResolvers.letsencrypt.acme.httpChallenge.entryPoint=web" + - "--certificatesResolvers.letsencrypt.acme.storage=/letsencrypt/acme.json" + - "--certificatesResolvers.letsencrypt.acme.email=gabriel@kyutai.org" + - "--certificatesResolvers.letsencrypt.acme.httpChallenge=true" + - "--log.level=TRACE" + ports: + - "80:80" + - "443:443" + volumes: + - "/var/run/docker.sock:/var/run/docker.sock:ro" + - letsencrypt:/letsencrypt + deploy: + placement: + constraints: + - node.role == manager + + backend-english: + image: rg.fr-par.scw.cloud/namespace-unruffled-tereshkova/pocket-tts-server-english:${COMMIT_SHA} + build: + context: . + volumes: + - hf-cache:/root/.cache/huggingface + deploy: + labels: + - "traefik.enable=true" + - "traefik.http.routers.server-english.rule=Host(`pocket-tts-api.kyutai.org`) && PathPrefix(`/english`)" + - "traefik.http.routers.server-english.middlewares=strip-api-english" + - "traefik.http.middlewares.strip-api-english.replacepathregex.regex=^/english/(.*)" + - "traefik.http.middlewares.strip-api-english.replacepathregex.replacement=/$$1" + - "traefik.http.routers.server-english.entrypoints=websecure" + - "traefik.http.routers.server-english.tls=true" + - "traefik.http.routers.server-english.tls.certresolver=letsencrypt" + - "traefik.http.services.server-english.loadbalancer.server.port=80" + - "traefik.http.routers.server-english.priority=100" + + backend-german: + image: rg.fr-par.scw.cloud/namespace-unruffled-tereshkova/pocket-tts-server-german:${COMMIT_SHA} + build: + context: . + volumes: + - hf-cache:/root/.cache/huggingface + deploy: + labels: + - "traefik.enable=true" + - "traefik.http.routers.server-german.rule=Host(`pocket-tts-api.kyutai.org`) && PathPrefix(`/german`)" + - "traefik.http.routers.server-german.middlewares=strip-api-german" + - "traefik.http.middlewares.strip-api-german.replacepathregex.regex=^/german/(.*)" + - "traefik.http.middlewares.strip-api-german.replacepathregex.replacement=/$$1" + - "traefik.http.routers.server-german.entrypoints=websecure" + - "traefik.http.routers.server-german.tls=true" + - "traefik.http.routers.server-german.tls.certresolver=letsencrypt" + - "traefik.http.services.server-german.loadbalancer.server.port=80" + - "traefik.http.routers.server-german.priority=100" + + +volumes: + letsencrypt: + portainer_data: + prometheus-data: + grafana-data: + +networks: + agent_network: + driver: overlay + attachable: true + driver_opts: + encrypted: "true" + default: + driver: overlay + attachable: true + driver_opts: + encrypted: "true" \ No newline at end of file diff --git a/tests/test_cli_generate.py b/tests/test_cli_generate.py index bc0bc68..ce8ba25 100644 --- a/tests/test_cli_generate.py +++ b/tests/test_cli_generate.py @@ -6,7 +6,6 @@ import pytest from typer.testing import CliRunner from pocket_tts.data.audio import audio_read -from pocket_tts.default_parameters import DEFAULT_VARIANT from pocket_tts.main import cli_app other_voice = "https://huggingface.co/kyutai/tts-voices/resolve/main/expresso/ex01-ex02_default_001_channel1_168s.wav" @@ -76,8 +75,8 @@ def test_generate_with_custom_parameters(tmp_path): "generate", "--text", "Testing custom parameters.", - "--config", - DEFAULT_VARIANT, + "--language", + "english", "--temperature", "0.8", "--lsd-decode-steps", diff --git a/tests/test_documentation_examples.py b/tests/test_documentation_examples.py index 723d9ae..acaabce 100644 --- a/tests/test_documentation_examples.py +++ b/tests/test_documentation_examples.py @@ -41,7 +41,9 @@ def test_load_model(): model = TTSModel.load_model() # Load with custom parameters - model = TTSModel.load_model(config="b6369a24", temp=0.5, lsd_decode_steps=5, eos_threshold=-3.0) + model = TTSModel.load_model( + language="english_2026-01", temp=0.5, lsd_decode_steps=5, eos_threshold=-3.0 + ) def test_device(): diff --git a/tests/test_split_sentences.py b/tests/test_split_sentences.py index 530208e..71fee22 100644 --- a/tests/test_split_sentences.py +++ b/tests/test_split_sentences.py @@ -13,14 +13,22 @@ def tokenizer(): def test_short_text_single_chunk(tokenizer): """Short text should produce a single chunk.""" - chunks = split_into_best_sentences(tokenizer, "Hello world.", 50) + chunks = split_into_best_sentences( + tokenizer, + "Hello world.", + 50, + pad_with_spaces_for_short_inputs=False, + remove_semicolons=False, + ) assert len(chunks) == 1 def test_multiple_sentences_split(tokenizer): """Multiple sentences should be split when they exceed max_tokens.""" text = "First sentence here. Second sentence here. Third sentence here. Fourth sentence here." - chunks = split_into_best_sentences(tokenizer, text, 10) + chunks = split_into_best_sentences( + tokenizer, text, 10, pad_with_spaces_for_short_inputs=False, remove_semicolons=False + ) assert len(chunks) > 1 @@ -34,7 +42,9 @@ def test_long_sentence_with_commas_is_split(tokenizer): "it was the season of Light, it was the season of Darkness, " "it was the spring of hope, it was the winter of despair" ) - chunks = split_into_best_sentences(tokenizer, text, 50) + chunks = split_into_best_sentences( + tokenizer, text, 50, pad_with_spaces_for_short_inputs=False, remove_semicolons=False + ) assert len(chunks) > 1, "Long comma-separated text should be split into multiple chunks" # Verify all content is preserved (no words should be lost in splitting) @@ -51,7 +61,9 @@ def test_long_sentence_with_commas_respects_max_tokens(tokenizer): "it was the epoch of belief, it was the epoch of incredulity" ) max_tokens = 20 - chunks = split_into_best_sentences(tokenizer, text, max_tokens) + chunks = split_into_best_sentences( + tokenizer, text, max_tokens, pad_with_spaces_for_short_inputs=False, remove_semicolons=False + ) for chunk in chunks: token_count = len(tokenizer(chunk.strip()).tokens[0].tolist()) # Allow some tolerance since comma clauses may vary in size @@ -69,14 +81,18 @@ def test_mixed_sentences_and_commas(tokenizer): "until it finally reaches a period. " "Another short one." ) - chunks = split_into_best_sentences(tokenizer, text, 20) + chunks = split_into_best_sentences( + tokenizer, text, 20, pad_with_spaces_for_short_inputs=False, remove_semicolons=False + ) assert len(chunks) >= 3 def test_no_commas_no_periods_stays_single_chunk(tokenizer): """Text with no splitting characters stays as a single chunk.""" text = "one two three four five six seven eight nine ten eleven twelve" - chunks = split_into_best_sentences(tokenizer, text, 5) + chunks = split_into_best_sentences( + tokenizer, text, 5, pad_with_spaces_for_short_inputs=False, remove_semicolons=False + ) # Should be 1 chunk since there are no split points assert len(chunks) == 1 @@ -87,14 +103,18 @@ def test_semicolons_and_colons_also_split(tokenizer): "First clause here; second clause here; third clause here; " "fourth clause here; fifth clause here; sixth clause here" ) - chunks = split_into_best_sentences(tokenizer, text, 15) + chunks = split_into_best_sentences( + tokenizer, text, 15, pad_with_spaces_for_short_inputs=False, remove_semicolons=False + ) assert len(chunks) > 1 def test_short_sentence_not_affected_by_comma_splitting(tokenizer): """Sentences under max_tokens should not be affected by comma logic.""" text = "Hello, world." - chunks = split_into_best_sentences(tokenizer, text, 50) + chunks = split_into_best_sentences( + tokenizer, text, 50, pad_with_spaces_for_short_inputs=False, remove_semicolons=False + ) assert len(chunks) == 1 assert "hello" in chunks[0].lower() assert "world" in chunks[0].lower() @@ -103,14 +123,18 @@ def test_short_sentence_not_affected_by_comma_splitting(tokenizer): def test_empty_string_raises(tokenizer): """Empty input should raise ValueError from prepare_text_prompt.""" with pytest.raises(ValueError, match="empty"): - split_into_best_sentences(tokenizer, "", 50) + split_into_best_sentences( + tokenizer, "", 50, pad_with_spaces_for_short_inputs=False, remove_semicolons=False + ) def test_oversized_clause_without_commas_still_returns(tokenizer): """A long clause with no split points should still be returned (not dropped).""" # 20 words with no punctuation at all - no way to split text = " ".join(f"word{i}" for i in range(20)) - chunks = split_into_best_sentences(tokenizer, text, 5) + chunks = split_into_best_sentences( + tokenizer, text, 5, pad_with_spaces_for_short_inputs=False, remove_semicolons=False + ) assert len(chunks) == 1 # prepare_text_prompt capitalizes the first char and adds a trailing period, # so compare case-insensitively and strip punctuation