insetgan.py

# Copyright (c) SenseTime Research. All rights reserved.


import torch
import torch.nn.functional as F
from tqdm import tqdm
from lpips import LPIPS
import numpy as np
from torch_utils.models import Generator as bodyGAN
from torch_utils.models_face import Generator as FaceGAN
import dlib
from utils.alignment import align_face_for_insetgan
from utils.util import visual,tensor_to_numpy, numpy_to_tensor
import legacy
import os
import click


class InsetGAN(torch.nn.Module):
    def __init__(self, stylebody_ckpt, styleface_ckpt):
        super().__init__()
        
        ## convert pkl to pth
        if not os.path.exists(stylebody_ckpt.replace('.pkl','.pth')):
            legacy.convert(stylebody_ckpt, stylebody_ckpt.replace('.pkl','.pth'))
        stylebody_ckpt = stylebody_ckpt.replace('.pkl','.pth') 
        
        if not os.path.exists(styleface_ckpt.replace('.pkl','.pth')):
            legacy.convert(styleface_ckpt, styleface_ckpt.replace('.pkl','.pth'))
        styleface_ckpt = styleface_ckpt.replace('.pkl','.pth') 
    
        # dual generator
        config = {"latent" : 512, "n_mlp" : 8, "channel_multiplier": 2}
        self.body_generator = bodyGAN(
                size = 1024,
                style_dim=config["latent"],
                n_mlp=config["n_mlp"],
                channel_multiplier=config["channel_multiplier"]
            )
        self.body_generator.load_state_dict(torch.load(stylebody_ckpt)['g_ema'])
        self.body_generator.eval().requires_grad_(False).cuda()

        self.face_generator = FaceGAN(
                size = 1024,
                style_dim=config["latent"],
                n_mlp=config["n_mlp"],
                channel_multiplier=config["channel_multiplier"]
            )
        self.face_generator.load_state_dict(torch.load(styleface_ckpt)['g_ema'])
        self.face_generator.eval().requires_grad_(False).cuda()
        # crop function
        self.dlib_predictor = dlib.shape_predictor('./pretrained_models/shape_predictor_68_face_landmarks.dat')
        self.dlib_cnn_face_detector = dlib.cnn_face_detection_model_v1("pretrained_models/mmod_human_face_detector.dat")

        # criterion
        self.lpips_loss = LPIPS(net='alex').cuda().eval()
        self.l1_loss = torch.nn.L1Loss(reduction='mean')
        
    def loss_coarse(self, A_face, B, p1=500, p2=0.05):
        A_face = F.interpolate(A_face, size=(64, 64), mode='area')
        B = F.interpolate(B, size=(64, 64), mode='area')
        loss_l1 = p1 * self.l1_loss(A_face, B)
        loss_lpips = p2 * self.lpips_loss(A_face, B)
        return loss_l1 + loss_lpips

    @staticmethod
    def get_border_mask(A, x, spec):
        mask = torch.zeros_like(A)
        mask[:, :, :x, ] = 1
        mask[:, :, -x:, ] = 1
        mask[:, :, :, :x ] = 1
        mask[:, :, :, -x:] = 1
        return mask
    
    @staticmethod
    def get_body_mask(A, crop, padding=4):
        mask = torch.ones_like(A)
        mask[:, :, crop[1]-padding:crop[3]+padding, crop[0]-padding:crop[2]+padding] = 0
        return mask

    def loss_border(self, A_face, B, p1=10000, p2=2, spec=None):
        mask = self.get_border_mask(A_face, 8, spec)
        loss_l1 = p1 * self.l1_loss(A_face*mask, B*mask)
        loss_lpips = p2 * self.lpips_loss(A_face*mask, B*mask)
        return loss_l1 + loss_lpips

    def loss_body(self, A, B, crop, p1=9000, p2=0.1):
        padding = int((crop[3] - crop[1]) / 20)
        mask = self.get_body_mask(A, crop, padding)
        loss_l1 = p1 * self.l1_loss(A*mask, B*mask)
        loss_lpips = p2 * self.lpips_loss(A*mask, B*mask)
        return loss_l1+loss_lpips

    def loss_face(self, A, B, crop, p1=5000, p2=1.75):
        mask = 1 - self.get_body_mask(A, crop)
        loss_l1 = p1 * self.l1_loss(A*mask, B*mask)
        loss_lpips = p2 * self.lpips_loss(A*mask, B*mask)
        return loss_l1+loss_lpips
    
    def loss_reg(self, w, w_mean, p1, w_plus_delta=None, p2=None):
        return p1 * torch.mean(((w - w_mean) ** 2)) + p2 * torch.mean(w_plus_delta ** 2)
   
    # FFHQ type 
    def detect_face_dlib(self, img):
        # tensor to numpy array rgb uint8
        img = tensor_to_numpy(img)
        aligned_image, crop, rect = align_face_for_insetgan(img=img, 
                                               detector=self.dlib_cnn_face_detector, 
                                               predictor=self.dlib_predictor, 
                                               output_size=256)

        aligned_image = np.array(aligned_image)
        aligned_image = numpy_to_tensor(aligned_image)
        return aligned_image, crop, rect
  
    # joint optimization
    def dual_optimizer(self, 
                       face_w,
                       body_w,
                       joint_steps=500,
                       face_initial_learning_rate=0.02,
                       body_initial_learning_rate=0.05,
                       lr_rampdown_length=0.25,
                       lr_rampup_length=0.05,
                       seed=None,
                       output_path=None,
                       video=0): 
        '''
        Given a face_w, optimize a body_w with suitable body pose & shape for face_w
        '''
        def visual_(path, synth_body, synth_face, body_crop, step, both=False, init_body_with_face=None):
            tmp = synth_body.clone().detach()
            tmp[:, :, body_crop[1]:body_crop[3], body_crop[0]:body_crop[2]] = synth_face
            if both:
                tmp = torch.cat([synth_body, tmp], dim=3)
            save_path = os.path.join(path, f"{step:04d}.jpg")
            visual(tmp, save_path)
            
        def forward(face_w_opt, 
                    body_w_opt, 
                    face_w_delta,
                    body_w_delta, 
                    body_crop, 
                    update_crop=False
                    ):
            if face_w_opt.shape[1] != 18:
                face_ws = (face_w_opt).repeat([1, 18, 1])
            else:
                face_ws = face_w_opt.clone()
            face_ws = face_ws + face_w_delta
            synth_face, _ = self.face_generator([face_ws], input_is_latent=True, randomize_noise=False)
            
            body_ws = (body_w_opt).repeat([1, 18, 1])
            body_ws = body_ws + body_w_delta
            synth_body, _ = self.body_generator([body_ws], input_is_latent=True, randomize_noise=False)
            
            if update_crop:
                old_r = (body_crop[3]-body_crop[1]) // 2, (body_crop[2]-body_crop[0]) // 2
                _, body_crop, _ = self.detect_face_dlib(synth_body)
                center = (body_crop[1] + body_crop[3]) // 2, (body_crop[0] + body_crop[2]) // 2
                body_crop = (center[1] - old_r[1], center[0] - old_r[0], center[1] + old_r[1], center[0] + old_r[0])
            
            synth_body_face = synth_body[:, :, body_crop[1]:body_crop[3], body_crop[0]:body_crop[2]]
            
            if synth_face.shape[2] > body_crop[3]-body_crop[1]:
                synth_face_resize = F.interpolate(synth_face, size=(body_crop[3]-body_crop[1], body_crop[2]-body_crop[0]), mode='area')
        
            return synth_body, synth_body_face, synth_face, synth_face_resize, body_crop
                
        def update_lr(init_lr, step, num_steps, lr_rampdown_length, lr_rampup_length):
            t = step / num_steps
            lr_ramp = min(1.0, (1.0 - t) / lr_rampdown_length)
            lr_ramp = 0.5 - 0.5 * np.cos(lr_ramp * np.pi)
            lr_ramp = lr_ramp * min(1.0, t / lr_rampup_length)
            lr = init_lr * lr_ramp
            return lr
        
        # update output_path
        output_path = os.path.join(output_path, seed)
        os.makedirs(output_path, exist_ok=True)
        
        # define optimized params
        body_w_mean = self.body_generator.mean_latent(10000).detach()
        face_w_opt = face_w.clone().detach().requires_grad_(True)
        body_w_opt = body_w.clone().detach().requires_grad_(True)
        face_w_delta = torch.zeros_like(face_w.repeat([1, 18, 1])).requires_grad_(True)
        body_w_delta = torch.zeros_like(body_w.repeat([1, 18, 1])).requires_grad_(True)
        # generate ref face & body
        ref_body, _ = self.body_generator([body_w.repeat([1, 18, 1])], input_is_latent=True, randomize_noise=False)
        # for inversion
        ref_face, _ = self.face_generator([face_w.repeat([1, 18, 1])], input_is_latent=True, randomize_noise=False)
        # get initilized crop
        _, body_crop, _ = self.detect_face_dlib(ref_body)
        _, _, face_crop = self.detect_face_dlib(ref_face) # NOTE: this is face rect only. no FFHQ type.
        # create optimizer
        face_optimizer = torch.optim.Adam([face_w_opt, face_w_delta], betas=(0.9, 0.999), lr=face_initial_learning_rate)
        body_optimizer = torch.optim.Adam([body_w_opt, body_w_delta], betas=(0.9, 0.999), lr=body_initial_learning_rate)
        
        global_step = 0
        # Stage1: remove background of face image
        face_steps = 25
        pbar = tqdm(range(face_steps))
        for step in pbar:
            face_lr = update_lr(face_initial_learning_rate / 2, step, face_steps, lr_rampdown_length, lr_rampup_length)
            for param_group in face_optimizer.param_groups:
                param_group['lr'] =face_lr
            synth_body, synth_body_face, synth_face_raw, synth_face, body_crop = forward(face_w_opt, 
                                                                                    body_w_opt, 
                                                                                    face_w_delta,
                                                                                    body_w_delta, 
                                                                                    body_crop)
            loss_face = self.loss_face(synth_face_raw, ref_face, face_crop, 5000, 1.75)
            loss_coarse = self.loss_coarse(synth_face, synth_body_face, 50, 0.05)
            loss_border = self.loss_border(synth_face, synth_body_face, 1000, 0.1)
            loss = loss_coarse + loss_border + loss_face
            face_optimizer.zero_grad()
            loss.backward()
            face_optimizer.step()
            # visualization
            if video:
                visual_(output_path, synth_body, synth_face, body_crop, global_step)
            pbar.set_description(
                (                
                    f"face: {step:.4f}, lr: {face_lr}, loss: {loss.item():.2f}, loss_coarse: {loss_coarse.item():.2f};"
                    f"loss_border: {loss_border.item():.2f}, loss_face: {loss_face.item():.2f};"
                )
            )
            global_step += 1
            
        # Stage2: find a suitable body
        body_steps = 150
        pbar = tqdm(range(body_steps))
        for step in pbar:
            body_lr = update_lr(body_initial_learning_rate, step, body_steps, lr_rampdown_length, lr_rampup_length)
            update_crop = True if (step % 50 == 0) else False
            # update_crop = False
            for param_group in body_optimizer.param_groups:
                param_group['lr'] =body_lr
            synth_body, synth_body_face, synth_face_raw, synth_face, body_crop = forward(face_w_opt, 
                                                                                    body_w_opt, 
                                                                                    face_w_delta, 
                                                                                    body_w_delta, 
                                                                                    body_crop,
                                                                                    update_crop=update_crop)
            loss_coarse = self.loss_coarse(synth_face, synth_body_face, 500, 0.05)
            loss_border = self.loss_border(synth_face, synth_body_face, 2500, 0)
            loss_body = self.loss_body(synth_body, ref_body, body_crop, 9000, 0.1)
            loss_reg = self.loss_reg(body_w_opt, body_w_mean, 15000, body_w_delta, 0)
            loss = loss_coarse + loss_border + loss_body + loss_reg
            body_optimizer.zero_grad()
            loss.backward()
            body_optimizer.step()     
            
            # visualization
            if video:
                visual_(output_path, synth_body, synth_face, body_crop, global_step)
            pbar.set_description(
                (
                    f"body: {step:.4f}, lr: {body_lr}, loss: {loss.item():.2f}, loss_coarse: {loss_coarse.item():.2f};"
                    f"loss_border: {loss_border.item():.2f}, loss_body: {loss_body.item():.2f}, loss_reg: {loss_reg:.2f}"
                )
            )    
            global_step += 1
        
        # Stage3: joint optimization
        interval = 50
        joint_face_steps = joint_steps // 2
        joint_body_steps = joint_steps // 2
        face_step = 0
        body_step = 0
        pbar = tqdm(range(joint_steps))
        flag = -1
        for step in pbar:
            if step % interval == 0: flag += 1
            text_flag = 'optimize_face' if flag % 2 == 0 else 'optimize_body'
            synth_body, synth_body_face, synth_face_raw, synth_face, body_crop = forward(face_w_opt, 
                                                                                    body_w_opt, 
                                                                                    face_w_delta, 
                                                                                    body_w_delta, 
                                                                                    body_crop)
            if text_flag == 'optimize_face':
                face_lr = update_lr(face_initial_learning_rate, face_step, joint_face_steps, lr_rampdown_length, lr_rampup_length)
                for param_group in face_optimizer.param_groups:
                    param_group['lr'] =face_lr
                loss_face = self.loss_face(synth_face_raw, ref_face, face_crop, 5000, 1.75)
                loss_coarse = self.loss_coarse(synth_face, synth_body_face, 500, 0.05)
                loss_border = self.loss_border(synth_face, synth_body_face, 25000, 0)
                loss = loss_coarse + loss_border + loss_face
                face_optimizer.zero_grad()
                loss.backward()
                face_optimizer.step()
                pbar.set_description(
                    (                
                        f"face: {step}, lr: {face_lr:.4f}, loss: {loss.item():.2f}, loss_coarse: {loss_coarse.item():.2f};"
                        f"loss_border: {loss_border.item():.2f}, loss_face: {loss_face.item():.2f};"
                    )
                )
                face_step += 1
            else:
                body_lr = update_lr(body_initial_learning_rate, body_step, joint_body_steps, lr_rampdown_length, lr_rampup_length)
                for param_group in body_optimizer.param_groups:
                    param_group['lr'] =body_lr
                loss_coarse = self.loss_coarse(synth_face, synth_body_face, 500, 0.05)
                loss_border = self.loss_border(synth_face, synth_body_face, 2500, 0)
                loss_body = self.loss_body(synth_body, ref_body, body_crop, 9000, 0.1)
                loss_reg = self.loss_reg(body_w_opt, body_w_mean, 25000, body_w_delta, 0)
                loss = loss_coarse + loss_border + loss_body + loss_reg
                body_optimizer.zero_grad()
                loss.backward()
                body_optimizer.step()
                pbar.set_description(
                    (
                        f"body: {step}, lr: {body_lr:.4f}, loss: {loss.item():.2f}, loss_coarse: {loss_coarse.item():.2f};"
                        f"loss_border: {loss_border.item():.2f}, loss_body: {loss_body.item():.2f}, loss_reg: {loss_reg:.2f}"
                    )
                )
                body_step += 1
            if video:
                visual_(output_path, synth_body, synth_face, body_crop, global_step)
            global_step += 1
        return face_w_opt.repeat([1, 18, 1])+face_w_delta, body_w_opt.repeat([1, 18, 1])+body_w_delta, body_crop




