Module Contents



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


Super resolution module for EG3D generator.


Synthesis block for EG3D's SuperResolutionModule.

class mmedit.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)[source]

Bases: mmedit.models.editors.stylegan2.StyleGAN2Generator

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

  • 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[source]

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

  • 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 registeried. Defaults to False.


Style codes after mapping.

Return type


synthesis(styles: torch.Tensor, *args, **kwargs) torch.Tensor[source]

Generate the Triplane feature.

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

  • *args – Arguments for StyleGAN2Generator’s forward.

  • **kwargs

    Arguments for StyleGAN2Generator’s forward.


The generated Triplane feature.

Return type


class mmedit.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)[source]

Bases: mmengine.model.BaseModule

Super resolution module for EG3D generator.

  • 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[source]

Forward function.

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


Image after super resolution.

Return type


class mmedit.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)[source]

Bases: mmengine.model.BaseModule

Synthesis block for EG3D’s SuperResolutionModule.

  • 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][source]

Forward Synthesis block.

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


Output feature and image.

Return type


Read the Docs v: latest
On Read the Docs
Project Home

Free document hosting provided by Read the Docs.