| 56 | |
| 57 | |
| 58 | class ResidualBlock(nn.Module): |
| 59 | def __init__(self, encoder_hidden, residual_channels, dilation): |
| 60 | super().__init__() |
| 61 | self.dilated_conv = Conv1d(residual_channels, 2 * residual_channels, 3, padding=dilation, dilation=dilation) |
| 62 | self.diffusion_projection = Linear(residual_channels, residual_channels) |
| 63 | self.conditioner_projection = Conv1d(encoder_hidden, 2 * residual_channels, 1) |
| 64 | self.output_projection = Conv1d(residual_channels, 2 * residual_channels, 1) |
| 65 | |
| 66 | def forward(self, x, conditioner, diffusion_step): |
| 67 | diffusion_step = self.diffusion_projection(diffusion_step).unsqueeze(-1) |
| 68 | conditioner = self.conditioner_projection(conditioner) |
| 69 | y = x + diffusion_step |
| 70 | |
| 71 | y = self.dilated_conv(y) + conditioner |
| 72 | |
| 73 | gate, filter = torch.chunk(y, 2, dim=1) |
| 74 | y = torch.sigmoid(gate) * torch.tanh(filter) |
| 75 | |
| 76 | y = self.output_projection(y) |
| 77 | residual, skip = torch.chunk(y, 2, dim=1) |
| 78 | return (x + residual) / sqrt(2.0), skip |
| 79 | |
| 80 | |
| 81 | class DiffNet(nn.Module): |