(self, size, style_dim, motion_dim, blur_kernel=[1, 3, 3, 1], channel_multiplier=1)
| 436 | |
| 437 | class Synthesis(nn.Module): |
| 438 | def __init__(self, size, style_dim, motion_dim, blur_kernel=[1, 3, 3, 1], channel_multiplier=1): |
| 439 | super(Synthesis, self).__init__() |
| 440 | |
| 441 | self.size = size |
| 442 | self.style_dim = style_dim |
| 443 | self.motion_dim = motion_dim |
| 444 | |
| 445 | self.direction = Direction(motion_dim) |
| 446 | |
| 447 | self.channels = { |
| 448 | 4: 512, |
| 449 | 8: 512, |
| 450 | 16: 512, |
| 451 | 32: 512, |
| 452 | 64: 256 * channel_multiplier, |
| 453 | 128: 128 * channel_multiplier, |
| 454 | 256: 64 * channel_multiplier, |
| 455 | 512: 32 * channel_multiplier, |
| 456 | 1024: 16 * channel_multiplier, |
| 457 | } |
| 458 | |
| 459 | self.input = ConstantInput(self.channels[4]) |
| 460 | self.conv1 = StyledConv(self.channels[4], self.channels[4], 3, style_dim, blur_kernel=blur_kernel) |
| 461 | self.to_rgb1 = ToRGB(self.channels[4], style_dim, upsample=False) |
| 462 | |
| 463 | self.log_size = int(math.log(size, 2)) |
| 464 | self.num_layers = (self.log_size - 2) * 2 + 1 |
| 465 | |
| 466 | self.convs = nn.ModuleList() |
| 467 | self.upsamples = nn.ModuleList() |
| 468 | self.to_rgbs = nn.ModuleList() |
| 469 | self.to_flows = nn.ModuleList() |
| 470 | |
| 471 | in_channel = self.channels[4] |
| 472 | |
| 473 | for i in range(3, self.log_size + 1): |
| 474 | out_channel = self.channels[2 ** i] |
| 475 | |
| 476 | self.convs.append(StyledConv(in_channel, out_channel, 3, style_dim, upsample=True, |
| 477 | blur_kernel=blur_kernel)) |
| 478 | self.convs.append(StyledConv(out_channel, out_channel, 3, style_dim, blur_kernel=blur_kernel)) |
| 479 | self.to_rgbs.append(ToRGB(out_channel, style_dim)) |
| 480 | |
| 481 | self.to_flows.append(ToFlow(out_channel, style_dim)) |
| 482 | |
| 483 | in_channel = out_channel |
| 484 | |
| 485 | self.n_latent = self.log_size * 2 - 2 |
| 486 | |
| 487 | def forward(self, wa, alpha, feats): |
| 488 | bs = wa.shape[0] |
nothing calls this directly
no test coverage detected