| import torch.nn as nn |
| import module as modules |
|
|
|
|
| class ResidualCouplingBlock(nn.Module): |
| def __init__( |
| self, |
| channels, |
| hidden_channels, |
| kernel_size, |
| dilation_rate, |
| n_layers, |
| n_flows=4, |
| gin_channels=0, |
| ): |
| super().__init__() |
| self.channels = channels |
| self.hidden_channels = hidden_channels |
| self.kernel_size = kernel_size |
| self.dilation_rate = dilation_rate |
| self.n_layers = n_layers |
| self.n_flows = n_flows |
| self.gin_channels = gin_channels |
|
|
| self.flows = nn.ModuleList() |
| for i in range(n_flows): |
| self.flows.append( |
| modules.ResidualCouplingLayer( |
| channels, |
| hidden_channels, |
| kernel_size, |
| dilation_rate, |
| n_layers, |
| gin_channels=gin_channels, |
| mean_only=True, |
| ) |
| ) |
| self.flows.append(modules.Flip()) |
|
|
| def forward(self, x, x_mask, g=None, reverse=False): |
| if not reverse: |
| for flow in self.flows: |
| x, _ = flow(x, x_mask, g=g, reverse=reverse) |
| else: |
| for flow in reversed(self.flows): |
| x = flow(x, x_mask, g=g, reverse=reverse) |
| return x |
|
|
