`mmagic.models.editors.stylegan2.stylegan2_discriminator`¶

Module Contents¶

Classes¶

`StyleGAN2Discriminator`	StyleGAN2 Discriminator.
`ADAStyleGAN2Discriminator`	StyleGAN2 Discriminator.
`ADAAug`	Data Augmentation Module for Adaptive Discriminator augmentation.

class mmagic.models.editors.stylegan2.stylegan2_discriminator.StyleGAN2Discriminator(in_size, img_channels=3, channel_multiplier=2, blur_kernel=[1, 3, 3, 1], mbstd_cfg=dict(group_size=4, channel_groups=1), cond_size=None, cond_mapping_channels=None, cond_mapping_layers=None, num_fp16_scales=0, fp16_enabled=False, out_fp32=True, convert_input_fp32=True, input_bgr2rgb=False, init_cfg=None, pretrained=None)[source]¶

Bases: mmengine.model.BaseModule

StyleGAN2 Discriminator.

The architecture of this discriminator is proposed in StyleGAN2. More details can be found in: Analyzing and Improving the Image Quality of StyleGAN CVPR2020.

You can load pretrained model through passing information into pretrained argument. We have already offered official weights as follows:

stylegan2-ffhq-config-f: https://download.openmmlab.com/mmediting/stylegan2/official_weights/stylegan2-ffhq-config-f-official_20210327_171224-bce9310c.pth # noqa
stylegan2-horse-config-f: https://download.openmmlab.com/mmediting/stylegan2/official_weights/stylegan2-horse-config-f-official_20210327_173203-ef3e69ca.pth # noqa
stylegan2-car-config-f: https://download.openmmlab.com/mmediting/stylegan2/official_weights/stylegan2-car-config-f-official_20210327_172340-8cfe053c.pth # noqa
stylegan2-cat-config-f: https://download.openmmlab.com/mmediting/stylegan2/official_weights/stylegan2-cat-config-f-official_20210327_172444-15bc485b.pth # noqa
stylegan2-church-config-f: https://download.openmmlab.com/mmediting/stylegan2/official_weights/stylegan2-church-config-f-official_20210327_172657-1d42b7d1.pth # noqa

If you want to load the ema model, you can just use following codes:

# ckpt_http is one of the valid path from http source
discriminator = StyleGAN2Discriminator(1024, 512,
                                       pretrained=dict(
                                           ckpt_path=ckpt_http,
                                           prefix='discriminator'))

Of course, you can also download the checkpoint in advance and set ckpt_path with local path.

Note that our implementation adopts BGR image as input, while the original StyleGAN2 provides RGB images to the discriminator. Thus, we provide bgr2rgb argument to convert the image space. If your images follow the RGB order, please set it to True accordingly.

Parameters

in_size (int) – The input size of images.
img_channels (int) – The number of channels of the input image. Defaults to 3.
channel_multiplier (int, optional) – The multiplier factor for the channel number. Defaults to 2.
blur_kernel (list, optional) – The blurry kernel. Defaults to [1, 3, 3, 1].
mbstd_cfg (dict, optional) – Configs for minibatch-stddev layer. Defaults to dict(group_size=4, channel_groups=1).
cond_size (int, optional) – The size of conditional input. If None or less than 1, no conditional mapping will be applied. Defaults to None.
cond_mapping_channels (int, optional) – The dimension of the output of conditional mapping. Only work when c_dim is larger than 0. If c_dim is larger than 0 and cmap_dim is None, will. Defaults to None.
cond_mapping_layers (int, optional) – The number of mapping layer used to map conditional input. Only work when c_dim is larger than 0. If cmapping_layer is None and c_dim is larger than 0, cmapping_layer will set as 8. Defaults to None.
num_fp16_scales (int, optional) – The number of resolutions to use auto fp16 training. Defaults to 0.
fp16_enabled (bool, optional) – Whether to use fp16 training in this module. Defaults to False.
out_fp32 (bool, optional) – Whether to convert the output feature map to torch.float32. Defaults to True.
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.
input_bgr2rgb (bool, optional) – Whether to reformat the input channels with order rgb. Since we provide several converted weights, whose input order is rgb. You can set this argument to True if you want to finetune on converted weights. Defaults to False.
pretrained (dict | None, optional) – Information for pretrained models. The necessary key is ‘ckpt_path’. Besides, you can also provide ‘prefix’ to load the generator part from the whole state dict. Defaults to None.

_load_pretrained_model(ckpt_path, prefix='', map_location='cpu', strict=True)[source]¶

forward(x: torch.Tensor, label: Optional[torch.Tensor] = None)[source]¶

Forward function.

Parameters

x (torch.Tensor) – Input image tensor.
label (torch.Tensor, optional) – The conditional input feed to mapping layer. Defaults to None.

Returns

Predict score for the input image.

Return type

torch.Tensor

