add nsf_univ and export ckpt code

Author: autumn-2-net

export_ckpt.py

@@ -0,0 +1,48 @@
+import importlib
+import pathlib
+
+import click
+import torch
+from tqdm import tqdm
+
+from utils import get_latest_checkpoint_path
+from utils.config_utils import read_full_config
+
+
+
+
+@click.command(help='Train a SOME model')
+@click.option('--exp_name', required=True, metavar='EXP', help='Name of the experiment')
+@click.option('--save_path', required=True, metavar='EXP', help='Name of the experiment')
+@click.option('--work_dir', required=False, metavar='DIR', help='Directory to save the experiment')
+def train( exp_name,save_path, work_dir):
+
+    # print_config(config)
+    if work_dir is None:
+        work_dir = pathlib.Path(__file__).parent / 'experiments'
+    else:
+        work_dir = pathlib.Path(work_dir)
+    work_dir = work_dir / exp_name
+    assert not work_dir.exists() or work_dir.is_dir(), f'Path \'{work_dir}\' is not a directory.'
+    work_dir.mkdir(parents=True, exist_ok=True)
+
+
+
+    ckp = {}
+
+    aaa2x = torch.load(get_latest_checkpoint_path(work_dir))['state_dict']
+    for i in tqdm(aaa2x):
+        i: str
+        if 'generator.' in i:
+            # print(i)
+            ckp[i.replace('generator.', '')] = aaa2x[i]
+
+    torch.save({'generator': ckp}, save_path)
+
+
+
+
+
+
+if __name__ == '__main__':
+    train()

modules/loss/nsf_univloss_msd.py

