Add new spectral codec definition (#14794)

* [TTS] Add new spectral codec definition

Signed-off-by: Ryan <[email protected]>

* Add codec MMD loss definitions

Signed-off-by: Ryan <[email protected]>

* Apply isort and black reformatting

Signed-off-by: rlangman <[email protected]>

---------

Signed-off-by: Ryan <[email protected]>
Signed-off-by: rlangman <[email protected]>
Signed-off-by: Jason <[email protected]>
Co-authored-by: Jason <[email protected]>
Co-authored-by: rlangman <[email protected]>

Author: rlangman

nemo/collections/common/parts/utils.py

@@ -159,12 +159,16 @@ def mask_sequence_tensor(tensor: torch.Tensor, lengths: torch.Tensor):
 
 class ClampActivation(nn.Module):
 
-    def __init__(self, min_value: float = -1.0, max_value: float = 1.0):
+    def __init__(self, min_value: float = -1.0, max_value: float = 1.0, clamp_training: bool = True):
         super().__init__()
         self.min_value = min_value
         self.max_value = max_value
+        self.clamp_training = clamp_training
 
     def forward(self, input: torch.Tensor) -> torch.Tensor:
+        if self.training and not self.clamp_training:
+            return input
+
         return torch.clamp(input, min=self.min_value, max=self.max_value)
 
 

nemo/collections/tts/losses/audio_codec_loss.py

@@ -512,3 +512,176 @@ def forward(self, disc_scores_real, disc_scores_gen):
         loss /= len(disc_scores_real)
 
         return loss
+
+
+class MMDLoss(Loss):
+    """
+    Maximum mean discrepancy (MMD) loss, as defined in https://arxiv.org/abs/2406.02315
+
+    Args:
+        kernel_radii: List of radii for Gaussian kernels
+        loss_scale: Constant to multiply loss by
+    """
+
+    def __init__(self, kernel_radii=(0.1, 1, 5, 10, 20, 50), loss_scale=1.0):
+        super().__init__()
+        self.kernel_radii = kernel_radii
+        self.loss_scale = loss_scale
+
+    @staticmethod
+    def _exp_kernel(dxx, r):
+        return torch.exp((-0.5 / r) * dxx).sum()
+
+    @staticmethod
+    def _shuffle_codebooks(x):
+        B, C, _ = x.size()
+        x_shuffled = torch.zeros_like(x)
+        for c in range(C):
+            batch_perm = torch.randperm(B, device=x.device)
+            x_shuffled[:, c, :] = x[batch_perm, c, :]
+        return x_shuffled
+
+    @property
+    def input_types(self):
+        return {
+            "inputs": [NeuralType(('B', 'C', 'D'), VoidType())],
+        }
+
+    @property
+    def output_types(self):
+        return {"loss": NeuralType(elements_type=LossType())}
+
+    @typecheck()
+    def forward(self, inputs):
+        B, C, D = inputs.size()
+
+        x = inputs
+        x_mean = x.mean(dim=(0,), keepdim=True)
+        x_stdev = torch.sqrt(x.var(dim=(0,), keepdim=True) + 1e-8)
+        x = (x - x_mean) / x_stdev
+        y = self._shuffle_codebooks(x)
+
+        # [B, C * D]
+        x = x.reshape([B, C * D])
+        y = y.reshape([B, C * D])
+
+        # [B, B]
+        xx = torch.mm(x, x.t())
+        yy = torch.mm(y, y.t())
+        zz = torch.mm(x, y.t())
+
+        rx = xx.diag().unsqueeze(0).expand_as(xx)
+        ry = yy.diag().unsqueeze(0).expand_as(yy)
+
+        dxx = rx.t() + rx - 2.0 * xx
+        dyy = ry.t() + ry - 2.0 * yy
+        dxy = rx.t() + ry - 2.0 * zz
+
+        loss = 0.0
+        coeff = -2.0 / B**2
+        denom = B * (B - 1)
+        for r in self.kernel_radii:
+            loss += (torch.utils.checkpoint.checkpoint(self._exp_kernel, dxx, r) - B) / denom
+            loss += coeff * torch.utils.checkpoint.checkpoint(self._exp_kernel, dxy, r)
+            loss += (torch.utils.checkpoint.checkpoint(self._exp_kernel, dyy, r) - B) / denom
+
+        loss = loss.clamp(min=0)
+        loss = self.loss_scale * loss
+        return loss
+
+
+class MMDCodebookLoss(Loss):
+    """
+    MMD loss which incentivizes independence between codebooks within each timestep.
+
+    Args:
+        num_codebooks: Number of codebooks.
+        codebook_dim: Dimension of a single codebook code.
+        loss_fn: MMDLoss instance.
+    """
+
+    def __init__(self, num_codebooks, codebook_dim, loss_fn):
+        super().__init__()
+        self.num_codebooks = num_codebooks
+        self.codebook_dim = codebook_dim
+        self.loss_fn = loss_fn
+
+    @property
+    def input_types(self):
+        return {
+            "inputs": [NeuralType(('B', 'D', 'T'), VoidType())],
+        }
+
+    @property
+    def output_types(self):
+        return {"loss": NeuralType(elements_type=LossType())}
+
+    @typecheck()
+    def forward(self, inputs):
+        B, D, T = inputs.size()
+
+        # [B, C, D / C, T]
+        x = inputs.reshape(B, self.num_codebooks, self.codebook_dim, T)
+        # [B*T, C, D / C]
+        x = rearrange(x, 'B C D T -> (B T) C D')
+        loss = self.loss_fn(inputs=x)
+        return loss
+
+
+class MMDEmbeddingLoss(Loss):
+    """
+    MMD loss which incentivizes independence between embedding values within each timestep.
+
+    Args:
+        loss_fn: MMDLoss instance.
+    """
+
+    def __init__(self, loss_fn):
+        super().__init__()
+        self.loss_fn = loss_fn
+
+    @property
+    def input_types(self):
+        return {
+            "inputs": [NeuralType(('B', 'D', 'T'), VoidType())],
+        }
+
+    @property
+    def output_types(self):
+        return {"loss": NeuralType(elements_type=LossType())}
+
+    @typecheck()
+    def forward(self, inputs):
+        # [B*T, 1, D]
+        x = rearrange(inputs, 'B D T -> (B T) D 1')
+        loss = self.loss_fn(inputs=x)
+        return loss
+
+
+class MMDTimeLoss(Loss):
+    """
+    MMD loss which incentivizes independence between different timesteps.
+
+    Args:
+        loss_fn: MMDLoss instance.
+    """
+
+    def __init__(self, loss_fn):
+        super().__init__()
+        self.loss_fn = loss_fn
+
+    @property
+    def input_types(self):
+        return {
+            "inputs": [NeuralType(('B', 'D', 'T'), VoidType())],
+        }
+
+    @property
+    def output_types(self):
+        return {"loss": NeuralType(elements_type=LossType())}
+
+    @typecheck()
+    def forward(self, inputs):
+        x = rearrange(inputs, 'B D T -> B T D')
+        loss = self.loss_fn(inputs=x)
+        return loss

nemo/collections/tts/models/audio_codec.py

@@ -143,6 +143,22 @@ def __init__(self, cfg: DictConfig, trainer: Trainer = None):
         self.gen_loss_fn = instantiate(cfg.generator_loss)
         self.disc_loss_fn = instantiate(cfg.discriminator_loss)
 
+        self.mmd_loss_start_epoch = cfg.get("mmd_loss_start_epoch", 0)
+
+        if "mmd_loss" in cfg:
+            self.mmd_loss_fn = instantiate(cfg.mmd_loss)
+            self.mmd_loss_scale = cfg.get("mmd_loss_scale", 1.0)
+        else:
+            self.mmd_loss_fn = None
+            self.mmd_loss_scale = None
+
+        if "mmd_time_loss" in cfg:
+            self.mmd_time_loss_fn = instantiate(cfg.mmd_time_loss)
+            self.mmd_time_loss_scale = cfg.get("mmd_time_loss_scale", 1.0)
+        else:
+            self.mmd_time_loss_fn = None
+            self.mmd_time_loss_scale = None
+
         feature_loss_type = cfg.get("feature_loss_type", "relative")
         if feature_loss_type == "relative":
             self.feature_loss_fn = RelativeFeatureMatchingLoss()
@@ -497,7 +513,7 @@ def _process_batch(self, batch):
         encoded = encoded.to(self.dtype)  # make sure vector quantizer output is in the model dtype
         audio_gen, _ = self.audio_decoder(inputs=encoded, input_len=encoded_len)
 
-        return audio, audio_len, audio_gen, commit_loss
+        return audio, audio_len, audio_gen, commit_loss, encoded
 
     @property
     def disc_update_prob(self) -> float:
@@ -514,7 +530,7 @@ def should_update_disc(self, batch_idx) -> bool:
     def training_step(self, batch, batch_idx):
         optim_gen, optim_disc = self.optimizers()
 
-        audio, audio_len, audio_gen, commit_loss = self._process_batch(batch)
+        audio, audio_len, audio_gen, commit_loss, codes = self._process_batch(batch)
 
         metrics = {
             "global_step": self.global_step,
@@ -578,6 +594,19 @@ def training_step(self, batch, batch_idx):
             metrics["g_loss_commit"] = commit_loss
             generator_losses.append(self.commit_loss_scale * commit_loss)
 
+        if self.mmd_loss_scale:
+            loss_mmd = self.mmd_loss_fn(inputs=codes)
+            metrics["g_loss_mmd"] = loss_mmd
+
+            if self.current_epoch >= self.mmd_loss_start_epoch:
+                generator_losses.append(self.mmd_loss_scale * loss_mmd)
+
+        if self.mmd_time_loss_scale:
+            loss_mmd_time = self.mmd_time_loss_fn(inputs=codes)
+            metrics["g_loss_mmd_time"] = loss_mmd_time
+            if self.current_epoch >= self.mmd_loss_start_epoch:
+                generator_losses.append(self.mmd_time_loss_scale * loss_mmd_time)
+
         # compute embeddings for speaker consistency loss
         if self.use_scl_loss:
             # concate generated and GT waveforms
@@ -623,7 +652,7 @@ def on_train_epoch_end(self):
         self.update_lr("epoch")
 
     def validation_step(self, batch, batch_idx):
-        audio, audio_len, audio_gen, _ = self._process_batch(batch)
+        audio, audio_len, audio_gen, _, _ = self._process_batch(batch)
 
         loss_mel_l1, loss_mel_l2 = self.mel_loss_fn(
             audio_real=audio.float(), audio_gen=audio_gen.float(), audio_len=audio_len

nemo/collections/tts/models/magpietts.py

@@ -36,6 +36,7 @@
 from nemo.collections.tts.models import AudioCodecModel
 from nemo.collections.tts.modules import transformer_2501
 from nemo.collections.tts.modules.aligner import AlignmentEncoder
+from nemo.collections.tts.modules.audio_codec_modules import VectorQuantizerIndexConverter
 from nemo.collections.tts.modules.magpietts_modules import (
     CharAwareSubwordEncoder,
     EOSDetectionMethod,
@@ -95,17 +96,32 @@ def __init__(self, cfg: DictConfig, trainer: 'Trainer' = None):
 
         # load codec
         codec_model = AudioCodecModel.restore_from(cfg.get('codecmodel_path'), strict=False)
+
         self.sample_rate = codec_model.sample_rate
+        self.codec_model_samples_per_frame = codec_model.samples_per_frame
         # del codec discriminator to free memory
         del codec_model.discriminator
 
-        # Set up codebook configuration
-        self.num_audio_codebooks = codec_model.num_codebooks
-        self.codec_model_samples_per_frame = codec_model.samples_per_frame
+        # When using FSQ tokens, the codebook structure can be changed at any time.
+        # An FSQ definition can be provided in `vector_quantizer` config to train with a codebook structure
+        # that is different than in the audio codec checkpoint.
+        vector_quantizer = cfg.get('vector_quantizer')
+        if vector_quantizer is not None:
+            vector_quantizer = instantiate(vector_quantizer)
+            self.num_audio_codebooks = vector_quantizer.num_codebooks
+            self.codebook_size = vector_quantizer.codebook_size
+            codec_converter = VectorQuantizerIndexConverter(
+                vector_quantizer_original=codec_model.vector_quantizer,
+                vector_quantizer_new=vector_quantizer,
+            )
+        else:
+            self.num_audio_codebooks = codec_model.num_codebooks
+            self.codebook_size = codec_model.codebook_size
+            codec_converter = None
+
         # Our codebooks start with actual audio codec tokens, followed by special tokens.
         # The `forced_*` options are for backward compatibility for models trained with older code.
-        num_audio_tokens = codec_model.codebook_size
-        get_token_index = partial(SpecialAudioToken.get_index, base_codebook_size=num_audio_tokens)
+        get_token_index = partial(SpecialAudioToken.get_index, base_codebook_size=self.codebook_size)
         self.audio_bos_id = cfg.get('forced_audio_bos_id', get_token_index(SpecialAudioToken.AUDIO_BOS))
         self.audio_eos_id = cfg.get('forced_audio_eos_id', get_token_index(SpecialAudioToken.AUDIO_EOS))
         self.context_audio_bos_id = cfg.get(
@@ -116,7 +132,7 @@ def __init__(self, cfg: DictConfig, trainer: 'Trainer' = None):
         )
         self.mask_token_id = cfg.get('forced_mask_token_id', get_token_index(SpecialAudioToken.MASK_TOKEN))
         self.num_all_tokens_per_codebook = cfg.get(
-            'forced_num_all_tokens_per_codebook', num_audio_tokens + len(SpecialAudioToken)
+            'forced_num_all_tokens_per_codebook', self.codebook_size + len(SpecialAudioToken)
         )
         self.use_bpe_char_tokenizer = cfg.get('use_bpe_char_tokenizer', False)
 
@@ -201,6 +217,7 @@ def __init__(self, cfg: DictConfig, trainer: 'Trainer' = None):
         # This needs to happen after super().__init__()
         self._codec_model = codec_model
         self._codec_model.freeze()  # Lightning does requires_grad = False and self.eval()
+        self._codec_converter = codec_converter
 
         audio_embeddings = []
         for _ in range(self.num_audio_codebooks * self.frame_stacking_factor):
@@ -450,6 +467,8 @@ def audio_to_codes(self, audio, audio_len, audio_type='target'):
         self._codec_model.eval()
         with torch.no_grad(), torch.autocast(device_type=audio.device.type, dtype=torch.float32):
             codes, codes_len = self._codec_model.encode(audio=audio, audio_len=audio_len)
+            if self._codec_converter is not None:
+                codes = self._codec_converter.convert_original_to_new(audio_tokens=codes, audio_lens=codes_len)
             # Add a timestep to begining and end of codes tensor
             bos_tensor = torch.full(
                 (codes.size(0), codes.size(1), 1), audio_bos_id, dtype=codes.dtype, device=codes.device
@@ -478,6 +497,10 @@ def codes_to_audio(self, codes, codes_len):
             codes_copy[codes == self.audio_bos_id] = 0  # zero is the padding token
             codes_copy[codes == self.audio_eos_id] = 0
             # Pass the modified integer token IDs
+            if self._codec_converter is not None:
+                codes_copy = self._codec_converter.convert_new_to_original(
+                    audio_tokens=codes_copy, audio_lens=codes_len
+                )
             audio, audio_len = self._codec_model.decode(tokens=codes_copy, tokens_len=codes_len)
             # audio: (B, T)
             # audio_len: (B,)
@@ -744,7 +767,7 @@ def clear_forbidden_logits(self, logits, forbid_audio_eos=False):
         logits[
             :,
             :,
-            SpecialAudioToken.get_forbidden_tokens(self._codec_model.codebook_size, forbid_audio_eos=forbid_audio_eos),
+            SpecialAudioToken.get_forbidden_tokens(self.codebook_size, forbid_audio_eos=forbid_audio_eos),
         ] = float('-inf')
         return logits
 
@@ -1276,6 +1299,10 @@ def prepare_context_tensors(self, batch):
             if 'context_audio_codes' in batch:
                 context_audio_codes = batch['context_audio_codes']
                 context_audio_codes_lens = batch['context_audio_codes_lens']
+                if self._codec_converter is not None:
+                    context_audio_codes = self._codec_converter.convert_original_to_new(
+                        audio_tokens=context_audio_codes, audio_lens=context_audio_codes_lens
+                    ).long()
             else:
                 context_audio_codes, context_audio_codes_lens = self.audio_to_codes(
                     batch['context_audio'], batch['context_audio_lens'], audio_type='context'
@@ -1498,6 +1525,10 @@ def process_batch(self, batch, mode="train"):
         else:
             audio_codes = batch['audio_codes']
             audio_codes_lens = batch['audio_codes_lens']
+            if self._codec_converter:
+                audio_codes = self._codec_converter.convert_original_to_new(
+                    audio_tokens=audio_codes, audio_lens=audio_codes_lens
+                ).long()
         if self.frame_stacking_factor > 1:
             # repeat the BOS token to frame_stacking_factor times. This is necessary since at inference
             # we need to start autoregressive generation from a full stack indicating BOS.
@@ -2326,6 +2357,7 @@ def infer_batch(
                 all_codes_next_argmax = self.sample_codes_from_logits(
                     all_code_logits_t,
                     temperature=0.01,
+                    topk=1,
                     unfinished_items=unfinished_items,
                     finished_items=finished_items,
                     forbid_audio_eos=forbid_audio_eos,