Use the new checkpoints structure (#10)

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:

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)
+        self.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.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,
+                device=self.in_proj.weight.device,
-                dtype=self.in_proj_weight.dtype,
+                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

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