mmagic.models.editors.eg3d.eg3d_modules

Module Contents

Classes

TriPlaneBackbone	Tr-plane backbone for EG3D generator. This class is a wrapper of
SuperResolutionModule	Super resolution module for EG3D generator.
SynthesisBlock	Synthesis block for EG3D's SuperResolutionModule.

class mmagic.models.editors.eg3d.eg3d_modules.TriPlaneBackbone(noise_size: int, out_size: int, out_channels: int, num_mlps: int = 8, style_channels: int = 512, cond_size: Optional[int] = 25, cond_mapping_channels: Optional[int] = None, cond_scale: float = 0, zero_cond_input: bool = False, *args, **kwargs)

Bases: mmagic.models.editors.stylegan2.StyleGAN2Generator

Tr-plane backbone for EG3D generator. This class is a wrapper of StyleGAN2Generator.

Parameters

• noise_size (int, optional) – The size of (number of channels) the input noise. If not passed, will be set the same value as style_channels. Defaults to None.

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

• out_channels (int) – The number of channels for output.

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

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

• cond_size (int, optional) – The size of the conditional input. If not passed or less than 1, no conditional embedding will be used. Defaults to None.

• cond_mapping_channels (int, optional) – The channels of the conditional mapping layers. If not passed, will use the same value as style_channels. Defaults to None.

• cond_scale (float) – The scale factor is multiple by the conditional input. Defaults to 1.

• zero_cond_input (bool) – Whether use ‘zero tensor’ as the conditional input. Defaults to False.

• *args – Arguments for StyleGAN2Generator.

• **kwargs –

  Arguments for StyleGAN2Generator.

mapping(noise: torch.Tensor, label: Optional[torch.Tensor] = None, truncation: float = 1, num_truncation_layer: Optional[int] = None, update_ws: bool = True) → torch.Tensor

Mapping input noise (z) to style space (w).

Parameters

• noise (torch.Tensor) – Noise input.

• label (Optional[torch.Tensor]) – Conditional inputs. Defaults to None.

• 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_ws (bool) – Whether update latent code with EMA. Only work when w_avg is registered. Defaults to False.

Returns

Style codes after mapping.

Return type

torch.Tensor

synthesis(styles: torch.Tensor, *args, **kwargs) → torch.Tensor

Generate the Triplane feature.

Parameters

• styles (torch.Tensor) – The input style code.

• *args – Arguments for StyleGAN2Generator’s forward.

• **kwargs –

  Arguments for StyleGAN2Generator’s forward.

Returns

The generated Triplane feature.

Return type

torch.Tensor

class mmagic.models.editors.eg3d.eg3d_modules.SuperResolutionModule(in_channels: int, in_size: Optional[int] = None, hidden_size: Optional[int] = None, out_size: Optional[int] = None, hidden_channels: int = 128, out_channels: int = 64, style_channels: Optional[int] = 512, sr_antialias: bool = True, fp16_enable: bool = False)

Bases: mmengine.model.BaseModule

Super resolution module for EG3D generator.

Parameters

• in_channels (int) – The channels of the input feature.

• in_size (int) – The size of the input feature.

• hidden_size (int) – The size of the hidden feature. Only support hidden size equals to in_size or in_size times two. Defaults to None.

• out_size (int) – The size of the output image. Defaults to None.

• hidden_channels (int) – The channels of the hidden feature. Defaults to 64.

• style_channels (int) – The channels of the style code. Defaults to 512.

• sr_antialias (bool) – Whether use antialias interpolation method in upsampling. Defaults to True.

• fp16_enable (bool) – Whether enable fp16 in this module. Defaults to False.

forward(img: torch.Tensor, feature: torch.Tensor, styles: Union[torch.Tensor, List[torch.Tensor]], add_noise: bool = False) → torch.Tensor

Forward function.

Parameters

• img (torch.Tensor) – Image to super resolution.

• x (torch.Tensor) – Feature map of the input image.

• styles (torch.Tensor) – Style codes in w space.

• add_noise (bool, optional) – Whether add noise to image. Defaults to False.

Returns

Image after super resolution.

Return type

torch.Tensor

class mmagic.models.editors.eg3d.eg3d_modules.SynthesisBlock(in_channels: int, out_channels: int, style_channels: int, img_channels: int, upsample: bool = True, conv_clamp: int = 256, fp16_enabled: bool = False)

Bases: mmengine.model.BaseModule

Synthesis block for EG3D’s SuperResolutionModule.

Parameters

• in_channels (int) – The number of channels for the input feature.

• out_channels (int) – The number of channels for the output feature.

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

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

• upsample (bool) – Whether do upsampling. Defaults to True.

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

• fp16_enabled (bool) – Whether enable fp16. Defaults to False.

forward(x: torch.Tensor, img: torch.Tensor, styles: torch.Tensor, add_noise: bool = False) → Tuple[torch.Tensor]

Forward Synthesis block.

Parameters

• x (torch.Tensor) – Input feature.

• img (torch.Tensor) – Input image.

• styles (torch.Tensor) – Input style code.

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

Returns

Output feature and image.

Return type

Tuple[torch.Tensor]