speechbrain · mravanelli · Nov 6, 2023 · Nov 4, 2023 · Nov 4, 2023 · Nov 5, 2023
diff --git a/conftest.py b/conftest.py
@@ -26,12 +26,22 @@ def pytest_generate_tests(metafunc):
 try:
     from transformers import Wav2Vec2Model  # noqa: F401
 except ModuleNotFoundError:
-    collect_ignore.append("speechbrain/lobes/models/huggingface_wav2vec.py")
+    collect_ignore.append(
+        "speechbrain/lobes/models/huggingface_transformers/wav2vec2.py"
+    )
 try:
     from transformers import WhisperModel  # noqa: F401
 except ModuleNotFoundError:
-    collect_ignore.append("speechbrain/lobes/models/huggingface_whisper.py")
+    collect_ignore.append(
+        "speechbrain/lobes/models/huggingface_transformers/whisper.py"
+    )
 try:
     import sacrebleu  # noqa: F401
 except ModuleNotFoundError:
     collect_ignore.append("speechbrain/utils/bleu.py")
+try:
+    import vocos  # noqa: F401
+except ModuleNotFoundError:
+    collect_ignore.append(
+        "speechbrain/lobes/models/huggingface_transformers/vocos.py"
+    )
diff --git a/speechbrain/lobes/models/huggingface_transformers/__init__.py b/speechbrain/lobes/models/huggingface_transformers/__init__.py
@@ -18,3 +18,4 @@
 from .wav2vec2 import *  # noqa
 from .wavlm import *  # noqa
 from .whisper import *  # noqa
