Spatially conditioned normalization as defined in https://arxiv.org/abs/2209.09002. Args: f_channels (`int`): The number of channels for input to group normalization layer, and output of the spatial norm layer. zq_channels (`int`): The number of chan
| 1673 | |
| 1674 | |
| 1675 | class SpatialNorm(nn.Module): |
| 1676 | """ |
| 1677 | Spatially conditioned normalization as defined in https://arxiv.org/abs/2209.09002. |
| 1678 | |
| 1679 | Args: |
| 1680 | f_channels (`int`): |
| 1681 | The number of channels for input to group normalization layer, and output of the spatial norm layer. |
| 1682 | zq_channels (`int`): |
| 1683 | The number of channels for the quantized vector as described in the paper. |
| 1684 | """ |
| 1685 | |
| 1686 | def __init__( |
| 1687 | self, |
| 1688 | f_channels: int, |
| 1689 | zq_channels: int, |
| 1690 | ): |
| 1691 | super().__init__() |
| 1692 | self.norm_layer = nn.GroupNorm(num_channels=f_channels, num_groups=32, eps=1e-6, affine=True) |
| 1693 | self.conv_y = nn.Conv2d(zq_channels, f_channels, kernel_size=1, stride=1, padding=0) |
| 1694 | self.conv_b = nn.Conv2d(zq_channels, f_channels, kernel_size=1, stride=1, padding=0) |
| 1695 | |
| 1696 | def forward(self, f: torch.FloatTensor, zq: torch.FloatTensor) -> torch.FloatTensor: |
| 1697 | f_size = f.shape[-2:] |
| 1698 | zq = F.interpolate(zq, size=f_size, mode="nearest") |
| 1699 | norm_f = self.norm_layer(f) |
| 1700 | new_f = norm_f * self.conv_y(zq) + self.conv_b(zq) |
| 1701 | return new_f |
| 1702 | |
| 1703 | |
| 1704 | ## Deprecated |