| import torch | |
| from torch import nn | |
| from modules.commons.conv import ConditionalConvBlocks | |
| from modules.commons.wavenet import WN | |
| class FlipLayer(nn.Module): | |
| def forward(self, x, *args, **kwargs): | |
| x = torch.flip(x, [1]) | |
| return x | |
| class CouplingLayer(nn.Module): | |
| def __init__(self, c_in, hidden_size, kernel_size, n_layers, p_dropout=0, c_in_g=0, nn_type='wn'): | |
| super().__init__() | |
| self.channels = c_in | |
| self.hidden_size = hidden_size | |
| self.kernel_size = kernel_size | |
| self.n_layers = n_layers | |
| self.c_half = c_in // 2 | |
| self.pre = nn.Conv1d(self.c_half, hidden_size, 1) | |
| if nn_type == 'wn': | |
| self.enc = WN(hidden_size, kernel_size, 1, n_layers, p_dropout=p_dropout, | |
| c_cond=c_in_g) | |
| elif nn_type == 'conv': | |
| self.enc = ConditionalConvBlocks( | |
| hidden_size, c_in_g, hidden_size, None, kernel_size, | |
| layers_in_block=1, is_BTC=False, num_layers=n_layers) | |
| self.post = nn.Conv1d(hidden_size, self.c_half, 1) | |
| def forward(self, x, nonpadding, cond=None, reverse=False): | |
| x0, x1 = x[:, :self.c_half], x[:, self.c_half:] | |
| x_ = self.pre(x0) * nonpadding | |
| x_ = self.enc(x_, nonpadding=nonpadding, cond=cond) | |
| m = self.post(x_) | |
| x1 = m + x1 if not reverse else x1 - m | |
| x = torch.cat([x0, x1], 1) | |
| return x * nonpadding | |
| class ResFlow(nn.Module): | |
| def __init__(self, | |
| c_in, | |
| hidden_size, | |
| kernel_size, | |
| n_flow_layers, | |
| n_flow_steps=4, | |
| c_cond=0, | |
| nn_type='wn'): | |
| super().__init__() | |
| self.flows = nn.ModuleList() | |
| for i in range(n_flow_steps): | |
| self.flows.append( | |
| CouplingLayer(c_in, hidden_size, kernel_size, n_flow_layers, c_in_g=c_cond, nn_type=nn_type)) | |
| self.flows.append(FlipLayer()) | |
| def forward(self, x, nonpadding, cond=None, reverse=False): | |
| for flow in (self.flows if not reverse else reversed(self.flows)): | |
| x = flow(x, nonpadding, cond=cond, reverse=reverse) | |
| return x | |