(
self,
input_channels: int,
vocos_dim: int,
vocos_intermediate_dim: int,
vocos_num_layers: int,
out_channels: int,
sample_ratios: List[int] = [1, 1],
)
| 27 | """Encoder module with convnext and downsampling blocks""" |
| 28 | |
| 29 | def __init__( |
| 30 | self, |
| 31 | input_channels: int, |
| 32 | vocos_dim: int, |
| 33 | vocos_intermediate_dim: int, |
| 34 | vocos_num_layers: int, |
| 35 | out_channels: int, |
| 36 | sample_ratios: List[int] = [1, 1], |
| 37 | ): |
| 38 | super().__init__() |
| 39 | """ |
| 40 | Encoder module with VocosBackbone and sampling blocks. |
| 41 | |
| 42 | Args: |
| 43 | sample_ratios (List[int]): sample ratios |
| 44 | example: [2, 2] means downsample by 2x and then upsample by 2x |
| 45 | """ |
| 46 | self.encoder = VocosBackbone( |
| 47 | input_channels=input_channels, |
| 48 | dim=vocos_dim, |
| 49 | intermediate_dim=vocos_intermediate_dim, |
| 50 | num_layers=vocos_num_layers, |
| 51 | condition_dim=None, |
| 52 | ) |
| 53 | |
| 54 | modules = [ |
| 55 | nn.Sequential( |
| 56 | SamplingBlock( |
| 57 | dim=vocos_dim, |
| 58 | groups=vocos_dim, |
| 59 | downsample_scale=ratio, |
| 60 | ), |
| 61 | VocosBackbone( |
| 62 | input_channels=vocos_dim, |
| 63 | dim=vocos_dim, |
| 64 | intermediate_dim=vocos_intermediate_dim, |
| 65 | num_layers=2, |
| 66 | condition_dim=None, |
| 67 | ), |
| 68 | ) |
| 69 | for ratio in sample_ratios |
| 70 | ] |
| 71 | |
| 72 | self.downsample = nn.Sequential(*modules) |
| 73 | |
| 74 | self.project = nn.Linear(vocos_dim, out_channels) |
| 75 | |
| 76 | def forward(self, x: torch.Tensor, *args): |
| 77 | """ |
nothing calls this directly
no test coverage detected