278 lines
9.6 KiB
Python
278 lines
9.6 KiB
Python
import os
|
|
import tensorflow as tf
|
|
|
|
from load_models import load_generator, load_discriminator
|
|
|
|
|
|
def handle_mapping(w_name, is_g_clone):
|
|
def extract_info(name):
|
|
splitted = name.split('/')
|
|
index = splitted.index('g_mapping')
|
|
indicator = splitted[index + 1]
|
|
val = indicator.split('_')[-1]
|
|
return val
|
|
|
|
level = extract_info(w_name)
|
|
o_prefix = f'G_mapping_1' if is_g_clone else f'G_mapping'
|
|
if 'w' in w_name:
|
|
official_var_name = f'{o_prefix}/Dense{level}/weight'
|
|
else:
|
|
official_var_name = f'{o_prefix}/Dense{level}/bias'
|
|
return official_var_name
|
|
|
|
|
|
def to_rgb_layer(name, r, o_prefix):
|
|
if 'conv/w:0' in name:
|
|
o_name = f'{o_prefix}/{r}x{r}/ToRGB/weight'
|
|
elif 'mod_dense/w:0' in name:
|
|
o_name = f'{o_prefix}/{r}x{r}/ToRGB/mod_weight'
|
|
elif 'mod_bias/b:0' in name:
|
|
o_name = f'{o_prefix}/{r}x{r}/ToRGB/mod_bias'
|
|
else: # if 'bias/b:0' in name:
|
|
o_name = f'{o_prefix}/{r}x{r}/ToRGB/bias'
|
|
return o_name
|
|
|
|
|
|
def handle_block_layer(name, r, o_prefix):
|
|
if 'conv_0/w:0' in name:
|
|
o_name = f'{o_prefix}/{r}x{r}/Conv0_up/weight'
|
|
elif 'conv_0/mod_dense/w:0' in name:
|
|
o_name = f'{o_prefix}/{r}x{r}/Conv0_up/mod_weight'
|
|
elif 'conv_0/mod_bias/b:0' in name:
|
|
o_name = f'{o_prefix}/{r}x{r}/Conv0_up/mod_bias'
|
|
elif 'noise_0/w:0' in name:
|
|
o_name = f'{o_prefix}/{r}x{r}/Conv0_up/noise_strength'
|
|
elif 'bias_0/b:0' in name:
|
|
o_name = f'{o_prefix}/{r}x{r}/Conv0_up/bias'
|
|
elif 'conv_1/w:0' in name:
|
|
o_name = f'{o_prefix}/{r}x{r}/Conv1/weight'
|
|
elif 'conv_1/mod_dense/w:0' in name:
|
|
o_name = f'{o_prefix}/{r}x{r}/Conv1/mod_weight'
|
|
elif 'conv_1/mod_bias/b:0' in name:
|
|
o_name = f'{o_prefix}/{r}x{r}/Conv1/mod_bias'
|
|
elif 'noise_1/w:0' in name:
|
|
o_name = f'{o_prefix}/{r}x{r}/Conv1/noise_strength'
|
|
else: # if 'bias_1/b:0' in name:
|
|
o_name = f'{o_prefix}/{r}x{r}/Conv1/bias'
|
|
return o_name
|
|
|
|
|
|
def handle_const_layer(name, o_prefix):
|
|
if 'const:0' in name:
|
|
o_name = f'{o_prefix}/4x4/Const/const'
|
|
elif 'conv/w:0' in name:
|
|
o_name = f'{o_prefix}/4x4/Conv/weight'
|
|
elif 'mod_dense/w:0' in name:
|
|
o_name = f'{o_prefix}/4x4/Conv/mod_weight'
|
|
elif 'mod_bias/b:0' in name:
|
|
o_name = f'{o_prefix}/4x4/Conv/mod_bias'
|
|
elif 'noise/w:0' in name:
|
|
o_name = f'{o_prefix}/4x4/Conv/noise_strength'
|
|
else:
|
|
o_name = f'{o_prefix}/4x4/Conv/bias'
|
|
return o_name
|
|
|
|
|
|
def handle_synthesis(w_name, is_g_clone):
|
|
def extract_info(name):
|
|
splitted = name.split('/')
|
|
index = splitted.index('g_synthesis')
|
|
indicator1 = splitted[index + 1]
|
|
indicator2 = splitted[index + 2]
|
|
r = indicator1.split('x')[1]
|
|
d = indicator2
|
|
return r, d
|
|
|
|
res, divider = extract_info(w_name)
|
|
o_prefix = f'G_synthesis_1' if is_g_clone else f'G_synthesis'
|
|
if divider == 'ToRGB':
|
|
official_var_name = to_rgb_layer(w_name, res, o_prefix)
|
|
elif divider == 'block':
|
|
official_var_name = handle_block_layer(w_name, res, o_prefix)
|
|
else: # const
|
|
official_var_name = handle_const_layer(w_name, o_prefix)
|
|
return official_var_name
|
|
|
|
|
|
def handle_discriminator_layer(w_name):
|
|
def extract_info(name):
|
|
splitted = name.split('/')
|
|
resolution = splitted[1]
|
|
resolution = resolution.split('x')[0]
|
|
return resolution
|
|
|
|
res = extract_info(w_name)
|
|
if 'last_dense' in w_name:
|
|
o_name = 'D/Output/weight'
|
|
elif 'last_bias' in w_name:
|
|
o_name = 'D/Output/bias'
|
|
elif 'FromRGB' in w_name:
|
|
o_name = 'D/1024x1024/FromRGB/weight' if 'conv_' in w_name else 'D/1024x1024/FromRGB/bias'
|
|
elif 'skip' in w_name:
|
|
o_name = f'D/{res}x{res}/Skip/weight'
|
|
elif 'dense_' in w_name:
|
|
o_name = f'D/4x4/Dense0/weight'
|
|
elif 'conv_' in w_name:
|
|
if res == '4':
|
|
o_name = f'D/{res}x{res}/Conv/weight'
|
|
else:
|
|
o_name = f'D/{res}x{res}/Conv0/weight' if '_0' in w_name else f'D/{res}x{res}/Conv1_down/weight'
|
|
elif 'bias_' in w_name:
|
|
if res == '4':
|
|
o_name = f'D/{res}x{res}/Conv/bias' if '_0' in w_name else f'D/{res}x{res}/Dense0/bias'
|
|
else:
|
|
o_name = f'D/{res}x{res}/Conv0/bias' if '_0' in w_name else f'D/{res}x{res}/Conv1_down/bias'
|
|
else:
|
|
raise ValueError('Something went wrong!!')
|
|
return o_name
|
|
|
|
|
|
def variable_name_mapper_g(g, is_g_clone):
|
|
name_mapper = dict()
|
|
for w in g.weights:
|
|
w_name, w_shape = w.name, w.shape
|
|
|
|
# mapping layer
|
|
if 'g_mapping' in w_name:
|
|
official_var_name = handle_mapping(w_name, is_g_clone)
|
|
elif 'g_synthesis' in w_name:
|
|
official_var_name = handle_synthesis(w_name, is_g_clone)
|
|
else:
|
|
# w_avg
|
|
official_var_name = 'Gs/dlatent_avg' if is_g_clone else 'G/dlatent_avg'
|
|
|
|
name_mapper[official_var_name] = w
|
|
return name_mapper
|
|
|
|
|
|
def variable_name_mapper_d(d):
|
|
name_mapper = dict()
|
|
for w in d.weights:
|
|
w_name, w_shape = w.name, w.shape
|
|
|
|
official_var_name = handle_discriminator_layer(w_name)
|
|
name_mapper[official_var_name] = w
|
|
return name_mapper
|
|
|
|
|
|
def check_shape(name_mapper, official_vars):
|
|
for official_name, v in name_mapper.items():
|
|
official_shape = [s for n, s in official_vars if n == official_name][0]
|
|
|
|
if official_shape == v.shape:
|
|
print('{}: shape matches'.format(official_name))
|
|
else:
|
|
# print(f'Official: {official_name} -> {official_shape}')
|
|
# print(f'Current: {v.name} -> {v.shape}')
|
|
raise ValueError('{}: wrong shape'.format(official_name))
|
|
return
|
|
|
|
|
|
def convert_official_generator_weights(ckpt_dir, is_g_clone, use_custom_cuda):
|
|
generator = load_generator(g_params=None, is_g_clone=is_g_clone, ckpt_dir=None, custom_cuda=use_custom_cuda)
|
|
|
|
# restore official ones to current implementation
|
|
official_checkpoint = tf.train.latest_checkpoint('./official-pretrained')
|
|
official_vars = tf.train.list_variables(official_checkpoint)
|
|
|
|
# get name mapper
|
|
name_mapper = variable_name_mapper_g(generator, is_g_clone=is_g_clone)
|
|
for name_g, tvar in name_mapper.items():
|
|
print(f'{name_g}: {tvar.name}')
|
|
|
|
# check shape
|
|
check_shape(name_mapper, official_vars)
|
|
|
|
# restore
|
|
tf.compat.v1.train.init_from_checkpoint(official_checkpoint, assignment_map=name_mapper)
|
|
|
|
# save
|
|
if is_g_clone:
|
|
ckpt = tf.train.Checkpoint(g_clone=generator)
|
|
else:
|
|
ckpt = tf.train.Checkpoint(generator=generator)
|
|
out_dir = os.path.join(ckpt_dir, 'g_clone' if is_g_clone else 'generator')
|
|
manager = tf.train.CheckpointManager(ckpt, out_dir, max_to_keep=1)
|
|
manager.save(checkpoint_number=0)
|
|
return
|
|
|
|
|
|
def convert_official_discriminator_weights(ckpt_dir, use_custom_cuda):
|
|
discriminator = load_discriminator(d_params=None, ckpt_dir=None, custom_cuda=use_custom_cuda)
|
|
|
|
# restore official ones
|
|
official_checkpoint = tf.train.latest_checkpoint('./official-pretrained')
|
|
official_vars = tf.train.list_variables(official_checkpoint)
|
|
|
|
# get name mapper
|
|
name_mapper = variable_name_mapper_d(discriminator)
|
|
for name_d, tvar in name_mapper.items():
|
|
print(f'{name_d}: {tvar.name}')
|
|
|
|
# check shape
|
|
check_shape(name_mapper, official_vars)
|
|
|
|
# restore
|
|
tf.compat.v1.train.init_from_checkpoint(official_checkpoint, assignment_map=name_mapper)
|
|
|
|
# save
|
|
ckpt = tf.train.Checkpoint(discriminator=discriminator)
|
|
out_dir = os.path.join(ckpt_dir, 'discriminator')
|
|
manager = tf.train.CheckpointManager(ckpt, out_dir, max_to_keep=1)
|
|
manager.save(checkpoint_number=0)
|
|
return
|
|
|
|
|
|
def convert_official_weights_together(ckpt_dir, use_custom_cuda):
|
|
# instantiate all models
|
|
discriminator = load_discriminator(d_params=None, ckpt_dir=None, custom_cuda=use_custom_cuda)
|
|
generator = load_generator(g_params=None, is_g_clone=False, ckpt_dir=None, custom_cuda=use_custom_cuda)
|
|
g_clone = load_generator(g_params=None, is_g_clone=True, ckpt_dir=None, custom_cuda=use_custom_cuda)
|
|
|
|
# restore official ones
|
|
official_checkpoint = tf.train.latest_checkpoint('./official-pretrained')
|
|
official_vars = tf.train.list_variables(official_checkpoint)
|
|
for name, shape in official_vars:
|
|
print(f'{name}: {shape}')
|
|
|
|
# get name mapper
|
|
name_mapper_d = variable_name_mapper_d(discriminator)
|
|
name_mapper_g1 = variable_name_mapper_g(generator, is_g_clone=False)
|
|
name_mapper_g2 = variable_name_mapper_g(g_clone, is_g_clone=True)
|
|
name_mapper = {**name_mapper_d, **name_mapper_g1, **name_mapper_g2}
|
|
|
|
# check shape
|
|
check_shape(name_mapper, official_vars)
|
|
|
|
# restore
|
|
tf.compat.v1.train.init_from_checkpoint(official_checkpoint, assignment_map=name_mapper)
|
|
|
|
# save
|
|
ckpt = tf.train.Checkpoint(discriminator=discriminator,
|
|
generator=generator,
|
|
g_clone=g_clone)
|
|
manager = tf.train.CheckpointManager(ckpt, ckpt_dir, max_to_keep=1)
|
|
manager.save(checkpoint_number=0)
|
|
return
|
|
|
|
|
|
def main():
|
|
from tf_utils import allow_memory_growth
|
|
|
|
allow_memory_growth()
|
|
|
|
ckpt_dir_base = './official-converted'
|
|
for use_custom_cuda in [True, False]:
|
|
ckpt_dir = os.path.join(ckpt_dir_base, 'cuda') if use_custom_cuda else os.path.join(ckpt_dir_base, 'ref')
|
|
|
|
convert_official_weights_together(ckpt_dir, use_custom_cuda)
|
|
|
|
# convert_official_discriminator_weights(ckpt_dir, use_custom_cuda)
|
|
# for is_g_clone in [True, False]:
|
|
# convert_official_generator_weights(ckpt_dir, is_g_clone, use_custom_cuda)
|
|
return
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|