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