new loss functions

custom_nodes/rvc_nodes.py

@@ -362,12 +362,14 @@ class RVCTrainParamsNode:
                 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)),
-                c_mfcc=("FLOAT",dict(default=0.,min=0.,max=100.,step=.1)),
-                c_lfcc=("FLOAT",dict(default=0.,min=0.,max=100.,step=.1)),
+                c_tefs=("FLOAT",dict(default=0.,min=0.,max=100.,step=.1)),
                 c_hd=("FLOAT",dict(default=0.,min=0.,max=100.,step=.1)),
-                c_sts=("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_balancer=("BOOLEAN",dict(default=False))
+                use_balancer=("BOOLEAN",dict(default=False)),
+                use_pareto=("BOOLEAN",dict(default=False)),
+                fast_mode=("BOOLEAN",dict(default=False)),
+                fp16_run=("BOOLEAN",dict(default=True)),
             )
         }
 

lib/audio.py

@@ -133,7 +133,7 @@ def load_audio(file, sr, **kwargs):
 
     return remix_audio((np.frombuffer(out, np.float32).flatten(), sr),**kwargs)
 
-def remix_audio(input_audio,target_sr=None,norm=False,to_int16=False,resample=False,axis=0,merge_type=None,max_volume=.99,**kwargs):
+def remix_audio(input_audio,target_sr=None,norm=False,to_int16=False,resample=False,axis=0,merge_type=None,max_volume=.95,**kwargs):
     audio = np.array(input_audio[0],dtype="float32")
     if target_sr is None: target_sr=input_audio[1]
 

lib/infer_pack/commons.py

@@ -67,6 +67,10 @@ def autocorrelation1d(log_magnitude, max_lag=None):
     
     return autocorr.nan_to_num(0)
 
+def minmax_scale(tensor: torch.Tensor, eps=1e-8):
+    tensor = (tensor - tensor.min()) / (tensor.max() - tensor.min() + eps)
+    return tensor
+
 def init_weights(m, mean=0.0, std=0.01):
     classname = m.__class__.__name__
     if classname.find("Conv") != -1:

lib/train/losses.py

@@ -1,14 +1,14 @@
 import torch
-import torchaudio
+import torchaudio.transforms as T
 import torch.nn.functional as F
 
 from ..utils import gc_collect
-from ..infer_pack.commons import median_pool1d, autocorrelation1d
+from ..infer_pack.commons import median_pool1d, minmax_scale
 
 class LossBalancer:
     model: torch.nn.Module
 
-    def __init__(self, model, initial_weights={}, historical_losses={}, ema_weights={}, epsilon=1e-8, weights_decay=0., loss_decay=.0, active=True, use_partial=False, use_pareto=True, use_norm=False):
+    def __init__(self, model, initial_weights={}, historical_losses={}, ema_weights={}, epsilon=1e-8, weights_decay=0., loss_decay=.0, active=True, use_pareto=True, use_norm=False):
         """
         Initializes the LossBalancer with optional initial weights, partial gradient computation, and EMA.
         
@@ -20,7 +20,6 @@ class LossBalancer:
         """
         self.model = model
         self.epsilon = epsilon
-        self.use_partial = use_partial
         self.weights_decay = weights_decay
         self.loss_decay = loss_decay
         self.initial_weights = initial_weights
@@ -34,7 +33,6 @@ class LossBalancer:
     def to_dict(self):
         return dict(
             epsilon = self.epsilon,
-            use_partial = self.use_partial,
             weights_decay = self.weights_decay,
             loss_decay = self.loss_decay,
             ema_weights = self.ema_weights,
@@ -53,33 +51,32 @@ class LossBalancer:
                 for k in new_weights
             }
     
-    def update_historical_losses(self, new_losses):
+    def update_historical_losses(self, new_losses: dict):
         """Updates the EMA of the weights."""
         if not self.historical_losses: self.historical_losses = new_losses
