| 449 | |
| 450 | |
| 451 | class SpatialNorm3D(nn.Module): |
| 452 | def __init__( |
| 453 | self, |
| 454 | f_channels, |
| 455 | zq_channels, |
| 456 | freeze_norm_layer=False, |
| 457 | add_conv=False, |
| 458 | pad_mode="constant", |
| 459 | gather=False, |
| 460 | **norm_layer_params, |
| 461 | ): |
| 462 | super().__init__() |
| 463 | if gather: |
| 464 | self.norm_layer = ContextParallelGroupNorm(num_channels=f_channels, **norm_layer_params) |
| 465 | else: |
| 466 | self.norm_layer = torch.nn.GroupNorm(num_channels=f_channels, **norm_layer_params) |
| 467 | # self.norm_layer = norm_layer(num_channels=f_channels, **norm_layer_params) |
| 468 | if freeze_norm_layer: |
| 469 | for p in self.norm_layer.parameters: |
| 470 | p.requires_grad = False |
| 471 | |
| 472 | self.add_conv = add_conv |
| 473 | if add_conv: |
| 474 | self.conv = ContextParallelCausalConv3d( |
| 475 | chan_in=zq_channels, |
| 476 | chan_out=zq_channels, |
| 477 | kernel_size=3, |
| 478 | ) |
| 479 | |
| 480 | self.conv_y = ContextParallelCausalConv3d( |
| 481 | chan_in=zq_channels, |
| 482 | chan_out=f_channels, |
| 483 | kernel_size=1, |
| 484 | ) |
| 485 | self.conv_b = ContextParallelCausalConv3d( |
| 486 | chan_in=zq_channels, |
| 487 | chan_out=f_channels, |
| 488 | kernel_size=1, |
| 489 | ) |
| 490 | |
| 491 | def forward(self, f, zq, clear_fake_cp_cache=True): |
| 492 | if f.shape[2] > 1 and f.shape[2] % 2 == 1: |
| 493 | f_first, f_rest = f[:, :, :1], f[:, :, 1:] |
| 494 | f_first_size, f_rest_size = f_first.shape[-3:], f_rest.shape[-3:] |
| 495 | zq_first, zq_rest = zq[:, :, :1], zq[:, :, 1:] |
| 496 | zq_first = torch.nn.functional.interpolate(zq_first, size=f_first_size, mode="nearest") |
| 497 | zq_rest = torch.nn.functional.interpolate(zq_rest, size=f_rest_size, mode="nearest") |
| 498 | zq = torch.cat([zq_first, zq_rest], dim=2) |
| 499 | else: |
| 500 | zq = torch.nn.functional.interpolate(zq, size=f.shape[-3:], mode="nearest") |
| 501 | |
| 502 | if self.add_conv: |
| 503 | zq = self.conv(zq, clear_cache=clear_fake_cp_cache) |
| 504 | |
| 505 | # f = conv_gather_from_context_parallel_region(f, dim=2, kernel_size=1) |
| 506 | norm_f = self.norm_layer(f) |
| 507 | # norm_f = conv_scatter_to_context_parallel_region(norm_f, dim=2, kernel_size=1) |
| 508 | |