Args: x: shape (b, t, c) x_lens: shape (b,)
(
self,
x: torch.Tensor,
x_lens: torch.Tensor,
)
| 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, |
nothing calls this directly
no test coverage detected