chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,185 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.nn import functional as F
|
||||
|
||||
from modules.commons.common_layers import (
|
||||
NormalInitEmbedding as Embedding,
|
||||
SinusoidalPosEmb,
|
||||
AdamWLinear,
|
||||
)
|
||||
from modules.fastspeech.tts_modules import FastSpeech2Encoder, mel2ph_to_dur, StretchRegulator
|
||||
from utils.hparams import hparams
|
||||
from utils.phoneme_utils import PAD_INDEX
|
||||
|
||||
|
||||
class FastSpeech2Acoustic(nn.Module):
|
||||
def __init__(self, vocab_size):
|
||||
super().__init__()
|
||||
self.txt_embed = Embedding(vocab_size, hparams['hidden_size'], PAD_INDEX)
|
||||
self.use_lang_id = hparams.get('use_lang_id', False)
|
||||
if self.use_lang_id:
|
||||
self.lang_embed = Embedding(hparams['num_lang'] + 1, hparams['hidden_size'], padding_idx=0)
|
||||
|
||||
self.use_stretch_embed = hparams.get('use_stretch_embed', False)
|
||||
if self.use_stretch_embed:
|
||||
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._stretch_embed_rnn_flattened = False
|
||||
|
||||
self.dur_embed = AdamWLinear(1, hparams['hidden_size'])
|
||||
self.use_mix_ln = hparams.get('use_mix_ln', False)
|
||||
if self.use_mix_ln:
|
||||
self.mix_ln_layer = hparams['mix_ln_layer']
|
||||
else:
|
||||
self.mix_ln_layer = []
|
||||
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), rope_interleaved=hparams.get('rope_interleaved', True),
|
||||
mix_ln_layer=self.mix_ln_layer
|
||||
)
|
||||
|
||||
self.pitch_embed = AdamWLinear(1, hparams['hidden_size'])
|
||||
self.variance_embed_list = []
|
||||
self.use_energy_embed = hparams.get('use_energy_embed', False)
|
||||
self.use_breathiness_embed = hparams.get('use_breathiness_embed', False)
|
||||
self.use_voicing_embed = hparams.get('use_voicing_embed', False)
|
||||
self.use_tension_embed = hparams.get('use_tension_embed', False)
|
||||
if self.use_energy_embed:
|
||||
self.variance_embed_list.append('energy')
|
||||
if self.use_breathiness_embed:
|
||||
self.variance_embed_list.append('breathiness')
|
||||
if self.use_voicing_embed:
|
||||
self.variance_embed_list.append('voicing')
|
||||
if self.use_tension_embed:
|
||||
self.variance_embed_list.append('tension')
|
||||
|
||||
self.use_variance_embeds = len(self.variance_embed_list) > 0
|
||||
if self.use_variance_embeds:
|
||||
self.variance_embeds = nn.ModuleDict({
|
||||
v_name: AdamWLinear(1, hparams['hidden_size'])
|
||||
for v_name in self.variance_embed_list
|
||||
})
|
||||
|
||||
self.use_variance_scaling = hparams.get('use_variance_scaling', False)
|
||||
if self.use_variance_scaling:
|
||||
self.variance_scaling_factor = {
|
||||
'energy': 1. / 96, # 96 dB — max dynamic range of 16-bit audio
|
||||
'breathiness': 1. / 96,
|
||||
'voicing': 1. / 96,
|
||||
'tension': 0.1, # 1 / 10; tension logits are roughly [-10, 10]
|
||||
'key_shift': 1. / 12, # one octave — max key shift in most editors
|
||||
'speed': 1.
|
||||
}
|
||||
else:
|
||||
self.variance_scaling_factor = {
|
||||
'energy': 1.,
|
||||
'breathiness': 1.,
|
||||
'voicing': 1.,
|
||||
'tension': 1.,
|
||||
'key_shift': 1.,
|
||||
'speed': 1.
|
||||
}
|
||||
|
||||
self.use_key_shift_embed = hparams.get('use_key_shift_embed', False)
|
||||
if self.use_key_shift_embed:
|
||||
self.key_shift_embed = AdamWLinear(1, hparams['hidden_size'])
|
||||
|
||||
self.use_speed_embed = hparams.get('use_speed_embed', False)
|
||||
if self.use_speed_embed:
|
||||
self.speed_embed = AdamWLinear(1, hparams['hidden_size'])
|
||||
|
||||
self.use_spk_id = hparams['use_spk_id']
|
||||
if self.use_spk_id:
|
||||
self.spk_embed = Embedding(hparams['num_spk'], hparams['hidden_size'])
|
||||
|
||||
def forward_variance_embedding(self, condition, key_shift=None, speed=None, **variances):
|
||||
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 self.use_key_shift_embed:
|
||||
key_shift_embed = self.key_shift_embed(key_shift[:, :, None] * self.variance_scaling_factor['key_shift'])
|
||||
condition += key_shift_embed
|
||||
|
||||
if self.use_speed_embed:
|
||||
speed_embed = self.speed_embed(speed[:, :, None] * self.variance_scaling_factor['speed'])
|
||||
condition += speed_embed
|
||||
|
||||
return condition
|
||||
|
||||
def forward(
|
||||
self, txt_tokens, mel2ph, f0,
|
||||
key_shift=None, speed=None,
|
||||
spk_embed_id=None, languages=None,
|
||||
**kwargs
|
||||
):
|
||||
spk_embed = None
|
||||
if self.use_spk_id:
|
||||
spk_mix_embed = kwargs.get('spk_mix_embed')
|
||||
if spk_mix_embed is not None:
|
||||
spk_embed = spk_mix_embed
|
||||
else:
|
||||
spk_embed = self.spk_embed(spk_embed_id)[:, None, :]
|
||||
txt_embed = self.txt_embed(txt_tokens)
|
||||
dur = mel2ph_to_dur(mel2ph, txt_tokens.shape[1])
|
||||
if self.use_variance_scaling:
|
||||
dur_embed = self.dur_embed(torch.log(1 + dur[:, :, None].float()))
|
||||
else:
|
||||
dur_embed = self.dur_embed(dur[:, :, None].float())
|
||||
if self.use_lang_id:
|
||||
lang_embed = self.lang_embed(languages)
|
||||
extra_embed = dur_embed + lang_embed
|
||||
else:
|
||||
extra_embed = dur_embed
|
||||
encoder_out = self.encoder(txt_embed, extra_embed, txt_tokens == 0, spk_embed)
|
||||
|
||||
encoder_out = F.pad(encoder_out, [0, 0, 1, 0])
|
||||
mel2ph_ = mel2ph[..., None].repeat([1, 1, encoder_out.shape[-1]])
|
||||
condition = torch.gather(encoder_out, 1, mel2ph_)
|
||||
|
||||
if self.use_stretch_embed:
|
||||
stretch = torch.round(1000 * self.sr(mel2ph, 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
|
||||
# flatten_parameters fuses the GRU weights into a contiguous buffer for cuDNN.
|
||||
# It only needs to happen once after weight init, device change, or load_state_dict.
|
||||
# We guard with a flag to avoid the redundant call on every forward.
|
||||
# Limitation: the flag lives on this module and is invisible to PyTorch. After
|
||||
# load_state_dict() or model.to(device) replaces the GRU weights, the flag stays
|
||||
# True and flatten_parameters is skipped — cuDNN will fall back to the slower path.
|
||||
# To restore the fast path, reset the flag manually: model._stretch_embed_rnn_flattened = False
|
||||
if not self._stretch_embed_rnn_flattened:
|
||||
self.stretch_embed_rnn.flatten_parameters()
|
||||
self._stretch_embed_rnn_flattened = True
|
||||
stretch_embed_rnn_out, _ = self.stretch_embed_rnn(condition)
|
||||
condition = condition + stretch_embed_rnn_out
|
||||
|
||||
if self.use_spk_id:
|
||||
condition += spk_embed
|
||||
|
||||
f0_mel = (1 + f0 / 700).log()
|
||||
pitch_embed = self.pitch_embed(f0_mel[:, :, None])
|
||||
condition += pitch_embed
|
||||
|
||||
condition = self.forward_variance_embedding(
|
||||
condition, key_shift=key_shift, speed=speed, **kwargs
|
||||
)
|
||||
|
||||
return condition
|
||||
@@ -0,0 +1,95 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
|
||||
import modules.compat as compat
|
||||
from modules.core.ddpm import MultiVarianceDiffusion
|
||||
from utils import filter_kwargs
|
||||
from utils.hparams import hparams
|
||||
|
||||
VARIANCE_CHECKLIST = ['energy', 'breathiness', 'voicing', 'tension']
|
||||
|
||||
|
||||
class ParameterAdaptorModule(torch.nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.variance_prediction_list = []
|
||||
self.predict_energy = hparams.get('predict_energy', False)
|
||||
self.predict_breathiness = hparams.get('predict_breathiness', False)
|
||||
self.predict_voicing = hparams.get('predict_voicing', False)
|
||||
self.predict_tension = hparams.get('predict_tension', False)
|
||||
if self.predict_energy:
|
||||
self.variance_prediction_list.append('energy')
|
||||
if self.predict_breathiness:
|
||||
self.variance_prediction_list.append('breathiness')
|
||||
if self.predict_voicing:
|
||||
self.variance_prediction_list.append('voicing')
|
||||
if self.predict_tension:
|
||||
self.variance_prediction_list.append('tension')
|
||||
self.predict_variances = len(self.variance_prediction_list) > 0
|
||||
|
||||
def build_adaptor(self, cls=MultiVarianceDiffusion):
|
||||
ranges = []
|
||||
clamps = []
|
||||
|
||||
if self.predict_energy:
|
||||
ranges.append((
|
||||
hparams['energy_db_min'],
|
||||
hparams['energy_db_max']
|
||||
))
|
||||
clamps.append((hparams['energy_db_min'], 0.))
|
||||
|
||||
if self.predict_breathiness:
|
||||
ranges.append((
|
||||
hparams['breathiness_db_min'],
|
||||
hparams['breathiness_db_max']
|
||||
))
|
||||
clamps.append((hparams['breathiness_db_min'], 0.))
|
||||
|
||||
if self.predict_voicing:
|
||||
ranges.append((
|
||||
hparams['voicing_db_min'],
|
||||
hparams['voicing_db_max']
|
||||
))
|
||||
clamps.append((hparams['voicing_db_min'], 0.))
|
||||
|
||||
if self.predict_tension:
|
||||
ranges.append((
|
||||
hparams['tension_logit_min'],
|
||||
hparams['tension_logit_max']
|
||||
))
|
||||
clamps.append((
|
||||
hparams['tension_logit_min'],
|
||||
hparams['tension_logit_max']
|
||||
))
|
||||
|
||||
variances_hparams = hparams['variances_prediction_args']
|
||||
total_repeat_bins = variances_hparams['total_repeat_bins']
|
||||
assert total_repeat_bins % len(self.variance_prediction_list) == 0, \
|
||||
f'Total number of repeat bins must be divisible by number of ' \
|
||||
f'variance parameters ({len(self.variance_prediction_list)}).'
|
||||
repeat_bins = total_repeat_bins // len(self.variance_prediction_list)
|
||||
backbone_type = compat.get_backbone_type(hparams, nested_config=variances_hparams)
|
||||
backbone_args = compat.get_backbone_args(variances_hparams, backbone_type=backbone_type)
|
||||
kwargs = filter_kwargs(
|
||||
{
|
||||
'ranges': ranges,
|
||||
'clamps': clamps,
|
||||
'repeat_bins': repeat_bins,
|
||||
'timesteps': hparams.get('timesteps'),
|
||||
'time_scale_factor': hparams.get('time_scale_factor'),
|
||||
'backbone_type': backbone_type,
|
||||
'backbone_args': backbone_args
|
||||
},
|
||||
cls
|
||||
)
|
||||
return cls(**kwargs)
|
||||
|
||||
def collect_variance_inputs(self, **kwargs) -> list:
|
||||
return [kwargs.get(name) for name in self.variance_prediction_list]
|
||||
|
||||
def collect_variance_outputs(self, variances: list | tuple) -> dict:
|
||||
return {
|
||||
name: pred
|
||||
for name, pred in zip(self.variance_prediction_list, variances)
|
||||
}
|
||||
@@ -0,0 +1,455 @@
|
||||
import math
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.nn import functional as F
|
||||
from modules.commons.rotary_embedding_torch import RotaryEmbedding
|
||||
from modules.commons.common_layers import SinusoidalPositionalEmbedding, EncSALayer, AdamWLinear
|
||||
from modules.commons.espnet_positional_embedding import RelPositionalEncoding
|
||||
|
||||
DEFAULT_MAX_SOURCE_POSITIONS = 2000
|
||||
DEFAULT_MAX_TARGET_POSITIONS = 2000
|
||||
|
||||
|
||||
class TransformerEncoderLayer(nn.Module):
|
||||
def __init__(self, hidden_size, dropout, kernel_size=None, act='gelu', num_heads=2, rotary_embed=None,
|
||||
layer_idx=None, mix_ln_layer=None):
|
||||
super().__init__()
|
||||
self.op = EncSALayer(
|
||||
hidden_size, num_heads, dropout=dropout,
|
||||
attention_dropout=0.0, relu_dropout=dropout,
|
||||
kernel_size=kernel_size,
|
||||
act=act, rotary_embed=rotary_embed,
|
||||
layer_idx=layer_idx, mix_ln_layer=mix_ln_layer
|
||||
)
|
||||
|
||||
def forward(self, x, **kwargs):
|
||||
return self.op(x, **kwargs)
|
||||
|
||||
|
||||
######################
|
||||
# fastspeech modules
|
||||
######################
|
||||
class LayerNorm(torch.nn.LayerNorm):
|
||||
"""Layer normalization module.
|
||||
:param int nout: output dim size
|
||||
:param int dim: dimension to be normalized
|
||||
"""
|
||||
|
||||
def __init__(self, nout, dim=-1):
|
||||
"""Construct an LayerNorm object."""
|
||||
super(LayerNorm, self).__init__(nout, eps=1e-12)
|
||||
self.dim = dim
|
||||
|
||||
def forward(self, x):
|
||||
"""Apply layer normalization.
|
||||
:param torch.Tensor x: input tensor
|
||||
:return: layer normalized tensor
|
||||
:rtype torch.Tensor
|
||||
"""
|
||||
if self.dim == -1:
|
||||
return super(LayerNorm, self).forward(x)
|
||||
return super(LayerNorm, self).forward(x.transpose(1, -1)).transpose(1, -1)
|
||||
|
||||
|
||||
class DurationPredictor(torch.nn.Module):
|
||||
"""Duration predictor module.
|
||||
This is a module of duration predictor described in `FastSpeech: Fast, Robust and Controllable Text to Speech`_.
|
||||
The duration predictor predicts a duration of each frame in log domain from the hidden embeddings of encoder.
|
||||
.. _`FastSpeech: Fast, Robust and Controllable Text to Speech`:
|
||||
https://arxiv.org/pdf/1905.09263.pdf
|
||||
Note:
|
||||
The calculation domain of outputs is different between in `forward` and in `inference`. In `forward`,
|
||||
the outputs are calculated in log domain but in `inference`, those are calculated in linear domain.
|
||||
"""
|
||||
|
||||
def __init__(self, in_dims, n_layers=2, n_chans=384, kernel_size=3,
|
||||
dropout_rate=0.1, offset=1.0, dur_loss_type='mse', arch='resnet'):
|
||||
"""Initialize duration predictor module.
|
||||
Args:
|
||||
in_dims (int): Input dimension.
|
||||
n_layers (int, optional): Number of convolutional layers.
|
||||
n_chans (int, optional): Number of channels of convolutional layers.
|
||||
kernel_size (int, optional): Kernel size of convolutional layers.
|
||||
dropout_rate (float, optional): Dropout rate.
|
||||
offset (float, optional): Offset value to avoid nan in log domain.
|
||||
"""
|
||||
super(DurationPredictor, self).__init__()
|
||||
self.offset = offset
|
||||
self.conv = torch.nn.ModuleList()
|
||||
self.kernel_size = kernel_size
|
||||
self.use_resnet = (arch == 'resnet')
|
||||
for idx in range(n_layers):
|
||||
in_chans = in_dims if idx == 0 else n_chans
|
||||
if self.use_resnet:
|
||||
self.conv.append(nn.Sequential(
|
||||
LayerNorm(in_chans, dim=1),
|
||||
nn.Conv1d(in_chans, n_chans, kernel_size, stride=1, padding=kernel_size // 2),
|
||||
nn.ReLU(),
|
||||
nn.Conv1d(n_chans, n_chans, 1),
|
||||
nn.Dropout(dropout_rate)
|
||||
))
|
||||
else:
|
||||
self.conv.append(nn.Sequential(
|
||||
nn.Identity(), # this is a placeholder for ConstantPad1d which is now merged into Conv1d
|
||||
nn.Conv1d(in_chans, n_chans, kernel_size, stride=1, padding=kernel_size // 2),
|
||||
nn.ReLU(),
|
||||
LayerNorm(n_chans, dim=1),
|
||||
nn.Dropout(dropout_rate)
|
||||
))
|
||||
if self.use_resnet and in_dims != n_chans:
|
||||
self.res_conv = nn.Conv1d(in_dims, n_chans, 1)
|
||||
else:
|
||||
self.res_conv = None
|
||||
self.loss_type = dur_loss_type
|
||||
if self.loss_type in ['mse', 'huber']:
|
||||
self.out_dims = 1
|
||||
# elif hparams['dur_loss_type'] == 'mog':
|
||||
# out_dims = 15
|
||||
# elif hparams['dur_loss_type'] == 'crf':
|
||||
# out_dims = 32
|
||||
# from torchcrf import CRF
|
||||
# self.crf = CRF(out_dims, batch_first=True)
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
self.linear = AdamWLinear(n_chans, self.out_dims)
|
||||
|
||||
def out2dur(self, xs):
|
||||
if self.loss_type in ['mse', 'huber']:
|
||||
# NOTE: calculate loss in log domain
|
||||
dur = xs.squeeze(-1).exp() - self.offset # (B, Tmax)
|
||||
# elif hparams['dur_loss_type'] == 'crf':
|
||||
# dur = torch.LongTensor(self.crf.decode(xs)).cuda()
|
||||
else:
|
||||
raise NotImplementedError()
|
||||
return dur
|
||||
|
||||
def forward(self, xs, x_masks=None, infer=True):
|
||||
"""Calculate forward propagation.
|
||||
Args:
|
||||
xs (Tensor): Batch of input sequences (B, Tmax, idim).
|
||||
x_masks (BoolTensor, optional): Batch of masks indicating padded part (B, Tmax).
|
||||
infer (bool): Whether inference
|
||||
Returns:
|
||||
(train) FloatTensor, (infer) LongTensor: Batch of predicted durations in linear domain (B, Tmax).
|
||||
"""
|
||||
xs = xs.transpose(1, -1) # (B, idim, Tmax)
|
||||
masks = 1 - x_masks.float()
|
||||
masks_ = masks[:, None, :]
|
||||
for idx, f in enumerate(self.conv):
|
||||
if self.use_resnet:
|
||||
residual = self.res_conv(xs) if idx == 0 and self.res_conv is not None else xs
|
||||
xs = residual + f(xs)
|
||||
else:
|
||||
xs = f(xs)
|
||||
if x_masks is not None:
|
||||
xs = xs * masks_
|
||||
xs = self.linear(xs.transpose(1, -1)) # [B, T, C]
|
||||
xs = xs * masks[:, :, None] # (B, T, C)
|
||||
|
||||
dur_pred = self.out2dur(xs)
|
||||
if infer:
|
||||
dur_pred = dur_pred.clamp(min=0.) # avoid negative value
|
||||
return dur_pred
|
||||
|
||||
|
||||
class VariancePredictor(torch.nn.Module):
|
||||
def __init__(self, vmin, vmax, in_dims,
|
||||
n_layers=5, n_chans=512, kernel_size=5,
|
||||
dropout_rate=0.1):
|
||||
"""Initialize variance predictor module.
|
||||
Args:
|
||||
in_dims (int): Input dimension.
|
||||
n_layers (int, optional): Number of convolutional layers.
|
||||
n_chans (int, optional): Number of channels of convolutional layers.
|
||||
kernel_size (int, optional): Kernel size of convolutional layers.
|
||||
dropout_rate (float, optional): Dropout rate.
|
||||
"""
|
||||
super(VariancePredictor, self).__init__()
|
||||
|
||||
self.vmin = vmin
|
||||
self.vmax = vmax
|
||||
self.conv = torch.nn.ModuleList()
|
||||
self.kernel_size = kernel_size
|
||||
for idx in range(n_layers):
|
||||
in_chans = in_dims if idx == 0 else n_chans
|
||||
self.conv.append(torch.nn.Sequential(
|
||||
torch.nn.Conv1d(in_chans, n_chans, kernel_size, stride=1, padding=kernel_size // 2),
|
||||
torch.nn.ReLU(),
|
||||
LayerNorm(n_chans, dim=1),
|
||||
torch.nn.Dropout(dropout_rate)
|
||||
))
|
||||
self.linear = torch.nn.Linear(n_chans, 1)
|
||||
self.embed_positions = SinusoidalPositionalEmbedding(in_dims, 0, init_size=4096)
|
||||
self.pos_embed_alpha = nn.Parameter(torch.Tensor([1]))
|
||||
|
||||
def out2value(self, xs):
|
||||
return (xs + 1) / 2 * (self.vmax - self.vmin) + self.vmin
|
||||
|
||||
def forward(self, xs, infer=True):
|
||||
"""
|
||||
:param xs: [B, T, H]
|
||||
:param infer: whether inference
|
||||
:return: [B, T]
|
||||
"""
|
||||
positions = self.pos_embed_alpha * self.embed_positions(xs[..., 0])
|
||||
xs = xs + positions
|
||||
xs = xs.transpose(1, -1) # (B, idim, Tmax)
|
||||
for f in self.conv:
|
||||
xs = f(xs) # (B, C, Tmax)
|
||||
xs = self.linear(xs.transpose(1, -1)).squeeze(-1) # (B, Tmax)
|
||||
if infer:
|
||||
xs = self.out2value(xs)
|
||||
return xs
|
||||
|
||||
|
||||
class PitchPredictor(torch.nn.Module):
|
||||
def __init__(self, vmin, vmax, num_bins, deviation,
|
||||
in_dims, n_layers=5, n_chans=384, kernel_size=5,
|
||||
dropout_rate=0.1):
|
||||
"""Initialize pitch predictor module.
|
||||
Args:
|
||||
in_dims (int): Input dimension.
|
||||
n_layers (int, optional): Number of convolutional layers.
|
||||
n_chans (int, optional): Number of channels of convolutional layers.
|
||||
kernel_size (int, optional): Kernel size of convolutional layers.
|
||||
dropout_rate (float, optional): Dropout rate.
|
||||
"""
|
||||
super(PitchPredictor, self).__init__()
|
||||
self.vmin = vmin
|
||||
self.vmax = vmax
|
||||
self.interval = (vmax - vmin) / (num_bins - 1) # align with centers of bins
|
||||
self.sigma = deviation / self.interval
|
||||
self.register_buffer('x', torch.arange(num_bins).float().reshape(1, 1, -1)) # [1, 1, N]
|
||||
|
||||
self.base_pitch_embed = torch.nn.Linear(1, in_dims)
|
||||
self.conv = torch.nn.ModuleList()
|
||||
self.kernel_size = kernel_size
|
||||
for idx in range(n_layers):
|
||||
in_chans = in_dims if idx == 0 else n_chans
|
||||
self.conv.append(torch.nn.Sequential(
|
||||
torch.nn.Conv1d(in_chans, n_chans, kernel_size, stride=1, padding=kernel_size // 2),
|
||||
torch.nn.ReLU(),
|
||||
LayerNorm(n_chans, dim=1),
|
||||
torch.nn.Dropout(dropout_rate)
|
||||
))
|
||||
self.linear = torch.nn.Linear(n_chans, num_bins)
|
||||
self.embed_positions = SinusoidalPositionalEmbedding(in_dims, 0, init_size=4096)
|
||||
self.pos_embed_alpha = nn.Parameter(torch.Tensor([1]))
|
||||
|
||||
def bins_to_values(self, bins):
|
||||
return bins * self.interval + self.vmin
|
||||
|
||||
def out2pitch(self, probs):
|
||||
logits = probs.sigmoid() # [B, T, N]
|
||||
# return logits
|
||||
# logits_sum = logits.sum(dim=2) # [B, T]
|
||||
bins = torch.sum(self.x * logits, dim=2) / torch.sum(logits, dim=2) # [B, T]
|
||||
pitch = self.bins_to_values(bins)
|
||||
# uv = logits_sum / (self.sigma * math.sqrt(2 * math.pi)) < 0.3
|
||||
# pitch[uv] = torch.nan
|
||||
return pitch
|
||||
|
||||
def forward(self, xs, base):
|
||||
"""
|
||||
:param xs: [B, T, H]
|
||||
:param base: [B, T]
|
||||
:return: [B, T, N]
|
||||
"""
|
||||
xs = xs + self.base_pitch_embed(base[..., None])
|
||||
positions = self.pos_embed_alpha * self.embed_positions(xs[..., 0])
|
||||
xs = xs + positions
|
||||
xs = xs.transpose(1, -1) # (B, idim, Tmax)
|
||||
for f in self.conv:
|
||||
xs = f(xs) # (B, C, Tmax)
|
||||
xs = self.linear(xs.transpose(1, -1)) # (B, Tmax, H)
|
||||
return self.out2pitch(xs) + base, xs
|
||||
|
||||
|
||||
class RhythmRegulator(torch.nn.Module):
|
||||
def __init__(self, eps=1e-5):
|
||||
super().__init__()
|
||||
self.eps = eps
|
||||
|
||||
def forward(self, ph_dur, ph2word, word_dur):
|
||||
"""
|
||||
Example (no batch dim version):
|
||||
1. ph_dur = [4,2,3,2]
|
||||
2. word_dur = [3,4,2], ph2word = [1,2,2,3]
|
||||
3. word_dur_in = [4,5,2]
|
||||
4. alpha_w = [0.75,0.8,1], alpha_ph = [0.75,0.8,0.8,1]
|
||||
5. ph_dur_out = [3,1.6,2.4,2]
|
||||
:param ph_dur: [B, T_ph]
|
||||
:param ph2word: [B, T_ph]
|
||||
:param word_dur: [B, T_w]
|
||||
"""
|
||||
ph_dur = ph_dur.float() * (ph2word > 0)
|
||||
word_dur = word_dur.float()
|
||||
word_dur_in = ph_dur.new_zeros(ph_dur.shape[0], ph2word.max() + 1).scatter_add(
|
||||
1, ph2word, ph_dur
|
||||
)[:, 1:] # [B, T_ph] => [B, T_w]
|
||||
alpha_w = word_dur / word_dur_in.clamp(min=self.eps) # avoid dividing by zero
|
||||
alpha_ph = torch.gather(F.pad(alpha_w, [1, 0]), 1, ph2word) # [B, T_w] => [B, T_ph]
|
||||
ph_dur_out = ph_dur * alpha_ph
|
||||
return ph_dur_out.round().long()
|
||||
|
||||
|
||||
class LengthRegulator(torch.nn.Module):
|
||||
# noinspection PyMethodMayBeStatic
|
||||
def forward(self, dur, dur_padding=None, alpha=None):
|
||||
"""
|
||||
Example (no batch dim version):
|
||||
1. dur = [2,2,3]
|
||||
2. token_idx = [[1],[2],[3]], dur_cumsum = [2,4,7], dur_cumsum_prev = [0,2,4]
|
||||
3. token_mask = [[1,1,0,0,0,0,0],
|
||||
[0,0,1,1,0,0,0],
|
||||
[0,0,0,0,1,1,1]]
|
||||
4. token_idx * token_mask = [[1,1,0,0,0,0,0],
|
||||
[0,0,2,2,0,0,0],
|
||||
[0,0,0,0,3,3,3]]
|
||||
5. (token_idx * token_mask).sum(0) = [1,1,2,2,3,3,3]
|
||||
|
||||
:param dur: Batch of durations of each frame (B, T_txt)
|
||||
:param dur_padding: Batch of padding of each frame (B, T_txt)
|
||||
:param alpha: duration rescale coefficient
|
||||
:return:
|
||||
mel2ph (B, T_speech)
|
||||
"""
|
||||
assert alpha is None or alpha > 0
|
||||
if alpha is not None:
|
||||
dur = torch.round(dur.float() * alpha).long()
|
||||
if dur_padding is not None:
|
||||
dur = dur * (1 - dur_padding.long())
|
||||
token_idx = torch.arange(1, dur.shape[1] + 1)[None, :, None].to(dur.device)
|
||||
dur_cumsum = torch.cumsum(dur, 1)
|
||||
dur_cumsum_prev = F.pad(dur_cumsum, [1, -1], mode='constant', value=0)
|
||||
|
||||
pos_idx = torch.arange(dur.sum(-1).max())[None, None].to(dur.device)
|
||||
token_mask = (pos_idx >= dur_cumsum_prev[:, :, None]) & (pos_idx < dur_cumsum[:, :, None])
|
||||
mel2ph = (token_idx * token_mask.long()).sum(1)
|
||||
return mel2ph
|
||||
|
||||
|
||||
class StretchRegulator(torch.nn.Module):
|
||||
# noinspection PyMethodMayBeStatic
|
||||
def forward(self, mel2ph, dur=None):
|
||||
"""
|
||||
Example (no batch dim version):
|
||||
1. dur = [2,4,3]
|
||||
2. mel2ph = [1,1,2,2,2,2,3,3,3]
|
||||
3. mel2dur = [2,2,4,4,4,4,3,3,3]
|
||||
4. bound_mask = [0,1,0,0,0,1,0,0,1]
|
||||
5. 1 - bound_mask * mel2dur = [1,-1,1,1,1,-3,1,1,-2] => pad => [0,1,-1,1,1,1,-3,1,1]
|
||||
6. stretch_denorm = [0,1,0,1,2,3,0,1,2]
|
||||
|
||||
:param dur: Batch of durations of each frame (B, T_txt)
|
||||
:param mel2ph: Batch of mel2ph (B, T_speech)
|
||||
:return:
|
||||
stretch (B, T_speech)
|
||||
"""
|
||||
if dur is None:
|
||||
dur = mel2ph_to_dur(mel2ph, mel2ph.max())
|
||||
dur = torch.cat([torch.ones_like(dur[:, :1]), dur], dim=1) # Avoid dividing by zero
|
||||
mel2dur = torch.gather(dur, 1, mel2ph)
|
||||
bound_mask = torch.gt(mel2ph[:, 1:], mel2ph[:, :-1])
|
||||
stretch_delta = 1 - bound_mask * mel2dur[:, :-1]
|
||||
stretch_delta = F.pad(stretch_delta, [1, 0])
|
||||
stretch_denorm = torch.cumsum(stretch_delta, dim=1)
|
||||
stretch = stretch_denorm.float() / mel2dur
|
||||
return stretch * (mel2ph > 0)
|
||||
|
||||
|
||||
def mel2ph_to_dur(mel2ph, T_txt, max_dur=None):
|
||||
B, _ = mel2ph.shape
|
||||
dur = mel2ph.new_zeros(B, T_txt + 1).scatter_add(1, mel2ph, torch.ones_like(mel2ph))
|
||||
dur = dur[:, 1:]
|
||||
if max_dur is not None:
|
||||
dur = dur.clamp(max=max_dur)
|
||||
return dur
|
||||
|
||||
|
||||
class FastSpeech2Encoder(nn.Module):
|
||||
def __init__(
|
||||
self, hidden_size, num_layers,
|
||||
ffn_kernel_size=9, ffn_act='gelu',
|
||||
dropout=None, num_heads=2, use_pos_embed=True, rel_pos=True,
|
||||
use_rope=False, rope_interleaved=True, mix_ln_layer=None
|
||||
):
|
||||
super().__init__()
|
||||
self.num_layers = num_layers
|
||||
embed_dim = self.hidden_size = hidden_size
|
||||
self.dropout = dropout
|
||||
self.use_pos_embed = use_pos_embed
|
||||
if use_pos_embed and use_rope:
|
||||
if embed_dim % (num_heads * 2) != 0:
|
||||
raise ValueError(
|
||||
"RoPE requires the hidden size to be multiple of "
|
||||
f"num_heads * 2 = {num_heads * 2}, but got {embed_dim}."
|
||||
)
|
||||
rotary_embed = RotaryEmbedding(dim=embed_dim // num_heads, interleaved=rope_interleaved)
|
||||
else:
|
||||
rotary_embed = None
|
||||
self.layers = nn.ModuleList([
|
||||
TransformerEncoderLayer(
|
||||
self.hidden_size, self.dropout,
|
||||
kernel_size=ffn_kernel_size, act=ffn_act,
|
||||
num_heads=num_heads, rotary_embed=rotary_embed,
|
||||
layer_idx=i, mix_ln_layer=mix_ln_layer
|
||||
)
|
||||
for i in range(self.num_layers)
|
||||
])
|
||||
self.layer_norm = nn.LayerNorm(embed_dim)
|
||||
|
||||
self.embed_scale = math.sqrt(hidden_size)
|
||||
self.padding_idx = 0
|
||||
self.rel_pos = rel_pos
|
||||
if use_rope:
|
||||
self.embed_positions = None
|
||||
elif self.rel_pos:
|
||||
self.embed_positions = RelPositionalEncoding(hidden_size, dropout_rate=0.0)
|
||||
else:
|
||||
self.embed_positions = SinusoidalPositionalEmbedding(
|
||||
hidden_size, self.padding_idx, init_size=DEFAULT_MAX_TARGET_POSITIONS,
|
||||
)
|
||||
|
||||
def forward_embedding(self, main_embed, extra_embed=None, padding_mask=None):
|
||||
# embed tokens and positions
|
||||
x = self.embed_scale * main_embed
|
||||
if extra_embed is not None:
|
||||
x = x + extra_embed
|
||||
if self.use_pos_embed and self.embed_positions is not None:
|
||||
if self.rel_pos:
|
||||
x = self.embed_positions(x)
|
||||
else:
|
||||
positions = self.embed_positions(~padding_mask)
|
||||
x = x + positions
|
||||
x = F.dropout(x, p=self.dropout, training=self.training)
|
||||
return x
|
||||
|
||||
def forward(self, main_embed, extra_embed, padding_mask, spk_embed=None, attn_mask=None, return_hiddens=False):
|
||||
x = self.forward_embedding(main_embed, extra_embed, padding_mask=padding_mask) # [B, T, H]
|
||||
nonpadding_mask_BT = 1 - padding_mask.float()[:, :, None] # [B, T, 1]
|
||||
|
||||
# NOTICE:
|
||||
# The following codes are commented out because
|
||||
# `self.use_pos_embed` is always False in the older versions,
|
||||
# and this argument did not compat with `hparams['use_pos_embed']`,
|
||||
# which defaults to True. The new version fixed this inconsistency,
|
||||
# resulting in temporary removal of pos_embed_alpha, which has actually
|
||||
# never been used before.
|
||||
|
||||
# if self.use_pos_embed:
|
||||
# positions = self.pos_embed_alpha * self.embed_positions(x[..., 0])
|
||||
# x = x + positions
|
||||
# x = F.dropout(x, p=self.dropout, training=self.training)
|
||||
|
||||
x = x * nonpadding_mask_BT
|
||||
hiddens = []
|
||||
for layer in self.layers:
|
||||
x = layer(x, encoder_padding_mask=padding_mask, cond=spk_embed, attn_mask=attn_mask) * nonpadding_mask_BT
|
||||
if return_hiddens:
|
||||
hiddens.append(x)
|
||||
x = self.layer_norm(x) * nonpadding_mask_BT
|
||||
if return_hiddens:
|
||||
x = torch.stack(hiddens, 0) # [L, B, T, C]
|
||||
return x
|
||||
@@ -0,0 +1,158 @@
|
||||
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,
|
||||
AdamWLinear,
|
||||
)
|
||||
from modules.fastspeech.tts_modules import FastSpeech2Encoder, DurationPredictor
|
||||
from utils.hparams import hparams
|
||||
from utils.phoneme_utils 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.use_lang_id = hparams['use_lang_id']
|
||||
self.use_variance_scaling = hparams.get('use_variance_scaling', False)
|
||||
self.txt_embed = Embedding(vocab_size, hparams['hidden_size'], PAD_INDEX)
|
||||
if self.use_lang_id:
|
||||
self.lang_embed = Embedding(hparams['num_lang'] + 1, hparams['hidden_size'], padding_idx=0)
|
||||
|
||||
if self.predict_dur:
|
||||
self.onset_embed = Embedding(2, hparams['hidden_size'])
|
||||
self.word_dur_embed = AdamWLinear(1, hparams['hidden_size'])
|
||||
else:
|
||||
self.ph_dur_embed = AdamWLinear(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), rope_interleaved=hparams.get('rope_interleaved', True)
|
||||
)
|
||||
|
||||
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'],
|
||||
arch=dur_hparams['arch']
|
||||
)
|
||||
|
||||
def forward(
|
||||
self, txt_tokens, midi, ph2word,
|
||||
ph_dur=None, word_dur=None,
|
||||
spk_embed=None, languages=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 languages (train, infer) [B, T_ph]
|
||||
: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])
|
||||
extra_embed = onset_embed + word_dur_embed
|
||||
elif self.use_variance_scaling:
|
||||
extra_embed = self.ph_dur_embed(torch.log(1 + ph_dur.float())[:, :, None])
|
||||
else:
|
||||
extra_embed = self.ph_dur_embed(ph_dur.float()[:, :, None])
|
||||
|
||||
if self.use_lang_id:
|
||||
lang_embed = self.lang_embed(languages)
|
||||
extra_embed += lang_embed
|
||||
encoder_out = self.encoder(txt_embed, extra_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.use_variance_scaling = hparams.get('use_variance_scaling', False)
|
||||
self.note_midi_embed = AdamWLinear(1, hidden_size)
|
||||
self.note_dur_embed = AdamWLinear(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'), rope_interleaved=hparams.get('rope_interleaved', True)
|
||||
)
|
||||
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]
|
||||
"""
|
||||
if self.use_variance_scaling:
|
||||
midi_embed = self.note_midi_embed(note_midi[:, :, None] / 128)
|
||||
dur_embed = self.note_dur_embed(torch.log(1 + note_dur.float())[:, :, None])
|
||||
else:
|
||||
midi_embed = self.note_midi_embed(note_midi[:, :, None])
|
||||
dur_embed = self.note_dur_embed(note_dur.float()[:, :, None])
|
||||
midi_embed *= ~note_rest[:, :, 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
|
||||
Reference in New Issue
Block a user