Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
convert_stylegan.py 3.38 KB
Newer Older
mashun1's avatar
anytext  
mashun1 committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90 
import torch

from basicsr.archs.stylegan2_arch import StyleGAN2Discriminator, StyleGAN2Generator


def convert_net_g(ori_net, crt_net):
    """Convert network generator."""

    for crt_k, crt_v in crt_net.items():
        if 'style_mlp' in crt_k:
            ori_k = crt_k.replace('style_mlp', 'style')
        elif 'constant_input.weight' in crt_k:
            ori_k = crt_k.replace('constant_input.weight', 'input.input')
        # style conv1
        elif 'style_conv1.modulated_conv' in crt_k:
            ori_k = crt_k.replace('style_conv1.modulated_conv', 'conv1.conv')
        elif 'style_conv1' in crt_k:
            if crt_v.shape == torch.Size([1]):
                ori_k = crt_k.replace('style_conv1', 'conv1.noise')
            else:
                ori_k = crt_k.replace('style_conv1', 'conv1')
        # style conv
        elif 'style_convs' in crt_k:
            ori_k = crt_k.replace('style_convs', 'convs').replace('modulated_conv', 'conv')
            if crt_v.shape == torch.Size([1]):
                ori_k = ori_k.replace('.weight', '.noise.weight')
        # to_rgb1
        elif 'to_rgb1.modulated_conv' in crt_k:
            ori_k = crt_k.replace('to_rgb1.modulated_conv', 'to_rgb1.conv')
        # to_rgbs
        elif 'to_rgbs' in crt_k:
            ori_k = crt_k.replace('modulated_conv', 'conv')
        elif 'noises' in crt_k:
            ori_k = crt_k.replace('.noise', '.noise_')
        else:
            ori_k = crt_k

        # replace
        if crt_net[crt_k].size() != ori_net[ori_k].size():
            raise ValueError('Wrong tensor size: \n'
                             f'crt_net: {crt_net[crt_k].size()}\n'
                             f'ori_net: {ori_net[ori_k].size()}')
        else:
            crt_net[crt_k] = ori_net[ori_k]

    return crt_net


def convert_net_d(ori_net, crt_net):
    """Convert network discriminator."""

    for crt_k, _ in crt_net.items():
        if 'conv_body' in crt_k:
            ori_k = crt_k.replace('conv_body', 'convs')
        else:
            ori_k = crt_k

        # replace
        if crt_net[crt_k].size() != ori_net[ori_k].size():
            raise ValueError('Wrong tensor size: \n'
                             f'crt_net: {crt_net[crt_k].size()}\n'
                             f'ori_net: {ori_net[ori_k].size()}')
        else:
            crt_net[crt_k] = ori_net[ori_k]
    return crt_net


if __name__ == '__main__':
    """Convert official stylegan2 weights from stylegan2-pytorch."""

    # configuration
    ori_net = torch.load('experiments/pretrained_models/stylegan2-ffhq.pth')
    save_path_g = 'experiments/pretrained_models/stylegan2_ffhq_config_f_1024_official.pth'  # noqa: E501
    save_path_d = 'experiments/pretrained_models/stylegan2_ffhq_config_f_1024_discriminator_official.pth'  # noqa: E501
    out_size = 1024
    channel_multiplier = 1

    # convert generator
    crt_net = StyleGAN2Generator(out_size, num_style_feat=512, num_mlp=8, channel_multiplier=channel_multiplier)
    crt_net = crt_net.state_dict()

    crt_net_params_ema = convert_net_g(ori_net['g_ema'], crt_net)
    torch.save(dict(params_ema=crt_net_params_ema, latent_avg=ori_net['latent_avg']), save_path_g)

    # convert discriminator
    crt_net = StyleGAN2Discriminator(out_size, channel_multiplier=channel_multiplier)
    crt_net = crt_net.state_dict()

    crt_net_params = convert_net_d(ori_net['d'], crt_net)
    torch.save(dict(params=crt_net_params), save_path_d)