"""
Jointly combine and optimize generated faces and bodies .
Examples:

\b
# Combine the generate human full-body image from the provided StyleGAN-Human pre-trained model
# and the generated face image from FFHQ model, optimize both latent codes to produce the coherent face-body image
python insetgan.py python insetgan.py --body_network=pretrained_models/stylegan_human_v2_1024.pkl --face_network=pretrained_models/ffhq.pkl \\
    --body_seed=82 --face_seed=43  --trunc=0.6 --outdir=outputs/insetgan/ --video 1 
"""

@click.command()
@click.pass_context
@click.option('--face_network', default="./pretrained_models/ffhq.pkl", help='Network pickle filename', required=True)
@click.option('--body_network', default='./pretrained_models/stylegan2_1024.pkl', help='Network pickle filename', required=True)
@click.option('--face_seed', type=int, default=82, help='selected random seed')
@click.option('--body_seed', type=int, default=43, help='selected random seed')
@click.option('--joint_steps', type=int, default=500, help='num steps for joint optimization')
@click.option('--trunc', 'truncation_psi', type=float, help='Truncation psi', default=0.6, show_default=True)
@click.option('--outdir', help='Where to save the output images', default= "outputs/insetgan/" , type=str, required=True, metavar='DIR')
@click.option('--video', help="set to 1 if want to save video", type=int, default=0)
def main(
        ctx: click.Context,
        face_network: str,
        body_network: str,
        face_seed: int,
        body_seed: int,
        joint_steps: int,
        truncation_psi: float,
        outdir: str,
        video: int):
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    insgan = InsetGAN(body_network, face_network)
    os.makedirs(outdir, exist_ok=True)
    face_z = np.random.RandomState(face_seed).randn(1, 512).astype(np.float32)
    face_mean = insgan.face_generator.mean_latent(3000)
    face_w = insgan.face_generator.get_latent(torch.from_numpy(face_z).to(device))  # [N, L, C]
    face_w = truncation_psi * face_w + (1-truncation_psi) * face_mean
    face_img, _ = insgan.face_generator([face_w], input_is_latent=True)

    body_z = np.random.RandomState(body_seed).randn(1, 512).astype(np.float32)
    body_mean = insgan.body_generator.mean_latent(3000)
    body_w = insgan.body_generator.get_latent(torch.from_numpy(body_z).to(device))  # [N, L, C]
    body_w = truncation_psi * body_w + (1-truncation_psi) * body_mean
    body_img, _ = insgan.body_generator([body_w], input_is_latent=True)

    _, body_crop, _ = insgan.detect_face_dlib(body_img)
    face_img = F.interpolate(face_img, size=(body_crop[3]-body_crop[1], body_crop[2]-body_crop[0]), mode='area')
    cp_body = body_img.clone()
    cp_body[:, :, body_crop[1]:body_crop[3], body_crop[0]:body_crop[2]] = face_img
    
    optim_face_w, optim_body_w, crop = insgan.dual_optimizer(
        face_w, 
        body_w,
        joint_steps=joint_steps,
        seed=f'{face_seed:04d}_{body_seed:04d}',
        output_path=outdir,
        video=video
    )
    
    if video:
        ffmpeg_cmd = f"ffmpeg -hide_banner -loglevel error -i ./{outdir}/{face_seed:04d}_{body_seed:04d}/%04d.jpg -c:v libx264 -vf fps=30 -pix_fmt yuv420p ./{outdir}/{face_seed:04d}_{body_seed:04d}.mp4"
        os.system(ffmpeg_cmd)
    new_face_img, _ = insgan.face_generator([optim_face_w], input_is_latent=True)
    new_shape = crop[3] - crop[1], crop[2] - crop[0]
    new_face_img_crop = F.interpolate(new_face_img, size=new_shape, mode='area')
    seamless_body, _ = insgan.body_generator([optim_body_w], input_is_latent=True)
    seamless_body[:, :, crop[1]:crop[3], crop[0]:crop[2]] = new_face_img_crop
    temp = torch.cat([cp_body, seamless_body], dim=3)
    visual(temp, f"{outdir}/{face_seed:04d}_{body_seed:04d}.png")

if __name__ == "__main__":
    main()