(
self,
embed_dim: int = 1024,
num_layers: int = 12,
num_heads: int = 16,
dropout: float = 0.1,
)
| 174 | # Nonstreaming Vocos backbone based on Transformer layers |
| 175 | class VocosBackbone(nn.Module): |
| 176 | def __init__( |
| 177 | self, |
| 178 | embed_dim: int = 1024, |
| 179 | num_layers: int = 12, |
| 180 | num_heads: int = 16, |
| 181 | dropout: float = 0.1, |
| 182 | ): |
| 183 | super().__init__() |
| 184 | self.in_proj = nn.Conv1d(embed_dim, embed_dim, kernel_size=7, padding=3) |
| 185 | self.prior_net = nn.Sequential( |
| 186 | ResnetBlock(embed_dim, embed_dim, dropout=dropout), |
| 187 | ResnetBlock(embed_dim, embed_dim, dropout=dropout), |
| 188 | ) |
| 189 | self.transformers = nn.ModuleList( |
| 190 | [WhisperEncoderLayer(embed_dim, num_heads) for _ in range(num_layers)] |
| 191 | ) |
| 192 | self.post_net = nn.Sequential( |
| 193 | ResnetBlock(embed_dim, embed_dim, dropout=dropout), |
| 194 | ResnetBlock(embed_dim, embed_dim, dropout=dropout), |
| 195 | ) |
| 196 | self.final_norm = nn.LayerNorm(embed_dim, eps=1e-6) |
| 197 | |
| 198 | def forward( |
| 199 | self, |
nothing calls this directly
no test coverage detected