Spaces:
Sleeping
Sleeping
| import torch | |
| import torch.nn as nn | |
| from torch.nn import functional as F | |
| from modules.commons.common_layers import ( | |
| NormalInitEmbedding as Embedding, | |
| XavierUniformInitLinear as Linear, | |
| ) | |
| from modules.fastspeech.tts_modules import FastSpeech2Encoder, DurationPredictor | |
| from utils.hparams import hparams | |
| from utils.text_encoder import PAD_INDEX | |
| class FastSpeech2Variance(nn.Module): | |
| def __init__(self, vocab_size): | |
| super().__init__() | |
| self.predict_dur = hparams['predict_dur'] | |
| self.linguistic_mode = 'word' if hparams['predict_dur'] else 'phoneme' | |
| self.txt_embed = Embedding(vocab_size, hparams['hidden_size'], PAD_INDEX) | |
| if self.predict_dur: | |
| self.onset_embed = Embedding(2, hparams['hidden_size']) | |
| self.word_dur_embed = Linear(1, hparams['hidden_size']) | |
| else: | |
| self.ph_dur_embed = Linear(1, hparams['hidden_size']) | |
| self.encoder = FastSpeech2Encoder( | |
| hidden_size=hparams['hidden_size'], num_layers=hparams['enc_layers'], | |
| ffn_kernel_size=hparams['enc_ffn_kernel_size'], ffn_act=hparams['ffn_act'], | |
| dropout=hparams['dropout'], num_heads=hparams['num_heads'], | |
| use_pos_embed=hparams['use_pos_embed'], rel_pos=hparams.get('rel_pos', False), | |
| use_rope=hparams.get('use_rope', False) | |
| ) | |
| dur_hparams = hparams['dur_prediction_args'] | |
| if self.predict_dur: | |
| self.midi_embed = Embedding(128, hparams['hidden_size']) | |
| self.dur_predictor = DurationPredictor( | |
| in_dims=hparams['hidden_size'], | |
| n_chans=dur_hparams['hidden_size'], | |
| n_layers=dur_hparams['num_layers'], | |
| dropout_rate=dur_hparams['dropout'], | |
| kernel_size=dur_hparams['kernel_size'], | |
| offset=dur_hparams['log_offset'], | |
| dur_loss_type=dur_hparams['loss_type'] | |
| ) | |
| def forward(self, txt_tokens, midi, ph2word, ph_dur=None, word_dur=None, spk_embed=None, infer=True): | |
| """ | |
| :param txt_tokens: (train, infer) [B, T_ph] | |
| :param midi: (train, infer) [B, T_ph] | |
| :param ph2word: (train, infer) [B, T_ph] | |
| :param ph_dur: (train, [infer]) [B, T_ph] | |
| :param word_dur: (infer) [B, T_w] | |
| :param spk_embed: (train) [B, T_ph, H] | |
| :param infer: whether inference | |
| :return: encoder_out, ph_dur_pred | |
| """ | |
| txt_embed = self.txt_embed(txt_tokens) | |
| if self.linguistic_mode == 'word': | |
| b = txt_tokens.shape[0] | |
| onset = torch.diff(ph2word, dim=1, prepend=ph2word.new_zeros(b, 1)) > 0 | |
| onset_embed = self.onset_embed(onset.long()) # [B, T_ph, H] | |
| if word_dur is None or not infer: | |
| word_dur = ph_dur.new_zeros(b, ph2word.max() + 1).scatter_add( | |
| 1, ph2word, ph_dur | |
| )[:, 1:] # [B, T_ph] => [B, T_w] | |
| word_dur = torch.gather(F.pad(word_dur, [1, 0], value=0), 1, ph2word) # [B, T_w] => [B, T_ph] | |
| word_dur_embed = self.word_dur_embed(word_dur.float()[:, :, None]) | |
| encoder_out = self.encoder(txt_embed, onset_embed + word_dur_embed, txt_tokens == 0) | |
| else: | |
| ph_dur_embed = self.ph_dur_embed(ph_dur.float()[:, :, None]) | |
| encoder_out = self.encoder(txt_embed, ph_dur_embed, txt_tokens == 0) | |
| if self.predict_dur: | |
| midi_embed = self.midi_embed(midi) # => [B, T_ph, H] | |
| dur_cond = encoder_out + midi_embed | |
| if spk_embed is not None: | |
| dur_cond += spk_embed | |
| ph_dur_pred = self.dur_predictor(dur_cond, x_masks=txt_tokens == PAD_INDEX, infer=infer) | |
| return encoder_out, ph_dur_pred | |
| else: | |
| return encoder_out, None | |
| class MelodyEncoder(nn.Module): | |
| def __init__(self, enc_hparams: dict): | |
| super().__init__() | |
| def get_hparam(key): | |
| return enc_hparams.get(key, hparams.get(key)) | |
| # MIDI inputs | |
| hidden_size = get_hparam('hidden_size') | |
| self.note_midi_embed = Linear(1, hidden_size) | |
| self.note_dur_embed = Linear(1, hidden_size) | |
| # ornament inputs | |
| self.use_glide_embed = hparams['use_glide_embed'] | |
| self.glide_embed_scale = hparams['glide_embed_scale'] | |
| if self.use_glide_embed: | |
| # 0: none, 1: up, 2: down | |
| self.note_glide_embed = Embedding(len(hparams['glide_types']) + 1, hidden_size, padding_idx=0) | |
| self.encoder = FastSpeech2Encoder( | |
| hidden_size=hidden_size, num_layers=get_hparam('enc_layers'), | |
| ffn_kernel_size=get_hparam('enc_ffn_kernel_size'), ffn_act=get_hparam('ffn_act'), | |
| dropout=get_hparam('dropout'), num_heads=get_hparam('num_heads'), | |
| use_pos_embed=get_hparam('use_pos_embed'), rel_pos=get_hparam('rel_pos'), | |
| use_rope=get_hparam('use_rope') | |
| ) | |
| self.out_proj = Linear(hidden_size, hparams['hidden_size']) | |
| def forward(self, note_midi, note_rest, note_dur, glide=None): | |
| """ | |
| :param note_midi: float32 [B, T_n], -1: padding | |
| :param note_rest: bool [B, T_n] | |
| :param note_dur: int64 [B, T_n] | |
| :param glide: int64 [B, T_n] | |
| :return: [B, T_n, H] | |
| """ | |
| midi_embed = self.note_midi_embed(note_midi[:, :, None]) * ~note_rest[:, :, None] | |
| dur_embed = self.note_dur_embed(note_dur.float()[:, :, None]) | |
| ornament_embed = 0 | |
| if self.use_glide_embed: | |
| ornament_embed += self.note_glide_embed(glide) * self.glide_embed_scale | |
| encoder_out = self.encoder( | |
| midi_embed, dur_embed + ornament_embed, | |
| padding_mask=note_midi < 0 | |
| ) | |
| encoder_out = self.out_proj(encoder_out) | |
| return encoder_out | |