wehub-resource-sync
236 lines
9.8 KiB
Python
236 lines
9.8 KiB
Python
import copy
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
from modules.commons.common_layers import NormalInitEmbedding as Embedding
|
|
from modules.fastspeech.acoustic_encoder import FastSpeech2Acoustic
|
|
from modules.fastspeech.variance_encoder import FastSpeech2Variance
|
|
from utils.hparams import hparams
|
|
from utils.phoneme_utils import PAD_INDEX
|
|
|
|
f0_bin = 256
|
|
f0_max = 1100.0
|
|
f0_min = 50.0
|
|
f0_mel_min = 1127 * np.log(1 + f0_min / 700)
|
|
f0_mel_max = 1127 * np.log(1 + f0_max / 700)
|
|
|
|
|
|
def uniform_attention_pooling(spk_embed, durations):
|
|
_, T_mel, _ = spk_embed.shape
|
|
ph_starts = torch.cumsum(torch.cat([torch.zeros_like(durations[:, :1]), durations[:, :-1]], dim=1), dim=1)
|
|
ph_ends = ph_starts + durations
|
|
mel_indices = torch.arange(T_mel, device=spk_embed.device).view(1, 1, T_mel)
|
|
phoneme_to_mel_mask = (mel_indices >= ph_starts.unsqueeze(-1)) & (mel_indices < ph_ends.unsqueeze(-1))
|
|
uniform_scores = phoneme_to_mel_mask.float()
|
|
sum_scores = uniform_scores.sum(dim=2, keepdim=True)
|
|
attn_weights = uniform_scores / (sum_scores + (sum_scores == 0).float()) # [B, T_ph, T_mel]
|
|
ph_spk_embed = torch.bmm(attn_weights, spk_embed)
|
|
|
|
return ph_spk_embed
|
|
|
|
|
|
def f0_to_coarse(f0):
|
|
f0_mel = 1127 * (1 + f0 / 700).log()
|
|
a = (f0_bin - 2) / (f0_mel_max - f0_mel_min)
|
|
b = f0_mel_min * a - 1.
|
|
f0_mel = torch.where(f0_mel > 0, f0_mel * a - b, f0_mel)
|
|
torch.clip_(f0_mel, min=1., max=float(f0_bin - 1))
|
|
f0_coarse = torch.round(f0_mel).long()
|
|
return f0_coarse
|
|
|
|
|
|
class LengthRegulator(nn.Module):
|
|
# noinspection PyMethodMayBeStatic
|
|
def forward(self, dur):
|
|
token_idx = torch.arange(1, dur.shape[1] + 1, device=dur.device)[None, :, None]
|
|
dur_cumsum = torch.cumsum(dur, dim=1)
|
|
dur_cumsum_prev = F.pad(dur_cumsum, (1, -1), mode='constant', value=0)
|
|
pos_idx = torch.arange(dur.sum(dim=1).max(), device=dur.device)[None, None]
|
|
token_mask = (pos_idx >= dur_cumsum_prev[:, :, None]) & (pos_idx < dur_cumsum[:, :, None])
|
|
mel2ph = (token_idx * token_mask).sum(dim=1)
|
|
return mel2ph
|
|
|
|
|
|
class FastSpeech2AcousticONNX(FastSpeech2Acoustic):
|
|
def __init__(self, vocab_size, cross_lingual_token_idx=None):
|
|
super().__init__(vocab_size=vocab_size)
|
|
self.register_buffer(
|
|
'cross_lingual_token_idx',
|
|
torch.LongTensor(cross_lingual_token_idx),
|
|
persistent=False
|
|
) # [N,]
|
|
if len(cross_lingual_token_idx) == 0:
|
|
self.use_lang_id = False
|
|
|
|
# for temporary compatibility; will be completely removed in the future
|
|
self.f0_embed_type = hparams.get('f0_embed_type', 'continuous')
|
|
if self.f0_embed_type == 'discrete':
|
|
self.pitch_embed = Embedding(300, hparams['hidden_size'], PAD_INDEX)
|
|
|
|
self.lr = LengthRegulator()
|
|
if hparams['use_key_shift_embed']:
|
|
self.shift_min, self.shift_max = hparams['augmentation_args']['random_pitch_shifting']['range']
|
|
if hparams['use_speed_embed']:
|
|
self.speed_min, self.speed_max = hparams['augmentation_args']['random_time_stretching']['range']
|
|
|
|
# noinspection PyMethodOverriding
|
|
def forward(
|
|
self, tokens, durations,
|
|
f0, variances: dict,
|
|
gender=None, velocity=None,
|
|
spk_embed=None,
|
|
languages=None
|
|
):
|
|
txt_embed = self.txt_embed(tokens)
|
|
durations = durations * (tokens > 0)
|
|
mel2ph = self.lr(durations)
|
|
_mel2ph = mel2ph
|
|
f0 = f0 * (mel2ph > 0)
|
|
mel2ph = mel2ph[..., None].repeat((1, 1, hparams['hidden_size']))
|
|
if self.use_variance_scaling:
|
|
dur_embed = self.dur_embed(torch.log(1 + durations.float())[:, :, None])
|
|
else:
|
|
dur_embed = self.dur_embed(durations.float()[:, :, None])
|
|
if self.use_lang_id:
|
|
lang_mask = torch.any(
|
|
tokens[..., None] == self.cross_lingual_token_idx[None, None],
|
|
dim=-1
|
|
)
|
|
lang_embed = self.lang_embed(languages * lang_mask)
|
|
extra_embed = dur_embed + lang_embed
|
|
else:
|
|
extra_embed = dur_embed
|
|
if hparams.get('use_mix_ln', False):
|
|
if hasattr(self, 'frozen_spk_embed'):
|
|
ph_spk_embed = self.frozen_spk_embed.repeat(1, tokens.shape[1], 1)
|
|
else:
|
|
ph_spk_embed = uniform_attention_pooling(spk_embed, durations)
|
|
else:
|
|
ph_spk_embed = None
|
|
encoded = self.encoder(txt_embed, extra_embed, tokens == PAD_INDEX, spk_embed=ph_spk_embed)
|
|
encoded = F.pad(encoded, (0, 0, 1, 0))
|
|
condition = torch.gather(encoded, 1, mel2ph)
|
|
|
|
if self.use_stretch_embed:
|
|
stretch = torch.round(1000 * self.sr(_mel2ph, durations))
|
|
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)
|
|
condition += stretch_embed
|
|
stretch_embed_rnn_out, _ = self.stretch_embed_rnn(condition)
|
|
condition += stretch_embed_rnn_out
|
|
|
|
if self.f0_embed_type == 'discrete':
|
|
pitch = f0_to_coarse(f0)
|
|
pitch_embed = self.pitch_embed(pitch)
|
|
else:
|
|
f0_mel = (1 + f0 / 700).log()
|
|
pitch_embed = self.pitch_embed(f0_mel[:, :, None])
|
|
condition += pitch_embed
|
|
|
|
if self.use_variance_embeds:
|
|
variance_embeds = torch.stack([
|
|
self.variance_embeds[v_name](variances[v_name][:, :, None] * self.variance_scaling_factor[v_name])
|
|
for v_name in self.variance_embed_list
|
|
], dim=-1).sum(-1)
|
|
condition += variance_embeds
|
|
|
|
if hparams['use_key_shift_embed']:
|
|
if hasattr(self, 'frozen_key_shift'):
|
|
key_shift_embed = self.key_shift_embed(self.frozen_key_shift[:, None, None] * self.variance_scaling_factor['key_shift'])
|
|
else:
|
|
gender = torch.clip(gender, min=-1., max=1.)
|
|
gender_mask = (gender < 0.).float()
|
|
key_shift = gender * ((1. - gender_mask) * self.shift_max + gender_mask * abs(self.shift_min))
|
|
key_shift_embed = self.key_shift_embed(key_shift[:, :, None] * self.variance_scaling_factor['key_shift'])
|
|
condition += key_shift_embed
|
|
|
|
if hparams['use_speed_embed']:
|
|
if velocity is not None:
|
|
velocity = torch.clip(velocity, min=self.speed_min, max=self.speed_max)
|
|
speed_embed = self.speed_embed(velocity[:, :, None] * self.variance_scaling_factor['speed'])
|
|
else:
|
|
speed_embed = self.speed_embed(torch.FloatTensor([1.]).to(condition.device)[:, None, None] * self.variance_scaling_factor['speed'])
|
|
condition += speed_embed
|
|
|
|
if hparams['use_spk_id']:
|
|
if hasattr(self, 'frozen_spk_embed'):
|
|
condition += self.frozen_spk_embed
|
|
else:
|
|
condition += spk_embed
|
|
return condition
|
|
|
|
|
|
class FastSpeech2VarianceONNX(FastSpeech2Variance):
|
|
def __init__(self, vocab_size, cross_lingual_token_idx=None):
|
|
super().__init__(vocab_size=vocab_size)
|
|
self.register_buffer(
|
|
'cross_lingual_token_idx',
|
|
torch.LongTensor(cross_lingual_token_idx),
|
|
persistent=False
|
|
)
|
|
if len(cross_lingual_token_idx) == 0:
|
|
self.use_lang_id = False
|
|
self.lr = LengthRegulator()
|
|
|
|
def forward_encoder_word(self, tokens, word_div, word_dur, languages=None):
|
|
txt_embed = self.txt_embed(tokens)
|
|
ph2word = self.lr(word_div)
|
|
onset = ph2word > F.pad(ph2word, [1, -1])
|
|
onset_embed = self.onset_embed(onset.long())
|
|
ph_word_dur = torch.gather(F.pad(word_dur, [1, 0]), 1, ph2word)
|
|
word_dur_embed = self.word_dur_embed(ph_word_dur.float()[:, :, None])
|
|
extra_embed = onset_embed + word_dur_embed
|
|
if self.use_lang_id:
|
|
lang_mask = torch.any(
|
|
tokens[..., None] == self.cross_lingual_token_idx[None, None],
|
|
dim=-1
|
|
)
|
|
lang_embed = self.lang_embed(languages * lang_mask)
|
|
extra_embed += lang_embed
|
|
x_masks = tokens == PAD_INDEX
|
|
return self.encoder(txt_embed, extra_embed, x_masks), x_masks
|
|
|
|
def forward_encoder_phoneme(self, tokens, ph_dur, languages=None):
|
|
txt_embed = self.txt_embed(tokens)
|
|
if self.use_variance_scaling:
|
|
ph_dur_embed = self.ph_dur_embed(torch.log(1 + ph_dur.float())[:, :, None])
|
|
else:
|
|
ph_dur_embed = self.ph_dur_embed(ph_dur.float()[:, :, None])
|
|
if self.use_lang_id:
|
|
lang_mask = torch.any(
|
|
tokens[..., None] == self.cross_lingual_token_idx[None, None],
|
|
dim=-1
|
|
)
|
|
lang_embed = self.lang_embed(languages * lang_mask)
|
|
extra_embed = ph_dur_embed + lang_embed
|
|
else:
|
|
extra_embed = ph_dur_embed
|
|
x_masks = tokens == PAD_INDEX
|
|
return self.encoder(txt_embed, extra_embed, x_masks), x_masks
|
|
|
|
def forward_dur_predictor(self, encoder_out, x_masks, ph_midi, spk_embed=None):
|
|
midi_embed = self.midi_embed(ph_midi)
|
|
dur_cond = encoder_out + midi_embed
|
|
if hparams['use_spk_id'] and spk_embed is not None:
|
|
dur_cond += spk_embed
|
|
ph_dur = self.dur_predictor(dur_cond, x_masks=x_masks)
|
|
return ph_dur
|
|
|
|
def view_as_encoder(self):
|
|
model = copy.deepcopy(self)
|
|
if self.predict_dur:
|
|
del model.dur_predictor
|
|
model.forward = model.forward_encoder_word
|
|
else:
|
|
model.forward = model.forward_encoder_phoneme
|
|
return model
|
|
|
|
def view_as_dur_predictor(self):
|
|
model = copy.deepcopy(self)
|
|
del model.encoder
|
|
model.forward = model.forward_dur_predictor
|
|
return model
|