mmagic.models.editors.stylegan2.stylegan2_modules

Module Contents

Classes

_FusedBiasLeakyReLU	Wrap FusedBiasLeakyReLU to support FP16 training.
UpsampleUpFIRDn	UpFIRDn for Upsampling.
DownsampleUpFIRDn	UpFIRDn for Downsampling.
ModulatedConv2d	Modulated Conv2d in StyleGANv2.
ModulatedStyleConv	Modulated Style Convolution.
ModulatedToRGB	To RGB layer.
ConvDownLayer	Convolution and Downsampling layer.
ResBlock	Residual block used in the discriminator of StyleGAN2.
ModMBStddevLayer	Modified MiniBatch Stddev Layer.

Attributes

conv2d

mmagic.models.editors.stylegan2.stylegan2_modules.conv2d

class mmagic.models.editors.stylegan2.stylegan2_modules._FusedBiasLeakyReLU(num_channels: int, negative_slope: float = 0.2, scale: float = 2 ** 0.5)

Bases: mmcv.ops.fused_bias_leakyrelu.FusedBiasLeakyReLU

Wrap FusedBiasLeakyReLU to support FP16 training.

forward(x)

Forward function.

参数

x (Tensor) – Input feature map with shape of (N, C, …).

返回

Output feature map.

返回类型

Tensor

class mmagic.models.editors.stylegan2.stylegan2_modules.UpsampleUpFIRDn(kernel, factor=2)

Bases: mmengine.model.BaseModule

UpFIRDn for Upsampling.

This module is used in the to_rgb layers in StyleGAN2 for upsampling the images.

参数

• kernel (Array) – Blur kernel/filter used in UpFIRDn.

• factor (int, optional) – Upsampling factor. Defaults to 2.

forward(x)

Forward function.

参数

x (Tensor) – Input feature map with shape of (N, C, H, W).

返回

Output feature map.

返回类型

Tensor

class mmagic.models.editors.stylegan2.stylegan2_modules.DownsampleUpFIRDn(kernel, factor=2)

Bases: mmengine.model.BaseModule

UpFIRDn for Downsampling.

This module is mentioned in StyleGAN2 for dowampling the feature maps.

参数

• kernel (Array) – Blur kernel/filter used in UpFIRDn.

• factor (int, optional) – Downsampling factor. Defaults to 2.

forward(input)

Forward function.

参数

input (Tensor) – Input feature map with shape of (N, C, H, W).

返回

Output feature map.

返回类型

Tensor

class mmagic.models.editors.stylegan2.stylegan2_modules.ModulatedConv2d(in_channels, out_channels, kernel_size, style_channels, demodulate=True, upsample=False, downsample=False, blur_kernel=[1, 3, 3, 1], equalized_lr_cfg=dict(mode='fan_in', lr_mul=1.0, gain=1.0), style_mod_cfg=dict(bias_init=1.0), style_bias=0.0, padding=None, eps=1e-08, fp16_enabled=False)

Bases: mmengine.model.BaseModule

Modulated Conv2d in StyleGANv2.

This module implements the modulated convolution layers proposed in StyleGAN2. Details can be found in Analyzing and Improving the Image Quality of StyleGAN, CVPR2020.

参数

• in_channels (int) – Input channels.

• out_channels (int) – Output channels.

• kernel_size (int) – Kernel size, same as nn.Con2d.

• style_channels (int) – Channels for the style codes.

• demodulate (bool, optional) – Whether to adopt demodulation. Defaults to True.

• upsample (bool, optional) – Whether to adopt upsampling in features. Defaults to False.

• downsample (bool, optional) – Whether to adopt downsampling in features. Defaults to False.

• blur_kernel (list[int], optional) – Blurry kernel. Defaults to [1, 3, 3, 1].

• equalized_lr_cfg (dict | None, optional) – Configs for equalized lr. Defaults to dict(mode=’fan_in’, lr_mul=1., gain=1.).

• style_mod_cfg (dict, optional) – Configs for style modulation module. Defaults to dict(bias_init=1.).

• style_bias (float, optional) – Bias value for style code. Defaults to 0..

• eps (float, optional) – Epsilon value to avoid computation error. Defaults to 1e-8.

forward(x, style, input_gain=None)

class mmagic.models.editors.stylegan2.stylegan2_modules.ModulatedStyleConv(in_channels, out_channels, kernel_size, style_channels, upsample=False, blur_kernel=[1, 3, 3, 1], demodulate=True, style_mod_cfg=dict(bias_init=1.0), style_bias=0.0, fp16_enabled=False, conv_clamp=256, fixed_noise=False)

Bases: mmengine.model.BaseModule

Modulated Style Convolution.

In this module, we integrate the modulated conv2d, noise injector and activation layers into together.

参数

• in_channels (int) – Input channels.

• out_channels (int) – Output channels.

• kernel_size (int) – Kernel size, same as nn.Con2d.

• style_channels (int) – Channels for the style codes.

• demodulate (bool, optional) – Whether to adopt demodulation. Defaults to True.

• upsample (bool, optional) – Whether to adopt upsampling in features. Defaults to False.

• downsample (bool, optional) – Whether to adopt downsampling in features. Defaults to False.

• blur_kernel (list[int], optional) – Blurry kernel. Defaults to [1, 3, 3, 1].

• equalized_lr_cfg (dict | None, optional) – Configs for equalized lr. Defaults to dict(mode=’fan_in’, lr_mul=1., gain=1.).

• style_mod_cfg (dict, optional) – Configs for style modulation module. Defaults to dict(bias_init=1.).

• style_bias (float, optional) – Bias value for style code. Defaults to 0..

