# python3.7 """Contains the implementation of generator described in StyleGAN2. Compared to that of StyleGAN, the generator in StyleGAN2 mainly introduces style demodulation, adds skip connections, increases model size, and disables progressive growth. This script ONLY supports config F in the original paper. Paper: https://arxiv.org/pdf/1912.04958.pdf Official TensorFlow implementation: https://github.com/NVlabs/stylegan2 """ import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from .sync_op import all_gather __all__ = ['StyleGAN2Generator'] # Resolutions allowed. _RESOLUTIONS_ALLOWED = [8, 16, 32, 64, 128, 256, 512, 1024] # Initial resolution. _INIT_RES = 4 # Architectures allowed. _ARCHITECTURES_ALLOWED = ['resnet', 'skip', 'origin'] # Default gain factor for weight scaling. _WSCALE_GAIN = 1.0 class StyleGAN2Generator(nn.Module): """Defines the generator network in StyleGAN2. NOTE: The synthesized images are with `RGB` channel order and pixel range [-1, 1]. Settings for the mapping network: (1) z_space_dim: Dimension of the input latent space, Z. (default: 512) (2) w_space_dim: Dimension of the outout latent space, W. (default: 512) (3) label_size: Size of the additional label for conditional generation. (default: 0) (4）mapping_layers: Number of layers of the mapping network. (default: 8) (5) mapping_fmaps: Number of hidden channels of the mapping network. (default: 512) (6) mapping_lr_mul: Learning rate multiplier for the mapping network. (default: 0.01) (7) repeat_w: Repeat w-code for different layers. Settings for the synthesis network: (1) resolution: The resolution of the output image. (2) image_channels: Number of channels of the output image. (default: 3) (3) final_tanh: Whether to use `tanh` to control the final pixel range. (default: False) (4) const_input: Whether to use a constant in the first convolutional layer. (default: True) (5) architecture: Type of architecture. Support `origin`, `skip`, and `resnet`. (default: `resnet`) (6) fused_modulate: Whether to fuse `style_modulate` and `conv2d` together. (default: True) (7) demodulate: Whether to perform style demodulation. (default: True) (8) use_wscale: Whether to use weight scaling. (default: True) (9) noise_type: Type of noise added to the convolutional results at each layer. (default: `spatial`) (10) fmaps_base: Factor to control number of feature maps for each layer. (default: 32 << 10) (11) fmaps_max: Maximum number of feature maps in each layer. (default: 512) """ def __init__(self, resolution, z_space_dim=512, w_space_dim=512, label_size=0, mapping_layers=8, mapping_fmaps=512, mapping_lr_mul=0.01, repeat_w=True, image_channels=3, final_tanh=False, const_input=True, architecture='skip', fused_modulate=True, demodulate=True, use_wscale=True, noise_type='spatial', fmaps_base=32 << 10, fmaps_max=512): """Initializes with basic settings. Raises: ValueError: If the `resolution` is not supported, or `architecture` is not supported. """ super().__init__() if resolution not in _RESOLUTIONS_ALLOWED: raise ValueError(f'Invalid resolution: `{resolution}`!\n' f'Resolutions allowed: {_RESOLUTIONS_ALLOWED}.') if architecture not in _ARCHITECTURES_ALLOWED: raise ValueError(f'Invalid architecture: `{architecture}`!\n' f'Architectures allowed: ' f'{_ARCHITECTURES_ALLOWED}.') self.init_res = _INIT_RES self.resolution = resolution self.z_space_dim = z_space_dim self.w_space_dim = w_space_dim self.label_size = label_size self.mapping_layers = mapping_layers self.mapping_fmaps = mapping_fmaps self.mapping_lr_mul = mapping_lr_mul self.repeat_w = repeat_w self.image_channels = image_channels self.final_tanh = final_tanh self.const_input = const_input self.architecture = architecture self.fused_modulate = fused_modulate self.demodulate = demodulate self.use_wscale = use_wscale self.noise_type = noise_type self.fmaps_base = fmaps_base self.fmaps_max = fmaps_max self.num_layers = int(np.log2(self.resolution // self.init_res * 2)) * 2 if self.repeat_w: self.mapping_space_dim = self.w_space_dim else: self.mapping_space_dim = self.w_space_dim * self.num_layers self.mapping = MappingModule(input_space_dim=self.z_space_dim, hidden_space_dim=self.mapping_fmaps, final_space_dim=self.mapping_space_dim, label_size=self.label_size, num_layers=self.mapping_layers, use_wscale=self.use_wscale, lr_mul=self.mapping_lr_mul) self.truncation = TruncationModule(w_space_dim=self.w_space_dim, num_layers=self.num_layers, repeat_w=self.repeat_w) self.synthesis = SynthesisModule(resolution=self.resolution, init_resolution=self.init_res, w_space_dim=self.w_space_dim, image_channels=self.image_channels, final_tanh=self.final_tanh, const_input=self.const_input, architecture=self.architecture, fused_modulate=self.fused_modulate, demodulate=self.demodulate, use_wscale=self.use_wscale, noise_type=self.noise_type, fmaps_base=self.fmaps_base, fmaps_max=self.fmaps_max) self.pth_to_tf_var_mapping = {} for key, val in self.mapping.pth_to_tf_var_mapping.items(): self.pth_to_tf_var_mapping[f'mapping.{key}'] = val for key, val in self.truncation.pth_to_tf_var_mapping.items(): self.pth_to_tf_var_mapping[f'truncation.{key}'] = val for key, val in self.synthesis.pth_to_tf_var_mapping.items(): self.pth_to_tf_var_mapping[f'synthesis.{key}'] = val def set_space_of_latent(self, space_of_latent='w'): """Sets the space to which the latent code belong. This function is particually used for choosing how to inject the latent code into the convolutional layers. The original generator will take a W-Space code and apply it for style modulation after an affine transformation. But, sometimes, it may need to directly feed an already affine-transformed code into the convolutional layer, e.g., when training an encoder for GAN inversion. We term the transformed space as Style Space (or Y-Space). This function is designed to tell the convolutional layers how to use the input code. Args: space_of_latent: The space to which the latent code belong. Case insensitive. (default: 'w') """ for module in self.modules(): if isinstance(module, ModulateConvBlock): setattr(module, 'space_of_latent', space_of_latent) def forward(self, z, label=None, w_moving_decay=0.995, style_mixing_prob=0.9, trunc_psi=None, trunc_layers=None, randomize_noise=False, **_unused_kwargs): mapping_results = self.mapping(z, label) w = mapping_results['w'] if self.training and w_moving_decay < 1: batch_w_avg = all_gather(w).mean(dim=0) self.truncation.w_avg.copy_( self.truncation.w_avg * w_moving_decay + batch_w_avg * (1 - w_moving_decay)) if self.training and style_mixing_prob > 0: new_z = torch.randn_like(z) new_w = self.mapping(new_z, label)['w'] if np.random.uniform() < style_mixing_prob: mixing_cutoff = np.random.randint(1, self.num_layers) w = self.truncation(w) new_w = self.truncation(new_w) w[:, :mixing_cutoff] = new_w[:, :mixing_cutoff] wp = self.truncation(w, trunc_psi, trunc_layers) synthesis_results = self.synthesis(wp, randomize_noise) return {**mapping_results, **synthesis_results} class MappingModule(nn.Module): """Implements the latent space mapping module. Basically, this module executes several dense layers in sequence. """ def __init__(self, input_space_dim=512, hidden_space_dim=512, final_space_dim=512, label_size=0, num_layers=8, normalize_input=True, use_wscale=True, lr_mul=0.01): super().__init__() self.input_space_dim = input_space_dim self.hidden_space_dim = hidden_space_dim self.final_space_dim = final_space_dim self.label_size = label_size self.num_layers = num_layers self.normalize_input = normalize_input self.use_wscale = use_wscale self.lr_mul = lr_mul self.norm = PixelNormLayer() if self.normalize_input else nn.Identity() self.pth_to_tf_var_mapping = {} for i in range(num_layers): dim_mul = 2 if label_size else 1 in_channels = (input_space_dim * dim_mul if i == 0 else hidden_space_dim) out_channels = (final_space_dim if i == (num_layers - 1) else hidden_space_dim) self.add_module(f'dense{i}', DenseBlock(in_channels=in_channels, out_channels=out_channels, use_wscale=self.use_wscale, lr_mul=self.lr_mul)) self.pth_to_tf_var_mapping[f'dense{i}.weight'] = f'Dense{i}/weight' self.pth_to_tf_var_mapping[f'dense{i}.bias'] = f'Dense{i}/bias' if label_size: self.label_weight = nn.Parameter( torch.randn(label_size, input_space_dim)) self.pth_to_tf_var_mapping[f'label_weight'] = f'LabelConcat/weight' def forward(self, z, label=None): if z.ndim != 2 or z.shape[1] != self.input_space_dim: raise ValueError(f'Input latent code should be with shape ' f'[batch_size, input_dim], where ' f'`input_dim` equals to {self.input_space_dim}!\n' f'But `{z.shape}` is received!') if self.label_size: if label is None: raise ValueError(f'Model requires an additional label ' f'(with size {self.label_size}) as input, ' f'but no label is received!') if label.ndim != 2 or label.shape != (z.shape[0], self.label_size): raise ValueError(f'Input label should be with shape ' f'[batch_size, label_size], where ' f'`batch_size` equals to that of ' f'latent codes ({z.shape[0]}) and ' f'`label_size` equals to {self.label_size}!\n' f'But `{label.shape}` is received!') embedding = torch.matmul(label, self.label_weight) z = torch.cat((z, embedding), dim=1) z = self.norm(z) w = z for i in range(self.num_layers): w = self.__getattr__(f'dense{i}')(w) results = { 'z': z, 'label': label, 'w': w, } if self.label_size: results['embedding'] = embedding return results class TruncationModule(nn.Module): """Implements the truncation module. Truncation is executed as follows: For layers in range [0, truncation_layers), the truncated w-code is computed as w_new = w_avg + (w - w_avg) * truncation_psi To disable truncation, please set (1) truncation_psi = 1.0 (None) OR (2) truncation_layers = 0 (None) NOTE: The returned tensor is layer-wise style codes. """ def __init__(self, w_space_dim, num_layers, repeat_w=True): super().__init__() self.num_layers = num_layers self.w_space_dim = w_space_dim self.repeat_w = repeat_w if self.repeat_w: self.register_buffer('w_avg', torch.zeros(w_space_dim)) else: self.register_buffer('w_avg', torch.zeros(num_layers * w_space_dim)) self.pth_to_tf_var_mapping = {'w_avg': 'dlatent_avg'} def forward(self, w, trunc_psi=None, trunc_layers=None): if w.ndim == 2: if self.repeat_w and w.shape[1] == self.w_space_dim: w = w.view(-1, 1, self.w_space_dim) wp = w.repeat(1, self.num_layers, 1) else: assert w.shape[1] == self.w_space_dim * self.num_layers wp = w.view(-1, self.num_layers, self.w_space_dim) else: wp = w assert wp.ndim == 3 assert wp.shape[1:] == (self.num_layers, self.w_space_dim) trunc_psi = 1.0 if trunc_psi is None else trunc_psi trunc_layers = 0 if trunc_layers is None else trunc_layers if trunc_psi < 1.0 and trunc_layers > 0: layer_idx = np.arange(self.num_layers).reshape(1, -1, 1) coefs = np.ones_like(layer_idx, dtype=np.float32) coefs[layer_idx < trunc_layers] *= trunc_psi coefs = torch.from_numpy(coefs).to(wp) w_avg = self.w_avg.view(1, -1, self.w_space_dim) wp = w_avg + (wp - w_avg) * coefs return wp class SynthesisModule(nn.Module): """Implements the image synthesis module. Basically, this module executes several convolutional layers in sequence. """ def __init__(self, resolution=1024, init_resolution=4, w_space_dim=512, image_channels=3, final_tanh=False, const_input=True, architecture='skip', fused_modulate=True, demodulate=True, use_wscale=True, noise_type='spatial', fmaps_base=32 << 10, fmaps_max=512): super().__init__() self.init_res = init_resolution self.init_res_log2 = int(np.log2(self.init_res)) self.resolution = resolution self.final_res_log2 = int(np.log2(self.resolution)) self.w_space_dim = w_space_dim self.image_channels = image_channels self.final_tanh = final_tanh self.const_input = const_input self.architecture = architecture self.fused_modulate = fused_modulate self.demodulate = demodulate self.use_wscale = use_wscale self.noise_type = noise_type self.fmaps_base = fmaps_base self.fmaps_max = fmaps_max self.num_layers = (self.final_res_log2 - self.init_res_log2 + 1) * 2 self.pth_to_tf_var_mapping = {} for res_log2 in range(self.init_res_log2, self.final_res_log2 + 1): res = 2 ** res_log2 block_idx = res_log2 - self.init_res_log2 # First convolution layer for each resolution. if res == self.init_res: if self.const_input: self.add_module(f'early_layer', InputBlock(init_resolution=self.init_res, channels=self.get_nf(res))) self.pth_to_tf_var_mapping[f'early_layer.const'] = ( f'{res}x{res}/Const/const') else: self.add_module(f'early_layer', DenseBlock(in_channels=self.w_space_dim, out_channels=self.get_nf(res), use_wscale=self.use_wscale)) self.pth_to_tf_var_mapping[f'early_layer.weight'] = ( f'{res}x{res}/Dense/weight') self.pth_to_tf_var_mapping[f'early_layer.bias'] = ( f'{res}x{res}/Dense/bias') else: layer_name = f'layer{2 * block_idx - 1}' self.add_module( layer_name, ModulateConvBlock(in_channels=self.get_nf(res // 2), out_channels=self.get_nf(res), resolution=res, w_space_dim=self.w_space_dim, scale_factor=2, fused_modulate=self.fused_modulate, demodulate=self.demodulate, use_wscale=self.use_wscale, noise_type=self.noise_type)) self.pth_to_tf_var_mapping[f'{layer_name}.weight'] = ( f'{res}x{res}/Conv0_up/weight') self.pth_to_tf_var_mapping[f'{layer_name}.bias'] = ( f'{res}x{res}/Conv0_up/bias') self.pth_to_tf_var_mapping[f'{layer_name}.style.weight'] = ( f'{res}x{res}/Conv0_up/mod_weight') self.pth_to_tf_var_mapping[f'{layer_name}.style.bias'] = ( f'{res}x{res}/Conv0_up/mod_bias') self.pth_to_tf_var_mapping[f'{layer_name}.noise_strength'] = ( f'{res}x{res}/Conv0_up/noise_strength') self.pth_to_tf_var_mapping[f'{layer_name}.noise'] = ( f'noise{2 * block_idx - 1}') if self.architecture == 'resnet': layer_name = f'layer{2 * block_idx - 1}' self.add_module( layer_name, ConvBlock(in_channels=self.get_nf(res // 2), out_channels=self.get_nf(res), kernel_size=1, add_bias=False, scale_factor=2, use_wscale=self.use_wscale, activation_type='linear')) self.pth_to_tf_var_mapping[f'{layer_name}.weight'] = ( f'{res}x{res}/Skip/weight') # Second convolution layer for each resolution. layer_name = f'layer{2 * block_idx}' self.add_module( layer_name, ModulateConvBlock(in_channels=self.get_nf(res), out_channels=self.get_nf(res), resolution=res, w_space_dim=self.w_space_dim, fused_modulate=self.fused_modulate, demodulate=self.demodulate, use_wscale=self.use_wscale, noise_type=self.noise_type)) tf_layer_name = 'Conv' if res == self.init_res else 'Conv1' self.pth_to_tf_var_mapping[f'{layer_name}.weight'] = ( f'{res}x{res}/{tf_layer_name}/weight') self.pth_to_tf_var_mapping[f'{layer_name}.bias'] = ( f'{res}x{res}/{tf_layer_name}/bias') self.pth_to_tf_var_mapping[f'{layer_name}.style.weight'] = ( f'{res}x{res}/{tf_layer_name}/mod_weight') self.pth_to_tf_var_mapping[f'{layer_name}.style.bias'] = ( f'{res}x{res}/{tf_layer_name}/mod_bias') self.pth_to_tf_var_mapping[f'{layer_name}.noise_strength'] = ( f'{res}x{res}/{tf_layer_name}/noise_strength') self.pth_to_tf_var_mapping[f'{layer_name}.noise'] = ( f'noise{2 * block_idx}') # Output convolution layer for each resolution (if needed). if res_log2 == self.final_res_log2 or self.architecture == 'skip': layer_name = f'output{block_idx}' self.add_module( layer_name, ModulateConvBlock(in_channels=self.get_nf(res), out_channels=image_channels, resolution=res, w_space_dim=self.w_space_dim, kernel_size=1, fused_modulate=self.fused_modulate, demodulate=False, use_wscale=self.use_wscale, add_noise=False, activation_type='linear')) self.pth_to_tf_var_mapping[f'{layer_name}.weight'] = ( f'{res}x{res}/ToRGB/weight') self.pth_to_tf_var_mapping[f'{layer_name}.bias'] = ( f'{res}x{res}/ToRGB/bias') self.pth_to_tf_var_mapping[f'{layer_name}.style.weight'] = ( f'{res}x{res}/ToRGB/mod_weight') self.pth_to_tf_var_mapping[f'{layer_name}.style.bias'] = ( f'{res}x{res}/ToRGB/mod_bias') if self.architecture == 'skip': self.upsample = UpsamplingLayer() self.final_activate = nn.Tanh() if final_tanh else nn.Identity() def get_nf(self, res): """Gets number of feature maps according to current resolution.""" return min(self.fmaps_base // res, self.fmaps_max) def forward(self, wp, randomize_noise=False, x=None, image=None, start=2, stop=None): stop = self.num_layers - 1 if stop is None else stop # if wp.ndim != 3 or wp.shape[1:] != (self.num_layers, self.w_space_dim): # raise ValueError(f'Input tensor should be with shape ' # f'[batch_size, num_layers, w_space_dim], where ' # f'`num_layers` equals to {self.num_layers}, and ' # f'`w_space_dim` equals to {self.w_space_dim}!\n' # f'But `{wp.shape}` is received!') results = {'wp': wp} x = self.early_layer(wp[:, 0]) if x is None else x if self.architecture == 'origin': for layer_idx in range(self.num_layers - 1): x, style = self.__getattr__(f'layer{layer_idx}')( x, wp[:, layer_idx], randomize_noise) results[f'style{layer_idx:02d}'] = style image, style = self.__getattr__(f'output{layer_idx // 2}')( x, wp[:, layer_idx + 1]) results[f'output_style{layer_idx // 2}'] = style elif self.architecture == 'skip': # for layer_idx in range(self.num_layers - 1): for layer_idx in range(start, stop): x, style = self.__getattr__(f'layer{layer_idx}')( x, wp[:, layer_idx], randomize_noise) results[f'style{layer_idx:02d}'] = style if layer_idx % 2 == 0: temp, style = self.__getattr__(f'output{layer_idx // 2}')( x, wp[:, layer_idx + 1]) results[f'output_style{layer_idx // 2}'] = style if layer_idx == 0 or image is None: image = temp else: image = temp + self.upsample(image) elif self.architecture == 'resnet': x, style = self.layer0(x) results[f'style00'] = style for layer_idx in range(1, self.num_layers - 1, 2): residual = self.__getattr__(f'skip_layer{layer_idx // 2}')(x) x, style = self.__getattr__(f'layer{layer_idx}')( x, wp[:, layer_idx], randomize_noise) results[f'style{layer_idx:02d}'] = style x, style = self.__getattr__(f'layer{layer_idx + 1}')( x, wp[:, layer_idx + 1], randomize_noise) results[f'style{layer_idx + 1:02d}'] = style x = (x + residual) / np.sqrt(2.0) image, style = self.__getattr__(f'output{layer_idx // 2 + 1}')( x, wp[:, layer_idx + 2]) results[f'output_style{layer_idx // 2}'] = style results['image'] = self.final_activate(image) if stop == self.num_layers -1 else image results['x'] = x return results class PixelNormLayer(nn.Module): """Implements pixel-wise feature vector normalization layer.""" def __init__(self, dim=1, epsilon=1e-8): super().__init__() self.dim = dim self.eps = epsilon def forward(self, x): norm = torch.sqrt( torch.mean(x ** 2, dim=self.dim, keepdim=True) + self.eps) return x / norm class UpsamplingLayer(nn.Module): """Implements the upsampling layer. This layer can also be used as filtering by setting `scale_factor` as 1. """ def __init__(self, scale_factor=2, kernel=(1, 3, 3, 1), extra_padding=0, kernel_gain=None): super().__init__() assert scale_factor >= 1 self.scale_factor = scale_factor if extra_padding != 0: assert scale_factor == 1 if kernel is None: kernel = np.ones((scale_factor), dtype=np.float32) else: kernel = np.array(kernel, dtype=np.float32) assert kernel.ndim == 1 kernel = np.outer(kernel, kernel) kernel = kernel / np.sum(kernel) if kernel_gain is None: kernel = kernel * (scale_factor ** 2) else: assert kernel_gain > 0 kernel = kernel * (kernel_gain ** 2) assert kernel.ndim == 2 assert kernel.shape[0] == kernel.shape[1] kernel = kernel[np.newaxis, np.newaxis] self.register_buffer('kernel', torch.from_numpy(kernel)) self.kernel = self.kernel.flip(0, 1) self.upsample_padding = (0, scale_factor - 1, # Width padding. 0, 0, # Width. 0, scale_factor - 1, # Height padding. 0, 0, # Height. 0, 0, # Channel. 0, 0) # Batch size. padding = kernel.shape[2] - scale_factor + extra_padding self.padding = ((padding + 1) // 2 + scale_factor - 1, padding // 2, (padding + 1) // 2 + scale_factor - 1, padding // 2) def forward(self, x): assert x.ndim == 4 channels = x.shape[1] if self.scale_factor > 1: x = x.view(-1, channels, x.shape[2], 1, x.shape[3], 1) x = F.pad(x, self.upsample_padding, mode='constant', value=0) x = x.view(-1, channels, x.shape[2] * self.scale_factor, x.shape[4] * self.scale_factor) x = x.view(-1, 1, x.shape[2], x.shape[3]) x = F.pad(x, self.padding, mode='constant', value=0) x = F.conv2d(x, self.kernel, stride=1) x = x.view(-1, channels, x.shape[2], x.shape[3]) return x class InputBlock(nn.Module): """Implements the input block. Basically, this block starts from a const input, which is with shape `(channels, init_resolution, init_resolution)`. """ def __init__(self, init_resolution, channels): super().__init__() self.const = nn.Parameter( torch.randn(1, channels, init_resolution, init_resolution)) def forward(self, w): x = self.const.repeat(w.shape[0], 1, 1, 1) return x class ConvBlock(nn.Module): """Implements the convolutional block (no style modulation). Basically, this block executes, convolutional layer, filtering layer (if needed), and activation layer in sequence. NOTE: This block is particularly used for skip-connection branch in the `resnet` structure. """ def __init__(self, in_channels, out_channels, kernel_size=3, add_bias=True, scale_factor=1, filtering_kernel=(1, 3, 3, 1), use_wscale=True, wscale_gain=_WSCALE_GAIN, lr_mul=1.0, activation_type='lrelu'): """Initializes with block settings. Args: in_channels: Number of channels of the input tensor. out_channels: Number of channels of the output tensor. kernel_size: Size of the convolutional kernels. (default: 3) add_bias: Whether to add bias onto the convolutional result. (default: True) scale_factor: Scale factor for upsampling. `1` means skip upsampling. (default: 1) filtering_kernel: Kernel used for filtering after upsampling. (default: (1, 3, 3, 1)) use_wscale: Whether to use weight scaling. (default: True) wscale_gain: Gain factor for weight scaling. (default: _WSCALE_GAIN) lr_mul: Learning multiplier for both weight and bias. (default: 1.0) activation_type: Type of activation. Support `linear` and `lrelu`. (default: `lrelu`) Raises: NotImplementedError: If the `activation_type` is not supported. """ super().__init__() if scale_factor > 1: self.use_conv2d_transpose = True extra_padding = scale_factor - kernel_size self.filter = UpsamplingLayer(scale_factor=1, kernel=filtering_kernel, extra_padding=extra_padding, kernel_gain=scale_factor) self.stride = scale_factor self.padding = 0 # Padding is done in `UpsamplingLayer`. else: self.use_conv2d_transpose = False assert kernel_size % 2 == 1 self.stride = 1 self.padding = kernel_size // 2 weight_shape = (out_channels, in_channels, kernel_size, kernel_size) fan_in = kernel_size * kernel_size * in_channels wscale = wscale_gain / np.sqrt(fan_in) if use_wscale: self.weight = nn.Parameter(torch.randn(*weight_shape) / lr_mul) self.wscale = wscale * lr_mul else: self.weight = nn.Parameter( torch.randn(*weight_shape) * wscale / lr_mul) self.wscale = lr_mul if add_bias: self.bias = nn.Parameter(torch.zeros(out_channels)) else: self.bias = None self.bscale = lr_mul if activation_type == 'linear': self.activate = nn.Identity() self.activate_scale = 1.0 elif activation_type == 'lrelu': self.activate = nn.LeakyReLU(negative_slope=0.2, inplace=True) self.activate_scale = np.sqrt(2.0) else: raise NotImplementedError(f'Not implemented activation function: ' f'`{activation_type}`!') def forward(self, x): weight = self.weight * self.wscale bias = self.bias * self.bscale if self.bias is not None else None if self.use_conv2d_transpose: weight = weight.permute(1, 0, 2, 3).flip(2, 3) x = F.conv_transpose2d(x, weight=weight, bias=bias, stride=self.scale_factor, padding=self.padding) x = self.filter(x) else: x = F.conv2d(x, weight=weight, bias=bias, stride=self.stride, padding=self.padding) x = self.activate(x) * self.activate_scale return x class ModulateConvBlock(nn.Module): """Implements the convolutional block with style modulation.""" def __init__(self, in_channels, out_channels, resolution, w_space_dim, kernel_size=3, add_bias=True, scale_factor=1, filtering_kernel=(1, 3, 3, 1), fused_modulate=True, demodulate=True, use_wscale=True, wscale_gain=_WSCALE_GAIN, lr_mul=1.0, add_noise=True, noise_type='spatial', activation_type='lrelu', epsilon=1e-8): """Initializes with block settings. Args: in_channels: Number of channels of the input tensor. out_channels: Number of channels of the output tensor. resolution: Resolution of the output tensor. w_space_dim: Dimension of W space for style modulation. kernel_size: Size of the convolutional kernels. (default: 3) add_bias: Whether to add bias onto the convolutional result. (default: True) scale_factor: Scale factor for upsampling. `1` means skip upsampling. (default: 1) filtering_kernel: Kernel used for filtering after upsampling. (default: (1, 3, 3, 1)) fused_modulate: Whether to fuse `style_modulate` and `conv2d` together. (default: True) demodulate: Whether to perform style demodulation. (default: True) use_wscale: Whether to use weight scaling. (default: True) wscale_gain: Gain factor for weight scaling. (default: _WSCALE_GAIN) lr_mul: Learning multiplier for both weight and bias. (default: 1.0) add_noise: Whether to add noise onto the output tensor. (default: True) noise_type: Type of noise added to the feature map after the convolution (if needed). Support `spatial` and `channel`. (default: `spatial`) activation_type: Type of activation. Support `linear` and `lrelu`. (default: `lrelu`) epsilon: Small number to avoid `divide by zero`. (default: 1e-8) Raises: NotImplementedError: If the `activation_type` is not supported. """ super().__init__() self.in_c = in_channels self.out_c = out_channels self.res = resolution self.w_space_dim = w_space_dim self.ksize = kernel_size self.eps = epsilon self.space_of_latent = 'w' if scale_factor > 1: self.use_conv2d_transpose = True extra_padding = scale_factor - kernel_size self.filter = UpsamplingLayer(scale_factor=1, kernel=filtering_kernel, extra_padding=extra_padding, kernel_gain=scale_factor) self.stride = scale_factor self.padding = 0 # Padding is done in `UpsamplingLayer`. else: self.use_conv2d_transpose = False assert kernel_size % 2 == 1 self.stride = 1 self.padding = kernel_size // 2 weight_shape = (out_channels, in_channels, kernel_size, kernel_size) fan_in = kernel_size * kernel_size * in_channels wscale = wscale_gain / np.sqrt(fan_in) if use_wscale: self.weight = nn.Parameter(torch.randn(*weight_shape) / lr_mul) self.wscale = wscale * lr_mul else: self.weight = nn.Parameter( torch.randn(*weight_shape) * wscale / lr_mul) self.wscale = lr_mul self.style = DenseBlock(in_channels=w_space_dim, out_channels=in_channels, additional_bias=1.0, use_wscale=use_wscale, activation_type='linear') self.fused_modulate = fused_modulate self.demodulate = demodulate if add_bias: self.bias = nn.Parameter(torch.zeros(out_channels)) else: self.bias = None self.bscale = lr_mul if activation_type == 'linear': self.activate = nn.Identity() self.activate_scale = 1.0 elif activation_type == 'lrelu': self.activate = nn.LeakyReLU(negative_slope=0.2, inplace=True) self.activate_scale = np.sqrt(2.0) else: raise NotImplementedError(f'Not implemented activation function: ' f'`{activation_type}`!') self.add_noise = add_noise if self.add_noise: self.noise_type = noise_type.lower() if self.noise_type == 'spatial': self.register_buffer('noise', torch.randn(1, 1, self.res, self.res)) elif self.noise_type == 'channel': self.register_buffer('noise', torch.randn(1, self.channels, 1, 1)) else: raise NotImplementedError(f'Not implemented noise type: ' f'`{self.noise_type}`!') self.noise_strength = nn.Parameter(torch.zeros(())) def forward_style(self, w): """Gets style code from the given input. More specifically, if the input is from W-Space, it will be projected by an affine transformation. If it is from the Style Space (Y-Space), no operation is required. NOTE: For codes from Y-Space, we use slicing to make sure the dimension is correct, in case that the code is padded before fed into this layer. """ if self.space_of_latent == 'w': if w.ndim != 2 or w.shape[1] != self.w_space_dim: raise ValueError(f'The input tensor should be with shape ' f'[batch_size, w_space_dim], where ' f'`w_space_dim` equals to ' f'{self.w_space_dim}!\n' f'But `{w.shape}` is received!') style = self.style(w) elif self.space_of_latent == 'y': if w.ndim != 2 or w.shape[1] < self.in_c: raise ValueError(f'The input tensor should be with shape ' f'[batch_size, y_space_dim], where ' f'`y_space_dim` equals to {self.in_c}!\n' f'But `{w.shape}` is received!') style = w[:, :self.in_c] return style def forward(self, x, w, randomize_noise=False): batch = x.shape[0] weight = self.weight * self.wscale weight = weight.permute(2, 3, 1, 0) # Style modulation. style = self.forward_style(w) _weight = weight.view(1, self.ksize, self.ksize, self.in_c, self.out_c) _weight = _weight * style.view(batch, 1, 1, self.in_c, 1) # Style demodulation. if self.demodulate: _weight_norm = torch.sqrt( torch.sum(_weight ** 2, dim=[1, 2, 3]) + self.eps) _weight = _weight / _weight_norm.view(batch, 1, 1, 1, self.out_c) if self.fused_modulate: x = x.view(1, batch * self.in_c, x.shape[2], x.shape[3]) weight = _weight.permute(1, 2, 3, 0, 4).reshape( self.ksize, self.ksize, self.in_c, batch * self.out_c) else: x = x * style.view(batch, self.in_c, 1, 1) if self.use_conv2d_transpose: weight = weight.flip(0, 1) if self.fused_modulate: weight = weight.view( self.ksize, self.ksize, self.in_c, batch, self.out_c) weight = weight.permute(0, 1, 4, 3, 2) weight = weight.reshape( self.ksize, self.ksize, self.out_c, batch * self.in_c) weight = weight.permute(3, 2, 0, 1) else: weight = weight.permute(2, 3, 0, 1) x = F.conv_transpose2d(x, weight=weight, bias=None, stride=self.stride, padding=self.padding, groups=(batch if self.fused_modulate else 1)) x = self.filter(x) else: weight = weight.permute(3, 2, 0, 1) x = F.conv2d(x, weight=weight, bias=None, stride=self.stride, padding=self.padding, groups=(batch if self.fused_modulate else 1)) if self.fused_modulate: x = x.view(batch, self.out_c, self.res, self.res) elif self.demodulate: x = x / _weight_norm.view(batch, self.out_c, 1, 1) if self.add_noise: if randomize_noise: if self.noise_type == 'spatial': noise = torch.randn(x.shape[0], 1, self.res, self.res).to(x) elif self.noise_type == 'channel': noise = torch.randn(x.shape[0], self.channels, 1, 1).to(x) else: noise = self.noise x = x + noise * self.noise_strength.view(1, 1, 1, 1) bias = self.bias * self.bscale if self.bias is not None else None if bias is not None: x = x + bias.view(1, -1, 1, 1) x = self.activate(x) * self.activate_scale return x, style class DenseBlock(nn.Module): """Implements the dense block. Basically, this block executes fully-connected layer and activation layer. NOTE: This layer supports adding an additional bias beyond the trainable bias parameter. This is specially used for the mapping from the w code to the style code. """ def __init__(self, in_channels, out_channels, add_bias=True, additional_bias=0, use_wscale=True, wscale_gain=_WSCALE_GAIN, lr_mul=1.0, activation_type='lrelu'): """Initializes with block settings. Args: in_channels: Number of channels of the input tensor. out_channels: Number of channels of the output tensor. add_bias: Whether to add bias onto the fully-connected result. (default: True) additional_bias: The additional bias, which is independent from the bias parameter. (default: 0.0) use_wscale: Whether to use weight scaling. (default: True) wscale_gain: Gain factor for weight scaling. (default: _WSCALE_GAIN) lr_mul: Learning multiplier for both weight and bias. (default: 1.0) activation_type: Type of activation. Support `linear` and `lrelu`. (default: `lrelu`) Raises: NotImplementedError: If the `activation_type` is not supported. """ super().__init__() weight_shape = (out_channels, in_channels) wscale = wscale_gain / np.sqrt(in_channels) if use_wscale: self.weight = nn.Parameter(torch.randn(*weight_shape) / lr_mul) self.wscale = wscale * lr_mul else: self.weight = nn.Parameter( torch.randn(*weight_shape) * wscale / lr_mul) self.wscale = lr_mul if add_bias: self.bias = nn.Parameter(torch.zeros(out_channels)) else: self.bias = None self.bscale = lr_mul self.additional_bias = additional_bias if activation_type == 'linear': self.activate = nn.Identity() self.activate_scale = 1.0 elif activation_type == 'lrelu': self.activate = nn.LeakyReLU(negative_slope=0.2, inplace=True) self.activate_scale = np.sqrt(2.0) else: raise NotImplementedError(f'Not implemented activation function: ' f'`{activation_type}`!') def forward(self, x): if x.ndim != 2: x = x.view(x.shape[0], -1) bias = self.bias * self.bscale if self.bias is not None else None x = F.linear(x, weight=self.weight * self.wscale, bias=bias) x = self.activate(x + self.additional_bias) * self.activate_scale return x