from typing import Dict import torch import torch.nn as nn import torch.nn.functional as F from torch import Tensor import modules.compat as compat from basics.base_module import CategorizedModule from modules.aux_decoder import AuxDecoderAdaptor from modules.commons.common_layers import ( NormalInitEmbedding as Embedding, SinusoidalPosEmb, AdamWLinear, ) from modules.core import ( GaussianDiffusion, PitchDiffusion, MultiVarianceDiffusion, RectifiedFlow, PitchRectifiedFlow, MultiVarianceRectifiedFlow ) from modules.fastspeech.acoustic_encoder import FastSpeech2Acoustic from modules.fastspeech.param_adaptor import ParameterAdaptorModule from modules.fastspeech.tts_modules import RhythmRegulator, LengthRegulator, StretchRegulator from modules.fastspeech.variance_encoder import FastSpeech2Variance, MelodyEncoder from utils.hparams import hparams class ShallowDiffusionOutput: def __init__(self, *, aux_out=None, diff_out=None): self.aux_out = aux_out self.diff_out = diff_out class DiffSingerAcoustic(CategorizedModule, ParameterAdaptorModule): @property def category(self): return 'acoustic' def __init__(self, vocab_size, out_dims): CategorizedModule.__init__(self) ParameterAdaptorModule.__init__(self) self.fs2 = FastSpeech2Acoustic( vocab_size=vocab_size ) self.use_shallow_diffusion = hparams.get('use_shallow_diffusion', False) self.shallow_args = hparams.get('shallow_diffusion_args', {}) if self.use_shallow_diffusion: self.train_aux_decoder = self.shallow_args['train_aux_decoder'] self.train_diffusion = self.shallow_args['train_diffusion'] self.aux_decoder_grad = self.shallow_args['aux_decoder_grad'] self.aux_decoder = AuxDecoderAdaptor( in_dims=hparams['hidden_size'], out_dims=out_dims, num_feats=1, spec_min=hparams['spec_min'], spec_max=hparams['spec_max'], aux_decoder_arch=self.shallow_args['aux_decoder_arch'], aux_decoder_args=self.shallow_args['aux_decoder_args'] ) self.diffusion_type = hparams.get('diffusion_type', 'ddpm') self.backbone_type = compat.get_backbone_type(hparams) self.backbone_args = compat.get_backbone_args(hparams, self.backbone_type) if self.diffusion_type == 'ddpm': self.diffusion = GaussianDiffusion( out_dims=out_dims, num_feats=1, timesteps=hparams['timesteps'], k_step=hparams['K_step'], backbone_type=self.backbone_type, backbone_args=self.backbone_args, spec_min=hparams['spec_min'], spec_max=hparams['spec_max'] ) elif self.diffusion_type == 'reflow': self.diffusion = RectifiedFlow( out_dims=out_dims, num_feats=1, t_start=hparams['T_start'], time_scale_factor=hparams['time_scale_factor'], backbone_type=self.backbone_type, backbone_args=self.backbone_args, spec_min=hparams['spec_min'], spec_max=hparams['spec_max'] ) else: raise NotImplementedError(self.diffusion_type) def forward( self, txt_tokens, mel2ph, f0, key_shift=None, speed=None, spk_embed_id=None, languages=None, gt_mel=None, infer=True, **kwargs ) -> ShallowDiffusionOutput: condition = self.fs2( txt_tokens, mel2ph, f0, key_shift=key_shift, speed=speed, spk_embed_id=spk_embed_id, languages=languages, **kwargs ) if infer: if self.use_shallow_diffusion: aux_mel_pred = self.aux_decoder(condition, infer=True) aux_mel_pred *= ((mel2ph > 0).float()[:, :, None]) if gt_mel is not None and self.shallow_args['val_gt_start']: src_mel = gt_mel else: src_mel = aux_mel_pred else: aux_mel_pred = src_mel = None mel_pred = self.diffusion(condition, src_spec=src_mel, infer=True) mel_pred *= ((mel2ph > 0).float()[:, :, None]) return ShallowDiffusionOutput(aux_out=aux_mel_pred, diff_out=mel_pred) else: if self.use_shallow_diffusion: if self.train_aux_decoder: aux_cond = condition * self.aux_decoder_grad + condition.detach() * (1 - self.aux_decoder_grad) aux_out = self.aux_decoder(aux_cond, infer=False) else: aux_out = None if self.train_diffusion: diff_out = self.diffusion(condition, gt_spec=gt_mel, infer=False) else: diff_out = None return ShallowDiffusionOutput(aux_out=aux_out, diff_out=diff_out) else: aux_out = None diff_out = self.diffusion(condition, gt_spec=gt_mel, infer=False) return ShallowDiffusionOutput(aux_out=aux_out, diff_out=diff_out) class DiffSingerVariance(CategorizedModule, ParameterAdaptorModule): @property def category(self): return 'variance' def __init__(self, vocab_size): CategorizedModule.__init__(self) ParameterAdaptorModule.__init__(self) self.predict_dur = hparams['predict_dur'] self.predict_pitch = hparams['predict_pitch'] self.use_stretch_embed = hparams.get('use_stretch_embed', None) assert self.use_stretch_embed is not None, "You may be loading an old version of the model checkpoint, which is incompatible with the new version due to some bug fixes. It is recommended to roll back to the old version (commit id: 6df0ee977c3728f14cb79c2db8b19df30b23a0bf)" if self.use_stretch_embed and (self.predict_pitch or self.predict_variances): self.sr = StretchRegulator() self.stretch_embed = nn.Sequential( SinusoidalPosEmb(hparams['hidden_size']), nn.Linear(hparams['hidden_size'], hparams['hidden_size'] * 4), nn.GELU(), nn.Linear(hparams['hidden_size'] * 4, hparams['hidden_size']), ) self.stretch_embed_rnn = nn.GRU(hparams['hidden_size'], hparams['hidden_size'], 1, batch_first=True) self.use_spk_id = hparams['use_spk_id'] if self.use_spk_id: self.spk_embed = Embedding(hparams['num_spk'], hparams['hidden_size']) self.fs2 = FastSpeech2Variance( vocab_size=vocab_size ) self.rr = RhythmRegulator() self.lr = LengthRegulator() self.diffusion_type = hparams.get('diffusion_type', 'ddpm') if self.predict_pitch: self.use_melody_encoder = hparams.get('use_melody_encoder', False) if self.use_melody_encoder: self.melody_encoder = MelodyEncoder(enc_hparams=hparams['melody_encoder_args']) self.delta_pitch_embed = AdamWLinear(1, hparams['hidden_size']) else: self.base_pitch_embed = AdamWLinear(1, hparams['hidden_size']) self.pitch_retake_embed = Embedding(2, hparams['hidden_size']) pitch_hparams = hparams['pitch_prediction_args'] self.pitch_backbone_type = compat.get_backbone_type(hparams, nested_config=pitch_hparams) self.pitch_backbone_args = compat.get_backbone_args(pitch_hparams, backbone_type=self.pitch_backbone_type) if self.diffusion_type == 'ddpm': self.pitch_predictor = PitchDiffusion( vmin=pitch_hparams['pitd_norm_min'], vmax=pitch_hparams['pitd_norm_max'], cmin=pitch_hparams['pitd_clip_min'], cmax=pitch_hparams['pitd_clip_max'], repeat_bins=pitch_hparams['repeat_bins'], timesteps=hparams['timesteps'], k_step=hparams['K_step'], backbone_type=self.pitch_backbone_type, backbone_args=self.pitch_backbone_args ) elif self.diffusion_type == 'reflow': self.pitch_predictor = PitchRectifiedFlow( vmin=pitch_hparams['pitd_norm_min'], vmax=pitch_hparams['pitd_norm_max'], cmin=pitch_hparams['pitd_clip_min'], cmax=pitch_hparams['pitd_clip_max'], repeat_bins=pitch_hparams['repeat_bins'], time_scale_factor=hparams['time_scale_factor'], backbone_type=self.pitch_backbone_type, backbone_args=self.pitch_backbone_args ) else: raise ValueError(f"Invalid diffusion type: {self.diffusion_type}") if self.predict_variances: self.pitch_embed = AdamWLinear(1, hparams['hidden_size']) self.variance_embeds = nn.ModuleDict({ v_name: AdamWLinear(1, hparams['hidden_size']) for v_name in self.variance_prediction_list }) if self.diffusion_type == 'ddpm': self.variance_predictor = self.build_adaptor(cls=MultiVarianceDiffusion) elif self.diffusion_type == 'reflow': self.variance_predictor = self.build_adaptor(cls=MultiVarianceRectifiedFlow) else: raise NotImplementedError(self.diffusion_type) self.use_variance_scaling = hparams.get('use_variance_scaling', False) self.custom_variance_scaling_factor = { 'energy': 1. / 96, 'breathiness': 1. / 96, 'voicing': 1. / 96, 'tension': 0.1, 'key_shift': 1. / 12, 'speed': 1. } self.default_variance_scaling_factor = { 'energy': 1., 'breathiness': 1., 'voicing': 1., 'tension': 1., 'key_shift': 1., 'speed': 1. } if self.use_variance_scaling: self.variance_retake_scaling = self.custom_variance_scaling_factor else: self.variance_retake_scaling = self.default_variance_scaling_factor def forward( self, txt_tokens, midi, ph2word, ph_dur=None, word_dur=None, mel2ph=None, note_midi=None, note_rest=None, note_dur=None, note_glide=None, mel2note=None, base_pitch=None, pitch=None, pitch_expr=None, pitch_retake=None, variance_retake: Dict[str, Tensor] = None, spk_id=None, languages=None, infer=True, **kwargs ): if self.use_spk_id: ph_spk_mix_embed = kwargs.get('ph_spk_mix_embed') spk_mix_embed = kwargs.get('spk_mix_embed') if ph_spk_mix_embed is not None and spk_mix_embed is not None: ph_spk_embed = ph_spk_mix_embed spk_embed = spk_mix_embed else: ph_spk_embed = spk_embed = self.spk_embed(spk_id)[:, None, :] # [B,] => [B, T=1, H] else: ph_spk_embed = spk_embed = None encoder_out, dur_pred_out = self.fs2( txt_tokens, midi=midi, ph2word=ph2word, ph_dur=ph_dur, word_dur=word_dur, spk_embed=ph_spk_embed, languages=languages, infer=infer ) if not self.predict_pitch and not self.predict_variances: return dur_pred_out, None, ({} if infer else None) if mel2ph is None and word_dur is not None: # inference from file dur_pred_align = self.rr(dur_pred_out, ph2word, word_dur) mel2ph = self.lr(dur_pred_align) mel2ph = F.pad(mel2ph, [0, base_pitch.shape[1] - mel2ph.shape[1]]) encoder_out = F.pad(encoder_out, [0, 0, 1, 0]) mel2ph_ = mel2ph[..., None].repeat([1, 1, hparams['hidden_size']]) condition = torch.gather(encoder_out, 1, mel2ph_) if self.use_stretch_embed: stretch = torch.round(1000 * self.sr(mel2ph, ph_dur)) if self.training and stretch.numel() > 1000: # construct a phoneme stretching index lookup table with a total of 1001 indexes (0~1000) table = self.stretch_embed(torch.arange(0, 1001, device=stretch.device)) stretch_embed = torch.index_select(table, 0, stretch.view(-1).long()).view_as(condition) else: stretch_embed = self.stretch_embed(stretch) condition += stretch_embed self.stretch_embed_rnn.flatten_parameters() stretch_embed_rnn_out, _ = self.stretch_embed_rnn(condition) condition = condition + stretch_embed_rnn_out if self.use_spk_id: condition += spk_embed if self.predict_pitch: if self.use_melody_encoder: melody_encoder_out = self.melody_encoder( note_midi, note_rest, note_dur, glide=note_glide ) melody_encoder_out = F.pad(melody_encoder_out, [0, 0, 1, 0]) mel2note_ = mel2note[..., None].repeat([1, 1, hparams['hidden_size']]) melody_condition = torch.gather(melody_encoder_out, 1, mel2note_) pitch_cond = condition + melody_condition else: pitch_cond = condition.clone() # preserve the original tensor to avoid further inplace operations retake_unset = pitch_retake is None if retake_unset: pitch_retake = torch.ones_like(mel2ph, dtype=torch.bool) if pitch_expr is None: pitch_retake_embed = self.pitch_retake_embed(pitch_retake.long()) else: retake_true_embed = self.pitch_retake_embed( torch.ones(1, 1, dtype=torch.long, device=txt_tokens.device) ) # [B=1, T=1] => [B=1, T=1, H] retake_false_embed = self.pitch_retake_embed( torch.zeros(1, 1, dtype=torch.long, device=txt_tokens.device) ) # [B=1, T=1] => [B=1, T=1, H] pitch_expr = (pitch_expr * pitch_retake)[:, :, None] # [B, T, 1] pitch_retake_embed = pitch_expr * retake_true_embed + (1. - pitch_expr) * retake_false_embed pitch_cond += pitch_retake_embed if self.use_melody_encoder: if retake_unset: # generate from scratch delta_pitch_in = torch.zeros_like(base_pitch) else: delta_pitch_in = (pitch - base_pitch) * ~pitch_retake if self.use_variance_scaling: pitch_cond += self.delta_pitch_embed(delta_pitch_in[:, :, None] / 12) else: pitch_cond += self.delta_pitch_embed(delta_pitch_in[:, :, None]) else: if not retake_unset: # retake base_pitch = base_pitch * pitch_retake + pitch * ~pitch_retake if self.use_variance_scaling: pitch_cond += self.base_pitch_embed(base_pitch[:, :, None] / 128) else: pitch_cond += self.base_pitch_embed(base_pitch[:, :, None]) if infer: pitch_pred_out = self.pitch_predictor(pitch_cond, infer=True) else: pitch_pred_out = self.pitch_predictor(pitch_cond, pitch - base_pitch, infer=False) else: pitch_pred_out = None if not self.predict_variances: return dur_pred_out, pitch_pred_out, ({} if infer else None) if pitch is None: pitch = base_pitch + pitch_pred_out if self.use_variance_scaling: var_cond = condition + self.pitch_embed(pitch[:, :, None] / 12) else: var_cond = condition + self.pitch_embed(pitch[:, :, None]) variance_inputs = self.collect_variance_inputs(**kwargs) if variance_retake is not None: variance_embeds = [ self.variance_embeds[v_name](v_input[:, :, None] * self.variance_retake_scaling[v_name]) * ~variance_retake[v_name][:, :, None] for v_name, v_input in zip(self.variance_prediction_list, variance_inputs) ] var_cond += torch.stack(variance_embeds, dim=-1).sum(-1) variance_outputs = self.variance_predictor(var_cond, variance_inputs, infer=infer) if infer: variances_pred_out = self.collect_variance_outputs(variance_outputs) else: variances_pred_out = variance_outputs return dur_pred_out, pitch_pred_out, variances_pred_out