(self, model, sample, return_output=False, infer=False)
| 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 losses |
| 53 | else: |
| 54 | return losses, output |
| 55 | |
| 56 | def _training_step(self, sample, batch_idx, _): |
| 57 | log_outputs = self.run_model(self.model, sample) |
no test coverage detected