valenti/spark-pocket-tts
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add french, italian, portuguese, spanish, german (#155)

Browse Source 
The french 6 layers variant wasn't included because it's not good enough
yet. Use the `french_24l` variant for now.
This commit is contained in:
Gabriel de Marmiesse
2026-04-17 17:52:16 +02:00
committed by GitHub
parent 9fc89c64f2
commit c90fc8ced6
31 changed files with 1165 additions and 131 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+23 -8
View File
@@ -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 * Uses only 2 CPU cores
* Python API and CLI * Python API and CLI
* Voice cloning * Voice cloning
* English only at the moment * Multi-language support: english, french, german, portuguese, italian, spanish
* Can handle infinitely long text inputs * Can handle infinitely long text inputs
* [Can run on client-side in the browser](#in-browser-implementations) * [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 ## Trying it from the website, without installing anything
@@ -48,18 +48,33 @@ uvx pocket-tts generate
pocket-tts generate pocket-tts generate
``` ```
Modify the voice with `--voice` and the text with `--text`. We provide a small catalog of voices. 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 You can take a look at [this page](https://huggingface.co/kyutai/tts-voices) which details the licenses
for each voice. for each voice.
* [alba](https://huggingface.co/kyutai/tts-voices/blob/main/alba-mackenna/casual.wav) * [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) * [anna](https://huggingface.co/kyutai/tts-voices/blob/main/vctk/p228_023_enhanced.wav)
* [javert](https://huggingface.co/kyutai/tts-voices/blob/main/voice-donations/Butter.wav) * [azelma](https://huggingface.co/kyutai/tts-voices/blob/main/vctk/p303_023_enhanced.wav)
* [jean](https://huggingface.co/kyutai/tts-voices/blob/main/ears/p010/freeform_speech_01.wav) * [bill_boerst](https://huggingface.co/kyutai/tts-voices/blob/main/voice-zero/bill_boerst.wav)
* [fantine](https://huggingface.co/kyutai/tts-voices/blob/main/vctk/p244_023.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) * [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) * [eponine](https://huggingface.co/kyutai/tts-voices/blob/main/vctk/p262_023_enhanced.wav)
* [azelma](https://huggingface.co/kyutai/tts-voices/blob/main/vctk/p303_023.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. 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). You can use your own or check out our [voice repository](https://huggingface.co/kyutai/tts-voices).
docs/API Reference/python-api.md
+7 -3
View File
@@ -37,16 +37,18 @@ The main class for text-to-speech generation.
#### Class Methods #### 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. Load and return a TTSModel instance with pre-trained weights.
**Parameters:** **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) - `temp` (float): Sampling temperature for generation (default: 0.7)
- `lsd_decode_steps` (int): Number of generation steps (default: 1) - `lsd_decode_steps` (int): Number of generation steps (default: 1)
- `noise_clamp` (float | None): Maximum value for noise sampling (default: None) - `noise_clamp` (float | None): Maximum value for noise sampling (default: None)
- `eos_threshold` (float): Threshold for end-of-sequence detection (default: -4.0) - `eos_threshold` (float): Threshold for end-of-sequence detection (default: -4.0)
- `quantize` (bool): Enable int8 quantization when loading the model (default: `False`)
**Returns:** **Returns:**
- `TTSModel`: Loaded model instance on CPU - `TTSModel`: Loaded model instance on CPU
@@ -59,7 +61,9 @@ from pocket_tts import TTSModel
model = TTSModel.load_model() model = TTSModel.load_model()
# Load with custom parameters # 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 #### Properties
docs/CLI Commands/export_voice.md
+2 -1
View File
@@ -24,7 +24,8 @@ Only the first 30 seconds of the audio file will be processed.
### Options ### Options
- `--quiet`: Do not print any output except errors. - `--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 ## Examples
docs/CLI Commands/generate.md
+4 -2
View File
@@ -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.") - `--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. - `--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") - `--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 ### 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) - `--lsd-decode-steps LSD_DECODE_STEPS`: Number of generation steps (default: 1)
- `--temperature TEMPERATURE`: Temperature for generation (default: 0.7) - `--temperature TEMPERATURE`: Temperature for generation (default: 0.7)
- `--noise-clamp NOISE_CLAMP`: Noise clamp value (default: None) - `--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 ### 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) - `--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 - `--quiet`, `-q`: Disable logging output
## Examples ## 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. 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. Then, use the --config option to point to your newly created config.
docs/CLI Commands/serve.md
+7 -12
View File
@@ -14,12 +14,12 @@ This starts a server on `http://localhost:8000` with the default voice model.
## Command Options ## 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") - `--host HOST`: Host to bind to (default: "localhost")
- `--port PORT`: Port to bind to (default: 8000) - `--port PORT`: Port to bind to (default: 8000)
- `--reload`: Enable auto-reload for development - `--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 ## Examples
### Basic Server ### Basic Server
@@ -32,21 +32,17 @@ pocket-tts serve
pocket-tts serve --host "localhost" --port 8080 pocket-tts serve --host "localhost" --port 8080
``` ```
### Custom Voice ### Custom Language
To select the default language model, pass `--language`:
```bash ```bash
# Use different voice pocket-tts serve --language french_24l
pocket-tts serve --voice "hf://kyutai/tts-voices/jessica-jian/casual.wav"
# Use local voice file
pocket-tts serve --voice "./my_voice.wav"
``` ```
### Custom Model Config ### Custom Model Config
If you'd like to override the paths from which the models are loaded, you can provide a custom YAML configuration. 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. 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. 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. For more advanced usage, see the [Python API documentation](python-api.md) for direct integration with the TTS model.
pocket_tts/config/english.yaml
+61
View File
@@ -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
pocket_tts/config/b6369a24.yaml → pocket_tts/config/english_2026-01.yaml
+8 -3
View File
@@ -1,9 +1,12 @@
# sig: b6369a24 # sig: b6369a24
weights_path: hf://kyutai/pocket-tts/tts_b6369a24.safetensors@427e3d61b276ed69fdd03de0d185fa8a8d97fc5b 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/tts_b6369a24.safetensors@d4fdd22ae8c8e1cb3634e150ebeff1dab2d16df3 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: flow_lm:
insert_bos_before_voice: false
dtype: float32 dtype: float32
flow: flow:
depth: 6 depth: 6
@@ -18,12 +21,14 @@ flow_lm:
dim: 1024 dim: 1024
n_bins: 4000 n_bins: 4000
tokenizer: sentencepiece 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 #weights_path: flow_lm_b6369a24.safetensors
mimi: mimi:
dtype: float32 dtype: float32
sample_rate: 24000 sample_rate: 24000
inner_dim: 512
outer_dim: 512
channels: 1 channels: 1
frame_rate: 12.5 frame_rate: 12.5
seanet: seanet:
pocket_tts/config/english_2026-04.yaml
+61
View File
@@ -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
pocket_tts/config/french_24l.yaml
+63
View File
@@ -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
pocket_tts/config/german.yaml
+61
View File
@@ -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
pocket_tts/config/german_24l.yaml
+61
View File
@@ -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
pocket_tts/config/italian.yaml
+61
View File
@@ -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
pocket_tts/config/italian_24l.yaml
+61
View File
@@ -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
pocket_tts/config/portuguese.yaml
+61
View File
@@ -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
pocket_tts/config/portuguese_24l.yaml
+61
View File
@@ -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
pocket_tts/config/spanish.yaml
+61
View File
@@ -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
pocket_tts/config/spanish_24l.yaml
+61
View File
@@ -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
pocket_tts/default_parameters.py
+2 -1
View File
@@ -1,8 +1,9 @@
DEFAULT_AUDIO_PROMPT = "alba" DEFAULT_AUDIO_PROMPT = "alba"
DEFAULT_VARIANT = "b6369a24" DEFAULT_LANGUAGE = "english"
DEFAULT_TEMPERATURE = 0.7 DEFAULT_TEMPERATURE = 0.7
DEFAULT_LSD_DECODE_STEPS = 1 DEFAULT_LSD_DECODE_STEPS = 1
DEFAULT_NOISE_CLAMP = None DEFAULT_NOISE_CLAMP = None
DEFAULT_EOS_THRESHOLD = -4.0 DEFAULT_EOS_THRESHOLD = -4.0
DEFAULT_FRAMES_AFTER_EOS = None DEFAULT_FRAMES_AFTER_EOS = None
# TODO: make this dynamic since english_2026-04 supports bigger chunks
MAX_TOKEN_PER_CHUNK = 50 MAX_TOKEN_PER_CHUNK = 50
pocket_tts/main.py
+73 -25
View File
@@ -22,12 +22,11 @@ from pocket_tts.default_parameters import (
DEFAULT_LSD_DECODE_STEPS, DEFAULT_LSD_DECODE_STEPS,
DEFAULT_NOISE_CLAMP, DEFAULT_NOISE_CLAMP,
DEFAULT_TEMPERATURE, DEFAULT_TEMPERATURE,
DEFAULT_VARIANT,
MAX_TOKEN_PER_CHUNK, MAX_TOKEN_PER_CHUNK,
) )
from pocket_tts.models.tts_model import TTSModel, export_model_state from pocket_tts.models.tts_model import TTSModel, export_model_state
from pocket_tts.utils.logging_utils import enable_logging 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__) logger = logging.getLogger(__name__)
@@ -42,7 +41,6 @@ cli_app = typer.Typer(
# Global model instance # Global model instance
tts_model: TTSModel | None = None tts_model: TTSModel | None = None
global_model_state = None
web_app = FastAPI( web_app = FastAPI(
title="Kyutai Pocket TTS API", description="Text-to-Speech generation API", version="1.0.0" title="Kyutai Pocket TTS API", description="Text-to-Speech generation API", version="1.0.0"
@@ -124,12 +122,15 @@ def text_to_speech(
Args: Args:
text: Text to convert to speech 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) voice_wav: Optional uploaded voice file (mutually exclusive with voice_url)
""" """
if not text.strip(): if not text.strip():
raise HTTPException(status_code=400, detail="Text cannot be empty") 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: 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") 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://") voice_url.startswith("http://")
or voice_url.startswith("https://") or voice_url.startswith("https://")
or voice_url.startswith("hf://") or voice_url.startswith("hf://")
or voice_url in PREDEFINED_VOICES or voice_url in _ORIGINS_OF_PREDEFINED_VOICES
): ):
raise HTTPException( raise HTTPException(
status_code=400, detail="voice_url must start with http://, https://, or hf://" status_code=400, detail="voice_url must start with http://, https://, or hf://"
@@ -161,8 +162,7 @@ def text_to_speech(
finally: finally:
os.unlink(temp_file_path) os.unlink(temp_file_path)
else: else:
# Use default global model state raise HTTPException(status_code=500, detail="This should never happen.")
model_state = global_model_state
return StreamingResponse( return StreamingResponse(
generate_data_with_state(text, model_state), generate_data_with_state(text, model_state),
@@ -176,30 +176,33 @@ def text_to_speech(
@cli_app.command() @cli_app.command()
def serve( 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", host: Annotated[str, typer.Option(help="Host to bind to")] = "localhost",
port: Annotated[int, typer.Option(help="Port to bind to")] = 8000, port: Annotated[int, typer.Option(help="Port to bind to")] = 8000,
reload: Annotated[bool, typer.Option(help="Enable auto-reload")] = False, reload: Annotated[bool, typer.Option(help="Enable auto-reload")] = False,
config: Annotated[ language: Annotated[
str, str | None,
typer.Option( 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[ quantize: Annotated[
bool, typer.Option(help="Apply int8 quantization to reduce memory usage") bool, typer.Option(help="Apply int8 quantization to reduce memory usage")
] = False, ] = False,
): ):
"""Start the FastAPI server.""" """Start the FastAPI server."""
global tts_model, global_model_state global tts_model
tts_model = TTSModel.load_model(config, quantize=quantize) tts_model = TTSModel.load_model(language=language, config=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")
uvicorn.run("pocket_tts.main:web_app", host=host, port=port, reload=reload) 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)") str, typer.Option(help="Path to audio conditioning file (voice to clone)")
] = DEFAULT_AUDIO_PROMPT, ] = DEFAULT_AUDIO_PROMPT,
quiet: Annotated[bool, typer.Option("-q", "--quiet", help="Disable logging output")] = False, 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[ config: Annotated[
str, typer.Option(help="Model signature or path to config .yaml file") str | None,
] = DEFAULT_VARIANT, 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[ lsd_decode_steps: Annotated[
int, typer.Option(help="Number of generation steps") int, typer.Option(help="Number of generation steps")
] = DEFAULT_LSD_DECODE_STEPS, ] = DEFAULT_LSD_DECODE_STEPS,
@@ -254,7 +276,13 @@ def generate(
logger.error("No input received from stdin.") logger.error("No input received from stdin.")
raise typer.Exit(code=1) raise typer.Exit(code=1)
tts_model = TTSModel.load_model( 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) tts_model.to(device)
@@ -293,13 +321,33 @@ def export_voice(
], ],
export_path: Annotated[str, typer.Argument(help="Output file or directory")], export_path: Annotated[str, typer.Argument(help="Output file or directory")],
quiet: Annotated[bool, typer.Option("-q", "--quiet", help="Disable logging output")] = False, 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""" """Convert and save audio to .safetensors file"""
log_level = logging.ERROR if quiet else logging.INFO log_level = logging.ERROR if quiet else logging.INFO
with enable_logging("pocket_tts", log_level): 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( model_state = tts_model.get_state_for_audio_prompt(
audio_conditioning=audio_path, truncate=True audio_conditioning=audio_path, truncate=True
) )
pocket_tts/models/flow_lm.py
+9 -1
View File
@@ -65,6 +65,7 @@ class FlowLMModel(nn.Module):
stats_ema_decay: float = 0.999, stats_ema_decay: float = 0.999,
text_padding_weight: float = 1.0, text_padding_weight: float = 1.0,
dtype=None, dtype=None,
insert_bos_before_voice: bool = False,
): ):
super().__init__() super().__init__()
self.conditioner = conditioner 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_std", torch.ones(ldim, dtype=dtype))
self.register_buffer("emb_mean", torch.zeros(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.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.input_linear = nn.Linear(self.ldim, dim, bias=False, dtype=dtype)
self.transformer = transformer self.transformer = transformer
@@ -185,7 +190,9 @@ class FlowLMModel(nn.Module):
return result return result
@classmethod @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 d_model = config.transformer.d_model
flow_mlp = SimpleMLPAdaLN.from_pydantic_config(config, latent_dim, d_model) flow_mlp = SimpleMLPAdaLN.from_pydantic_config(config, latent_dim, d_model)
@@ -205,4 +212,5 @@ class FlowLMModel(nn.Module):
conditioner=conditioner, conditioner=conditioner,
ldim=latent_dim, ldim=latent_dim,
dtype=getattr(torch, config.dtype), dtype=getattr(torch, config.dtype),
insert_bos_before_voice=insert_bos_before_voice,
) )
pocket_tts/models/mimi.py
+10 -2
View File
@@ -22,6 +22,8 @@ class MimiModel(nn.Module):
encoder_frame_rate: float, encoder_frame_rate: float,
sample_rate: int, sample_rate: int,
channels: int, channels: int,
inner_dim: int | None,
outer_dim: int | None,
encoder_transformer: ProjectedTransformer, encoder_transformer: ProjectedTransformer,
decoder_transformer: ProjectedTransformer, decoder_transformer: ProjectedTransformer,
): ):
@@ -50,8 +52,14 @@ class MimiModel(nn.Module):
assert downsample_stride == int(downsample_stride), ( assert downsample_stride == int(downsample_stride), (
f"Only integer strides are supported, got {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 @property
def frame_size(self) -> int: def frame_size(self) -> int:
pocket_tts/models/tts_model.py
+136 -29
View File
@@ -11,6 +11,7 @@ from pathlib import Path
import safetensors import safetensors
import safetensors.torch import safetensors.torch
import scipy.io.wavfile
import torch import torch
from torch import nn from torch import nn
from torch.nn import functional as F 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.data.audio_utils import convert_audio
from pocket_tts.default_parameters import ( from pocket_tts.default_parameters import (
DEFAULT_EOS_THRESHOLD, DEFAULT_EOS_THRESHOLD,
DEFAULT_LANGUAGE,
DEFAULT_LSD_DECODE_STEPS, DEFAULT_LSD_DECODE_STEPS,
DEFAULT_NOISE_CLAMP, DEFAULT_NOISE_CLAMP,
DEFAULT_TEMPERATURE, DEFAULT_TEMPERATURE,
DEFAULT_VARIANT,
MAX_TOKEN_PER_CHUNK, MAX_TOKEN_PER_CHUNK,
) )
from pocket_tts.models.flow_lm import FlowLMModel 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.seanet import SEANetDecoder, SEANetEncoder
from pocket_tts.modules.stateful_module import StatefulModule, increment_steps, init_states 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.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 ( from pocket_tts.utils.utils import (
PREDEFINED_VOICES, _ORIGINS_OF_PREDEFINED_VOICES,
DEBUG_MIMI,
display_execution_time, display_execution_time,
download_if_necessary, download_if_necessary,
get_predefined_voice,
size_of_dict, size_of_dict,
) )
from pocket_tts.utils.weights_loading import get_flow_lm_state_dict, get_mimi_state_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 = ( VOICE_CLONING_UNSUPPORTED = (
f"We could not download the weights for the model with voice cloning, " f"We could not download the weights for the model with voice cloning, "
f"but you're trying to use 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"If you want access to the model with voice cloning, go to "
f"https://huggingface.co/kyutai/pocket-tts and accept the terms, " 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`." 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, noise_clamp: float | None,
eos_threshold, eos_threshold,
config: Config, 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__() super().__init__()
self.flow_lm = flow_lm self.flow_lm = flow_lm
@@ -77,6 +84,10 @@ class TTSModel(nn.Module):
self.eos_threshold = eos_threshold self.eos_threshold = eos_threshold
self.config = config self.config = config
self.has_voice_cloning = True 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 @property
def device(self) -> str: def device(self) -> str:
@@ -88,23 +99,54 @@ class TTSModel(nn.Module):
@classmethod @classmethod
def _from_pydantic_config( 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: ) -> Self:
flow_lm = FlowLMModel.from_pydantic_config( 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 return tts_model
@classmethod @classmethod
def _from_pydantic_config_with_weights( 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: ) -> Self:
tts_model = cls._from_pydantic_config( 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( 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.flow_lm.weights_path is not None:
if config.mimi.weights_path is None: if config.mimi.weights_path is None:
@@ -137,6 +179,8 @@ class TTSModel(nn.Module):
sample_rate=mimi_config["sample_rate"], sample_rate=mimi_config["sample_rate"],
frame_rate=mimi_config["frame_rate"], frame_rate=mimi_config["frame_rate"],
encoder_frame_rate=mimi_config["sample_rate"] / encoder.hop_length, 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, encoder_transformer=encoder_transformer,
decoder_transformer=decoder_transformer, decoder_transformer=decoder_transformer,
).to(device="cpu") ).to(device="cpu")
@@ -153,11 +197,7 @@ class TTSModel(nn.Module):
tts_model.mimi.load_state_dict(mimi_state, strict=True) tts_model.mimi.load_state_dict(mimi_state, strict=True)
tts_model.mimi.eval() 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: if config.weights_path is not None:
logger.info(f"Loading TTSModel weights from {config.weights_path}") logger.info(f"Loading TTSModel weights from {config.weights_path}")
try: try:
@@ -174,6 +214,10 @@ class TTSModel(nn.Module):
"No weights_path specified for FlowLM or TTSModel, model is uninitialized!" "No weights_path specified for FlowLM or TTSModel, model is uninitialized!"
) )
size_in_mb = size_of_dict(tts_model.state_dict()) // 1e6 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") 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__ # TODO: move this in the __init__ and make self.mimi in __init__
@@ -188,7 +232,8 @@ class TTSModel(nn.Module):
@classmethod @classmethod
def load_model( def load_model(
cls, cls,
config: str | Path = DEFAULT_VARIANT, language: str | None = None,
config: str | Path | None = None,
temp: float | int = DEFAULT_TEMPERATURE, temp: float | int = DEFAULT_TEMPERATURE,
lsd_decode_steps: int = DEFAULT_LSD_DECODE_STEPS, lsd_decode_steps: int = DEFAULT_LSD_DECODE_STEPS,
noise_clamp: float | int | None = DEFAULT_NOISE_CLAMP, 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. with the specified generation parameters and ready for inference.
Args: Args:
config: a path to a custom YAML config file saved locally (e.g., C://pocket_tts/pocket_tts_config.yaml) language: Optional language identifier to select a predefined config. Incompatible with
or a model variant identifier (e.g., '610b0b2c'; must match a YAML file in the config directory). 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 temp: Sampling temperature for generation. Higher values produce more
diverse but potentially lower quality output. diverse but potentially lower quality output.
lsd_decode_steps: Number of steps for Lagrangian Self Distillation lsd_decode_steps: Number of steps for Lagrangian Self Distillation
@@ -238,15 +286,23 @@ class TTSModel(nn.Module):
model = TTSModel.load_model(quantize=True) model = TTSModel.load_model(quantize=True)
``` ```
""" """
if str(config).endswith(".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_path = Path(config)
config = load_config(config_path) config = load_config(config_path)
logger.info(f"Loading model from config at {config_path}...") logger.info(f"Loading model from config at {config_path}...")
else:
config = load_config(Path(__file__).parents[1] / f"config/{config}.yaml")
tts_model = TTSModel._from_pydantic_config_with_weights( 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: if quantize:
@@ -306,8 +362,23 @@ class TTSModel(nn.Module):
) )
return output_embeddings[:, None, :], is_eos 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: def _encode_audio(self, audio: torch.Tensor) -> torch.Tensor:
encoded = self.mimi.encode_to_latent(audio) 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) latents = encoded.transpose(-1, -2).to(torch.float32)
conditioning = F.linear(latents, self.flow_lm.speaker_proj_weight) conditioning = F.linear(latents, self.flow_lm.speaker_proj_weight)
return conditioning return conditioning
@@ -525,17 +596,25 @@ class TTSModel(nn.Module):
and audio decoding. Generation performance is logged including and audio decoding. Generation performance is logged including
real-time factor (RTF) metrics. 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 # 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. # 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 # TODO: add the teacher forcing method for long texts where we use the audio of one chunk
# as conditioning for the next chunk. # as conditioning for the next chunk.
chunks = split_into_best_sentences( 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: 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 frames_after_eos_guess += 2
effective_frames = ( effective_frames = (
frames_after_eos if frames_after_eos is not None else frames_after_eos_guess 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) 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. # 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)): if not self.has_voice_cloning and isinstance(audio_conditioning, (str, Path)):
raise ValueError(VOICE_CLONING_UNSUPPORTED) raise ValueError(VOICE_CLONING_UNSUPPORTED)
@@ -793,6 +885,9 @@ class TTSModel(nn.Module):
with display_execution_time("Encoding audio prompt"): with display_execution_time("Encoding audio prompt"):
prompt = self._encode_audio(audio_conditioning.unsqueeze(0).to(self.device)) 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]) model_state = init_states(self.flow_lm, batch_size=1, sequence_length=prompt.shape[1])
with display_execution_time("Prompting audio"): with display_execution_time("Prompting audio"):
@@ -810,11 +905,15 @@ class TTSModel(nn.Module):
return math.ceil(gen_len_sec * frame_rate) 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() text = text.strip()
if text == "": if text == "":
raise ValueError("Text prompt cannot be empty") raise ValueError("Text prompt cannot be empty")
text = text.replace("\n", " ").replace("\r", " ").replace(" ", " ") text = text.replace("\n", " ").replace("\r", " ").replace(" ", " ")
if remove_semicolons:
text = text.replace(";", ",")
number_of_words = len(text.split()) number_of_words = len(text.split())
if number_of_words <= 4: if number_of_words <= 4:
frames_after_eos_guess = 3 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 # 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. # 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 text = " " * 8 + text
return text, frames_after_eos_guess return text, frames_after_eos_guess
@@ -871,8 +970,16 @@ def _segments_from_boundaries(
return segments return segments
def split_into_best_sentences(tokenizer, text_to_generate: str, max_tokens: int) -> list[str]: def split_into_best_sentences(
text_to_generate, _ = prepare_text_prompt(text_to_generate) 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() text_to_generate = text_to_generate.strip()
tokens = tokenizer(text_to_generate) tokens = tokenizer(text_to_generate)
list_of_tokens = tokens.tokens[0].tolist() list_of_tokens = tokens.tokens[0].tolist()
pocket_tts/modules/resample.py
+9 -4
View File
@@ -10,11 +10,14 @@ class ConvDownsample1d(nn.Module):
with a kernel size of twice the stride. 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__() super().__init__()
if out_dimension is None:
out_dimension = dimension
self.conv = StreamingConv1d( self.conv = StreamingConv1d(
dimension, dimension,
dimension, out_dimension,
kernel_size=2 * stride, kernel_size=2 * stride,
stride=stride, stride=stride,
groups=1, groups=1,
@@ -31,10 +34,12 @@ class ConvTrUpsample1d(nn.Module):
Upsample by some integer amount `stride` using transposed convolutions. 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__() super().__init__()
if in_dimension is None:
in_dimension = dimension
self.convtr = StreamingConvTranspose1d( self.convtr = StreamingConvTranspose1d(
dimension, in_dimension,
dimension, dimension,
kernel_size=2 * stride, kernel_size=2 * stride,
stride=stride, stride=stride,
pocket_tts/utils/config.py
+14 -1
View File
@@ -5,6 +5,8 @@ from pathlib import Path
import yaml import yaml
from pydantic import BaseModel, ConfigDict from pydantic import BaseModel, ConfigDict
CONFIGS_DIR = Path(__file__).parent.parent / "config"
class StrictModel(BaseModel): class StrictModel(BaseModel):
model_config = ConfigDict(extra="forbid") model_config = ConfigDict(extra="forbid")
@@ -45,6 +47,7 @@ class FlowLMConfig(StrictModel):
# conditioning # conditioning
lookup_table: LookupTable lookup_table: LookupTable
weights_path: str | None = None weights_path: str | None = None
insert_bos_before_voice: bool = False
# SEANet configuration # SEANet configuration
@@ -101,6 +104,8 @@ class MimiConfig(StrictModel):
# Quantizer # Quantizer
quantizer: QuantizerConfig quantizer: QuantizerConfig
weights_path: str | None = None weights_path: str | None = None
inner_dim: int | None = None
outer_dim: int | None = None
class Config(StrictModel): class Config(StrictModel):
@@ -108,13 +113,21 @@ class Config(StrictModel):
mimi: MimiConfig mimi: MimiConfig
weights_path: str | None = None weights_path: str | None = None
weights_path_without_voice_cloning: 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: def load_config(yaml_path: str | Path) -> Config:
yaml_path = Path(yaml_path) yaml_path = Path(yaml_path)
if not yaml_path.exists(): 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: with open(yaml_path, "r") as f:
config_dict = yaml.safe_load(f) config_dict = yaml.safe_load(f)
pocket_tts/utils/utils.py
+5 -17
View File
@@ -1,15 +1,16 @@
import hashlib import hashlib
import logging import logging
import os
import time import time
from pathlib import Path from pathlib import Path
import requests import requests
import safetensors.torch
import torch import torch
from huggingface_hub import hf_hub_download from huggingface_hub import hf_hub_download
from torch import nn from torch import nn
PROJECT_ROOT = Path(__file__).parent.parent.parent PROJECT_ROOT = Path(__file__).parent.parent.parent
DEBUG_MIMI = os.environ.get("DEBUG_MIMI", "0") == "1"
_ORIGINS_OF_PREDEFINED_VOICES = { _ORIGINS_OF_PREDEFINED_VOICES = {
"cosette": "hf://kyutai/tts-voices/expresso/ex04-ex02_confused_001_channel1_499s.wav", "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", "caro_davy": "hf://kyutai/tts-voices/voice-zero/caro_davy.wav",
} }
PREDEFINED_VOICES = {
# don't forget to change this def get_predefined_voice(language: str, name: str) -> str:
x: f"hf://kyutai/pocket-tts-without-voice-cloning/embeddings_v3/{x}.safetensors@075c0abfe7e41450521b0200b5168cfbc16bc77b" return f"hf://kyutai/pocket-tts-without-voice-cloning/languages/{language}/embeddings/{name}.safetensors@d29db7978e464fb90cb3359ee0c69a273b9142cc"
for x in _ORIGINS_OF_PREDEFINED_VOICES
}
def make_cache_directory() -> Path: def make_cache_directory() -> Path:
@@ -113,14 +112,3 @@ def download_if_necessary(file_path: str) -> Path:
return Path(cached_file) return Path(cached_file)
else: else:
return Path(file_path) 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"]
pocket_tts/utils/weights_loading.py
+46 -2
View File
@@ -1,6 +1,7 @@
from pathlib import Path from pathlib import Path
import safetensors import safetensors
import torch
def get_flow_lm_state_dict(path: Path) -> dict: 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.") key.startswith("flow.w_s_t.")
or key == "condition_provider.conditioners.transcript_in_segment.learnt_padding" or key == "condition_provider.conditioners.transcript_in_segment.learnt_padding"
or key == "condition_provider.conditioners.speaker_wavs.learnt_padding" or key == "condition_provider.conditioners.speaker_wavs.learnt_padding"
or key == "num_ema_updates"
): ):
# skip lookup table weights # skip lookup table weights
continue continue
@@ -19,6 +21,11 @@ def get_flow_lm_state_dict(path: Path) -> dict:
new_name = "conditioner.embed.weight" new_name = "conditioner.embed.weight"
if key == "condition_provider.conditioners.speaker_wavs.output_proj.weight": if key == "condition_provider.conditioners.speaker_wavs.output_proj.weight":
new_name = "speaker_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) state_dict[new_name] = f.get_tensor(key)
return state_dict return state_dict
@@ -27,9 +34,46 @@ def get_mimi_state_dict(path: Path) -> dict:
state_dict = {} state_dict = {}
with safetensors.safe_open(path, framework="pt", device="cpu") as f: with safetensors.safe_open(path, framework="pt", device="cpu") as f:
for key in f.keys(): 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 # skip vq weights
continue 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 return state_dict
scripts/generate_default_voices.py
+5 -2
View File
@@ -1,9 +1,11 @@
import sys
import scipy.io.wavfile import scipy.io.wavfile
from pocket_tts import TTSModel, export_model_state from pocket_tts import TTSModel, export_model_state
from pocket_tts.utils.utils import _ORIGINS_OF_PREDEFINED_VOICES 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(): for voice_name, voice_origin in _ORIGINS_OF_PREDEFINED_VOICES.items():
print(f"Processing voice: {voice_name} from origin: {voice_origin}") 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( 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( scipy.io.wavfile.write(
f"./built-in-voices-generated/{voice_name}.wav", model.sample_rate, audio.numpy() f"./built-in-voices-generated/{voice_name}.wav", model.sample_rate, audio.numpy()
swarm-config.yaml
+89
View File
@@ -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"
tests/test_cli_generate.py
+2 -3
View File
@@ -6,7 +6,6 @@ import pytest
from typer.testing import CliRunner from typer.testing import CliRunner
from pocket_tts.data.audio import audio_read from pocket_tts.data.audio import audio_read
from pocket_tts.default_parameters import DEFAULT_VARIANT
from pocket_tts.main import cli_app 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" 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", "generate",
"--text", "--text",
"Testing custom parameters.", "Testing custom parameters.",
"--config", "--language",
DEFAULT_VARIANT, "english",
"--temperature", "--temperature",
"0.8", "0.8",
"--lsd-decode-steps", "--lsd-decode-steps",
tests/test_documentation_examples.py
+3 -1
View File
@@ -41,7 +41,9 @@ def test_load_model():
model = TTSModel.load_model() model = TTSModel.load_model()
# Load with custom parameters # 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(): def test_device():
tests/test_split_sentences.py
+34 -10
View File
@@ -13,14 +13,22 @@ def tokenizer():
def test_short_text_single_chunk(tokenizer): def test_short_text_single_chunk(tokenizer):
"""Short text should produce a single chunk.""" """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 assert len(chunks) == 1
def test_multiple_sentences_split(tokenizer): def test_multiple_sentences_split(tokenizer):
"""Multiple sentences should be split when they exceed max_tokens.""" """Multiple sentences should be split when they exceed max_tokens."""
text = "First sentence here. Second sentence here. Third sentence here. Fourth sentence here." 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 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 season of Light, it was the season of Darkness, "
"it was the spring of hope, it was the winter of despair" "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" 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) # 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" "it was the epoch of belief, it was the epoch of incredulity"
) )
max_tokens = 20 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: for chunk in chunks:
token_count = len(tokenizer(chunk.strip()).tokens[0].tolist()) token_count = len(tokenizer(chunk.strip()).tokens[0].tolist())
# Allow some tolerance since comma clauses may vary in size # 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. " "until it finally reaches a period. "
"Another short one." "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 assert len(chunks) >= 3
def test_no_commas_no_periods_stays_single_chunk(tokenizer): def test_no_commas_no_periods_stays_single_chunk(tokenizer):
"""Text with no splitting characters stays as a single chunk.""" """Text with no splitting characters stays as a single chunk."""
text = "one two three four five six seven eight nine ten eleven twelve" 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 # Should be 1 chunk since there are no split points
assert len(chunks) == 1 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; " "First clause here; second clause here; third clause here; "
"fourth clause here; fifth clause here; sixth 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 assert len(chunks) > 1
def test_short_sentence_not_affected_by_comma_splitting(tokenizer): def test_short_sentence_not_affected_by_comma_splitting(tokenizer):
"""Sentences under max_tokens should not be affected by comma logic.""" """Sentences under max_tokens should not be affected by comma logic."""
text = "Hello, world." 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 len(chunks) == 1
assert "hello" in chunks[0].lower() assert "hello" in chunks[0].lower()
assert "world" 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): def test_empty_string_raises(tokenizer):
"""Empty input should raise ValueError from prepare_text_prompt.""" """Empty input should raise ValueError from prepare_text_prompt."""
with pytest.raises(ValueError, match="empty"): 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): def test_oversized_clause_without_commas_still_returns(tokenizer):
"""A long clause with no split points should still be returned (not dropped).""" """A long clause with no split points should still be returned (not dropped)."""
# 20 words with no punctuation at all - no way to split # 20 words with no punctuation at all - no way to split
text = " ".join(f"word{i}" for i in range(20)) 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 assert len(chunks) == 1
# prepare_text_prompt capitalizes the first char and adds a trailing period, # prepare_text_prompt capitalizes the first char and adds a trailing period,
# so compare case-insensitively and strip punctuation # so compare case-insensitively and strip punctuation