mmagic.models.editors.stylegan3.stylegan3_generator

Module Contents

Classes

StyleGAN3Generator

StyleGAN3 Generator.

class mmagic.models.editors.stylegan3.stylegan3_generator.StyleGAN3Generator(out_size, style_channels, img_channels, noise_size=512, rgb2bgr=False, pretrained=None, synthesis_cfg=dict(type='SynthesisNetwork'), mapping_cfg=dict(type='MappingNetwork'))

Bases: mmengine.model.BaseModule

StyleGAN3 Generator.

In StyleGAN3, we make several changes to StyleGANv2’s generator which include transformed fourier features, filtered nonlinearity and non-critical sampling, etc. More details can be found in: Alias-Free Generative Adversarial Networks NeurIPS’2021.

Ref: https://github.com/NVlabs/stylegan3

Parameters

• out_size (int) – The output size of the StyleGAN3 generator.

• style_channels (int) – The number of channels for style code.

• img_channels (int) – The number of output’s channels.

• noise_size (int, optional) – Size of the input noise vector. Defaults to 512.

• rgb2bgr (bool, optional) – Whether to reformat the output channels with order bgr. We provide several pre-trained StyleGAN3 weights whose output channels order is rgb. You can set this argument to True to use the weights.

• pretrained (str | dict, optional) – Path for the pretrained model or dict containing information for pretrained models whose necessary key is ‘ckpt_path’. Besides, you can also provide ‘prefix’ to load the generator part from the whole state dict. Defaults to None.

• synthesis_cfg (dict, optional) – Config for synthesis network. Defaults to dict(type=’SynthesisNetwork’).

• mapping_cfg (dict, optional) – Config for mapping network. Defaults to dict(type=’MappingNetwork’).

_load_pretrained_model(ckpt_path, prefix='', map_location='cpu', strict=True)

forward(noise, num_batches=0, input_is_latent=False, truncation=1, num_truncation_layer=None, update_emas=False, force_fp32=True, return_noise=False, return_latents=False)

Forward Function for stylegan3.

Parameters

• noise (torch.Tensor | callable | None) – You can directly give a batch of noise through a torch.Tensor or offer a callable function to sample a batch of noise data. Otherwise, the None indicates to use the default noise sampler.

• num_batches (int, optional) – The number of batch size. Defaults to 0.

• input_is_latent (bool, optional) – If True, the input tensor is the latent tensor. Defaults to False.

• truncation (float, optional) – Truncation factor. Give value less than 1., the truncation trick will be adopted. Defaults to 1.

• num_truncation_layer (int, optional) – Number of layers use truncated latent. Defaults to None.

• update_emas (bool, optional) – Whether update moving average of mean latent. Defaults to False.

• force_fp32 (bool, optional) – Force fp32 ignore the weights. Defaults to True.

• return_noise (bool, optional) – If True, noise_batch will be returned in a dict with fake_img. Defaults to False.

• return_latents (bool, optional) – If True, latent will be returned in a dict with fake_img. Defaults to False.

Returns

Generated image tensor or dictionary containing more data.

Return type

torch.Tensor | dict

get_mean_latent(num_samples=4096, **kwargs)

Get mean latent of W space in this generator.

Parameters

num_samples (int, optional) – Number of sample times. Defaults to 4096.

Returns

Mean latent of this generator.

Return type

Tensor

get_training_kwargs(phase)

Get training kwargs. In StyleGANv3, we enable fp16, and update magnitude ema during training of discriminator. This function is used to pass related arguments.

Parameters

phase (str) – Current training phase.

Returns

Training kwargs.

Return type

dict