lambda/generative-music-visualizer
4
1# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.2#3# NVIDIA CORPORATION and its licensors retain all intellectual property4# and proprietary rights in and to this software, related documentation5# and any modifications thereto. Any use, reproduction, disclosure or6# distribution of this software and related documentation without an express7# license agreement from NVIDIA CORPORATION is strictly prohibited.8 9"""Converting legacy network pickle into the new format."""10 11import click12import pickle13import re14import copy15import numpy as np16import torch17import dnnlib18from torch_utils import misc19 20#----------------------------------------------------------------------------21 22def load_network_pkl(f, force_fp16=False):23 data = _LegacyUnpickler(f).load()24 25 # Legacy TensorFlow pickle => convert.26 if isinstance(data, tuple) and len(data) == 3 and all(isinstance(net, _TFNetworkStub) for net in data):27 tf_G, tf_D, tf_Gs = data28 G = convert_tf_generator(tf_G)29 D = convert_tf_discriminator(tf_D)30 G_ema = convert_tf_generator(tf_Gs)31 data = dict(G=G, D=D, G_ema=G_ema)32 33 # Add missing fields.34 if 'training_set_kwargs' not in data:35 data['training_set_kwargs'] = None36 if 'augment_pipe' not in data:37 data['augment_pipe'] = None38 39 # Validate contents.40 assert isinstance(data['G'], torch.nn.Module)41 assert isinstance(data['D'], torch.nn.Module)42 assert isinstance(data['G_ema'], torch.nn.Module)43 assert isinstance(data['training_set_kwargs'], (dict, type(None)))44 assert isinstance(data['augment_pipe'], (torch.nn.Module, type(None)))45 46 # Force FP16.47 if force_fp16:48 for key in ['G', 'D', 'G_ema']:49 old = data[key]50 kwargs = copy.deepcopy(old.init_kwargs)51 fp16_kwargs = kwargs.get('synthesis_kwargs', kwargs)52 fp16_kwargs.num_fp16_res = 453 fp16_kwargs.conv_clamp = 25654 if kwargs != old.init_kwargs:55 new = type(old)(**kwargs).eval().requires_grad_(False)56 misc.copy_params_and_buffers(old, new, require_all=True)57 data[key] = new58 return data59 60#----------------------------------------------------------------------------61 62class _TFNetworkStub(dnnlib.EasyDict):63 pass64 65class _LegacyUnpickler(pickle.Unpickler):66 def find_class(self, module, name):67 if module == 'dnnlib.tflib.network' and name == 'Network':68 return _TFNetworkStub69 return super().find_class(module, name)70 71#----------------------------------------------------------------------------72 73def _collect_tf_params(tf_net):74 # pylint: disable=protected-access75 tf_params = dict()76 def recurse(prefix, tf_net):77 for name, value in tf_net.variables:78 tf_params[prefix + name] = value79 for name, comp in tf_net.components.items():80 recurse(prefix + name + '/', comp)81 recurse('', tf_net)82 return tf_params83 84#----------------------------------------------------------------------------85 86def _populate_module_params(module, *patterns):87 for name, tensor in misc.named_params_and_buffers(module):88 found = False89 value = None90 for pattern, value_fn in zip(patterns[0::2], patterns[1::2]):91 match = re.fullmatch(pattern, name)92 if match:93 found = True94 if value_fn is not None:95 value = value_fn(*match.groups())96 break97 try:98 assert found99 if value is not None:100 tensor.copy_(torch.from_numpy(np.array(value)))101 except:102 print(name, list(tensor.shape))103 raise104 105#----------------------------------------------------------------------------106 107def convert_tf_generator(tf_G):108 if tf_G.version < 4:109 raise ValueError('TensorFlow pickle version too low')110 111 # Collect kwargs.112 tf_kwargs = tf_G.static_kwargs113 known_kwargs = set()114 def kwarg(tf_name, default=None, none=None):115 known_kwargs.add(tf_name)116 val = tf_kwargs.get(tf_name, default)117 return val if val is not None else none118 119 # Convert kwargs.120 from training import networks_stylegan2121 network_class = networks_stylegan2.Generator122 kwargs = dnnlib.EasyDict(123 z_dim = kwarg('latent_size', 512),124 c_dim = kwarg('label_size', 0),125 w_dim = kwarg('dlatent_size', 512),126 img_resolution = kwarg('resolution', 1024),127 img_channels = kwarg('num_channels', 3),128 channel_base = kwarg('fmap_base', 16384) * 2,129 channel_max = kwarg('fmap_max', 512),130 num_fp16_res = kwarg('num_fp16_res', 0),131 conv_clamp = kwarg('conv_clamp', None),132 architecture = kwarg('architecture', 'skip'),133 resample_filter = kwarg('resample_kernel', [1,3,3,1]),134 use_noise = kwarg('use_noise', True),135 activation = kwarg('nonlinearity', 'lrelu'),136 mapping_kwargs = dnnlib.EasyDict(137 num_layers = kwarg('mapping_layers', 8),138 embed_features = kwarg('label_fmaps', None),139 layer_features = kwarg('mapping_fmaps', None),140 activation = kwarg('mapping_nonlinearity', 'lrelu'),141 lr_multiplier = kwarg('mapping_lrmul', 0.01),142 w_avg_beta = kwarg('w_avg_beta', 0.995, none=1),143 ),144 )145 146 # Check for unknown kwargs.147 kwarg('truncation_psi')148 kwarg('truncation_cutoff')149 kwarg('style_mixing_prob')150 kwarg('structure')151 kwarg('conditioning')152 kwarg('fused_modconv')153 unknown_kwargs = list(set(tf_kwargs.keys()) - known_kwargs)154 if len(unknown_kwargs) > 0:155 raise ValueError('Unknown TensorFlow kwarg', unknown_kwargs[0])156 157 # Collect params.158 tf_params = _collect_tf_params(tf_G)159 for name, value in list(tf_params.items()):160 match = re.fullmatch(r'ToRGB_lod(\d+)/(.*)', name)161 if match:162 r = kwargs.img_resolution // (2 ** int(match.group(1)))163 tf_params[f'{r}x{r}/ToRGB/{match.group(2)}'] = value164 kwargs.synthesis.kwargs.architecture = 'orig'165 #for name, value in tf_params.items(): print(f'{name:<50s}{list(value.shape)}')166 167 # Convert params.168 G = network_class(**kwargs).eval().requires_grad_(False)169 # pylint: disable=unnecessary-lambda170 # pylint: disable=f-string-without-interpolation171 _populate_module_params(G,172 r'mapping\.w_avg', lambda: tf_params[f'dlatent_avg'],173 r'mapping\.embed\.weight', lambda: tf_params[f'mapping/LabelEmbed/weight'].transpose(),174 r'mapping\.embed\.bias', lambda: tf_params[f'mapping/LabelEmbed/bias'],175 r'mapping\.fc(\d+)\.weight', lambda i: tf_params[f'mapping/Dense{i}/weight'].transpose(),176 r'mapping\.fc(\d+)\.bias', lambda i: tf_params[f'mapping/Dense{i}/bias'],177 r'synthesis\.b4\.const', lambda: tf_params[f'synthesis/4x4/Const/const'][0],178 r'synthesis\.b4\.conv1\.weight', lambda: tf_params[f'synthesis/4x4/Conv/weight'].transpose(3, 2, 0, 1),179 r'synthesis\.b4\.conv1\.bias', lambda: tf_params[f'synthesis/4x4/Conv/bias'],180 r'synthesis\.b4\.conv1\.noise_const', lambda: tf_params[f'synthesis/noise0'][0, 0],181 r'synthesis\.b4\.conv1\.noise_strength', lambda: tf_params[f'synthesis/4x4/Conv/noise_strength'],182 r'synthesis\.b4\.conv1\.affine\.weight', lambda: tf_params[f'synthesis/4x4/Conv/mod_weight'].transpose(),183 r'synthesis\.b4\.conv1\.affine\.bias', lambda: tf_params[f'synthesis/4x4/Conv/mod_bias'] + 1,184 r'synthesis\.b(\d+)\.conv0\.weight', lambda r: tf_params[f'synthesis/{r}x{r}/Conv0_up/weight'][::-1, ::-1].transpose(3, 2, 0, 1),185 r'synthesis\.b(\d+)\.conv0\.bias', lambda r: tf_params[f'synthesis/{r}x{r}/Conv0_up/bias'],186 r'synthesis\.b(\d+)\.conv0\.noise_const', lambda r: tf_params[f'synthesis/noise{int(np.log2(int(r)))*2-5}'][0, 0],187 r'synthesis\.b(\d+)\.conv0\.noise_strength', lambda r: tf_params[f'synthesis/{r}x{r}/Conv0_up/noise_strength'],188 r'synthesis\.b(\d+)\.conv0\.affine\.weight', lambda r: tf_params[f'synthesis/{r}x{r}/Conv0_up/mod_weight'].transpose(),189 r'synthesis\.b(\d+)\.conv0\.affine\.bias', lambda r: tf_params[f'synthesis/{r}x{r}/Conv0_up/mod_bias'] + 1,190 r'synthesis\.b(\d+)\.conv1\.weight', lambda r: tf_params[f'synthesis/{r}x{r}/Conv1/weight'].transpose(3, 2, 0, 1),191 r'synthesis\.b(\d+)\.conv1\.bias', lambda r: tf_params[f'synthesis/{r}x{r}/Conv1/bias'],192 r'synthesis\.b(\d+)\.conv1\.noise_const', lambda r: tf_params[f'synthesis/noise{int(np.log2(int(r)))*2-4}'][0, 0],193 r'synthesis\.b(\d+)\.conv1\.noise_strength', lambda r: tf_params[f'synthesis/{r}x{r}/Conv1/noise_strength'],194 r'synthesis\.b(\d+)\.conv1\.affine\.weight', lambda r: tf_params[f'synthesis/{r}x{r}/Conv1/mod_weight'].transpose(),195 r'synthesis\.b(\d+)\.conv1\.affine\.bias', lambda r: tf_params[f'synthesis/{r}x{r}/Conv1/mod_bias'] + 1,196 r'synthesis\.b(\d+)\.torgb\.weight', lambda r: tf_params[f'synthesis/{r}x{r}/ToRGB/weight'].transpose(3, 2, 0, 1),197 r'synthesis\.b(\d+)\.torgb\.bias', lambda r: tf_params[f'synthesis/{r}x{r}/ToRGB/bias'],198 r'synthesis\.b(\d+)\.torgb\.affine\.weight', lambda r: tf_params[f'synthesis/{r}x{r}/ToRGB/mod_weight'].transpose(),199 r'synthesis\.b(\d+)\.torgb\.affine\.bias', lambda r: tf_params[f'synthesis/{r}x{r}/ToRGB/mod_bias'] + 1,200 r'synthesis\.b(\d+)\.skip\.weight', lambda r: tf_params[f'synthesis/{r}x{r}/Skip/weight'][::-1, ::-1].transpose(3, 2, 0, 1),201 r'.*\.resample_filter', None,202 r'.*\.act_filter', None,203 )204 return G205 206#----------------------------------------------------------------------------207 208def convert_tf_discriminator(tf_D):209 if tf_D.version < 4:210 raise ValueError('TensorFlow pickle version too low')211 212 # Collect kwargs.213 tf_kwargs = tf_D.static_kwargs214 known_kwargs = set()215 def kwarg(tf_name, default=None):216 known_kwargs.add(tf_name)217 return tf_kwargs.get(tf_name, default)218 219 # Convert kwargs.220 kwargs = dnnlib.EasyDict(221 c_dim = kwarg('label_size', 0),222 img_resolution = kwarg('resolution', 1024),223 img_channels = kwarg('num_channels', 3),224 architecture = kwarg('architecture', 'resnet'),225 channel_base = kwarg('fmap_base', 16384) * 2,226 channel_max = kwarg('fmap_max', 512),227 num_fp16_res = kwarg('num_fp16_res', 0),228 conv_clamp = kwarg('conv_clamp', None),229 cmap_dim = kwarg('mapping_fmaps', None),230 block_kwargs = dnnlib.EasyDict(231 activation = kwarg('nonlinearity', 'lrelu'),232 resample_filter = kwarg('resample_kernel', [1,3,3,1]),233 freeze_layers = kwarg('freeze_layers', 0),234 ),235 mapping_kwargs = dnnlib.EasyDict(236 num_layers = kwarg('mapping_layers', 0),237 embed_features = kwarg('mapping_fmaps', None),238 layer_features = kwarg('mapping_fmaps', None),239 activation = kwarg('nonlinearity', 'lrelu'),240 lr_multiplier = kwarg('mapping_lrmul', 0.1),241 ),242 epilogue_kwargs = dnnlib.EasyDict(243 mbstd_group_size = kwarg('mbstd_group_size', None),244 mbstd_num_channels = kwarg('mbstd_num_features', 1),245 activation = kwarg('nonlinearity', 'lrelu'),246 ),247 )248 249 # Check for unknown kwargs.250 kwarg('structure')251 kwarg('conditioning')252 unknown_kwargs = list(set(tf_kwargs.keys()) - known_kwargs)253 if len(unknown_kwargs) > 0:254 raise ValueError('Unknown TensorFlow kwarg', unknown_kwargs[0])255 256 # Collect params.257 tf_params = _collect_tf_params(tf_D)258 for name, value in list(tf_params.items()):259 match = re.fullmatch(r'FromRGB_lod(\d+)/(.*)', name)260 if match:261 r = kwargs.img_resolution // (2 ** int(match.group(1)))262 tf_params[f'{r}x{r}/FromRGB/{match.group(2)}'] = value263 kwargs.architecture = 'orig'264 #for name, value in tf_params.items(): print(f'{name:<50s}{list(value.shape)}')265 266 # Convert params.267 from training import networks_stylegan2268 D = networks_stylegan2.Discriminator(**kwargs).eval().requires_grad_(False)269 # pylint: disable=unnecessary-lambda270 # pylint: disable=f-string-without-interpolation271 _populate_module_params(D,272 r'b(\d+)\.fromrgb\.weight', lambda r: tf_params[f'{r}x{r}/FromRGB/weight'].transpose(3, 2, 0, 1),273 r'b(\d+)\.fromrgb\.bias', lambda r: tf_params[f'{r}x{r}/FromRGB/bias'],274 r'b(\d+)\.conv(\d+)\.weight', lambda r, i: tf_params[f'{r}x{r}/Conv{i}{["","_down"][int(i)]}/weight'].transpose(3, 2, 0, 1),275 r'b(\d+)\.conv(\d+)\.bias', lambda r, i: tf_params[f'{r}x{r}/Conv{i}{["","_down"][int(i)]}/bias'],276 r'b(\d+)\.skip\.weight', lambda r: tf_params[f'{r}x{r}/Skip/weight'].transpose(3, 2, 0, 1),277 r'mapping\.embed\.weight', lambda: tf_params[f'LabelEmbed/weight'].transpose(),278 r'mapping\.embed\.bias', lambda: tf_params[f'LabelEmbed/bias'],279 r'mapping\.fc(\d+)\.weight', lambda i: tf_params[f'Mapping{i}/weight'].transpose(),280 r'mapping\.fc(\d+)\.bias', lambda i: tf_params[f'Mapping{i}/bias'],281 r'b4\.conv\.weight', lambda: tf_params[f'4x4/Conv/weight'].transpose(3, 2, 0, 1),282 r'b4\.conv\.bias', lambda: tf_params[f'4x4/Conv/bias'],283 r'b4\.fc\.weight', lambda: tf_params[f'4x4/Dense0/weight'].transpose(),284 r'b4\.fc\.bias', lambda: tf_params[f'4x4/Dense0/bias'],285 r'b4\.out\.weight', lambda: tf_params[f'Output/weight'].transpose(),286 r'b4\.out\.bias', lambda: tf_params[f'Output/bias'],287 r'.*\.resample_filter', None,288 )289 return D290 291#----------------------------------------------------------------------------292 293@click.command()294@click.option('--source', help='Input pickle', required=True, metavar='PATH')295@click.option('--dest', help='Output pickle', required=True, metavar='PATH')296@click.option('--force-fp16', help='Force the networks to use FP16', type=bool, default=False, metavar='BOOL', show_default=True)297def convert_network_pickle(source, dest, force_fp16):298 """Convert legacy network pickle into the native PyTorch format.299 300 The tool is able to load the main network configurations exported using the TensorFlow version of StyleGAN2 or StyleGAN2-ADA.301 It does not support e.g. StyleGAN2-ADA comparison methods, StyleGAN2 configs A-D, or StyleGAN1 networks.302 303 Example:304 305 \b306 python legacy.py \\307 --source=https://nvlabs-fi-cdn.nvidia.com/stylegan2/networks/stylegan2-cat-config-f.pkl \\308 --dest=stylegan2-cat-config-f.pkl309 """310 print(f'Loading "{source}"...')311 with dnnlib.util.open_url(source) as f:312 data = load_network_pkl(f, force_fp16=force_fp16)313 print(f'Saving "{dest}"...')314 with open(dest, 'wb') as f:315 pickle.dump(data, f)316 print('Done.')317 318#----------------------------------------------------------------------------319 320if __name__ == "__main__":321 convert_network_pickle() # pylint: disable=no-value-for-parameter322 323#----------------------------------------------------------------------------324