| 529 | |
| 530 | |
| 531 | class Upsample3D(nn.Module): |
| 532 | def __init__( |
| 533 | self, |
| 534 | in_channels, |
| 535 | with_conv, |
| 536 | compress_time=False, |
| 537 | ): |
| 538 | super().__init__() |
| 539 | self.with_conv = with_conv |
| 540 | if self.with_conv: |
| 541 | self.conv = torch.nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1) |
| 542 | self.compress_time = compress_time |
| 543 | |
| 544 | def forward(self, x): |
| 545 | if self.compress_time and x.shape[2] > 1: |
| 546 | if x.shape[2] % 2 == 1: |
| 547 | # split first frame |
| 548 | x_first, x_rest = x[:, :, 0], x[:, :, 1:] |
| 549 | |
| 550 | x_first = torch.nn.functional.interpolate(x_first, scale_factor=2.0, mode="nearest") |
| 551 | x_rest = torch.nn.functional.interpolate(x_rest, scale_factor=2.0, mode="nearest") |
| 552 | x = torch.cat([x_first[:, :, None, :, :], x_rest], dim=2) |
| 553 | else: |
| 554 | x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest") |
| 555 | |
| 556 | else: |
| 557 | # only interpolate 2D |
| 558 | t = x.shape[2] |
| 559 | x = rearrange(x, "b c t h w -> (b t) c h w") |
| 560 | x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest") |
| 561 | x = rearrange(x, "(b t) c h w -> b c t h w", t=t) |
| 562 | |
| 563 | if self.with_conv: |
| 564 | t = x.shape[2] |
| 565 | x = rearrange(x, "b c t h w -> (b t) c h w") |
| 566 | x = self.conv(x) |
| 567 | x = rearrange(x, "(b t) c h w -> b c t h w", t=t) |
| 568 | return x |
| 569 | |
| 570 | |
| 571 | class DownSample3D(nn.Module): |