# python3.7 """Contains the implementation of generator described in StyleGAN3. Compared to that of StyleGAN2, the generator in StyleGAN3 controls the frequency flow along with the convolutional layers growing. Paper: https://arxiv.org/pdf/2106.12423.pdf Official implementation: https://github.com/NVlabs/stylegan3 """ import numpy as np import scipy.signal import torch import torch.nn as nn import torch.nn.functional as F from third_party.stylegan3_official_ops import bias_act from third_party.stylegan3_official_ops import filtered_lrelu from third_party.stylegan3_official_ops import conv2d_gradfix from .utils.ops import all_gather __all__ = ['StyleGAN3Generator'] # Resolutions allowed. _RESOLUTIONS_ALLOWED = [8, 16, 32, 64, 128, 256, 512, 1024] # pylint: disable=missing-function-docstring class StyleGAN3Generator(nn.Module): """Defines the generator network in StyleGAN3. NOTE: The synthesized images are with `RGB` channel order and pixel range [-1, 1]. Settings for the mapping network: (1) z_dim: Dimension of the input latent space, Z. (default: 512) (2) w_dim: Dimension of the output latent space, W. (default: 512) (3) repeat_w: Repeat w-code for different layers. (default: True) (4) normalize_z: Whether to normalize the z-code. (default: True) (5) mapping_layers: Number of layers of the mapping network. (default: 2) (6) mapping_fmaps: Number of hidden channels of the mapping network. (default: 512) (7) mapping_lr_mul: Learning rate multiplier for the mapping network. (default: 0.01) Settings for conditional generation: (1) label_dim: Dimension of the additional label for conditional generation. In one-hot conditioning case, it is equal to the number of classes. If set to 0, conditioning training will be disabled. (default: 0) (2) embedding_dim: Dimension of the embedding space, if needed. (default: 512) (3) embedding_bias: Whether to add bias to embedding learning. (default: True) (4) embedding_lr_mul: Learning rate multiplier for the embedding learning. (default: 1.0) (5) normalize_embedding: Whether to normalize the embedding. (default: True) (6) normalize_embedding_latent: Whether to normalize the embedding together with the latent. (default: False) Settings for the synthesis network: (1) resolution: The resolution of the output image. (default: -1) (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) output_scale: Factor to scaling the output image. (default: 0.25) (5) num_layers: Number of synthesis layers, excluding the first positional encoding layer and the last ToRGB layer. (default: 14) (6) num_critical: Number of synthesis layers with critical sampling. These layers are always set as top (with highest resolution) ones. (7) fmaps_base: Factor to control number of feature maps for each layer. (default: 32 << 10) (8) fmaps_max: Maximum number of feature maps in each layer. (default: 512) (9) kernel_size: Size of convolutional kernels. (default: 1) (10) conv_clamp: A threshold to clamp the output of convolution layers to avoid overflow under FP16 training. (default: None) (11) first_cutoff: Cutoff frequency of the first layer. (default: 2) (12) first_stopband: Stopband of the first layer. (default: 2 ** 2.1) (13) last_stopband_rel: Stopband of the last layer, relative to the last cutoff, which is `resolution / 2`. Concretely, `last_stopband` will be equal to `resolution / 2 * last_stopband_rel`. (default: 2 ** 0.3) (14) margin_size: Size of margin for each feature map. (default: 10) (15) filter_size: Size of filter for upsampling and downsampling around the activation. (default: 6) (16) act_upsampling: Factor used to upsample the feature map before activation for anti-aliasing. (default: 2) (17) use_radial_filter: Whether to use radial filter for downsampling after the activation. (default: False) (18) eps: A small value to avoid divide overflow. (default: 1e-8) Runtime settings: (1) w_moving_decay: Decay factor for updating `w_avg`, which is used for training only. Set `None` to disable. (default: 0.998) (2) sync_w_avg: Synchronizing the stats of `w_avg` across replicas. If set as `True`, the stats will be more accurate, yet the speed maybe a little bit slower. (default: False) (3) style_mixing_prob: Probability to perform style mixing as a training regularization. Set `None` to disable. (default: None) (4) trunc_psi: Truncation psi, set `None` to disable. (default: None) (5) trunc_layers: Number of layers to perform truncation. (default: None) (6) magnitude_moving_decay: Decay factor for updating `magnitude_ema` in each `SynthesisLayer`, which is used for training only. Set `None` to disable. (default: 0.999) (7) update_ema: Whether to update `w_avg` in the `MappingNetwork` and `magnitude_ema` in each `SynthesisLayer`. This field only takes effect in `training` model. (default: False) (8) fp16_res: Layers at resolution higher than (or equal to) this field will use `float16` precision for computation. This is merely used for acceleration. If set as `None`, all layers will use `float32` by default. (default: None) (9) impl: Implementation mode of some particular ops, e.g., `filtering`, `bias_act`, etc. `cuda` means using the official CUDA implementation from StyleGAN3, while `ref` means using the native PyTorch ops. (default: `cuda`) """ def __init__(self, # Settings for mapping network. z_dim=512, w_dim=512, repeat_w=True, normalize_z=True, mapping_layers=2, mapping_fmaps=512, mapping_lr_mul=0.01, # Settings for conditional generation. label_dim=0, embedding_dim=512, embedding_bias=True, embedding_lr_mul=1.0, normalize_embedding=True, normalize_embedding_latent=False, # Settings for synthesis network. resolution=-1, image_channels=3, final_tanh=False, output_scale=0.25, num_layers=14, num_critical=2, fmaps_base=32 << 10, fmaps_max=512, kernel_size=1, conv_clamp=256, first_cutoff=2, first_stopband=2 ** 2.1, last_stopband_rel=2 ** 0.3, margin_size=10, filter_size=6, act_upsampling=2, use_radial_filter=False, eps=1e-8): """Initializes with basic settings.""" super().__init__() if resolution not in _RESOLUTIONS_ALLOWED: raise ValueError(f'Invalid resolution: `{resolution}`!\n' f'Resolutions allowed: {_RESOLUTIONS_ALLOWED}.') self.z_dim = z_dim self.w_dim = w_dim self.repeat_w = repeat_w self.normalize_z = normalize_z self.mapping_layers = mapping_layers self.mapping_fmaps = mapping_fmaps self.mapping_lr_mul = mapping_lr_mul self.label_dim = label_dim self.embedding_dim = embedding_dim self.embedding_bias = embedding_bias self.embedding_lr_mul = embedding_lr_mul self.normalize_embedding = normalize_embedding self.normalize_embedding_latent = normalize_embedding_latent self.resolution = resolution self.image_channels = image_channels self.final_tanh = final_tanh self.output_scale = output_scale self.num_layers = num_layers + 2 # Including InputLayer and ToRGBLayer. self.num_critical = num_critical self.fmaps_base = fmaps_base self.fmaps_max = fmaps_max self.kernel_size = kernel_size self.conv_clamp = conv_clamp self.first_cutoff = first_cutoff self.first_stopband = first_stopband self.last_stopband_rel = last_stopband_rel self.margin_size = margin_size self.filter_size = filter_size self.act_upsampling = act_upsampling self.use_radial_filter = use_radial_filter self.eps = eps # Dimension of latent space, which is convenient for sampling. self.latent_dim = (z_dim,) self.mapping = MappingNetwork( input_dim=z_dim, output_dim=w_dim, num_outputs=self.num_layers, repeat_output=repeat_w, normalize_input=normalize_z, num_layers=mapping_layers, hidden_dim=mapping_fmaps, lr_mul=mapping_lr_mul, label_dim=label_dim, embedding_dim=embedding_dim, embedding_bias=embedding_bias, embedding_lr_mul=embedding_lr_mul, normalize_embedding=normalize_embedding, normalize_embedding_latent=normalize_embedding_latent, eps=eps) # This is used for truncation trick. if self.repeat_w: self.register_buffer('w_avg', torch.zeros(w_dim)) else: self.register_buffer('w_avg', torch.zeros(self.num_layers * w_dim)) self.synthesis = SynthesisNetwork(resolution=resolution, w_dim=w_dim, image_channels=image_channels, final_tanh=final_tanh, output_scale=output_scale, num_layers=num_layers, num_critical=num_critical, fmaps_base=fmaps_base, fmaps_max=fmaps_max, kernel_size=kernel_size, conv_clamp=conv_clamp, first_cutoff=first_cutoff, first_stopband=first_stopband, last_stopband_rel=last_stopband_rel, margin_size=margin_size, filter_size=filter_size, act_upsampling=act_upsampling, use_radial_filter=use_radial_filter, eps=eps) self.var_mapping = {'w_avg': 'mapping.w_avg'} for key, val in self.mapping.var_mapping.items(): self.var_mapping[f'mapping.{key}'] = f'mapping.{val}' for key, val in self.synthesis.var_mapping.items(): self.var_mapping[f'synthesis.{key}'] = f'synthesis.{val}' def set_space_of_latent(self, space_of_latent): """Sets the space to which the latent code belong. See `SynthesisNetwork` for more details. """ self.synthesis.set_space_of_latent(space_of_latent) def forward(self, z, label=None, w_moving_decay=0.998, sync_w_avg=False, style_mixing_prob=None, trunc_psi=None, trunc_layers=None, magnitude_moving_decay=0.999, update_ema=False, fp16_res=None, impl='cuda'): """Connects mapping network and synthesis network. This forward function will also update the average `w_code`, perform style mixing as a training regularizer, and do truncation trick, which is specially designed for inference. Concretely, the truncation trick acts as follows: For layers in range [0, truncation_layers), the truncated w-code is computed as w_new = w_avg + (w - w_avg) * trunc_psi To disable truncation, please set (1) trunc_psi = 1.0 (None) OR (2) trunc_layers = 0 (None) """ mapping_results = self.mapping(z, label, impl=impl) w = mapping_results['w'] if self.training and update_ema and w_moving_decay is not None: if sync_w_avg: batch_w_avg = all_gather(w.detach()).mean(dim=0) else: batch_w_avg = w.detach().mean(dim=0) self.w_avg.copy_(batch_w_avg.lerp(self.w_avg, w_moving_decay)) wp = mapping_results.pop('wp') if self.training and style_mixing_prob is not None: if np.random.uniform() < style_mixing_prob: new_z = torch.randn_like(z) new_wp = self.mapping(new_z, label, impl=impl)['wp'] mixing_cutoff = np.random.randint(1, self.num_layers) wp[:, mixing_cutoff:] = new_wp[:, mixing_cutoff:] if not self.training: 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: w_avg = self.w_avg.reshape(1, -1, self.w_dim)[:, :trunc_layers] wp[:, :trunc_layers] = w_avg.lerp( wp[:, :trunc_layers], trunc_psi) synthesis_results = self.synthesis( wp, magnitude_moving_decay=magnitude_moving_decay, update_ema=update_ema, fp16_res=fp16_res, impl=impl) return {**mapping_results, **synthesis_results} class MappingNetwork(nn.Module): """Implements the latent space mapping network. Basically, this network executes several dense layers in sequence, and the label embedding if needed. """ def __init__(self, input_dim, output_dim, num_outputs, repeat_output, normalize_input, num_layers, hidden_dim, lr_mul, label_dim, embedding_dim, embedding_bias, embedding_lr_mul, normalize_embedding, normalize_embedding_latent, eps): super().__init__() self.input_dim = input_dim self.output_dim = output_dim self.num_outputs = num_outputs self.repeat_output = repeat_output self.normalize_input = normalize_input self.num_layers = num_layers self.hidden_dim = hidden_dim self.lr_mul = lr_mul self.label_dim = label_dim self.embedding_dim = embedding_dim self.embedding_bias = embedding_bias self.embedding_lr_mul = embedding_lr_mul self.normalize_embedding = normalize_embedding self.normalize_embedding_latent = normalize_embedding_latent self.eps = eps self.var_mapping = {} self.norm = PixelNormLayer(dim=1, eps=eps) if self.label_dim > 0: input_dim = input_dim + embedding_dim self.embedding = DenseLayer(in_channels=label_dim, out_channels=embedding_dim, init_weight_std=1.0, add_bias=embedding_bias, init_bias=0.0, lr_mul=embedding_lr_mul, activation_type='linear') self.var_mapping['embedding.weight'] = 'embed.weight' if self.embedding_bias: self.var_mapping['embedding.bias'] = 'embed.bias' if num_outputs is not None and not repeat_output: output_dim = output_dim * num_outputs for i in range(num_layers): in_channels = (input_dim if i == 0 else hidden_dim) out_channels = (output_dim if i == (num_layers - 1) else hidden_dim) self.add_module(f'dense{i}', DenseLayer(in_channels=in_channels, out_channels=out_channels, init_weight_std=1.0, add_bias=True, init_bias=0.0, lr_mul=lr_mul, activation_type='lrelu')) self.var_mapping[f'dense{i}.weight'] = f'fc{i}.weight' self.var_mapping[f'dense{i}.bias'] = f'fc{i}.bias' def forward(self, z, label=None, impl='cuda'): if z.ndim != 2 or z.shape[1] != self.input_dim: raise ValueError(f'Input latent code should be with shape ' f'[batch_size, input_dim], where ' f'`input_dim` equals to {self.input_dim}!\n' f'But `{z.shape}` is received!') if self.normalize_input: z = self.norm(z) if self.label_dim > 0: if label is None: raise ValueError(f'Model requires an additional label ' f'(with dimension {self.label_dim}) as input, ' f'but no label is received!') if label.ndim != 2 or label.shape != (z.shape[0], self.label_dim): raise ValueError(f'Input label should be with shape ' f'[batch_size, label_dim], where ' f'`batch_size` equals to that of ' f'latent codes ({z.shape[0]}) and ' f'`label_dim` equals to {self.label_dim}!\n' f'But `{label.shape}` is received!') label = label.to(dtype=torch.float32) embedding = self.embedding(label, impl=impl) if self.normalize_embedding: embedding = self.norm(embedding) w = torch.cat((z, embedding), dim=1) else: w = z if self.label_dim > 0 and self.normalize_embedding_latent: w = self.norm(w) for i in range(self.num_layers): w = getattr(self, f'dense{i}')(w, impl=impl) wp = None if self.num_outputs is not None: if self.repeat_output: wp = w.unsqueeze(1).repeat((1, self.num_outputs, 1)) else: wp = w.reshape(-1, self.num_outputs, self.output_dim) results = { 'z': z, 'label': label, 'w': w, 'wp': wp, } if self.label_dim > 0: results['embedding'] = embedding return results class SynthesisNetwork(nn.Module): """Implements the image synthesis network. Basically, this network executes several convolutional layers in sequence. """ def __init__(self, resolution, w_dim, image_channels, final_tanh, output_scale, num_layers, num_critical, fmaps_base, fmaps_max, kernel_size, conv_clamp, first_cutoff, first_stopband, last_stopband_rel, margin_size, filter_size, act_upsampling, use_radial_filter, eps): super().__init__() self.resolution = resolution self.w_dim = w_dim self.image_channels = image_channels self.final_tanh = final_tanh self.output_scale = output_scale self.num_layers = num_layers self.num_critical = num_critical self.fmaps_base = fmaps_base self.fmaps_max = fmaps_max self.kernel_size = kernel_size self.conv_clamp = conv_clamp self.first_cutoff = first_cutoff self.first_stopband = first_stopband self.last_stopband_rel = last_stopband_rel self.margin_size = margin_size self.filter_size = filter_size self.act_upsampling = act_upsampling self.use_radial_filter = use_radial_filter self.eps = eps self.var_mapping = {} # Get layer settings. last_cutoff = resolution / 2 last_stopband = last_cutoff * last_stopband_rel layer_indices = np.arange(num_layers + 1) exponents = np.minimum(layer_indices / (num_layers - num_critical), 1) cutoffs = first_cutoff * (last_cutoff / first_cutoff) ** exponents stopbands = ( first_stopband * (last_stopband / first_stopband) ** exponents) sampling_rates = np.exp2(np.ceil(np.log2( np.minimum(stopbands * 2, self.resolution)))) sampling_rates = np.int64(sampling_rates) half_widths = np.maximum(stopbands, sampling_rates / 2) - cutoffs sizes = sampling_rates + margin_size * 2 sizes[-2:] = resolution sizes = np.int64(sizes) channels = np.rint(np.minimum((fmaps_base / 2) / cutoffs, fmaps_max)) channels[-1] = image_channels channels = np.int64(channels) self.cutoffs = cutoffs self.stopbands = stopbands self.sampling_rates = sampling_rates self.half_widths = half_widths self.sizes = sizes self.channels = channels # Input layer, with positional encoding. self.early_layer = InputLayer(w_dim=w_dim, channels=channels[0], size=sizes[0], sampling_rate=sampling_rates[0], cutoff=cutoffs[0]) self.var_mapping['early_layer.weight'] = 'input.weight' self.var_mapping['early_layer.affine.weight'] = 'input.affine.weight' self.var_mapping['early_layer.affine.bias'] = 'input.affine.bias' self.var_mapping['early_layer.transform'] = 'input.transform' self.var_mapping['early_layer.frequency'] = 'input.freqs' self.var_mapping['early_layer.phase'] = 'input.phases' # Convolutional layers. for idx in range(num_layers + 1): # Position related settings. if idx < num_layers: kernel_size = self.kernel_size demodulate = True act_upsampling = self.act_upsampling else: # ToRGB layer. kernel_size = 1 demodulate = False act_upsampling = 1 if idx < num_layers - num_critical: # Non-critical sampling. use_radial_filter = self.use_radial_filter else: # Critical sampling. use_radial_filter = False prev_idx = max(idx - 1, 0) layer_name = f'layer{idx}' official_layer_name = f'L{idx}_{sizes[idx]}_{channels[idx]}' self.add_module( layer_name, SynthesisLayer(in_channels=channels[prev_idx], out_channels=channels[idx], w_dim=w_dim, kernel_size=kernel_size, demodulate=demodulate, eps=eps, conv_clamp=conv_clamp, in_size=sizes[prev_idx], out_size=sizes[idx], in_sampling_rate=sampling_rates[prev_idx], out_sampling_rate=sampling_rates[idx], in_cutoff=cutoffs[prev_idx], out_cutoff=cutoffs[idx], in_half_width=half_widths[prev_idx], out_half_width=half_widths[idx], filter_size=filter_size, use_radial_filter=use_radial_filter, act_upsampling=act_upsampling)) self.var_mapping[f'{layer_name}.magnitude_ema'] = ( f'{official_layer_name}.magnitude_ema') self.var_mapping[f'{layer_name}.conv.weight'] = ( f'{official_layer_name}.weight') self.var_mapping[f'{layer_name}.conv.style.weight'] = ( f'{official_layer_name}.affine.weight') self.var_mapping[f'{layer_name}.conv.style.bias'] = ( f'{official_layer_name}.affine.bias') self.var_mapping[f'{layer_name}.filter.bias'] = ( f'{official_layer_name}.bias') if idx < num_layers: # ToRGB layer does not need filters. self.var_mapping[f'{layer_name}.filter.up_filter'] = ( f'{official_layer_name}.up_filter') self.var_mapping[f'{layer_name}.filter.down_filter'] = ( f'{official_layer_name}.down_filter') def set_space_of_latent(self, space_of_latent): """Sets the space to which the latent code belong. This function is particularly 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. Support `W` and `Y`. """ space_of_latent = space_of_latent.upper() for module in self.modules(): if isinstance(module, ModulateConvLayer): setattr(module, 'space_of_latent', space_of_latent) def forward(self, wp, magnitude_moving_decay=0.999, update_ema=False, fp16_res=None, impl='cuda'): results = {'wp': wp} x = self.early_layer(wp[:, 0]) for idx, sampling_rate in enumerate(self.sampling_rates): if fp16_res is not None and sampling_rate >= fp16_res: x = x.to(torch.float16) layer = getattr(self, f'layer{idx}') x, style = layer(x, wp[:, idx + 1], magnitude_moving_decay=magnitude_moving_decay, update_ema=update_ema, impl=impl) results[f'style{idx}'] = style if self.output_scale != 1: x = x * self.output_scale x = x.to(torch.float32) if self.final_tanh: x = torch.tanh(x) results['image'] = x return results class PixelNormLayer(nn.Module): """Implements pixel-wise feature vector normalization layer.""" def __init__(self, dim, eps): super().__init__() self.dim = dim self.eps = eps def extra_repr(self): return f'dim={self.dim}, epsilon={self.eps}' def forward(self, x): scale = (x.square().mean(dim=self.dim, keepdim=True) + self.eps).rsqrt() return x * scale class InputLayer(nn.Module): """Implements the input layer with positional encoding. Basically, this block outputs a feature map with shape `(channels, size, size)` based on the coordinate information. `sampling_rate` and `cutoff` are used to control the coordinate range and strength respectively. For a low-pass filter, `cutoff` is the same as the `bandwidth`. The initial frequency of the starting feature map is controlled by the positional encoding `sin(2 * pi * x)`, where `x = trans(coord) * frequency + phase`. We would like to introduce rich information (i.e. frequencies), but keep all frequencies lower than stopband, which is `sampling_rate / 2`. Besides, this layer also supports learning a transformation from the latent code w, and providing a customized transformation for inference. Please use the buffer `transform`. NOTE: `size` is different from `sampling_rate`. `sampling_rate` is the actual size of the current stage, which determines the maximum frequency that the feature maps can hold. `size` is the actual height and width of the current feature map, including the extended border. """ def __init__(self, w_dim, channels, size, sampling_rate, cutoff): super().__init__() self.w_dim = w_dim self.channels = channels self.size = size self.sampling_rate = sampling_rate self.cutoff = cutoff # Coordinate of the entire feature map, with resolution (size, size). # The coordinate range for the central (sampling_rate, sampling_rate) # region is set as (-0.0, 0.5), which extends to the remaining region. theta = torch.eye(2, 3) theta[0, 0] = 0.5 / sampling_rate * size theta[1, 1] = 0.5 / sampling_rate * size grid = F.affine_grid(theta=theta.unsqueeze(0), size=(1, 1, size, size), align_corners=False) self.register_buffer('grid', grid) # Draw random frequency from a uniform 2D disc for each channel # regarding X and Y dimension. And also draw a random phase for each # channel. Accordingly, each channel has three pre-defined parameters, # which are X-frequency, Y-frequency, and phase. frequency = torch.randn(channels, 2) radius = frequency.square().sum(dim=1, keepdim=True).sqrt() frequency = frequency / (radius * radius.square().exp().pow(0.25)) frequency = frequency * cutoff self.register_buffer('frequency', frequency) phase = torch.rand(channels) - 0.5 self.register_buffer('phase', phase) # This layer is used to map the latent code w to transform factors, # with order: cos(angle), sin(angle), transpose_x, transpose_y. self.affine = DenseLayer(in_channels=w_dim, out_channels=4, init_weight_std=0.0, add_bias=True, init_bias=(1, 0, 0, 0), lr_mul=1.0, activation_type='linear') # It is possible to use this buffer to customize the transform of the # output synthesis. self.register_buffer('transform', torch.eye(3)) # Use 1x1 conv to convert positional encoding to features. self.weight = nn.Parameter(torch.randn(channels, channels)) self.weight_scale = 1 / np.sqrt(channels) def extra_repr(self): return (f'channels={self.channels}, ' f'size={self.size}, ' f'sampling_rate={self.sampling_rate}, ' f'cutoff={self.cutoff:.3f}, ') def forward(self, w): batch = w.shape[0] # Get transformation matrix. # Factor controlled by latent code. transformation_factor = self.affine(w) # Ensure the range of cosine and sine value (first two dimension). _norm = transformation_factor[:, :2].norm(dim=1, keepdim=True) transformation_factor = transformation_factor / _norm # Rotation. rotation = torch.eye(3, device=w.device).unsqueeze(0) rotation = rotation.repeat((batch, 1, 1)) rotation[:, 0, 0] = transformation_factor[:, 0] rotation[:, 0, 1] = -transformation_factor[:, 1] rotation[:, 1, 0] = transformation_factor[:, 1] rotation[:, 1, 1] = transformation_factor[:, 0] # Translation. translation = torch.eye(3, device=w.device).unsqueeze(0) translation = translation.repeat((batch, 1, 1)) translation[:, 0, 2] = -transformation_factor[:, 2] translation[:, 1, 2] = -transformation_factor[:, 3] # Customized transformation. transform = rotation @ translation @ self.transform.unsqueeze(0) # Transform frequency and shift, which is equivalent to transforming # the coordinate. For example, given a coordinate, X, we would like to # first transform it with the rotation matrix, R, and the translation # matrix, T, as X' = RX + T. Then, we will apply frequency, f, and # phase, p, with sin(2 * pi * (fX' + p)). Natively, we have # fX' + p = f(RX + T) + p = (fR)X + (fT + p) frequency = self.frequency.unsqueeze(0) @ transform[:, :2, :2] # [NC2] phase = self.frequency.unsqueeze(0) @ transform[:, :2, 2:] # [NC] phase = phase.squeeze(2) + self.phase.unsqueeze(0) # [NC] # Positional encoding. x = self.grid # [NHW2] x = x.unsqueeze(3) # [NHW12] x = x @ frequency.transpose(1, 2).unsqueeze(1).unsqueeze(2) # [NHW1C] x = x.squeeze(3) # [NHWC] x = x + phase.unsqueeze(1).unsqueeze(2) # [NHWC] x = torch.sin(2 * np.pi * x) # [NHWC] # Dampen out-of-band frequency that may be introduced by the customized # transform `self.transform`. frequency_norm = frequency.norm(dim=2) stopband = self.sampling_rate / 2 factor = (frequency_norm - self.cutoff) / (stopband - self.cutoff) amplitude = (1 - factor).clamp(0, 1) # [NC] x = x * amplitude.unsqueeze(1).unsqueeze(2) # [NHWC] # Project positional encoding to features. weight = self.weight * self.weight_scale x = x @ weight.t() return x.permute(0, 3, 1, 2).contiguous() class SynthesisLayer(nn.Module): """Implements the synthesis layer. Each synthesis layer (including ToRGB layer) consists of a `ModulateConvLayer` and a `FilteringActLayer`. Besides, this layer will trace the magnitude (norm) of the input feature map, and update the statistic with `magnitude_moving_decay`. """ def __init__(self, # Settings for modulated convolution. in_channels, out_channels, w_dim, kernel_size, demodulate, eps, conv_clamp, # Settings for filtering activation. in_size, out_size, in_sampling_rate, out_sampling_rate, in_cutoff, out_cutoff, in_half_width, out_half_width, filter_size, use_radial_filter, act_upsampling): """Initializes with layer settings. Args: in_channels: Number of channels of the input tensor. out_channels: Number of channels of the output tensor. w_dim: Dimension of W space for style modulation. kernel_size: Size of the convolutional kernels. demodulate: Whether to perform style demodulation. eps: A small value to avoid divide overflow. conv_clamp: A threshold to clamp the output of convolution layers to avoid overflow under FP16 training. in_size: Size of the input feature map, i.e., height and width. out_size: Size of the output feature map, i.e., height and width. in_sampling_rate: Sampling rate of the input feature map. Different from `in_size` that includes extended border, this field controls the actual maximum frequency that can be represented by the feature map. out_sampling_rate: Sampling rate of the output feature map. in_cutoff: Cutoff frequency of the input feature map. out_cutoff: Cutoff frequency of the output feature map. in_half_width: Half-width of the transition band of the input feature map. out_half_width: Half-width of the transition band of the output feature map. filter_size: Size of the filter used in this layer. use_radial_filter: Whether to use radial filter. act_upsampling: Upsampling factor used before the activation. `1` means do not wrap upsampling and downsampling around the activation. """ super().__init__() self.in_channels = in_channels self.out_channels = out_channels self.w_dim = w_dim self.kernel_size = kernel_size self.demodulate = demodulate self.eps = eps self.conv_clamp = conv_clamp self.in_size = in_size self.out_size = out_size self.in_sampling_rate = in_sampling_rate self.out_sampling_rate = out_sampling_rate self.in_cutoff = in_cutoff self.out_cutoff = out_cutoff self.in_half_width = in_half_width self.out_half_width = out_half_width self.filter_size = filter_size self.use_radial_filter = use_radial_filter self.act_upsampling = act_upsampling self.conv = ModulateConvLayer(in_channels=in_channels, out_channels=out_channels, w_dim=w_dim, kernel_size=kernel_size, demodulate=demodulate, eps=eps) self.register_buffer('magnitude_ema', torch.ones(())) self.filter = FilteringActLayer(out_channels=out_channels, in_size=in_size, out_size=out_size, in_sampling_rate=in_sampling_rate, out_sampling_rate=out_sampling_rate, in_cutoff=in_cutoff, out_cutoff=out_cutoff, in_half_width=in_half_width, out_half_width=out_half_width, filter_size=filter_size, use_radial_filter=use_radial_filter, conv_padding=self.conv.padding, act_upsampling=act_upsampling) def extra_repr(self): return f'conv_clamp={self.conv_clamp}' def forward(self, x, w, magnitude_moving_decay=0.999, update_ema=False, impl='cuda'): if self.training and update_ema and magnitude_moving_decay is not None: magnitude = x.detach().to(torch.float32).square().mean() self.magnitude_ema.copy_( magnitude.lerp(self.magnitude_ema, magnitude_moving_decay)) input_gain = self.magnitude_ema.rsqrt() x, style = self.conv(x, w, gain=input_gain, impl=impl) if self.act_upsampling > 1: x = self.filter(x, np.sqrt(2), 0.2, self.conv_clamp, impl=impl) else: x = self.filter(x, 1, 1, self.conv_clamp, impl=impl) return x, style class ModulateConvLayer(nn.Module): """Implements the convolutional layer with style modulation. Different from the one introduced in StyleGAN2, this layer has following changes: (1) fusing `conv` and `style modulation` into one op by default (2) NOT adding a noise onto the output feature map. (3) NOT activating the feature map, which is moved to `FilteringActLayer`. """ def __init__(self, in_channels, out_channels, w_dim, kernel_size, demodulate, eps): """Initializes with layer settings. Args: in_channels: Number of channels of the input tensor. out_channels: Number of channels of the output tensor. w_dim: Dimension of W space for style modulation. kernel_size: Size of the convolutional kernels. demodulate: Whether to perform style demodulation. eps: A small value to avoid divide overflow. """ super().__init__() self.in_channels = in_channels self.out_channels = out_channels self.w_dim = w_dim self.kernel_size = kernel_size self.demodulate = demodulate self.eps = eps self.space_of_latent = 'W' # Set up weight. weight_shape = (out_channels, in_channels, kernel_size, kernel_size) self.weight = nn.Parameter(torch.randn(*weight_shape)) self.wscale = 1.0 / np.sqrt(kernel_size * kernel_size * in_channels) self.padding = kernel_size - 1 # Set up style. self.style = DenseLayer(in_channels=w_dim, out_channels=in_channels, init_weight_std=1.0, add_bias=True, init_bias=1.0, lr_mul=1.0, activation_type='linear') def extra_repr(self): return (f'in_ch={self.in_channels}, ' f'out_ch={self.out_channels}, ' f'ksize={self.kernel_size}, ' f'demodulate={self.demodulate}') def forward_style(self, w, impl='cuda'): """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. """ space_of_latent = self.space_of_latent.upper() if space_of_latent == 'W': if w.ndim != 2 or w.shape[1] != self.w_dim: raise ValueError(f'The input tensor should be with shape ' f'[batch_size, w_dim], where ' f'`w_dim` equals to {self.w_dim}!\n' f'But `{w.shape}` is received!') style = self.style(w, impl=impl) elif space_of_latent == 'Y': if w.ndim != 2 or w.shape[1] < self.in_channels: raise ValueError(f'The input tensor should be with shape ' f'[batch_size, y_dim], where ' f'`y_dim` equals to {self.in_channels}!\n' f'But `{w.shape}` is received!') style = w[:, :self.in_channels] else: raise NotImplementedError(f'Not implemented `space_of_latent`: ' f'`{space_of_latent}`!') return style def forward(self, x, w, gain=None, impl='cuda'): dtype = x.dtype N, C, H, W = x.shape # Affine on `w`. style = self.forward_style(w, impl=impl) if not self.demodulate: _style = style * self.wscale # Equivalent to scaling weight. else: _style = style weight = self.weight out_ch, in_ch, kh, kw = weight.shape assert in_ch == C # Pre-normalize inputs. if self.demodulate: weight = (weight * weight.square().mean(dim=(1, 2, 3), keepdim=True).rsqrt()) _style = _style * _style.square().mean().rsqrt() weight = weight.unsqueeze(0) weight = weight * _style.reshape(N, 1, in_ch, 1, 1) # modulation if self.demodulate: decoef = (weight.square().sum(dim=(2, 3, 4)) + self.eps).rsqrt() weight = weight * decoef.reshape(N, out_ch, 1, 1, 1) # demodulation if gain is not None: gain = gain.expand(N, in_ch) weight = weight * gain.reshape(N, 1, in_ch, 1, 1) # Fuse `conv` and `style modulation` as one op, using group convolution. x = x.reshape(1, N * in_ch, H, W) w = weight.reshape(N * out_ch, in_ch, kh, kw).to(dtype) x = conv2d_gradfix.conv2d( x, w, padding=self.padding, groups=N, impl=impl) x = x.reshape(N, out_ch, x.shape[2], x.shape[3]) assert x.dtype == dtype assert style.dtype == torch.float32 return x, style class FilteringActLayer(nn.Module): """Implements the activation, wrapped with upsampling and downsampling. Basically, this layer executes the following operations in order: (1) Apply bias. (2) Upsample the feature map to increase sampling rate. (3) Apply non-linearity as activation. (4) Downsample the feature map to target size. This layer is mostly borrowed from the official implementation: https://github.com/NVlabs/stylegan3/blob/main/training/networks_stylegan3.py """ def __init__(self, out_channels, in_size, out_size, in_sampling_rate, out_sampling_rate, in_cutoff, out_cutoff, in_half_width, out_half_width, filter_size, use_radial_filter, conv_padding, act_upsampling): """Initializes with layer settings. Args: out_channels: Number of output channels, which is used for `bias`. in_size: Size of the input feature map, i.e., height and width. out_size: Size of the output feature map, i.e., height and width. in_sampling_rate: Sampling rate of the input feature map. Different from `in_size` that includes extended border, this field controls the actual maximum frequency that can be represented by the feature map. out_sampling_rate: Sampling rate of the output feature map. in_cutoff: Cutoff frequency of the input feature map. out_cutoff: Cutoff frequency of the output feature map. in_half_width: Half-width of the transition band of the input feature map. out_half_width: Half-width of the transition band of the output feature map. filter_size: Size of the filter used in this layer. use_radial_filter: Whether to use radial filter. conv_padding: The padding used in the previous convolutional layer. act_upsampling: Upsampling factor used before the activation. `1` means do not wrap upsampling and downsampling around the activation. """ super().__init__() self.out_channels = out_channels self.in_size = in_size self.out_size = out_size self.in_sampling_rate = in_sampling_rate self.out_sampling_rate = out_sampling_rate self.in_cutoff = in_cutoff self.out_cutoff = out_cutoff self.in_half_width = in_half_width self.out_half_width = out_half_width self.filter_size = filter_size self.use_radial_filter = use_radial_filter self.conv_padding = conv_padding self.act_upsampling = act_upsampling # Define bias. self.bias = nn.Parameter(torch.zeros(out_channels)) # This sampling rate describes the upsampled feature map before # activation. temp_sampling_rate = max(in_sampling_rate, out_sampling_rate) temp_sampling_rate = temp_sampling_rate * act_upsampling # Design upsampling filter. up_factor = int(np.rint(temp_sampling_rate / in_sampling_rate)) assert in_sampling_rate * up_factor == temp_sampling_rate if up_factor > 1: self.up_factor = up_factor self.up_taps = filter_size * up_factor else: self.up_factor = 1 self.up_taps = 1 # No filtering. self.register_buffer( 'up_filter', self.design_lowpass_filter(numtaps=self.up_taps, cutoff=in_cutoff, width=in_half_width * 2, fs=temp_sampling_rate, radial=False)) # Design downsampling filter. down_factor = int(np.rint(temp_sampling_rate / out_sampling_rate)) assert out_sampling_rate * down_factor == temp_sampling_rate if down_factor > 1: self.down_factor = down_factor self.down_taps = filter_size * down_factor else: self.down_factor = 1 self.down_taps = 1 # No filtering. self.register_buffer( 'down_filter', self.design_lowpass_filter(numtaps=self.down_taps, cutoff=out_cutoff, width=out_half_width * 2, fs=temp_sampling_rate, radial=use_radial_filter)) # Compute padding. # Desired output size before downsampling. pad_total = (out_size - 1) * self.down_factor + 1 # Input size after upsampling. pad_total = pad_total - (in_size + conv_padding) * self.up_factor # Size reduction caused by the filters. pad_total = pad_total + self.up_taps + self.down_taps - 2 # Shift sample locations according to the symmetric interpretation. pad_lo = (pad_total + self.up_factor) // 2 pad_hi = pad_total - pad_lo self.padding = list(map(int, (pad_lo, pad_hi, pad_lo, pad_hi))) def extra_repr(self): return (f'in_size={self.in_size}, ' f'out_size={self.out_size}, ' f'in_srate={self.in_sampling_rate}, ' f'out_srate={self.out_sampling_rate}, ' f'in_cutoff={self.in_cutoff:.3f}, ' f'out_cutoff={self.out_cutoff:.3f}, ' f'in_half_width={self.in_half_width:.3f}, ' f'out_half_width={self.out_half_width:.3f}, ' f'up_factor={self.up_factor}, ' f'up_taps={self.up_taps}, ' f'down_factor={self.down_factor}, ' f'down_taps={self.down_taps}, ' f'filter_size={self.filter_size}, ' f'radial_filter={self.use_radial_filter}, ' f'conv_padding={self.conv_padding}, ' f'act_upsampling={self.act_upsampling}') @staticmethod def design_lowpass_filter(numtaps, cutoff, width, fs, radial=False): """Designs a low-pass filter. Args: numtaps: Length of the filter (number of coefficients, i.e., the filter order + 1). cutoff: Cutoff frequency of the output filter. width: Width of the transition region. fs: Sampling frequency. radial: Whether to use radially symmetric jinc-based filter. (default: False) """ if numtaps == 1: return None assert numtaps > 1 if not radial: # Separable Kaiser low-pass filter. f = scipy.signal.firwin(numtaps=numtaps, cutoff=cutoff, width=width, fs=fs) else: # Radially symmetric jinc-based filter. x = (np.arange(numtaps) - (numtaps - 1) / 2) / fs r = np.hypot(*np.meshgrid(x, x)) f = scipy.special.j1(2 * cutoff * (np.pi * r)) / (np.pi * r) beta = scipy.signal.kaiser_beta( scipy.signal.kaiser_atten(numtaps, width / (fs / 2))) w = np.kaiser(numtaps, beta) f = f * np.outer(w, w) f = f / np.sum(f) return torch.as_tensor(f, dtype=torch.float32) def forward(self, x, gain, slope, clamp, impl='cuda'): dtype = x.dtype x = filtered_lrelu.filtered_lrelu(x=x, fu=self.up_filter, fd=self.down_filter, b=self.bias.to(dtype), up=self.up_factor, down=self.down_factor, padding=self.padding, gain=gain, slope=slope, clamp=clamp, impl=impl) assert x.dtype == dtype return x class DenseLayer(nn.Module): """Implements the dense layer.""" def __init__(self, in_channels, out_channels, init_weight_std, add_bias, init_bias, lr_mul, activation_type): """Initializes with layer settings. Args: in_channels: Number of channels of the input tensor. out_channels: Number of channels of the output tensor. init_weight_std: The initial standard deviation of weight. add_bias: Whether to add bias onto the fully-connected result. init_bias: The initial bias value before training. lr_mul: Learning multiplier for both weight and bias. activation_type: Type of activation. """ super().__init__() self.in_channels = in_channels self.out_channels = out_channels self.init_weight_std = init_weight_std self.add_bias = add_bias self.init_bias = init_bias self.lr_mul = lr_mul self.activation_type = activation_type weight_shape = (out_channels, in_channels) self.weight = nn.Parameter( torch.randn(*weight_shape) * init_weight_std / lr_mul) self.wscale = lr_mul / np.sqrt(in_channels) if add_bias: init_bias = np.float32(np.float32(init_bias) / lr_mul) if isinstance(init_bias, np.float32): self.bias = nn.Parameter(torch.full([out_channels], init_bias)) else: assert isinstance(init_bias, np.ndarray) self.bias = nn.Parameter(torch.from_numpy(init_bias)) self.bscale = lr_mul else: self.bias = None def extra_repr(self): return (f'in_ch={self.in_channels}, ' f'out_ch={self.out_channels}, ' f'init_weight_std={self.init_weight_std}, ' f'bias={self.add_bias}, ' f'init_bias={self.init_bias}, ' f'lr_mul={self.lr_mul:.3f}, ' f'act={self.activation_type}') def forward(self, x, impl='cuda'): dtype = x.dtype if x.ndim != 2: x = x.flatten(start_dim=1) weight = self.weight.to(dtype) * self.wscale bias = None if self.bias is not None: bias = self.bias.to(dtype) if self.bscale != 1.0: bias = bias * self.bscale # Fast pass for linear activation. if self.activation_type == 'linear' and bias is not None: x = torch.addmm(bias.unsqueeze(0), x, weight.t()) else: x = x.matmul(weight.t()) x = bias_act.bias_act(x, bias, act=self.activation_type, impl=impl) assert x.dtype == dtype return x # pylint: enable=missing-function-docstring