gfpgan_model.py

#   Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserve.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import math
import sys
import paddle
from paddle.nn import functional as F
from paddle import autograd

from .base_model import BaseModel
from .builder import MODELS
from .generators.builder import build_generator
from .discriminators.builder import build_discriminator
from .criterions.builder import build_criterion
from ..modules.init import init_weights
from collections import OrderedDict
from ..solver import build_lr_scheduler, build_optimizer
from ppgan.utils.visual import *
from ppgan.models.generators.gfpganv1_arch import FacialComponentDiscriminator
from ppgan.utils.download import get_path_from_url


@MODELS.register()
class GFPGANModel(BaseModel):
    """ This class implements the gfpgan model.

    """
    def __init__(self, **opt):

        super(GFPGANModel, self).__init__()
        self.opt = opt
        train_opt = opt
        if 'image_visual' in self.opt['path']:
            self.image_paths = self.opt['path']['image_visual']
        self.current_iter = 0
        self.nets['net_g'] = build_generator(opt['network_g'])
        self.log_size = int(math.log(self.opt['network_g']['out_size'], 2))
        # define networks (both generator and discriminator)
        self.nets['net_g_ema'] = build_generator(self.opt['network_g'])
        self.nets['net_d'] = build_discriminator(self.opt['network_d'])
        self.nets['net_g_ema'].eval()
        pretrain_network_g = self.opt['path'].get('pretrain_network_g', None)
        if pretrain_network_g != None:
            t_weight = get_path_from_url(pretrain_network_g)
            t_weight = paddle.load(t_weight)
            if 'net_g' in t_weight:
                self.nets['net_g'].set_state_dict(t_weight['net_g'])
                self.nets['net_g_ema'].set_state_dict(t_weight['net_g_ema'])
            else:
                self.nets['net_g'].set_state_dict(t_weight)
                self.nets['net_g_ema'].set_state_dict(t_weight)

            del t_weight

        self.nets['net_d'].train()
        self.nets['net_g'].train()
        if ('network_d_left_eye' in self.opt
                and 'network_d_right_eye' in self.opt
                and 'network_d_mouth' in self.opt):
            self.use_facial_disc = True
        else:
            self.use_facial_disc = False

        if self.use_facial_disc:
            # left eye
            self.nets['net_d_left_eye'] = FacialComponentDiscriminator()
            self.nets['net_d_right_eye'] = FacialComponentDiscriminator()
            self.nets['net_d_mouth'] = FacialComponentDiscriminator()
            load_path = self.opt['path'].get('pretrain_network_d_left_eye')
            if load_path is not None:
                load_val = get_path_from_url(load_path)
                load_val = paddle.load(load_val)
                self.nets['net_d_left_eye'].set_state_dict(load_val)
                self.nets['net_d_right_eye'].set_state_dict(load_val)
                self.nets['net_d_mouth'].set_state_dict(load_val)
                del load_val
            self.nets['net_d_left_eye'].train()
            self.nets['net_d_right_eye'].train()
            self.nets['net_d_mouth'].train()
            self.cri_component = build_criterion(train_opt['gan_component_opt'])

        if train_opt.get('pixel_opt'):
            self.cri_pix = build_criterion(train_opt['pixel_opt'])
        else:
            self.cri_pix = None

        # perceptual loss
        if train_opt.get('perceptual_opt'):
            self.cri_perceptual = build_criterion(train_opt['perceptual_opt'])
        else:
            self.cri_perceptual = None

        # L1 loss is used in pyramid loss, component style loss and identity loss
        self.cri_l1 = build_criterion(train_opt['L1_opt'])

        # gan loss (wgan)
        self.cri_gan = build_criterion(train_opt['gan_opt'])

        # ----------- define identity loss ----------- #
        if 'network_identity' in self.opt:
            self.use_identity = True
        else:
            self.use_identity = False

        if self.use_identity:
            # define identity network
            self.network_identity = build_discriminator(
                self.opt['network_identity'])
            load_path = self.opt['path'].get('pretrain_network_identity')
            if load_path is not None:
                load_val = get_path_from_url(load_path)
                load_val = paddle.load(load_val)
                self.network_identity.set_state_dict(load_val)
                del load_val
            self.network_identity.eval()
            for param in self.network_identity.parameters():
                param.stop_gradient = True

        # regularization weights
        self.r1_reg_weight = train_opt['r1_reg_weight']  # for discriminator
        self.net_d_iters = train_opt.get('net_d_iters', 1)
        self.net_d_init_iters = train_opt.get('net_d_init_iters', 0)
        self.net_d_reg_every = train_opt['net_d_reg_every']

    def setup_input(self, data):
        self.lq = data['lq']

        if 'gt' in data:
            self.gt = data['gt']

        if 'loc_left_eye' in data:
            # get facial component locations, shape (batch, 4)
            self.loc_left_eyes = data['loc_left_eye'].astype('float32')
            self.loc_right_eyes = data['loc_right_eye'].astype('float32')
            self.loc_mouths = data['loc_mouth'].astype('float32')

    def forward(self, test_mode=False, regularize=False):
        pass

    def train_iter(self, optimizers=None):
        # optimize nets['net_g']
        for p in self.nets['net_d'].parameters():
            p.stop_gradient = True
        self.optimizers['optim_g'].clear_grad(set_to_zero=False)

        # do not update facial component net_d
        if self.use_facial_disc:
            for p in self.nets['net_d_left_eye'].parameters():
                p.stop_gradient = True
            for p in self.nets['net_d_right_eye'].parameters():
                p.stop_gradient = True
            for p in self.nets['net_d_mouth'].parameters():
                p.stop_gradient = True

        # image pyramid loss weight
        pyramid_loss_weight = self.opt.get('pyramid_loss_weight', 0)
        if pyramid_loss_weight > 0 and self.current_iter > self.opt.get(
                'remove_pyramid_loss', float('inf')):
            pyramid_loss_weight = 1e-12  # very small weight to avoid unused param error
        if pyramid_loss_weight > 0:
            self.output, out_rgbs = self.nets['net_g'](self.lq, return_rgb=True)
            pyramid_gt = self.construct_img_pyramid()
        else:
            self.output, out_rgbs = self.nets['net_g'](self.lq,
                                                       return_rgb=False)

        # get roi-align regions
        if self.use_facial_disc:
            self.get_roi_regions(eye_out_size=80, mouth_out_size=120)
        l_g_total = 0
        if (self.current_iter % self.net_d_iters == 0
                and self.current_iter > self.net_d_init_iters):
            # pixel loss
            if self.cri_pix:
                l_g_pix = self.cri_pix(self.output, self.gt)
                l_g_total += l_g_pix
                self.losses['l_g_pix'] = l_g_pix

            # image pyramid loss
            if pyramid_loss_weight > 0:
                for i in range(0, self.log_size - 2):
                    l_pyramid = self.cri_l1(out_rgbs[i],
                                            pyramid_gt[i]) * pyramid_loss_weight
                    l_g_total += l_pyramid
                    self.losses[f'l_p_{2**(i+3)}'] = l_pyramid

            # perceptual loss
            if self.cri_perceptual:
                l_g_percep, l_g_style = self.cri_perceptual(
                    self.output, self.gt)
                if l_g_percep is not None:
                    l_g_total += l_g_percep
                    self.losses['l_g_percep'] = l_g_percep
                if l_g_style is not None:
                    l_g_total += l_g_style
                    self.losses['l_g_style'] = l_g_style

            # gan loss
            fake_g_pred = self.nets['net_d'](self.output)
            l_g_gan = self.cri_gan(fake_g_pred, True, is_disc=False)
            l_g_total += l_g_gan
            self.losses['l_g_gan'] = l_g_gan

            # facial component loss
            if self.use_facial_disc:
                # left eye
                fake_left_eye, fake_left_eye_feats = self.nets[
                    'net_d_left_eye'](self.left_eyes, return_feats=True)
                l_g_gan = self.cri_component(fake_left_eye, True, is_disc=False)
                l_g_total += l_g_gan
                self.losses['l_g_gan_left_eye'] = l_g_gan
                # right eye
                fake_right_eye, fake_right_eye_feats = self.nets[
                    'net_d_right_eye'](self.right_eyes, return_feats=True)
                l_g_gan = self.cri_component(fake_right_eye,
                                             True,
                                             is_disc=False)
                l_g_total += l_g_gan
                self.losses['l_g_gan_right_eye'] = l_g_gan
                # mouth
                fake_mouth, fake_mouth_feats = self.nets['net_d_mouth'](
                    self.mouths, return_feats=True)
                l_g_gan = self.cri_component(fake_mouth, True, is_disc=False)
                l_g_total += l_g_gan
                self.losses['l_g_gan_mouth'] = l_g_gan

                if self.opt.get('comp_style_weight', 0) > 0:
                    # get gt feat
                    _, real_left_eye_feats = self.nets['net_d_left_eye'](
                        self.left_eyes_gt, return_feats=True)
                    _, real_right_eye_feats = self.nets['net_d_right_eye'](
                        self.right_eyes_gt, return_feats=True)
                    _, real_mouth_feats = self.nets['net_d_mouth'](
                        self.mouths_gt, return_feats=True)

                    def _comp_style(feat, feat_gt, criterion):
                        return criterion(self._gram_mat(
                            feat[0]), self._gram_mat(
                                feat_gt[0].detach())) * 0.5 + criterion(
                                    self._gram_mat(feat[1]),
                                    self._gram_mat(feat_gt[1].detach()))

                    # facial component style loss
                    comp_style_loss = 0
                    comp_style_loss += _comp_style(fake_left_eye_feats,
                                                   real_left_eye_feats,
                                                   self.cri_l1)
                    comp_style_loss += _comp_style(fake_right_eye_feats,
                                                   real_right_eye_feats,
                                                   self.cri_l1)
                    comp_style_loss += _comp_style(fake_mouth_feats,
                                                   real_mouth_feats,
                                                   self.cri_l1)
                    comp_style_loss = comp_style_loss * self.opt[
                        'comp_style_weight']
                    l_g_total += comp_style_loss
                    self.losses['l_g_comp_style_loss'] = comp_style_loss

            # identity loss
            if self.use_identity:
                identity_weight = self.opt['identity_weight']
                # get gray images and resize
                out_gray = self.gray_resize_for_identity(self.output)
                gt_gray = self.gray_resize_for_identity(self.gt)

                identity_gt = self.network_identity(gt_gray).detach()
                identity_out = self.network_identity(out_gray)
                l_identity = self.cri_l1(identity_out,
                                         identity_gt) * identity_weight
                l_g_total += l_identity
                self.losses['l_identity'] = l_identity

            l_g_total.backward()
            self.optimizers['optim_g'].step()
        # EMA
        self.accumulate(self.nets['net_g_ema'],
                        self.nets['net_g'],
                        decay=0.5**(32 / (10 * 1000)))

        # ----------- optimize net_d ----------- #
        for p in self.nets['net_d'].parameters():
            p.stop_gradient = False
        self.optimizers['optim_d'].clear_grad(set_to_zero=False)
        if self.use_facial_disc:
            for p in self.nets['net_d_left_eye'].parameters():
                p.stop_gradient = False
            for p in self.nets['net_d_right_eye'].parameters():
                p.stop_gradient = False
            for p in self.nets['net_d_mouth'].parameters():
                p.stop_gradient = False
            self.optimizers['optim_net_d_left_eye'].clear_grad(
                set_to_zero=False)
            self.optimizers['optim_net_d_right_eye'].clear_grad(
                set_to_zero=False)
            self.optimizers['optim_net_d_mouth'].clear_grad(set_to_zero=False)
        fake_d_pred = self.nets['net_d'](self.output.detach())
        real_d_pred = self.nets['net_d'](self.gt)

        l_d = self.cri_gan(real_d_pred, True, is_disc=True) + self.cri_gan(
            fake_d_pred, False, is_disc=True)
        self.losses['l_d'] = l_d
        # In WGAN, real_score should be positive and fake_score should be negative
        self.losses['real_score'] = real_d_pred.detach().mean()
        self.losses['fake_score'] = fake_d_pred.detach().mean()
        l_d.backward()
        if self.current_iter % self.net_d_reg_every == 0:
            self.gt.stop_gradient = False
            real_pred = self.nets['net_d'](self.gt)
            l_d_r1 = r1_penalty(real_pred, self.gt)
            l_d_r1 = (self.r1_reg_weight / 2 * l_d_r1 * self.net_d_reg_every +
                      0 * real_pred[0])
            self.losses['l_d_r1'] = l_d_r1.detach().mean()
            l_d_r1.backward()

        self.optimizers['optim_d'].step()

        # optimize facial component discriminators
        if self.use_facial_disc:
            # left eye
            fake_d_pred, _ = self.nets['net_d_left_eye'](
                self.left_eyes.detach())
            real_d_pred, _ = self.nets['net_d_left_eye'](self.left_eyes_gt)
            l_d_left_eye = self.cri_component(
                real_d_pred, True, is_disc=True) + self.cri_gan(
                    fake_d_pred, False, is_disc=True)
            self.losses['l_d_left_eye'] = l_d_left_eye
            l_d_left_eye.backward()
            # right eye
            fake_d_pred, _ = self.nets['net_d_right_eye'](
                self.right_eyes.detach())
            real_d_pred, _ = self.nets['net_d_right_eye'](self.right_eyes_gt)
            l_d_right_eye = self.cri_component(
                real_d_pred, True, is_disc=True) + self.cri_gan(
                    fake_d_pred, False, is_disc=True)
            self.losses['l_d_right_eye'] = l_d_right_eye
            l_d_right_eye.backward()
            # mouth
            fake_d_pred, _ = self.nets['net_d_mouth'](self.mouths.detach())
            real_d_pred, _ = self.nets['net_d_mouth'](self.mouths_gt)
            l_d_mouth = self.cri_component(real_d_pred, True,
                                           is_disc=True) + self.cri_gan(
                                               fake_d_pred, False, is_disc=True)
            self.losses['l_d_mouth'] = l_d_mouth
            l_d_mouth.backward()

            self.optimizers['optim_net_d_left_eye'].step()
            self.optimizers['optim_net_d_right_eye'].step()
            self.optimizers['optim_net_d_mouth'].step()
        # if self.current_iter%1000==0:

    def test_iter(self, metrics=None):
        self.nets['net_g_ema'].eval()
        self.fake_img, _ = self.nets['net_g_ema'](self.lq)
        self.visual_items['cur_fake'] = self.fake_img[0]
        self.visual_items['cur_gt'] = self.gt[0]
        self.visual_items['cur_lq'] = self.lq[0]
        with paddle.no_grad():
            if metrics is not None:
                for metric in metrics.values():
                    metric.update(self.fake_img.detach().numpy(),
                                  self.gt.detach().numpy())

    def setup_lr_schedulers(self, cfg):
        self.lr_scheduler = OrderedDict()
        self.lr_scheduler['_g'] = build_lr_scheduler(cfg)
        self.lr_scheduler['_component'] = build_lr_scheduler(cfg)
        cfg_d = cfg.copy()
        net_d_reg_ratio = self.net_d_reg_every / (self.net_d_reg_every + 1)
        cfg_d['learning_rate'] *= net_d_reg_ratio
        self.lr_scheduler['_d'] = build_lr_scheduler(cfg_d)
        return self.lr_scheduler

    def setup_optimizers(self, lr, cfg):
        # ----------- optimizer g ----------- #
        net_g_reg_ratio = 1
        parameters = []
        parameters += self.nets['net_g'].parameters()
        cfg['optim_g']['beta1'] = 0**net_g_reg_ratio
        cfg['optim_g']['beta2'] = 0.99**net_g_reg_ratio

        self.optimizers['optim_g'] = build_optimizer(cfg['optim_g'],
                                                     self.lr_scheduler['_g'],
                                                     parameters)

        # ----------- optimizer d ----------- #
        net_d_reg_ratio = self.net_d_reg_every / (self.net_d_reg_every + 1)
        parameters = []
        parameters += self.nets['net_d'].parameters()
        cfg['optim_d']['beta1'] = 0**net_d_reg_ratio
        cfg['optim_d']['beta2'] = 0.99**net_d_reg_ratio

        self.optimizers['optim_d'] = build_optimizer(cfg['optim_d'],
                                                     self.lr_scheduler['_d'],
                                                     parameters)

        # ----------- optimizers for facial component networks ----------- #
        if self.use_facial_disc:
            parameters = []
            parameters += self.nets['net_d_left_eye'].parameters()

            self.optimizers['optim_net_d_left_eye'] = build_optimizer(
                cfg['optim_component'], self.lr_scheduler['_component'],
                parameters)

            parameters = []
            parameters += self.nets['net_d_right_eye'].parameters()

            self.optimizers['optim_net_d_right_eye'] = build_optimizer(
                cfg['optim_component'], self.lr_scheduler['_component'],
                parameters)

            parameters = []
            parameters += self.nets['net_d_mouth'].parameters()

            self.optimizers['optim_net_d_mouth'] = build_optimizer(
                cfg['optim_component'], self.lr_scheduler['_component'],
                parameters)

        return self.optimizers

    def construct_img_pyramid(self):
        """Construct image pyramid for intermediate restoration loss"""
        pyramid_gt = [self.gt]
        down_img = self.gt
        for _ in range(0, self.log_size - 3):
            down_img = F.interpolate(down_img,
                                     scale_factor=0.5,
                                     mode='bilinear',
                                     align_corners=False)
            pyramid_gt.insert(0, down_img)
        return pyramid_gt

    def get_roi_regions(self, eye_out_size=80, mouth_out_size=120):
        from paddle.vision.ops import roi_align
        face_ratio = int(self.opt['network_g']['out_size'] / 512)
        eye_out_size *= face_ratio
        mouth_out_size *= face_ratio

        rois_eyes = []
        rois_mouths = []
        num_eye = []
        num_mouth = []
        for b in range(self.loc_left_eyes.shape[0]):  # loop for batch size
            # left eye and right eye

            img_inds = paddle.ones([2, 1], dtype=self.loc_left_eyes.dtype) * b
            bbox = paddle.stack(
                [self.loc_left_eyes[b, :], self.loc_right_eyes[b, :]],
                axis=0)  # shape: (2, 4)
            # rois = paddle.concat([img_inds, bbox], axis=-1)  # shape: (2, 5)
            rois_eyes.append(bbox)
            # mouse
            img_inds = paddle.ones([1, 1], dtype=self.loc_left_eyes.dtype) * b
            num_eye.append(2)
            num_mouth.append(1)
            # rois = paddle.concat([img_inds, self.loc_mouths[b:b + 1, :]], axis=-1)  # shape: (1, 5)
            rois_mouths.append(self.loc_mouths[b:b + 1, :])
        rois_eyes = paddle.concat(rois_eyes, 0)
        rois_mouths = paddle.concat(rois_mouths, 0)
        # real images
        num_eye = paddle.to_tensor(num_eye, dtype='int32')
        num_mouth = paddle.to_tensor(num_mouth, dtype='int32')

        all_eyes = roi_align(self.gt,
                             boxes=rois_eyes,
                             boxes_num=num_eye,
                             output_size=eye_out_size,
                             aligned=False) * face_ratio
        self.left_eyes_gt = all_eyes[0::2, :, :, :]
        self.right_eyes_gt = all_eyes[1::2, :, :, :]
        self.mouths_gt = roi_align(self.gt,
                                   boxes=rois_mouths,
                                   boxes_num=num_mouth,
                                   output_size=mouth_out_size,
                                   aligned=False) * face_ratio
        # output
        all_eyes = roi_align(self.output,
                             boxes=rois_eyes,
                             boxes_num=num_eye,
                             output_size=eye_out_size,
                             aligned=False) * face_ratio
        self.left_eyes = all_eyes[0::2, :, :, :]
        self.right_eyes = all_eyes[1::2, :, :, :]
        self.mouths = roi_align(self.output,
                                boxes=rois_mouths,
                                boxes_num=num_mouth,
                                output_size=mouth_out_size,
                                aligned=False) * face_ratio

    def _gram_mat(self, x):
        """Calculate Gram matrix.

        Args:
            x (paddle.Tensor): Tensor with shape of (n, c, h, w).

        Returns:
            paddle.Tensor: Gram matrix.
        """
        n, c, h, w = x.shape
        features = x.reshape((n, c, w * h))
        features_t = features.transpose([0, 2, 1])
        gram = features.bmm(features_t) / (c * h * w)
        return gram

    def gray_resize_for_identity(self, out, size=128):
        out_gray = (0.2989 * out[:, 0, :, :] + 0.5870 * out[:, 1, :, :] +
                    0.1140 * out[:, 2, :, :])
        out_gray = out_gray.unsqueeze(1)
        out_gray = F.interpolate(out_gray, (size, size),
                                 mode='bilinear',
                                 align_corners=False)
        return out_gray

    def accumulate(self, model1, model2, decay=0.999):
        par1 = dict(model1.state_dict())
        par2 = dict(model2.state_dict())

        for k in par1.keys():
            par1[k] = par1[k] * decay + par2[k] * (1 - decay)

        model1.load_dict(par1)


def r1_penalty(real_pred, real_img):
    """R1 regularization for discriminator. The core idea is to
        penalize the gradient on real data alone: when the
        generator distribution produces the true data distribution
        and the discriminator is equal to 0 on the data manifold, the
        gradient penalty ensures that the discriminator cannot create
        a non-zero gradient orthogonal to the data manifold without
        suffering a loss in the GAN game.

        Ref:
        Eq. 9 in Which training methods for GANs do actually converge.
        """
    grad_real = paddle.grad(outputs=real_pred.sum(),
                            inputs=real_img,
                            create_graph=True)[0]
    grad_penalty = grad_real.pow(2).reshape(
        (grad_real.shape[0], -1)).sum(1).mean()
    return grad_penalty