mmagic.models.editors.glean.glean_styleganv2
¶
Module Contents¶
Classes¶
GLEAN (using StyleGANv2) architecture for super-resolution. |
|
Feature extractor composed of Residual-in-Residual Dense Blocks (RRDBs). |
- class mmagic.models.editors.glean.glean_styleganv2.GLEANStyleGANv2(in_size, out_size, img_channels=3, rrdb_channels=64, num_rrdbs=23, style_channels=512, 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, init_cfg=None, fp16_enabled=False, bgr2rgb=False)[source]¶
Bases:
mmengine.model.BaseModule
GLEAN (using StyleGANv2) architecture for super-resolution.
- Paper:
GLEAN: Generative Latent Bank for Large-Factor Image Super-Resolution, CVPR, 2021
This method makes use of StyleGAN2 and hence the arguments mostly follow that in ‘StyleGAN2v2Generator’.
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.
You can load pretrained model through passing information into
pretrained
argument. We have already offered official weights as follows:stylegan2-ffhq-config-f: http://download.openmmlab.com/mmediting/stylegan2/official_weights/stylegan2-ffhq-config-f-official_20210327_171224-bce9310c.pth # noqa
stylegan2-horse-config-f: http://download.openmmlab.com/mmediting/stylegan2/official_weights/stylegan2-horse-config-f-official_20210327_173203-ef3e69ca.pth # noqa
stylegan2-car-config-f: http://download.openmmlab.com/mmediting/stylegan2/official_weights/stylegan2-car-config-f-official_20210327_172340-8cfe053c.pth # noqa
stylegan2-cat-config-f: http://download.openmmlab.com/mmediting/stylegan2/official_weights/stylegan2-cat-config-f-official_20210327_172444-15bc485b.pth # noqa
stylegan2-church-config-f: http://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 generator = StyleGANv2Generator(1024, 512, pretrained=dict( ckpt_path=ckpt_http, prefix='generator_ema'))
Of course, you can also download the checkpoint in advance and set
ckpt_path
with local path. If you just want to load the original generator (not the ema model), please set the prefix with ‘generator’.Note that our implementation allows to generate BGR image, while the original StyleGAN2 outputs RGB images by default. Thus, we provide
bgr2rgb
argument to convert the image space.- Parameters
in_size (int) – The size of the input image.
out_size (int) – The output size of the StyleGAN2 generator.
img_channels (int) – Number of channels of the input images. 3 for RGB image and 1 for grayscale image. Default: 3.
rrdb_channels (int) – Number of channels of the RRDB features. Default: 64.
num_rrdbs (int) – Number of RRDB blocks in the encoder. Default: 23.
style_channels (int) – The number of channels for style code. Default: 512.
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.
init_cfg (dict, optional) – Initialization config dict. Default: None.
fp16_enabled (bool, optional) – Whether to use fp16 training in this module. Defaults to False.
bgr2rgb (bool, optional) – Whether to flip the image channel dimension. Defaults to False.
- class mmagic.models.editors.glean.glean_styleganv2.RRDBFeatureExtractor(in_channels=3, mid_channels=64, num_blocks=23, growth_channels=32)[source]¶
Bases:
torch.nn.Module
Feature extractor composed of Residual-in-Residual Dense Blocks (RRDBs).
It is equivalent to ESRGAN with the upsampling module removed.
- Parameters
in_channels (int) – Channel number of inputs. Default: 3.
mid_channels (int) – Channel number of intermediate features. Default: 64
num_blocks (int) – Block number in the trunk network. Default: 23.
growth_channels (int) – Channels for each growth. Default: 32.