| 398 | |
| 399 | class CouplingBlock(nn.Module): |
| 400 | def __init__(self, in_channels, hidden_channels, kernel_size, dilation_rate, n_layers, |
| 401 | gin_channels=0, p_dropout=0, sigmoid_scale=False, |
| 402 | share_cond_layers=False, wn=None): |
| 403 | super().__init__() |
| 404 | self.in_channels = in_channels |
| 405 | self.hidden_channels = hidden_channels |
| 406 | self.kernel_size = kernel_size |
| 407 | self.dilation_rate = dilation_rate |
| 408 | self.n_layers = n_layers |
| 409 | self.gin_channels = gin_channels |
| 410 | self.p_dropout = p_dropout |
| 411 | self.sigmoid_scale = sigmoid_scale |
| 412 | |
| 413 | start = torch.nn.Conv1d(in_channels // 2, hidden_channels, 1) |
| 414 | start = torch.nn.utils.weight_norm(start) |
| 415 | self.start = start |
| 416 | # Initializing last layer to 0 makes the affine coupling layers |
| 417 | # do nothing at first. This helps with training stability |
| 418 | end = torch.nn.Conv1d(hidden_channels, in_channels, 1) |
| 419 | end.weight.data.zero_() |
| 420 | end.bias.data.zero_() |
| 421 | self.end = end |
| 422 | self.wn = WN(hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels, |
| 423 | p_dropout, share_cond_layers) |
| 424 | if wn is not None: |
| 425 | self.wn.in_layers = wn.in_layers |
| 426 | self.wn.res_skip_layers = wn.res_skip_layers |
| 427 | |
| 428 | def forward(self, x, x_mask=None, reverse=False, g=None, **kwargs): |
| 429 | if x_mask is None: |