| 464 | return x |
| 465 | |
| 466 | class AttnBlock(nn.Module): |
| 467 | def __init__(self, |
| 468 | in_channels |
| 469 | ): |
| 470 | super().__init__() |
| 471 | |
| 472 | self.norm = BaseGroupNorm(num_groups=32, num_channels=in_channels) |
| 473 | self.q = CausalConvChannelLast(in_channels, in_channels, kernel_size=1) |
| 474 | self.k = CausalConvChannelLast(in_channels, in_channels, kernel_size=1) |
| 475 | self.v = CausalConvChannelLast(in_channels, in_channels, kernel_size=1) |
| 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() |
| 493 | h = self.attention(x) |
| 494 | x = self.proj_out(h, residual=x) |
| 495 | x = x.permute(0,4,1,2,3) |
| 496 | return x |
| 497 | |
| 498 | class Resnet3DBlock(nn.Module): |
| 499 | def __init__(self, |