Spaces:
Build error
Build error
File size: 13,228 Bytes
9206300 |
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 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 |
import torch
from modules.GenerSpeech.model.glow_modules import Glow
from modules.fastspeech.tts_modules import PitchPredictor
import random
from modules.GenerSpeech.model.prosody_util import ProsodyAligner, LocalStyleAdaptor
from utils.pitch_utils import f0_to_coarse, denorm_f0
from modules.commons.common_layers import *
import torch.distributions as dist
from utils.hparams import hparams
from modules.GenerSpeech.model.mixstyle import MixStyle
from modules.fastspeech.fs2 import FastSpeech2
import json
from modules.fastspeech.tts_modules import DEFAULT_MAX_SOURCE_POSITIONS, DEFAULT_MAX_TARGET_POSITIONS
class GenerSpeech(FastSpeech2):
'''
GenerSpeech: Towards Style Transfer for Generalizable Out-Of-Domain Text-to-Speech
https://arxiv.org/abs/2205.07211
'''
def __init__(self, dictionary, out_dims=None):
super().__init__(dictionary, out_dims)
# Mixstyle
self.norm = MixStyle(p=0.5, alpha=0.1, eps=1e-6, hidden_size=self.hidden_size)
# emotion embedding
self.emo_embed_proj = Linear(256, self.hidden_size, bias=True)
# build prosody extractor
## frame level
self.prosody_extractor_utter = LocalStyleAdaptor(self.hidden_size, hparams['nVQ'], self.padding_idx)
self.l1_utter = nn.Linear(self.hidden_size * 2, self.hidden_size)
self.align_utter = ProsodyAligner(num_layers=2)
## phoneme level
self.prosody_extractor_ph = LocalStyleAdaptor(self.hidden_size, hparams['nVQ'], self.padding_idx)
self.l1_ph = nn.Linear(self.hidden_size * 2, self.hidden_size)
self.align_ph = ProsodyAligner(num_layers=2)
## word level
self.prosody_extractor_word = LocalStyleAdaptor(self.hidden_size, hparams['nVQ'], self.padding_idx)
self.l1_word = nn.Linear(self.hidden_size * 2, self.hidden_size)
self.align_word = ProsodyAligner(num_layers=2)
self.pitch_inpainter_predictor = PitchPredictor(
self.hidden_size, n_chans=self.hidden_size,
n_layers=3, dropout_rate=0.1, odim=2,
padding=hparams['ffn_padding'], kernel_size=hparams['predictor_kernel'])
# build attention layer
self.max_source_positions = DEFAULT_MAX_SOURCE_POSITIONS
self.embed_positions = SinusoidalPositionalEmbedding(
self.hidden_size, self.padding_idx,
init_size=self.max_source_positions + self.padding_idx + 1,
)
# build post flow
cond_hs = 80
if hparams.get('use_txt_cond', True):
cond_hs = cond_hs + hparams['hidden_size']
cond_hs = cond_hs + hparams['hidden_size'] * 3 # for emo, spk embedding and prosody embedding
self.post_flow = Glow(
80, hparams['post_glow_hidden'], hparams['post_glow_kernel_size'], 1,
hparams['post_glow_n_blocks'], hparams['post_glow_n_block_layers'],
n_split=4, n_sqz=2,
gin_channels=cond_hs,
share_cond_layers=hparams['post_share_cond_layers'],
share_wn_layers=hparams['share_wn_layers'],
sigmoid_scale=hparams['sigmoid_scale']
)
self.prior_dist = dist.Normal(0, 1)
def forward(self, txt_tokens, mel2ph=None, ref_mel2ph=None, ref_mel2word=None, spk_embed=None, emo_embed=None, ref_mels=None,
f0=None, uv=None, skip_decoder=False, global_steps=0, infer=False, **kwargs):
ret = {}
encoder_out = self.encoder(txt_tokens) # [B, T, C]
src_nonpadding = (txt_tokens > 0).float()[:, :, None]
# add spk/emo embed
spk_embed = self.spk_embed_proj(spk_embed)[:, None, :]
emo_embed = self.emo_embed_proj(emo_embed)[:, None, :]
# add dur
dur_inp = (encoder_out + spk_embed + emo_embed) * src_nonpadding
mel2ph = self.add_dur(dur_inp, mel2ph, txt_tokens, ret)
tgt_nonpadding = (mel2ph > 0).float()[:, :, None]
decoder_inp = self.expand_states(encoder_out, mel2ph)
decoder_inp = self.norm(decoder_inp, spk_embed + emo_embed)
# add prosody VQ
ret['ref_mel2ph'] = ref_mel2ph
ret['ref_mel2word'] = ref_mel2word
prosody_utter_mel = self.get_prosody_utter(decoder_inp, ref_mels, ret, infer, global_steps)
prosody_ph_mel = self.get_prosody_ph(decoder_inp, ref_mels, ret, infer, global_steps)
prosody_word_mel = self.get_prosody_word(decoder_inp, ref_mels, ret, infer, global_steps)
# add pitch embed
pitch_inp_domain_agnostic = decoder_inp * tgt_nonpadding
pitch_inp_domain_specific = (decoder_inp + spk_embed + emo_embed + prosody_utter_mel + prosody_ph_mel + prosody_word_mel) * tgt_nonpadding
predicted_pitch = self.inpaint_pitch(pitch_inp_domain_agnostic, pitch_inp_domain_specific, f0, uv, mel2ph, ret)
# decode
decoder_inp = decoder_inp + spk_embed + emo_embed + predicted_pitch + prosody_utter_mel + prosody_ph_mel + prosody_word_mel
ret['decoder_inp'] = decoder_inp = decoder_inp * tgt_nonpadding
if skip_decoder:
return ret
ret['mel_out'] = self.run_decoder(decoder_inp, tgt_nonpadding, ret, infer=infer, **kwargs)
# postflow
is_training = self.training
ret['x_mask'] = tgt_nonpadding
ret['spk_embed'] = spk_embed
ret['emo_embed'] = emo_embed
ret['ref_prosody'] = prosody_utter_mel + prosody_ph_mel + prosody_word_mel
self.run_post_glow(ref_mels, infer, is_training, ret)
return ret
def get_prosody_ph(self, encoder_out, ref_mels, ret, infer=False, global_steps=0):
# get VQ prosody
if global_steps > hparams['vq_start'] or infer:
prosody_embedding, loss, ppl = self.prosody_extractor_ph(ref_mels, ret['ref_mel2ph'], no_vq=False)
ret['vq_loss_ph'] = loss
ret['ppl_ph'] = ppl
else:
prosody_embedding = self.prosody_extractor_ph(ref_mels, ret['ref_mel2ph'], no_vq=True)
# add positional embedding
positions = self.embed_positions(prosody_embedding[:, :, 0])
prosody_embedding = self.l1_ph(torch.cat([prosody_embedding, positions], dim=-1))
# style-to-content attention
src_key_padding_mask = encoder_out[:, :, 0].eq(self.padding_idx).data
prosody_key_padding_mask = prosody_embedding[:, :, 0].eq(self.padding_idx).data
if global_steps < hparams['forcing']:
output, guided_loss, attn_emo = self.align_ph(encoder_out.transpose(0, 1), prosody_embedding.transpose(0, 1),
src_key_padding_mask, prosody_key_padding_mask, forcing=True)
else:
output, guided_loss, attn_emo = self.align_ph(encoder_out.transpose(0, 1), prosody_embedding.transpose(0, 1),
src_key_padding_mask, prosody_key_padding_mask, forcing=False)
ret['gloss_ph'] = guided_loss
ret['attn_ph'] = attn_emo
return output.transpose(0, 1)
def get_prosody_word(self, encoder_out, ref_mels, ret, infer=False, global_steps=0):
# get VQ prosody
if global_steps > hparams['vq_start'] or infer:
prosody_embedding, loss, ppl = self.prosody_extractor_word(ref_mels, ret['ref_mel2word'], no_vq=False)
ret['vq_loss_word'] = loss
ret['ppl_word'] = ppl
else:
prosody_embedding = self.prosody_extractor_word(ref_mels, ret['ref_mel2word'], no_vq=True)
# add positional embedding
positions = self.embed_positions(prosody_embedding[:, :, 0])
prosody_embedding = self.l1_word(torch.cat([prosody_embedding, positions], dim=-1))
# style-to-content attention
src_key_padding_mask = encoder_out[:, :, 0].eq(self.padding_idx).data
prosody_key_padding_mask = prosody_embedding[:, :, 0].eq(self.padding_idx).data
if global_steps < hparams['forcing']:
output, guided_loss, attn_emo = self.align_word(encoder_out.transpose(0, 1), prosody_embedding.transpose(0, 1),
src_key_padding_mask, prosody_key_padding_mask, forcing=True)
else:
output, guided_loss, attn_emo = self.align_word(encoder_out.transpose(0, 1), prosody_embedding.transpose(0, 1),
src_key_padding_mask, prosody_key_padding_mask, forcing=False)
ret['gloss_word'] = guided_loss
ret['attn_word'] = attn_emo
return output.transpose(0, 1)
def get_prosody_utter(self, encoder_out, ref_mels, ret, infer=False, global_steps=0):
# get VQ prosody
if global_steps > hparams['vq_start'] or infer:
prosody_embedding, loss, ppl = self.prosody_extractor_utter(ref_mels, no_vq=False)
ret['vq_loss_utter'] = loss
ret['ppl_utter'] = ppl
else:
prosody_embedding = self.prosody_extractor_utter(ref_mels, no_vq=True)
# add positional embedding
positions = self.embed_positions(prosody_embedding[:, :, 0])
prosody_embedding = self.l1_utter(torch.cat([prosody_embedding, positions], dim=-1))
# style-to-content attention
src_key_padding_mask = encoder_out[:, :, 0].eq(self.padding_idx).data
prosody_key_padding_mask = prosody_embedding[:, :, 0].eq(self.padding_idx).data
if global_steps < hparams['forcing']:
output, guided_loss, attn_emo = self.align_utter(encoder_out.transpose(0, 1), prosody_embedding.transpose(0, 1),
src_key_padding_mask, prosody_key_padding_mask, forcing=True)
else:
output, guided_loss, attn_emo = self.align_utter(encoder_out.transpose(0, 1), prosody_embedding.transpose(0, 1),
src_key_padding_mask, prosody_key_padding_mask, forcing=False)
ret['gloss_utter'] = guided_loss
ret['attn_utter'] = attn_emo
return output.transpose(0, 1)
def inpaint_pitch(self, pitch_inp_domain_agnostic, pitch_inp_domain_specific, f0, uv, mel2ph, ret):
if hparams['pitch_type'] == 'frame':
pitch_padding = mel2ph == 0
if hparams['predictor_grad'] != 1:
pitch_inp_domain_agnostic = pitch_inp_domain_agnostic.detach() + hparams['predictor_grad'] * (pitch_inp_domain_agnostic - pitch_inp_domain_agnostic.detach())
pitch_inp_domain_specific = pitch_inp_domain_specific.detach() + hparams['predictor_grad'] * (pitch_inp_domain_specific - pitch_inp_domain_specific.detach())
pitch_domain_agnostic = self.pitch_predictor(pitch_inp_domain_agnostic)
pitch_domain_specific = self.pitch_inpainter_predictor(pitch_inp_domain_specific)
pitch_pred = pitch_domain_agnostic + pitch_domain_specific
ret['pitch_pred'] = pitch_pred
use_uv = hparams['pitch_type'] == 'frame' and hparams['use_uv']
if f0 is None:
f0 = pitch_pred[:, :, 0] # [B, T]
if use_uv:
uv = pitch_pred[:, :, 1] > 0 # [B, T]
f0_denorm = denorm_f0(f0, uv if use_uv else None, hparams, pitch_padding=pitch_padding)
pitch = f0_to_coarse(f0_denorm) # start from 0 [B, T_txt]
ret['f0_denorm'] = f0_denorm
ret['f0_denorm_pred'] = denorm_f0(pitch_pred[:, :, 0], (pitch_pred[:, :, 1] > 0) if use_uv else None, hparams, pitch_padding=pitch_padding)
if hparams['pitch_type'] == 'ph':
pitch = torch.gather(F.pad(pitch, [1, 0]), 1, mel2ph)
ret['f0_denorm'] = torch.gather(F.pad(ret['f0_denorm'], [1, 0]), 1, mel2ph)
ret['f0_denorm_pred'] = torch.gather(F.pad(ret['f0_denorm_pred'], [1, 0]), 1, mel2ph)
pitch_embed = self.pitch_embed(pitch)
return pitch_embed
def run_post_glow(self, tgt_mels, infer, is_training, ret):
x_recon = ret['mel_out'].transpose(1, 2)
g = x_recon
B, _, T = g.shape
if hparams.get('use_txt_cond', True):
g = torch.cat([g, ret['decoder_inp'].transpose(1, 2)], 1)
g_spk_embed = ret['spk_embed'].repeat(1, T, 1).transpose(1, 2)
g_emo_embed = ret['emo_embed'].repeat(1, T, 1).transpose(1, 2)
l_ref_prosody = ret['ref_prosody'].transpose(1, 2)
g = torch.cat([g, g_spk_embed, g_emo_embed, l_ref_prosody], dim=1)
prior_dist = self.prior_dist
if not infer:
if is_training:
self.train()
x_mask = ret['x_mask'].transpose(1, 2)
y_lengths = x_mask.sum(-1)
g = g.detach()
tgt_mels = tgt_mels.transpose(1, 2)
z_postflow, ldj = self.post_flow(tgt_mels, x_mask, g=g)
ldj = ldj / y_lengths / 80
ret['z_pf'], ret['ldj_pf'] = z_postflow, ldj
ret['postflow'] = -prior_dist.log_prob(z_postflow).mean() - ldj.mean()
else:
x_mask = torch.ones_like(x_recon[:, :1, :])
z_post = prior_dist.sample(x_recon.shape).to(g.device) * hparams['noise_scale']
x_recon_, _ = self.post_flow(z_post, x_mask, g, reverse=True)
x_recon = x_recon_
ret['mel_out'] = x_recon.transpose(1, 2) |