| 223 | |
| 224 | # Streaming Vocos backbone based on Transformer layers |
| 225 | class CausalVocosBackbone(nn.Module): |
| 226 | def __init__( |
| 227 | self, |
| 228 | embed_dim: int = 1024, |
| 229 | num_layers: int = 12, |
| 230 | num_heads: int = 16, |
| 231 | dropout: float = 0.1, |
| 232 | ): |
| 233 | super().__init__() |
| 234 | self.in_proj = CausalConv1d(embed_dim, embed_dim, kernel_size=7) |
| 235 | self.prior_net = nn.Sequential( |
| 236 | CausalResnetBlock(embed_dim, embed_dim, dropout=dropout), |
| 237 | CausalResnetBlock(embed_dim, embed_dim, dropout=dropout), |
| 238 | ) |
| 239 | self.transformers = nn.ModuleList( |
| 240 | [WhisperEncoderLayer(embed_dim, num_heads) for _ in range(num_layers)] |
| 241 | ) |
| 242 | self.post_net = nn.Sequential( |
| 243 | CausalResnetBlock(embed_dim, embed_dim, dropout=dropout), |
| 244 | CausalResnetBlock(embed_dim, embed_dim, dropout=dropout), |
| 245 | ) |
| 246 | self.final_norm = nn.LayerNorm(embed_dim, eps=1e-6) |
| 247 | |
| 248 | def forward( |
| 249 | self, |
| 250 | x: torch.Tensor, |
| 251 | x_lens: torch.Tensor, |
| 252 | ): |
| 253 | """ |
| 254 | Args: |
| 255 | x: shape (b, t, c) |
| 256 | x_lens: shape (b,) |
| 257 | """ |
| 258 | x = x.transpose(1, 2) |
| 259 | x = self.in_proj(x) |
| 260 | x = self.prior_net(x) |
| 261 | x = x.transpose(1, 2) |
| 262 | |
| 263 | # NOTE(sfy): We have no padding in training, so safe for sdpa attention, no Nan. |
| 264 | # Also, 1 token(12.5Hz) -> 4 latents(50Hz) -> 8 latents(100Hz), |
| 265 | # so we design a 8 block causal attention mask instead of fully causal to improve performance |
| 266 | attention_mask = make_block_causal_mask(x_lens, chunk_size=8) |
| 267 | for layer in self.transformers: |
| 268 | x = layer(x, attention_mask) |
| 269 | |
| 270 | x = x.transpose(1, 2) |
| 271 | x = self.post_net(x) |
| 272 | x = x.transpose(1, 2) |
| 273 | x = self.final_norm(x) |
| 274 | return x |
| 275 | |
| 276 | def forward_chunk( |
| 277 | self, |
| 278 | x: torch.Tensor, |
| 279 | conv_cache1: torch.Tensor = None, |
| 280 | conv_cache2: torch.Tensor = None, |
| 281 | kv_cache: torch.Tensor = None, |
| 282 | ): |