import matplotlib matplotlib.use('Agg') from tasks.tts.tts_base import TTSBaseTask from vocoders.base_vocoder import get_vocoder_cls from tasks.tts.dataset_utils import FastSpeechDataset from modules.commons.ssim import ssim import os from modules.fastspeech.tts_modules import mel2ph_to_dur from utils.hparams import hparams from utils.plot import spec_to_figure, dur_to_figure, f0_to_figure from utils.pitch_utils import denorm_f0 from modules.fastspeech.fs2 import FastSpeech2 import torch import torch.optim import torch.utils.data import torch.nn.functional as F import utils import torch.distributions import numpy as np class FastSpeech2Task(TTSBaseTask): def __init__(self): super(FastSpeech2Task, self).__init__() self.dataset_cls = FastSpeechDataset self.mse_loss_fn = torch.nn.MSELoss() mel_losses = hparams['mel_loss'].split("|") self.loss_and_lambda = {} for i, l in enumerate(mel_losses): if l == '': continue if ':' in l: l, lbd = l.split(":") lbd = float(lbd) else: lbd = 1.0 self.loss_and_lambda[l] = lbd print("| Mel losses:", self.loss_and_lambda) self.sil_ph = self.phone_encoder.sil_phonemes() f0_stats_fn = f'{hparams["binary_data_dir"]}/train_f0s_mean_std.npy' if os.path.exists(f0_stats_fn): hparams['f0_mean'], hparams['f0_std'] = np.load(f0_stats_fn) hparams['f0_mean'] = float(hparams['f0_mean']) hparams['f0_std'] = float(hparams['f0_std']) def build_tts_model(self): self.model = FastSpeech2(self.phone_encoder) def build_model(self): self.build_tts_model() if hparams['load_ckpt'] != '': self.load_ckpt(hparams['load_ckpt'], strict=False) utils.print_arch(self.model) return self.model def _training_step(self, sample, batch_idx, _): loss_output = self.run_model(self.model, sample) total_loss = sum([v for v in loss_output.values() if isinstance(v, torch.Tensor) and v.requires_grad]) loss_output['batch_size'] = sample['txt_tokens'].size()[0] return total_loss, loss_output def validation_step(self, sample, batch_idx): outputs = {} outputs['losses'] = {} outputs['losses'], model_out = self.run_model(self.model, sample, return_output=True) outputs['total_loss'] = sum(outputs['losses'].values()) outputs['nsamples'] = sample['nsamples'] mel_out = self.model.out2mel(model_out['mel_out']) outputs = utils.tensors_to_scalars(outputs) if self.global_step % hparams['valid_infer_interval'] == 0 \ and batch_idx < hparams['num_valid_plots']: vmin = hparams['mel_vmin'] vmax = hparams['mel_vmax'] self.plot_mel(batch_idx, sample['mels'], mel_out) self.plot_dur(batch_idx, sample, model_out) if hparams['use_pitch_embed']: self.plot_pitch(batch_idx, sample, model_out) if self.vocoder is None: self.vocoder = get_vocoder_cls(hparams)() if self.global_step > 0: spk_embed = sample.get('spk_embed') if not hparams['use_spk_id'] else sample.get('spk_ids') # with gt duration model_out = self.model(sample['txt_tokens'], mel2ph=sample['mel2ph'], spk_embed=spk_embed, infer=True) wav_pred = self.vocoder.spec2wav(model_out['mel_out'][0].cpu()) self.logger.add_audio(f'wav_gtdur_{batch_idx}', wav_pred, self.global_step, hparams['audio_sample_rate']) self.logger.add_figure( f'mel_gtdur_{batch_idx}', spec_to_figure(model_out['mel_out'][0], vmin, vmax), self.global_step) # with pred duration model_out = self.model(sample['txt_tokens'], spk_embed=spk_embed, infer=True) self.logger.add_figure( f'mel_{batch_idx}', spec_to_figure(model_out['mel_out'][0], vmin, vmax), self.global_step) wav_pred = self.vocoder.spec2wav(model_out['mel_out'][0].cpu()) self.logger.add_audio(f'wav_{batch_idx}', wav_pred, self.global_step, hparams['audio_sample_rate']) # gt wav if self.global_step <= hparams['valid_infer_interval']: mel_gt = sample['mels'][0].cpu() wav_gt = self.vocoder.spec2wav(mel_gt) self.logger.add_audio(f'wav_gt_{batch_idx}', wav_gt, self.global_step, 22050) return outputs def run_model(self, model, sample, return_output=False): txt_tokens = sample['txt_tokens'] # [B, T_t] target = sample['mels'] # [B, T_s, 80] mel2ph = sample['mel2ph'] # [B, T_s] f0 = sample['f0'] uv = sample['uv'] energy = sample['energy'] spk_embed = sample.get('spk_embed') if not hparams['use_spk_id'] else sample.get('spk_ids') output = model(txt_tokens, mel2ph=mel2ph, spk_embed=spk_embed, ref_mels=target, f0=f0, uv=uv, energy=energy, tgt_mels=target, infer=False) losses = {} self.add_mel_loss(output['mel_out'], target, losses) self.add_dur_loss(output['dur'], mel2ph, txt_tokens, losses=losses) if hparams['use_pitch_embed']: self.add_pitch_loss(output, sample, losses) if not return_output: return losses else: return losses, output ############ # losses ############ def add_mel_loss(self, mel_out, target, losses, postfix='', mel_mix_loss=None): nonpadding = target.abs().sum(-1).ne(0).float() for loss_name, lbd in self.loss_and_lambda.items(): if 'l1' == loss_name: l = self.l1_loss(mel_out, target) elif 'mse' == loss_name: l = self.mse_loss(mel_out, target) elif 'ssim' == loss_name: l = self.ssim_loss(mel_out, target) elif 'gdl' == loss_name: l = self.gdl_loss_fn(mel_out, target, nonpadding) \ * self.loss_and_lambda['gdl'] losses[f'{loss_name}{postfix}'] = l * lbd def l1_loss(self, decoder_output, target): # decoder_output : B x T x n_mel # target : B x T x n_mel l1_loss = F.l1_loss(decoder_output, target, reduction='none') weights = self.weights_nonzero_speech(target) l1_loss = (l1_loss * weights).sum() / weights.sum() return l1_loss def add_energy_loss(self, energy_pred, energy, losses): nonpadding = (energy != 0).float() loss = (F.mse_loss(energy_pred, energy, reduction='none') * nonpadding).sum() / nonpadding.sum() loss = loss * hparams['lambda_energy'] losses['e'] = loss def mse_loss(self, decoder_output, target): # decoder_output : B x T x n_mel # target : B x T x n_mel assert decoder_output.shape == target.shape mse_loss = F.mse_loss(decoder_output, target, reduction='none') weights = self.weights_nonzero_speech(target) mse_loss = (mse_loss * weights).sum() / weights.sum() return mse_loss def ssim_loss(self, decoder_output, target, bias=6.0): # decoder_output : B x T x n_mel # target : B x T x n_mel assert decoder_output.shape == target.shape weights = self.weights_nonzero_speech(target) decoder_output = decoder_output[:, None] + bias target = target[:, None] + bias ssim_loss = 1 - ssim(decoder_output, target, size_average=False) ssim_loss = (ssim_loss * weights).sum() / weights.sum() return ssim_loss def add_dur_loss(self, dur_pred, mel2ph, txt_tokens, losses=None): """ :param dur_pred: [B, T], float, log scale :param mel2ph: [B, T] :param txt_tokens: [B, T] :param losses: :return: """ B, T = txt_tokens.shape nonpadding = (txt_tokens != 0).float() dur_gt = mel2ph_to_dur(mel2ph, T).float() * nonpadding is_sil = torch.zeros_like(txt_tokens).bool() for p in self.sil_ph: is_sil = is_sil | (txt_tokens == self.phone_encoder.encode(p)[0]) is_sil = is_sil.float() # [B, T_txt] losses['pdur'] = F.mse_loss(dur_pred, (dur_gt + 1).log(), reduction='none') losses['pdur'] = (losses['pdur'] * nonpadding).sum() / nonpadding.sum() losses['pdur'] = losses['pdur'] * hparams['lambda_ph_dur'] dur_pred = (dur_pred.exp() - 1).clamp(min=0) # use linear scale for sent and word duration if hparams['lambda_word_dur'] > 0: word_id = (is_sil.cumsum(-1) * (1 - is_sil)).long() word_dur_p = dur_pred.new_zeros([B, word_id.max() + 1]).scatter_add(1, word_id, dur_pred)[:, 1:] word_dur_g = dur_gt.new_zeros([B, word_id.max() + 1]).scatter_add(1, word_id, dur_gt)[:, 1:] wdur_loss = F.mse_loss((word_dur_p + 1).log(), (word_dur_g + 1).log(), reduction='none') word_nonpadding = (word_dur_g > 0).float() wdur_loss = (wdur_loss * word_nonpadding).sum() / word_nonpadding.sum() losses['wdur'] = wdur_loss * hparams['lambda_word_dur'] if hparams['lambda_sent_dur'] > 0: sent_dur_p = dur_pred.sum(-1) sent_dur_g = dur_gt.sum(-1) sdur_loss = F.mse_loss((sent_dur_p + 1).log(), (sent_dur_g + 1).log(), reduction='mean') losses['sdur'] = sdur_loss.mean() * hparams['lambda_sent_dur'] def add_pitch_loss(self, output, sample, losses): mel2ph = sample['mel2ph'] # [B, T_s] f0 = sample['f0'] uv = sample['uv'] nonpadding = (mel2ph != 0).float() if hparams['pitch_type'] == 'frame' \ else (sample['txt_tokens'] != 0).float() self.add_f0_loss(output['pitch_pred'], f0, uv, losses, nonpadding=nonpadding) # output['pitch_pred']: [B, T, 2], f0: [B, T], uv: [B, T] def add_f0_loss(self, p_pred, f0, uv, losses, nonpadding, postfix=''): assert p_pred[..., 0].shape == f0.shape if hparams['use_uv'] and hparams['pitch_type'] == 'frame': assert p_pred[..., 1].shape == uv.shape, (p_pred.shape, uv.shape) losses[f'uv{postfix}'] = (F.binary_cross_entropy_with_logits( p_pred[:, :, 1], uv, reduction='none') * nonpadding).sum() \ / nonpadding.sum() * hparams['lambda_uv'] nonpadding = nonpadding * (uv == 0).float() f0_pred = p_pred[:, :, 0] pitch_loss_fn = F.l1_loss if hparams['pitch_loss'] == 'l1' else F.mse_loss losses[f'f0{postfix}'] = (pitch_loss_fn(f0_pred, f0, reduction='none') * nonpadding).sum() \ / nonpadding.sum() * hparams['lambda_f0'] ############ # validation plots ############ def plot_dur(self, batch_idx, sample, model_out): T_txt = sample['txt_tokens'].shape[1] dur_gt = mel2ph_to_dur(sample['mel2ph'], T_txt)[0] dur_pred = model_out['dur'] if hasattr(self.model, 'out2dur'): dur_pred = self.model.out2dur(model_out['dur']).float() txt = self.phone_encoder.decode(sample['txt_tokens'][0].cpu().numpy()) txt = txt.split(" ") self.logger.add_figure( f'dur_{batch_idx}', dur_to_figure(dur_gt, dur_pred, txt), self.global_step) def plot_pitch(self, batch_idx, sample, model_out): self.logger.add_figure( f'f0_{batch_idx}', f0_to_figure(model_out['f0_denorm'][0], None, model_out['f0_denorm_pred'][0]), self.global_step) ############ # inference ############ def test_step(self, sample, batch_idx): spk_embed = sample.get('spk_embed') if not hparams['use_spk_id'] else sample.get('spk_ids') txt_tokens = sample['txt_tokens'] mel2ph, uv, f0 = None, None, None ref_mels = sample['mels'] if hparams['use_gt_dur']: mel2ph = sample['mel2ph'] if hparams['use_gt_f0']: f0 = sample['f0'] uv = sample['uv'] run_model = lambda: self.model( txt_tokens, spk_embed=spk_embed, mel2ph=mel2ph, f0=f0, uv=uv, ref_mels=ref_mels, infer=True) if hparams['profile_infer']: mel2ph, uv, f0 = sample['mel2ph'], sample['uv'], sample['f0'] with utils.Timer('fs', enable=True): outputs = run_model() if 'gen_wav_time' not in self.stats: self.stats['gen_wav_time'] = 0 wav_time = float(outputs["mels_out"].shape[1]) * hparams['hop_size'] / hparams["audio_sample_rate"] self.stats['gen_wav_time'] += wav_time print(f'[Timer] wav total seconds: {self.stats["gen_wav_time"]}') from pytorch_memlab import LineProfiler with LineProfiler(self.model.forward) as prof: run_model() prof.print_stats() else: outputs = run_model() sample['outputs'] = self.model.out2mel(outputs['mel_out']) sample['mel2ph_pred'] = outputs['mel2ph'] if hparams['use_pitch_embed']: sample['f0'] = denorm_f0(sample['f0'], sample['uv'], hparams) if hparams['pitch_type'] == 'ph': sample['f0'] = torch.gather(F.pad(sample['f0'], [1, 0]), 1, sample['mel2ph']) sample['f0_pred'] = outputs.get('f0_denorm') return self.after_infer(sample)