Kangarroar's picture
Upload 154 files
ed1cdd1
raw
history blame
5.59 kB
from modules.commons.common_layers import *
from utils.hparams import hparams
from modules.fastspeech.tts_modules import PitchPredictor
from utils.pitch_utils import denorm_f0
class Prenet(nn.Module):
def __init__(self, in_dim=80, out_dim=256, kernel=5, n_layers=3, strides=None):
super(Prenet, self).__init__()
padding = kernel // 2
self.layers = []
self.strides = strides if strides is not None else [1] * n_layers
for l in range(n_layers):
self.layers.append(nn.Sequential(
nn.Conv1d(in_dim, out_dim, kernel_size=kernel, padding=padding, stride=self.strides[l]),
nn.ReLU(),
nn.BatchNorm1d(out_dim)
))
in_dim = out_dim
self.layers = nn.ModuleList(self.layers)
self.out_proj = nn.Linear(out_dim, out_dim)
def forward(self, x):
"""
:param x: [B, T, 80]
:return: [L, B, T, H], [B, T, H]
"""
# padding_mask = x.abs().sum(-1).eq(0).data # [B, T]
padding_mask = x.abs().sum(-1).eq(0).detach()
nonpadding_mask_TB = 1 - padding_mask.float()[:, None, :] # [B, 1, T]
x = x.transpose(1, 2)
hiddens = []
for i, l in enumerate(self.layers):
nonpadding_mask_TB = nonpadding_mask_TB[:, :, ::self.strides[i]]
x = l(x) * nonpadding_mask_TB
hiddens.append(x)
hiddens = torch.stack(hiddens, 0) # [L, B, H, T]
hiddens = hiddens.transpose(2, 3) # [L, B, T, H]
x = self.out_proj(x.transpose(1, 2)) # [B, T, H]
x = x * nonpadding_mask_TB.transpose(1, 2)
return hiddens, x
class ConvBlock(nn.Module):
def __init__(self, idim=80, n_chans=256, kernel_size=3, stride=1, norm='gn', dropout=0):
super().__init__()
self.conv = ConvNorm(idim, n_chans, kernel_size, stride=stride)
self.norm = norm
if self.norm == 'bn':
self.norm = nn.BatchNorm1d(n_chans)
elif self.norm == 'in':
self.norm = nn.InstanceNorm1d(n_chans, affine=True)
elif self.norm == 'gn':
self.norm = nn.GroupNorm(n_chans // 16, n_chans)
elif self.norm == 'ln':
self.norm = LayerNorm(n_chans // 16, n_chans)
elif self.norm == 'wn':
self.conv = torch.nn.utils.weight_norm(self.conv.conv)
self.dropout = nn.Dropout(dropout)
self.relu = nn.ReLU()
def forward(self, x):
"""
:param x: [B, C, T]
:return: [B, C, T]
"""
x = self.conv(x)
if not isinstance(self.norm, str):
if self.norm == 'none':
pass
elif self.norm == 'ln':
x = self.norm(x.transpose(1, 2)).transpose(1, 2)
else:
x = self.norm(x)
x = self.relu(x)
x = self.dropout(x)
return x
class ConvStacks(nn.Module):
def __init__(self, idim=80, n_layers=5, n_chans=256, odim=32, kernel_size=5, norm='gn',
dropout=0, strides=None, res=True):
super().__init__()
self.conv = torch.nn.ModuleList()
self.kernel_size = kernel_size
self.res = res
self.in_proj = Linear(idim, n_chans)
if strides is None:
strides = [1] * n_layers
else:
assert len(strides) == n_layers
for idx in range(n_layers):
self.conv.append(ConvBlock(
n_chans, n_chans, kernel_size, stride=strides[idx], norm=norm, dropout=dropout))
self.out_proj = Linear(n_chans, odim)
def forward(self, x, return_hiddens=False):
"""
:param x: [B, T, H]
:return: [B, T, H]
"""
x = self.in_proj(x)
x = x.transpose(1, -1) # (B, idim, Tmax)
hiddens = []
for f in self.conv:
x_ = f(x)
x = x + x_ if self.res else x_ # (B, C, Tmax)
hiddens.append(x)
x = x.transpose(1, -1)
x = self.out_proj(x) # (B, Tmax, H)
if return_hiddens:
hiddens = torch.stack(hiddens, 1) # [B, L, C, T]
return x, hiddens
return x
class PitchExtractor(nn.Module):
def __init__(self, n_mel_bins=80, conv_layers=2):
super().__init__()
self.hidden_size = hparams['hidden_size']
self.predictor_hidden = hparams['predictor_hidden'] if hparams['predictor_hidden'] > 0 else self.hidden_size
self.conv_layers = conv_layers
self.mel_prenet = Prenet(n_mel_bins, self.hidden_size, strides=[1, 1, 1])
if self.conv_layers > 0:
self.mel_encoder = ConvStacks(
idim=self.hidden_size, n_chans=self.hidden_size, odim=self.hidden_size, n_layers=self.conv_layers)
self.pitch_predictor = PitchPredictor(
self.hidden_size, n_chans=self.predictor_hidden,
n_layers=5, dropout_rate=0.1, odim=2,
padding=hparams['ffn_padding'], kernel_size=hparams['predictor_kernel'])
def forward(self, mel_input=None):
ret = {}
mel_hidden = self.mel_prenet(mel_input)[1]
if self.conv_layers > 0:
mel_hidden = self.mel_encoder(mel_hidden)
ret['pitch_pred'] = pitch_pred = self.pitch_predictor(mel_hidden)
pitch_padding = mel_input.abs().sum(-1) == 0
use_uv = hparams['pitch_type'] == 'frame' #and hparams['use_uv']
ret['f0_denorm_pred'] = denorm_f0(
pitch_pred[:, :, 0], (pitch_pred[:, :, 1] > 0) if use_uv else None,
hparams, pitch_padding=pitch_padding)
return ret