Repository: slp-rl/slamkit
Pull request #4: Add vocoder support locally
Author: login: cantabile-kwok
Created: 2025-03-04T08:11:41Z | Merged: 2025-03-05T07:36:21Z

Isolate vocoder files from textlesslib so that users are not required to install textlesslib, if they want to generate waveforms given prompts.


Diff:
diff --git a/README.md b/README.md
index 7494a9d..428fde1 100644
--- a/README.md
+++ b/README.md
@@ -124,10 +124,18 @@ An example of running it:
 python cli/eval.py tokeniser=unit_hubert_25 metric=tstorycloze metric.data_path=<TSC_PATH> batch_size=32 model.pretrained_model=<TRAINED_MODEL_PATH>
 ```
 
+Another example of running generation given a prompt, using a pretrained vocoder:
+```bash
+python cli/eval.py tokeniser=unit_hubert_25 metric=generate batch_size=32 model.pretrained_model=slprl/slam metric.data_path=/some/path/*.wav vocoder=vocoder_hubert_25
+# then, find generated files at `generated/`.
+```
+
 
 ❗Note that the `model.pretrained_model` needs to point to a folder of a specific step within the output folder `training_args.output_dir` from training or a model from hf.
 
-❗Note that we currently only support running one metric with each run of `eval.py`. You can use [hydra multirun](https://hydra.cc/docs/tutorials/basic/running_your_app/multi-run/) to run them in sucession or in parallel
+❗Note that we currently only support running one metric with each run of `eval.py`. You can use [hydra multirun](https://hydra.cc/docs/tutorials/basic/running_your_app/multi-run/) to run them in sucession or in parallel.
+
+❗Note that you can set `TEXTLESS_CHECKPOINT_ROOT` to specify the download location of the vocoder. It defaults to `~/.textless/`.
 
 ## Results
 We provide some results for our pre-trained models, compared to other SLMs.
@@ -171,6 +179,9 @@ We welcome contributions to this repository. Want to add support for new tokenis
 Add support for even more efficient implementations? If you are interested in building this open source 
 project - open an Issue, and one of the maintainers will guide you in opening a PR!
 
+## Acknowledgements
+We isolate vocoder-related code from [textlesslib](https://github.com/facebookresearch/textlesslib), in `slamkit/vocoder/textless_*`.
+
 ## Citation
 If you found this repository useful, please cite our work:
 ```bibtex
diff --git a/cli/eval.py b/cli/eval.py
index 4465a0c..ec0775a 100644
--- a/cli/eval.py
+++ b/cli/eval.py
@@ -42,6 +42,8 @@ def main(cfg: DictConfig):
         elif cfg.metric.metric_type == 'salmon':
             res = salmon(model, path, used_token_modality, mean_nll, cfg.metric.parts, cfg.batch_size, cfg.num_workers, cfg.pin_memory)
         elif cfg.metric.metric_type == 'generate':
+            if cfg.vocoder is None:
+                logger.warning("You are currently trying to run generation without a vocoder, which will generate tokens, but has no effect. You can use a vocoder by, e.g. setting `vocoder=vocoder_hubert_25`")
             res = generate(model, path, cfg.batch_size, used_token_modality,
                         cfg.metric.prompt_length, tokeniser.fe_sample_rate, cfg.metric.num_files,
                         cfg.num_workers, cfg.pin_memory, **cfg.metric.get("generate_kwargs", {}))
@@ -56,7 +58,7 @@ def main(cfg: DictConfig):
             if key == "generate" or key == "prompts":
                 continue
             print(f"{key}: {val}")
-    if cfg.metric.get("out_path", False) and "generate" in res:
+    if cfg.metric.get("out_path", False) and "generate" in res and cfg.vocoder is not None:
         import torchaudio
         os.makedirs(cfg.metric.out_path, exist_ok=True)
         for i, gen in enumerate(res["generate"]):
@@ -73,7 +75,7 @@ def main(cfg: DictConfig):
         if cfg.logger.run_id is None:
             raise ValueError('No run_id specified for wandb logging')
         wandb.init(project=cfg.logger.project, entity=cfg.logger.entity, id=cfg.logger.run_id, resume="must")
-        if "generate" in res and "prompts" in res:
+        if "generate" in res and "prompts" in res and cfg.vocoder is not None:
             logs = {}
             for i, (gen, prompt) in enumerate(zip(res["generate"], res["prompts"])):
                 if i == cfg.metric.get("num_log", -1):
diff --git a/slamkit/vocoder/hifi_gan_vocoder.py b/slamkit/vocoder/hifi_gan_vocoder.py
index 256b2a0..97fde81 100644
--- a/slamkit/vocoder/hifi_gan_vocoder.py
+++ b/slamkit/vocoder/hifi_gan_vocoder.py
@@ -11,7 +11,7 @@ def __init__(self,
                  style_meta: str = None,
                  ):
         super().__init__()
-        from textless.vocoders.hifigan.vocoder import CodeHiFiGANVocoder
+        from .hifigan.vocoder import CodeHiFiGANVocoder
         self.vocoder = CodeHiFiGANVocoder.by_name(dense_model_name,
                                                   quantizer_model_name,
                                                   vocab_size,
diff --git a/slamkit/vocoder/hifigan/__init__.py b/slamkit/vocoder/hifigan/__init__.py
new file mode 100644
index 0000000..54f929f
--- /dev/null
+++ b/slamkit/vocoder/hifigan/__init__.py
@@ -0,0 +1,8 @@
+# Copied from textlesslib (https://github.com/facebookresearch/textlesslib)
+# Please refer to the MIT license there (https://github.com/facebookresearch/textlesslib?tab=MIT-1-ov-file)
+# Below is the original header from the file
+
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
diff --git a/slamkit/vocoder/hifigan/generator.py b/slamkit/vocoder/hifigan/generator.py
new file mode 100644
index 0000000..46ef9d7
--- /dev/null
+++ b/slamkit/vocoder/hifigan/generator.py
@@ -0,0 +1,238 @@
+# Copied from textlesslib (https://github.com/facebookresearch/textlesslib)
+# Please refer to the MIT license there (https://github.com/facebookresearch/textlesslib?tab=MIT-1-ov-file)
+# Below is the original header from the file
+
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Hifi-gan Generator and CodeGenerator modules
+# Adapted from:
+# - https://github.com/jik876/hifi-gan
+# - https://github.com/facebookresearch/fairseq/tree/main/fairseq/models/text_to_speech
+# - https://github.com/facebookresearch/speech-resynthesis/blob/main/examples/speech_to_speech_translation/models.py
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn import Conv1d, ConvTranspose1d
+from torch.nn.utils import remove_weight_norm, weight_norm
+from .resblock import ResBlock, init_weights, LRELU_SLOPE
+
+
+class Generator(nn.Module):
+    def __init__(self, cfg):
+        super(Generator, self).__init__()
+        self.num_kernels = len(cfg["resblock_kernel_sizes"])
+        self.num_upsamples = len(cfg["upsample_rates"])
+        self.conv_pre = weight_norm(
+            Conv1d(
+                cfg.get("model_in_dim", 80),
+                cfg["upsample_initial_channel"],
+                7,
+                1,
+                padding=3,
+            )
+        )
+
+        self.ups = nn.ModuleList()
+        for i, (u, k) in enumerate(
+            zip(cfg["upsample_rates"], cfg["upsample_kernel_sizes"])
+        ):
+            self.ups.append(
+                weight_norm(
+                    ConvTranspose1d(
+                        cfg["upsample_initial_channel"] // (2**i),
+                        cfg["upsample_initial_channel"] // (2 ** (i + 1)),
+                        k,
+                        u,
+                        padding=(k - u) // 2,
+                    )
+                )
+            )
+
+        self.resblocks = nn.ModuleList()
+        for i in range(len(self.ups)):
+            ch = cfg["upsample_initial_channel"] // (2 ** (i + 1))
+            for k, d in zip(
+                cfg["resblock_kernel_sizes"], cfg["resblock_dilation_sizes"]
+            ):
+                self.resblocks.append(ResBlock(ch, k, d))
+
+        self.conv_post = weight_norm(Conv1d(ch, 1, 7, 1, padding=3))
+        self.ups.apply(init_weights)
+        self.conv_post.apply(init_weights)
+
+    def forward(self, x):
+        x = self.conv_pre(x)
+        for i in range(self.num_upsamples):
+            x = F.leaky_relu(x, LRELU_SLOPE)
+            x = self.ups[i](x)
+            xs = None
+            for j in range(self.num_kernels):
+                if xs is None:
+                    xs = self.resblocks[i * self.num_kernels + j](x)
+                else:
+                    xs += self.resblocks[i * self.num_kernels + j](x)
+            x = xs / self.num_kernels
+        x = F.leaky_relu(x)
+        x = self.conv_post(x)
+        x = torch.tanh(x)
+
+        return x
+
+    def remove_weight_norm(self):
+        print("Removing weight norm...")
+        for layer in self.ups:
+            remove_weight_norm(layer)
+        for layer in self.resblocks:
+            layer.remove_weight_norm()
+        remove_weight_norm(self.conv_pre)
+        remove_weight_norm(self.conv_post)
+
+
+class CodeGenerator(Generator):
+    def __init__(self, cfg):
+        super().__init__(cfg)
+        self.dict = nn.Embedding(cfg["num_embeddings"], cfg["embedding_dim"])
+        self.multispkr = cfg.get("multispkr", None)
+        self.embedder = cfg.get("embedder_params", None)
+
+        self.multistyle = cfg.get("multistyle", None)
+
+        if self.multispkr and not self.embedder:
+            self.spkr = nn.Embedding(cfg.get("num_speakers", 200), cfg["embedding_dim"])
+        elif self.embedder:
+            self.spkr = nn.Linear(cfg.get("embedder_dim", 256), cfg["embedding_dim"])
+
+        if self.multistyle:
+            self.style = nn.Embedding(cfg.get("num_styles", 100), cfg["embedding_dim"])
+
+        self.dur_predictor = None
+        if cfg.get("dur_predictor_params", None):
+            self.dur_predictor = VariancePredictor(
+                cfg["dur_predictor_params"]["encoder_embed_dim"],
+                cfg["dur_predictor_params"]["var_pred_hidden_dim"],
+                cfg["dur_predictor_params"]["var_pred_kernel_size"],
+                cfg["dur_predictor_params"]["var_pred_dropout"],
+            )
+
+        self.f0 = cfg.get("f0", None)
+        n_f0_bin = cfg.get("f0_quant_num_bin", 0)
+        self.f0_quant_embed = (
+            None if n_f0_bin <= 0 else nn.Embedding(n_f0_bin, cfg["embedding_dim"])
+        )
+
+    @staticmethod
+    def _upsample(signal, max_frames):
+        if signal.dim() == 3:
+            bsz, channels, cond_length = signal.size()
+        elif signal.dim() == 2:
+            signal = signal.unsqueeze(2)
+            bsz, channels, cond_length = signal.size()
+        else:
+            signal = signal.view(-1, 1, 1)
+            bsz, channels, cond_length = signal.size()
+
+        signal = signal.unsqueeze(3).repeat(1, 1, 1, max_frames // cond_length)
+
+        # pad zeros as needed (if signal's shape does not divide completely with max_frames)
+        reminder = (max_frames - signal.shape[2] * signal.shape[3]) // signal.shape[3]
+        if reminder > 0:
+            raise NotImplementedError(
+                "Padding condition signal - misalignment between condition features."
+            )
+
+        signal = signal.view(bsz, channels, max_frames)
+        return signal
+
+    def forward(self, **kwargs):
+        x = self.dict(kwargs["code"]).transpose(1, 2)
+
+        if self.dur_predictor and kwargs.get("dur_prediction", False):
+            assert x.size(0) == 1, "only support single sample"
+            log_dur_pred = self.dur_predictor(x.transpose(1, 2))
+            dur_out = torch.clamp(
+                torch.round((torch.exp(log_dur_pred) - 1)).long(), min=1
+            )
+            # B x C x T
+            x = torch.repeat_interleave(x, dur_out.view(-1), dim=2)
+
+        if self.f0:
+            if self.f0_quant_embed:
+                kwargs["f0"] = self.f0_quant_embed(kwargs["f0"].long()).transpose(1, 2)
+            else:
+                kwargs["f0"] = kwargs["f0"].unsqueeze(1)
+
+            if x.shape[-1] < kwargs["f0"].shape[-1]:
+                x = self._upsample(x, kwargs["f0"].shape[-1])
+            elif x.shape[-1] > kwargs["f0"].shape[-1]:
+                kwargs["f0"] = self._upsample(kwargs["f0"], x.shape[-1])
+            x = torch.cat([x, kwargs["f0"]], dim=1)
+
+        if self.multispkr:
+            assert (
+                "spkr" in kwargs
+            ), 'require "spkr" input for multispeaker CodeHiFiGAN vocoder'
+            spkr = self.spkr(kwargs["spkr"]).transpose(1, 2)
+            spkr = self._upsample(spkr, x.shape[-1])
+            x = torch.cat([x, spkr], dim=1)
+
+        if self.multistyle:
+            assert (
+                "style" in kwargs
+            ), 'require "style" input for multispeaker CodeHiFiGAN vocoder'
+            style = self.style(kwargs["style"]).transpose(1, 2)
+            style = self._upsample(style, x.shape[-1])
+            x = torch.cat([x, style], dim=1)
+
+        for k, feat in kwargs.items():
+            if k in ["spkr", "code", "f0", "dur_prediction", "style"]:
+                continue
+
+            feat = self._upsample(feat, x.shape[-1])
+            x = torch.cat([x, feat], dim=1)
+
+        return super().forward(x)
+
+
+class VariancePredictor(nn.Module):
+    def __init__(
+        self,
+        encoder_embed_dim,
+        var_pred_hidden_dim,
+        var_pred_kernel_size,
+        var_pred_dropout,
+    ):
+        super().__init__()
+        self.conv1 = nn.Sequential(
+            nn.Conv1d(
+                encoder_embed_dim,
+                var_pred_hidden_dim,
+                kernel_size=var_pred_kernel_size,
+                padding=(var_pred_kernel_size - 1) // 2,
+            ),
+            nn.ReLU(),
+        )
+        self.ln1 = nn.LayerNorm(var_pred_hidden_dim)
+        self.dropout = var_pred_dropout
+        self.conv2 = nn.Sequential(
+            nn.Conv1d(
+                var_pred_hidden_dim,
+                var_pred_hidden_dim,
+                kernel_size=var_pred_kernel_size,
+                padding=1,
+            ),
+            nn.ReLU(),
+        )
+        self.ln2 = nn.LayerNorm(var_pred_hidden_dim)
+        self.proj = nn.Linear(var_pred_hidden_dim, 1)
+
+    def forward(self, x):
+        # Input: B x T x C; Output: B x T
+        x = self.conv1(x.transpose(1, 2)).transpose(1, 2)
+        x = F.dropout(self.ln1(x), p=self.dropout, training=self.training)
+        x = self.conv2(x.transpose(1, 2)).transpose(1, 2)
+        x = F.dropout(self.ln2(x), p=self.dropout, training=self.training)
+        return self.proj(x).squeeze(dim=2)
diff --git a/slamkit/vocoder/hifigan/resblock.py b/slamkit/vocoder/hifigan/resblock.py
new file mode 100644
index 0000000..b1fd6c9
--- /dev/null
+++ b/slamkit/vocoder/hifigan/resblock.py
@@ -0,0 +1,118 @@
+# Copied from textlesslib (https://github.com/facebookresearch/textlesslib)
+# Please refer to the MIT license there (https://github.com/facebookresearch/textlesslib?tab=MIT-1-ov-file)
+# Below is the original header from the file
+
+# Residual block from the Hifi-gan model:
+# https://github.com/jik876/hifi-gan
+
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn import Conv1d
+from torch.nn.utils import remove_weight_norm, weight_norm
+
+
+LRELU_SLOPE = 0.1
+
+
+def init_weights(m, mean=0.0, std=0.01):
+    classname = m.__class__.__name__
+    if classname.find("Conv") != -1:
+        m.weight.data.normal_(mean, std)
+
+
+def get_padding(kernel_size, dilation=1):
+    return (kernel_size * dilation - dilation) // 2
+
+
+class ResBlock(nn.Module):
+    def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5)):
+        super(ResBlock, self).__init__()
+
+        self.convs1 = nn.ModuleList(
+            [
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[0],
+                        padding=get_padding(kernel_size, dilation[0]),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[1],
+                        padding=get_padding(kernel_size, dilation[1]),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[2],
+                        padding=get_padding(kernel_size, dilation[2]),
+                    )
+                ),
+            ]
+        )
+
+        self.convs1.apply(init_weights)
+
+        self.convs2 = nn.ModuleList(
+            [
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=1,
+                        padding=get_padding(kernel_size, 1),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=1,
+                        padding=get_padding(kernel_size, 1),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=1,
+                        padding=get_padding(kernel_size, 1),
+                    )
+                ),
+            ]
+        )
+
+        self.convs2.apply(init_weights)
+
+    def forward(self, x):
+        for c1, c2 in zip(self.convs1, self.convs2):
+            xt = F.leaky_relu(x, LRELU_SLOPE)
+            xt = c1(xt)
+            xt = F.leaky_relu(xt, LRELU_SLOPE)
+            xt = c2(xt)
+            x = xt + x
+        return x
+
+    def remove_weight_norm(self):
+        for layer in self.convs1:
+            remove_weight_norm(layer)
+        for layer in self.convs2:
+            remove_weight_norm(layer)
diff --git a/slamkit/vocoder/hifigan/vocoder.py b/slamkit/vocoder/hifigan/vocoder.py
new file mode 100644
index 0000000..e493cd0
--- /dev/null
+++ b/slamkit/vocoder/hifigan/vocoder.py
@@ -0,0 +1,181 @@
+# Copied from textlesslib (https://github.com/facebookresearch/textlesslib)
+# Please refer to the MIT license there (https://github.com/facebookresearch/textlesslib?tab=MIT-1-ov-file)
+# Below is the original header from the file
+
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+import torch.nn as nn
+import numpy as np
+from .generator import CodeGenerator
+
+import json
+from typing import Union
+from pathlib import Path
+
+from ..textless_checkpoint_manager import CHECKPOINT_MANAGER
+
+
+class CodeHiFiGANVocoder(nn.Module):
+    def __init__(
+        self,
+        hifigan_model_path: str,
+        hifigan_config_path: str,
+        hifigan_speaker_path: str = None,
+        hifigan_style_path: str = None,
+        fp16: str = False,
+    ) -> None:
+        super().__init__()
+
+        # Load hifigan config
+        with open(hifigan_config_path) as f:
+            cfg = json.load(f)
+        self.cfg = cfg
+
+        # Load hifigan model
+        self.model = CodeGenerator(cfg)
+        state_dict = torch.load(hifigan_model_path, map_location="cpu")
+        self.model.load_state_dict(state_dict["generator"])
+        self.model.eval()
+        if fp16:
+            self.model.half()
+        self.model.remove_weight_norm()
+        print("CodeHiFiGAN model loaded!")
+
+        # Load hifigan metadata (if exists)
+        self.speakers, self.styles = load_vocoder_meta(
+            speakers_path=hifigan_speaker_path, styles_path=hifigan_style_path
+        )
+
+        # Useful for detecting device
+        self.register_buffer("_float_tensor", torch.tensor([0], dtype=torch.float))
+
+    def forward(
+        self,
+        code: torch.Tensor,
+        speaker_id: int = 0,
+        style_id: int = 0,
+        f0: torch.Tensor = None,
+        dur_prediction: bool = False,
+    ) -> torch.Tensor:
+        x = dict()
+        x["code"] = code
+        if self.model.multispkr:
+            x["spkr"] = torch.LongTensor([speaker_id]).view(1, 1).to(code.device)
+        if self.model.multistyle:
+            x["style"] = torch.LongTensor([style_id]).view(1, 1).to(code.device)
+        if self.model.f0:
+            assert f0 is not None
+            x["f0"] = f0.to(code.device)
+
+        x["dur_prediction"] = dur_prediction
+        if dur_prediction:
+            assert (
+                self.model.dur_predictor is not None
+            ), "This CodeHiFiGAN model doesn't support duration prediction!"
+
+        # remove invalid code
+        mask = x["code"] >= 0
+        x["code"] = x["code"][mask].unsqueeze(dim=0)
+        if "f0" in x:
+            f0_up_ratio = x["f0"].size(1) // x["code"].size(1)
+            mask = mask.unsqueeze(2).repeat(1, 1, f0_up_ratio).view(-1, x["f0"].size(1))
+            x["f0"] = x["f0"][mask].unsqueeze(dim=0)
+
+        return self.model(**x).detach().squeeze()
+
+    @property
+    def device(self) -> torch.device:
+        return self._float_tensor.device
+
+    @property
+    def output_sample_rate(self) -> int:
+        return self.cfg.get("sampling_rate", 16_000)
+
+    @classmethod
+    def by_name(
+        cls,
+        dense_model_name: str,
+        quantizer_model_name: str,
+        vocab_size: int,
+        vocoder_suffix: str = None,
+        speaker_meta: bool = False,
+        style_meta: bool = False,
+    ):
+
+        # Get hifigan checkpoint name and path
+        hifigan_checkpoint_name = (
+            f"{dense_model_name}-{quantizer_model_name}-{vocab_size}-hifigan"
+        )
+        if vocoder_suffix is not None:
+            hifigan_checkpoint_name += "-" + vocoder_suffix
+        hifigan_checkpoint_path = CHECKPOINT_MANAGER.get_by_name(
+            hifigan_checkpoint_name
+        )
+
+        # Get hifigan config name and path
+        hifigan_config_name = f"{hifigan_checkpoint_name}-config"
+        hifigan_config_path = CHECKPOINT_MANAGER.get_by_name(hifigan_config_name)
+
+        # Get hifigan speaker metadata name and path
+        hifigan_speaker_path = None
+        if speaker_meta:
+            hifigan_speaker_name = f"{hifigan_checkpoint_name}-speakers"
+            hifigan_speaker_path = CHECKPOINT_MANAGER.get_by_name(hifigan_speaker_name)
+
+        # Get hifigan style metadata name and path
+        hifigan_style_path = None
+        if style_meta:
+            hifigan_style_name = f"{hifigan_checkpoint_name}-styles"
+            hifigan_style_path = CHECKPOINT_MANAGER.get_by_name(hifigan_style_name)
+
+        return cls(
+            hifigan_checkpoint_path,
+            hifigan_config_path,
+            hifigan_speaker_path,
+            hifigan_style_path,
+        )
+
+
+def load_vocoder_meta(speakers_path=None, styles_path=None):
+    """
+    Load speakers and styles of a vocoder from text files.
+    """
+    speakers, styles = None, None
+
+    if speakers_path is not None and Path(speakers_path).exists():
+        with open(speakers_path) as f:
+            speakers = [line.strip() for line in f]
+        print(f"Loaded {len(speakers)} speakers. First few speakers: {speakers[:10]}")
+
+    if styles_path is not None and Path(styles_path).exists():
+        with open(styles_path) as f:
+            styles = [line.strip() for line in f]
+        print(f"Loaded {len(styles)} styles. First few styles: {styles[:10]}")
+
+    return speakers, styles
+
+
+def preprocess_code(
+    code: Union[str, list, torch.Tensor], deduplicate_code: bool = False
+) -> torch.Tensor:
+    """
+    Convert the code to a long tensor.
+    The code can be one of the following forms:
+        - String of integers: "1 2 3"
+        - List/Array of (string of) integers: [1, 2, 3]
+    """
+    if isinstance(code, str):
+        code = code.split()
+    if isinstance(code, list):
+        code = list(map(int, code))
+        code = torch.tensor(code)
+    elif isinstance(code, np.ndarray):
+        code = torch.from_numpy(code)
+    code = code.long()
+    if deduplicate_code:
+        code = torch.unique_consecutive(code)
+    return code.view(1, -1)
diff --git a/slamkit/vocoder/textless_checkpoint_manager/__init__.py b/slamkit/vocoder/textless_checkpoint_manager/__init__.py
new file mode 100644
index 0000000..f133b58
--- /dev/null
+++ b/slamkit/vocoder/textless_checkpoint_manager/__init__.py
@@ -0,0 +1,353 @@
+# Copied from textlesslib (https://github.com/facebookresearch/textlesslib)
+# Please refer to the MIT license there (https://github.com/facebookresearch/textlesslib?tab=MIT-1-ov-file)
+# Below is the original header from the file
+
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+from .manager import Checkpoint, CheckpointManager
+
+root = os.environ.get("TEXTLESS_CHECKPOINT_ROOT", "~/.textless/")
+CHECKPOINT_MANAGER: CheckpointManager = CheckpointManager(disk_root=root)
+
+
+def populate_checkpoints():
+    global CHECKPOINT_MANAGER
+
+    # HuBERT based
+    checkpoints = [
+        # Dense model
+        Checkpoint(
+            name="hubert-base-ls960",
+            remote_path="https://dl.fbaipublicfiles.com/hubert/hubert_base_ls960.pt",
+            fname="hubert_base_ls960.pt",
+            sha256="1703cf8d2cdc76f8c046f5f6a9bcd224e0e6caf4744cad1a1f4199c32cac8c8d",
+        ),
+        Checkpoint(
+            name="mhubert-base-vp_en_es_fr",
+            remote_path="https://dl.fbaipublicfiles.com/hubert/mhubert_base_vp_en_es_fr_it3.pt",
+            fname="mhubert_base_vp_en_es_fr_it3.pt",
+            sha256="23527cd6311a6617bfa199f574e2eed478cd6a58a634f4007165c10ec0019832",
+        ),
+        Checkpoint(
+            name="mhubert-base-vp_mls_cv_8lang",
+            remote_path="https://dl.fbaipublicfiles.com/hubert/mhubert_base_vp_mls_cv_8lang_it3.pt",
+            fname="mhubert_base_vp_mls_cv_8lang_it3.pt",
+            sha256="fa82eaccb557465f4732662e3a744c30c953059147ac6a4c92f1f3e0daeda61e",
+        ),
+        Checkpoint(
+            name="mhubert-base-25hz",
+            remote_path="https://dl.fbaipublicfiles.com/textless_nlp/twist/speech_tokenizer/mhubert_base_25hz_cp_mls_cv_sp_fisher.pt",
+            fname="mhubert_base_25hz_cp_mls_cv_sp_fisher.pt",
+            sha256="edf0f5be2581a568a236aeb2781a5607c5af25806e12f40f7f2391e861fe3cf0",
+        ),
+        # Quantizers
+        Checkpoint(
+            name="hubert-base-ls960-kmeans-50",
+            remote_path="https://dl.fbaipublicfiles.com/textless_nlp/gslm/hubert/km50/km.bin",
+            fname="hubert_base_ls960_km50.pt",
+            sha256="d01a7d5bc2c54b7b5f25f321ba525b4d230b06e3927f90bad0394198bc89f494",
+        ),
+        Checkpoint(
+            name="hubert-base-ls960-kmeans-100",
+            remote_path="https://dl.fbaipublicfiles.com/textless_nlp/gslm/hubert/km100/km.bin",
+            fname="hubert_base_ls960_km100.pt",
+            sha256="f14a3104615485381fc489701d6761c9abbdbb0d43607cf55518a1e2891023fe",
+        ),
+        Checkpoint(
+            name="hubert-base-ls960-kmeans-200",
+            remote_path="https://dl.fbaipublicfiles.com/textless_nlp/gslm/hubert/km200/km.bin",
+            fname="hubert_base_ls960_km200.pt",
+            sha256="b3c46c9cdd1707ad852dd53c359aa73942e79d53c432a9a8a419ed046408024b",
+        ),
+        Checkpoint(
+            name="hubert-base-ls960-kmeans-500",
+            remote_path="https://dl.fbaipublicfiles.com/textless_nlp/gslm/hubert/km500/km.bin",
+            fname="hubert_base_ls960_km500.pt",
+            sha256="411c8668e1314751404f58636f935fc73540a6793890435da4a8ffadf157398e",
+        ),
+        Checkpoint(
+            name="mhubert-base-vp_en_es_fr-layer-11-kmeans-1000",
+            remote_path="https://dl.fbaipublicfiles.com/hubert/mhubert_base_vp_en_es_fr_it3_L11_km1000.bin",
+            fname="mhubert_base_vp_en_es_fr_it3_L11_km1000.bin",
+            sha256="c70f5416c38044d0c20e33fdeedf4533974a8ac4029985cdef3b755038381a0b",
+        ),
+        Checkpoint(
+            name="hubert-base-ls960-layer-9-kmeans-500",
+            remote_path="https://dl.fbaipublicfiles.com/hubert/hubert_base_ls960_L9_km500.bin",
+            fname="hubert_base_ls960_L9_km500.pt",
+            sha256="4450dcbc90aa7d0553306a4b0fafef35c9b3c840468f02969273dfbb9db89a23",
+        ),
+        Checkpoint(
+            name="hubert-base-ls960-layer-9-kmeans-expresso-2000",
+            remote_path="https://dl.fbaipublicfiles.com/textless_nlp/expresso/checkpoints/hubert_base_ls960_L9_km2000_expresso.bin",
+            fname="hubert_base_ls960_L9_km2000_expresso.pt",
+            sha256="ae73e3cc7ea0426faf43c8bb67fc3846a0f8708232526d045a101836e8c7929f",
+        ),
+        Checkpoint(
+            name="mhubert-base-vp_mls_cv_8lang-kmeans-2000",
+            remote_path="https://dl.fbaipublicfiles.com/hubert/mhubert_base_vp_mls_cv_8lang_it3_L12_km2000.bin",
+            fname="mhubert_base_vp_mls_cv_8lang_it3_L12_km2000.pt",
+            sha256="05de230870427dcdce334d29922ccce354c7e7d0574edc645f5daed47b1ccc18",
+        ),
+        Checkpoint(
+            name="mhubert-base-vp_mls_cv_8lang-kmeans-expresso-2000",
+            remote_path="https://dl.fbaipublicfiles.com/textless_nlp/expresso/checkpoints/mhubert_base_vp_mls_cv_8lang_it3_L12_km2000_expresso.bin",
+            fname="mhubert_base_vp_mls_cv_8lang_it3_L12_km2000_expresso.pt",
+            sha256="bda2dbf5fe42bf59359799b281498a9d7fd0a9ec703d6dc54a4a4a0ff70a4e22",
+        ),
+        Checkpoint(
+            name="mhubert-base-25hz-kmeans-500",
+            remote_path="https://dl.fbaipublicfiles.com/textless_nlp/twist/speech_tokenizer/mhubert_base_25hz_cp_mls_cv_sp_fisher_L11_km500.bin",
+            fname="mhubert_base_25hz_cp_mls_cv_sp_fisher_L11_km500.bin",
+            sha256="03cc04a9c24fec4285e73e709c485756d8f116aa8e724eac555de6a7cf8d28ad",
+        ),
+        # Tacotron2
+        Checkpoint(
+            name="hubert-base-ls960-kmeans-50-tacotron",
+            remote_path="https://dl.fbaipublicfiles.com/textless_nlp/gslm/hubert/tts_km50/tts_checkpoint_best.pt",
+            fname="hubert_base_ls960_kmeans_50_tacotron.pt",
+            sha256="335e881a897cfa3389804110de8ac3909159d4de395880fbf1d3167a9477451e",
+        ),
+        Checkpoint(
+            name="hubert-base-ls960-kmeans-100-tacotron",
+            remote_path="https://dl.fbaipublicfiles.com/textless_nlp/gslm/hubert/tts_km100/tts_checkpoint_best.pt",
+            fname="hubert_base_ls960_kmeans_100_tacotron.pt",
+            sha256="b208f8d6433eb5524405aa29d2b5fdacddb63a182d9830b629232e63b3543e4d",
+        ),
+        Checkpoint(
+            name="hubert-base-ls960-kmeans-200-tacotron",
+            remote_path="https://dl.fbaipublicfiles.com/textless_nlp/gslm/hubert/tts_km200/tts_checkpoint_best.pt",
+            fname="hubert_base_ls960_kmeans_200_tacotron.pt",
+            sha256="093f009935a4dadd692db85e859246e369cb17be3aecd22038fb70af4d5b0590",
+        ),
+        # NB: must be named '*-codes'
+        Checkpoint(
+            name="hubert-base-ls960-kmeans-50-tacotron-codes",
+            remote_path="https://dl.fbaipublicfiles.com/textless_nlp/gslm/hubert/tts_km50/code_dict",
+            fname="hubert_base_ls960_kmeans_50_tacotron_codes.pt",
+            sha256="5f01dd57fd3b4044fac93aaac2589bf49e34cbe1dc0713254c0f339ba2123bce",
+        ),
+        Checkpoint(
+            name="hubert-base-ls960-kmeans-100-tacotron-codes",
+            remote_path="https://dl.fbaipublicfiles.com/textless_nlp/gslm/hubert/tts_km100/code_dict",
+            fname="hubert_base_ls960_kmeans_100_tacotron_codes.pt",
+            sha256="6d506216aa5bad159f167e2535293b4e5ec8e1073b64449d30b66b460ebf6da0",
+        ),
+        Checkpoint(
+            name="hubert-base-ls960-kmeans-200-tacotron-codes",
+            remote_path="https://dl.fbaipublicfiles.com/textless_nlp/gslm/hubert/tts_km200/code_dict",
+            fname="hubert_base_ls960_kmeans_200_tacotron_codes.pt",
+            sha256="ea01ba3592e27c871b63b32e37d6532234edf7eee7077bdcc094061ee72922e6",
+        ),
+        # Hifigan
+        Checkpoint(
+            name="mhubert-base-25hz-kmeans-500-hifigan",
+            remote_path="https://dl.fbaipublicfiles.com/textless_nlp/twist/speech_tokenizer/hifigan_lj_mhubert_base_25hz.pt",
+            fname="hifigan_lj_mhubert_base_25hz.pt",
+            sha256="d88224e95c501e2cd59a6e4014753169cfe060fb7ade3cc0da03c809fef73b79",
+        ),
+        Checkpoint(
+            name="mhubert-base-25hz-kmeans-500-hifigan-config",
+            remote_path="https://dl.fbaipublicfiles.com/textless_nlp/twist/speech_tokenizer/hifigan_lj_mhubert_base_25hz_config.json",
+            fname="hifigan_lj_mhubert_base_25hz_config.json",
+            sha256="116dc39be4970cd393e562acd0a0ec70a86aaba7e9ca469b8b836cb3fe46afc7",
+        ),
+    ]
+
+    # Expresso Hifigan models
+    checkpoints += [
+        # hubert_base_km500_L9
+        Checkpoint(
+            name="hubert-base-ls960-layer-9-kmeans-500-hifigan",
+            remote_path="https://dl.fbaipublicfiles.com/textless_nlp/expresso/checkpoints/hifigan_expresso_lj_vctk_hubert_base_ls960_L9_km500/generator.pt",
+            fname="hifigan_expresso_lj_vctk_hubert_base_ls960_L9_km500_generator.pt",
+            sha256="d50f4704e961c29cb386924ef32cd0c3a4d1089a3f9be856f8a3b6dcfdd6147a",
+        ),
+        Checkpoint(
+            name="hubert-base-ls960-layer-9-kmeans-500-hifigan-config",
+            remote_path="https://dl.fbaipublicfiles.com/textless_nlp/expresso/checkpoints/hifigan_expresso_lj_vctk_hubert_base_ls960_L9_km500/config.json",
+            fname="hifigan_expresso_lj_vctk_hubert_base_ls960_L9_km500_config.json",
+            sha256="31e65ef5e100cb43de08f2c8d225496ace1ed87a49fcb8ee44bd6805f079b1f0",
+        ),
+        Checkpoint(
+            name="hubert-base-ls960-layer-9-kmeans-500-hifigan-speakers",
+            remote_path="https://dl.fbaipublicfiles.com/textless_nlp/expresso/checkpoints/hifigan_expresso_lj_vctk_hubert_base_ls960_L9_km500/speakers.txt",
+            fname="hifigan_expresso_lj_vctk_hubert_base_ls960_L9_km500_speakers.txt",
+            sha256="2efabb7a4ad786da34fe6fbe99c3906b7abd40b2e01c81b7efde76142ae3d552",
+        ),
+        Checkpoint(
+            name="hubert-base-ls960-layer-9-kmeans-500-hifigan-styles",
+            remote_path="https://dl.fbaipublicfiles.com/textless_nlp/expresso/checkpoints/hifigan_expresso_lj_vctk_hubert_base_ls960_L9_km500/styles.txt",
+            fname="hifigan_expresso_lj_vctk_hubert_base_ls960_L9_km500_styles.txt",
+            sha256="a6f63206518448277c63c3a251ac13484784829cf2e1b0a23b0860dd32543520",
+        ),
+        # hubert_base_km2000_L9_expresso
+        Checkpoint(
+            name="hubert-base-ls960-layer-9-kmeans-expresso-2000-hifigan",
+            remote_path="https://dl.fbaipublicfiles.com/textless_nlp/expresso/checkpoints/hifigan_expresso_lj_vctk_hubert_base_ls960_L9_km2000_expresso/generator.pt",
+            fname="hifigan_expresso_lj_vctk_hubert_base_ls960_L9_km2000_expresso_generator.pt",
+            sha256="b892fdcca889f4bc6d7f0a0f407ae5d01eca8aaa66b9be333ef33a1c50f2ef52",
+        ),
+        Checkpoint(
+            name="hubert-base-ls960-layer-9-kmeans-expresso-2000-hifigan-config",
+            remote_path="https://dl.fbaipublicfiles.com/textless_nlp/expresso/checkpoints/hifigan_expresso_lj_vctk_hubert_base_ls960_L9_km2000_expresso/config.json",
+            fname="hifigan_expresso_lj_vctk_hubert_base_ls960_L9_km2000_expresso_config.json",
+            sha256="73932563cedc43cd1dcddfab36ce3b2c3cb76eef3423d25d1fb536ce1ec513bb",
+        ),
+        Checkpoint(
+            name="hubert-base-ls960-layer-9-kmeans-expresso-2000-hifigan-speakers",
+            remote_path="https://dl.fbaipublicfiles.com/textless_nlp/expresso/checkpoints/hifigan_expresso_lj_vctk_hubert_base_ls960_L9_km2000_expresso/speakers.txt",
+            fname="hifigan_expresso_lj_vctk_hubert_base_ls960_L9_km2000_expresso_speakers.txt",
+            sha256="2efabb7a4ad786da34fe6fbe99c3906b7abd40b2e01c81b7efde76142ae3d552",
+        ),
+        Checkpoint(
+            name="hubert-base-ls960-layer-9-kmeans-expresso-2000-hifigan-styles",
+            remote_path="https://dl.fbaipublicfiles.com/textless_nlp/expresso/checkpoints/hifigan_expresso_lj_vctk_hubert_base_ls960_L9_km2000_expresso/styles.txt",
+            fname="hifigan_expresso_lj_vctk_hubert_base_ls960_L9_km2000_expresso_styles.txt",
+            sha256="a6f63206518448277c63c3a251ac13484784829cf2e1b0a23b0860dd32543520",
+        ),
+        # mhubert_base_vp_mls_cv_8lang_it3_L12_km2000
+        Checkpoint(
+            name="mhubert-base-vp_mls_cv_8lang-kmeans-2000-hifigan",
+            remote_path="https://dl.fbaipublicfiles.com/textless_nlp/expresso/checkpoints/hifigan_expresso_lj_vctk_mhubert_base_vp_mls_cv_8lang_it3_L12_km2000/generator.pt",
+            fname="hifigan_expresso_lj_vctk_mhubert_base_vp_mls_cv_8lang_it3_L12_km2000_generator.pt",
+            sha256="0d160c97519d3ed644e2ec636581b79e56767d60325308849fbb389a494c1466",
+        ),
+        Checkpoint(
+            name="mhubert-base-vp_mls_cv_8lang-kmeans-2000-hifigan-config",
+            remote_path="https://dl.fbaipublicfiles.com/textless_nlp/expresso/checkpoints/hifigan_expresso_lj_vctk_mhubert_base_vp_mls_cv_8lang_it3_L12_km2000/config.json",
+            fname="hifigan_expresso_lj_vctk_mhubert_base_vp_mls_cv_8lang_it3_L12_km2000_config.json",
+            sha256="73932563cedc43cd1dcddfab36ce3b2c3cb76eef3423d25d1fb536ce1ec513bb",
+        ),
+        Checkpoint(
+            name="mhubert-base-vp_mls_cv_8lang-kmeans-2000-hifigan-speakers",
+            remote_path="https://dl.fbaipublicfiles.com/textless_nlp/expresso/checkpoints/hifigan_expresso_lj_vctk_mhubert_base_vp_mls_cv_8lang_it3_L12_km2000/speakers.txt",
+            fname="hifigan_expresso_lj_vctk_mhubert_base_vp_mls_cv_8lang_it3_L12_km2000_speakers.txt",
+            sha256="2efabb7a4ad786da34fe6fbe99c3906b7abd40b2e01c81b7efde76142ae3d552",
+        ),
+        Checkpoint(
+            name="mhubert-base-vp_mls_cv_8lang-kmeans-2000-hifigan-styles",
+            remote_path="https://dl.fbaipublicfiles.com/textless_nlp/expresso/checkpoints/hifigan_expresso_lj_vctk_mhubert_base_vp_mls_cv_8lang_it3_L12_km2000/styles.txt",
+            fname="hifigan_expresso_lj_vctk_mhubert_base_vp_mls_cv_8lang_it3_L12_km2000_styles.txt",
+            sha256="a6f63206518448277c63c3a251ac13484784829cf2e1b0a23b0860dd32543520",
+        ),
+        # mhubert_base_vp_mls_cv_8lang_it3_L12_km2000_expresso
+        Checkpoint(
+            name="mhubert-base-vp_mls_cv_8lang-kmeans-expresso-2000-hifigan",
+            remote_path="https://dl.fbaipublicfiles.com/textless_nlp/expresso/checkpoints/hifigan_expresso_lj_vctk_mhubert_base_vp_mls_cv_8lang_it3_L12_km2000_expresso/generator.pt",
+            fname="hifigan_expresso_lj_vctk_mhubert_base_vp_mls_cv_8lang_it3_L12_km2000_expresso_generator.pt",
+            sha256="8d1585f20edd9b0982f80acbb387a892adf19e4aa719b8dea795c8fa96285cbc",
+        ),
+        Checkpoint(
+            name="mhubert-base-vp_mls_cv_8lang-kmeans-expresso-2000-hifigan-config",
+            remote_path="https://dl.fbaipublicfiles.com/textless_nlp/expresso/checkpoints/hifigan_expresso_lj_vctk_mhubert_base_vp_mls_cv_8lang_it3_L12_km2000_expresso/config.json",
+            fname="hifigan_expresso_lj_vctk_mhubert_base_vp_mls_cv_8lang_it3_L12_km2000_expresso_config.json",
+            sha256="73932563cedc43cd1dcddfab36ce3b2c3cb76eef3423d25d1fb536ce1ec513bb",
+        ),
+        Checkpoint(
+            name="mhubert-base-vp_mls_cv_8lang-kmeans-expresso-2000-hifigan-speakers",
+            remote_path="https://dl.fbaipublicfiles.com/textless_nlp/expresso/checkpoints/hifigan_expresso_lj_vctk_mhubert_base_vp_mls_cv_8lang_it3_L12_km2000_expresso/speakers.txt",
+            fname="hifigan_expresso_lj_vctk_mhubert_base_vp_mls_cv_8lang_it3_L12_km2000_expresso_speakers.txt",
+            sha256="2efabb7a4ad786da34fe6fbe99c3906b7abd40b2e01c81b7efde76142ae3d552",
+        ),
+        Checkpoint(
+            name="mhubert-base-vp_mls_cv_8lang-kmeans-expresso-2000-hifigan-styles",
+            remote_path="https://dl.fbaipublicfiles.com/textless_nlp/expresso/checkpoints/hifigan_expresso_lj_vctk_mhubert_base_vp_mls_cv_8lang_it3_L12_km2000_expresso/styles.txt",
+            fname="hifigan_expresso_lj_vctk_mhubert_base_vp_mls_cv_8lang_it3_L12_km2000_expresso_styles.txt",
+            sha256="a6f63206518448277c63c3a251ac13484784829cf2e1b0a23b0860dd32543520",
+        ),
+    ]
+
+    # CPC-based stuff
+    checkpoints += [
+        # dense model
+        Checkpoint(
+            name="cpc-big-ll6k",
+            remote_path="https://dl.fbaipublicfiles.com/textless_nlp/gslm/cpc/cpc_big_ll6kh_top_ctc.pt",
+            fname="cpc_big_ll6kh_top_ctc.pt",
+            sha256="73155dad5d7c986fe7b7f548050060a8e9cc9a0ffd111a22932f38c3e617c5b8",
+        ),
+        # Quantizers
+        Checkpoint(
+            name="cpc-big-ll6k-kmeans-50",
+            remote_path="https://dl.fbaipublicfiles.com/textless_nlp/gslm/cpc/km50/km.bin",
+            fname="cpc_big_ll6k_km50.pt",
+            sha256="c48be5717aebc08169aa5165d58267449b8c8568624a346bb9f4b26eac3b0240",
+        ),
+        Checkpoint(
+            name="cpc-big-ll6k-kmeans-100",
+            remote_path="https://dl.fbaipublicfiles.com/textless_nlp/gslm/cpc/km100/km.bin",
+            fname="cpc_big_ll6k_km100.pt",
+            sha256="077b96e010b1e87be627ef2bef0f5e5cdaa1c01722aa363b2dc0bb6a638e8b26",
+        ),
+        Checkpoint(
+            name="cpc-big-ll6k-kmeans-200",
+            remote_path="https://dl.fbaipublicfiles.com/textless_nlp/gslm/cpc/km200/km.bin",
+            fname="cpc_big_ll6k_km200.pt",
+            sha256="2d863d1c6f251d19e667998248c87642ecb5e3ffaa410e9ebfa64341db4f6de2",
+        ),
+        Checkpoint(
+            name="cpc-big-ll6k-kmeans-500",
+            remote_path="https://dl.fbaipublicfiles.com/textless_nlp/gslm/cpc/km500/km.bin",
+            fname="cpc_big_ll6k_km500.pt",
+            sha256="35398f665dff06801a7bf1e595bc176711ed751d554e64d999dc42f1dd106561",
+        ),
+        # Tacotron2
+        Checkpoint(
+            name="cpc-big-ll6k-kmeans-50-tacotron",
+            remote_path="https://dl.fbaipublicfiles.com/textless_nlp/gslm/cpc/tts_km50/tts_checkpoint_best.pt",
+            fname="cpc_big_ll6k_kmeans_50_tacotron.pt",
+            sha256="e80a46561d1ded73bbe6e7272fcd6b9943fb7607229040b06324f654b66396fd",
+        ),
+        Checkpoint(
+            name="cpc-big-ll6k-kmeans-100-tacotron",
+            remote_path="https://dl.fbaipublicfiles.com/textless_nlp/gslm/cpc/tts_km100/tts_checkpoint_best.pt",
+            fname="cpc_big_ll6k_kmeans_100_tacotron.pt",
+            sha256="6cfe1ce4bbdfd0f531189f4013cefc9e1296aec62f8bc0172d68d088ab344a50",
+        ),
+        Checkpoint(
+            name="cpc-big-ll6k-kmeans-200-tacotron",
+            remote_path="https://dl.fbaipublicfiles.com/textless_nlp/gslm/cpc/tts_km200/tts_checkpoint_best.pt",
+            fname="cpc_big_ll6k_kmeans_200_tacotron.pt",
+            sha256="5845b1c0a82d9176b9d8c33a32f098795c946077e741a9da2b665bafda17b7ff",
+        ),
+        Checkpoint(
+            name="cpc-big-ll6k-kmeans-50-tacotron-codes",
+            remote_path="https://dl.fbaipublicfiles.com/textless_nlp/gslm/cpc/tts_km50/code_dict",
+            fname="cpc_big_ll6k_kmeans_50_tacotron_codes.pt",
+            sha256="5f01dd57fd3b4044fac93aaac2589bf49e34cbe1dc0713254c0f339ba2123bce",
+        ),
+        Checkpoint(
+            name="cpc-big-ll6k-kmeans-100-tacotron-codes",
+            remote_path="https://dl.fbaipublicfiles.com/textless_nlp/gslm/cpc/tts_km100/code_dict",
+            fname="cpc_big_ll6k_kmeans_100_tacotron_codes.pt",
+            sha256="a343085a83c1acfc96a2bfecf5098b380e860cdbb832cbcbbb0bf1358bd7c932",
+        ),
+        Checkpoint(
+            name="cpc-big-ll6k-kmeans-200-tacotron-codes",
+            remote_path="https://dl.fbaipublicfiles.com/textless_nlp/gslm/cpc/tts_km200/code_dict",
+            fname="cpc_big_ll6k_kmeans_200_tacotron_codes.pt",
+            sha256="48dae50f9f2d4561c4a84fa09ba748395808d7d3323b1d5754af71669117e2aa",
+        ),
+    ]
+
+    # Common
+    checkpoints += [
+        # Vocoder models
+        Checkpoint(
+            name="waveglow",
+            remote_path="https://dl.fbaipublicfiles.com/textless_nlp/gslm/waveglow_256channels_standalone.pt",
+            fname="waveglow_256channels_standalone.pt",
+            sha256="f383c7fd785502fc6a6bffd604fc14cb35d6155cdde30c53faaaeafa8a904dab",
+        ),
+    ]
+
+    for checkpoint in checkpoints:
+        CHECKPOINT_MANAGER.add_checkpoint(checkpoint)
+
+
+populate_checkpoints()
diff --git a/slamkit/vocoder/textless_checkpoint_manager/manager.py b/slamkit/vocoder/textless_checkpoint_manager/manager.py
new file mode 100644
index 0000000..b9e5207
--- /dev/null
+++ b/slamkit/vocoder/textless_checkpoint_manager/manager.py
@@ -0,0 +1,61 @@
+# Copied from textlesslib (https://github.com/facebookresearch/textlesslib)
+# Please refer to the MIT license there (https://github.com/facebookresearch/textlesslib?tab=MIT-1-ov-file)
+# Below is the original header from the file
+
+# Copyright (c) Facebook, Inc. and its affiliates.
+
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Union
+
+from dataclasses import dataclass
+from torch.hub import download_url_to_file
+import pathlib
+
+
+@dataclass
+class Checkpoint:
+    name: str
+    remote_path: str
+    fname: str
+    sha256: str
+
+
+class CheckpointManager:
+    def __init__(self, disk_root: Union[str, pathlib.Path] = "~/.textless/"):
+        self.disk_root = pathlib.Path(disk_root).expanduser().resolve()
+        if not self.disk_root.exists():
+            self.disk_root.mkdir()
+
+        self.storage: dict[str, Checkpoint] = {}
+
+    def add_checkpoint(self, checkpoint: Checkpoint) -> None:
+        name = checkpoint.name
+        assert name not in self.storage
+        self.storage[name] = checkpoint
+
+    def download_by_name(self, name: str) -> None:
+        checkpoint = self.storage[name]
+        download_url_to_file(
+            url=checkpoint.remote_path,
+            dst=(self.disk_root / checkpoint.fname),
+            hash_prefix=checkpoint.sha256,
+            progress=True,
+        )
+
+    def get_by_name(self, name: str, download_if_needed: bool = True) -> pathlib.Path:
+        checkpoint = self.storage[name]
+        disk_name = self.disk_root / checkpoint.fname
+
+        if not disk_name.exists():
+            if download_if_needed:
+                self.download_by_name(name)
+            else:
+                raise FileNotFoundError(
+                    f"Checkpoint {checkpoint} was not found locally at {disk_name}, please set `allow_download` flag"
+                )
+        return disk_name
+
+    def set_root(self, new_root: Union[str, pathlib.Path]) -> None:
+        self.disk_root = pathlib.Path(new_root)


Reviews:
user: login: gallilmaimon
body: Thanks for the quick PR! All in all this looks really good to me. I had 2 small comments.

Also I'd appreciate it if you could run some tests to show that the results are identical to before, for instance take a couple of audio files and then encode them and vocode them and see that the textless version and the local version produce the same results. While this should be the case this integration test is important for a key part of the pipeline :)

I will also get @avishaiElmakies to take a look to see if he has further comments. 

Thanks again!
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-04T08:57:12Z
user: login: avishaiElmakies
body: @cantabile-kwok  
First of all, thanks for the PR! It is very much appreciated!
I left one comment about the name of `textless_hifigan`. 
after that it looks good to me
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-04T18:42:21Z
user: login: cantabile-kwok
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-05T03:27:19Z

Review comments:
diff_hunk: @@ -56,7 +56,7 @@ def main(cfg: DictConfig):
             if key == "generate" or key == "prompts":
                 continue
             print(f"{key}: {val}")
-    if cfg.metric.get("out_path", False) and "generate" in res:
+    if cfg.metric.get("out_path", False) and "generate" in res and cfg.vocoder is not None:
path: cli/eval.py
user: login: gallilmaimon
body: Thanks for adding this, we noticed this unexpected behaviour when there is no vocoder specified.
I think it is worth adding a logger.warning message saying something like:
"You are currently trying to run generation without a vocoder, which will generate tokens, but have no effect. You probably want to set `vocoder=vocoder_hubert_25`"

Or something like that so that this doesn't just fail silently without an indication. We could also have this as an assertion, but okay to leave as a warning in case we would want this in the future. 
created_at: 2025-03-04T08:38:23Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,114 @@
+# Residual block from the Hifi-gan model:
path: slamkit/vocoder/textless_hifigan/resblock.py
user: login: gallilmaimon
body: I think you missed the general "copied from textlesslib" message here, right?
created_at: 2025-03-04T08:42:49Z
author_association: COLLABORATOR
diff_hunk: @@ -11,7 +11,7 @@ def __init__(self,
                  style_meta: str = None,
                  ):
         super().__init__()
-        from textless.vocoders.hifigan.vocoder import CodeHiFiGANVocoder
+        from .textless_hifigan.vocoder import CodeHiFiGANVocoder
path: slamkit/vocoder/hifi_gan_vocoder.py
user: login: avishaiElmakies
body: I think that the name `textless_hifigan` is redundant. can we change it to `hifigan`? 
created_at: 2025-03-04T18:42:12Z
author_association: COLLABORATOR
diff_hunk: @@ -11,7 +11,7 @@ def __init__(self,
                  style_meta: str = None,
                  ):
         super().__init__()
-        from textless.vocoders.hifigan.vocoder import CodeHiFiGANVocoder
+        from .textless_hifigan.vocoder import CodeHiFiGANVocoder
path: slamkit/vocoder/hifi_gan_vocoder.py
user: login: cantabile-kwok
body: Yeah, I think you are right. Changed
created_at: 2025-03-05T03:27:19Z
author_association: CONTRIBUTOR

PR comments:
user: login: cantabile-kwok
created_at: 2025-03-04T09:03:09Z
body: > Thanks for the quick PR! All in all this looks really good to me. I had 2 small comments.
> 
> Also I'd appreciate it if you could run some tests to show that the results are identical to before, for instance take a couple of audio files and then encode them and vocode them and see that the textless version and the local version produce the same results. While this should be the case this integration test is important for a key part of the pipeline :)
> 
> I will also get @avishaiElmakies to take a look to see if he has further comments.
> 
> Thanks again!

Of course, I will take a look to see whether this produces identical results.
author_association: CONTRIBUTOR
user: login: cantabile-kwok
created_at: 2025-03-04T09:42:14Z
body: @gallilmaimon I have tested the behavior and it went all well. My test code is as follows. I did not commit this testing script as it is not necessary.

```python
from slamkit.vocoder.textless_hifigan.vocoder import CodeHiFiGANVocoder as CodeHiFiGANVocoder_local
import torch
from tqdm import trange

vocoder_local = CodeHiFiGANVocoder_local.by_name("mhubert-base-25hz", "kmeans", 500)

from textless.vocoders.hifigan.vocoder import CodeHiFiGANVocoder as CodeHiFiGANVocoder_online
vocoder_online = CodeHiFiGANVocoder_online.by_name("mhubert-base-25hz", "kmeans", 500)

# test 1. load encoded waveforms, vocode, and save to audio
import json
import torchaudio
path = "example_data/features.jsonl"
with open(path) as f:
    data = [json.loads(line) for line in f]
    for i, sample in enumerate(data):
        tokens = torch.LongTensor(sample["units"]).unsqueeze(0)
        wave = vocoder_local(tokens, dur_prediction=(vocoder_local.model.dur_predictor is not None))
        torchaudio.save(f"local_{i}.wav", wave.unsqueeze(0), 16000)
        wave_gt = vocoder_online(tokens, dur_prediction=(vocoder_online.model.dur_predictor is not None))
        torchaudio.save(f"online_{i}.wav", wave_gt.unsqueeze(0), 16000)

# test 2. vocode random indexes and compare with online vocoder
device = "mps"
vocoder_local.to(device).eval()
vocoder_online.to(device).eval()

with torch.no_grad():
    for _ in trange(50):
        L = torch.randint(0, 1000, (1,))
        tokens = torch.randint(0, 500, (1, L.item())).to(device)
        res_local = vocoder_local(tokens, dur_prediction=(vocoder_local.model.dur_predictor is not None))
        res_online = vocoder_online(tokens, dur_prediction=(vocoder_online.model.dur_predictor is not None))
        assert torch.allclose(res_local, res_online, atol=1e-5), f"Local and online vocoder results are different"
```

I've also added the missing header, and the warning message in `cli/eval.py`. Please check if the warning is placed at a proper position.
author_association: CONTRIBUTOR
user: login: gallilmaimon
created_at: 2025-03-05T07:35:47Z
body: @cantabile-kwok Awesome, thanks again for this PR! merging:)
author_association: COLLABORATOR