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
* 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).
docs/API Reference/python-api.md
+8 -4
View File
@@ -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).
For more information about the command-line interface, see the [Generate Documentation](generate.md) or [Serve Documentation](serve.md).
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
- `--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
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.")
- `--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.
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
- `--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.
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
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:
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_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
pocket_tts/main.py
+73 -25
View File
@@ -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
)
pocket_tts/models/flow_lm.py
+9 -1
View File
@@ -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,
)
pocket_tts/models/mimi.py
+10 -2
View File
@@ -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:
pocket_tts/models/tts_model.py
+139 -32
View File
@@ -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()
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.
"""
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,
pocket_tts/utils/config.py
+14 -1
View File
@@ -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)
pocket_tts/utils/utils.py
+5 -17
View File
@@ -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"]
pocket_tts/utils/weights_loading.py
+46 -2
View File
@@ -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
scripts/generate_default_voices.py
+5 -2
View File
@@ -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()
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 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",
tests/test_documentation_examples.py
+3 -1
View File
@@ -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():
tests/test_split_sentences.py
+34 -10
View File
@@ -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