Upload 2 files

wavvae 2.zip

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6e65acb0f186a7d25a66b8211773cf5a1030d1edd972f11882cf5ca2135521dd
+size 43138

wavvae3.py

@@ -0,0 +1,268 @@
+import argparse
+import filecmp
+import multiprocessing
+import os
+import subprocess
+import librosa
+from functools import partial
+from multiprocessing import Pool, Process
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.distributed as dist
+from torch.optim import AdamW
+
+from modules.vocoder.commons.stft_loss import MultiResolutionSTFTLoss
+from modules.vocoder.hifigan.hifigan import MultiPeriodDiscriminator, MultiScaleDiscriminator, \
+    generator_loss, feature_loss, discriminator_loss
+from modules.vocoder.hifigan.mel_utils import mel_spectrogram
+from modules.vocoder.univnet.mrd import MultiResolutionDiscriminator
+from modules.tts.wavvae.decoder.wavvae_v3 import WavVAE_V3
+from tasks.tts.utils.audio import torch_wav2spec
+from tasks.tts.utils.audio.align import mel2token_to_dur
+from utils.commons.ckpt_utils import load_ckpt
+from utils.commons.hparams import hparams
+
+from attrdict import AttrDict
+from tasks.tts.dataset_mixin import TTSDatasetMixin
+from utils.commons.base_task import BaseTask
+from utils.commons.import_utils import import_module_bystr
+from utils.nn.schedulers import WarmupSchedule, CosineSchedule
+
+
+class WavVAETask(TTSDatasetMixin, BaseTask):
+    def __init__(self):
+        super().__init__()
+        self.dataset_cls = import_module_bystr(hparams['dataset_cls'])
+        self.val_dataset_cls = import_module_bystr(hparams['val_dataset_cls'])
+        self.processer_fn = import_module_bystr(hparams['processer_fn'])
+        self.build_fast_dataloader = import_module_bystr(hparams['build_fast_dataloader'])
+        self.hparams = hparams
+        self.config = AttrDict(hparams)
+
+        # Online load mel with GPU
+        sample_rate = hparams["audio_sample_rate"]
+        fft_size = hparams["win_size"]
+        win_size = hparams["win_size"]
+        hop_size = hparams["hop_size"]
+        num_mels = hparams["audio_num_mel_bins"]
+        fmin = hparams["fmin"]
+        fmax = hparams["fmax"]
+        mel_basis = librosa.filters.mel(
+            sr=sample_rate, n_fft=fft_size, n_mels=num_mels, fmin=fmin, fmax=fmax
+        )
+        self.torch_wav2spec_ = partial(
+            torch_wav2spec, mel_basis=mel_basis, fft_size=fft_size, hop_size=hop_size, win_length=win_size,
+        )
+
+    def build_model(self):
+        self.model_gen = WavVAE_V3(hparams=hparams)
+
+        self.model_disc = torch.nn.ModuleDict()
+        self.model_disc['mpd'] = MultiPeriodDiscriminator(hparams['mpd'], use_cond=hparams['use_cond_disc'])
+        self.model_disc['msd'] = MultiScaleDiscriminator(use_cond=hparams['use_cond_disc'])
+        if hparams['use_mrd']:
+            self.model_disc['mrd'] = MultiResolutionDiscriminator(hparams)
+        self.stft_loss = MultiResolutionSTFTLoss()
+
+        load_ckpt(self.model_gen.encoder, './checkpoints/1231_megatts3_wavvae_v2_25hz', 'model.module.encoder', strict=False)
+        load_ckpt(self.model_gen.decoder, './checkpoints/1117_melgan-nsf_full_1', 'model_gen', force=True, strict=True)
+        load_ckpt(self.model_disc, './checkpoints/1117_melgan-nsf_full_1', 'model_disc', force=True, strict=True)
+        return {'trainable': [self.model_gen, self.model_disc['mpd'], self.model_disc['msd'], self.model_disc['mrd']], 'others': []}
+
+    def load_model(self):
+        if hparams.get('load_ckpt', '') != '':
+            load_ckpt(self.model, hparams['load_ckpt'], 'model', strict=False)
+
+    def build_optimizer(self):
+        optimizer_gen = torch.optim.AdamW(self.model_gen.parameters(), lr=hparams['lr'],
+                                          betas=[hparams['adam_b1'], hparams['adam_b2']])
+        optimizer_disc = torch.optim.AdamW(self.model_disc.parameters(),
+                                           lr=hparams.get('disc_lr', hparams['lr']),
+                                           betas=[hparams['adam_b1'], hparams['adam_b2']])
+        return [optimizer_gen, optimizer_disc]
+
+    def build_scheduler(self, optimizer):
+        return None
+
+    def _training_step(self, sample, batch_idx, optimizer_idx):
+        log_outputs = {}
+        loss_weights = {}
+        sample['wavs'] = sample['wavs'].float()
+        # return None, {}
+
+        if self.global_step % 100 == 0:
+            devices = os.environ.get('CUDA_VISIBLE_DEVICES', '').split(",")
+            for d in devices:
+                os.system(f'pkill -f "voidgpu{d}"')
+
+        y = sample['wavs']
+        loss_output = {}
+        if optimizer_idx == 0:
+            #######################
+            #      Generator      #
+            #######################
+            y_, posterior = self.model_gen(y)
+            y = y.unsqueeze(1)
+            y_mel = mel_spectrogram(y.squeeze(1), hparams).transpose(1, 2)
+            y_hat_mel = mel_spectrogram(y_.squeeze(1), hparams).transpose(1, 2)
+            loss_output['mel'] = F.l1_loss(y_hat_mel, y_mel) * hparams['lambda_mel']
+            if self.training:
+                _, y_p_hat_g, fmap_f_r, fmap_f_g = self.model_disc['mpd'](y, y_, None)
+                _, y_s_hat_g, fmap_s_r, fmap_s_g = self.model_disc['msd'](y, y_, None)
+                loss_output['a_p'] = generator_loss(y_p_hat_g) * hparams['lambda_adv'] * hparams.get('lambda_mpd', 1.0)
+                loss_output['a_s'] = generator_loss(y_s_hat_g) * hparams['lambda_adv'] * hparams.get('lambda_msd', 1.0)
+                if hparams['use_mrd']:
+                    y_r_hat_g = [x[1] for x in self.model_disc['mrd'](y_)]
+                    loss_output['a_r'] = generator_loss(y_r_hat_g) \
+                                         * hparams['lambda_adv'] * hparams.get('lambda_mrd', 1.0)
+                if hparams['use_ms_stft']:
+                    loss_output['sc'], loss_output['mag'] = self.stft_loss(y.squeeze(1), y_.squeeze(1))
+                loss_output['kl_loss'] = posterior.kl().mean() * hparams.get('lambda_kl', 1.0)
+            self.y_ = y_.detach()
+        else:
+            #######################
+            #    Discriminator    #
+            #######################
+            if not self.training:
+                return None
+            y = y.unsqueeze(1)
+            y_ = self.y_
+            # MPD
+            y_p_hat_r, y_p_hat_g, _, _ = self.model_disc['mpd'](y, y_.detach(), None)
+            loss_output['r_p'], loss_output['f_p'] = discriminator_loss(y_p_hat_r, y_p_hat_g)
+            # MSD
+            y_s_hat_r, y_s_hat_g, _, _ = self.model_disc['msd'](y, y_.detach(), None)
+            loss_output['r_s'], loss_output['f_s'] = discriminator_loss(y_s_hat_r, y_s_hat_g)
+            # MRD
+            if hparams['use_mrd']:
+                y_r_hat_r = [x[1] for x in self.model_disc['mrd'](y)]
+                y_r_hat_g = [x[1] for x in self.model_disc['mrd'](y_.detach())]
+                loss_output['r_r'], loss_output['f_r'] = discriminator_loss(y_r_hat_r, y_r_hat_g)
+        total_loss = sum(loss_output.values())
+        loss_output['bs'] = sample['wavs'].shape[0]
+        return total_loss, loss_output
+
+    def save_valid_result(self, sample, batch_idx, model_out):
+        sr = hparams['audio_sample_rate']
+        mel_out = model_out.get('mel_out')
+        f0 = sample.get('f0')
+        f0_gt = sample.get('f0')
+        if f0 is not None:
+            f0_gt = f0_gt.cpu()[-1]
+        if mel_out is not None:
+            f0_pred = self.predict_f0(sample['mels'])
+            self.plot_mel(batch_idx, sample['mels'], mel_out, f0s={'f0': f0_pred, 'f0g': f0_gt})
+        # gt wav
+        if self.global_step <= hparams['valid_infer_interval']:
+            mel_gt = sample['mels'][-1].cpu()
+            f0 = self.predict_f0(sample['mels'][-1:])
+            wav_gt = self.vocoder.spec2wav(mel_gt, f0=f0)
+            self.logger.add_audio(f'wav_gt_{batch_idx}', wav_gt, self.global_step, sr)
+
+        if self.global_step >= 0:
+            # with gt duration
+            model_out = self.run_model(sample, infer=True, infer_use_gt_dur=True)
+            # dur_info = self.get_plot_dur_info(sample, model_out)
+            # del dur_info['dur_pred']
+            dur_info = None
+
+            f0 = self.predict_f0(model_out['mel_out'])
+            wav_pred = self.vocoder.spec2wav(model_out['mel_out'][-1].cpu(), f0=f0)
+            self.logger.add_audio(f'wav_gdur_{batch_idx}', wav_pred, self.global_step, sr)
+            self.plot_mel(batch_idx, sample['mels'][-1:], model_out['mel_out'][-1], f'mel_gdur_{batch_idx}',
+                          dur_info=dur_info, f0s={'f0': f0, 'f0g': f0_gt})
+
+            # with pred duration
+            if not hparams['use_gt_dur'] and not hparams['use_gt_latent']:
+                model_out = self.run_model(sample, infer=True, infer_use_gt_dur=False)
+                # dur_info = self.get_plot_dur_info(sample, model_out)
+                dur_info = None
+                f0 = self.predict_f0(model_out['mel_out'])
+                self.plot_mel(
+                    batch_idx, sample['mels'], model_out['mel_out'][-1], f'mel_pdur_{batch_idx}',
+                    dur_info=dur_info, f0s={'f0': f0, 'f0g': f0_gt})
+                wav_pred = self.vocoder.spec2wav(model_out['mel_out'][-1].cpu(), f0=f0)
+                self.logger.add_audio(f'wav_pdur_{batch_idx}', wav_pred, self.global_step, sr)
+
+    def get_plot_dur_info(self, sample, model_out):
+        T_txt = sample['txt_tokens'].shape[1]
+        dur_gt = mel2token_to_dur(sample['mel2ph'], T_txt)[-1]
+        dur_pred = model_out['dur'] if 'dur' in model_out else dur_gt
+        txt = self.token_encoder.decode(sample['txt_tokens'][-1].cpu().numpy())
+        txt = txt.split(" ")
+        return {'dur_gt': dur_gt, 'dur_pred': dur_pred, 'txt': txt}
+
+    def on_before_optimization(self, opt_idx):
+        if opt_idx == 0:
+            nn.utils.clip_grad_norm_(self.model_gen.parameters(), hparams['generator_grad_norm'])
+        else:
+            nn.utils.clip_grad_norm_(self.model_disc.parameters(), hparams["discriminator_grad_norm"])
+
+    def to(self, device=None, dtype=None):
+        super().to(device=device, dtype=dtype)
+        # trainer doesn't move ema to device automatically, we do it mannually
+        if hparams.get('use_ema', False):
+            self.ema.to(device=device, dtype=dtype)
+
+    def cuda(self,device):
+        super().cuda(device)
+        if hparams.get('use_ema', False):
+            self.ema.to(device=device)
+    
+    @torch.no_grad()
+    def validation_step(self, sample, batch_idx):
+        infer_steps = self.hparams.get('infer_steps', 12)
+        outputs = self._validation_step(sample, batch_idx, infer_steps)
+        return outputs
+
+    def _validation_step(self, sample, batch_idx, infer_steps):
+        outputs = {}
+        if self.trainer.proc_rank == 0:
+            # self.vae.eval()
+            # with torch.inference_mode():
+            #     with torch.cuda.amp.autocast(dtype=torch.bfloat16, enabled=True):
+            #         lat = self.vae.get_latent(sample["mels"])
+            #         lat_lens = latent_lengths.clamp(max=lat.size(1))
+            # mel = self.vae.decode(lat)
+            pass
+            # outputs['losses'], _ = self.run_model(sample)
+            # _, model_out = self.run_model(sample, infer=True, infer_steps=infer_steps)
+            # outputs = tensors_to_scalars(outputs)
+            # output_ldm = model_out['ldm_out']
+            # T = output_ldm.shape[1]
+            # ldm = sample['kps'][:, :T]  # [B, T, nkp, kp_dim] [0, 1]
+            # B, T, nkp, kp_dim = ldm.shape
+            # output_ldm = self.denormalize_ldm(output_ldm)
+            # recon_ldm = model_out['recon_ldm']
+            # recon_ldm = self.denormalize_ldm(recon_ldm)
+
+            # results_dir = f"{hparams['work_dir']}/results/{self.global_step}_infersteps{infer_steps}_cfg{hparams['cfg_w']}"
+            # os.makedirs(results_dir, exist_ok=True)
+            # n_ctx = model_out['ctx_mask'][0, :, 0].sum().long().item()
+            # writer_kp = imageio.get_writer(f"{results_dir}/{batch_idx:06d}_kp.sil.mp4", fps=25)
+            # writer_gt = imageio.get_writer(f"{results_dir}/{batch_idx:06d}_gt.sil.mp4", fps=25)
+            # writer_pred = imageio.get_writer(f"{results_dir}/{batch_idx:06d}_pred.sil.mp4", fps=25)
+            # for i in range(T):
+            #     img = self.draw_ldm(recon_ldm[0, i])
+            #     writer_gt.append_data(img)
+            #     img = self.draw_ldm(ldm[0, i])
+            #     writer_kp.append_data(img)
+            #     if i < n_ctx:
+            #         writer_pred.append_data(img)
+            #     else:
+            #         img = self.draw_ldm(
+            #             output_ldm[0, i], color=(255, 255, 0),
+            #         )
+            #         writer_pred.append_data(img)
+            # writer_gt.close()
+            # writer_kp.close()
+            # writer_pred.close()
+        return outputs
+
+    @torch.no_grad()
+    def test_step(self, sample, batch_idx):
+        infer_steps = hparams['infer_steps']
+        return self._validation_step(sample, batch_idx, infer_steps)