editor.py

'''Refer to https://github.com/rosinality/stylegan2-pytorch/blob/master/closed_form_factorization.py and '''
'''         https://github.com/rosinality/stylegan2-pytorch/blob/master/apply_factor.py'''

import numpy as np
import paddle

from utils import arg_type, func_args, make_image, get_generator

from PIL import Image
from tqdm import tqdm


svd_vector = None
def get_factorized_style_vector(generator, index):
    global svd_vector
    if svd_vector is None:
        named_parameters = dict(generator.named_parameters())
        modulate = {
            k: v
            for k, v in named_parameters.items()
            if "modulation" in k and "to_rgbs" not in k and "weight" in k
        }

        weight_mat = []
        for _, v in modulate.items():
            weight_mat.append(v.t())

        W = paddle.concat(weight_mat, 0)
        svd_vector = paddle.to_tensor(np.linalg.svd(W.numpy())[-1].T.astype('float32'))

    return svd_vector[:,index].unsqueeze(0).unsqueeze(0).tile([1,1,1])


def edit(
    latent_codes: paddle.Tensor = None,
    generator = None,
    ckpt: arg_type(str, help="path to the model checkpoint") = None,
    model_type: arg_type(str, help="inner model type. `ffhq-config-f` for default genrator and `ffhq-inversion` for pSp") = None,
    size: arg_type(int, help="original output image resolution") = 1024,
    style_dim: arg_type(int, help="dimensions of style z") = 512,
    n_mlp: arg_type(int, help="the number of multi-layer perception layers for style z") = 8,
    channel_multiplier: arg_type(int, help="channel product, affect model size and the quality of generated pictures") = 2,
    style_vector_file: arg_type(str, help="path to the style vector") = 'docs/stylegan2directions/age.npy',
    start_offset: arg_type(float, help="the first offset of the style vector relative to the latent code") = -5,
    final_offset: arg_type(float, help="the last offset of the style vector relative to the latent code") = 5,
    offset_step_size: arg_type(float, help="the offset step size from the start offset to the final offset") = 0.1,
    use_factorizer: arg_type(
       'edit:use_factorizer', action="store_true", 
       help="whether to use the style vectors of the factorizer"
    ) = False,
    factorizer_style_index: arg_type(int, help="the index of one style vector in the style vectors of the factorizer") = 0,
):
    
    generator = generator if generator is not None else get_generator(
        weight_path=None if ckpt is None else ckpt,
        model_type='ffhq-config-f' if model_type is None else model_type,
        size=size,
        style_dim=style_dim,
        n_mlp=n_mlp,
        channel_multiplier=channel_multiplier
    )
    
    if use_factorizer:
        style_vector = get_factorized_style_vector(generator, factorizer_style_index)
    else:
        style_vector = paddle.to_tensor(np.load(style_vector_file)).astype('float32').unsqueeze(0)
        
    degrees = []
    degree = start_offset
    while degree <= final_offset + 1e-8:
        degrees.append(degree)
        degree += offset_step_size

    frames = []
    latent_codes_seqs = []
    for degree in tqdm(degrees):
        with paddle.no_grad():
            latent_in = latent_codes + degree * style_vector
            img_gen, _ = generator([latent_in], input_is_latent=True, randomize_noise=False)
        frames.append(make_image(img_gen))
        latent_codes_seqs.append(latent_in)

    imgs_seq = [[] for _ in range(img_gen.shape[0])]
    latents_seq = [[] for _ in range(latent_in.shape[0])]
    for i in range(img_gen.shape[0]):
        for frame, latent_codes_seq in zip(frames, latent_codes_seqs):
            imgs_seq[i].append(frame[i])
            latents_seq[i].append(latent_codes_seq[i])

    return imgs_seq, latents_seq


if __name__ == "__main__":
    import argparse
    import os

    from utils import save_video

    parser = argparse.ArgumentParser(
        description="Image attributes editor in stylegan2 latent space"
    )
    parser, arg_names = func_args(parser, edit)
    parser.add_argument(
        "--save_mp4", action="store_true", help="saving training progress images as mp4 videos"
    )
    parser.add_argument(
        "files", metavar="FILES", nargs="+", help="path to image latent codes to be edited"
    )
    parser.add_argument(
        "--output", type=str, default="./output", help="output directory"
    )

    args = parser.parse_args()

    latent_codes = []
    for path in args.files:
        latent_code = paddle.load(path)['latent_code']
        latent_codes.append(latent_code)
    latent_codes = paddle.stack(latent_codes, 0)

    imgs_seq, latents_seq = edit(latent_codes, **{arg_name: getattr(args, arg_name) for arg_name in arg_names})

    os.makedirs(args.output, exist_ok=True)
    for i, input_name in enumerate(args.files):

        for j, latent_code in enumerate(latents_seq[i]):
            code_name = os.path.join(
                args.output, 
                os.path.splitext(os.path.basename(input_name))[0] + "-" + str(j).zfill(3) + ".pd"
            )
            latent_file = {
                "latent_code": latent_code,
            }
            paddle.save(latent_file, code_name)

        img_name = os.path.join(
            args.output, 
            os.path.splitext(os.path.basename(input_name))[0] + "-start-offset.png"
        )
        pil_img = Image.fromarray(imgs_seq[i][0])
        pil_img.save(img_name)

        img_name = os.path.join(
            args.output, 
            os.path.splitext(os.path.basename(input_name))[0] + "-final-offset.png"
        )
        pil_img = Image.fromarray(imgs_seq[i][-1])
        pil_img.save(img_name)

        if args.save_mp4:
            fps = 30
            duration = 5
            save_video(
                imgs_seq[i],
                os.path.join(
                    args.output, 
                    os.path.splitext(os.path.basename(input_name))[0] + "-offset.mp4"
                ),
                fps, duration
            )