+from .encodec import *  # noqa
diff --git a/speechbrain/lobes/models/huggingface_transformers/encodec.py b/speechbrain/lobes/models/huggingface_transformers/encodec.py
@@ -0,0 +1,233 @@
+"""This lobe enables the integration of huggingface pretrained EnCodec.
+
+EnCodec makes it possible to compress audio into a sequence of discrete tokens
+at different bandwidths - and to reconstruct audio from such sequences, with
+some loss of quality depending on the bandwidth.
+
+Note that while encodec can be used to reconstruct speech data, for a
+high-quality reconstruction, it is recommended to use a specially trained
+vocoder, such as Vocos (speechbrain.lobes.models.huggingface_transformers.vocos)
+
+Repository: https://huggingface.co/docs/transformers/v4.31.0/en/model_doc/encodec
+Paper: https://arxiv.org/abs/2210.13438
+
+Authors
+ * Artem Ploujnikov 2023
+"""
+
+import torch
+import logging
+from torch.nn import functional as F
+from speechbrain.dataio.dataio import length_to_mask
+from speechbrain.lobes.models.huggingface_transformers.huggingface import (
+    HFTransformersInterface,
+)
+
+DEFAULT_SAMPLE_RATE = 24000
+
+logger = logging.getLogger(__name__)
+
+
+class Encodec(HFTransformersInterface):
+    """An wrapper for the HuggingFace encodec model
+
+    Arguments
+    ---------
+    source : str
+        A HuggingFace repository identifier or a path
+    save_path : str
+        The location where the pretrained model will be saved
+    sample_rate : int
+        The audio sampling rate
+    bandwidth : float
+        The encoding bandwidth, in kbps (optional)
+        Supported bandwidths:
+        1.5, 3.0, 6.0, 12.0, 24.0
+    flat_embeddings : bool
+        If set to True, embeddings will be flattened into
+        (Batch x Length x (Heads * Embedding))
+    freeze : bool
+        whether the model will be frozen (e.g. not trainable if used
+        as part of training another model)
+
+    Example
+    -------
+    >>> model_hub = "facebook/encodec_24khz"
+    >>> save_path = "savedir"
+    >>> model = Encodec(model_hub, save_path)
+    >>> audio = torch.randn(4, 1000)
+    >>> length = torch.tensor([1.0, .5, .75, 1.0])
+    >>> tokens, emb = model.encode(audio, length)
+    >>> tokens.shape
+    torch.Size([4, 4, 2])
+    >>> emb.shape
+    torch.Size([4, 4, 2, 128])
+    >>> rec = model.decode(tokens, length)
+    >>> rec.shape
+    torch.Size([4, 1, 1280])
+    >>> rec_emb = model.decode_emb(emb, length)
+    >>> rec_emb.shape
+    torch.Size([4, 1, 1280])
+    """
+
+    def __init__(
+        self,
+        source,
+        save_path=None,
+        sample_rate=None,
+        bandwidth=1.5,
+        flat_embeddings=False,
+        freeze=True,
+    ):
+        super().__init__(source=source, save_path=save_path, freeze=freeze)
+        if not sample_rate:
+            sample_rate = DEFAULT_SAMPLE_RATE
+        self.sample_rate = sample_rate
+        self.bandwidth = bandwidth
+        self.flat_embeddings = flat_embeddings
+        self.num_heads = self.model.quantizer.get_num_quantizers_for_bandwidth(
+            bandwidth
+        )
+        quantizer_layers = self.model.quantizer.layers[: self.num_heads]
+        self.vocabulary = torch.stack(
+            [layer.codebook.embed for layer in quantizer_layers]
+        )
+        _, self.num_tokens, self.emb_dim = self.vocabulary.shape
+        self.vocabulary_flat = self.vocabulary.reshape(
+            self.num_heads * self.num_tokens, self.emb_dim
+        )
+        self.token_index_offsets = (
+            torch.arange(self.num_heads)[None, None, :] * self.num_tokens
+        )
+        if self.freeze:
+            logger.warning("huggingface_Encodec - Encodec is frozen.")
+            for param in self.model.parameters():
+                param.requires_grad = False
+
+    def forward(self, inputs, length):
+        """Encodes the input audio as tokens
+
+        Arguments
+        ---------
+        inputs : torch.Tensor
+            A (Batch x Samples) or (Batch x Channel x Samples)
+            tensor of audio
+        length : torch.Tensor
+            A tensor of relative lengths
+
+        Returns
+        -------
+        tokens : torch.Tensor
+            a (Batch X Tokens) tensor of audio tokens
+        """
+        return self.encode(inputs, length)
+
+    def encode(self, inputs, length):
+        """Encodes the input audio as tokens
+
+        Arguments
+        ---------
+        inputs : torch.Tensor
+            A (Batch x Samples) or (Batch x Channel x Samples)
+            tensor of audio
+        length : torch.Tensor
+            A tensor of relative lengths
+
+        Returns
+        -------
+        tokens : torch.Tensor
+            A (Batch x Tokens x Heads) tensor of audio tokens
+        emb : torch.Tensor
+            Raw vector embeddings from the model's
+            quantizers
+        """
+        with torch.set_grad_enabled(not self.freeze):
+            if inputs.dim() == 2:
+                inputs = inputs.unsqueeze(1)
+            max_len = inputs.size(-1)
+            mask = length_to_mask(
+                length * max_len, max_len, device=inputs.device
+            ).unsqueeze(1)
+            result = self.model.encode(inputs, mask, bandwidth=self.bandwidth)
+            tokens = result.audio_codes.squeeze(0).transpose(-1, -2)
+            emb = self.embeddings(tokens)
+            return tokens, emb
+
+    def embeddings(self, tokens):
+        """Converts token indexes to vector embeddings
+
+        Arguments
+        ---------
+        tokens : torch.Tensor
+            a (Batch x Length x Heads) tensor of token indexes
+
+        Returns
+        -------
+        emb : torch.Tensor
+            a (Batch x Length x Heads x Embedding) tensor
+            of raw vector embeddings from the model's
+            quantizer codebooks
+        """
+        idx = tokens + self.token_index_offsets
+        emb = F.embedding(idx, self.vocabulary_flat)
+        if self.flat_embeddings:
+            batch_size, max_len, num_heads, emb_dim = emb.shape
+            emb = emb.reshape(batch_size, max_len, num_heads * emb_dim)
+        return emb
+
+    def decode(self, tokens, length=None):
+        """Decodes audio from tokens
+
+        Arguments
+        ---------
+        tokens : torch.Tensor
+            A (Batch x Length x Heads) tensor of audio tokens
+        length : torch.Tensor
+            A 1-D tensor of relative lengths
+
+        Returns
+        -------
+        audio : torch.Tensor
+            the reconstructed audio
+        """
+        with torch.set_grad_enabled(not self.freeze):
+            result = self.model.decode(
+                tokens.unsqueeze(0).transpose(-1, -2), [None]
+            )
+            audio = result.audio_values
+            if length is not None:
+                max_len = audio.size(-1)
+                mask = length_to_mask(
+                    length * max_len, max_len, device=tokens.device
+                ).unsqueeze(1)
+                audio = audio * mask
+            return audio
+
+    def decode_emb(self, emb, length):
+        """Decodes raw vector embeddings into audio
+
+        Arguments
+        ---------
+        emb : torch.Tensor
+            A (Batch x Length x Heads x Embedding) tensor of
+            raw vector embeddings
+
+        Returns
+        -------
+        audio : torch.Tensor
+            the reconstructed audio
+        """
+        with torch.set_grad_enabled(not self.freeze):
+            if self.flat_embeddings:
+                batch_size, max_len, _ = emb.shape
+                emb = emb.reshape(
+                    batch_size, max_len, self.num_heads, self.emb_dim
+                )
+            scaled_states = emb.pow(2).sum(-1, keepdim=True)
+            vocab = self.vocabulary.transpose(-1, -2).unsqueeze(0)
+            emb_perm = emb.permute(0, 2, 1, 3)
+            emb_vocab_prod = (emb_perm @ vocab).moveaxis(1, 2)
+            vocab_sum = vocab.pow(2).sum(-2, keepdim=True).moveaxis(1, 2)
+            dist = -(scaled_states - 2 * emb_vocab_prod + vocab_sum)
+            tokens = dist.max(dim=-1).indices
+            return self.decode(tokens, length)