| 433 | |
| 434 | |
| 435 | class TimeAttention(Attention): |
| 436 | def forward(self, x, *args, **kwargs): |
| 437 | x = rearrange(x, "b c t h w -> b h w t c") |
| 438 | x, batch_ps = pack_one(x, "* t c") |
| 439 | |
| 440 | x = super().forward(x, *args, **kwargs) |
| 441 | |
| 442 | x = unpack_one(x, batch_ps, "* t c") |
| 443 | return rearrange(x, "b h w t c -> b c t h w") |
| 444 | |
| 445 | |
| 446 | class GEGLU(Module): |