(
self,
input_channels: int,
vocos_dim: int,
vocos_intermediate_dim: int,
vocos_num_layers: int,
out_channels: int,
condition_dim: int = None,
sample_ratios: List[int] = [1, 1],
use_tanh_at_final: bool = False,
)
| 32 | """ |
| 33 | |
| 34 | def __init__( |
| 35 | self, |
| 36 | input_channels: int, |
| 37 | vocos_dim: int, |
| 38 | vocos_intermediate_dim: int, |
| 39 | vocos_num_layers: int, |
| 40 | out_channels: int, |
| 41 | condition_dim: int = None, |
| 42 | sample_ratios: List[int] = [1, 1], |
| 43 | use_tanh_at_final: bool = False, |
| 44 | ): |
| 45 | super().__init__() |
| 46 | |
| 47 | self.linear_pre = nn.Linear(input_channels, vocos_dim) |
| 48 | modules = [ |
| 49 | nn.Sequential( |
| 50 | SamplingBlock( |
| 51 | dim=vocos_dim, |
| 52 | groups=vocos_dim, |
| 53 | upsample_scale=ratio, |
| 54 | ), |
| 55 | VocosBackbone( |
| 56 | input_channels=vocos_dim, |
| 57 | dim=vocos_dim, |
| 58 | intermediate_dim=vocos_intermediate_dim, |
| 59 | num_layers=2, |
| 60 | condition_dim=None, |
| 61 | ), |
| 62 | ) |
| 63 | for ratio in sample_ratios |
| 64 | ] |
| 65 | |
| 66 | self.downsample = nn.Sequential(*modules) |
| 67 | |
| 68 | self.vocos_backbone = VocosBackbone( |
| 69 | input_channels=vocos_dim, |
| 70 | dim=vocos_dim, |
| 71 | intermediate_dim=vocos_intermediate_dim, |
| 72 | num_layers=vocos_num_layers, |
| 73 | condition_dim=condition_dim, |
| 74 | ) |
| 75 | self.linear = nn.Linear(vocos_dim, out_channels) |
| 76 | self.use_tanh_at_final = use_tanh_at_final |
| 77 | |
| 78 | def forward(self, x: torch.Tensor, c: torch.Tensor = None): |
| 79 | """encoder forward. |
nothing calls this directly
no test coverage detected