| 362 | |
| 363 | |
| 364 | class ToRGB(nn.Module): |
| 365 | def __init__(self, in_channel, style_dim, upsample=True, blur_kernel=[1, 3, 3, 1]): |
| 366 | super().__init__() |
| 367 | |
| 368 | if upsample: |
| 369 | self.upsample = Upsample(blur_kernel) |
| 370 | |
| 371 | self.conv = ConvLayer(in_channel, 3, 1) |
| 372 | self.bias = nn.Parameter(torch.zeros(1, 3, 1, 1)) |
| 373 | |
| 374 | def forward(self, input, skip=None): |
| 375 | out = self.conv(input) |
| 376 | out = out + self.bias |
| 377 | |
| 378 | if skip is not None: |
| 379 | skip = self.upsample(skip) |
| 380 | out = out + skip |
| 381 | |
| 382 | return out |
| 383 | |
| 384 | |
| 385 | class ToFlow(nn.Module): |