models/stylegan3_generator.py

# 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) * truncation_psi

        To disable truncation, please set

        (1) truncation_psi = 1.0 (None) OR
        (2) truncation_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