| 296 | |
| 297 | |
| 298 | class StyledConv(nn.Module): |
| 299 | def __init__(self, in_channel, out_channel, kernel_size, style_dim, upsample=False, blur_kernel=[1, 3, 3, 1], |
| 300 | demodulate=True): |
| 301 | super().__init__() |
| 302 | |
| 303 | self.conv = ModulatedConv2d( |
| 304 | in_channel, |
| 305 | out_channel, |
| 306 | kernel_size, |
| 307 | style_dim, |
| 308 | upsample=upsample, |
| 309 | blur_kernel=blur_kernel, |
| 310 | demodulate=demodulate, |
| 311 | ) |
| 312 | |
| 313 | self.noise = NoiseInjection() |
| 314 | self.activate = FusedLeakyReLU(out_channel) |
| 315 | |
| 316 | def forward(self, input, style, noise=None): |
| 317 | out = self.conv(input, style) |
| 318 | out = self.noise(out, noise=noise) |
| 319 | out = self.activate(out) |
| 320 | |
| 321 | return out |
| 322 | |
| 323 | |
| 324 | class ConvLayer(nn.Sequential): |