| 27 | |
| 28 | |
| 29 | class DecoderBlock(nn.Module): |
| 30 | def __init__( |
| 31 | self, |
| 32 | input_dim: int = 16, |
| 33 | output_dim: int = 8, |
| 34 | kernel_size: int = 2, |
| 35 | stride: int = 1, |
| 36 | ): |
| 37 | super().__init__() |
| 38 | self.block = nn.Sequential( |
| 39 | Snake1d(input_dim), |
| 40 | WNConvTranspose1d( |
| 41 | input_dim, |
| 42 | output_dim, |
| 43 | kernel_size=kernel_size, |
| 44 | stride=stride, |
| 45 | padding=(kernel_size - stride) // 2, |
| 46 | ), |
| 47 | ResidualUnit(output_dim, dilation=1), |
| 48 | ResidualUnit(output_dim, dilation=3), |
| 49 | ResidualUnit(output_dim, dilation=9), |
| 50 | ) |
| 51 | |
| 52 | def forward(self, x): |
| 53 | return self.block(x) |
| 54 | |
| 55 | |
| 56 | class WaveGenerator(nn.Module): |