Update TTS/tts/layers/xtts/dvae.py

TTS/tts/layers/xtts/dvae.py

@@ -24,9 +24,7 @@ def eval_decorator(fn):
     return inner
 
 
-def dvae_wav_to_mel(
-    wav, mel_norms_file="../experiments/clips_mel_norms.pth", mel_norms=None, device=torch.device("cpu")
-):
+def dvae_wav_to_mel(wav, mel_norms_file="../experiments/clips_mel_norms.pth", mel_norms=None, device=torch.device("cpu")):
     mel_stft = torchaudio.transforms.MelSpectrogram(
         n_fft=1024,
         hop_length=256,
@@ -44,7 +42,7 @@ def dvae_wav_to_mel(
     mel = torch.log(torch.clamp(mel, min=1e-5))
     if mel_norms is None:
         mel_norms = torch.load(mel_norms_file, map_location=device)
-    mel = mel / mel_norms.unsqueeze(0).unsqueeze(-1)
+    mel = mel / (mel_norms.unsqueeze(0).unsqueeze(-1) + 1e-8)  # Adicionando um valor pequeno para evitar divisão por zero
     return mel
 
 
@@ -112,7 +110,7 @@ class Quantize(nn.Module):
             self.embed_avg.data.mul_(self.decay).add_(embed_sum, alpha=1 - self.decay)
             n = self.cluster_size.sum()
             cluster_size = (self.cluster_size + self.eps) / (n + self.n_embed * self.eps) * n
-            embed_normalized = self.embed_avg / cluster_size.unsqueeze(0)
+            embed_normalized = self.embed_avg / (cluster_size.unsqueeze(0) + self.eps)  # Adicionando eps para evitar divisão por zero
             self.embed.data.copy_(embed_normalized)
 
         diff = (quantize.detach() - input).pow(2).mean()
@@ -198,6 +196,7 @@ class UpsampledConv(nn.Module):
 
 # DiscreteVAE partially derived from lucidrains DALLE implementation
 # Credit: https://github.com/lucidrains/DALLE-pytorch
+
 class DiscreteVAE(nn.Module):
     def __init__(
         self,
@@ -215,7 +214,7 @@ class DiscreteVAE(nn.Module):
         activation="relu",
         smooth_l1_loss=False,
         straight_through=False,
-        normalization=None,  # ((0.5,) * 3, (0.5,) * 3),
+        normalization=None,
         record_codes=False,
         discretization_loss_averaging_steps=100,
         lr_quantizer_args={},
@@ -231,7 +230,7 @@ class DiscreteVAE(nn.Module):
             num_tokens, 2, 1 / (num_tokens * 2), discretization_loss_averaging_steps
         )
 
-        assert positional_dims > 0 and positional_dims < 3  # This VAE only supports 1d and 2d inputs for now.
+        assert positional_dims > 0 and positional_dims < 3
         if positional_dims == 2:
             conv = nn.Conv2d
             conv_transpose = nn.ConvTranspose2d
@@ -246,7 +245,7 @@ class DiscreteVAE(nn.Module):
         elif activation == "silu":
             act = nn.SiLU
         else:
-            assert NotImplementedError()
+            raise NotImplementedError()
 
         enc_layers = []
         dec_layers = []
@@ -293,7 +292,6 @@ class DiscreteVAE(nn.Module):
         self.loss_fn = F.smooth_l1_loss if smooth_l1_loss else F.mse_loss
         self.codebook = Quantize(codebook_dim, num_tokens, new_return_order=True)
 
-        # take care of normalization within class
         self.normalization = normalization
         self.record_codes = record_codes
         if record_codes:
@@ -303,19 +301,18 @@ class DiscreteVAE(nn.Module):
         self.internal_step = 0
 
     def norm(self, images):
-        if not self.normalization is not None:
+        if self.normalization is None:
             return images
 
         means, stds = map(lambda t: torch.as_tensor(t).to(images), self.normalization)
         arrange = "c -> () c () ()" if self.positional_dims == 2 else "c -> () c ()"
         means, stds = map(lambda t: rearrange(t, arrange), (means, stds))
         images = images.clone()
-        images.sub_(means).div_(stds)
+        images.sub_(means).div_(stds + 1e-8)  # Adicionando eps para evitar divisão por zero
         return images
 
     def get_debug_values(self, step, __):
         if self.record_codes and self.total_codes > 0:
-            # Report annealing schedule
             return {"histogram_codes": self.codes[: self.total_codes]}
         else:
             return {}
@@ -356,9 +353,6 @@ class DiscreteVAE(nn.Module):
         sampled, codes, commitment_loss = self.codebook(logits)
         return self.decode(codes)
 
-    # Note: This module is not meant to be run in forward() except while training. It has special logic which performs
-    # evaluation using quantized values when it detects that it is being run in eval() mode, which will be substantially
-    # more lossy (but useful for determining network performance).
     def forward(self, img):
         img = self.norm(img)
         logits = self.encoder(img).permute((0, 2, 3, 1) if len(img.shape) == 4 else (0, 2, 1))
@@ -371,16 +365,13 @@ class DiscreteVAE(nn.Module):
                 out = d(out)
             self.log_codes(codes)
         else:
-            # This is non-differentiable, but gives a better idea of how the network is actually performing.
             out, _ = self.decode(codes)
 
-        # reconstruction loss
         recon_loss = self.loss_fn(img, out, reduction="none")
 
         return recon_loss, commitment_loss, out
 
     def log_codes(self, codes):
-        # This is so we can debug the distribution of codes being learned.
         if self.record_codes and self.internal_step % 10 == 0:
             codes = codes.flatten()
             l = codes.shape[0]
@@ -391,3 +382,4 @@ class DiscreteVAE(nn.Module):
                 self.code_ind = 0
             self.total_codes += 1
         self.internal_step += 1
+