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
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