| 692 | |
| 693 | |
| 694 | class SpatialDownsample2x(Module): |
| 695 | def __init__(self, dim, dim_out=None, kernel_size=3, antialias=False): |
| 696 | super().__init__() |
| 697 | dim_out = default(dim_out, dim) |
| 698 | self.maybe_blur = Blur() if antialias else identity |
| 699 | self.conv = nn.Conv2d(dim, dim_out, kernel_size, stride=2, padding=kernel_size // 2) |
| 700 | |
| 701 | def forward(self, x): |
| 702 | x = self.maybe_blur(x, space_only=True) |
| 703 | |
| 704 | x = rearrange(x, "b c t h w -> b t c h w") |
| 705 | x, ps = pack_one(x, "* c h w") |
| 706 | |
| 707 | out = self.conv(x) |
| 708 | |
| 709 | out = unpack_one(out, ps, "* c h w") |
| 710 | out = rearrange(out, "b t c h w -> b c t h w") |
| 711 | return out |
| 712 | |
| 713 | |
| 714 | class TimeDownsample2x(Module): |