-        else: self.historical_losses = {
-                k: self.loss_decay * self.historical_losses.get(k, 0.) + (1 - self.loss_decay) * new_losses[k]
-                for k in 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]
+                
 
     def calculate_loss_slope(self, key, current_loss):
         """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 = abs((current_loss-ema_loss)/ema_loss)  # relative loss change
-        if slope>1: slope*=slope # punish large changes
+        slope = abs(current_loss-ema_loss)/ema_loss  # relative loss change
         return slope
     
-    def calculate_gradients(self, current_loss):
+    def calculate_gradients(self, key, current_loss):
         """Calculates the gradient norm of the current loss wrt the model params."""
         self.model.zero_grad()  # Clear previous gradients
 
-        # Backward pass with partial or full gradient computation
-        model_params = list(self.model.parameters())[-1] if self.use_partial else self.model.parameters()
-        output_params = torch.autograd.grad([current_loss], [model_params], retain_graph=True, only_inputs=True)
+        # Backward pass to output layer
+        model_params = list(self.model.parameters())[-1]
+        output_params = torch.autograd.grad(current_loss, model_params, grad_outputs=torch.ones_like(current_loss, device=model_params.device), retain_graph=True, only_inputs=True)[0]
 
         # Compute L2 gradient norm
-        grad_norm = torch.stack([param.norm() for param in output_params if hasattr(param, "norm")]).norm(2)
-        del output_params
-        return grad_norm
+        grad_norm = output_params.view(output_params.size(0),-1).norm(2, dim=-1)
+        ema_loss = self.historical_losses.get(key, current_loss)+self.epsilon
+        return grad_norm/ema_loss
 
     def pareto_normalizer(self, data: dict):
         # Sort the data in descending order based on the values
@@ -136,7 +133,7 @@ class LossBalancer:
 
             if self.use_norm:
                 # Compute L2 gradient norm
-                grad_norm = self.calculate_gradients(loss)
+                grad_norm = self.calculate_gradients(key, loss)
                 gradients[key] = grad_norm.item()
             else:
                 # Use loss plateau detection with EMA
@@ -155,10 +152,12 @@ class LossBalancer:
 
         # Update EMA weights
         self.update_ema_weights(normalized_weights)
-        balanced_loss = sum(self.ema_weights[k] * loss for k, loss in valid_losses.items())
-        
-        # update historical losses
-        self.update_historical_losses({k: v.item() for k,v in valid_losses.items()})
+        balanced_loss = 0
+        for k, loss in valid_losses.items():
+            balanced_loss += self.ema_weights[k] * loss
+            
+            # update historical losses
+            self.update_historical_losses({k: loss.item()})
         
         return balanced_loss
 
@@ -175,57 +174,99 @@ class LossBalancer:
         gc_collect()
 
 
-def compute_phase_loss(original_log_magnitude, generated_log_magnitude, max_lag=None):
+def compute_tsi_loss(original_log_magnitude: torch.Tensor, generated_log_magnitude: torch.Tensor, dim=-1, eps=1e-8):
     """
-    Computes a time-shift invariant loss based on autocorrelation for log-magnitude data.
-    :param original_log_magnitude: Original log-magnitude spectrogram (batch, channels, time)
-    :param generated_log_magnitude: Generated log-magnitude spectrogram (batch, channels, time)
-    :param max_lag: Maximum lag for which to compute autocorrelation.
-    :return: Loss value.
+    Computes the correlation loss between the original and generated log-magnitude spectrograms.
+    :param original_envelope: Original log magnitude spectrogram (batch, time)
+    :param generated_envelope: Generated log magnitude spectrogram (batch, time)
+    :return: Correlation loss.
     """
-    # Compute autocorrelations
-    original_autocorr = autocorrelation1d(original_log_magnitude, max_lag)
-    generated_autocorr = autocorrelation1d(generated_log_magnitude, max_lag)
+    original_envelope = compute_envelope(original_log_magnitude, eps=eps, dim=dim)
+    generated_envelope = compute_envelope(generated_log_magnitude, eps=eps, dim=dim)
 
-    # Compute the loss between autocorrelations
-    loss = F.smooth_l1_loss(generated_autocorr, original_autocorr)
+    # Normalize the envelope
+    original_envelope = minmax_scale(original_envelope, eps=eps)
+    generated_envelope = minmax_scale(generated_envelope, eps=eps)
+    
+    # Compute the correlation
+    numerator = (original_envelope * generated_envelope).sum(dim=-1)
+    denominator = torch.sqrt((original_envelope ** 2).sum(dim=-1) * (generated_envelope ** 2).sum(dim=-1)+eps)
+    correlation = numerator / denominator
+    
+    # Compute the loss as negative correlation
+    loss = 1-correlation.mean()
     
     return loss
 
-def compute_envelope(log_magnitude: torch.Tensor, dim=-1, kernel_size=3):
+def compute_envelope(log_magnitude: torch.Tensor, dim=-1, kernel_size=3, eps=1e-8):
+    """
+    Compute the envelope of the log-magnitude spectrum of an audio signal.
 
