def convert_stylegan2ada_pth_weight(src_weight_path,
                                    dst_weight_path,
                                    test_num=10,
                                    save_test_image=False,
                                    verbose=False):
    """Converts the pre-trained StyleGAN2-ADA-PyTorch weights.

    Args:
        src_weight_path: Path to the source model to load weights from.
        dst_weight_path: Path to the target model to save converted weights.
        test_num: Number of samples used to test the conversion. (default: 10)
        save_test_image: Whether to save the test images. (default: False)
        verbose: Whether to print verbose log message. (default: False)
    """

    print(f'========================================')
    print(f'Loading source weights from `{src_weight_path}` ...')
    sys.path.insert(0, CODE_PATH)
    with open(src_weight_path, 'rb') as f:
        model = pickle.load(f)
    sys.path.pop(0)
    print(f'Successfully loaded!')
    print(f'--------------------')

    z_space_dim = model['G'].z_dim
    label_size = model['G'].c_dim
    w_space_dim = model['G'].w_dim
    image_channels = model['G'].img_channels
    resolution = model['G'].img_resolution
    repeat_w = True
    fmaps_base = model['G']._init_kwargs['synthesis_kwargs'].channel_base
    mapping_layers = model['G']._init_kwargs['mapping_kwargs'].num_layers


    print(f'Converting source weights (G) to target ...')
    G_vars = dict(model['G'].named_parameters())
    G_vars.update(dict(model['G'].named_buffers()))
    G = build_model(gan_type=GAN_TPYE,
                    module='generator',
                    resolution=resolution,
                    z_space_dim=z_space_dim,
                    w_space_dim=w_space_dim,
                    label_size=label_size,
                    repeat_w=repeat_w,
                    image_channels=image_channels,
                    fmaps_base=fmaps_base,
                    mapping_layers=mapping_layers)
    G_state_dict = G.state_dict()
    official_tf_to_pth_var_mapping = {}
    for name in G_vars.keys():
        for pattern, fn in G_PTH_TO_TF_VAR_MAPPING_PATTERN.items():
            match = re.fullmatch(pattern, name)
            if match:
                if fn is not None:
                    official_tf_to_pth_var_mapping[fn(*match.groups())] = name
                break
    for dst_var_name, tf_var_name in G.pth_to_tf_var_mapping.items():
        assert tf_var_name in official_tf_to_pth_var_mapping
        assert dst_var_name in G_state_dict
        src_var_name = official_tf_to_pth_var_mapping[tf_var_name]
        assert src_var_name in G_vars
        if verbose:
            print(f'    Converting `{src_var_name}` to `{dst_var_name}`.')
        var = G_vars[src_var_name].data
        if 'weight' in tf_var_name:
            if 'Conv0_up/weight' in tf_var_name:
                var = var.flip(2, 3)
            elif 'Skip' in tf_var_name:
                var = var.flip(2, 3)
        if 'bias' in tf_var_name:
            if 'mod_bias' in tf_var_name:
                var = var - 1
        if 'Const' in tf_var_name:
            var = var.unsqueeze(0)
        if 'noise' in tf_var_name and 'noise_' not in tf_var_name:
            var = var.unsqueeze(0).unsqueeze(0)
        G_state_dict[dst_var_name] = var
    print(f'Successfully converted!')
    print(f'--------------------')

    print(f'Converting source weights (Gs) to target ...')
    Gs_vars = dict(model['G_ema'].named_parameters())
    Gs_vars.update(dict(model['G_ema'].named_buffers()))
    Gs = build_model(gan_type=GAN_TPYE,
                     module='generator',
                     resolution=resolution,
                     z_space_dim=z_space_dim,
                     w_space_dim=w_space_dim,
                     label_size=label_size,
                     repeat_w=repeat_w,
                     image_channels=image_channels,
                     fmaps_base=fmaps_base,
                     mapping_layers=mapping_layers)
    Gs_state_dict = Gs.state_dict()
    official_tf_to_pth_var_mapping = {}
    for name in Gs_vars.keys():
        for pattern, fn in G_PTH_TO_TF_VAR_MAPPING_PATTERN.items():
            match = re.fullmatch(pattern, name)
            if match:
                if fn is not None:
                    official_tf_to_pth_var_mapping[fn(*match.groups())] = name
                break
    for dst_var_name, tf_var_name in Gs.pth_to_tf_var_mapping.items():
        assert tf_var_name in official_tf_to_pth_var_mapping
        assert dst_var_name in Gs_state_dict
        src_var_name = official_tf_to_pth_var_mapping[tf_var_name]
        assert src_var_name in Gs_vars
        if verbose:
            print(f'    Converting `{src_var_name}` to `{dst_var_name}`.')
        var = Gs_vars[src_var_name].data
        if 'weight' in tf_var_name:
            if 'Conv0_up/weight' in tf_var_name:
                var = var.flip(2, 3)
            elif 'Skip' in tf_var_name:
                var = var.flip(2, 3)
        if 'bias' in tf_var_name:
            if 'mod_bias' in tf_var_name:
                var = var - 1
        if 'Const' in tf_var_name:
            var = var.unsqueeze(0)
        if 'noise' in tf_var_name and 'noise_' not in tf_var_name:
            var = var.unsqueeze(0).unsqueeze(0)
        Gs_state_dict[dst_var_name] = var
    print(f'Successfully converted!')
    print(f'--------------------')

    print(f'Converting source weights (D) to target ...')
    D_vars = dict(model['D'].named_parameters())
    D_vars.update(dict(model['D'].named_buffers()))
    D = build_model(gan_type=GAN_TPYE,
                    module='discriminator',
                    resolution=resolution,
                    label_size=label_size,
                    image_channels=image_channels,
                    fmaps_base=fmaps_base)
    D_state_dict = D.state_dict()
    official_tf_to_pth_var_mapping = {}
    for name in D_vars.keys():
        for pattern, fn in D_PTH_TO_TF_VAR_MAPPING_PATTERN.items():
            match = re.fullmatch(pattern, name)
            if match:
                if fn is not None:
                    official_tf_to_pth_var_mapping[fn(*match.groups())] = name
                break
    for dst_var_name, tf_var_name in D.pth_to_tf_var_mapping.items():
        assert tf_var_name in official_tf_to_pth_var_mapping
        assert dst_var_name in D_state_dict
        src_var_name = official_tf_to_pth_var_mapping[tf_var_name]
        assert src_var_name in D_vars
        if verbose:
            print(f'    Converting `{src_var_name}` to `{dst_var_name}`.')
        var = D_vars[src_var_name].data
        D_state_dict[dst_var_name] = var
    print(f'Successfully converted!')
    print(f'--------------------')

    print(f'Saving target weights to `{dst_weight_path}` ...')
    state_dict = {
        'generator': G_state_dict,
        'discriminator': D_state_dict,
        'generator_smooth': Gs_state_dict,
    }
    torch.save(state_dict, dst_weight_path)
    print(f'Successfully saved!')
    print(f'--------------------')

    # Start testing if needed.
    if test_num <= 0:
        warnings.warn(f'Skip testing the converted weights!')
        return

    if save_test_image:
        html = HtmlPageVisualizer(num_rows=test_num, num_cols=3)
        html.set_headers(['Index', 'Before Conversion', 'After Conversion'])
        for i in range(test_num):
            html.set_cell(i, 0, text=f'{i}')

    print(f'Testing conversion results ...')
    G.load_state_dict(G_state_dict)
    D.load_state_dict(D_state_dict)
    Gs.load_state_dict(Gs_state_dict)
    G.eval().cuda()
    D.eval().cuda()
    Gs.eval().cuda()
    model['G'].eval().cuda()
    model['D'].eval().cuda()
    model['G_ema'].eval().cuda()

    gs_distance = 0.0
    dg_distance = 0.0
    for i in tqdm(range(test_num)):
        # Test Gs(z).
        code = np.random.randn(1, z_space_dim)
        code = torch.from_numpy(code).type(torch.FloatTensor).cuda()
        if label_size:
            label_id = np.random.randint(label_size)
            label = np.zeros((1, label_size), np.float32)
            label[0, label_id] = 1.0
            label = torch.from_numpy(label).type(torch.FloatTensor).cuda()
        else:
            label_id = 0
            label = None

        src_output = model['G_ema'](code,
                                    label,
                                    truncation_psi=TRUNC_PSI,
                                    truncation_cutoff=TRUNC_LAYERS,
                                    noise_mode=NOISE_MODE)
        src_output = src_output.detach().cpu().numpy()
        dst_output = Gs(code,
                        label=label,
                        trunc_psi=TRUNC_PSI,
                        trunc_layers=TRUNC_LAYERS,
                        randomize_noise=RANDOMIZE_NOISE)['image']
        dst_output = dst_output.detach().cpu().numpy()
        distance = np.average(np.abs(src_output - dst_output))
        if verbose:
            print(f'    Test {i:03d}: Gs distance {distance:.6e}.')
        gs_distance += distance

        if save_test_image:
            html.set_cell(i, 1, image=postprocess_image(src_output)[0])
            html.set_cell(i, 2, image=postprocess_image(dst_output)[0])

        # Test D(G(z)).
        code = np.random.randn(1, z_space_dim)
        code = torch.from_numpy(code).type(torch.FloatTensor).cuda()
        if label_size:
            label_id = np.random.randint(label_size)
            label = np.zeros((1, label_size), np.float32)
            label[0, label_id] = 1.0
            label = torch.from_numpy(label).type(torch.FloatTensor).cuda()
        else:
            label_id = 0
            label = None
        src_image = model['G'](code,
                               label,
                               truncation_psi=TRUNC_PSI,
                               truncation_cutoff=TRUNC_LAYERS,
                               noise_mode=NOISE_MODE)
        src_output = model['D'](src_image, label)
        src_output = src_output.detach().cpu().numpy()
        dst_image = G(code,
                      label=label,
                      trunc_psi=TRUNC_PSI,
                      trunc_layers=TRUNC_LAYERS,
                      randomize_noise=RANDOMIZE_NOISE)['image']
        dst_output = D(dst_image, label)
        dst_output = dst_output.detach().cpu().numpy()
        distance = np.average(np.abs(src_output - dst_output))
        if verbose:
            print(f'    Test {i:03d}: D(G) distance {distance:.6e}.')
        dg_distance += distance

    print(f'Average Gs distance is {gs_distance / test_num:.6e}.')
    print(f'Average D(G) distance is {dg_distance / test_num:.6e}.')
    print(f'========================================')

    if save_test_image:
        html.save(f'{dst_weight_path}.conversion_test.html')
Converts the pre-trained StyleGAN2-ADA-PyTorch weights.

    Args:
        src_weight_path: Path to the source model to load weights from.
        dst_weight_path: Path to the target model to save converted weights.
        test_num: Number of samples used to test the conversion. (default: 10)
        save_test_image: Whether to save the test images. (default: False)
        verbose: Whether to print verbose log message. (default: False)
    
converters/stylegan2ada_pth_converter.py