class mmagic.models.editors.stylegan2.stylegan2_discriminator.ADAStyleGAN2Discriminator(in_size, *args, data_aug=None, **kwargs)[source]¶

Bases: StyleGAN2Discriminator

StyleGAN2 Discriminator.

The architecture of this discriminator is proposed in StyleGAN2. More details can be found in: Analyzing and Improving the Image Quality of StyleGAN CVPR2020.

You can load pretrained model through passing information into pretrained argument. We have already offered official weights as follows:

stylegan2-ffhq-config-f: https://download.openmmlab.com/mmediting/stylegan2/official_weights/stylegan2-ffhq-config-f-official_20210327_171224-bce9310c.pth # noqa
stylegan2-horse-config-f: https://download.openmmlab.com/mmediting/stylegan2/official_weights/stylegan2-horse-config-f-official_20210327_173203-ef3e69ca.pth # noqa
stylegan2-car-config-f: https://download.openmmlab.com/mmediting/stylegan2/official_weights/stylegan2-car-config-f-official_20210327_172340-8cfe053c.pth # noqa
stylegan2-cat-config-f: https://download.openmmlab.com/mmediting/stylegan2/official_weights/stylegan2-cat-config-f-official_20210327_172444-15bc485b.pth # noqa
stylegan2-church-config-f: https://download.openmmlab.com/mmediting/stylegan2/official_weights/stylegan2-church-config-f-official_20210327_172657-1d42b7d1.pth # noqa

If you want to load the ema model, you can just use following codes:

# ckpt_http is one of the valid path from http source
discriminator = StyleGAN2Discriminator(1024, 512,
                                       pretrained=dict(
                                           ckpt_path=ckpt_http,
                                           prefix='discriminator'))

Of course, you can also download the checkpoint in advance and set ckpt_path with local path.

Note that our implementation adopts BGR image as input, while the original StyleGAN2 provides RGB images to the discriminator. Thus, we provide bgr2rgb argument to convert the image space. If your images follow the RGB order, please set it to True accordingly.

Parameters

in_size (int) – The input size of images.
img_channels (int) – The number of channels of the input image. Defaults to 3.
channel_multiplier (int, optional) – The multiplier factor for the channel number. Defaults to 2.
blur_kernel (list, optional) – The blurry kernel. Defaults to [1, 3, 3, 1].
mbstd_cfg (dict, optional) – Configs for minibatch-stddev layer. Defaults to dict(group_size=4, channel_groups=1).
cond_size (int, optional) – The size of conditional input. If None or less than 1, no conditional mapping will be applied. Defaults to None.
cond_mapping_channels (int, optional) – The dimension of the output of conditional mapping. Only work when c_dim is larger than 0. If c_dim is larger than 0 and cmap_dim is None, will. Defaults to None.
cond_mapping_layers (int, optional) – The number of mapping layer used to map conditional input. Only work when c_dim is larger than 0. If cmapping_layer is None and c_dim is larger than 0, cmapping_layer will set as 8. Defaults to None.
num_fp16_scales (int, optional) – The number of resolutions to use auto fp16 training. Defaults to 0.
fp16_enabled (bool, optional) – Whether to use fp16 training in this module. Defaults to False.
out_fp32 (bool, optional) – Whether to convert the output feature map to torch.float32. Defaults to True.
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.
input_bgr2rgb (bool, optional) – Whether to reformat the input channels with order rgb. Since we provide several converted weights, whose input order is rgb. You can set this argument to True if you want to finetune on converted weights. Defaults to False.
pretrained (dict | None, optional) – Information for pretrained models. The necessary key is ‘ckpt_path’. Besides, you can also provide ‘prefix’ to load the generator part from the whole state dict. Defaults to None.

forward(x)[source]¶: Forward function.

class mmagic.models.editors.stylegan2.stylegan2_discriminator.ADAAug(aug_pipeline=None, update_interval=4, augment_initial_p=0.0, ada_target=0.6, ada_kimg=500)[source]¶

Bases: mmengine.model.BaseModule

Data Augmentation Module for Adaptive Discriminator augmentation.

Parameters

aug_pipeline (dict, optional) – Config for augmentation pipeline. Defaults to None.
update_interval (int, optional) – Interval for updating augmentation probability. Defaults to 4.
augment_initial_p (float, optional) – Initial augmentation probability. Defaults to 0..
ada_target (float, optional) – ADA target. Defaults to 0.6.
ada_kimg (int, optional) – ADA training duration. Defaults to 500.

update(iteration=0, num_batches=0)[source]¶

Update Augment probability.

Parameters

iteration (int, optional) – Training iteration. Defaults to 0.
num_batches (int, optional) – The number of reals batches. Defaults to 0.

mmagic.models.editors.stylegan2.stylegan2_discriminator¶

Module Contents¶

Classes¶

`mmagic.models.editors.stylegan2.stylegan2_discriminator`¶