| 431 | |
| 432 | |
| 433 | class Upsample3D(nn.Module): |
| 434 | def __init__( |
| 435 | self, |
| 436 | in_channels, |
| 437 | with_conv, |
| 438 | compress_time=False, |
| 439 | ): |
| 440 | super().__init__() |
| 441 | self.with_conv = with_conv |
| 442 | if self.with_conv: |
| 443 | self.conv = torch.nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1) |
| 444 | self.compress_time = compress_time |
| 445 | |
| 446 | def forward(self, x): |
| 447 | if self.compress_time: |
| 448 | if x.shape[2] == 1 and not _USE_CP: |
| 449 | x = torch.nn.functional.interpolate(x[:, :, 0], scale_factor=2.0, mode="nearest")[:, :, None, :, :] |
| 450 | elif get_context_parallel_rank() == 0: |
| 451 | # split first frame |
| 452 | x_first, x_rest = x[:, :, 0], x[:, :, 1:] |
| 453 | |
| 454 | x_first = torch.nn.functional.interpolate(x_first, scale_factor=2.0, mode="nearest") |
| 455 | x_rest = torch.nn.functional.interpolate(x_rest, scale_factor=2.0, mode="nearest") |
| 456 | x = torch.cat([x_first[:, :, None, :, :], x_rest], dim=2) |
| 457 | else: |
| 458 | x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest") |
| 459 | |
| 460 | else: |
| 461 | # only interpolate 2D |
| 462 | t = x.shape[2] |
| 463 | x = rearrange(x, "b c t h w -> (b t) c h w") |
| 464 | x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest") |
| 465 | x = rearrange(x, "(b t) c h w -> b c t h w", t=t) |
| 466 | |
| 467 | if self.with_conv: |
| 468 | t = x.shape[2] |
| 469 | x = rearrange(x, "b c t h w -> (b t) c h w") |
| 470 | x = self.conv(x) |
| 471 | x = rearrange(x, "(b t) c h w -> b c t h w", t=t) |
| 472 | return x |
| 473 | |
| 474 | |
| 475 | class DownSample3D(nn.Module): |