(self, x, index, condition)
| 340 | return x |
| 341 | |
| 342 | def forward(self, x, index, condition): |
| 343 | residual = x.type(torch.float32) |
| 344 | x_ln = self.self_attn_ln(x) |
| 345 | x2 = self.window_attention(x_ln, self.self_attn, index) |
| 346 | x = residual + self.drop_path(x2) |
| 347 | |
| 348 | x = rearrange(x, 'b (p n) d -> (b p) n d', p=2) # split back to frame_0 and frame_1 |
| 349 | residual = x.type(torch.float32) |
| 350 | x_ln = self.cross_attn_ln(x) |
| 351 | x2 = self.cross_attn(x_ln, condition, condition, need_weights=False)[0] |
| 352 | x = residual + self.drop_path(x2) |
| 353 | x = rearrange(x, '(b p) n d -> b (p n) d', p=2) # combine frame_0 and frame_1 |
| 354 | |
| 355 | residual = x.type(torch.float32) |
| 356 | x_ln = self.mlp_ln(x) |
| 357 | x2 = self.mlp(x_ln) |
| 358 | x = residual + self.drop_path(x2) |
| 359 | |
| 360 | return x |
| 361 | |
| 362 | class SinusoidalPositionalEncoding(nn.Module): |
| 363 | def __init__(self, max_len, d_model): |
nothing calls this directly
no test coverage detected