@@ -0,0 +1,134 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from modules.ddsp.loss import HybridLoss
+from modules.loss.stft_loss import warp_stft
+from utils.wav2mel import PitchAdjustableMelSpectrogram
+
+
+class nsf_univloss_msd(nn.Module):
+    def __init__(self, config: dict):
+        super().__init__()
+        self.mel = PitchAdjustableMelSpectrogram(sample_rate=config['audio_sample_rate'],
+                                                 n_fft=config['fft_size'],
+                                                 win_length=config['win_size'],
+                                                 hop_length=config['hop_size'],
+                                                 f_min=config['fmin'],
+                                                 f_max=config['fmax_for_loss'],
+                                                 n_mels=config['audio_num_mel_bins'], )
+        self.L1loss = nn.L1Loss()
+        self.labauxloss = config.get('lab_aux_loss', 45)
+        self.labddsploss=config.get('lab_ddsp_loss', 2)
+        # self.stft=warp_stft({'fft_sizes':[1024, 2048, 512,],'hop_sizes':[120, 240, 50,],'win_lengths':[600, 1200, 240,]})
+
+        # self.stft = warp_stft(
+        #     {'fft_sizes': [2048, 2048, 4096, 1024, 512, 256, 128], 'hop_sizes': [512, 240, 480, 100, 50, 25, 12],
+        #      'win_lengths': [2048, 1200, 2400, 480, 240, 120, 60]})
+        self.stft = warp_stft({'fft_sizes': config['loss_fft_sizes'], 'hop_sizes': config['loss_hop_sizes'],
+                           'win_lengths': config['loss_win_lengths']})
+
+        self.deuv = config.get('detuv', 2000)
+
+        # self.ddsploss = HybridLoss(block_size=config['hop_size'], fft_min=config['ddsp_fftmin'],
+        #                            fft_max=config['ddsp_fftmax'], n_scale=config['ddsp_nscale'],
+        #                            lambda_uv=config['ddsp_lambdauv'], device='cuda')
+        # fft_sizes = [2048, 4096, 1024, 512, 256, 128],
+        # hop_sizes = [240, 480, 100, 50, 25, 12],
+        # win_lengths = [1200, 2400, 480, 240, 120, 60]
+
+    def discriminator_loss(self, disc_real_outputs, disc_generated_outputs):
+        loss = 0
+        rlosses = 0
+        glosses = 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
+            rlosses += r_loss.item()
+            glosses += g_loss.item()
+            r_losses.append(r_loss.item())
+            g_losses.append(g_loss.item())
+
+        return loss, rlosses, glosses, r_losses, g_losses
+
+    def Dloss(self, Dfake, Dtrue):
+
+        (Fmrd_out, _), (Fmpd_out, _) = Dfake
+        (Tmrd_out, _), (Tmpd_out, _) = Dtrue
+        mrdloss, mrdrlosses, mrdglosses, _, _ = self.discriminator_loss(Tmrd_out, Fmrd_out)
+        mpdloss, mpdrlosses, mpdglosses, _, _ = self.discriminator_loss(Tmpd_out, Fmpd_out)
+        loss = mrdloss + mpdloss
+        return loss, {'DmrdlossF': mrdglosses, 'DmrdlossT': mrdrlosses, 'DmpdlossT': mpdrlosses,
+                      'DmpdlossF': mpdglosses}
+
+    def feature_loss(self, 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 GDloss(self, GDfake, GDtrue):
+        loss = 0
+        gen_losses = []
+        mrd_losses = 0
+        mpd_losses = 0
+        (mrd_out, Fmrd_featrue), (mpd_out, Fmpd_featrue) = GDfake
+        (_, Tmrd_featrue), (_, Tmpd_featrue) = GDtrue
+        for dg in mrd_out:
+            l = torch.mean((1 - dg) ** 2)
+            gen_losses.append(l.item())
+            # loss += l
+            mrd_losses = l + mrd_losses
+
+        for dg in mpd_out:
+            l = torch.mean((1 - dg) ** 2)
+            gen_losses.append(l.item())
+            # loss += l
+            mpd_losses = l + mpd_losses
+
+        mrd_featrue_loss = self.feature_loss(Tmrd_featrue, Fmrd_featrue)
+        mpd_featrue_loss = self.feature_loss(Tmpd_featrue, Fmpd_featrue)
+        # loss +=msd_featrue_loss
+        # loss +=mpd_featrue_loss
+        loss =  mpd_featrue_loss + mpd_losses + mrd_losses+mrd_featrue_loss
+        # (msd_losses, mpd_losses), (msd_featrue_loss, mpd_featrue_loss), gen_losses
+        return loss, {'Gmrdloss': mrd_losses, 'Gmpdloss': mpd_losses, 'Gmrd_featrue_loss': mrd_featrue_loss,
+                      'Dmpd_featrue_loss': mpd_featrue_loss}
+
+    # def Auxloss(self,Goutput, sample):
+    #
+    #     Gmel=self.mel.dynamic_range_compression_torch(self.mel(Goutput['audio'].squeeze(1)))
+    #     # Rmel=sample['mel']
+    #     Rmel = self.mel.dynamic_range_compression_torch(self.mel(sample['audio'].squeeze(1)))
+    #     loss=self.L1loss(Gmel, Rmel)*self.labauxloss
+    #     return loss,{'auxloss':loss}
+
+    def Auxloss(self, Goutput, sample, step):
+
+        # Gmel=self.mel.dynamic_range_compression_torch(self.mel(Goutput['audio'].squeeze(1)))
+        # # Rmel=sample['mel']
+        # Rmel = self.mel.dynamic_range_compression_torch(self.mel(sample['audio'].squeeze(1)))
+        detach_uv = False
+        if step < self.deuv:
+            detach_uv = True
+
+        #
+        # lossddsp, (loss_rss, loss_uv) = self.ddsploss(Goutput['ddspwav'].squeeze(1), Goutput['s_h'],
+        #                                         sample['audio'].squeeze(1),sample['uv'].float(),
+        #                                           detach_uv=detach_uv,
+        #                                           uv_tolerance=0.15)
+
+        # lossddsp=0
+        # loss_rss=0
+        # loss_uv=0
+
+
+        sc_loss, mag_loss = self.stft.stft(Goutput['audio'].squeeze(1), sample['audio'].squeeze(1))
+        loss = (sc_loss + mag_loss) * self.labauxloss
+        return loss, {'auxloss': loss, 'auxloss_sc_loss': sc_loss, 'auxloss_mag_loss': mag_loss,}