| 476 | self.proj_out = CausalConvChannelLast(in_channels, in_channels, kernel_size=1) |
| 477 | |
| 478 | def attention(self, x, is_init=True): |
| 479 | x = self.norm(x, act_silu=False, channel_last=True) |
| 480 | q = self.q(x, is_init) |
| 481 | k = self.k(x, is_init) |
| 482 | v = self.v(x, is_init) |
| 483 | |
| 484 | b, t, h, w, c = q.shape |
| 485 | q, k, v = map(lambda x: rearrange(x, "b t h w c -> b 1 (t h w) c"), (q, k, v)) |
| 486 | x = nn.functional.scaled_dot_product_attention(q, k, v, is_causal=True) |
| 487 | x = rearrange(x, "b 1 (t h w) c -> b t h w c", t=t, h=h, w=w) |
| 488 | |
| 489 | return x |
| 490 | |
| 491 | def forward(self, x): |
| 492 | x = x.permute(0,2,3,4,1).contiguous() |