Silentlin/DiffSinger
89
1import torch2 3import utils4from utils.hparams import hparams5from .diff.net import DiffNet6from .diff.shallow_diffusion_tts import GaussianDiffusion, OfflineGaussianDiffusion7from .diffspeech_task import DiffSpeechTask8from vocoders.base_vocoder import get_vocoder_cls, BaseVocoder9from modules.fastspeech.pe import PitchExtractor10from modules.fastspeech.fs2 import FastSpeech211from modules.diffsinger_midi.fs2 import FastSpeech2MIDI12from modules.fastspeech.tts_modules import mel2ph_to_dur13 14from usr.diff.candidate_decoder import FFT15from utils.pitch_utils import denorm_f016from tasks.tts.fs2_utils import FastSpeechDataset17from tasks.tts.fs2 import FastSpeech2Task18 19import numpy as np20import os21import torch.nn.functional as F22 23DIFF_DECODERS = {24 'wavenet': lambda hp: DiffNet(hp['audio_num_mel_bins']),25 'fft': lambda hp: FFT(26 hp['hidden_size'], hp['dec_layers'], hp['dec_ffn_kernel_size'], hp['num_heads']),27}28 29 30class DiffSingerTask(DiffSpeechTask):31 def __init__(self):32 super(DiffSingerTask, self).__init__()33 self.dataset_cls = FastSpeechDataset34 self.vocoder: BaseVocoder = get_vocoder_cls(hparams)()35 if hparams.get('pe_enable') is not None and hparams['pe_enable']:36 self.pe = PitchExtractor().cuda()37 utils.load_ckpt(self.pe, hparams['pe_ckpt'], 'model', strict=True)38 self.pe.eval()39 40 def build_tts_model(self):41 # import torch42 # from tqdm import tqdm43 # v_min = torch.ones([80]) * 10044 # v_max = torch.ones([80]) * -10045 # for i, ds in enumerate(tqdm(self.dataset_cls('train'))):46 # v_max = torch.max(torch.max(ds['mel'].reshape(-1, 80), 0)[0], v_max)47 # v_min = torch.min(torch.min(ds['mel'].reshape(-1, 80), 0)[0], v_min)48 # if i % 100 == 0:49 # print(i, v_min, v_max)50 # print('final', v_min, v_max)51 mel_bins = hparams['audio_num_mel_bins']52 self.model = GaussianDiffusion(53 phone_encoder=self.phone_encoder,54 out_dims=mel_bins, denoise_fn=DIFF_DECODERS[hparams['diff_decoder_type']](hparams),55 timesteps=hparams['timesteps'],56 K_step=hparams['K_step'],57 loss_type=hparams['diff_loss_type'],58 spec_min=hparams['spec_min'], spec_max=hparams['spec_max'],59 )60 if hparams['fs2_ckpt'] != '':61 utils.load_ckpt(self.model.fs2, hparams['fs2_ckpt'], 'model', strict=True)62 # self.model.fs2.decoder = None63 for k, v in self.model.fs2.named_parameters():64 v.requires_grad = False65 66 def validation_step(self, sample, batch_idx):67 outputs = {}68 txt_tokens = sample['txt_tokens'] # [B, T_t]69 70 target = sample['mels'] # [B, T_s, 80]71 energy = sample['energy']72 # fs2_mel = sample['fs2_mels']73 spk_embed = sample.get('spk_embed') if not hparams['use_spk_id'] else sample.get('spk_ids')74 mel2ph = sample['mel2ph']75 f0 = sample['f0']76 uv = sample['uv']77 78 outputs['losses'] = {}79 80 outputs['losses'], model_out = self.run_model(self.model, sample, return_output=True, infer=False)81 82 83 outputs['total_loss'] = sum(outputs['losses'].values())84 outputs['nsamples'] = sample['nsamples']85 outputs = utils.tensors_to_scalars(outputs)86 if batch_idx < hparams['num_valid_plots']:87 model_out = self.model(88 txt_tokens, spk_embed=spk_embed, mel2ph=mel2ph, f0=f0, uv=uv, energy=energy, ref_mels=None, infer=True)89 90 if hparams.get('pe_enable') is not None and hparams['pe_enable']:91 gt_f0 = self.pe(sample['mels'])['f0_denorm_pred'] # pe predict from GT mel92 pred_f0 = self.pe(model_out['mel_out'])['f0_denorm_pred'] # pe predict from Pred mel93 else:94 gt_f0 = denorm_f0(sample['f0'], sample['uv'], hparams)95 pred_f0 = model_out.get('f0_denorm')96 self.plot_wav(batch_idx, sample['mels'], model_out['mel_out'], is_mel=True, gt_f0=gt_f0, f0=pred_f0)97 self.plot_mel(batch_idx, sample['mels'], model_out['mel_out'], name=f'diffmel_{batch_idx}')98 self.plot_mel(batch_idx, sample['mels'], model_out['fs2_mel'], name=f'fs2mel_{batch_idx}')99 return outputs100 101 102class ShallowDiffusionOfflineDataset(FastSpeechDataset):103 def __getitem__(self, index):104 sample = super(ShallowDiffusionOfflineDataset, self).__getitem__(index)105 item = self._get_item(index)106 107 if self.prefix != 'train' and hparams['fs2_ckpt'] != '':108 fs2_ckpt = os.path.dirname(hparams['fs2_ckpt'])109 item_name = item['item_name']110 fs2_mel = torch.Tensor(np.load(f'{fs2_ckpt}/P_mels_npy/{item_name}.npy')) # ~M generated by FFT-singer.111 sample['fs2_mel'] = fs2_mel112 return sample113 114 def collater(self, samples):115 batch = super(ShallowDiffusionOfflineDataset, self).collater(samples)116 if self.prefix != 'train' and hparams['fs2_ckpt'] != '':117 batch['fs2_mels'] = utils.collate_2d([s['fs2_mel'] for s in samples], 0.0)118 return batch119 120 121class DiffSingerOfflineTask(DiffSingerTask):122 def __init__(self):123 super(DiffSingerOfflineTask, self).__init__()124 self.dataset_cls = ShallowDiffusionOfflineDataset125 126 def build_tts_model(self):127 mel_bins = hparams['audio_num_mel_bins']128 self.model = OfflineGaussianDiffusion(129 phone_encoder=self.phone_encoder,130 out_dims=mel_bins, denoise_fn=DIFF_DECODERS[hparams['diff_decoder_type']](hparams),131 timesteps=hparams['timesteps'],132 K_step=hparams['K_step'],133 loss_type=hparams['diff_loss_type'],134 spec_min=hparams['spec_min'], spec_max=hparams['spec_max'],135 )136 # if hparams['fs2_ckpt'] != '':137 # utils.load_ckpt(self.model.fs2, hparams['fs2_ckpt'], 'model', strict=True)138 # self.model.fs2.decoder = None139 140 def run_model(self, model, sample, return_output=False, infer=False):141 txt_tokens = sample['txt_tokens'] # [B, T_t]142 target = sample['mels'] # [B, T_s, 80]143 mel2ph = sample['mel2ph'] # [B, T_s]144 f0 = sample['f0']145 uv = sample['uv']146 energy = sample['energy']147 fs2_mel = None #sample['fs2_mels']148 spk_embed = sample.get('spk_embed') if not hparams['use_spk_id'] else sample.get('spk_ids')149 if hparams['pitch_type'] == 'cwt':150 cwt_spec = sample[f'cwt_spec']151 f0_mean = sample['f0_mean']152 f0_std = sample['f0_std']153 sample['f0_cwt'] = f0 = model.cwt2f0_norm(cwt_spec, f0_mean, f0_std, mel2ph)154 155 output = model(txt_tokens, mel2ph=mel2ph, spk_embed=spk_embed,156 ref_mels=[target, fs2_mel], f0=f0, uv=uv, energy=energy, infer=infer)157 158 losses = {}159 if 'diff_loss' in output:160 losses['mel'] = output['diff_loss']161 # self.add_dur_loss(output['dur'], mel2ph, txt_tokens, losses=losses)162 # if hparams['use_pitch_embed']:163 # self.add_pitch_loss(output, sample, losses)164 if hparams['use_energy_embed']:165 self.add_energy_loss(output['energy_pred'], energy, losses)166 167 if not return_output:168 return losses169 else:170 return losses, output171 172 def validation_step(self, sample, batch_idx):173 outputs = {}174 txt_tokens = sample['txt_tokens'] # [B, T_t]175 176 target = sample['mels'] # [B, T_s, 80]177 energy = sample['energy']178 # fs2_mel = sample['fs2_mels']179 spk_embed = sample.get('spk_embed') if not hparams['use_spk_id'] else sample.get('spk_ids')180 mel2ph = sample['mel2ph']181 f0 = sample['f0']182 uv = sample['uv']183 184 outputs['losses'] = {}185 186 outputs['losses'], model_out = self.run_model(self.model, sample, return_output=True, infer=False)187 188 189 outputs['total_loss'] = sum(outputs['losses'].values())190 outputs['nsamples'] = sample['nsamples']191 outputs = utils.tensors_to_scalars(outputs)192 if batch_idx < hparams['num_valid_plots']:193 fs2_mel = sample['fs2_mels']194 model_out = self.model(195 txt_tokens, spk_embed=spk_embed, mel2ph=mel2ph, f0=f0, uv=uv, energy=energy,196 ref_mels=[None, fs2_mel], infer=True)197 if hparams.get('pe_enable') is not None and hparams['pe_enable']:198 gt_f0 = self.pe(sample['mels'])['f0_denorm_pred'] # pe predict from GT mel199 pred_f0 = self.pe(model_out['mel_out'])['f0_denorm_pred'] # pe predict from Pred mel200 else:201 gt_f0 = denorm_f0(sample['f0'], sample['uv'], hparams)202 pred_f0 = model_out.get('f0_denorm')203 self.plot_wav(batch_idx, sample['mels'], model_out['mel_out'], is_mel=True, gt_f0=gt_f0, f0=pred_f0)204 self.plot_mel(batch_idx, sample['mels'], model_out['mel_out'], name=f'diffmel_{batch_idx}')205 self.plot_mel(batch_idx, sample['mels'], fs2_mel, name=f'fs2mel_{batch_idx}')206 return outputs207 208 def test_step(self, sample, batch_idx):209 spk_embed = sample.get('spk_embed') if not hparams['use_spk_id'] else sample.get('spk_ids')210 txt_tokens = sample['txt_tokens']211 energy = sample['energy']212 if hparams['profile_infer']:213 pass214 else:215 mel2ph, uv, f0 = None, None, None216 if hparams['use_gt_dur']:217 mel2ph = sample['mel2ph']218 if hparams['use_gt_f0']:219 f0 = sample['f0']220 uv = sample['uv']221 fs2_mel = sample['fs2_mels']222 outputs = self.model(223 txt_tokens, spk_embed=spk_embed, mel2ph=mel2ph, f0=f0, uv=uv, ref_mels=[None, fs2_mel], energy=energy,224 infer=True)225 sample['outputs'] = self.model.out2mel(outputs['mel_out'])226 sample['mel2ph_pred'] = outputs['mel2ph']227 228 if hparams.get('pe_enable') is not None and hparams['pe_enable']:229 sample['f0'] = self.pe(sample['mels'])['f0_denorm_pred'] # pe predict from GT mel230 sample['f0_pred'] = self.pe(sample['outputs'])['f0_denorm_pred'] # pe predict from Pred mel231 else:232 sample['f0'] = denorm_f0(sample['f0'], sample['uv'], hparams)233 sample['f0_pred'] = outputs.get('f0_denorm')234 return self.after_infer(sample)235 236 237class MIDIDataset(FastSpeechDataset):238 def __getitem__(self, index):239 sample = super(MIDIDataset, self).__getitem__(index)240 item = self._get_item(index)241 sample['f0_midi'] = torch.FloatTensor(item['f0_midi'])242 sample['pitch_midi'] = torch.LongTensor(item['pitch_midi'])[:hparams['max_frames']]243 244 return sample245 246 def collater(self, samples):247 batch = super(MIDIDataset, self).collater(samples)248 batch['f0_midi'] = utils.collate_1d([s['f0_midi'] for s in samples], 0.0)249 batch['pitch_midi'] = utils.collate_1d([s['pitch_midi'] for s in samples], 0)250 # print((batch['pitch_midi'] == f0_to_coarse(batch['f0_midi'])).all())251 return batch252 253 254class OpencpopDataset(FastSpeechDataset):255 def __getitem__(self, index):256 sample = super(OpencpopDataset, self).__getitem__(index)257 item = self._get_item(index)258 sample['pitch_midi'] = torch.LongTensor(item['pitch_midi'])[:hparams['max_frames']]259 sample['midi_dur'] = torch.FloatTensor(item['midi_dur'])[:hparams['max_frames']]260 sample['is_slur'] = torch.LongTensor(item['is_slur'])[:hparams['max_frames']]261 sample['word_boundary'] = torch.LongTensor(item['word_boundary'])[:hparams['max_frames']]262 return sample263 264 def collater(self, samples):265 batch = super(OpencpopDataset, self).collater(samples)266 batch['pitch_midi'] = utils.collate_1d([s['pitch_midi'] for s in samples], 0)267 batch['midi_dur'] = utils.collate_1d([s['midi_dur'] for s in samples], 0)268 batch['is_slur'] = utils.collate_1d([s['is_slur'] for s in samples], 0)269 batch['word_boundary'] = utils.collate_1d([s['word_boundary'] for s in samples], 0)270 return batch271 272 273class DiffSingerMIDITask(DiffSingerTask):274 def __init__(self):275 super(DiffSingerMIDITask, self).__init__()276 # self.dataset_cls = MIDIDataset277 self.dataset_cls = OpencpopDataset278 279 def run_model(self, model, sample, return_output=False, infer=False):280 txt_tokens = sample['txt_tokens'] # [B, T_t]281 target = sample['mels'] # [B, T_s, 80]282 # mel2ph = sample['mel2ph'] if hparams['use_gt_dur'] else None # [B, T_s]283 mel2ph = sample['mel2ph']284 if hparams.get('switch_midi2f0_step') is not None and self.global_step > hparams['switch_midi2f0_step']:285 f0 = None286 uv = None287 else:288 f0 = sample['f0']289 uv = sample['uv']290 energy = sample['energy']291 292 spk_embed = sample.get('spk_embed') if not hparams['use_spk_id'] else sample.get('spk_ids')293 if hparams['pitch_type'] == 'cwt':294 cwt_spec = sample[f'cwt_spec']295 f0_mean = sample['f0_mean']296 f0_std = sample['f0_std']297 sample['f0_cwt'] = f0 = model.cwt2f0_norm(cwt_spec, f0_mean, f0_std, mel2ph)298 299 output = model(txt_tokens, mel2ph=mel2ph, spk_embed=spk_embed,300 ref_mels=target, f0=f0, uv=uv, energy=energy, infer=infer, pitch_midi=sample['pitch_midi'],301 midi_dur=sample.get('midi_dur'), is_slur=sample.get('is_slur'))302 303 losses = {}304 if 'diff_loss' in output:305 losses['mel'] = output['diff_loss']306 self.add_dur_loss(output['dur'], mel2ph, txt_tokens, sample['word_boundary'], losses=losses)307 if hparams['use_pitch_embed']:308 self.add_pitch_loss(output, sample, losses)309 if hparams['use_energy_embed']:310 self.add_energy_loss(output['energy_pred'], energy, losses)311 if not return_output:312 return losses313 else:314 return losses, output315 316 def validation_step(self, sample, batch_idx):317 outputs = {}318 txt_tokens = sample['txt_tokens'] # [B, T_t]319 320 target = sample['mels'] # [B, T_s, 80]321 energy = sample['energy']322 # fs2_mel = sample['fs2_mels']323 spk_embed = sample.get('spk_embed') if not hparams['use_spk_id'] else sample.get('spk_ids')324 mel2ph = sample['mel2ph']325 326 outputs['losses'] = {}327 328 outputs['losses'], model_out = self.run_model(self.model, sample, return_output=True, infer=False)329 330 outputs['total_loss'] = sum(outputs['losses'].values())331 outputs['nsamples'] = sample['nsamples']332 outputs = utils.tensors_to_scalars(outputs)333 if batch_idx < hparams['num_valid_plots']:334 model_out = self.model(335 txt_tokens, spk_embed=spk_embed, mel2ph=mel2ph, f0=None, uv=None, energy=energy, ref_mels=None, infer=True,336 pitch_midi=sample['pitch_midi'], midi_dur=sample.get('midi_dur'), is_slur=sample.get('is_slur'))337 338 if hparams.get('pe_enable') is not None and hparams['pe_enable']:339 gt_f0 = self.pe(sample['mels'])['f0_denorm_pred'] # pe predict from GT mel340 pred_f0 = self.pe(model_out['mel_out'])['f0_denorm_pred'] # pe predict from Pred mel341 else:342 gt_f0 = denorm_f0(sample['f0'], sample['uv'], hparams)343 pred_f0 = model_out.get('f0_denorm')344 self.plot_wav(batch_idx, sample['mels'], model_out['mel_out'], is_mel=True, gt_f0=gt_f0, f0=pred_f0)345 self.plot_mel(batch_idx, sample['mels'], model_out['mel_out'], name=f'diffmel_{batch_idx}')346 self.plot_mel(batch_idx, sample['mels'], model_out['fs2_mel'], name=f'fs2mel_{batch_idx}')347 if hparams['use_pitch_embed']:348 self.plot_pitch(batch_idx, sample, model_out)349 return outputs350 351 def add_dur_loss(self, dur_pred, mel2ph, txt_tokens, wdb, losses=None):352 """353 :param dur_pred: [B, T], float, log scale354 :param mel2ph: [B, T]355 :param txt_tokens: [B, T]356 :param losses:357 :return:358 """359 B, T = txt_tokens.shape360 nonpadding = (txt_tokens != 0).float()361 dur_gt = mel2ph_to_dur(mel2ph, T).float() * nonpadding362 is_sil = torch.zeros_like(txt_tokens).bool()363 for p in self.sil_ph:364 is_sil = is_sil | (txt_tokens == self.phone_encoder.encode(p)[0])365 is_sil = is_sil.float() # [B, T_txt]366 367 # phone duration loss368 if hparams['dur_loss'] == 'mse':369 losses['pdur'] = F.mse_loss(dur_pred, (dur_gt + 1).log(), reduction='none')370 losses['pdur'] = (losses['pdur'] * nonpadding).sum() / nonpadding.sum()371 dur_pred = (dur_pred.exp() - 1).clamp(min=0)372 else:373 raise NotImplementedError374 375 # use linear scale for sent and word duration376 if hparams['lambda_word_dur'] > 0:377 idx = F.pad(wdb.cumsum(axis=1), (1, 0))[:, :-1]378 # word_dur_g = dur_gt.new_zeros([B, idx.max() + 1]).scatter_(1, idx, midi_dur) # midi_dur can be implied by add gt-ph_dur379 word_dur_p = dur_pred.new_zeros([B, idx.max() + 1]).scatter_add(1, idx, dur_pred)380 word_dur_g = dur_gt.new_zeros([B, idx.max() + 1]).scatter_add(1, idx, dur_gt)381 wdur_loss = F.mse_loss((word_dur_p + 1).log(), (word_dur_g + 1).log(), reduction='none')382 word_nonpadding = (word_dur_g > 0).float()383 wdur_loss = (wdur_loss * word_nonpadding).sum() / word_nonpadding.sum()384 losses['wdur'] = wdur_loss * hparams['lambda_word_dur']385 if hparams['lambda_sent_dur'] > 0:386 sent_dur_p = dur_pred.sum(-1)387 sent_dur_g = dur_gt.sum(-1)388 sdur_loss = F.mse_loss((sent_dur_p + 1).log(), (sent_dur_g + 1).log(), reduction='mean')389 losses['sdur'] = sdur_loss.mean() * hparams['lambda_sent_dur']390 391 392class AuxDecoderMIDITask(FastSpeech2Task):393 def __init__(self):394 super().__init__()395 # self.dataset_cls = MIDIDataset396 self.dataset_cls = OpencpopDataset397 398 def build_tts_model(self):399 if hparams.get('use_midi') is not None and hparams['use_midi']:400 self.model = FastSpeech2MIDI(self.phone_encoder)401 else:402 self.model = FastSpeech2(self.phone_encoder)403 404 def run_model(self, model, sample, return_output=False):405 txt_tokens = sample['txt_tokens'] # [B, T_t]406 target = sample['mels'] # [B, T_s, 80]407 mel2ph = sample['mel2ph'] # [B, T_s]408 f0 = sample['f0']409 uv = sample['uv']410 energy = sample['energy']411 412 spk_embed = sample.get('spk_embed') if not hparams['use_spk_id'] else sample.get('spk_ids')413 if hparams['pitch_type'] == 'cwt':414 cwt_spec = sample[f'cwt_spec']415 f0_mean = sample['f0_mean']416 f0_std = sample['f0_std']417 sample['f0_cwt'] = f0 = model.cwt2f0_norm(cwt_spec, f0_mean, f0_std, mel2ph)418 419 output = model(txt_tokens, mel2ph=mel2ph, spk_embed=spk_embed,420 ref_mels=target, f0=f0, uv=uv, energy=energy, infer=False, pitch_midi=sample['pitch_midi'],421 midi_dur=sample.get('midi_dur'), is_slur=sample.get('is_slur'))422 423 losses = {}424 self.add_mel_loss(output['mel_out'], target, losses)425 self.add_dur_loss(output['dur'], mel2ph, txt_tokens, sample['word_boundary'], losses=losses)426 if hparams['use_pitch_embed']:427 self.add_pitch_loss(output, sample, losses)428 if hparams['use_energy_embed']:429 self.add_energy_loss(output['energy_pred'], energy, losses)430 if not return_output:431 return losses432 else:433 return losses, output434 435 def add_dur_loss(self, dur_pred, mel2ph, txt_tokens, wdb, losses=None):436 """437 :param dur_pred: [B, T], float, log scale438 :param mel2ph: [B, T]439 :param txt_tokens: [B, T]440 :param losses:441 :return:442 """443 B, T = txt_tokens.shape444 nonpadding = (txt_tokens != 0).float()445 dur_gt = mel2ph_to_dur(mel2ph, T).float() * nonpadding446 is_sil = torch.zeros_like(txt_tokens).bool()447 for p in self.sil_ph:448 is_sil = is_sil | (txt_tokens == self.phone_encoder.encode(p)[0])449 is_sil = is_sil.float() # [B, T_txt]450 451 # phone duration loss452 if hparams['dur_loss'] == 'mse':453 losses['pdur'] = F.mse_loss(dur_pred, (dur_gt + 1).log(), reduction='none')454 losses['pdur'] = (losses['pdur'] * nonpadding).sum() / nonpadding.sum()455 dur_pred = (dur_pred.exp() - 1).clamp(min=0)456 else:457 raise NotImplementedError458 459 # use linear scale for sent and word duration460 if hparams['lambda_word_dur'] > 0:461 idx = F.pad(wdb.cumsum(axis=1), (1, 0))[:, :-1]462 # word_dur_g = dur_gt.new_zeros([B, idx.max() + 1]).scatter_(1, idx, midi_dur) # midi_dur can be implied by add gt-ph_dur463 word_dur_p = dur_pred.new_zeros([B, idx.max() + 1]).scatter_add(1, idx, dur_pred)464 word_dur_g = dur_gt.new_zeros([B, idx.max() + 1]).scatter_add(1, idx, dur_gt)465 wdur_loss = F.mse_loss((word_dur_p + 1).log(), (word_dur_g + 1).log(), reduction='none')466 word_nonpadding = (word_dur_g > 0).float()467 wdur_loss = (wdur_loss * word_nonpadding).sum() / word_nonpadding.sum()468 losses['wdur'] = wdur_loss * hparams['lambda_word_dur']469 if hparams['lambda_sent_dur'] > 0:470 sent_dur_p = dur_pred.sum(-1)471 sent_dur_g = dur_gt.sum(-1)472 sdur_loss = F.mse_loss((sent_dur_p + 1).log(), (sent_dur_g + 1).log(), reduction='mean')473 losses['sdur'] = sdur_loss.mean() * hparams['lambda_sent_dur']474 475 def validation_step(self, sample, batch_idx):476 outputs = {}477 outputs['losses'] = {}478 outputs['losses'], model_out = self.run_model(self.model, sample, return_output=True)479 outputs['total_loss'] = sum(outputs['losses'].values())480 outputs['nsamples'] = sample['nsamples']481 mel_out = self.model.out2mel(model_out['mel_out'])482 outputs = utils.tensors_to_scalars(outputs)483 # if sample['mels'].shape[0] == 1:484 # self.add_laplace_var(mel_out, sample['mels'], outputs)485 if batch_idx < hparams['num_valid_plots']:486 self.plot_mel(batch_idx, sample['mels'], mel_out)487 self.plot_dur(batch_idx, sample, model_out)488 if hparams['use_pitch_embed']:489 self.plot_pitch(batch_idx, sample, model_out)490 return outputs