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

Use the new checkpoints structure (#10)

Browse Source
This commit is contained in:
Gabriel de Marmiesse
2026-01-12 12:56:13 +01:00
committed by GitHub
parent d4004e513e
commit becd3137a1
3 changed files with 10 additions and 87 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pocket_tts/models/tts_model.py
-79
View File
@@ -44,84 +44,6 @@ torch.set_num_threads(1)
logger = logging.getLogger(__name__)
KEYS_TO_MODIFY = [
"mimi.decoder.model.0.conv.conv.bias",
"mimi.decoder.model.0.conv.conv.weight",
"mimi.decoder.model.11.conv.conv.bias",
"mimi.decoder.model.11.conv.conv.weight",
"mimi.decoder.model.2.convtr.convtr.bias",
"mimi.decoder.model.2.convtr.convtr.weight",
"mimi.decoder.model.3.block.1.conv.conv.bias",
"mimi.decoder.model.3.block.1.conv.conv.weight",
"mimi.decoder.model.3.block.3.conv.conv.bias",
"mimi.decoder.model.3.block.3.conv.conv.weight",
"mimi.decoder.model.5.convtr.convtr.bias",
"mimi.decoder.model.5.convtr.convtr.weight",
"mimi.decoder.model.6.block.1.conv.conv.bias",
"mimi.decoder.model.6.block.1.conv.conv.weight",
"mimi.decoder.model.6.block.3.conv.conv.bias",
"mimi.decoder.model.6.block.3.conv.conv.weight",
"mimi.decoder.model.8.convtr.convtr.bias",
"mimi.decoder.model.8.convtr.convtr.weight",
"mimi.decoder.model.9.block.1.conv.conv.bias",
"mimi.decoder.model.9.block.1.conv.conv.weight",
"mimi.decoder.model.9.block.3.conv.conv.bias",
"mimi.decoder.model.9.block.3.conv.conv.weight",
"mimi.downsample.conv.conv.conv.weight",
"mimi.encoder.model.0.conv.conv.bias",
"mimi.encoder.model.0.conv.conv.weight",
"mimi.encoder.model.1.block.1.conv.conv.bias",
"mimi.encoder.model.1.block.1.conv.conv.weight",
"mimi.encoder.model.1.block.3.conv.conv.bias",
"mimi.encoder.model.1.block.3.conv.conv.weight",
"mimi.encoder.model.11.conv.conv.bias",
"mimi.encoder.model.11.conv.conv.weight",
"mimi.encoder.model.3.conv.conv.bias",
"mimi.encoder.model.3.conv.conv.weight",
"mimi.encoder.model.4.block.1.conv.conv.bias",
"mimi.encoder.model.4.block.1.conv.conv.weight",
"mimi.encoder.model.4.block.3.conv.conv.bias",
"mimi.encoder.model.4.block.3.conv.conv.weight",
"mimi.encoder.model.6.conv.conv.bias",
"mimi.encoder.model.6.conv.conv.weight",
"mimi.encoder.model.7.block.1.conv.conv.bias",
"mimi.encoder.model.7.block.1.conv.conv.weight",
"mimi.encoder.model.7.block.3.conv.conv.bias",
"mimi.encoder.model.7.block.3.conv.conv.weight",
"mimi.encoder.model.9.conv.conv.bias",
"mimi.encoder.model.9.conv.conv.weight",
"mimi.upsample.convtr.convtr.convtr.weight",
]
KEY_WITH_REMOVED_ZERO = [
"mimi.encoder_transformer.transformer.layers.0.self_attn.in_projs.0.weight",
"mimi.encoder_transformer.transformer.layers.0.self_attn.out_projs.0.weight",
"mimi.encoder_transformer.transformer.layers.1.self_attn.in_projs.0.weight",
"mimi.encoder_transformer.transformer.layers.1.self_attn.out_projs.0.weight",
"mimi.decoder_transformer.transformer.layers.0.self_attn.in_projs.0.weight",
"mimi.decoder_transformer.transformer.layers.0.self_attn.out_projs.0.weight",
"mimi.decoder_transformer.transformer.layers.1.self_attn.in_projs.0.weight",
"mimi.decoder_transformer.transformer.layers.1.self_attn.out_projs.0.weight",
]
def adapt_state_dict_of_TTSModel(state_dict: dict) -> dict:
# should be removed before release
new_state_dict = {}
for key, value in state_dict.items():
if key == "mimi.quantizer.input_proj.weight":
continue
if key in KEYS_TO_MODIFY:
new_key = key.replace(".conv.conv.", ".conv.").replace(".convtr.convtr.", ".convtr.")
new_state_dict[new_key] = value
elif key in KEY_WITH_REMOVED_ZERO:
new_key = key.replace("s.0.weight", ".weight")
new_state_dict[new_key] = value
else:
new_state_dict[key] = value
return new_state_dict
class TTSModel(nn.Module):
def __init__(
self,
@@ -230,7 +152,6 @@ class TTSModel(nn.Module):
weights_file = download_if_necessary(config.weights_path_without_voice_cloning)
state_dict = safetensors.torch.load_file(weights_file)
state_dict = adapt_state_dict_of_TTSModel(state_dict)
tts_model.load_state_dict(state_dict, strict=True)
if config.flow_lm.weights_path is None and config.weights_path is None:
pocket_tts/modules/transformer.py
+5 -7
View File
@@ -56,9 +56,7 @@ class StreamingMultiheadAttention(StatefulModule):
kv_dim = (embed_dim // num_heads) * num_kv
out_dim += 2 * kv_dim
mult = 1
in_proj = nn.Linear(embed_dim, mult * out_dim, bias=False)
# We try to follow the default PyTorch MHA convention, to easily compare results.
self.in_proj_weight = in_proj.weight
self.in_proj = nn.Linear(embed_dim, mult * out_dim, bias=False)
self.out_proj = nn.Linear(embed_dim, mult * embed_dim, bias=False)
def _get_mask(self, shape: tuple[int, int], shift: int, device: torch.device) -> torch.Tensor:
@@ -66,14 +64,14 @@ class StreamingMultiheadAttention(StatefulModule):
def init_state(self, batch_size: int, sequence_length: int) -> dict[str, torch.Tensor]:
dim_per_head = self.embed_dim // self.num_heads
initial_current_end = torch.zeros((0,)).to(self.in_proj_weight.device)
initial_current_end = torch.zeros((0,)).to(self.in_proj.weight.device)
return dict(
current_end=initial_current_end,
cache=torch.full(
(2, batch_size, sequence_length, self.num_heads, dim_per_head),
float("NaN"),
device=self.in_proj_weight.device,
dtype=self.in_proj_weight.dtype,
device=self.in_proj.weight.device,
dtype=self.in_proj.weight.dtype,
),
)
@@ -101,7 +99,7 @@ class StreamingMultiheadAttention(StatefulModule):
def forward(self, query: torch.Tensor, model_state: dict | None):
state = self.check_model_state(model_state)
projected = nn.functional.linear(query, self.in_proj_weight)
projected = self.in_proj(query)
# Reshape from (b, t, p*h*d) to (b, t, p, h, d) where p=3, h=num_heads
b, t, _ = projected.shape
d = self.embed_dim // self.num_heads
pocket_tts/utils/utils.py
+5 -1
View File
@@ -85,7 +85,11 @@ def download_if_necessary(file_path: str) -> Path:
splitted = file_path.split("/")
repo_id = "/".join(splitted[:2])
filename = "/".join(splitted[2:])
cached_file = hf_hub_download(repo_id=repo_id, filename=filename)
if "@" in filename:
filename, revision = filename.split("@")
else:
revision = None
cached_file = hf_hub_download(repo_id=repo_id, filename=filename, revision=revision)
return Path(cached_file)
else:
return Path(file_path)