# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Modified from espnet(https://github.com/espnet/espnet)
"""StyleMelGAN's TADEResBlock Modules."""
from functools import partial
import paddle.nn.functional as F
from paddle import nn
[docs]class TADELayer(nn.Layer):
"""TADE Layer module."""
def __init__(
self,
in_channels: int=64,
aux_channels: int=80,
kernel_size: int=9,
bias: bool=True,
upsample_factor: int=2,
upsample_mode: str="nearest", ):
"""Initilize TADE layer."""
super().__init__()
self.norm = nn.InstanceNorm1D(
in_channels,
momentum=0.1,
data_format="NCL",
weight_attr=False,
bias_attr=False)
self.aux_conv = nn.Sequential(
nn.Conv1D(
aux_channels,
in_channels,
kernel_size,
1,
bias_attr=bias,
padding=(kernel_size - 1) // 2, ), )
self.gated_conv = nn.Sequential(
nn.Conv1D(
in_channels,
in_channels * 2,
kernel_size,
1,
bias_attr=bias,
padding=(kernel_size - 1) // 2, ), )
self.upsample = nn.Upsample(
scale_factor=upsample_factor, mode=upsample_mode)
[docs] def forward(self, x, c):
"""Calculate forward propagation.
Args:
x (Tensor):
Input tensor (B, in_channels, T).
c (Tensor):
Auxiliary input tensor (B, aux_channels, T).
Returns:
Tensor:
Output tensor (B, in_channels, T * upsample_factor).
Tensor:
Upsampled aux tensor (B, in_channels, T * upsample_factor).
"""
x = self.norm(x)
# 'bilinear', 'bicubic' and 'nearest' only support 4-D tensor.
c = self.upsample(c.unsqueeze(-1))
c = c[:, :, :, 0]
c = self.aux_conv(c)
cg = self.gated_conv(c)
cg1, cg2 = cg.split(2, axis=1)
# 'bilinear', 'bicubic' and 'nearest' only support 4-D tensor.
y = cg1 * self.upsample(x.unsqueeze(-1))[:, :, :, 0] + cg2
return y, c
[docs]class TADEResBlock(nn.Layer):
"""TADEResBlock module."""
def __init__(
self,
in_channels: int=64,
aux_channels: int=80,
kernel_size: int=9,
dilation: int=2,
bias: bool=True,
upsample_factor: int=2,
# this is a diff in paddle, the mode only can be "linear" when input is 3D
upsample_mode: str="nearest",
gated_function: str="softmax", ):
"""Initialize TADEResBlock module."""
super().__init__()
self.tade1 = TADELayer(
in_channels=in_channels,
aux_channels=aux_channels,
kernel_size=kernel_size,
bias=bias,
upsample_factor=1,
upsample_mode=upsample_mode, )
self.gated_conv1 = nn.Conv1D(
in_channels,
in_channels * 2,
kernel_size,
1,
bias_attr=bias,
padding=(kernel_size - 1) // 2, )
self.tade2 = TADELayer(
in_channels=in_channels,
aux_channels=in_channels,
kernel_size=kernel_size,
bias=bias,
upsample_factor=upsample_factor,
upsample_mode=upsample_mode, )
self.gated_conv2 = nn.Conv1D(
in_channels,
in_channels * 2,
kernel_size,
1,
bias_attr=bias,
dilation=dilation,
padding=(kernel_size - 1) // 2 * dilation, )
self.upsample = nn.Upsample(
scale_factor=upsample_factor, mode=upsample_mode)
if gated_function == "softmax":
self.gated_function = partial(F.softmax, axis=1)
elif gated_function == "sigmoid":
self.gated_function = F.sigmoid
else:
raise ValueError(f"{gated_function} is not supported.")
[docs] def forward(self, x, c):
"""Calculate forward propagation.
Args:
x (Tensor):
Input tensor (B, in_channels, T).
c (Tensor):
Auxiliary input tensor (B, aux_channels, T).
Returns:
Tensor:
Output tensor (B, in_channels, T * upsample_factor).
Tensor:
Upsampled auxirialy tensor (B, in_channels, T * upsample_factor).
"""
residual = x
x, c = self.tade1(x, c)
x = self.gated_conv1(x)
xa, xb = x.split(2, axis=1)
x = self.gated_function(xa) * F.tanh(xb)
x, c = self.tade2(x, c)
x = self.gated_conv2(x)
xa, xb = x.split(2, axis=1)
x = self.gated_function(xa) * F.tanh(xb)
# 'bilinear', 'bicubic' and 'nearest' only support 4-D tensor.
return self.upsample(residual.unsqueeze(-1))[:, :, :, 0] + x, c