|
import torch |
|
from torch import nn |
|
import torch.nn.functional as F |
|
import torchaudio |
|
from transformers import AutoModel |
|
|
|
|
|
class SpectralConvergengeLoss(torch.nn.Module): |
|
"""Spectral convergence loss module.""" |
|
|
|
def __init__(self): |
|
"""Initilize spectral convergence loss module.""" |
|
super(SpectralConvergengeLoss, self).__init__() |
|
|
|
def forward(self, x_mag, y_mag): |
|
"""Calculate forward propagation. |
|
Args: |
|
x_mag (Tensor): Magnitude spectrogram of predicted signal (B, #frames, #freq_bins). |
|
y_mag (Tensor): Magnitude spectrogram of groundtruth signal (B, #frames, #freq_bins). |
|
Returns: |
|
Tensor: Spectral convergence loss value. |
|
""" |
|
return torch.norm(y_mag - x_mag, p=1) / torch.norm(y_mag, p=1) |
|
|
|
|
|
class STFTLoss(torch.nn.Module): |
|
"""STFT loss module.""" |
|
|
|
def __init__( |
|
self, fft_size=1024, shift_size=120, win_length=600, window=torch.hann_window |
|
): |
|
"""Initialize STFT loss module.""" |
|
super(STFTLoss, self).__init__() |
|
self.fft_size = fft_size |
|
self.shift_size = shift_size |
|
self.win_length = win_length |
|
self.to_mel = torchaudio.transforms.MelSpectrogram( |
|
sample_rate=24000, |
|
n_fft=fft_size, |
|
win_length=win_length, |
|
hop_length=shift_size, |
|
window_fn=window, |
|
) |
|
|
|
self.spectral_convergenge_loss = SpectralConvergengeLoss() |
|
|
|
def forward(self, x, y): |
|
"""Calculate forward propagation. |
|
Args: |
|
x (Tensor): Predicted signal (B, T). |
|
y (Tensor): Groundtruth signal (B, T). |
|
Returns: |
|
Tensor: Spectral convergence loss value. |
|
Tensor: Log STFT magnitude loss value. |
|
""" |
|
x_mag = self.to_mel(x) |
|
mean, std = -4, 4 |
|
x_mag = (torch.log(1e-5 + x_mag) - mean) / std |
|
|
|
y_mag = self.to_mel(y) |
|
mean, std = -4, 4 |
|
y_mag = (torch.log(1e-5 + y_mag) - mean) / std |
|
|
|
sc_loss = self.spectral_convergenge_loss(x_mag, y_mag) |
|
return sc_loss |
|
|
|
|
|
class MultiResolutionSTFTLoss(torch.nn.Module): |
|
"""Multi resolution STFT loss module.""" |
|
|
|
def __init__( |
|
self, |
|
fft_sizes=[1024, 2048, 512], |
|
hop_sizes=[120, 240, 50], |
|
win_lengths=[600, 1200, 240], |
|
window=torch.hann_window, |
|
): |
|
"""Initialize Multi resolution STFT loss module. |
|
Args: |
|
fft_sizes (list): List of FFT sizes. |
|
hop_sizes (list): List of hop sizes. |
|
win_lengths (list): List of window lengths. |
|
window (str): Window function type. |
|
""" |
|
super(MultiResolutionSTFTLoss, self).__init__() |
|
assert len(fft_sizes) == len(hop_sizes) == len(win_lengths) |
|
self.stft_losses = torch.nn.ModuleList() |
|
for fs, ss, wl in zip(fft_sizes, hop_sizes, win_lengths): |
|
self.stft_losses += [STFTLoss(fs, ss, wl, window)] |
|
|
|
def forward(self, x, y): |
|
"""Calculate forward propagation. |
|
Args: |
|
x (Tensor): Predicted signal (B, T). |
|
y (Tensor): Groundtruth signal (B, T). |
|
Returns: |
|
Tensor: Multi resolution spectral convergence loss value. |
|
Tensor: Multi resolution log STFT magnitude loss value. |
|
""" |
|
sc_loss = 0.0 |
|
for f in self.stft_losses: |
|
sc_l = f(x, y) |
|
sc_loss += sc_l |
|
sc_loss /= len(self.stft_losses) |
|
|
|
return sc_loss |
|
|
|
|
|
def feature_loss(fmap_r, fmap_g): |
|
loss = 0 |
|
for dr, dg in zip(fmap_r, fmap_g): |
|
for rl, gl in zip(dr, dg): |
|
loss += torch.mean(torch.abs(rl - gl)) |
|
|
|
return loss * 2 |
|
|
|
|
|
def discriminator_loss(disc_real_outputs, disc_generated_outputs): |
|
loss = 0 |
|
r_losses = [] |
|
g_losses = [] |
|
for dr, dg in zip(disc_real_outputs, disc_generated_outputs): |
|
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()) |
|
|
|
return loss, r_losses, g_losses |
|
|
|
|
|
def generator_loss(disc_outputs): |
|
loss = 0 |
|
gen_losses = [] |
|
for dg in disc_outputs: |
|
l = torch.mean((1 - dg) ** 2) |
|
gen_losses.append(l) |
|
loss += l |
|
|
|
return loss, gen_losses |
|
|
|
|
|
""" https://dl.acm.org/doi/abs/10.1145/3573834.3574506 """ |
|
|
|
|
|
def discriminator_TPRLS_loss(disc_real_outputs, disc_generated_outputs): |
|
loss = 0 |
|
for dr, dg in zip(disc_real_outputs, disc_generated_outputs): |
|
tau = 0.04 |
|
m_DG = torch.median((dr - dg)) |
|
L_rel = torch.mean((((dr - dg) - m_DG) ** 2)[dr < dg + m_DG]) |
|
loss += tau - F.relu(tau - L_rel) |
|
return loss |
|
|
|
|
|
def generator_TPRLS_loss(disc_real_outputs, disc_generated_outputs): |
|
loss = 0 |
|
for dg, dr in zip(disc_real_outputs, disc_generated_outputs): |
|
tau = 0.04 |
|
m_DG = torch.median((dr - dg)) |
|
L_rel = torch.mean((((dr - dg) - m_DG) ** 2)[dr < dg + m_DG]) |
|
loss += tau - F.relu(tau - L_rel) |
|
return loss |
|
|
|
|
|
class GeneratorLoss(torch.nn.Module): |
|
def __init__(self, mpd, msd): |
|
super(GeneratorLoss, self).__init__() |
|
self.mpd = mpd |
|
self.msd = msd |
|
|
|
def forward(self, y, y_hat): |
|
y_df_hat_r, y_df_hat_g, fmap_f_r, fmap_f_g = self.mpd(y, y_hat) |
|
y_ds_hat_r, y_ds_hat_g, fmap_s_r, fmap_s_g = self.msd(y, y_hat) |
|
loss_fm_f = feature_loss(fmap_f_r, fmap_f_g) |
|
loss_fm_s = feature_loss(fmap_s_r, fmap_s_g) |
|
loss_gen_f, losses_gen_f = generator_loss(y_df_hat_g) |
|
loss_gen_s, losses_gen_s = generator_loss(y_ds_hat_g) |
|
|
|
loss_rel = generator_TPRLS_loss(y_df_hat_r, y_df_hat_g) + generator_TPRLS_loss( |
|
y_ds_hat_r, y_ds_hat_g |
|
) |
|
|
|
loss_gen_all = loss_gen_s + loss_gen_f + loss_fm_s + loss_fm_f + loss_rel |
|
|
|
return loss_gen_all.mean() |
|
|
|
|
|
class DiscriminatorLoss(torch.nn.Module): |
|
def __init__(self, mpd, msd): |
|
super(DiscriminatorLoss, self).__init__() |
|
self.mpd = mpd |
|
self.msd = msd |
|
|
|
def forward(self, y, y_hat): |
|
|
|
y_df_hat_r, y_df_hat_g, _, _ = self.mpd(y, y_hat) |
|
loss_disc_f, losses_disc_f_r, losses_disc_f_g = discriminator_loss( |
|
y_df_hat_r, y_df_hat_g |
|
) |
|
|
|
y_ds_hat_r, y_ds_hat_g, _, _ = self.msd(y, y_hat) |
|
loss_disc_s, losses_disc_s_r, losses_disc_s_g = discriminator_loss( |
|
y_ds_hat_r, y_ds_hat_g |
|
) |
|
|
|
loss_rel = discriminator_TPRLS_loss( |
|
y_df_hat_r, y_df_hat_g |
|
) + discriminator_TPRLS_loss(y_ds_hat_r, y_ds_hat_g) |
|
|
|
d_loss = loss_disc_s + loss_disc_f + loss_rel |
|
|
|
return d_loss.mean() |
|
|
|
|
|
class WavLMLoss(torch.nn.Module): |
|
def __init__(self, model, wd, model_sr, slm_sr=16000): |
|
super(WavLMLoss, self).__init__() |
|
self.wavlm = AutoModel.from_pretrained(model) |
|
self.wd = wd |
|
self.resample = torchaudio.transforms.Resample(model_sr, slm_sr) |
|
|
|
def forward(self, wav, y_rec): |
|
with torch.no_grad(): |
|
wav_16 = self.resample(wav) |
|
wav_embeddings = self.wavlm( |
|
input_values=wav_16, output_hidden_states=True |
|
).hidden_states |
|
y_rec_16 = self.resample(y_rec) |
|
y_rec_embeddings = self.wavlm( |
|
input_values=y_rec_16.squeeze(), output_hidden_states=True |
|
).hidden_states |
|
|
|
floss = 0 |
|
for er, eg in zip(wav_embeddings, y_rec_embeddings): |
|
floss += torch.mean(torch.abs(er - eg)) |
|
|
|
return floss.mean() |
|
|
|
def generator(self, y_rec): |
|
y_rec_16 = self.resample(y_rec) |
|
y_rec_embeddings = self.wavlm( |
|
input_values=y_rec_16, output_hidden_states=True |
|
).hidden_states |
|
y_rec_embeddings = ( |
|
torch.stack(y_rec_embeddings, dim=1) |
|
.transpose(-1, -2) |
|
.flatten(start_dim=1, end_dim=2) |
|
) |
|
y_df_hat_g = self.wd(y_rec_embeddings) |
|
loss_gen = torch.mean((1 - y_df_hat_g) ** 2) |
|
|
|
return loss_gen |
|
|
|
def discriminator(self, wav, y_rec): |
|
with torch.no_grad(): |
|
wav_16 = self.resample(wav) |
|
wav_embeddings = self.wavlm( |
|
input_values=wav_16, output_hidden_states=True |
|
).hidden_states |
|
y_rec_16 = self.resample(y_rec) |
|
y_rec_embeddings = self.wavlm( |
|
input_values=y_rec_16, output_hidden_states=True |
|
).hidden_states |
|
|
|
y_embeddings = ( |
|
torch.stack(wav_embeddings, dim=1) |
|
.transpose(-1, -2) |
|
.flatten(start_dim=1, end_dim=2) |
|
) |
|
y_rec_embeddings = ( |
|
torch.stack(y_rec_embeddings, dim=1) |
|
.transpose(-1, -2) |
|
.flatten(start_dim=1, end_dim=2) |
|
) |
|
|
|
y_d_rs = self.wd(y_embeddings) |
|
y_d_gs = self.wd(y_rec_embeddings) |
|
|
|
y_df_hat_r, y_df_hat_g = y_d_rs, y_d_gs |
|
|
|
r_loss = torch.mean((1 - y_df_hat_r) ** 2) |
|
g_loss = torch.mean((y_df_hat_g) ** 2) |
|
|
|
loss_disc_f = r_loss + g_loss |
|
|
|
return loss_disc_f.mean() |
|
|
|
def discriminator_forward(self, wav): |
|
with torch.no_grad(): |
|
wav_16 = self.resample(wav) |
|
wav_embeddings = self.wavlm( |
|
input_values=wav_16, output_hidden_states=True |
|
).hidden_states |
|
y_embeddings = ( |
|
torch.stack(wav_embeddings, dim=1) |
|
.transpose(-1, -2) |
|
.flatten(start_dim=1, end_dim=2) |
|
) |
|
|
|
y_d_rs = self.wd(y_embeddings) |
|
|
|
return y_d_rs |
|
|
