mmagic.models.editors.mspie.mspie_stylegan2_generator

Module Contents

Classes

MSStyleGANv2Generator

StyleGAN2 Generator.

class mmagic.models.editors.mspie.mspie_stylegan2_generator.MSStyleGANv2Generator(out_size, style_channels, num_mlps=8, channel_multiplier=2, blur_kernel=[1, 3, 3, 1], lr_mlp=0.01, default_style_mode='mix', eval_style_mode='single', mix_prob=0.9, no_pad=False, deconv2conv=False, interp_pad=None, up_config=dict(scale_factor=2, mode='nearest'), up_after_conv=False, head_pos_encoding=None, head_pos_size=(4, 4), interp_head=False)

Bases: mmengine.model.BaseModule

StyleGAN2 Generator.

In StyleGAN2, we use a static architecture composing of a style mapping module and number of convolutional style blocks. More details can be found in: Analyzing and Improving the Image Quality of StyleGAN CVPR2020.

Parameters

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

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

• num_mlps (int, optional) – The number of MLP layers. Defaults to 8.

• 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].

• lr_mlp (float, optional) – The learning rate for the style mapping layer. Defaults to 0.01.

• default_style_mode (str, optional) – The default mode of style mixing. In training, we adopt mixing style mode in default. However, in the evaluation, we use ‘single’ style mode. [‘mix’, ‘single’] are currently supported. Defaults to ‘mix’.

• eval_style_mode (str, optional) – The evaluation mode of style mixing. Defaults to ‘single’.

• mix_prob (float, optional) – Mixing probability. The value should be in range of [0, 1]. Defaults to 0.9.

train(mode=True)

Set train/eval mode.

Parameters

mode (bool, optional) – Whether set train mode. Defaults to True.

make_injected_noise(chosen_scale=0)

make noises that will be injected into feature maps.

Parameters

chosen_scale (int, optional) – Chosen scale. Defaults to 0.

Returns

List of layer-wise noise tensor.

Return type

list[Tensor]

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

style_mixing(n_source, n_target, inject_index=1, truncation_latent=None, truncation=0.7, chosen_scale=0)

Generating style mixing images.

Parameters

• n_source (int) – Number of source images.

• n_target (int) – Number of target images.

• inject_index (int, optional) – Index from which replace with source latent. Defaults to 1.

• truncation_latent (torch.Tensor, optional) – Mean truncation latent. Defaults to None.

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

• curr_scale (int) – Current image scale. Defaults to -1.

• transition_weight (float, optional) – The weight used in resolution transition. Defaults to 1.0.

• chosen_scale (int, optional) – Chosen scale. Defaults to 0.

Returns

Table of style-mixing images.

Return type

torch.Tensor

forward(styles, num_batches=- 1, return_noise=False, return_latents=False, inject_index=None, truncation=1, truncation_latent=None, input_is_latent=False, injected_noise=None, randomize_noise=True, chosen_scale=0)

Forward function.

This function has been integrated with the truncation trick. Please refer to the usage of truncation and truncation_latent.

Parameters

• styles (torch.Tensor | list[torch.Tensor] | callable | None) – In StyleGAN2, you can provide noise tensor or latent tensor. Given a list containing more than one noise or latent tensors, style mixing trick will be used in training. Of course, 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.

• 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.

• inject_index (int | None, optional) – The index number for mixing style codes. Defaults to None.

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

• truncation_latent (torch.Tensor, optional) – Mean truncation latent. Defaults to None.

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

• injected_noise (torch.Tensor | None, optional) – Given a tensor, the random noise will be fixed as this input injected noise. Defaults to None.

• randomize_noise (bool, optional) – If False, images are sampled with the buffered noise tensor injected to the style conv block. Defaults to True.

Returns

Generated image tensor or dictionary containing more data.

Return type

torch.Tensor | dict