Spaces:
Runtime error
Runtime error
| import numpy as np | |
| from torch import nn | |
| from torch.nn import Conv2d, Sequential, Module | |
| from models.hyperstyle.encoders.helpers import SeparableBlock | |
| from models.stylegan2.model import EqualLinear | |
| # layer_idx: [kernel_size, in_channels, out_channels] | |
| PARAMETERS = { | |
| 0: [3, 512, 512], | |
| 1: [1, 512, 3], | |
| 2: [3, 512, 512], | |
| 3: [3, 512, 512], | |
| 4: [1, 512, 3], | |
| 5: [3, 512, 512], | |
| 6: [3, 512, 512], | |
| 7: [1, 512, 3], | |
| 8: [3, 512, 512], | |
| 9: [3, 512, 512], | |
| 10: [1, 512, 3], | |
| 11: [3, 512, 512], | |
| 12: [3, 512, 512], | |
| 13: [1, 512, 3], | |
| 14: [3, 512, 256], | |
| 15: [3, 256, 256], | |
| 16: [1, 256, 3], | |
| 17: [3, 256, 128], | |
| 18: [3, 128, 128], | |
| 19: [1, 128, 3], | |
| 20: [3, 128, 64], | |
| 21: [3, 64, 64], | |
| 22: [1, 64, 3], | |
| 23: [3, 64, 32], | |
| 24: [3, 32, 32], | |
| 25: [1, 32, 3] | |
| } | |
| TO_RGB_LAYERS = [1, 4, 7, 10, 13, 16, 19, 22, 25] | |
| class RefinementBlock(Module): | |
| def __init__(self, layer_idx, opts, n_channels=512, inner_c=256, spatial=16): | |
| super(RefinementBlock, self).__init__() | |
| self.layer_idx = layer_idx | |
| self.opts = opts | |
| self.kernel_size, self.in_channels, self.out_channels = PARAMETERS[self.layer_idx] | |
| self.spatial = spatial | |
| self.n_channels = n_channels | |
| self.inner_c = inner_c | |
| self.out_c = 512 | |
| num_pools = int(np.log2(self.spatial)) - 1 | |
| if self.kernel_size == 3: | |
| num_pools = num_pools - 1 | |
| self.modules = [] | |
| self.modules += [Conv2d(self.n_channels, self.inner_c, kernel_size=3, stride=2, padding=1), nn.LeakyReLU()] | |
| for i in range(num_pools - 1): | |
| self.modules += [Conv2d(self.inner_c, self.inner_c, kernel_size=3, stride=2, padding=1), nn.LeakyReLU()] | |
| self.modules += [Conv2d(self.inner_c, self.out_c, kernel_size=3, stride=2, padding=1), nn.LeakyReLU()] | |
| self.convs = nn.Sequential(*self.modules) | |
| if layer_idx in TO_RGB_LAYERS: | |
| self.output = Sequential( | |
| Conv2d(self.out_c, self.in_channels * self.out_channels, kernel_size=1, stride=1, padding=0)) | |
| else: | |
| self.output = Sequential(nn.AdaptiveAvgPool2d((1, 1)), | |
| Conv2d(self.out_c, self.in_channels * self.out_channels, kernel_size=1, stride=1, | |
| padding=0)) | |
| def forward(self, x): | |
| x = self.convs(x) | |
| x = self.output(x) | |
| if self.layer_idx in TO_RGB_LAYERS: | |
| x = x.view(-1, self.out_channels, self.in_channels, self.kernel_size, self.kernel_size) | |
| else: | |
| x = x.view(-1, self.out_channels, self.in_channels) | |
| x = x.unsqueeze(3).repeat(1, 1, 1, self.kernel_size).unsqueeze(4).repeat(1, 1, 1, 1, self.kernel_size) | |
| return x | |
| class HyperRefinementBlock(Module): | |
| def __init__(self, hypernet, n_channels=512, inner_c=128, spatial=16): | |
| super(HyperRefinementBlock, self).__init__() | |
| self.n_channels = n_channels | |
| self.inner_c = inner_c | |
| self.out_c = 512 | |
| num_pools = int(np.log2(spatial)) | |
| modules = [Conv2d(self.n_channels, self.inner_c, kernel_size=3, stride=1, padding=1), nn.LeakyReLU()] | |
| for i in range(num_pools - 1): | |
| modules += [Conv2d(self.inner_c, self.inner_c, kernel_size=3, stride=2, padding=1), nn.LeakyReLU()] | |
| modules += [Conv2d(self.inner_c, self.out_c, kernel_size=3, stride=2, padding=1), nn.LeakyReLU()] | |
| self.convs = nn.Sequential(*modules) | |
| self.linear = EqualLinear(self.out_c, self.out_c, lr_mul=1) | |
| self.hypernet = hypernet | |
| def forward(self, features): | |
| code = self.convs(features) | |
| code = code.view(-1, self.out_c) | |
| code = self.linear(code) | |
| weight_delta = self.hypernet(code) | |
| return weight_delta | |
| class RefinementBlockSeparable(Module): | |
| def __init__(self, layer_idx, opts, n_channels=512, inner_c=256, spatial=16): | |
| super(RefinementBlockSeparable, self).__init__() | |
| self.layer_idx = layer_idx | |
| self.kernel_size, self.in_channels, self.out_channels = PARAMETERS[self.layer_idx] | |
| self.spatial = spatial | |
| self.n_channels = n_channels | |
| self.inner_c = inner_c | |
| self.out_c = 512 | |
| num_pools = int(np.log2(self.spatial)) - 1 | |
| self.modules = [] | |
| self.modules += [Conv2d(self.n_channels, self.inner_c, kernel_size=3, stride=2, padding=1), nn.LeakyReLU()] | |
| for i in range(num_pools - 1): | |
| self.modules += [Conv2d(self.inner_c, self.inner_c, kernel_size=3, stride=2, padding=1), nn.LeakyReLU()] | |
| self.modules += [Conv2d(self.inner_c, self.out_c, kernel_size=3, stride=2, padding=1), nn.LeakyReLU()] | |
| self.convs = nn.Sequential(*self.modules) | |
| self.opts = opts | |
| if self.layer_idx in TO_RGB_LAYERS: | |
| self.output = Sequential(Conv2d(self.out_c, self.in_channels * self.out_channels, | |
| kernel_size=1, stride=1, padding=0)) | |
| else: | |
| self.output = Sequential(SeparableBlock(input_size=self.out_c, | |
| kernel_channels_in=self.in_channels, | |
| kernel_channels_out=self.out_channels, | |
| kernel_size=self.kernel_size)) | |
| def forward(self, x): | |
| x = self.convs(x) | |
| x = self.output(x) | |
| if self.layer_idx in TO_RGB_LAYERS: | |
| x = x.view(-1, self.out_channels, self.in_channels, self.kernel_size, self.kernel_size) | |
| return x |