| from math import sqrt |
|
|
| import torch |
| from torch import nn |
| import numpy as np |
|
|
|
|
| class PixelNorm(nn.Module): |
| def __init__(self): |
| super().__init__() |
|
|
| def forward(self, input): |
| return input / torch.sqrt(torch.mean(input ** 2, dim=1, keepdim=True) + 1e-6) |
|
|
|
|
| class FullyConnectedLayer(nn.Module): |
| def __init__(self, in_features, out_features, bias=True, |
| activation='linear', lr_multiplier=1, bias_init=0): |
| super().__init__() |
| self.activation = activation |
| self.weight = nn.Parameter(torch.randn([out_features, in_features]) / lr_multiplier) |
| self.bias = nn.Parameter(torch.full([out_features], np.float32(bias_init))) if bias else None |
| self.weight_gain = lr_multiplier / np.sqrt(in_features) |
| self.bias_gain = lr_multiplier |
|
|
| def forward(self, x): |
| w = self.weight.to(x.dtype) * self.weight_gain |
| b = self.bias |
| if b is not None: |
| b = b.to(x.dtype) |
| if self.bias_gain != 1: |
| b = b * self.bias_gain |
| x = torch.addmm(b.unsqueeze(0), x, w.t()) |
| return x |
|
|
|
|
| class EqualLinear(nn.Module): |
| def __init__(self, in_dim, out_dim): |
| super().__init__() |
|
|
| linear = nn.Linear(in_dim, out_dim) |
| linear.bias.data.zero_() |
|
|
| self.linear = linear |
|
|
| def forward(self, input): |
| return self.linear(input) |
|
|
| def normalize_2nd_moment(x, dim=1, eps=1e-8): |
| return x * (x.square().mean(dim=dim, keepdim=True) + eps).rsqrt() |
|
|
|
|
| class MappingNetowrk(nn.Module): |
| def __init__(self, code_dim=512, n_mlp=8, mapping_lr_multiplier=1.0): |
| super().__init__() |
|
|
| layers = [PixelNorm()] |
| for i in range(n_mlp): |
| layers.append(FullyConnectedLayer(code_dim, code_dim, |
| lr_multiplier=mapping_lr_multiplier)) |
| layers.append(nn.LeakyReLU(0.2)) |
|
|
| self.style = nn.Sequential(*layers) |
|
|
| def forward( |
| self, |
| input, |
| noise=None, |
| step=0, |
| alpha=-1, |
| mean_style=None, |
| style_weight=0, |
| mixing_range=(-1, -1), |
| ): |
| styles = [] |
|
|
| |
|
|
| if type(input) not in (list, tuple): |
| input = [input] |
|
|
| for i in input: |
| x = self.style(i) |
| styles.append(x) |
|
|
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| |
| |
| |
| |
| |
| |
| |
|
|
| return styles |
|
|
| def forward_w_trajectory(self, z): |
| """One style tensor per MLP block (after each FC + activation), for video/strip viz.""" |
| if type(z) not in (list, tuple): |
| z = [z] |
| z0 = z[0] |
| x = self.style[0](z0) |
| out_list = [] |
| i = 1 |
| while i < len(self.style): |
| x = self.style[i](x) |
| i += 1 |
| if i < len(self.style): |
| x = self.style[i](x) |
| i += 1 |
| out_list.append(x) |
| return out_list |
|
|
| |
| |
| |
| |
|
|
|
|
| class AdaptiveInstanceNorm(nn.Module): |
| def __init__(self, in_channel, style_dim): |
| super().__init__() |
|
|
| self.norm = nn.InstanceNorm2d(in_channel, eps=1e-3) |
| self.style = EqualLinear(style_dim, in_channel * 2) |
|
|
| nn.init.zeros_(self.style.linear.bias) |
|
|
| def forward(self, input, style): |
| style = self.style(style).unsqueeze(2).unsqueeze(3) |
| gamma, beta = style.chunk(2, 1) |
|
|
| if input.shape[-1] > 1: |
| out = self.norm(input) |
| else: |
| out = input |
|
|
| out = (1 + gamma) * out + beta |
| return out |
|
|
|
|
| class NoiseInjection(nn.Module): |
| def __init__(self, channel): |
| super().__init__() |
|
|
| self.weight = nn.Parameter(torch.randn(1, channel, 1, 1), requires_grad=False) |
|
|
| def forward(self, image, spatial_noise): |
| return image |
|
|
