From becd3137a100ad9e3279b4f8e8a798fbf260f74a Mon Sep 17 00:00:00 2001 From: Gabriel de Marmiesse Date: Mon, 12 Jan 2026 12:56:13 +0100 Subject: [PATCH] Use the new checkpoints structure (#10) --- pocket_tts/models/tts_model.py | 79 ------------------------------- pocket_tts/modules/transformer.py | 12 ++--- pocket_tts/utils/utils.py | 6 ++- 3 files changed, 10 insertions(+), 87 deletions(-) diff --git a/pocket_tts/models/tts_model.py b/pocket_tts/models/tts_model.py index 8488185..0970a1c 100644 --- a/pocket_tts/models/tts_model.py +++ b/pocket_tts/models/tts_model.py @@ -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: diff --git a/pocket_tts/modules/transformer.py b/pocket_tts/modules/transformer.py index a63f241..2b8278c 100644 --- a/pocket_tts/modules/transformer.py +++ b/pocket_tts/modules/transformer.py @@ -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 diff --git a/pocket_tts/utils/utils.py b/pocket_tts/utils/utils.py index 1df6e18..d6e32ab 100644 --- a/pocket_tts/utils/utils.py +++ b/pocket_tts/utils/utils.py @@ -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)