Modules/slmadv.py

import torch
import numpy as np
import torch.nn.functional as F

class SLMAdversarialLoss(torch.nn.Module):

    def __init__(self, model, wl, sampler, min_len, max_len, batch_percentage=0.5, skip_update=10, sig=1.5):
        super(SLMAdversarialLoss, self).__init__()
        self.model = model
        self.wl = wl
        self.sampler = sampler
        
        self.min_len = min_len
        self.max_len = max_len
        self.batch_percentage = batch_percentage
        
        self.sig = sig
        self.skip_update = skip_update
        
    def forward(self, iters, y_rec_gt, y_rec_gt_pred, waves, mel_input_length, ref_text, ref_lengths, use_ind, s_trg, ref_s=None):
        text_mask = length_to_mask(ref_lengths).to(ref_text.device)
        bert_dur = self.model.bert(ref_text, attention_mask=(~text_mask).int())
        d_en = self.model.bert_encoder(bert_dur).transpose(-1, -2) 
        
        if use_ind and np.random.rand() < 0.5:
            s_preds = s_trg
        else:
            num_steps = np.random.randint(3, 5)
            if ref_s is not None:
                s_preds = self.sampler(noise = torch.randn_like(s_trg).unsqueeze(1).to(ref_text.device), 
                      embedding=bert_dur,
                      embedding_scale=1,
                               features=ref_s, # reference from the same speaker as the embedding
                         embedding_mask_proba=0.1,
                         num_steps=num_steps).squeeze(1)
            else:
                s_preds = self.sampler(noise = torch.randn_like(s_trg).unsqueeze(1).to(ref_text.device), 
                      embedding=bert_dur,
                      embedding_scale=1,
                         embedding_mask_proba=0.1,
                         num_steps=num_steps).squeeze(1)
            
        s_dur = s_preds[:, 128:]
        s = s_preds[:, :128]
        
        d, _ = self.model.predictor(d_en, s_dur, 
                                                ref_lengths, 
                                                torch.randn(ref_lengths.shape[0], ref_lengths.max(), 2).to(ref_text.device), 
                                                text_mask)
        
        bib = 0

        output_lengths = []
        attn_preds = []
        
        # differentiable duration modeling
        for _s2s_pred, _text_length in zip(d, ref_lengths):

            _s2s_pred_org = _s2s_pred[:_text_length, :]

            _s2s_pred = torch.sigmoid(_s2s_pred_org)
            _dur_pred = _s2s_pred.sum(axis=-1)

            l = int(torch.round(_s2s_pred.sum()).item())
            t = torch.arange(0, l).expand(l)

            t = torch.arange(0, l).unsqueeze(0).expand((len(_s2s_pred), l)).to(ref_text.device)
            loc = torch.cumsum(_dur_pred, dim=0) - _dur_pred / 2

            h = torch.exp(-0.5 * torch.square(t - (l - loc.unsqueeze(-1))) / (self.sig)**2)

            out = torch.nn.functional.conv1d(_s2s_pred_org.unsqueeze(0), 
                                         h.unsqueeze(1), 
                                         padding=h.shape[-1] - 1, groups=int(_text_length))[..., :l]
            attn_preds.append(F.softmax(out.squeeze(), dim=0))

            output_lengths.append(l)

        max_len = max(output_lengths)
        
        with torch.no_grad():
            t_en = self.model.text_encoder(ref_text, ref_lengths, text_mask)
            
        s2s_attn = torch.zeros(len(ref_lengths), int(ref_lengths.max()), max_len).to(ref_text.device)
        for bib in range(len(output_lengths)):
            s2s_attn[bib, :ref_lengths[bib], :output_lengths[bib]] = attn_preds[bib]

        asr_pred = t_en @ s2s_attn

        _, p_pred = self.model.predictor(d_en, s_dur, 
                                                ref_lengths, 
                                                s2s_attn, 
                                                text_mask)
        
        mel_len = max(int(min(output_lengths) / 2 - 1), self.min_len // 2)
        mel_len = min(mel_len, self.max_len // 2)
        
        # get clips
        
        en = []
        p_en = []
        sp = []
        
        F0_fakes = []
        N_fakes = []
        
        wav = []

        for bib in range(len(output_lengths)):
            mel_length_pred = output_lengths[bib]
            mel_length_gt = int(mel_input_length[bib].item() / 2)
            if mel_length_gt <= mel_len or mel_length_pred <= mel_len:
                continue

            sp.append(s_preds[bib])

            random_start = np.random.randint(0, mel_length_pred - mel_len)
            en.append(asr_pred[bib, :, random_start:random_start+mel_len])
            p_en.append(p_pred[bib, :, random_start:random_start+mel_len])

            # get ground truth clips
            random_start = np.random.randint(0, mel_length_gt - mel_len)
            y = waves[bib][(random_start * 2) * 300:((random_start+mel_len) * 2) * 300]
            wav.append(torch.from_numpy(y).to(ref_text.device))
            
            if len(wav) >= self.batch_percentage * len(waves): # prevent OOM due to longer lengths
                break

        if len(sp) <= 1:
            return None
            
        sp = torch.stack(sp)
        wav = torch.stack(wav).float()
        en = torch.stack(en)
        p_en = torch.stack(p_en)
        
        F0_fake, N_fake = self.model.predictor.F0Ntrain(p_en, sp[:, 128:])
        y_pred = self.model.decoder(en, F0_fake, N_fake, sp[:, :128])
        
        # discriminator loss
        if (iters + 1) % self.skip_update == 0:
            if np.random.randint(0, 2) == 0:
                wav = y_rec_gt_pred
                use_rec = True
            else:
                use_rec = False

            crop_size = min(wav.size(-1), y_pred.size(-1))
            if use_rec: # use reconstructed (shorter lengths), do length invariant regularization
                if wav.size(-1) > y_pred.size(-1):
                    real_GP = wav[:, : , :crop_size]
                    out_crop = self.wl.discriminator_forward(real_GP.detach().squeeze())
                    out_org = self.wl.discriminator_forward(wav.detach().squeeze())
                    loss_reg = F.l1_loss(out_crop, out_org[..., :out_crop.size(-1)])

                    if np.random.randint(0, 2) == 0:
                        d_loss = self.wl.discriminator(real_GP.detach().squeeze(), y_pred.detach().squeeze()).mean()
                    else:
                        d_loss = self.wl.discriminator(wav.detach().squeeze(), y_pred.detach().squeeze()).mean()
                else:
                    real_GP = y_pred[:, : , :crop_size]
                    out_crop = self.wl.discriminator_forward(real_GP.detach().squeeze())
                    out_org = self.wl.discriminator_forward(y_pred.detach().squeeze())
                    loss_reg = F.l1_loss(out_crop, out_org[..., :out_crop.size(-1)])

                    if np.random.randint(0, 2) == 0:
                        d_loss = self.wl.discriminator(wav.detach().squeeze(), real_GP.detach().squeeze()).mean()
                    else:
                        d_loss = self.wl.discriminator(wav.detach().squeeze(), y_pred.detach().squeeze()).mean()
                
                # regularization (ignore length variation)
                d_loss += loss_reg

                out_gt = self.wl.discriminator_forward(y_rec_gt.detach().squeeze())
                out_rec = self.wl.discriminator_forward(y_rec_gt_pred.detach().squeeze())

                # regularization (ignore reconstruction artifacts)
                d_loss += F.l1_loss(out_gt, out_rec)

            else:
                d_loss = self.wl.discriminator(wav.detach().squeeze(), y_pred.detach().squeeze()).mean()
        else:
            d_loss = 0
            
        # generator loss
        gen_loss = self.wl.generator(y_pred.squeeze())
        
        gen_loss = gen_loss.mean()
        
        return d_loss, gen_loss, y_pred.detach().cpu().numpy()
    
def length_to_mask(lengths):
    mask = torch.arange(lengths.max()).unsqueeze(0).expand(lengths.shape[0], -1).type_as(lengths)
    mask = torch.gt(mask+1, lengths.unsqueeze(1))
    return mask


# import torch
# import numpy as np
# import torch.nn.functional as F

# from accelerate import Accelerator, DistributedDataParallelKwargs
# from accelerate.utils import tqdm, ProjectConfiguration

# class SLMAdversarialLoss(torch.nn.Module):

#     def __init__(self, model, wl, sampler, min_len, max_len, batch_percentage=0.5, skip_update=10, sig=1.5):
#         super(SLMAdversarialLoss, self).__init__()
#         self.model = model
#         self.wl = wl
#         self.sampler = sampler
        
#         self.min_len = min_len
#         self.max_len = max_len
#         self.batch_percentage = batch_percentage
        
#         self.sig = sig
#         self.skip_update = skip_update
        
#     def forward(self, iters, accelerator, y_rec_gt, y_rec_gt_pred, waves, mel_input_length, ref_text, ref_lengths, use_ind, s_trg, ref_s=None):
#         text_mask = length_to_mask(ref_lengths).to(ref_text.device)
#         bert_dur = self.model.bert(ref_text, attention_mask=(~text_mask).int())
#         d_en = self.model.bert_encoder(bert_dur).transpose(-1, -2) 

#         if use_ind and np.random.rand() < 0.5:
#             s_preds = s_trg
#         else:
#             num_steps = np.random.randint(3, 5)
#             if ref_s is not None:
#                 s_preds = self.sampler(noise = torch.randn_like(s_trg).unsqueeze(1).to(ref_text.device), 
#                       embedding=bert_dur,
#                       embedding_scale=1,
#                                features=ref_s, # reference from the same speaker as the embedding
#                          embedding_mask_proba=0.1,
#                          num_steps=num_steps).squeeze(1)
#             else:
#                 s_preds = self.sampler(noise = torch.randn_like(s_trg).unsqueeze(1).to(ref_text.device), 
#                       embedding=bert_dur,
#                       embedding_scale=1,
#                          embedding_mask_proba=0.1,
#                          num_steps=num_steps).squeeze(1)
            
#         s_dur = s_preds[:, 128:]
#         s = s_preds[:, :128]
        
#         d, _ = self.model.predictor(d_en, s_dur, 
#                                                 ref_lengths, 
#                                                 torch.randn(ref_lengths.shape[0], ref_lengths.max(), 2).to(ref_text.device), 
#                                                 text_mask)
        
#         bib = 0

#         output_lengths = []
#         attn_preds = []
        
#         # differentiable duration modeling
#         for _s2s_pred, _text_length in zip(d, ref_lengths):

#             _s2s_pred_org = _s2s_pred[:_text_length, :]

#             _s2s_pred = torch.sigmoid(_s2s_pred_org)
#             _dur_pred = _s2s_pred.sum(axis=-1)

#             l = int(torch.round(_s2s_pred.sum()).item())
#             t = torch.arange(0, l).expand(l)

#             t = torch.arange(0, l).unsqueeze(0).expand((len(_s2s_pred), l)).to(ref_text.device)
#             loc = torch.cumsum(_dur_pred, dim=0) - _dur_pred / 2

#             h = torch.exp(-0.5 * torch.square(t - (l - loc.unsqueeze(-1))) / (self.sig)**2)

#             out = torch.nn.functional.conv1d(_s2s_pred_org.unsqueeze(0), 
#                                          h.unsqueeze(1), 
#                                          padding=h.shape[-1] - 1, groups=int(_text_length))[..., :l]
#             attn_preds.append(F.softmax(out.squeeze(), dim=0))

#             output_lengths.append(l)

#         max_len = max(output_lengths)
        
#         with torch.no_grad():
#             t_en = self.model.text_encoder(ref_text, ref_lengths, text_mask)
            
#         s2s_attn = torch.zeros(len(ref_lengths), int(ref_lengths.max()), max_len).to(ref_text.device)
#         for bib in range(len(output_lengths)):
#             s2s_attn[bib, :ref_lengths[bib], :output_lengths[bib]] = attn_preds[bib]

#         asr_pred = t_en @ s2s_attn

#         _, p_pred = self.model.predictor(d_en, s_dur, 
#                                                 ref_lengths, 
#                                                 s2s_attn, 
#                                                 text_mask)
        
#         mel_len = max(int(min(output_lengths) / 2 - 1), self.min_len // 2)
#         mel_len = min(mel_len, self.max_len // 2)
        
#         # get clips
        
#         en = []
#         p_en = []
#         sp = []
        
#         F0_fakes = []
#         N_fakes = []
        
#         wav = []

#         for bib in range(len(output_lengths)):
#             mel_length_pred = output_lengths[bib]
#             mel_length_gt = int(mel_input_length[bib].item() / 2)
#             if mel_length_gt <= mel_len or mel_length_pred <= mel_len:
#                 continue

#             sp.append(s_preds[bib])

#             random_start = np.random.randint(0, mel_length_pred - mel_len)
#             en.append(asr_pred[bib, :, random_start:random_start+mel_len])
#             p_en.append(p_pred[bib, :, random_start:random_start+mel_len])

#             # get ground truth clips
#             random_start = np.random.randint(0, mel_length_gt - mel_len)
#             y = waves[bib][(random_start * 2) * 300:((random_start+mel_len) * 2) * 300]
#             wav.append(torch.from_numpy(y).to(ref_text.device))
            
#             if len(wav) >= self.batch_percentage * len(waves): # prevent OOM due to longer lengths
#                 break
            
#         # global_min_batch = accelerator.gather(torch.tensor([len(wav)], device=ref_text.device)).min().item()
#         # if global_min_batch <= 1:
#         #     raise ValueError("skip slmadv")

#         if len(sp) <= 1:
#             return None
            
#         sp = torch.stack(sp)
#         wav = torch.stack(wav).float()
#         en = torch.stack(en)
#         p_en = torch.stack(p_en)

#         F0_fake, N_fake = self.model.predictor(texts=p_en, style=sp[:, 128:], f0=True)
#         y_pred = self.model.decoder(en, F0_fake, N_fake, sp[:, :128])
        
#         # discriminator loss
#         if (iters + 1) % self.skip_update == 0:
#             if np.random.randint(0, 2) == 0:
#                 wav = y_rec_gt_pred
#                 use_rec = True
#             else:
#                 use_rec = False

#             crop_size = min(wav.size(-1), y_pred.size(-1))
#             if use_rec: # use reconstructed (shorter lengths), do length invariant regularization
#                 if wav.size(-1) > y_pred.size(-1):
#                     real_GP = wav[:, : , :crop_size]
#                     out_crop = self.wl(wav = real_GP.detach().squeeze(),y_rec=None, discriminator_forward=True)
#                     out_org = self.wl(wav = wav.detach().squeeze(),y_rec=None, discriminator_forward=True)
#                     loss_reg = F.l1_loss(out_crop, out_org[..., :out_crop.size(-1)])

#                     if np.random.randint(0, 2) == 0:
#                         d_loss = self.wl(wav = real_GP.detach().squeeze(),y_rec= y_pred.detach().squeeze(), discriminator=True).mean()
#                     else:
#                         d_loss = self.wl(wav = wav.detach().squeeze(), y_rec = y_pred.detach().squeeze(), discriminator=True).mean()
#                 else:
#                     real_GP = y_pred[:, : , :crop_size]
#                     out_crop = self.wl(wav = real_GP.detach().squeeze(), y_rec=None, discriminator_forward=True)
#                     out_org = self.wl(wav = y_pred.detach().squeeze(),y_rec=None, discriminator_forward=True)
#                     loss_reg = F.l1_loss(out_crop, out_org[..., :out_crop.size(-1)])

#                     if np.random.randint(0, 2) == 0:
#                         d_loss = self.wl(wav = wav.detach().squeeze(), y_rec = real_GP.detach().squeeze(), discriminator=True ).mean()
#                     else:
#                         d_loss = self.wl(wav = wav.detach().squeeze(), y_rec = y_pred.detach().squeeze(), discriminator=True).mean()
                
#                 # regularization (ignore length variation)
#                 d_loss += loss_reg

#                 out_gt = self.wl(wav = y_rec_gt.detach().squeeze(),y_rec=None, discriminator_forward=True)
#                 out_rec = self.wl(wav = y_rec_gt_pred.detach().squeeze(), y_rec=None, discriminator_forward=True)

#                 # regularization (ignore reconstruction artifacts)
#                 d_loss += F.l1_loss(out_gt, out_rec)

#             else:
#                 d_loss = self.wl(wav = wav.detach().squeeze(),y_rec= y_pred.detach().squeeze(), discriminator=True).mean()
#         else:
#             d_loss = 0
            
#         # generator loss
#         gen_loss = self.wl(wav = None, y_rec = y_pred.squeeze(), generator=True)
        
#         gen_loss = gen_loss.mean()
        
        
#         return d_loss, gen_loss, y_pred.detach().cpu().numpy()
    
# def length_to_mask(lengths):
#     mask = torch.arange(lengths.max()).unsqueeze(0).expand(lengths.shape[0], -1).type_as(lengths)
#     mask = torch.gt(mask+1, lengths.unsqueeze(1))
#     return mask