(self, x)
| 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): |
no test coverage detected