Spaces:
Running
on
T4
Running
on
T4
| import random | |
| import torch | |
| from torch import nn | |
| from model.stylegan.model import ConvLayer, PixelNorm, EqualLinear, Generator | |
| class AdaptiveInstanceNorm(nn.Module): | |
| def __init__(self, fin, style_dim=512): | |
| super().__init__() | |
| self.norm = nn.InstanceNorm2d(fin, affine=False) | |
| self.style = nn.Linear(style_dim, fin * 2) | |
| self.style.bias.data[:fin] = 1 | |
| self.style.bias.data[fin:] = 0 | |
| def forward(self, input, style): | |
| style = self.style(style).unsqueeze(2).unsqueeze(3) | |
| gamma, beta = style.chunk(2, 1) | |
| out = self.norm(input) | |
| out = gamma * out + beta | |
| return out | |
| # modulative residual blocks (ModRes) | |
| class AdaResBlock(nn.Module): | |
| def __init__(self, fin, style_dim=512, dilation=1): # modified | |
| super().__init__() | |
| self.conv = ConvLayer(fin, fin, 3, dilation=dilation) # modified | |
| self.conv2 = ConvLayer(fin, fin, 3, dilation=dilation) # modified | |
| self.norm = AdaptiveInstanceNorm(fin, style_dim) | |
| self.norm2 = AdaptiveInstanceNorm(fin, style_dim) | |
| # model initialization | |
| # the convolution filters are set to values close to 0 to produce negligible residual features | |
| self.conv[0].weight.data *= 0.01 | |
| self.conv2[0].weight.data *= 0.01 | |
| def forward(self, x, s, w=1): | |
| skip = x | |
| if w == 0: | |
| return skip | |
| out = self.conv(self.norm(x, s)) | |
| out = self.conv2(self.norm2(out, s)) | |
| out = out * w + skip | |
| return out | |
| class DualStyleGAN(nn.Module): | |
| def __init__(self, size, style_dim, n_mlp, channel_multiplier=2, twoRes=True, res_index=6): | |
| super().__init__() | |
| layers = [PixelNorm()] | |
| for i in range(n_mlp-6): | |
| layers.append(EqualLinear(512, 512, lr_mul=0.01, activation="fused_lrelu")) | |
| # color transform blocks T_c | |
| self.style = nn.Sequential(*layers) | |
| # StyleGAN2 | |
| self.generator = Generator(size, style_dim, n_mlp, channel_multiplier) | |
| # The extrinsic style path | |
| self.res = nn.ModuleList() | |
| self.res_index = res_index//2 * 2 | |
| self.res.append(AdaResBlock(self.generator.channels[2 ** 2])) # for conv1 | |
| for i in range(3, self.generator.log_size + 1): | |
| out_channel = self.generator.channels[2 ** i] | |
| if i < 3 + self.res_index//2: | |
| # ModRes | |
| self.res.append(AdaResBlock(out_channel)) | |
| self.res.append(AdaResBlock(out_channel)) | |
| else: | |
| # structure transform block T_s | |
| self.res.append(EqualLinear(512, 512)) | |
| # FC layer is initialized with identity matrices, meaning no changes to the input latent code | |
| self.res[-1].weight.data = torch.eye(512) * 512.0**0.5 + torch.randn(512, 512) * 0.01 | |
| self.res.append(EqualLinear(512, 512)) | |
| self.res[-1].weight.data = torch.eye(512) * 512.0**0.5 + torch.randn(512, 512) * 0.01 | |
| self.res.append(EqualLinear(512, 512)) # for to_rgb7 | |
| self.res[-1].weight.data = torch.eye(512) * 512.0**0.5 + torch.randn(512, 512) * 0.01 | |
| self.size = self.generator.size | |
| self.style_dim = self.generator.style_dim | |
| self.log_size = self.generator.log_size | |
| self.num_layers = self.generator.num_layers | |
| self.n_latent = self.generator.n_latent | |
| self.channels = self.generator.channels | |
| def forward( | |
| self, | |
| styles, # intrinsic style code | |
| exstyles, # extrinsic style code | |
| return_latents=False, | |
| return_feat=False, | |
| inject_index=None, | |
| truncation=1, | |
| truncation_latent=None, | |
| input_is_latent=False, | |
| noise=None, | |
| randomize_noise=True, | |
| z_plus_latent=False, # intrinsic style code is z+ or z | |
| use_res=True, # whether to use the extrinsic style path | |
| fuse_index=18, # layers > fuse_index do not use the extrinsic style path | |
| interp_weights=[1]*18, # weight vector for style combination of two paths | |
| ): | |
| if not input_is_latent: | |
| if not z_plus_latent: | |
| styles = [self.generator.style(s) for s in styles] | |
| else: | |
| styles = [self.generator.style(s.reshape(s.shape[0]*s.shape[1], s.shape[2])).reshape(s.shape) for s in styles] | |
| if noise is None: | |
| if randomize_noise: | |
| noise = [None] * self.generator.num_layers | |
| else: | |
| noise = [ | |
| getattr(self.generator.noises, f"noise_{i}") for i in range(self.generator.num_layers) | |
| ] | |
| if truncation < 1: | |
| style_t = [] | |
| for style in styles: | |
| style_t.append( | |
| truncation_latent + truncation * (style - truncation_latent) | |
| ) | |
| styles = style_t | |
| if len(styles) < 2: | |
| inject_index = self.generator.n_latent | |
| if styles[0].ndim < 3: | |
| latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1) | |
| else: | |
| latent = styles[0] | |
| else: | |
| if inject_index is None: | |
| inject_index = random.randint(1, self.generator.n_latent - 1) | |
| if styles[0].ndim < 3: | |
| latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1) | |
| latent2 = styles[1].unsqueeze(1).repeat(1, self.generator.n_latent - inject_index, 1) | |
| latent = torch.cat([latent, latent2], 1) | |
| else: | |
| latent = torch.cat([styles[0][:,0:inject_index], styles[1][:,inject_index:]], 1) | |
| if use_res: | |
| if exstyles.ndim < 3: | |
| resstyles = self.style(exstyles).unsqueeze(1).repeat(1, self.generator.n_latent, 1) | |
| adastyles = exstyles.unsqueeze(1).repeat(1, self.generator.n_latent, 1) | |
| else: | |
| nB, nL, nD = exstyles.shape | |
| resstyles = self.style(exstyles.reshape(nB*nL, nD)).reshape(nB, nL, nD) | |
| adastyles = exstyles | |
| out = self.generator.input(latent) | |
| out = self.generator.conv1(out, latent[:, 0], noise=noise[0]) | |
| if use_res and fuse_index > 0: | |
| out = self.res[0](out, resstyles[:, 0], interp_weights[0]) | |
| skip = self.generator.to_rgb1(out, latent[:, 1]) | |
| i = 1 | |
| for conv1, conv2, noise1, noise2, to_rgb in zip( | |
| self.generator.convs[::2], self.generator.convs[1::2], noise[1::2], noise[2::2], self.generator.to_rgbs): | |
| if use_res and fuse_index >= i and i > self.res_index: | |
| out = conv1(out, interp_weights[i] * self.res[i](adastyles[:, i]) + | |
| (1-interp_weights[i]) * latent[:, i], noise=noise1) | |
| else: | |
| out = conv1(out, latent[:, i], noise=noise1) | |
| if use_res and fuse_index >= i and i <= self.res_index: | |
| out = self.res[i](out, resstyles[:, i], interp_weights[i]) | |
| if use_res and fuse_index >= (i+1) and i > self.res_index: | |
| out = conv2(out, interp_weights[i+1] * self.res[i+1](adastyles[:, i+1]) + | |
| (1-interp_weights[i+1]) * latent[:, i+1], noise=noise2) | |
| else: | |
| out = conv2(out, latent[:, i + 1], noise=noise2) | |
| if use_res and fuse_index >= (i+1) and i <= self.res_index: | |
| out = self.res[i+1](out, resstyles[:, i+1], interp_weights[i+1]) | |
| if use_res and fuse_index >= (i+2) and i >= self.res_index-1: | |
| skip = to_rgb(out, interp_weights[i+2] * self.res[i+2](adastyles[:, i+2]) + | |
| (1-interp_weights[i+2]) * latent[:, i + 2], skip) | |
| else: | |
| skip = to_rgb(out, latent[:, i + 2], skip) | |
| i += 2 | |
| if i > self.res_index and return_feat: | |
| return out, skip | |
| image = skip | |
| if return_latents: | |
| return image, latent | |
| else: | |
| return image, None | |
| def make_noise(self): | |
| return self.generator.make_noise() | |
| def mean_latent(self, n_latent): | |
| return self.generator.mean_latent(n_latent) | |
| def get_latent(self, input): | |
| return self.generator.style(input) |