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