Spaces:
Runtime error
Runtime error
File size: 5,444 Bytes
53fa903 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 |
import torch
from modules.tts.diffspeech.shallow_diffusion_tts import GaussianDiffusion
from tasks.tts.fs2_orig import FastSpeech2OrigTask
import utils
from utils.commons.hparams import hparams
from utils.commons.ckpt_utils import load_ckpt
from utils.audio.pitch.utils import denorm_f0
class DiffSpeechTask(FastSpeech2OrigTask):
def build_tts_model(self):
# get min and max
# import torch
# from tqdm import tqdm
# v_min = torch.ones([80]) * 100
# v_max = torch.ones([80]) * -100
# for i, ds in enumerate(tqdm(self.dataset_cls('train'))):
# v_max = torch.max(torch.max(ds['mel'].reshape(-1, 80), 0)[0], v_max)
# v_min = torch.min(torch.min(ds['mel'].reshape(-1, 80), 0)[0], v_min)
# if i % 100 == 0:
# print(i, v_min, v_max)
# print('final', v_min, v_max)
dict_size = len(self.token_encoder)
self.model = GaussianDiffusion(dict_size, hparams)
if hparams['fs2_ckpt'] != '':
load_ckpt(self.model.fs2, hparams['fs2_ckpt'], 'model', strict=True)
for k, v in self.model.fs2.named_parameters():
if 'predictor' not in k:
v.requires_grad = False
# or
# for k, v in self.model.fs2.named_parameters():
# v.requires_grad = False
def build_optimizer(self, model):
self.optimizer = optimizer = torch.optim.AdamW(
filter(lambda p: p.requires_grad, model.parameters()),
lr=hparams['lr'],
betas=(hparams['optimizer_adam_beta1'], hparams['optimizer_adam_beta2']),
weight_decay=hparams['weight_decay'])
return optimizer
def build_scheduler(self, optimizer):
return torch.optim.lr_scheduler.StepLR(optimizer, hparams['decay_steps'], gamma=0.5)
def run_model(self, sample, infer=False, *args, **kwargs):
txt_tokens = sample['txt_tokens'] # [B, T_t]
spk_embed = sample.get('spk_embed')
spk_id = sample.get('spk_ids')
if not infer:
target = sample['mels'] # [B, T_s, 80]
mel2ph = sample['mel2ph'] # [B, T_s]
f0 = sample.get('f0')
uv = sample.get('uv')
output = self.model(txt_tokens, mel2ph=mel2ph, spk_embed=spk_embed, spk_id=spk_id,
ref_mels=target, f0=f0, uv=uv, infer=False)
losses = {}
if 'diff_loss' in output:
losses['mel'] = output['diff_loss']
self.add_dur_loss(output['dur'], mel2ph, txt_tokens, losses=losses)
if hparams['use_pitch_embed']:
self.add_pitch_loss(output, sample, losses)
return losses, output
else:
use_gt_dur = kwargs.get('infer_use_gt_dur', hparams['use_gt_dur'])
use_gt_f0 = kwargs.get('infer_use_gt_f0', hparams['use_gt_f0'])
mel2ph, uv, f0 = None, None, None
if use_gt_dur:
mel2ph = sample['mel2ph']
if use_gt_f0:
f0 = sample['f0']
uv = sample['uv']
output = self.model(txt_tokens, mel2ph=mel2ph, spk_embed=spk_embed, spk_id=spk_id,
ref_mels=None, f0=f0, uv=uv, infer=True)
return output
def save_valid_result(self, sample, batch_idx, model_out):
sr = hparams['audio_sample_rate']
f0_gt = None
# mel_out = model_out['mel_out']
if sample.get('f0') is not None:
f0_gt = denorm_f0(sample['f0'][0].cpu(), sample['uv'][0].cpu())
# self.plot_mel(batch_idx, sample['mels'], mel_out, f0s=f0_gt)
if self.global_step > 0:
# wav_pred = self.vocoder.spec2wav(mel_out[0].cpu(), f0=f0_gt)
# self.logger.add_audio(f'wav_val_{batch_idx}', wav_pred, self.global_step, sr)
# 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']
wav_pred = self.vocoder.spec2wav(model_out['mel_out'][0].cpu(), f0=f0_gt)
self.logger.add_audio(f'wav_gdur_{batch_idx}', wav_pred, self.global_step, sr)
self.plot_mel(batch_idx, sample['mels'], model_out['mel_out'][0], f'diffmel_gdur_{batch_idx}',
dur_info=dur_info, f0s=f0_gt)
self.plot_mel(batch_idx, sample['mels'], model_out['fs2_mel'][0], f'fs2mel_gdur_{batch_idx}',
dur_info=dur_info, f0s=f0_gt) # gt mel vs. fs2 mel
# with pred duration
if not hparams['use_gt_dur']:
model_out = self.run_model(sample, infer=True, infer_use_gt_dur=False)
dur_info = self.get_plot_dur_info(sample, model_out)
self.plot_mel(batch_idx, sample['mels'], model_out['mel_out'][0], f'mel_pdur_{batch_idx}',
dur_info=dur_info, f0s=f0_gt)
wav_pred = self.vocoder.spec2wav(model_out['mel_out'][0].cpu(), f0=f0_gt)
self.logger.add_audio(f'wav_pdur_{batch_idx}', wav_pred, self.global_step, sr)
# gt wav
if self.global_step <= hparams['valid_infer_interval']:
mel_gt = sample['mels'][0].cpu()
wav_gt = self.vocoder.spec2wav(mel_gt, f0=f0_gt)
self.logger.add_audio(f'wav_gt_{batch_idx}', wav_gt, self.global_step, sr)
|