added multiscale mel

custom_nodes/rvc_nodes.py

@@ -359,6 +359,7 @@ class RVCTrainParamsNode:
         return {
             "optional": dict(
                 batch_size=("INT",dict(default=4,min=1,max=64,step=1)),
+                c_gen=("FLOAT",dict(default=1.,min=0.,max=100.,step=.1)),
                 c_mel=("FLOAT",dict(default=45.,min=0.,max=100.,step=.1)),
                 c_kl=("FLOAT",dict(default=1.,min=0.,max=100.,step=.1)),
                 c_fm=("FLOAT",dict(default=2.,min=0.,max=100.,step=.1)),
@@ -366,6 +367,7 @@ class RVCTrainParamsNode:
                 c_hd=("FLOAT",dict(default=0.,min=0.,max=100.,step=.1)),
                 c_tsi=("FLOAT",dict(default=0.,min=0.,max=100.,step=.1)),
                 c_gp=("FLOAT",dict(default=0.,min=0.,max=100.,step=.1)),
+                use_multiscale=("BOOLEAN",dict(default=False)),
                 use_balancer=("BOOLEAN",dict(default=False)),
                 use_pareto=("BOOLEAN",dict(default=False)),
                 fast_mode=("BOOLEAN",dict(default=False)),

lib/train/losses.py

@@ -1,4 +1,10 @@
-from typing import Optional
+from collections import namedtuple
+import functools
+from typing import Callable, List, Optional
+import typing
+from librosa.filters import mel as librosa_mel_fn
+import numpy as np
+from scipy import signal
 import torch
 import torchaudio.transforms as T
 import torch.nn.functional as F
@@ -43,7 +49,7 @@ class LossBalancer:
             use_norm = self.use_norm
         )
 
-    def update_ema_weights(self, new_weights):
+    def update_ema_weights(self, new_weights: dict):
         """Updates the EMA of the weights."""
         if not self.ema_weights: self.ema_weights = dict(**new_weights)
         else: self.ema_weights = {
@@ -56,19 +62,21 @@ class LossBalancer:
         """Updates the EMA of the weights."""
         if not self.historical_losses: self.historical_losses = dict(**new_losses)
         else: 
-            for k in new_losses:
-                self.historical_losses[k]=self.loss_decay*self.historical_losses.get(k, 0.) + (1 - self.loss_decay)*new_losses[k]
+            for k,v in new_losses.items():
+                if np.nan_to_num(v,nan=-1)==-1: print(f"{k=} {v=}")
+                self.historical_losses[k]=self.loss_decay*self.historical_losses.get(k, v) + (1 - self.loss_decay)*v
         return dict(**self.historical_losses)
 
-    def calculate_loss_slope(self, key, current_loss):
+    def calculate_loss_slope(self, key: str, current_loss: torch.Tensor):
         """Calculates the slope of the loss using the current loss and its EMA."""
         ema_loss = self.historical_losses.get(key, current_loss)+self.epsilon
         slope = (current_loss-ema_loss)/ema_loss  # relative loss change
         return slope
     
-    def calculate_gradients(self, key, current_loss, input):
+    def calculate_gradients(self, key: str, current_loss: torch.Tensor, input: torch.Tensor):
         """Calculates the gradient norm of the current loss wrt the model params."""
         self.model.zero_grad()  # Clear previous gradients
+        input.requires_grad_(True)
 
         # Backward pass to output layer
         output_params = torch.autograd.grad(current_loss, [input], retain_graph=True, allow_unused=True, materialize_grads=True)[0]
@@ -89,7 +97,7 @@ class LossBalancer:
         top_20_count = max(int(0.2 * len(sorted_data)),1)
         
         # Calculate the weights
-        weights = [0] * len(sorted_data)
+        weights = np.zeros(len(sorted_data))
         total_weight = 1.0
         top_20_weight = 0.8 * total_weight
         remaining_weight = total_weight - top_20_weight
@@ -141,14 +149,13 @@ class LossBalancer:
 
             if self.use_norm and input is not None:
                 # Compute L2 gradient norm
-                input.requires_grad_(True)
                 grad_norm = self.calculate_gradients(key, weighted_loss, input)
                 gradients[key] = max(grad_norm.item(), self.epsilon)
             else:
                 # Use loss plateau detection with EMA
                 loss_slope = self.calculate_loss_slope(key, weighted_loss)
                 gradients[key] = max(loss_slope.item(), self.epsilon)
-            valid_losses[key] = weighted_loss
+            valid_losses[key] = weighted_loss.nan_to_num(self.epsilon)
 
         if not valid_losses: return torch.tensor(0.0)  # If all losses are skipped
 
@@ -163,10 +170,10 @@ class LossBalancer:
         self.update_ema_weights(normalized_weights)
         balanced_loss = 0
         for k, loss in valid_losses.items():
-            balanced_loss += self.ema_weights[k] * loss
+            balanced_loss += self.ema_weights.get(k, 1.) * loss
             
-            # update historical losses
-            self.update_historical_losses({k: loss.item()})
+        # update historical losses
+        self.update_historical_losses({k: loss.item() for k, loss in valid_losses.items()})
         
         return balanced_loss
 
@@ -185,7 +192,7 @@ class LossBalancer:
 
     @property
     def weighted_ema_loss(self):
-        return sum(v*self.ema_weights.get(k,1./len(self.historical_losses)) for k,v in self.historical_losses.items())
+        return sum(v*self.ema_weights.get(k,1.) for k,v in self.historical_losses.items())
 
 def compute_tsi_loss(original_log_magnitude: torch.Tensor, generated_log_magnitude: torch.Tensor, dim=-1, eps=1e-8):
     """
@@ -229,7 +236,7 @@ def compute_envelope(log_magnitude: torch.Tensor, dim=-1, kernel_size=3, eps=1e-
     max_filtered_tensor = F.max_pool1d(log_magnitude, kernel_size=kernel_size, stride=1, padding=kernel_size//2)
 
     # Replace NaNs with zeros and sum along the second-to-last dimension
-    return max_filtered_tensor.nan_to_num(0).sum(dim)
+    return max_filtered_tensor.nan_to_num(eps).sum(dim)
 
 def compute_tefs(audio_signal: torch.Tensor, eps=1e-8):
     """
@@ -274,7 +281,7 @@ def compute_tefs(audio_signal: torch.Tensor, eps=1e-8):
     # Compute the instantaneous phase from the analytic signal
     phase = torch.angle(analytic_signal).cos()
 
-    return envelope.nan_to_num(0), phase.nan_to_num(0)
+    return envelope.nan_to_num(eps), phase.nan_to_num(eps)
 
 def compute_harmonics(mag: torch.Tensor, harmonic_kernel_sizes=[11,17,23], percussive_kernel_sizes=[3,7,13], eps=1e-8):
     # Calculate log-magnitude spectrograms
@@ -283,10 +290,10 @@ def compute_harmonics(mag: torch.Tensor, harmonic_kernel_sizes=[11,17,23], percu
     percussive_list = []
 
     for kernel_size in harmonic_kernel_sizes:
-        harmonic_list.append(median_pool1d(mag, kernel_size=kernel_size, dim=-1).nan_to_num(0).view(mag.size(0),-1))
+        harmonic_list.append(median_pool1d(mag, kernel_size=kernel_size, dim=-1).nan_to_num(eps).view(mag.size(0),-1))
 
     for kernel_size in percussive_kernel_sizes:
-        percussive_list.append(median_pool1d(mag, kernel_size=kernel_size, dim=-2).nan_to_num(0).view(mag.size(0),-1))
+        percussive_list.append(median_pool1d(mag, kernel_size=kernel_size, dim=-2).nan_to_num(eps).view(mag.size(0),-1))
 
     # Concatenate the results
     harmonic = torch.cat(harmonic_list, dim=-1)
@@ -296,7 +303,7 @@ def compute_harmonics(mag: torch.Tensor, harmonic_kernel_sizes=[11,17,23], percu
     harmonic = minmax_scale(harmonic, eps=eps)
     percussive = minmax_scale(percussive, eps=eps)
 
-    return harmonic.nan_to_num(0), percussive.nan_to_num(0)
+    return harmonic.nan_to_num(eps), percussive.nan_to_num(eps)
 
 def combined_aux_loss(
         original_audio: torch.Tensor, generated_audio: torch.Tensor,
@@ -392,41 +399,223 @@ def gradient_norm_loss(original_audio: torch.Tensor, generated_audio: torch.Tens
         loss += torch.log1p((grad_norm - 1) ** 2)
     return loss/len(disc_interpolated_output)
 
-def feature_loss(fmap_r, fmap_g):
+# Adapted from https://github.com/NVIDIA/BigVGAN/blob/main/loss.py
+# LICENSE: https://github.com/NVIDIA/BigVGAN/blob/main/LICENSE
+class MultiScaleMelSpectrogramLoss(torch.nn.Module):
+    """Compute distance between mel spectrograms. Can be used
+    in a multi-scale way.
+
+    Parameters
+    ----------
+    n_mels : List[int]
+        Number of mels per STFT, by default [5, 10, 20, 40, 80, 160, 320],
+    window_lengths : List[int], optional
+        Length of each window of each STFT, by default [32, 64, 128, 256, 512, 1024, 2048]
+    loss_fn : typing.Callable, optional
+        How to compare each loss, by default nn.L1Loss()
+    clamp_eps : float, optional
+        Clamp on the log magnitude, below, by default 1e-5
+    mag_weight : float, optional
+        Weight of raw magnitude portion of loss, by default 0.0 (no ampliciation on mag part)
+    log_weight : float, optional
+        Weight of log magnitude portion of loss, by default 1.0
+    pow : float, optional
+        Power to raise magnitude to before taking log, by default 1.0
+    weight : float, optional
+        Weight of this loss, by default 1.0
+    match_stride : bool, optional
+        Whether to match the stride of convolutional layers, by default False
+
+    Implementation copied from: https://github.com/descriptinc/lyrebird-audiotools/blob/961786aa1a9d628cca0c0486e5885a457fe70c1a/audiotools/metrics/spectral.py
+    Additional code copied and modified from https://github.com/descriptinc/audiotools/blob/master/audiotools/core/audio_signal.py
+    """
+
+    def __init__(
+        self,
+        sampling_rate: int,
+        n_mels: List[int] = [5, 10, 20, 40, 80, 160, 320],
+        window_lengths: List[int] = [32, 64, 128, 256, 512, 1024, 2048],
+        loss_fn: Callable = torch.nn.L1Loss(),
+        clamp_eps: float = 1e-5,
+        mag_weight: float = 0.0,
+        log_weight: float = 1.0,
+        pow: float = 1.0,
+        weight: float = 1.0,
+        match_stride: bool = False,
+        mel_fmin: List[float] = [0, 0, 0, 0, 0, 0, 0],
+        mel_fmax: List[float] = [None, None, None, None, None, None, None],
+        window_type: str = "hann",
+    ):
+        super().__init__()
+        self.sampling_rate = sampling_rate
+
+        STFTParams = namedtuple(
+            "STFTParams",
+            ["window_length", "hop_length", "window_type", "match_stride"],
+        )
+
+        self.stft_params = [
+            STFTParams(
+                window_length=w,
+                hop_length=w // 4,
+                match_stride=match_stride,
+                window_type=window_type,
+            )
+            for w in window_lengths
+        ]
+        self.n_mels = n_mels
+        self.loss_fn = loss_fn
+        self.clamp_eps = clamp_eps
+        self.log_weight = log_weight
+        self.mag_weight = mag_weight
+        self.weight = weight
+        self.mel_fmin = mel_fmin
+        self.mel_fmax = mel_fmax
+        self.pow = pow
+
+    @staticmethod
+    @functools.lru_cache(None)
+    def get_window(
+        window_type,
+        window_length,
+    ):
+        return signal.get_window(window_type, window_length)
+
+    @staticmethod
+    @functools.lru_cache(None)
+    def get_mel_filters(sr, n_fft, n_mels, fmin, fmax):
+        return librosa_mel_fn(sr=sr, n_fft=n_fft, n_mels=n_mels, fmin=fmin, fmax=fmax)
+
+    def mel_spectrogram(
+        self,
+        wav,
+        n_mels,
+        fmin,
+        fmax,
+        window_length,
+        hop_length,
+        match_stride,
+        window_type,
+    ):
+        """
+        Mirrors AudioSignal.mel_spectrogram used by BigVGAN-v2 training from: 
+        https://github.com/descriptinc/audiotools/blob/master/audiotools/core/audio_signal.py
+        """
+        B, C, T = wav.shape
+
+        if match_stride:
+            assert (
+                hop_length == window_length // 4
+            ), "For match_stride, hop must equal n_fft // 4"
+            right_pad = np.ceil(T / hop_length) * hop_length - T
+            pad = (window_length - hop_length) // 2
+        else:
+            right_pad = 0
+            pad = 0
+
+        wav = torch.nn.functional.pad(wav, (pad, pad + right_pad), mode="reflect")
+
+        window = self.get_window(window_type, window_length)
+        window = torch.from_numpy(window).to(wav.device).float()
+
+        stft = torch.stft(
+            wav.reshape(-1, T),
+            n_fft=window_length,
+            hop_length=hop_length,
+            window=window,
+            return_complex=True,
+            center=True,
+        )
+        _, nf, nt = stft.shape
+        stft = stft.reshape(B, C, nf, nt)
+        if match_stride:
+            """
+            Drop first two and last two frames, which are added, because of padding. Now num_frames * hop_length = num_samples.
+            """
+            stft = stft[..., 2:-2]
+        magnitude = torch.abs(stft)
+
+        nf = magnitude.shape[2]
+        mel_basis = self.get_mel_filters(
+            self.sampling_rate, 2 * (nf - 1), n_mels, fmin, fmax
+        )
+        mel_basis = torch.from_numpy(mel_basis).to(wav.device)
+        mel_spectrogram = magnitude.transpose(2, -1) @ mel_basis.T
+        mel_spectrogram = mel_spectrogram.transpose(-1, 2)
+
+        return mel_spectrogram
+
+    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+        """Computes mel loss between an estimate and a reference
+        signal.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Estimate signal
+        y : torch.Tensor
+            Reference signal
+
+        Returns
+        -------
+        torch.Tensor
+            Mel loss.
+        """
+
+        loss = 0.0
+        for n_mels, fmin, fmax, s in zip(
+            self.n_mels, self.mel_fmin, self.mel_fmax, self.stft_params
+        ):
+            kwargs = {
+                "n_mels": n_mels,
+                "fmin": fmin,
+                "fmax": fmax,
+                "window_length": s.window_length,
+                "hop_length": s.hop_length,
+                "match_stride": s.match_stride,
+                "window_type": s.window_type,
+            }
+
+            x_mels = self.mel_spectrogram(x, **kwargs)
+            y_mels = self.mel_spectrogram(y, **kwargs)
+            x_logmels = torch.log10(x_mels.pow(self.pow)+self.clamp_eps).nan_to_num(self.clamp_eps)
+            y_logmels = torch.log10(y_mels.pow(self.pow)+self.clamp_eps).nan_to_num(self.clamp_eps)
+
+            if self.log_weight!=0: loss += self.log_weight * self.loss_fn(x_logmels, y_logmels)
+            if self.mag_weight!=0: loss += self.mag_weight * self.loss_fn(x_mels, y_mels)
+
+        return loss
+    
+def feature_loss(fmap_r: List[List[torch.Tensor]], fmap_g: List[List[torch.Tensor]]):
     loss = 0
     for dr, dg in zip(fmap_r, fmap_g):
         for rl, gl in zip(dr, dg):
-            rl = rl.float()
-            gl = gl.float()
             loss += torch.mean(torch.abs(rl - gl))
-
     return loss
 
-def discriminator_loss(disc_real_outputs, disc_generated_outputs):
+def discriminator_loss(
+        disc_real_outputs: List[torch.Tensor],
+        disc_generated_outputs: List[torch.Tensor]
+        ):
     loss = 0
     r_losses = []
     g_losses = []
     for dr, dg in zip(disc_real_outputs, disc_generated_outputs):
-        dr = dr.float()
-        dg = dg.float()
         r_loss = torch.mean((1 - dr) ** 2)
         g_loss = torch.mean(dg**2)
         loss += r_loss + g_loss
         r_losses.append(r_loss.item())
         g_losses.append(g_loss.item())
-    # loss /= len(disc_generated_outputs) #average
     return loss, r_losses, g_losses
 
 
-def generator_loss(disc_outputs):
+def generator_loss(disc_outputs: List[torch.Tensor]):
     loss = 0
     gen_losses = []
     for dg in disc_outputs:
-        dg = dg.float()
         l = torch.mean((1 - dg) ** 2)
         gen_losses.append(l)
         loss += l
-    # loss /= len(disc_outputs) #average
     return loss, gen_losses
 
 

training_cli.py

@@ -31,7 +31,7 @@ from .lib.train.data_utils import (
     TextAudioCollate,
     DistributedBucketSampler,
 )
-from .lib.train.losses import LossBalancer, combined_aux_loss, generator_loss, discriminator_loss, feature_loss, gradient_norm_loss, kl_loss
+from .lib.train.losses import LossBalancer, MultiScaleMelSpectrogramLoss, combined_aux_loss, generator_loss, discriminator_loss, feature_loss, gradient_norm_loss, kl_loss
 from .lib.train.mel_processing import mel_spectrogram_torch, spec_to_mel_torch
 
 def save_checkpoint(ckpt, name, epoch, hps, model_path=None):
@@ -101,7 +101,6 @@ def train_model(hps: "utils.HParams"):
         n_gpus = 1
     
     gpu_devices = hps.gpus.split("-") if hps.gpus else range(n_gpus)
-    
 
     children = {}
     for i, device in enumerate(gpu_devices):
@@ -122,7 +121,8 @@ def train_model(hps: "utils.HParams"):
 
 def run(rank, n_gpus, hps, device):
     print(f"{__name__=}")
-    global global_step, least_loss, loss_file, best_model_name, gradient_clip_value
+    global global_step, least_loss, loss_file, best_model_name, gradient_clip_value, MultiscaleMelLoss
+    MultiscaleMelLoss = MultiScaleMelSpectrogramLoss(hps.data.sampling_rate)
     global_step = 0
     loss_file = os.path.join(hps.model_dir,"losses.json")
     gradient_clip_value = np.log(hps.train.c_gp/hps.train.learning_rate) if hps.train.get("c_gp",0.)>0 else None
@@ -297,7 +297,7 @@ def run(rank, n_gpus, hps, device):
             use_pareto=hps.train.get("use_pareto",False),
             use_norm=not hps.train.get("fast_mode",False),
             initial_weights=dict(
-                loss_gen=1.,
+                loss_gen=hps.train.get("c_gen",1.),
                 loss_fm=hps.train.get("c_fm",2.),
                 loss_mel=hps.train.get("c_mel",45.),
                 loss_kl=hps.train.get("c_kl",1.),
@@ -351,7 +351,7 @@ def train_and_evaluate(
     if writers is not None:
         writer, _ = writers
 
-    global global_step, least_loss, loss_file, best_model_name, gradient_clip_value
+    global global_step, least_loss, loss_file, best_model_name, gradient_clip_value, MultiscaleMelLoss
 
     net_g.train()
     net_d.train()
@@ -532,9 +532,10 @@ def train_and_evaluate(
             # Generator
             y_d_hat_r, y_d_hat_g, fmap_r, fmap_g = net_d(wave, y_hat)
             with autocast(enabled=False):
-                loss_mel = F.l1_loss(y_mel, y_hat_mel) #*hps.train.get("c_mel",0.)
-                loss_kl = kl_loss(z_p, logs_q, m_p, logs_p, z_mask) #*hps.train.get("c_kl",0.)
-                loss_fm = feature_loss(fmap_r, fmap_g) #*hps.train.get("c_fm",0.)
+                if hps.train.get("use_multiscale"): loss_mel = MultiscaleMelLoss(y_hat, wave)
+                else: loss_mel = F.l1_loss(y_mel, y_hat_mel)
+                loss_kl = kl_loss(z_p, logs_q, m_p, logs_p, z_mask)
+                loss_fm = feature_loss(fmap_r, fmap_g)
                 harmonic_loss, tefs_loss, tsi_loss = combined_aux_loss(
                     wave, y_hat,n_mels=hps.data.n_mel_channels,sample_rate=hps.data.sampling_rate,
                     c_tefs=hps.train.get("c_tefs",0.),
@@ -595,7 +596,7 @@ def train_and_evaluate(
                     **{f"loss/g/{i}": v for i, v in enumerate(losses_gen)},
                     **{f"loss/d_r/{i}": v for i, v in enumerate(losses_disc_r)},
                     **{f"loss/d_g/{i}": v for i, v in enumerate(losses_disc_g)},
-                    **{f"balancer/weights/{k}": v for k, v in balancer.ema_weights.items()}
+                    **{f"balancer/weights/{k}": v for k, v in balancer.ema_weights.items()},
                     **{f"balancer/losses/{k}": v for k, v in balancer.historical_losses.items()}
                 }
 
@@ -603,7 +604,7 @@ def train_and_evaluate(
                     "slice/mel_org": utils.plot_spectrogram_to_numpy(y_mel[0].data.cpu().numpy()),
                     "slice/mel_gen": utils.plot_spectrogram_to_numpy(y_hat_mel[0].data.cpu().numpy()),
                     "all/mel": utils.plot_spectrogram_to_numpy(mel[0].data.cpu().numpy()),
-                    "slice/diff": utils.plot_spectrogram_to_numpy((y_mel[0]-y_hat_mel[0]).abs().data.cpu().numpy(), cmap="hot")
+                    "slice/diff^2": utils.plot_spectrogram_to_numpy((y_mel[0]-y_hat_mel[0]).pow(2).data.cpu().numpy(), cmap="hot")
                 }
 
                 with torch.no_grad():
@@ -643,7 +644,7 @@ def train_and_evaluate(
             optim_d,
             hps.train.learning_rate,
             epoch,
-            os.path.join(hps.model_dir, f"D_{saved_epoch}.pth"),
+            os.path.join(hps.model_dir, f"D_{saved_epoch}.pth")
         )
         if hps.save_every_weights:
             ckpt = net_g.module.state_dict() if hasattr(net_g, "module") else net_g.state_dict()