| 189 | |
| 190 | |
| 191 | class ModulatedConv2d(nn.Module): |
| 192 | def __init__(self, in_channel, out_channel, kernel_size, style_dim, demodulate=True, upsample=False, |
| 193 | downsample=False, blur_kernel=[1, 3, 3, 1], ): |
| 194 | super().__init__() |
| 195 | |
| 196 | self.eps = 1e-8 |
| 197 | self.kernel_size = kernel_size |
| 198 | self.in_channel = in_channel |
| 199 | self.out_channel = out_channel |
| 200 | self.upsample = upsample |
| 201 | self.downsample = downsample |
| 202 | |
| 203 | if upsample: |
| 204 | factor = 2 |
| 205 | p = (len(blur_kernel) - factor) - (kernel_size - 1) |
| 206 | pad0 = (p + 1) // 2 + factor - 1 |
| 207 | pad1 = p // 2 + 1 |
| 208 | |
| 209 | self.blur = Blur(blur_kernel, pad=(pad0, pad1), upsample_factor=factor) |
| 210 | |
| 211 | if downsample: |
| 212 | factor = 2 |
| 213 | p = (len(blur_kernel) - factor) + (kernel_size - 1) |
| 214 | pad0 = (p + 1) // 2 |
| 215 | pad1 = p // 2 |
| 216 | |
| 217 | self.blur = Blur(blur_kernel, pad=(pad0, pad1)) |
| 218 | |
| 219 | fan_in = in_channel * kernel_size ** 2 |
| 220 | self.scale = 1 / math.sqrt(fan_in) |
| 221 | self.padding = kernel_size // 2 |
| 222 | |
| 223 | self.weight = nn.Parameter(torch.randn(1, out_channel, in_channel, kernel_size, kernel_size)) |
| 224 | |
| 225 | self.modulation = EqualLinear(style_dim, in_channel, bias_init=1) |
| 226 | self.demodulate = demodulate |
| 227 | |
| 228 | def __repr__(self): |
| 229 | return ( |
| 230 | f'{self.__class__.__name__}({self.in_channel}, {self.out_channel}, {self.kernel_size}, ' |
| 231 | f'upsample={self.upsample}, downsample={self.downsample})' |
| 232 | ) |
| 233 | |
| 234 | def forward(self, input, style): |
| 235 | batch, in_channel, height, width = input.shape |
| 236 | |
| 237 | style = self.modulation(style).view(batch, 1, in_channel, 1, 1) |
| 238 | weight = self.scale * self.weight * style |
| 239 | |
| 240 | if self.demodulate: |
| 241 | demod = torch.rsqrt(weight.pow(2).sum([2, 3, 4]) + 1e-8) |
| 242 | weight = weight * demod.view(batch, self.out_channel, 1, 1, 1) |
| 243 | |
| 244 | weight = weight.view(batch * self.out_channel, in_channel, self.kernel_size, self.kernel_size) |
| 245 | |
| 246 | if self.upsample: |
| 247 | input = input.view(1, batch * in_channel, height, width) |
| 248 | weight = weight.view(batch, self.out_channel, in_channel, self.kernel_size, self.kernel_size) |