fix First Order Motion (#727)

* fix First Order Motion

* fix bug

README.md

@@ -123,6 +123,8 @@ GAN-Generative Adversarial Network, was praised by "the Father of Convolutional
 * [InvDN](./docs/en_US/tutorials/invdn.md)
 * [AOT-GAN](./docs/en_US/tutorials/aotgan.md)
 * [NAFNet](./docs/en_US/tutorials/nafnet.md)
+* [GFPGan](./docs/en_US/tutorials/gfpgan.md)
+* [GPEN](./docs/en_US/tutorials/gpen.md)
 
 
 ## Composite Application

README_cn.md

@@ -132,6 +132,7 @@ GAN--生成对抗网络，被“卷积网络之父”**Yann LeCun（杨立昆）
   * 人脸编码：[Pixel2Style2Pixel](./docs/zh_CN/tutorials/pixel2style2pixel.md)
   * 人脸增强：[FaceEnhancement](./docs/zh_CN/tutorials/face_enhancement.md)
   * 人脸解析：[FaceParsing](./docs/zh_CN/tutorials/face_parse.md)
+  * 盲人脸修复：[GFPGan](./docs/zh_CN/tutorials/gfpgan.md)、[GPEN](./docs/zh_CN/tutorials/gpen.md)
 * 分辨率提升
   * 单张图片超分：[Single Image Super Resolution(SISR)](./docs/zh_CN/tutorials/single_image_super_resolution.md)
     * 包含模型：RealSR、ESRGAN、LESRCNN、PAN、DRN

configs/gpen_256_ffhq.yaml

@@ -6,7 +6,7 @@ find_unused_parameters: True
 model:
   name: GPENModel
   generator:
-    name: GPEN
+    name: GPENGenerator
     size: 256
     style_dim: 512
     n_mlp: 8

docs/zh_CN/tutorials/face_enhancement.md

@@ -20,9 +20,9 @@ img = faceenhancer.enhance_from_image(img)
 
 注意：请将图片转为float类型输入，目前不支持int8类型
 
-### 训练(TODO)
+### 训练
 
-未来还将添加训练脚本方便用户训练出更多类型的 GPEN 人脸增强。
+[详见](../../zh_CN/tutorials/gpen.md)
 
 ## 人脸增强结果展示
 

ppgan/apps/gpen_predictor.py

@@ -21,7 +21,7 @@ import sys
 sys.path.append(".")
 from .base_predictor import BasePredictor
 from ppgan.datasets.gpen_dataset import GFPGAN_degradation
-from ppgan.models.generators import GPEN
+from ppgan.models.generators import GPENGenerator
 from ppgan.metrics.fid import FID
 from ppgan.utils.download import get_path_from_url
 import cv2
@@ -103,7 +103,7 @@ class GPENPredictor(BasePredictor):
             checkpoint = paddle.load(weight_path)
 
         warnings.filterwarnings("always")
-        self.generator = GPEN(size, style_dim, n_mlp, channel_multiplier,
+        self.generator = GPENGenerator(size, style_dim, n_mlp, channel_multiplier,
                               narrow)
         self.generator.set_state_dict(checkpoint)
         self.generator.eval()

ppgan/models/generators/__init__.py

@@ -41,7 +41,7 @@ from .msvsr import MSVSR
 from .generator_singan import SinGANGenerator
 from .rcan import RCAN
 from .prenet import PReNet
-from .gpen import GPEN
+from .generator_gpen import GPENGenerator
 from .swinir import SwinIR
 from .gfpganv1_clean_arch import GFPGANv1Clean
 from .gfpganv1_arch import GFPGANv1, StyleGAN2DiscriminatorGFPGAN

ppgan/models/generators/generator_gpen.py

@@ -18,10 +18,13 @@
 
 import math
 import random
+import itertools
 import paddle
 import paddle.nn as nn
 import paddle.nn.functional as F
 
+from ppgan.models.generators.builder import GENERATORS
+from ppgan.models.discriminators.discriminator_styleganv2 import ConvLayer
 from ppgan.modules.equalized import EqualLinear_gpen as EqualLinear
 from ppgan.modules.fused_act import FusedLeakyReLU
 from ppgan.modules.upfirdn2d import Upfirdn2dUpsample, Upfirdn2dBlur
@@ -451,3 +454,85 @@ class StyleGANv2Generator(nn.Layer):
 
         else:
             return image, None
+
+@GENERATORS.register()
+class GPENGenerator(nn.Layer):
+
+    def __init__(
+        self,
+        size,
+        style_dim,
+        n_mlp,
+        channel_multiplier=2,
+        narrow=1,
+        blur_kernel=[1, 3, 3, 1],
+        lr_mlp=0.01,
+        is_concat=True,
+    ):
+        super(GPENGenerator, self).__init__()
+        channels = {
+            4: int(512 * narrow),
+            8: int(512 * narrow),
+            16: int(512 * narrow),
+            32: int(512 * narrow),
+            64: int(256 * channel_multiplier * narrow),
+            128: int(128 * channel_multiplier * narrow),
+            256: int(64 * channel_multiplier * narrow),
+            512: int(32 * channel_multiplier * narrow),
+            1024: int(16 * channel_multiplier * narrow),
+            2048: int(8 * channel_multiplier * narrow)
+        }
+        self.log_size = int(math.log(size, 2))
+        self.generator = StyleGANv2Generator(
+            size,
+            style_dim,
+            n_mlp,
+            channel_multiplier=channel_multiplier,
+            narrow=narrow,
+            blur_kernel=blur_kernel,
+            lr_mlp=lr_mlp,
+            is_concat=is_concat)
+
+        conv = [ConvLayer(3, channels[size], 1)]
+        self.ecd0 = nn.Sequential(*conv)
+        in_channel = channels[size]
+
+        self.names = ['ecd%d' % i for i in range(self.log_size - 1)]
+        for i in range(self.log_size, 2, -1):
+            out_channel = channels[2**(i - 1)]
+            conv = [ConvLayer(in_channel, out_channel, 3, downsample=True)]
+            setattr(self, self.names[self.log_size - i + 1],
+                    nn.Sequential(*conv))
+            in_channel = out_channel
+        self.final_linear = nn.Sequential(
+            EqualLinear(channels[4] * 4 * 4,
+                        style_dim,
+                        activation='fused_lrelu'))
+
+    def forward(
+        self,
+        inputs,
+        return_latents=False,
+        inject_index=None,
+        truncation=1,
+        truncation_latent=None,
+        input_is_latent=False,
+    ):
+        noise = []
+        for i in range(self.log_size - 1):
+            ecd = getattr(self, self.names[i])
+            inputs = ecd(inputs)
+            noise.append(inputs)
+        inputs = inputs.reshape([inputs.shape[0], -1])
+        outs = self.final_linear(inputs)
+        noise = list(
+            itertools.chain.from_iterable(
+                itertools.repeat(x, 2) for x in noise))[::-1]
+        outs = self.generator([outs],
+                              return_latents,
+                              inject_index,
+                              truncation,
+                              truncation_latent,
+                              input_is_latent,
+                              noise=noise[1:])
+        return outs

ppgan/models/generators/gpen.py

@@ -14,16 +14,14 @@
 
 # code was heavily based on code was heavily based on https://github.com/yangxy/GPEN
 
-import itertools
-import paddle.nn as nn
 import math
+import paddle
+import paddle.nn as nn
 from ppgan.models.generators.builder import GENERATORS
-from ppgan.modules.equalized import EqualLinear_gpen as EqualLinear
-from ppgan.models.generators.generator_gpen import StyleGANv2Generator
+from ppgan.models.generators import StyleGANv2Generator
 from ppgan.models.discriminators.discriminator_styleganv2 import ConvLayer
+from ppgan.modules.equalized import EqualLinear
 
-
-@GENERATORS.register()
 class GPEN(nn.Layer):
 
     def __init__(
@@ -32,23 +30,21 @@ class GPEN(nn.Layer):
         style_dim,
         n_mlp,
         channel_multiplier=2,
-        narrow=1,
         blur_kernel=[1, 3, 3, 1],
         lr_mlp=0.01,
         is_concat=True,
     ):
         super(GPEN, self).__init__()
         channels = {
-            4: int(512 * narrow),
-            8: int(512 * narrow),
-            16: int(512 * narrow),
-            32: int(512 * narrow),
-            64: int(256 * channel_multiplier * narrow),
-            128: int(128 * channel_multiplier * narrow),
-            256: int(64 * channel_multiplier * narrow),
-            512: int(32 * channel_multiplier * narrow),
-            1024: int(16 * channel_multiplier * narrow),
-            2048: int(8 * channel_multiplier * narrow)
+            4: 512,
+            8: 512,
+            16: 512,
+            32: 512,
+            64: 256 * channel_multiplier,
+            128: 128 * channel_multiplier,
+            256: 64 * channel_multiplier,
+            512: 32 * channel_multiplier,
+            1024: 16 * channel_multiplier,
         }
         self.log_size = int(math.log(size, 2))
         self.generator = StyleGANv2Generator(
@@ -56,7 +52,6 @@ class GPEN(nn.Layer):
             style_dim,
             n_mlp,
             channel_multiplier=channel_multiplier,
-            narrow=narrow,
             blur_kernel=blur_kernel,
             lr_mlp=lr_mlp,
             is_concat=is_concat)
@@ -93,14 +88,7 @@ class GPEN(nn.Layer):
             noise.append(inputs)
         inputs = inputs.reshape([inputs.shape[0], -1])
         outs = self.final_linear(inputs)
-        noise = list(
-            itertools.chain.from_iterable(
-                itertools.repeat(x, 2) for x in noise))[::-1]
-        outs = self.generator([outs],
-                              return_latents,
-                              inject_index,
-                              truncation,
-                              truncation_latent,
-                              input_is_latent,
-                              noise=noise[1:])
+        outs = self.generator([outs], return_latents, inject_index, truncation,
+                              truncation_latent, input_is_latent, 
+                              noise=noise[::-1])
         return outs