justyoung/DiffSinger

1import torch2 3import utils4from .diff.diffusion import GaussianDiffusion5from .diff.net import DiffNet6from tasks.tts.fs2 import FastSpeech2Task7from utils.hparams import hparams8 9 10DIFF_DECODERS = {11    'wavenet': lambda hp: DiffNet(hp['audio_num_mel_bins']),12}13 14 15class DiffFsTask(FastSpeech2Task):16    def build_tts_model(self):17        mel_bins = hparams['audio_num_mel_bins']18        self.model = GaussianDiffusion(19            phone_encoder=self.phone_encoder,20            out_dims=mel_bins, denoise_fn=DIFF_DECODERS[hparams['diff_decoder_type']](hparams),21            timesteps=hparams['timesteps'],22            loss_type=hparams['diff_loss_type'],23            spec_min=hparams['spec_min'], spec_max=hparams['spec_max'],24        )25 26    def run_model(self, model, sample, return_output=False, infer=False):27        txt_tokens = sample['txt_tokens']  # [B, T_t]28        target = sample['mels']  # [B, T_s, 80]29        mel2ph = sample['mel2ph']  # [B, T_s]30        f0 = sample['f0']31        uv = sample['uv']32        energy = sample['energy']33        spk_embed = sample.get('spk_embed') if not hparams['use_spk_id'] else sample.get('spk_ids')34        if hparams['pitch_type'] == 'cwt':35            cwt_spec = sample[f'cwt_spec']36            f0_mean = sample['f0_mean']37            f0_std = sample['f0_std']38            sample['f0_cwt'] = f0 = model.cwt2f0_norm(cwt_spec, f0_mean, f0_std, mel2ph)39 40        output = model(txt_tokens, mel2ph=mel2ph, spk_embed=spk_embed,41                       ref_mels=target, f0=f0, uv=uv, energy=energy, infer=infer)42 43        losses = {}44        if 'diff_loss' in output:45            losses['mel'] = output['diff_loss']46        self.add_dur_loss(output['dur'], mel2ph, txt_tokens, losses=losses)47        if hparams['use_pitch_embed']:48            self.add_pitch_loss(output, sample, losses)49        if hparams['use_energy_embed']:50            self.add_energy_loss(output['energy_pred'], energy, losses)51        if not return_output:52            return losses53        else:54            return losses, output55 56    def _training_step(self, sample, batch_idx, _):57        log_outputs = self.run_model(self.model, sample)58        total_loss = sum([v for v in log_outputs.values() if isinstance(v, torch.Tensor) and v.requires_grad])59        log_outputs['batch_size'] = sample['txt_tokens'].size()[0]60        log_outputs['lr'] = self.scheduler.get_lr()[0]61        return total_loss, log_outputs62 63    def validation_step(self, sample, batch_idx):64        outputs = {}65        outputs['losses'] = {}66        outputs['losses'], model_out = self.run_model(self.model, sample, return_output=True, infer=False)67        outputs['total_loss'] = sum(outputs['losses'].values())68        outputs['nsamples'] = sample['nsamples']69        outputs = utils.tensors_to_scalars(outputs)70        if batch_idx < hparams['num_valid_plots']:71            _, model_out = self.run_model(self.model, sample, return_output=True, infer=True)72            self.plot_mel(batch_idx, sample['mels'], model_out['mel_out'])73        return outputs74 75    def build_scheduler(self, optimizer):76        return torch.optim.lr_scheduler.StepLR(optimizer, hparams['decay_steps'], gamma=0.5)77 78    def optimizer_step(self, epoch, batch_idx, optimizer, optimizer_idx):79        if optimizer is None:80            return81        optimizer.step()82        optimizer.zero_grad()83        if self.scheduler is not None:84            self.scheduler.step(self.global_step // hparams['accumulate_grad_batches'])85