| 712 | |
| 713 | |
| 714 | class TimeDownsample2x(Module): |
| 715 | def __init__(self, dim, dim_out=None, kernel_size=3, antialias=False): |
| 716 | super().__init__() |
| 717 | dim_out = default(dim_out, dim) |
| 718 | self.maybe_blur = Blur() if antialias else identity |
| 719 | self.time_causal_padding = (kernel_size - 1, 0) |
| 720 | self.conv = nn.Conv1d(dim, dim_out, kernel_size, stride=2) |
| 721 | |
| 722 | def forward(self, x): |
| 723 | x = self.maybe_blur(x, time_only=True) |
| 724 | |
| 725 | x = rearrange(x, "b c t h w -> b h w c t") |
| 726 | x, ps = pack_one(x, "* c t") |
| 727 | |
| 728 | x = F.pad(x, self.time_causal_padding) |
| 729 | out = self.conv(x) |
| 730 | |
| 731 | out = unpack_one(out, ps, "* c t") |
| 732 | out = rearrange(out, "b h w c t -> b c t h w") |
| 733 | return out |
| 734 | |
| 735 | |
| 736 | # depth to space upsamples |