-    # average over axis
-    pooled_magnitude = log_magnitude.mean(dim)
+    Args:
+        log_magnitude (torch.Tensor): The log-magnitude spectrum of the audio signal.
+        dim (int): The dimension along which to average. Default is -1.
+        kernel_size (int): The size of the max pooling kernel. Default is 3.
 
-    # use max pool to get the peaks
-    max_filtered_tensor = F.max_pool1d(pooled_magnitude, kernel_size=kernel_size, stride=1, padding=kernel_size//2)
+    Returns:
+        torch.Tensor: The computed envelope of the log-magnitude spectrum.
+    """
+
+    # Normalize the pooled magnitude
+    log_magnitude = F.normalize(log_magnitude, dim=dim, eps=eps)
+
+    # Use max pooling to get the peaks
+    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)
+
+def compute_tefs(audio_signal: torch.Tensor, eps=1e-8):
+    """
+    Computes the Hilbert transform of a batched mono audio signal to extract the envelope.
     
-    return max_filtered_tensor.nan_to_num(0).sum(-2)
+    Args:
+        audio_signal (torch.Tensor): Batched mono audio signal of shape (batch_size, signal_length).
+        
+    Returns:
+        torch.Tensor: The envelope of the audio signal of shape (batch_size, signal_length).
+    """
+    # Perform FFT on the input signal
+    audio_fft = torch.fft.fft(audio_signal.float(), dim=-1)
 
-def compute_sts_loss(gen_stft: torch.Tensor, org_stft: torch.Tensor,eps=1e-6):
+    # Create a filter to zero out negative frequencies
+    signal_length = audio_signal.shape[-1]
+    h = torch.zeros(signal_length, device=audio_signal.device)
     
-    # B?N?T
-    org_mag = torch.log10(org_stft.abs()+eps)
-    gen_mag = torch.log10(gen_stft.abs()+eps)
+    # First component is left untouched
+    h[0] = 1
+    
+    if signal_length % 2 == 0:
+        # For even signal length
+        h[1:signal_length // 2] = 2
+        h[signal_length // 2] = 1
+    else:
+        # For odd signal length
+        h[1:(signal_length + 1) // 2] = 2
 
-    # temporal invariant phase
-    phase_loss = compute_phase_loss(org_mag.sum(-2,keepdim=True),gen_mag.sum(-2,keepdim=True))
+    # Apply the filter in the frequency domain
+    hilbert_fft = audio_fft * h
 
-    # temporal envelope
-    gen_temp = compute_envelope(gen_mag,dim=-2,kernel_size=3)
-    org_temp = compute_envelope(org_mag,dim=-2,kernel_size=3)
-    temporal_loss = F.smooth_l1_loss(gen_temp, org_temp)
+    # Perform inverse FFT to obtain the analytic signal
+    analytic_signal = torch.fft.ifft(hilbert_fft, dim=-1)
 
-    # spectral
-    gen_spec = compute_envelope(gen_mag,dim=-1,kernel_size=3)
-    org_spec = compute_envelope(org_mag,dim=-1,kernel_size=3)
-    spectral_loss = F.smooth_l1_loss(gen_spec, org_spec)
+    # Extract the envelope by taking the magnitude of the analytic signal
+    envelope = torch.abs(analytic_signal)
 
-    return temporal_loss, spectral_loss, phase_loss
+    # Normalize the envelope
+    envelope = minmax_scale(envelope, eps=eps)
 
-def compute_harmonics(stft_matrix: torch.Tensor, harmonic_kernel_sizes=[11,17,23], percussive_kernel_sizes=[3,7,13], eps=1e-6):
+    # Compute the instantaneous phase from the analytic signal
+    phase = torch.angle(analytic_signal)
+
+    return envelope, phase
+
+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
-    mag = torch.log10(stft_matrix.abs()+eps)
 
     harmonic_list = []
     percussive_list = []
@@ -240,21 +281,24 @@ def compute_harmonics(stft_matrix: torch.Tensor, harmonic_kernel_sizes=[11,17,23
     harmonic = torch.cat(harmonic_list, dim=-1)
     percussive = torch.cat(percussive_list, dim=-1)
 
+    # normalize values
+    harmonic = minmax_scale(harmonic, eps=eps)
+    percussive = minmax_scale(percussive, eps=eps)
+
     return harmonic, percussive
 
 def combined_aux_loss(
-        original_audio: torch.Tensor, generated_audio: torch.Tensor, sample_rate: int,
-        c_mfcc=1., c_lfcc=1., c_hd=1., c_sts=1.,
+        original_audio: torch.Tensor, generated_audio: torch.Tensor,
+        c_tefs=1., c_hd=1., c_tsi=1., n_mels=128, sample_rate=40000,
         n_fft=1024, hop_length=320, win_length=1024,
-        n_mfcc=13, n_lfcc=13, n_filter=80,
-        f_min=0., f_max=None, eps=None):
+        eps=None):
 
     kernel_size = n_fft//hop_length+1
     if kernel_size % 2 == 0: kernel_size+=1 #enforce odd kernel size
     if eps is None: eps = torch.finfo(original_audio.dtype).eps
 
     # Compute STFT once
-    if c_hd+c_sts+c_mfcc+c_lfcc>0:
+    if c_hd+c_tsi>0:
         
         hann_window = torch.hann_window(win_length).to(
             dtype=original_audio.dtype, device=original_audio.device
@@ -278,65 +322,41 @@ def combined_aux_loss(
             onesided=True,
             center=False)
         
-        if c_sts+c_mfcc+c_lfcc>0:
-            org_mag = torch.log10(original_stft.abs()+eps)
-            gen_mag = torch.log10(generated_stft.abs()+eps)
+        MelScaler = T.MelScale(n_mels=n_mels, sample_rate=sample_rate, n_stft=n_fft // 2 + 1).to(original_audio.device)
+        org_mag = MelScaler(original_stft.abs()+eps)
+        gen_mag = MelScaler(generated_stft.abs()+eps)
     
     # Harmonic Loss
     if c_hd>0:
-        original_harmonics, original_percussives = compute_harmonics(original_stft,eps=eps)
-        generated_harmonics, generated_percussives = compute_harmonics(generated_stft,eps=eps)
+        original_harmonics, original_percussives = compute_harmonics(org_mag, eps=eps)
+        generated_harmonics, generated_percussives = compute_harmonics(gen_mag, eps=eps)
         # Define loss terms
-        harmonic_loss = F.smooth_l1_loss(generated_harmonics, original_harmonics)
-        harmonic_loss += F.smooth_l1_loss(generated_percussives, original_percussives)
+        harmonic_loss = F.l1_loss(generated_harmonics, original_harmonics)
+        harmonic_loss += F.l1_loss(generated_percussives, original_percussives)
         harmonic_loss *= c_hd
-        harmonic_loss
     else: harmonic_loss = 0
 
-    # Spectral Temporal Smoothness Loss
-    if c_sts>0:
-        # temporal invariant phase
-        phase_loss = compute_phase_loss(org_mag.sum(-2,keepdim=True),gen_mag.sum(-2,keepdim=True)) * c_sts
-        
-    else: phase_loss = 0
+    # temporal invariant phase
+    if c_tsi>0:
+        freq_tsi = compute_tsi_loss(org_mag,gen_mag,dim=-1, eps=eps)
+        temp_tsi = compute_tsi_loss(org_mag,gen_mag,dim=-2, eps=eps)
+        tsi_loss = (freq_tsi+temp_tsi) * c_tsi
+    else: tsi_loss = 0
 
-    # MFCC Loss
-    if c_mfcc>0:
+    # Temperol Envelope and Fine Structure Loss
+    if c_tefs>0:
         # temporal envelope
-        gen_temp = compute_envelope(gen_mag,dim=-2,kernel_size=3)
-        org_temp = compute_envelope(org_mag,dim=-2,kernel_size=3)
-        temporal_loss = F.smooth_l1_loss(gen_temp, org_temp) * c_mfcc
+        gen_te, gen_tfs = compute_tefs(generated_audio, eps=eps)
+        org_te, org_tfs = compute_tefs(original_audio, eps=eps)
+        phase_diff = torch.remainder(gen_tfs - org_tfs, 2 * torch.pi)
 
-        # mfcc_transform = torchaudio.transforms.MFCC(
-        #     sample_rate=sample_rate, n_mfcc=n_mfcc, log_mels=True,
-        #     melkwargs=dict(n_fft=n_fft, hop_length=hop_length, win_length=win_length, n_mels=n_filter, f_min=f_min, f_max=f_max)
-        # ).to(original_audio.device)
-        # original_mfcc = mfcc_transform(original_audio)
-        # generated_mfcc = mfcc_transform(generated_audio)
-        # mfcc_loss = F.smooth_l1_loss(generated_mfcc, original_mfcc) * c_mfcc
-        
-    else: temporal_loss = 0
-
-    # LFCC Loss
-    if c_lfcc>0:
-        # spectral
-        gen_spec = compute_envelope(gen_mag,dim=-1,kernel_size=3)
-        org_spec = compute_envelope(org_mag,dim=-1,kernel_size=3)
-        spectral_loss = F.smooth_l1_loss(gen_spec, org_spec) * c_lfcc
-
-        # lfcc_transform = torchaudio.transforms.LFCC(
-        #     sample_rate=sample_rate, n_lfcc=n_lfcc, n_filter=n_filter, f_min=f_min, f_max=f_max,log_lf=True,
-        #     speckwargs=dict(n_fft=n_fft, hop_length=hop_length, win_length=win_length)
-        # ).to(original_audio.device)
-        # original_lfcc = lfcc_transform(original_audio)
-        # generated_lfcc = lfcc_transform(generated_audio)
-        # lfcc_loss = F.smooth_l1_loss(generated_lfcc, original_lfcc) * c_lfcc
-    else: spectral_loss = 0
-    mfcc_loss=lfcc_loss=0
+        tefs_loss = F.l1_loss(gen_te, org_te) +  phase_diff.sin().abs().mean() #F.l1_loss(gen_tfs, org_tfs)
+        tefs_loss*=c_tefs
+    else: tefs_loss = 0
     
-    return harmonic_loss, temporal_loss, spectral_loss, mfcc_loss, lfcc_loss, phase_loss
+    return harmonic_loss, tefs_loss, tsi_loss
 
-def gradient_norm_loss(original_audio: torch.Tensor, generated_audio: torch.Tensor, net_d: torch.nn.Module):
+def gradient_norm_loss(original_audio: torch.Tensor, generated_audio: torch.Tensor, net_d: torch.nn.Module, eps=1e-8):
     loss=0
     # Compute the gradient penalty
     # Randomly interpolate between real and generated data
@@ -353,12 +373,12 @@ def gradient_norm_loss(original_audio: torch.Tensor, generated_audio: torch.Tens
         gradients = torch.autograd.grad(
             outputs=output,
             inputs=interpolated,
-            grad_outputs=torch.ones(output.size(), device=original_audio.device),
+            grad_outputs=torch.ones_like(output, device=original_audio.device),
             retain_graph=True,
             only_inputs=True
-        )
-        grad_norm = torch.stack([param.norm() for param in gradients if hasattr(param, "norm")]).norm(2)
-        loss += torch.log10(grad_norm).abs().mean() # force gradnorm=1
+        )[0]
+        grad_norm = gradients.view(gradients.size(0), -1).norm(2, dim=-1)+eps
+        loss += torch.log1p((grad_norm - 1) ** 2).mean()
     return loss/len(disc_interpolated_output)
 
 def feature_loss(fmap_r, fmap_g):

lib/train/utils.py

@@ -16,10 +16,9 @@ DEFAULT_TRAINING_PARAMS = dict(
     c_kl=1.,
     c_fm=2.,
     c_gp=0.,
-    c_mfcc=0.,
-    c_lfcc=0.,
+    c_tefs=0.,
     c_hd=0.,
-    c_sts=0.
+    c_tsi=0.
 )
 
 logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)

training_cli.py

@@ -273,14 +273,13 @@ def run(rank, n_gpus, hps, device):
     except Exception as e:
         logger.error(f"Failed to load balancer state: {e}")
         balancer = LossBalancer(
-            net_g,
-            use_partial=True,
+            net_g,            
             weights_decay=.5 / (1 + np.exp(-10 * (epoch_str / hps.total_epoch - 0.16)))+.5, #sigmoid scaled ema .8 at 20% epoch
             loss_decay=.8,
             epsilon=hps.train.eps,
             active=hps.train.get("use_balancer",False),
-            use_pareto=True,
-            use_norm=hps.train.get("c_gp",0)>0
+            use_pareto=hps.train.get("use_pareto",False),
+            use_norm=not hps.train.get("fast_mode",False)
             )
     cache = []
     for epoch in range(epoch_str, hps.train.epochs + 1):
@@ -495,7 +494,7 @@ def train_and_evaluate(
             y_d_hat_r, y_d_hat_g, _, _ = net_d(wave, gen_wave)
             
             with autocast(enabled=False):
-                gradient_penalty = gradient_norm_loss(wave,gen_wave, net_d)*hps.train.c_gp if hps.train.get("c_gp",0.)>0 else 0
+                gradient_penalty = gradient_norm_loss(wave,gen_wave, net_d, eps=hps.train.eps)*hps.train.c_gp if hps.train.get("c_gp",0.)>0 else 0
                 loss_disc, losses_disc_r, losses_disc_g = discriminator_loss(y_d_hat_r, y_d_hat_g)
                 loss_disc_all = loss_disc+gradient_penalty
 
@@ -512,22 +511,18 @@ def train_and_evaluate(
                 loss_mel = F.l1_loss(y_mel, y_hat_mel)*hps.train.c_mel if hps.train.get("c_mel",0.)>0 else 0 #default 45
                 loss_kl = kl_loss(z_p, logs_q, m_p, logs_p, z_mask)*hps.train.c_kl if hps.train.get("c_kl",0.)>0 else 0 #default 1
                 loss_fm = feature_loss(fmap_r, fmap_g)*hps.train.c_fm if hps.train.get("c_fm",0.)>0 else 0 #default 2
-                harmonic_loss, temporal_loss, spectral_loss, mfcc_loss, lfcc_loss, phase_loss = combined_aux_loss(
-                    wave, y_hat, SR_MAP[hps.sample_rate],
-                    c_lfcc=hps.train.c_lfcc if hps.train.get("c_lfcc",0.)>0 else 0.,
-                    c_mfcc=hps.train.c_mfcc if hps.train.get("c_mfcc",0.)>0 else 0.,
+                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.c_tefs if hps.train.get("c_tefs",0.)>0 else 0.,
                     c_hd=hps.train.c_hd if hps.train.get("c_hd",0.)>0 else 0.,
-                    c_sts=hps.train.c_sts if hps.train.get("c_sts",0.)>0 else 0.,
+                    c_tsi=hps.train.c_tsi if hps.train.get("c_tsi",0.)>0 else 0.,
                     n_fft=hps.data.filter_length,
                     hop_length=hps.data.hop_length,
                     win_length=hps.data.win_length,
-                    n_filter=hps.data.n_mel_channels,
-                    f_min=hps.data.mel_fmin,
-                    f_max=hps.data.mel_fmax,
                     eps=hps.train.eps
                 )
                 loss_gen, losses_gen = generator_loss(y_d_hat_g)
-                aux_loss = harmonic_loss + temporal_loss + spectral_loss + mfcc_loss + lfcc_loss + phase_loss
+                aux_loss = harmonic_loss + tefs_loss + tsi_loss
                 # loss_gen_all = loss_gen + loss_fm + loss_mel + loss_kl + aux_loss
                 loss_gen_all = balancer.on_train_batch_start(dict(
                     loss_gen=loss_gen,
@@ -535,11 +530,8 @@ def train_and_evaluate(
                     loss_mel=loss_mel,
                     loss_kl=loss_kl,
                     harmonic_loss=harmonic_loss,
-                    temporal_loss=temporal_loss,
-                    spectral_loss=spectral_loss,
-                    phase_loss=phase_loss,
-                    mfcc_loss=mfcc_loss,
-                    lfcc_loss=lfcc_loss
+                    tsi_loss=tsi_loss,
+                    tefs_loss=tefs_loss,
                     ))
                 
             optim_g.zero_grad()
@@ -554,7 +546,7 @@ def train_and_evaluate(
             if loss_gen + loss_fm + loss_mel + loss_kl + aux_loss<least_loss:
                 least_loss = loss_gen + loss_fm + loss_mel + loss_kl + aux_loss
                 logger.info(f"[lowest loss] {least_loss=:.3f}: {loss_disc=:.3f} {gradient_penalty=:.3f} <==> {loss_gen=:.3f} {loss_fm=:.3f} {loss_mel=:.3f} {loss_kl=:.3f}")
-                logger.info(f"[aux] {aux_loss=:.3f}: {harmonic_loss=:.3f}, {temporal_loss=:.3f}, {spectral_loss=:.3f}, {mfcc_loss=:.3f}, {lfcc_loss=:.3f} {phase_loss=:.3f}")
+                logger.info(f"[aux] {aux_loss=:.3f}: {harmonic_loss=:.3f}, {tefs_loss=:.3f}, {tsi_loss=:.3f}")
 
                 if hps.save_best_model:
                     if hasattr(net_g, "module"): ckpt = net_g.module.state_dict()
@@ -578,11 +570,8 @@ def train_and_evaluate(
                                         "total/loss/mel": commons.serialize_tensor(loss_mel),
                                         "total/loss/kl": commons.serialize_tensor(loss_kl),
                                         "aux/loss/harmonic": commons.serialize_tensor(harmonic_loss),
-                                        "aux/loss/temporal": commons.serialize_tensor(temporal_loss),
-                                        "aux/loss/spectral": commons.serialize_tensor(spectral_loss),
-                                        "aux/loss/mfcc": commons.serialize_tensor(mfcc_loss),
-                                        "aux/loss/lfcc": commons.serialize_tensor(lfcc_loss),
-                                        "aux/loss/phase_loss": commons.serialize_tensor(phase_loss),
+                                        "aux/loss/tefs": commons.serialize_tensor(tefs_loss),
+                                        "aux/loss/tsi": commons.serialize_tensor(tsi_loss),
                                         "gradient/grad_norm_disc": commons.serialize_tensor(grad_norm_d),
                                         "gradient/grad_norm_gen": commons.serialize_tensor(grad_norm_g),
                                         "gradient/gradient_penalty": commons.serialize_tensor(gradient_penalty),
@@ -608,11 +597,8 @@ def train_and_evaluate(
                     "total/loss/mel": loss_mel,
                     "total/loss/kl": loss_kl,
                     "aux/loss/harmonic": harmonic_loss,
-                    "aux/loss/temporal": temporal_loss,
-                    "aux/loss/spectral": spectral_loss,
-                    "aux/loss/mfcc": mfcc_loss,
-                    "aux/loss/lfcc": lfcc_loss,
-                    "aux/loss/phase_loss": phase_loss,
+                    "aux/loss/tefs": tefs_loss,
+                    "aux/loss/tsi": tsi_loss,
                     "gradient/lr": lr,
                     "gradient/grad_norm_disc": grad_norm_d,
                     "gradient/grad_norm_gen": grad_norm_g,
@@ -677,7 +663,7 @@ def train_and_evaluate(
         logger.info(f"====> Epoch {epoch} (Total Loss={(loss_disc_all+loss_gen_all).item():.3f}): {global_step=} {lr=:.2E} {epoch_recorder.record()}")
         logger.info(f"|| {loss_disc_all=:.3f}: {loss_disc=:.3f}, {gradient_penalty=:.3f}")
         logger.info(f"|| {loss_gen_all=:.3f}: {loss_gen=:.3f}, {loss_fm=:.3f}, {loss_mel=:.3f}, {loss_kl=:.3f}")
-        logger.info(f"|| {harmonic_loss=:.3f}, {temporal_loss=:.3f}, {spectral_loss=:.3f}, {mfcc_loss=:.3f}, {lfcc_loss=:.3f}, {phase_loss=:.3f}")
+        logger.info(f"|| {harmonic_loss=:.3f}, {tefs_loss=:.3f}, {tsi_loss=:.3f}")
         balancer.on_epoch_end(.5 / (1 + np.exp(-10 * (epoch / hps.total_epoch - 0.16)))+.5) #sigmoid scaling of ema
 
     if epoch >= hps.total_epoch and rank == 0: