|
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).detach() |
|
nonpadding_mask_TB = 1 - padding_mask.float()[:, None, :] |
|
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) |
|
hiddens = hiddens.transpose(2, 3) |
|
x = self.out_proj(x.transpose(1, 2)) |
|
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) |
|
hiddens = [] |
|
for f in self.conv: |
|
x_ = f(x) |
|
x = x + x_ if self.res else x_ |
|
hiddens.append(x) |
|
x = x.transpose(1, -1) |
|
x = self.out_proj(x) |
|
if return_hiddens: |
|
hiddens = torch.stack(hiddens, 1) |
|
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' |
|
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 |