Spatially conditioned normalization as defined in https://huggingface.co/papers/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
| 4177 | |
| 4178 | |
| 4179 | class SpatialNorm(nn.Module): |
| 4180 | """ |
| 4181 | Spatially conditioned normalization as defined in https://huggingface.co/papers/2209.09002. |
| 4182 | |
| 4183 | Args: |
| 4184 | f_channels (`int`): |
| 4185 | The number of channels for input to group normalization layer, and output of the spatial norm layer. |
| 4186 | zq_channels (`int`): |
| 4187 | The number of channels for the quantized vector as described in the paper. |
| 4188 | """ |
| 4189 | |
| 4190 | def __init__( |
| 4191 | self, |
| 4192 | f_channels: int, |
| 4193 | zq_channels: int, |
| 4194 | ): |
| 4195 | super().__init__() |
| 4196 | self.norm_layer = nn.GroupNorm(num_channels=f_channels, num_groups=32, eps=1e-6, affine=True) |
| 4197 | self.conv_y = nn.Conv2d(zq_channels, f_channels, kernel_size=1, stride=1, padding=0) |
| 4198 | self.conv_b = nn.Conv2d(zq_channels, f_channels, kernel_size=1, stride=1, padding=0) |
| 4199 | |
| 4200 | def forward(self, f: torch.Tensor, zq: torch.Tensor) -> torch.Tensor: |
| 4201 | f_size = f.shape[-2:] |
| 4202 | zq = F.interpolate(zq, size=f_size, mode="nearest") |
| 4203 | norm_f = self.norm_layer(f) |
| 4204 | new_f = norm_f * self.conv_y(zq) + self.conv_b(zq) |
| 4205 | return new_f |
| 4206 | |
| 4207 | |
| 4208 | class IPAdapterAttnProcessor(nn.Module): |