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
| 3578 | |
| 3579 | |
| 3580 | class SpatialNorm(nn.Module): |
| 3581 | """ |
| 3582 | Spatially conditioned normalization as defined in https://arxiv.org/abs/2209.09002. |
| 3583 | |
| 3584 | Args: |
| 3585 | f_channels (`int`): |
| 3586 | The number of channels for input to group normalization layer, and output of the spatial norm layer. |
| 3587 | zq_channels (`int`): |
| 3588 | The number of channels for the quantized vector as described in the paper. |
| 3589 | """ |
| 3590 | |
| 3591 | def __init__( |
| 3592 | self, |
| 3593 | f_channels: int, |
| 3594 | zq_channels: int, |
| 3595 | ): |
| 3596 | super().__init__() |
| 3597 | self.norm_layer = nn.GroupNorm(num_channels=f_channels, num_groups=32, eps=1e-6, affine=True) |
| 3598 | self.conv_y = nn.Conv2d(zq_channels, f_channels, kernel_size=1, stride=1, padding=0) |
| 3599 | self.conv_b = nn.Conv2d(zq_channels, f_channels, kernel_size=1, stride=1, padding=0) |
| 3600 | |
| 3601 | def forward(self, f: torch.Tensor, zq: torch.Tensor) -> torch.Tensor: |
| 3602 | f_size = f.shape[-2:] |
| 3603 | zq = F.interpolate(zq, size=f_size, mode="nearest") |
| 3604 | norm_f = self.norm_layer(f) |
| 3605 | new_f = norm_f * self.conv_y(zq) + self.conv_b(zq) |
| 3606 | return new_f |
| 3607 | |
| 3608 | |
| 3609 | class IPAdapterAttnProcessor(nn.Module): |