| 697 | |
| 698 | |
| 699 | class EncoderBlock(nn.Module): |
| 700 | def __init__(self, dim: int = 16, stride: int = 1): |
| 701 | super().__init__() |
| 702 | self.block = nn.Sequential( |
| 703 | ResidualUnit(dim // 2, dilation=1), |
| 704 | ResidualUnit(dim // 2, dilation=3), |
| 705 | ResidualUnit(dim // 2, dilation=9), |
| 706 | Snake1d(dim // 2), |
| 707 | WNConv1d( |
| 708 | dim // 2, |
| 709 | dim, |
| 710 | kernel_size=2 * stride, |
| 711 | stride=stride, |
| 712 | padding=math.ceil(stride / 2), |
| 713 | ), |
| 714 | ) |
| 715 | |
| 716 | def forward(self, x): |
| 717 | return self.block(x) |
| 718 | |
| 719 | |
| 720 | class Encoder(nn.Module): |