wehub-resource-sync
243 lines
10 KiB
Python
243 lines
10 KiB
Python
import matplotlib
|
|
import torch
|
|
import torch.distributions
|
|
import torch.optim
|
|
import torch.utils.data
|
|
|
|
import utils
|
|
import utils.infer_utils
|
|
from basics.base_dataset import BaseDataset
|
|
from basics.base_task import BaseTask
|
|
from basics.base_vocoder import BaseVocoder
|
|
from modules.aux_decoder import build_aux_loss
|
|
from modules.losses import DiffusionLoss, RectifiedFlowLoss
|
|
from modules.toplevel import DiffSingerAcoustic, ShallowDiffusionOutput
|
|
from modules.vocoders.registry import get_vocoder_cls
|
|
from utils.hparams import hparams
|
|
from utils.plot import spec_to_figure
|
|
|
|
matplotlib.use('Agg')
|
|
|
|
|
|
class AcousticDataset(BaseDataset):
|
|
def __init__(self, prefix, preload=False):
|
|
super(AcousticDataset, self).__init__(prefix, hparams['dataset_size_key'], preload)
|
|
self.required_variances = {} # key: variance name, value: padding value
|
|
if hparams['use_energy_embed']:
|
|
self.required_variances['energy'] = 0.0
|
|
if hparams['use_breathiness_embed']:
|
|
self.required_variances['breathiness'] = 0.0
|
|
if hparams['use_voicing_embed']:
|
|
self.required_variances['voicing'] = 0.0
|
|
if hparams['use_tension_embed']:
|
|
self.required_variances['tension'] = 0.0
|
|
|
|
self.need_key_shift = hparams['use_key_shift_embed']
|
|
self.need_speed = hparams['use_speed_embed']
|
|
self.need_spk_id = hparams['use_spk_id']
|
|
self.need_lang_id = hparams['use_lang_id']
|
|
|
|
def collater(self, samples):
|
|
batch = super().collater(samples)
|
|
if batch['size'] == 0:
|
|
return batch
|
|
|
|
tokens = utils.collate_nd([s['tokens'] for s in samples], 0)
|
|
f0 = utils.collate_nd([s['f0'] for s in samples], 0.0)
|
|
mel2ph = utils.collate_nd([s['mel2ph'] for s in samples], 0)
|
|
mel = utils.collate_nd([s['mel'] for s in samples], 0.0)
|
|
batch.update({
|
|
'tokens': tokens,
|
|
'mel2ph': mel2ph,
|
|
'mel': mel,
|
|
'f0': f0,
|
|
})
|
|
for v_name, v_pad in self.required_variances.items():
|
|
batch[v_name] = utils.collate_nd([s[v_name] for s in samples], v_pad)
|
|
if self.need_key_shift:
|
|
batch['key_shift'] = torch.FloatTensor([s['key_shift'] for s in samples])[:, None]
|
|
if self.need_speed:
|
|
batch['speed'] = torch.FloatTensor([s['speed'] for s in samples])[:, None]
|
|
if self.need_spk_id:
|
|
spk_ids = torch.LongTensor([s['spk_id'] for s in samples])
|
|
batch['spk_ids'] = spk_ids
|
|
if self.need_lang_id:
|
|
languages = utils.collate_nd([s['languages'] for s in samples], 0)
|
|
batch['languages'] = languages
|
|
return batch
|
|
|
|
|
|
class AcousticTask(BaseTask):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.dataset_cls = AcousticDataset
|
|
self.diffusion_type = hparams['diffusion_type']
|
|
assert self.diffusion_type in ['ddpm', 'reflow'], f"Unknown diffusion type: {self.diffusion_type}"
|
|
self.use_shallow_diffusion = hparams['use_shallow_diffusion']
|
|
if self.use_shallow_diffusion:
|
|
self.shallow_args = hparams['shallow_diffusion_args']
|
|
self.train_aux_decoder = self.shallow_args['train_aux_decoder']
|
|
self.train_diffusion = self.shallow_args['train_diffusion']
|
|
|
|
self.use_vocoder = hparams['infer'] or hparams['val_with_vocoder']
|
|
if self.use_vocoder:
|
|
self.vocoder: BaseVocoder = get_vocoder_cls(hparams)()
|
|
self.logged_gt_wav = set()
|
|
self.required_variances = []
|
|
if hparams['use_energy_embed']:
|
|
self.required_variances.append('energy')
|
|
if hparams['use_breathiness_embed']:
|
|
self.required_variances.append('breathiness')
|
|
if hparams['use_voicing_embed']:
|
|
self.required_variances.append('voicing')
|
|
if hparams['use_tension_embed']:
|
|
self.required_variances.append('tension')
|
|
super()._finish_init()
|
|
|
|
def _build_model(self):
|
|
return DiffSingerAcoustic(
|
|
vocab_size=len(self.phoneme_dictionary),
|
|
out_dims=hparams['audio_num_mel_bins']
|
|
)
|
|
|
|
# noinspection PyAttributeOutsideInit
|
|
def build_losses_and_metrics(self):
|
|
if self.use_shallow_diffusion:
|
|
self.aux_mel_loss = build_aux_loss(self.shallow_args['aux_decoder_arch'])
|
|
self.lambda_aux_mel_loss = hparams['lambda_aux_mel_loss']
|
|
self.register_validation_loss('aux_mel_loss')
|
|
if self.diffusion_type == 'ddpm':
|
|
self.mel_loss = DiffusionLoss(loss_type=hparams['main_loss_type'])
|
|
elif self.diffusion_type == 'reflow':
|
|
self.mel_loss = RectifiedFlowLoss(
|
|
loss_type=hparams['main_loss_type'], log_norm=hparams['main_loss_log_norm']
|
|
)
|
|
else:
|
|
raise ValueError(f"Unknown diffusion type: {self.diffusion_type}")
|
|
self.register_validation_loss('mel_loss')
|
|
|
|
def run_model(self, sample, infer=False):
|
|
txt_tokens = sample['tokens'] # [B, T_ph]
|
|
target = sample['mel'] # [B, T_s, M]
|
|
mel2ph = sample['mel2ph'] # [B, T_s]
|
|
f0 = sample['f0']
|
|
variances = {
|
|
v_name: sample[v_name]
|
|
for v_name in self.required_variances
|
|
}
|
|
key_shift = sample.get('key_shift')
|
|
speed = sample.get('speed')
|
|
|
|
if hparams['use_spk_id']:
|
|
spk_embed_id = sample['spk_ids']
|
|
else:
|
|
spk_embed_id = None
|
|
if hparams['use_lang_id']:
|
|
languages = sample['languages']
|
|
else:
|
|
languages = None
|
|
output: ShallowDiffusionOutput = self.model(
|
|
txt_tokens, mel2ph=mel2ph, f0=f0, **variances,
|
|
key_shift=key_shift, speed=speed,
|
|
spk_embed_id=spk_embed_id, languages=languages,
|
|
gt_mel=target, infer=infer
|
|
)
|
|
|
|
if infer:
|
|
return output
|
|
else:
|
|
losses = {}
|
|
|
|
if output.aux_out is not None:
|
|
aux_out = output.aux_out
|
|
norm_gt = self.model.aux_decoder.norm_spec(target)
|
|
aux_mel_loss = self.lambda_aux_mel_loss * self.aux_mel_loss(aux_out, norm_gt)
|
|
losses['aux_mel_loss'] = aux_mel_loss
|
|
|
|
non_padding = (mel2ph > 0).unsqueeze(-1).float()
|
|
if output.diff_out is not None:
|
|
if self.diffusion_type == 'ddpm':
|
|
x_recon, x_noise = output.diff_out
|
|
mel_loss = self.mel_loss(x_recon, x_noise, non_padding=non_padding)
|
|
elif self.diffusion_type == 'reflow':
|
|
v_pred, v_gt, t = output.diff_out
|
|
mel_loss = self.mel_loss(v_pred, v_gt, t=t, non_padding=non_padding)
|
|
else:
|
|
raise ValueError(f"Unknown diffusion type: {self.diffusion_type}")
|
|
losses['mel_loss'] = mel_loss
|
|
|
|
return losses
|
|
|
|
def on_train_start(self):
|
|
if self.use_vocoder and self.vocoder.get_device() != self.device:
|
|
self.vocoder.to_device(self.device)
|
|
|
|
def _on_validation_start(self):
|
|
if self.use_vocoder and self.vocoder.get_device() != self.device:
|
|
self.vocoder.to_device(self.device)
|
|
|
|
def _validation_step(self, sample, batch_idx):
|
|
losses = self.run_model(sample, infer=False)
|
|
if sample['size'] > 0 and min(sample['indices']) < hparams['num_valid_plots']:
|
|
mel_out: ShallowDiffusionOutput = self.run_model(sample, infer=True)
|
|
for i in range(len(sample['indices'])):
|
|
data_idx = sample['indices'][i].item()
|
|
if data_idx < hparams['num_valid_plots']:
|
|
if self.use_vocoder:
|
|
self.plot_wav(
|
|
data_idx, sample['mel'][i],
|
|
mel_out.aux_out[i] if mel_out.aux_out is not None else None,
|
|
mel_out.diff_out[i],
|
|
sample['f0'][i]
|
|
)
|
|
if mel_out.aux_out is not None:
|
|
self.plot_mel(data_idx, sample['mel'][i], mel_out.aux_out[i], 'auxmel')
|
|
if mel_out.diff_out is not None:
|
|
self.plot_mel(data_idx, sample['mel'][i], mel_out.diff_out[i], 'diffmel')
|
|
return losses, sample['size']
|
|
|
|
############
|
|
# validation plots
|
|
############
|
|
def plot_wav(self, data_idx, gt_mel, aux_mel, diff_mel, f0):
|
|
f0_len = self.valid_dataset.metadata['f0'][data_idx]
|
|
mel_len = self.valid_dataset.metadata['mel'][data_idx]
|
|
gt_mel = gt_mel[:mel_len].unsqueeze(0)
|
|
if aux_mel is not None:
|
|
aux_mel = aux_mel[:mel_len].unsqueeze(0)
|
|
if diff_mel is not None:
|
|
diff_mel = diff_mel[:mel_len].unsqueeze(0)
|
|
f0 = f0[:f0_len].unsqueeze(0)
|
|
if data_idx not in self.logged_gt_wav:
|
|
gt_wav = self.vocoder.spec2wav_torch(gt_mel, f0=f0)
|
|
self.logger.all_rank_experiment.add_audio(
|
|
f'gt_{data_idx}', gt_wav,
|
|
sample_rate=hparams['audio_sample_rate'],
|
|
global_step=self.global_step
|
|
)
|
|
self.logged_gt_wav.add(data_idx)
|
|
if aux_mel is not None:
|
|
aux_wav = self.vocoder.spec2wav_torch(aux_mel, f0=f0)
|
|
self.logger.all_rank_experiment.add_audio(
|
|
f'aux_{data_idx}', aux_wav,
|
|
sample_rate=hparams['audio_sample_rate'],
|
|
global_step=self.global_step
|
|
)
|
|
if diff_mel is not None:
|
|
diff_wav = self.vocoder.spec2wav_torch(diff_mel, f0=f0)
|
|
self.logger.all_rank_experiment.add_audio(
|
|
f'diff_{data_idx}', diff_wav,
|
|
sample_rate=hparams['audio_sample_rate'],
|
|
global_step=self.global_step
|
|
)
|
|
|
|
def plot_mel(self, data_idx, gt_spec, out_spec, name_prefix='mel'):
|
|
vmin = hparams['mel_vmin']
|
|
vmax = hparams['mel_vmax']
|
|
mel_len = self.valid_dataset.metadata['mel'][data_idx]
|
|
spec_cat = torch.cat([(out_spec - gt_spec).abs() + vmin, gt_spec, out_spec], -1)
|
|
title_text = f"{self.valid_dataset.metadata['spk_names'][data_idx]} - {self.valid_dataset.metadata['names'][data_idx]}"
|
|
self.logger.all_rank_experiment.add_figure(f'{name_prefix}_{data_idx}', spec_to_figure(
|
|
spec_cat[:mel_len], vmin, vmax, title_text
|
|
), global_step=self.global_step)
|