• fp16_enabled (bool, optional) – Whether to use fp16 training in this module. Defaults to False.

• conv_clamp (float, optional) – Clamp the convolutional layer results to avoid gradient overflow. Defaults to 256.0.

forward(x, style, noise=None, add_noise=True, return_noise=False)

Forward Function.

参数

• x ([Tensor) – Input features with shape of (N, C, H, W).

• style (Tensor) – Style latent with shape of (N, C).

• noise (Tensor, optional) – Noise for injection. Defaults to None.

• add_noise (bool, optional) – Whether apply noise injection to feature. Defaults to True.

• return_noise (bool, optional) – Whether to return noise tensors. Defaults to False.

返回

Output features with shape of (N, C, H, W)

返回类型

Tensor

class mmagic.models.editors.stylegan2.stylegan2_modules.ModulatedToRGB(in_channels, style_channels, out_channels=3, upsample=True, blur_kernel=[1, 3, 3, 1], style_mod_cfg=dict(bias_init=1.0), style_bias=0.0, fp16_enabled=False, conv_clamp=256, out_fp32=True)

Bases: mmengine.model.BaseModule

To RGB layer.

This module is designed to output image tensor in StyleGAN2.

参数

• in_channels (int) – Input channels.

• style_channels (int) – Channels for the style codes.

• out_channels (int, optional) – Output channels. Defaults to 3.

• upsample (bool, optional) – Whether to adopt upsampling in features. Defaults to False.

• blur_kernel (list[int], optional) – Blurry kernel. Defaults to [1, 3, 3, 1].

• style_mod_cfg (dict, optional) – Configs for style modulation module. Defaults to dict(bias_init=1.).

• style_bias (float, optional) – Bias value for style code. Defaults to 0..

• fp16_enabled (bool, optional) – Whether to use fp16 training in this module. Defaults to False.

• conv_clamp (float, optional) – Clamp the convolutional layer results to avoid gradient overflow. Defaults to 256.0.

• out_fp32 (bool, optional) – Whether to convert the output feature map to torch.float32. Defaults to True.

forward(x, style, skip=None)

Forward Function.

参数

• x ([Tensor) – Input features with shape of (N, C, H, W).

• style (Tensor) – Style latent with shape of (N, C).

• skip (Tensor, optional) – Tensor for skip link. Defaults to None.

返回

Output features with shape of (N, C, H, W)

返回类型

Tensor

class mmagic.models.editors.stylegan2.stylegan2_modules.ConvDownLayer(in_channels, out_channels, kernel_size, downsample=False, blur_kernel=[1, 3, 3, 1], bias=True, act_cfg=dict(type='fused_bias'), fp16_enabled=False, conv_clamp=256.0)

Bases: torch.nn.Sequential

Convolution and Downsampling layer.

参数

• in_channels (int) – Input channels.

• out_channels (int) – Output channels.

• kernel_size (int) – Kernel size, same as nn.Con2d.

• downsample (bool, optional) – Whether to adopt downsampling in features. Defaults to False.

• blur_kernel (list[int], optional) – Blurry kernel. Defaults to [1, 3, 3, 1].

• bias (bool, optional) – Whether to use bias parameter. Defaults to True.

• act_cfg (dict, optional) – Activation configs. Defaults to dict(type=’fused_bias’).

• fp16_enabled (bool, optional) – Whether to use fp16 training in this module. Defaults to False.

• conv_clamp (float, optional) – Clamp the convolutional layer results to avoid gradient overflow. Defaults to 256.0.

forward(x)

class mmagic.models.editors.stylegan2.stylegan2_modules.ResBlock(in_channels, out_channels, blur_kernel=[1, 3, 3, 1], fp16_enabled=False, convert_input_fp32=True)

Bases: mmengine.model.BaseModule

Residual block used in the discriminator of StyleGAN2.

参数

• in_channels (int) – Input channels.

• out_channels (int) – Output channels.

• kernel_size (int) – Kernel size, same as nn.Con2d.

• fp16_enabled (bool, optional) – Whether to use fp16 training in this module. Defaults to False.

• convert_input_fp32 (bool, optional) – Whether to convert input type to fp32 if not fp16_enabled. This argument is designed to deal with the cases where some modules are run in FP16 and others in FP32. Defaults to True.

forward(input)

Forward function.

参数

input (Tensor) – Input feature map with shape of (N, C, H, W).

返回

Output feature map.

返回类型

Tensor

class mmagic.models.editors.stylegan2.stylegan2_modules.ModMBStddevLayer(group_size=4, channel_groups=1, sync_std=False, sync_groups=None, eps=1e-08)

Bases: mmengine.model.BaseModule

Modified MiniBatch Stddev Layer.

This layer is modified from MiniBatchStddevLayer used in PGGAN. In StyleGAN2, the authors add a new feature, channel_groups, into this layer.

Note that to accelerate the training procedure, we also add a new feature of sync_std to achieve multi-nodes/machine training. This feature is still in beta version and we have tested it on 256 scales.

参数

• group_size (int, optional) – The size of groups in batch dimension. Defaults to 4.

• channel_groups (int, optional) – The size of groups in channel dimension. Defaults to 1.

• sync_std (bool, optional) – Whether to use synchronized std feature. Defaults to False.

• sync_groups (int | None, optional) – The size of groups in node dimension. Defaults to None.

• eps (float, optional) – Epsilon value to avoid computation error. Defaults to 1e-8.

forward(x)

Forward function.

参数

x (Tensor) – Input feature map with shape of (N, C, H, W).

返回

Output feature map with shape of (N, C+1, H, W).

返回类型

Tensor