| 72 | |
| 73 | |
| 74 | class ConvBlock(nn.Module): |
| 75 | def __init__(self, in_channels, out_channels, use_pixelnorm=True): |
| 76 | super(ConvBlock, self).__init__() |
| 77 | self.use_pn = use_pixelnorm |
| 78 | self.conv1 = WSConv2d(in_channels, out_channels) |
| 79 | self.conv2 = WSConv2d(out_channels, out_channels) |
| 80 | self.leaky = nn.LeakyReLU(0.2) |
| 81 | self.pn = PixelNorm() |
| 82 | |
| 83 | def forward(self, x): |
| 84 | x = self.leaky(self.conv1(x)) |
| 85 | x = self.pn(x) if self.use_pn else x |
| 86 | x = self.leaky(self.conv2(x)) |
| 87 | x = self.pn(x) if self.use_pn else x |
| 88 | return x |
| 89 | |
| 90 | |
| 91 | class Generator(nn.Module): |