update

components/Generator_256.py

@@ -0,0 +1,304 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+#############################################################
+# File: Generator_Invobn_config1.py
+# Created Date: Saturday February 26th 2022
+# Author: Chen Xuanhong
+# Email: chenxuanhongzju@outlook.com
+# Last Modified:  Tuesday, 19th April 2022 7:03:46 pm
+# Modified By: Chen Xuanhong
+# Copyright (c) 2022 Shanghai Jiao Tong University
+#############################################################
+
+import os
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+import math
+
+
+
+class ResBlk(nn.Module):
+    def __init__(self, dim_in, dim_out, actv=nn.LeakyReLU(0.2),
+                 normalize="in", downsample=False):
+        super().__init__()
+        self.actv = actv
+        self.normalize = normalize
+        self.downsample = downsample
+        self.learned_sc = dim_in != dim_out
+        self.equal_var  = math.sqrt(2)
+        self._build_weights(dim_in, dim_out)
+
+    def _build_weights(self, dim_in, dim_out):
+        self.conv1 = nn.Conv2d(dim_in, dim_in, 3, 1, 1)
+        self.conv2 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        if self.normalize.lower() == "in":
+            self.norm1 = nn.InstanceNorm2d(dim_in, affine=True)
+            self.norm2 = nn.InstanceNorm2d(dim_in, affine=True)
+        elif self.normalize.lower() == "bn":
+            self.norm1 = nn.BatchNorm2d(dim_in)
+            self.norm2 = nn.BatchNorm2d(dim_in)
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        if self.downsample:
+            x = F.avg_pool2d(x, 2)
+        return x
+
+    def _residual(self, x):
+        if self.normalize:
+            x = self.norm1(x)
+        x = self.actv(x)
+        x = self.conv1(x)
+        if self.downsample:
+            x = F.avg_pool2d(x, 2)
+        if self.normalize:
+            x = self.norm2(x)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x):
+        x = self._shortcut(x) + self._residual(x)
+        return x /self.equal_var   # unit variance
+
+class AdaIN(nn.Module):
+    def __init__(self, style_dim, num_features):
+        super().__init__()
+        self.norm = nn.InstanceNorm2d(num_features, affine=False)
+        self.fc = nn.Linear(style_dim, num_features*2)
+
+    def forward(self, x, s):
+        h = self.fc(s)
+        h = h.view(h.size(0), h.size(1), 1, 1)
+        gamma, beta = torch.chunk(h, chunks=2, dim=1)
+        return (1 + gamma) * self.norm(x) + beta
+
+class ResUpBlk(nn.Module):
+    def __init__(self, dim_in, dim_out,actv=nn.LeakyReLU(0.2),normalize="in"):
+        super().__init__()
+        self.actv = actv
+        self.normalize = normalize
+        self.learned_sc = dim_in != dim_out
+        self.equal_var = math.sqrt(2)
+        self._build_weights(dim_in, dim_out)
+
+    def _build_weights(self, dim_in, dim_out):
+        self.conv1 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        self.conv2 = nn.Conv2d(dim_out, dim_out, 3, 1, 1)
+        if self.normalize.lower() == "in":
+            self.norm1 = nn.InstanceNorm2d(dim_in, affine=True)
+            self.norm2 = nn.InstanceNorm2d(dim_out, affine=True)
+        elif self.normalize.lower() == "bn":
+            self.norm1 = nn.BatchNorm2d(dim_in)
+            self.norm2 = nn.BatchNorm2d(dim_out)
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        x = F.interpolate(x, scale_factor=2, mode='nearest')
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        return x
+
+    def _residual(self, x):
+        x = self.norm1(x)
+        x = self.actv(x)
+        x = F.interpolate(x, scale_factor=2, mode='nearest')
+        x = self.conv1(x)
+        x = self.norm2(x)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x):
+        out = self._residual(x)
+        out = (out + self._shortcut(x)) / self.equal_var
+        return out
+
+class AdainResBlk(nn.Module):
+    def __init__(self, dim_in, dim_out, style_dim=512,
+                 actv=nn.LeakyReLU(0.2), upsample=False):
+        super().__init__()
+        self.actv = actv
+        self.upsample = upsample
+        self.learned_sc = dim_in != dim_out
+        self.equal_var = math.sqrt(2)
+        self._build_weights(dim_in, dim_out, style_dim)
+
+    def _build_weights(self, dim_in, dim_out, style_dim=64):
+        self.conv1 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        self.conv2 = nn.Conv2d(dim_out, dim_out, 3, 1, 1)
+        self.norm1 = AdaIN(style_dim, dim_in)
+        self.norm2 = AdaIN(style_dim, dim_out)
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        if self.upsample:
+            x = F.interpolate(x, scale_factor=2, mode='nearest')
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        return x
+
+    def _residual(self, x, s):
+        x = self.norm1(x, s)
+        x = self.actv(x)
+        if self.upsample:
+            x = F.interpolate(x, scale_factor=2, mode='nearest')
+        x = self.conv1(x)
+        x = self.norm2(x, s)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x, s):
+        out = self._residual(x, s)
+        out = (out + self._shortcut(x)) / self.equal_var
+        return out
+
+
+class Generator(nn.Module):
+    def __init__(
+                self,
+                **kwargs
+                ):
+        super().__init__()
+
+        id_dim      = kwargs["id_dim"]
+        k_size      = kwargs["g_kernel_size"]
+        res_num     = kwargs["res_num"]
+        in_channel  = kwargs["in_channel"]
+        up_mode     = kwargs["up_mode"]
+        norm        = kwargs["norm"].lower()
+        
+        aggregator  = kwargs["aggregator"]
+        res_mode    = kwargs["res_mode"]
+
+        padding_size= int((k_size -1)/2)
+        padding_type= 'reflect'
+
+        if norm.lower() == "in":
+            norm_out = nn.InstanceNorm2d(in_channel, affine=True)
+            norm_mask= nn.InstanceNorm2d
+        elif norm.lower() == "bn":
+            norm_out = nn.BatchNorm2d(in_channel)
+            norm_mask = nn.BatchNorm2d
+        
+        
+        activation = nn.LeakyReLU(0.2)
+        # activation = nn.ReLU()
+
+        self.from_rgb = nn.Conv2d(3, in_channel, 1, 1, 0)
+        # self.first_layer = nn.Sequential(nn.Conv2d(3, 64, kernel_size=3, padding=1, bias=False),
+        #                         nn.BatchNorm2d(64), activation)
+        ### downsample
+        self.down1 = ResBlk(in_channel, in_channel, normalize=norm, downsample=True)# 128
+
+        self.down2 = ResBlk(in_channel, in_channel*2, normalize=norm, downsample=True)# 64
+                                
+        self.down3 = ResBlk(in_channel*2, in_channel*4,normalize=norm, downsample=True)# 32
+
+        self.down4 = ResBlk(in_channel*4, in_channel*8, normalize=norm, downsample=True)# 16
+
+        # self.down6 = ResBlk(in_channel*8, in_channel*8, normalize=True, downsample=True)# 8
+        
+        
+        ### resnet blocks
+        BN = []
+        for i in range(res_num):
+            BN += [
+                AdainResBlk(in_channel*8, in_channel*8, style_dim=id_dim, upsample=False)]
+        self.BottleNeck = nn.Sequential(*BN)
+
+        # self.up6 = AdainResBlk(in_channel*8, in_channel*8, style_dim=id_dim, upsample=True) # 16
+        
+        # self.up5 = AdainResBlk(in_channel*8, in_channel*8, style_dim=id_dim, upsample=True) # 1
+
+        self.up4 = AdainResBlk(in_channel*8, in_channel*4, style_dim=id_dim, upsample=True) # 32
+
+        self.up3 = AdainResBlk(in_channel*4, in_channel*2, style_dim=id_dim, upsample=True) # 64
+
+        # self.maskhead = nn.Sequential(
+        #                         nn.Conv2d(in_channel*2, in_channel//2, kernel_size=3, stride=1, padding=1, bias=False),
+        #                         norm_mask, # 64
+        #                         activation,
+        #                         nn.Conv2d(in_channel//2, 1, kernel_size=3, stride=1, padding=1),
+        #                         nn.Sigmoid())
+        self.maskhead_lr = nn.Sequential(
+                                nn.UpsamplingNearest2d(scale_factor = 2),
+                                nn.Conv2d(in_channel*8, in_channel, kernel_size=3, stride=1, padding=1,bias=False),
+                                norm_mask(in_channel, affine=True), # 32
+                                activation,
+                                nn.UpsamplingNearest2d(scale_factor = 2),
+                                nn.Conv2d(in_channel, in_channel//4, kernel_size=3, stride=1, padding=1, bias=False),
+                                norm_mask(in_channel//4, affine=True), # 64
+                                activation
+        )
+        self.maskhead_hr = nn.Sequential(
+                                nn.UpsamplingNearest2d(scale_factor = 2),
+                                nn.Conv2d(in_channel//4, in_channel//16, kernel_size=3, stride=1, padding=1,bias=False),
+                                norm_mask(in_channel//16, affine=True), # 128
+                                activation,
+                                nn.UpsamplingNearest2d(scale_factor = 2),
+                                nn.Conv2d(in_channel//16, 1, kernel_size=3, stride=1, padding=1),
+                                nn.Sigmoid() # 256
+        )
+        self.maskhead_out = nn.Sequential(nn.Conv2d(in_channel//4, 1, kernel_size=1, stride=1),
+                                nn.Sigmoid())
+        
+        # self.up2 = AdainResBlk(in_channel*2, in_channel, style_dim=id_dim, upsample=True)
+        # self.up2 = AdainResBlk(in_channel*2, in_channel, style_dim=id_dim, upsample=True)
+        self.up2 = ResUpBlk(in_channel*2, in_channel, normalize=norm)
+
+        # self.up1 = AdainResBlk(in_channel, in_channel, style_dim=id_dim, upsample=True)
+        self.up1 = ResUpBlk(in_channel, in_channel, normalize=norm)
+        # ResUpBlk(in_channel, in_channel, normalize="in") # 512
+
+        
+        
+        self.to_rgb = nn.Sequential(
+            norm_out,
+            activation,
+            nn.Conv2d(in_channel, 3, 3, 1, 1))
+
+        # self.last_layer = nn.Sequential(nn.ReflectionPad2d(3),
+        #             nn.Conv2d(64, 3, kernel_size=7, padding=0))
+
+
+    #     self.__weights_init__()
+
+    # def __weights_init__(self):
+    #     for layer in self.encoder:
+    #         if isinstance(layer,nn.Conv2d):
+    #             nn.init.xavier_uniform_(layer.weight)
+
+    #     for layer in self.encoder2:
+    #         if isinstance(layer,nn.Conv2d):
+    #             nn.init.xavier_uniform_(layer.weight)
+
+    def forward(self, img, id):
+        res = self.from_rgb(img)
+        res = self.down1(res)
+        skip = self.down2(res)
+        res = self.down3(skip)
+        res = self.down4(res)
+        mask_feat= self.maskhead_lr(res)
+        
+        for i in range(len(self.BottleNeck)):
+            res = self.BottleNeck[i](res, id)
+        res = self.up4(res,id)
+        res = self.up3(res,id)
+        mask_lr= self.maskhead_out(mask_feat)
+        # res = (1-mask) * self.sigma(skip) + mask * res
+        res = (1-mask_lr) * skip + mask_lr * res
+        res = self.up2(res) #  + skip
+        res = self.up1(res)
+        res = self.to_rgb(res)
+        mask_hr=self.maskhead_hr(mask_feat)
+        res = (1-mask_hr) * img + mask_hr * res
+        return res, mask_lr, mask_hr

components/Generator_2mask.py

@@ -0,0 +1,312 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+#############################################################
+# File: Generator_Invobn_config1.py
+# Created Date: Saturday February 26th 2022
+# Author: Chen Xuanhong
+# Email: chenxuanhongzju@outlook.com
+# Last Modified:  Friday, 15th April 2022 12:30:27 am
+# Modified By: Chen Xuanhong
+# Copyright (c) 2022 Shanghai Jiao Tong University
+#############################################################
+
+import os
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+import math
+
+
+
+class ResBlk(nn.Module):
+    def __init__(self, dim_in, dim_out, actv=nn.LeakyReLU(0.2),
+                 normalize="in", downsample=False):
+        super().__init__()
+        self.actv = actv
+        self.normalize = normalize
+        self.downsample = downsample
+        self.learned_sc = dim_in != dim_out
+        self.equal_var  = math.sqrt(2)
+        self._build_weights(dim_in, dim_out)
+
+    def _build_weights(self, dim_in, dim_out):
+        self.conv1 = nn.Conv2d(dim_in, dim_in, 3, 1, 1)
+        self.conv2 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        if self.normalize.lower() == "in":
+            self.norm1 = nn.InstanceNorm2d(dim_in, affine=True)
+            self.norm2 = nn.InstanceNorm2d(dim_in, affine=True)
+        elif self.normalize.lower() == "bn":
+            self.norm1 = nn.BatchNorm2d(dim_in)
+            self.norm2 = nn.BatchNorm2d(dim_in)
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        if self.downsample:
+            x = F.avg_pool2d(x, 2)
+        return x
+
+    def _residual(self, x):
+        if self.normalize:
+            x = self.norm1(x)
+        x = self.actv(x)
+        x = self.conv1(x)
+        if self.downsample:
+            x = F.avg_pool2d(x, 2)
+        if self.normalize:
+            x = self.norm2(x)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x):
+        x = self._shortcut(x) + self._residual(x)
+        return x /self.equal_var   # unit variance
+
+class AdaIN(nn.Module):
+    def __init__(self, style_dim, num_features):
+        super().__init__()
+        self.norm = nn.InstanceNorm2d(num_features, affine=False)
+        self.fc = nn.Linear(style_dim, num_features*2)
+
+    def forward(self, x, s):
+        h = self.fc(s)
+        h = h.view(h.size(0), h.size(1), 1, 1)
+        gamma, beta = torch.chunk(h, chunks=2, dim=1)
+        return (1 + gamma) * self.norm(x) + beta
+
+class ResUpBlk(nn.Module):
+    def __init__(self, dim_in, dim_out,actv=nn.LeakyReLU(0.2),normalize="in"):
+        super().__init__()
+        self.actv = actv
+        self.normalize = normalize
+        self.learned_sc = dim_in != dim_out
+        self.equal_var = math.sqrt(2)
+        self._build_weights(dim_in, dim_out)
+
+    def _build_weights(self, dim_in, dim_out):
+        self.conv1 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        self.conv2 = nn.Conv2d(dim_out, dim_out, 3, 1, 1)
+        if self.normalize.lower() == "in":
+            self.norm1 = nn.InstanceNorm2d(dim_in, affine=True)
+            self.norm2 = nn.InstanceNorm2d(dim_out, affine=True)
+        elif self.normalize.lower() == "bn":
+            self.norm1 = nn.BatchNorm2d(dim_in)
+            self.norm2 = nn.BatchNorm2d(dim_out)
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        x = F.interpolate(x, scale_factor=2, mode='nearest')
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        return x
+
+    def _residual(self, x):
+        x = self.norm1(x)
+        x = self.actv(x)
+        x = F.interpolate(x, scale_factor=2, mode='nearest')
+        x = self.conv1(x)
+        x = self.norm2(x)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x):
+        out = self._residual(x)
+        out = (out + self._shortcut(x)) / self.equal_var
+        return out
+
+class AdainResBlk(nn.Module):
+    def __init__(self, dim_in, dim_out, style_dim=512,
+                 actv=nn.LeakyReLU(0.2), upsample=False):
+        super().__init__()
+        self.actv = actv
+        self.upsample = upsample
+        self.learned_sc = dim_in != dim_out
+        self.equal_var = math.sqrt(2)
+        self._build_weights(dim_in, dim_out, style_dim)
+
+    def _build_weights(self, dim_in, dim_out, style_dim=64):
+        self.conv1 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        self.conv2 = nn.Conv2d(dim_out, dim_out, 3, 1, 1)
+        self.norm1 = AdaIN(style_dim, dim_in)
+        self.norm2 = AdaIN(style_dim, dim_out)
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        if self.upsample:
+            x = F.interpolate(x, scale_factor=2, mode='nearest')
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        return x
+
+    def _residual(self, x, s):
+        x = self.norm1(x, s)
+        x = self.actv(x)
+        if self.upsample:
+            x = F.interpolate(x, scale_factor=2, mode='nearest')
+        x = self.conv1(x)
+        x = self.norm2(x, s)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x, s):
+        out = self._residual(x, s)
+        out = (out + self._shortcut(x)) / self.equal_var
+        return out
+
+
+class Generator(nn.Module):
+    def __init__(
+                self,
+                **kwargs
+                ):
+        super().__init__()
+
+        id_dim      = kwargs["id_dim"]
+        k_size      = kwargs["g_kernel_size"]
+        res_num     = kwargs["res_num"]
+        in_channel  = kwargs["in_channel"]
+        up_mode     = kwargs["up_mode"]
+        norm        = kwargs["norm"].lower()
+        
+        aggregator  = kwargs["aggregator"]
+        res_mode    = kwargs["res_mode"]
+
+        padding_size= int((k_size -1)/2)
+        padding_type= 'reflect'
+
+        if norm.lower() == "in":
+            norm_out = nn.InstanceNorm2d(in_channel, affine=True)
+            norm_mask= nn.InstanceNorm2d
+        elif norm.lower() == "bn":
+            norm_out = nn.BatchNorm2d(in_channel)
+            norm_mask = nn.BatchNorm2d
+        
+        
+        activation = nn.LeakyReLU(0.2)
+        # activation = nn.ReLU()
+
+        self.from_rgb = nn.Conv2d(3, in_channel, 1, 1, 0)
+        # self.first_layer = nn.Sequential(nn.Conv2d(3, 64, kernel_size=3, padding=1, bias=False),
+        #                         nn.BatchNorm2d(64), activation)
+        ### downsample
+        self.down1 = ResBlk(in_channel, in_channel, normalize=norm, downsample=True)# 256
+
+        self.down2 = ResBlk(in_channel, in_channel*2, normalize=norm, downsample=True)# 128
+                                
+        self.down3 = ResBlk(in_channel*2, in_channel*4,normalize=norm, downsample=True)# 64
+
+        self.down4 = ResBlk(in_channel*4, in_channel*8, normalize=norm, downsample=True)# 32
+
+        self.down5 = ResBlk(in_channel*8, in_channel*8, normalize=norm, downsample=True)# 16
+
+        # self.down6 = ResBlk(in_channel*8, in_channel*8, normalize=True, downsample=True)# 8
+        
+        
+        ### resnet blocks
+        BN = []
+        for i in range(res_num):
+            BN += [
+                AdainResBlk(in_channel*8, in_channel*8, style_dim=id_dim, upsample=False)]
+        self.BottleNeck = nn.Sequential(*BN)
+
+        # self.up6 = AdainResBlk(in_channel*8, in_channel*8, style_dim=id_dim, upsample=True) # 16
+        
+        self.up5 = AdainResBlk(in_channel*8, in_channel*8, style_dim=id_dim, upsample=True) # 32
+
+        self.up4 = AdainResBlk(in_channel*8, in_channel*4, style_dim=id_dim, upsample=True) # 64
+
+        self.up3 = AdainResBlk(in_channel*4, in_channel*2, style_dim=id_dim, upsample=True) # 128
+
+        # self.maskhead = nn.Sequential(
+        #                         nn.Conv2d(in_channel*2, in_channel//2, kernel_size=3, stride=1, padding=1, bias=False),
+        #                         norm_mask, # 64
+        #                         activation,
+        #                         nn.Conv2d(in_channel//2, 1, kernel_size=3, stride=1, padding=1),
+        #                         nn.Sigmoid())
+        self.maskhead_lr = nn.Sequential(
+                                nn.UpsamplingNearest2d(scale_factor = 2),
+                                nn.Conv2d(in_channel*8, in_channel, kernel_size=3, stride=1, padding=1,bias=False),
+                                norm_mask(in_channel, affine=True), # 32
+                                activation,
+                                nn.UpsamplingNearest2d(scale_factor = 2),
+                                nn.Conv2d(in_channel, in_channel//4, kernel_size=3, stride=1, padding=1, bias=False),
+                                norm_mask(in_channel//4, affine=True), # 64
+                                activation,
+                                nn.UpsamplingNearest2d(scale_factor = 2),
+                                nn.Conv2d(in_channel//4, in_channel//8, kernel_size=3, stride=1, padding=1),
+                                norm_mask(in_channel//8, affine=True), # 128
+                                activation,
+        )
+        self.maskhead_hr = nn.Sequential(
+                                nn.UpsamplingNearest2d(scale_factor = 2),
+                                nn.Conv2d(in_channel//8, in_channel//16, kernel_size=3, stride=1, padding=1,bias=False),
+                                norm_mask(in_channel//16, affine=True), # 256
+                                activation,
+                                nn.UpsamplingNearest2d(scale_factor = 2),
+                                nn.Conv2d(in_channel//16, 1, kernel_size=3, stride=1, padding=1),
+                                nn.Sigmoid() # 512
+        )
+        self.maskhead_out = nn.Sequential(nn.Conv2d(in_channel//8, 1, kernel_size=1, stride=1),
+                                nn.Sigmoid())
+        
+        # self.up2 = AdainResBlk(in_channel*2, in_channel, style_dim=id_dim, upsample=True)
+        # self.up2 = AdainResBlk(in_channel*2, in_channel, style_dim=id_dim, upsample=True)
+        self.up2 = ResUpBlk(in_channel*2, in_channel, normalize=norm)
+
+        # self.up1 = AdainResBlk(in_channel, in_channel, style_dim=id_dim, upsample=True)
+        self.up1 = ResUpBlk(in_channel, in_channel, normalize=norm)
+        # ResUpBlk(in_channel, in_channel, normalize="in") # 512
+
+        
+        
+        self.to_rgb = nn.Sequential(
+            norm_out,
+            activation,
+            nn.Conv2d(in_channel, 3, 3, 1, 1))
+
+        # self.last_layer = nn.Sequential(nn.ReflectionPad2d(3),
+        #             nn.Conv2d(64, 3, kernel_size=7, padding=0))
+
+
+    #     self.__weights_init__()
+
+    # def __weights_init__(self):
+    #     for layer in self.encoder:
+    #         if isinstance(layer,nn.Conv2d):
+    #             nn.init.xavier_uniform_(layer.weight)
+
+    #     for layer in self.encoder2:
+    #         if isinstance(layer,nn.Conv2d):
+    #             nn.init.xavier_uniform_(layer.weight)
+
+    def forward(self, img, id):
+        res = self.from_rgb(img)
+        res = self.down1(res)
+        skip = self.down2(res)
+        res = self.down3(skip)
+        res = self.down4(res)
+        res = self.down5(res)
+        mask_feat= self.maskhead_lr(res)
+        
+        for i in range(len(self.BottleNeck)):
+            res = self.BottleNeck[i](res, id)
+        res = self.up5(res,id)
+        res = self.up4(res,id)
+        res = self.up3(res,id)
+        mask_lr= self.maskhead_out(mask_feat)
+        # res = (1-mask) * self.sigma(skip) + mask * res
+        res = (1-mask_lr) * skip + mask_lr * res
+        res = self.up2(res) #  + skip
+        res = self.up1(res)
+        res = self.to_rgb(res)
+        mask_hr=self.maskhead_hr(mask_feat)
+        res = (1-mask_hr) * img + mask_hr * res
+        return res, mask_lr, mask_hr

components/Generator_2mask2.py

@@ -0,0 +1,312 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+#############################################################
+# File: Generator_Invobn_config1.py
+# Created Date: Saturday February 26th 2022
+# Author: Chen Xuanhong
+# Email: chenxuanhongzju@outlook.com
+# Last Modified:  Tuesday, 19th April 2022 12:45:55 am
+# Modified By: Chen Xuanhong
+# Copyright (c) 2022 Shanghai Jiao Tong University
+#############################################################
+
+import os
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+import math
+
+
+
+class ResBlk(nn.Module):
+    def __init__(self, dim_in, dim_out, actv=nn.LeakyReLU(0.2),
+                 normalize="in", downsample=False):
+        super().__init__()
+        self.actv = actv
+        self.normalize = normalize
+        self.downsample = downsample
+        self.learned_sc = dim_in != dim_out
+        self.equal_var  = math.sqrt(2)
+        self._build_weights(dim_in, dim_out)
+
+    def _build_weights(self, dim_in, dim_out):
+        self.conv1 = nn.Conv2d(dim_in, dim_in, 3, 1, 1)
+        self.conv2 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        if self.normalize.lower() == "in":
+            self.norm1 = nn.InstanceNorm2d(dim_in, affine=True)
+            self.norm2 = nn.InstanceNorm2d(dim_in, affine=True)
+        elif self.normalize.lower() == "bn":
+            self.norm1 = nn.BatchNorm2d(dim_in)
+            self.norm2 = nn.BatchNorm2d(dim_in)
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        if self.downsample:
+            x = F.avg_pool2d(x, 2)
+        return x
+
+    def _residual(self, x):
+        if self.normalize:
+            x = self.norm1(x)
+        x = self.actv(x)
+        x = self.conv1(x)
+        if self.downsample:
+            x = F.avg_pool2d(x, 2)
+        if self.normalize:
+            x = self.norm2(x)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x):
+        x = self._shortcut(x) + self._residual(x)
+        return x /self.equal_var   # unit variance
+
+class AdaIN(nn.Module):
+    def __init__(self, style_dim, num_features):
+        super().__init__()
+        self.norm = nn.InstanceNorm2d(num_features, affine=False)
+        self.fc = nn.Linear(style_dim, num_features*2)
+
+    def forward(self, x, s):
+        h = self.fc(s)
+        h = h.view(h.size(0), h.size(1), 1, 1)
+        gamma, beta = torch.chunk(h, chunks=2, dim=1)
+        return (1 + gamma) * self.norm(x) + beta
+
+class ResUpBlk(nn.Module):
+    def __init__(self, dim_in, dim_out,actv=nn.LeakyReLU(0.2),normalize="in"):
+        super().__init__()
+        self.actv = actv
+        self.normalize = normalize
+        self.learned_sc = dim_in != dim_out
+        self.equal_var = math.sqrt(2)
+        self._build_weights(dim_in, dim_out)
+
+    def _build_weights(self, dim_in, dim_out):
+        self.conv1 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        self.conv2 = nn.Conv2d(dim_out, dim_out, 3, 1, 1)
+        if self.normalize.lower() == "in":
+            self.norm1 = nn.InstanceNorm2d(dim_in, affine=True)
+            self.norm2 = nn.InstanceNorm2d(dim_out, affine=True)
+        elif self.normalize.lower() == "bn":
+            self.norm1 = nn.BatchNorm2d(dim_in)
+            self.norm2 = nn.BatchNorm2d(dim_out)
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        x = F.interpolate(x, scale_factor=2, mode='nearest')
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        return x
+
+    def _residual(self, x):
+        x = self.norm1(x)
+        x = self.actv(x)
+        x = F.interpolate(x, scale_factor=2, mode='nearest')
+        x = self.conv1(x)
+        x = self.norm2(x)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x):
+        out = self._residual(x)
+        out = (out + self._shortcut(x)) / self.equal_var
+        return out
+
+class AdainResBlk(nn.Module):
+    def __init__(self, dim_in, dim_out, style_dim=512,
+                 actv=nn.LeakyReLU(0.2), upsample=False):
+        super().__init__()
+        self.actv = actv
+        self.upsample = upsample
+        self.learned_sc = dim_in != dim_out
+        self.equal_var = math.sqrt(2)
+        self._build_weights(dim_in, dim_out, style_dim)
+
+    def _build_weights(self, dim_in, dim_out, style_dim=64):
+        self.conv1 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        self.conv2 = nn.Conv2d(dim_out, dim_out, 3, 1, 1)
+        self.norm1 = AdaIN(style_dim, dim_in)
+        self.norm2 = AdaIN(style_dim, dim_out)
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        if self.upsample:
+            x = F.interpolate(x, scale_factor=2, mode='nearest')
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        return x
+
+    def _residual(self, x, s):
+        x = self.norm1(x, s)
+        x = self.actv(x)
+        if self.upsample:
+            x = F.interpolate(x, scale_factor=2, mode='nearest')
+        x = self.conv1(x)
+        x = self.norm2(x, s)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x, s):
+        out = self._residual(x, s)
+        out = (out + self._shortcut(x)) / self.equal_var
+        return out
+
+
+class Generator(nn.Module):
+    def __init__(
+                self,
+                **kwargs
+                ):
+        super().__init__()
+
+        id_dim      = kwargs["id_dim"]
+        k_size      = kwargs["g_kernel_size"]
+        res_num     = kwargs["res_num"]
+        in_channel  = kwargs["in_channel"]
+        up_mode     = kwargs["up_mode"]
+        norm        = kwargs["norm"].lower()
+        
+        aggregator  = kwargs["aggregator"]
+        res_mode    = kwargs["res_mode"]
+
+        padding_size= int((k_size -1)/2)
+        padding_type= 'reflect'
+
+        if norm.lower() == "in":
+            norm_out = nn.InstanceNorm2d(in_channel, affine=True)
+            norm_mask= nn.InstanceNorm2d
+        elif norm.lower() == "bn":
+            norm_out = nn.BatchNorm2d(in_channel)
+            norm_mask = nn.BatchNorm2d
+        
+        
+        activation = nn.LeakyReLU(0.2)
+        # activation = nn.ReLU()
+
+        self.from_rgb = nn.Conv2d(3, in_channel, 1, 1, 0)
+        # self.first_layer = nn.Sequential(nn.Conv2d(3, 64, kernel_size=3, padding=1, bias=False),
+        #                         nn.BatchNorm2d(64), activation)
+        ### downsample
+        self.down1 = ResBlk(in_channel, in_channel, normalize=norm, downsample=True)# 256
+
+        self.down2 = ResBlk(in_channel, in_channel*2, normalize=norm, downsample=True)# 128
+                                
+        self.down3 = ResBlk(in_channel*2, in_channel*4,normalize=norm, downsample=True)# 64
+
+        self.down4 = ResBlk(in_channel*4, in_channel*8, normalize=norm, downsample=True)# 32
+
+        self.down5 = ResBlk(in_channel*8, in_channel*8, normalize=norm, downsample=True)# 16
+
+        # self.down6 = ResBlk(in_channel*8, in_channel*8, normalize=True, downsample=True)# 8
+        
+        
+        ### resnet blocks
+        BN = []
+        for i in range(res_num):
+            BN += [
+                AdainResBlk(in_channel*8, in_channel*8, style_dim=id_dim, upsample=False)]
+        self.BottleNeck = nn.Sequential(*BN)
+
+        # self.up6 = AdainResBlk(in_channel*8, in_channel*8, style_dim=id_dim, upsample=True) # 16
+        
+        self.up5 = AdainResBlk(in_channel*8, in_channel*8, style_dim=id_dim, upsample=True) # 32
+
+        self.up4 = AdainResBlk(in_channel*8, in_channel*4, style_dim=id_dim, upsample=True) # 64
+
+        self.up3 = AdainResBlk(in_channel*4, in_channel*2, style_dim=id_dim, upsample=True) # 128
+
+        # self.maskhead = nn.Sequential(
+        #                         nn.Conv2d(in_channel*2, in_channel//2, kernel_size=3, stride=1, padding=1, bias=False),
+        #                         norm_mask, # 64
+        #                         activation,
+        #                         nn.Conv2d(in_channel//2, 1, kernel_size=3, stride=1, padding=1),
+        #                         nn.Sigmoid())
+        self.maskhead_lr = nn.Sequential(
+                                nn.UpsamplingNearest2d(scale_factor = 2),
+                                nn.Conv2d(in_channel*8, in_channel, kernel_size=3, stride=1, padding=1,bias=False),
+                                norm_mask(in_channel, affine=True), # 32
+                                activation,
+                                nn.UpsamplingNearest2d(scale_factor = 2),
+                                nn.Conv2d(in_channel, in_channel//4, kernel_size=3, stride=1, padding=1, bias=False),
+                                norm_mask(in_channel//4, affine=True), # 64
+                                activation,
+                                nn.UpsamplingNearest2d(scale_factor = 2),
+                                nn.Conv2d(in_channel//4, in_channel//8, kernel_size=3, stride=1, padding=1),
+                                norm_mask(in_channel//8, affine=True), # 128
+                                activation,
+        )
+        self.maskhead_hr = nn.Sequential(
+                                nn.UpsamplingNearest2d(scale_factor = 2),
+                                nn.Conv2d(in_channel//8, in_channel//16, kernel_size=3, stride=1, padding=1,bias=False),
+                                norm_mask(in_channel//16, affine=True), # 256
+                                activation,
+                                nn.UpsamplingNearest2d(scale_factor = 2),
+                                nn.Conv2d(in_channel//16, 1, kernel_size=3, stride=1, padding=1),
+                                nn.Sigmoid() # 512
+        )
+        self.maskhead_out = nn.Sequential(nn.Conv2d(in_channel//8, 1, kernel_size=1, stride=1),
+                                nn.Sigmoid())
+        
+        # self.up2 = AdainResBlk(in_channel*2, in_channel, style_dim=id_dim, upsample=True)
+        self.up2 = AdainResBlk(in_channel*2, in_channel, style_dim=id_dim, upsample=True)
+        # self.up2 = ResUpBlk(in_channel*2, in_channel, normalize=norm)
+
+        self.up1 = AdainResBlk(in_channel, in_channel, style_dim=id_dim, upsample=True)
+        # self.up1 = ResUpBlk(in_channel, in_channel, normalize=norm)
+        # ResUpBlk(in_channel, in_channel, normalize="in") # 512
+
+        
+        
+        self.to_rgb = nn.Sequential(
+            norm_out,
+            activation,
+            nn.Conv2d(in_channel, 3, 3, 1, 1))
+
+        # self.last_layer = nn.Sequential(nn.ReflectionPad2d(3),
+        #             nn.Conv2d(64, 3, kernel_size=7, padding=0))
+
+
+    #     self.__weights_init__()
+
+    # def __weights_init__(self):
+    #     for layer in self.encoder:
+    #         if isinstance(layer,nn.Conv2d):
+    #             nn.init.xavier_uniform_(layer.weight)
+
+    #     for layer in self.encoder2:
+    #         if isinstance(layer,nn.Conv2d):
+    #             nn.init.xavier_uniform_(layer.weight)
+
+    def forward(self, img, id):
+        res = self.from_rgb(img)
+        res = self.down1(res)
+        skip = self.down2(res)
+        res = self.down3(skip)
+        res = self.down4(res)
+        res = self.down5(res)
+        mask_feat= self.maskhead_lr(res)
+        
+        for i in range(len(self.BottleNeck)):
+            res = self.BottleNeck[i](res, id)
+        res = self.up5(res,id)
+        res = self.up4(res,id)
+        res = self.up3(res,id)
+        mask_lr= self.maskhead_out(mask_feat)
+        # res = (1-mask) * self.sigma(skip) + mask * res
+        res = (1-mask_lr) * skip + mask_lr * res
+        res = self.up2(res,id) #  + skip
+        res = self.up1(res,id)
+        res = self.to_rgb(res)
+        mask_hr=self.maskhead_hr(mask_feat)
+        res = (1-mask_hr) * img + mask_hr * res
+        return res, mask_lr, mask_hr

components/Generator_2mask_DWConv.py

@@ -0,0 +1,453 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+#############################################################
+# File: Generator_Invobn_config1.py
+# Created Date: Saturday February 26th 2022
+# Author: Chen Xuanhong
+# Email: chenxuanhongzju@outlook.com
+# Last Modified:  Monday, 18th April 2022 10:20:12 pm
+# Modified By: Chen Xuanhong
+# Copyright (c) 2022 Shanghai Jiao Tong University
+#############################################################
+
+import math
+
+import torch
+from   torch import nn
+import torch.nn.functional as F
+
+from   components.ModulatedDWConv import ModulatedDWConv2d
+
+
+class ResBlk(nn.Module):
+    def __init__(self, dim_in, dim_out, actv=nn.LeakyReLU(0.2),
+                 normalize="in", downsample=False):
+        super().__init__()
+        self.actv = actv
+        self.normalize = normalize
+        self.downsample = downsample
+        self.learned_sc = dim_in != dim_out
+        self.equal_var  = math.sqrt(2)
+        self._build_weights(dim_in, dim_out)
+
+    def _build_weights(self, dim_in, dim_out):
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(dim_in, dim_in, 3, 1, 1, groups=dim_in),
+            nn.Conv2d(dim_in, dim_in, 1, 1)
+        )
+        
+        self.conv2 = nn.Sequential(
+            nn.Conv2d(dim_in, dim_in, 3, 1, 1, groups=dim_in),
+            nn.Conv2d(dim_in, dim_out, 1, 1)
+        )
+        # self.conv1 = nn.Conv2d(dim_in, dim_in, 3, 1, 1)
+        # self.conv2 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        if self.normalize.lower() == "in":
+            self.norm1 = nn.InstanceNorm2d(dim_in, affine=True)
+            self.norm2 = nn.InstanceNorm2d(dim_in, affine=True)
+        elif self.normalize.lower() == "bn":
+            self.norm1 = nn.BatchNorm2d(dim_in)
+            self.norm2 = nn.BatchNorm2d(dim_in)
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        if self.downsample:
+            x = F.avg_pool2d(x, 2)
+        return x
+
+    def _residual(self, x):
+        if self.normalize:
+            x = self.norm1(x)
+        x = self.actv(x)
+        x = self.conv1(x)
+        if self.downsample:
+            x = F.avg_pool2d(x, 2)
+        if self.normalize:
+            x = self.norm2(x)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x):
+        x = self._shortcut(x) + self._residual(x)
+        return x /self.equal_var   # unit variance
+
+class AdaIN(nn.Module):
+    def __init__(self, style_dim, num_features):
+        super().__init__()
+        self.norm = nn.InstanceNorm2d(num_features, affine=False)
+        self.fc = nn.Linear(style_dim, num_features*2)
+
+    def forward(self, x, s):
+        h = self.fc(s)
+        h = h.view(h.size(0), h.size(1), 1, 1)
+        gamma, beta = torch.chunk(h, chunks=2, dim=1)
+        return (1 + gamma) * self.norm(x) + beta
+
+class ResUpBlk(nn.Module):
+    def __init__(self, dim_in, dim_out,actv=nn.LeakyReLU(0.2),normalize="in"):
+        super().__init__()
+        self.actv = actv
+        self.normalize = normalize
+        self.learned_sc = dim_in != dim_out
+        self.equal_var = math.sqrt(2)
+        self._build_weights(dim_in, dim_out)
+
+    def _build_weights(self, dim_in, dim_out):
+        self.conv1 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        self.conv2 = nn.Conv2d(dim_out, dim_out, 3, 1, 1)
+        if self.normalize.lower() == "in":
+            self.norm1 = nn.InstanceNorm2d(dim_in, affine=True)
+            self.norm2 = nn.InstanceNorm2d(dim_out, affine=True)
+        elif self.normalize.lower() == "bn":
+            self.norm1 = nn.BatchNorm2d(dim_in)
+            self.norm2 = nn.BatchNorm2d(dim_out)
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        x = F.interpolate(x, scale_factor=2, mode='nearest')
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        return x
+
+    def _residual(self, x):
+        x = self.norm1(x)
+        x = self.actv(x)
+        x = F.interpolate(x, scale_factor=2, mode='nearest')
+        x = self.conv1(x)
+        x = self.norm2(x)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x):
+        out = self._residual(x)
+        out = (out + self._shortcut(x)) / self.equal_var
+        return out
+
+# class ResUpBlk(nn.Module):
+#     def __init__(self, dim_in, dim_out,actv=nn.LeakyReLU(0.2),normalize="in"):
+#         super().__init__()
+#         self.actv = actv
+#         self.normalize = normalize
+#         self.learned_sc = dim_in != dim_out
+#         self.equal_var = math.sqrt(2)
+#         self._build_weights(dim_in, dim_out)
+
+#     def _build_weights(self, dim_in, dim_out):
+#         self.conv1 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+#         self.conv2 = nn.Conv2d(dim_out, dim_out, 3, 1, 1)
+#         if self.normalize.lower() == "in":
+#             self.norm1 = nn.InstanceNorm2d(dim_in, affine=True)
+#             self.norm2 = nn.InstanceNorm2d(dim_out, affine=True)
+#         elif self.normalize.lower() == "bn":
+#             self.norm1 = nn.BatchNorm2d(dim_in)
+#             self.norm2 = nn.BatchNorm2d(dim_out)
+#         if self.learned_sc:
+#             self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+#     def _shortcut(self, x):
+#         x = F.interpolate(x, scale_factor=2, mode='nearest')
+#         if self.learned_sc:
+#             x = self.conv1x1(x)
+#         return x
+
+#     def _residual(self, x):
+#         x = self.norm1(x)
+#         x = self.actv(x)
+#         x = F.interpolate(x, scale_factor=2, mode='nearest')
+#         x = self.conv1(x)
+#         x = self.norm2(x)
+#         x = self.actv(x)
+#         x = self.conv2(x)
+#         return x
+
+#     def forward(self, x):
+#         out = self._residual(x)
+#         out = (out + self._shortcut(x)) / self.equal_var
+#         return out
+
+class AdainResBlk(nn.Module):
+    def __init__(self, dim_in, dim_out, style_dim=512,
+                 actv=nn.LeakyReLU(0.2), upsample=False):
+        super().__init__()
+        self.actv = actv
+        self.upsample = upsample
+        self.learned_sc = dim_in != dim_out
+        self.equal_var = math.sqrt(2)
+        self._build_weights(dim_in, dim_out, style_dim)
+
+    def _build_weights(self, dim_in, dim_out, style_dim=64):
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(dim_in, dim_in, 3, 1, 1,groups=dim_in),
+            nn.Conv2d(dim_in, dim_out, 1, 1)
+        )
+
+        self.conv2 = nn.Sequential(
+            nn.Conv2d(dim_out, dim_out, 3, 1, 1,groups=dim_out),
+            nn.Conv2d(dim_out, dim_out, 1)
+        )
+
+        # self.conv1 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        # self.conv2 = nn.Conv2d(dim_out, dim_out, 3, 1, 1)
+
+        self.norm1 = AdaIN(style_dim, dim_in)
+        self.norm2 = AdaIN(style_dim, dim_out)
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        if self.upsample:
+            x = F.interpolate(x, scale_factor=2, mode='nearest')
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        return x
+
+    def _residual(self, x, s):
+        x = self.norm1(x, s)
+        x = self.actv(x)
+        if self.upsample:
+            x = F.interpolate(x, scale_factor=2, mode='nearest')
+        x = self.conv1(x)
+        x = self.norm2(x, s)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x, s):
+        out = self._residual(x, s)
+        out = (out + self._shortcut(x)) / self.equal_var
+        return out
+
+class ModulatedResBlk(nn.Module):
+    def __init__(self,
+                dim_in,
+                dim_out,
+                style_dim=512,
+                actv=nn.LeakyReLU(0.2),
+                upsample=False):
+        super().__init__()
+        self.actv = actv
+        self.upsample = upsample
+        self.learned_sc = dim_in != dim_out
+        self.equal_var = math.sqrt(2)
+        self._build_weights(dim_in, dim_out, style_dim)
+
+    def _build_weights(self, dim_in, dim_out, style_dim=64):
+        self.conv1 = ModulatedDWConv2d(dim_in, dim_out, style_dim, 3)
+        self.conv2 = ModulatedDWConv2d(dim_out, dim_out, style_dim, 3)
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        if self.upsample:
+            x = F.interpolate(x, scale_factor=2, mode='nearest')
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        return x
+
+    def _residual(self, x, s):
+        x = self.actv(x)
+        if self.upsample:
+            x = F.interpolate(x, scale_factor=2, mode='nearest')
+        x = self.conv1(x,s)
+        x = self.actv(x)
+        x = self.conv2(x,s)
+        return x
+
+    def forward(self, x, s):
+        out = self._residual(x, s)
+        out = (out + self._shortcut(x)) / self.equal_var
+        return out
+
+
+class Generator(nn.Module):
+    def __init__(
+                self,
+                **kwargs
+                ):
+        super().__init__()
+
+        id_dim      = kwargs["id_dim"]
+        k_size      = kwargs["g_kernel_size"]
+        res_num     = kwargs["res_num"]
+        in_channel  = kwargs["in_channel"]
+        up_mode     = kwargs["up_mode"]
+        norm        = kwargs["norm"].lower()
+        
+        aggregator  = kwargs["aggregator"]
+        res_mode    = kwargs["res_mode"]
+
+        padding_size= int((k_size -1)/2)
+        padding_type= 'reflect'
+
+        if norm.lower() == "in":
+            norm_out = nn.InstanceNorm2d(in_channel, affine=True)
+            norm_mask= nn.InstanceNorm2d
+        elif norm.lower() == "bn":
+            norm_out = nn.BatchNorm2d(in_channel)
+            norm_mask = nn.BatchNorm2d
+        
+        
+        activation = nn.LeakyReLU(0.2)
+        # activation = nn.ReLU()
+
+        self.from_rgb = nn.Conv2d(3, in_channel, 1, 1, 0)
+        # self.first_layer = nn.Sequential(nn.Conv2d(3, 64, kernel_size=3, padding=1, bias=False),
+        #                         nn.BatchNorm2d(64), activation)
+        ### downsample
+        self.down1 = ResBlk(in_channel, in_channel, normalize=norm, downsample=True)# 256
+
+        self.down2 = ResBlk(in_channel, in_channel*2, normalize=norm, downsample=True)# 128
+                                
+        self.down3 = ResBlk(in_channel*2, in_channel*4,normalize=norm, downsample=True)# 64
+
+        self.down4 = ResBlk(in_channel*4, in_channel*8, normalize=norm, downsample=True)# 32
+
+        self.down5 = ResBlk(in_channel*8, in_channel*8, normalize=norm, downsample=True)# 16
+
+        # self.down6 = ResBlk(in_channel*8, in_channel*8, normalize=True, downsample=True)# 8
+        
+        
+        ### resnet blocks
+        BN = []
+        for i in range(res_num):
+            BN += [
+                AdainResBlk(in_channel*8, in_channel*8, style_dim=id_dim, upsample=False)]
+        self.BottleNeck = nn.Sequential(*BN)
+
+        # self.up6 = AdainResBlk(in_channel*8, in_channel*8, style_dim=id_dim, upsample=True) # 16
+        
+        self.up5 = AdainResBlk(in_channel*8, in_channel*8, style_dim=id_dim, upsample=True) # 32
+
+        self.up4 = AdainResBlk(in_channel*8, in_channel*4, style_dim=id_dim, upsample=True) # 64
+
+        self.up3 = AdainResBlk(in_channel*4, in_channel*2, style_dim=id_dim, upsample=True) # 128
+
+        # self.maskhead = nn.Sequential(
+        #                         nn.Conv2d(in_channel*2, in_channel//2, kernel_size=3, stride=1, padding=1, bias=False),
+        #                         norm_mask, # 64
+        #                         activation,
+        #                         nn.Conv2d(in_channel//2, 1, kernel_size=3, stride=1, padding=1),
+        #                         nn.Sigmoid())
+        self.maskhead_lr = nn.Sequential(
+                                nn.UpsamplingNearest2d(scale_factor = 2),
+                                nn.Conv2d(in_channel*8, in_channel, kernel_size=3, stride=1, padding=1,bias=False),
+                                norm_mask(in_channel, affine=True), # 32
+                                activation,
+                                nn.UpsamplingNearest2d(scale_factor = 2),
+                                nn.Conv2d(in_channel, in_channel//4, kernel_size=3, stride=1, padding=1, bias=False),
+                                norm_mask(in_channel//4, affine=True), # 64
+                                activation,
+                                nn.UpsamplingNearest2d(scale_factor = 2),
+                                nn.Conv2d(in_channel//4, in_channel//8, kernel_size=3, stride=1, padding=1),
+                                norm_mask(in_channel//8, affine=True), # 128
+                                activation,
+        )
+
+        self.maskhead_hr = nn.Sequential(
+                                nn.UpsamplingNearest2d(scale_factor = 2),
+                                nn.Conv2d(in_channel//8, in_channel//16, kernel_size=3, stride=1, padding=1,bias=False),
+                                norm_mask(in_channel//16, affine=True), # 256
+                                activation,
+                                nn.UpsamplingNearest2d(scale_factor = 2),
+                                nn.Conv2d(in_channel//16, 1, kernel_size=3, stride=1, padding=1),
+                                nn.Sigmoid() # 512
+        )
+
+        self.maskhead_out = nn.Sequential(nn.Conv2d(in_channel//8, 1, kernel_size=1, stride=1),
+                                nn.Sigmoid())
+
+        # self.maskhead_lr = nn.Sequential(
+        #                         nn.UpsamplingNearest2d(scale_factor = 2),
+        #                         nn.Conv2d(in_channel*8, in_channel*8, 3, 1, 1,groups=in_channel*8),
+        #                         nn.Conv2d(in_channel*8, in_channel, 1, bias=False),
+        #                         # nn.Conv2d(in_channel*8, in_channel, kernel_size=3, stride=1, padding=1,bias=False),
+        #                         norm_mask(in_channel, affine=True), # 32
+        #                         activation,
+        #                         nn.UpsamplingNearest2d(scale_factor = 2),
+        #                         nn.Conv2d(in_channel, in_channel, 3, 1, 1,groups=in_channel),
+        #                         nn.Conv2d(in_channel, in_channel//4, 1, bias=False),
+        #                         # nn.Conv2d(in_channel, in_channel//4, kernel_size=3, stride=1, padding=1, bias=False),
+        #                         norm_mask(in_channel//4, affine=True), # 64
+        #                         activation,
+        #                         nn.UpsamplingNearest2d(scale_factor = 2),
+        #                         # nn.Conv2d(in_channel//4, in_channel//8, kernel_size=3, stride=1, padding=1),
+        #                         nn.Conv2d(in_channel//4, in_channel//4, 3, 1, 1,groups=in_channel//4),
+        #                         nn.Conv2d(in_channel//4, in_channel//8, 1, bias=False),
+        #                         norm_mask(in_channel//8, affine=True), # 128
+        #                         activation,
+        # )
+        # self.maskhead_hr = nn.Sequential(
+        #                         nn.UpsamplingNearest2d(scale_factor = 2),
+        #                         # nn.Conv2d(in_channel//8, in_channel//16, kernel_size=3, stride=1, padding=1,bias=False),
+        #                         nn.Conv2d(in_channel//8, in_channel//8, 3, 1, 1,groups=in_channel//8),
+        #                         nn.Conv2d(in_channel//8, in_channel//16, 1, bias=False),
+        #                         norm_mask(in_channel//16, affine=True), # 256
+        #                         activation,
+        #                         nn.UpsamplingNearest2d(scale_factor = 2),
+        #                         nn.Conv2d(in_channel//16, 1, kernel_size=3, stride=1, padding=1),
+        #                         nn.Sigmoid() # 512
+        # )
+        # self.maskhead_out = nn.Sequential(nn.Conv2d(in_channel//8, 1, kernel_size=1, stride=1),
+        #                         nn.Sigmoid())
+        
+        # self.up2 = AdainResBlk(in_channel*2, in_channel, style_dim=id_dim, upsample=True)
+        # self.up2 = AdainResBlk(in_channel*2, in_channel, style_dim=id_dim, upsample=True)
+        self.up2 = AdainResBlk(in_channel*2, in_channel, style_dim=id_dim, upsample=True)
+        # ResUpBlk(in_channel*2, in_channel, normalize=norm)
+
+        # self.up1 = AdainResBlk(in_channel, in_channel, style_dim=id_dim, upsample=True)
+        self.up1 = AdainResBlk(in_channel, in_channel, style_dim=id_dim, upsample=True)
+        # ResUpBlk(in_channel, in_channel, normalize=norm)
+        # ResUpBlk(in_channel, in_channel, normalize="in") # 512
+
+        
+        
+        self.to_rgb = nn.Sequential(
+            norm_out,
+            activation,
+            nn.Conv2d(in_channel, 3, 3, 1, 1))
+
+        # self.last_layer = nn.Sequential(nn.ReflectionPad2d(3),
+        #             nn.Conv2d(64, 3, kernel_size=7, padding=0))
+
+
+    #     self.__weights_init__()
+
+    # def __weights_init__(self):
+    #     for layer in self.encoder:
+    #         if isinstance(layer,nn.Conv2d):
+    #             nn.init.xavier_uniform_(layer.weight)
+
+    #     for layer in self.encoder2:
+    #         if isinstance(layer,nn.Conv2d):
+    #             nn.init.xavier_uniform_(layer.weight)
+
+    def forward(self, img, id):
+        res = self.from_rgb(img)
+        res = self.down1(res)
+        skip = self.down2(res)
+        res = self.down3(skip)
+        res = self.down4(res)
+        res = self.down5(res)
+        mask_feat= self.maskhead_lr(res)
+        
+        for i in range(len(self.BottleNeck)):
+            res = self.BottleNeck[i](res, id)
+        res = self.up5(res,id)
+        res = self.up4(res,id)
+        res = self.up3(res,id)
+        mask_lr= self.maskhead_out(mask_feat)
+        # res = (1-mask) * self.sigma(skip) + mask * res
+        res = (1-mask_lr) * skip + mask_lr * res
+        res = self.up2(res,id) #  + skip
+        res = self.up1(res,id)
+        res = self.to_rgb(res)
+        mask_hr=self.maskhead_hr(mask_feat)
+        res = (1-mask_hr) * img + mask_hr * res
+        return res, mask_lr, mask_hr

components/Generator_2maskhead_config copy.py

@@ -0,0 +1,293 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+#############################################################
+# File: Generator_Invobn_config1.py
+# Created Date: Saturday February 26th 2022
+# Author: Chen Xuanhong
+# Email: chenxuanhongzju@outlook.com
+# Last Modified:  Wednesday, 13th April 2022 10:22:52 am
+# Modified By: Chen Xuanhong
+# Copyright (c) 2022 Shanghai Jiao Tong University
+#############################################################
+
+import os
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+import math
+
+
+
+class ResBlk(nn.Module):
+    def __init__(self, dim_in, dim_out, actv=nn.LeakyReLU(0.2),
+                 normalize="bn", downsample=False):
+        super().__init__()
+        self.actv = actv
+        self.normalize = normalize
+        self.downsample = downsample
+        self.learned_sc = dim_in != dim_out
+        self.equal_var  = math.sqrt(2)
+        self._build_weights(dim_in, dim_out)
+
+    def _build_weights(self, dim_in, dim_out):
+        self.conv1 = nn.Conv2d(dim_in, dim_in, 3, 1, 1)
+        self.conv2 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        if self.normalize.lower() == "in":
+            self.norm1 = nn.InstanceNorm2d(dim_in, affine=True)
+            self.norm2 = nn.InstanceNorm2d(dim_in, affine=True)
+        elif self.normalize.lower() == "bn":
+            self.norm1 = nn.BatchNorm2d(dim_in)
+            self.norm2 = nn.BatchNorm2d(dim_in)
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        if self.downsample:
+            x = F.avg_pool2d(x, 2)
+        return x
+
+    def _residual(self, x):
+        if self.normalize:
+            x = self.norm1(x)
+        x = self.actv(x)
+        x = self.conv1(x)
+        if self.downsample:
+            x = F.avg_pool2d(x, 2)
+        if self.normalize:
+            x = self.norm2(x)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x):
+        x = self._shortcut(x) + self._residual(x)
+        return x /self.equal_var   # unit variance
+
+class AdaIN(nn.Module):
+    def __init__(self, style_dim, num_features):
+        super().__init__()
+        self.norm = nn.InstanceNorm2d(num_features, affine=False)
+        self.fc = nn.Linear(style_dim, num_features*2)
+
+    def forward(self, x, s):
+        h = self.fc(s)
+        h = h.view(h.size(0), h.size(1), 1, 1)
+        gamma, beta = torch.chunk(h, chunks=2, dim=1)
+        return (1 + gamma) * self.norm(x) + beta
+
+class ResUpBlk(nn.Module):
+    def __init__(self, dim_in, dim_out,actv=nn.LeakyReLU(0.2),normalize="in"):
+        super().__init__()
+        self.actv = actv
+        self.normalize = normalize
+        self.learned_sc = dim_in != dim_out
+        self.equal_var = math.sqrt(2)
+        self._build_weights(dim_in, dim_out)
+
+    def _build_weights(self, dim_in, dim_out):
+        self.conv1 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        self.conv2 = nn.Conv2d(dim_out, dim_out, 3, 1, 1)
+        if self.normalize.lower() == "in":
+            self.norm1 = nn.InstanceNorm2d(dim_in, affine=True)
+            self.norm2 = nn.InstanceNorm2d(dim_out, affine=True)
+        elif self.normalize.lower() == "bn":
+            self.norm1 = nn.BatchNorm2d(dim_in)
+            self.norm2 = nn.BatchNorm2d(dim_out)
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        x = F.interpolate(x, scale_factor=2, mode='nearest')
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        return x
+
+    def _residual(self, x):
+        x = self.norm1(x)
+        x = self.actv(x)
+        x = F.interpolate(x, scale_factor=2, mode='nearest')
+        x = self.conv1(x)
+        x = self.norm2(x)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x):
+        out = self._residual(x)
+        out = (out + self._shortcut(x)) / self.equal_var
+        return out
+
+class AdainResBlk(nn.Module):
+    def __init__(self, dim_in, dim_out, style_dim=512,
+                 actv=nn.LeakyReLU(0.2), upsample=False):
+        super().__init__()
+        self.actv = actv
+        self.upsample = upsample
+        self.learned_sc = dim_in != dim_out
+        self.equal_var = math.sqrt(2)
+        self._build_weights(dim_in, dim_out, style_dim)
+
+    def _build_weights(self, dim_in, dim_out, style_dim=64):
+        self.conv1 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        self.conv2 = nn.Conv2d(dim_out, dim_out, 3, 1, 1)
+        self.norm1 = AdaIN(style_dim, dim_in)
+        self.norm2 = AdaIN(style_dim, dim_out)
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        if self.upsample:
+            x = F.interpolate(x, scale_factor=2, mode='nearest')
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        return x
+
+    def _residual(self, x, s):
+        x = self.norm1(x, s)
+        x = self.actv(x)
+        if self.upsample:
+            x = F.interpolate(x, scale_factor=2, mode='nearest')
+        x = self.conv1(x)
+        x = self.norm2(x, s)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x, s):
+        out = self._residual(x, s)
+        out = (out + self._shortcut(x)) / self.equal_var
+        return out
+
+
+class Generator(nn.Module):
+    def __init__(
+                self,
+                **kwargs
+                ):
+        super().__init__()
+
+        id_dim      = kwargs["id_dim"]
+        k_size      = kwargs["g_kernel_size"]
+        res_num     = kwargs["res_num"]
+        in_channel  = kwargs["in_channel"]
+        up_mode     = kwargs["up_mode"]
+        norm        = kwargs["norm"]
+        norm        = norm.lower()
+        
+        aggregator  = kwargs["aggregator"]
+        res_mode    = kwargs["res_mode"]
+
+        padding_size= int((k_size -1)/2)
+        padding_type= 'reflect'
+        
+        
+        activation = nn.LeakyReLU(0.2)
+        # activation = nn.ReLU()
+
+        self.from_rgb = nn.Conv2d(3, in_channel, 1, 1, 0)
+        # self.first_layer = nn.Sequential(nn.Conv2d(3, 64, kernel_size=3, padding=1, bias=False),
+        #                         nn.BatchNorm2d(64), activation)
+        ### downsample
+        self.down1 = ResBlk(in_channel, in_channel, normalize=norm, downsample=True)# 256
+
+        self.down2 = ResBlk(in_channel, in_channel*2, normalize=norm, downsample=True)# 128
+
+        # self.sigma = ResBlk(in_channel*2,in_channel*2)
+                                
+        self.down3 = ResBlk(in_channel*2, in_channel*4,normalize=norm, downsample=True)# 64
+
+        self.down4 = ResBlk(in_channel*4, in_channel*8, normalize=norm, downsample=True)# 32
+
+        self.down5 = ResBlk(in_channel*8, in_channel*8, normalize=norm, downsample=True)# 16
+
+        # self.down6 = ResBlk(in_channel*8, in_channel*8, normalize=True, downsample=True)# 8
+        
+        
+        ### resnet blocks
+        BN = []
+        for i in range(res_num):
+            BN += [
+                AdainResBlk(in_channel*8, in_channel*8, style_dim=id_dim, upsample=False)]
+        self.BottleNeck = nn.Sequential(*BN)
+
+        # self.up6 = AdainResBlk(in_channel*8, in_channel*8, style_dim=id_dim, upsample=True) # 16
+        
+        self.up5 = AdainResBlk(in_channel*8, in_channel*8, style_dim=id_dim, upsample=True) # 32
+
+        self.up4 = AdainResBlk(in_channel*8, in_channel*4, style_dim=id_dim, upsample=True) # 64
+
+        self.up3 = AdainResBlk(in_channel*4, in_channel*2, style_dim=id_dim, upsample=True) # 128
+
+        self.maskhead_lr = nn.Sequential(
+                                nn.Conv2d(in_channel*2, in_channel, kernel_size=3, stride=1, padding=1, bias=False),
+                                nn.BatchNorm2d(in_channel), # 64
+                                activation,
+                                nn.Conv2d(in_channel, 1, kernel_size=3, stride=1, padding=1),
+                                nn.Sigmoid()
+        )
+        
+        # self.up2 = AdainResBlk(in_channel*2, in_channel, style_dim=id_dim, upsample=True)
+        # self.up2 = AdainResBlk(in_channel*2, in_channel, style_dim=id_dim, upsample=True)
+        self.up2 = ResUpBlk(in_channel*2, in_channel, normalize=norm)
+
+        # self.up1 = AdainResBlk(in_channel, in_channel, style_dim=id_dim, upsample=True)
+        self.up1 = ResUpBlk(in_channel, in_channel, normalize=norm)
+        # ResUpBlk(in_channel, in_channel, normalize="in") # 512
+        self.maskhead_hr = nn.Sequential(
+                                nn.Conv2d(in_channel, in_channel//8, kernel_size=3, stride=1, padding=1, bias=False),
+                                nn.BatchNorm2d(in_channel//8), # 64
+                                activation,
+                                nn.Conv2d(in_channel//8, 1, kernel_size=3, stride=1, padding=1),
+                                nn.Sigmoid()
+        )
+
+        if norm.lower() == "in":
+            norm_out = nn.InstanceNorm2d(in_channel, affine=True)
+        elif norm.lower() == "bn":
+            norm_out = nn.BatchNorm2d(in_channel)
+        
+        self.to_rgb = nn.Sequential(
+            norm_out,
+            activation,
+            nn.Conv2d(in_channel, 3, 3, 1, 1))
+
+        # self.last_layer = nn.Sequential(nn.ReflectionPad2d(3),
+        #             nn.Conv2d(64, 3, kernel_size=7, padding=0))
+
+
+    #     self.__weights_init__()
+
+    # def __weights_init__(self):
+    #     for layer in self.encoder:
+    #         if isinstance(layer,nn.Conv2d):
+    #             nn.init.xavier_uniform_(layer.weight)
+
+    #     for layer in self.encoder2:
+    #         if isinstance(layer,nn.Conv2d):
+    #             nn.init.xavier_uniform_(layer.weight)
+
+    def forward(self, img, id):
+        res = self.from_rgb(img)
+        res = self.down1(res)
+        skip = self.down2(res)
+        res = self.down3(skip)
+        res = self.down4(res)
+        res = self.down5(res)
+        
+        for i in range(len(self.BottleNeck)):
+            res = self.BottleNeck[i](res, id)
+        res = self.up5(res,id)
+        res = self.up4(res,id)
+        res = self.up3(res,id)
+        mask= self.maskhead_lr(res)
+        # res = (1-mask) * self.sigma(skip) + mask * res
+        res = (1-mask) * skip + mask * res
+        res = self.up2(res) #  + skip
+        res = self.up1(res)
+        mask_hr = self.maskhead_hr(res)
+        res = self.to_rgb(res)
+        res = (1-mask_hr)*img + mask_hr*res
+        return res, mask, mask_hr

components/Generator_2maskhead_config.py

@@ -0,0 +1,293 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+#############################################################
+# File: Generator_Invobn_config1.py
+# Created Date: Saturday February 26th 2022
+# Author: Chen Xuanhong
+# Email: chenxuanhongzju@outlook.com
+# Last Modified:  Wednesday, 13th April 2022 10:22:52 am
+# Modified By: Chen Xuanhong
+# Copyright (c) 2022 Shanghai Jiao Tong University
+#############################################################
+
+import os
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+import math
+
+
+
+class ResBlk(nn.Module):
+    def __init__(self, dim_in, dim_out, actv=nn.LeakyReLU(0.2),
+                 normalize="bn", downsample=False):
+        super().__init__()
+        self.actv = actv
+        self.normalize = normalize
+        self.downsample = downsample
+        self.learned_sc = dim_in != dim_out
+        self.equal_var  = math.sqrt(2)
+        self._build_weights(dim_in, dim_out)
+
+    def _build_weights(self, dim_in, dim_out):
+        self.conv1 = nn.Conv2d(dim_in, dim_in, 3, 1, 1)
+        self.conv2 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        if self.normalize.lower() == "in":
+            self.norm1 = nn.InstanceNorm2d(dim_in, affine=True)
+            self.norm2 = nn.InstanceNorm2d(dim_in, affine=True)
+        elif self.normalize.lower() == "bn":
+            self.norm1 = nn.BatchNorm2d(dim_in)
+            self.norm2 = nn.BatchNorm2d(dim_in)
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        if self.downsample:
+            x = F.avg_pool2d(x, 2)
+        return x
+
+    def _residual(self, x):
+        if self.normalize:
+            x = self.norm1(x)
+        x = self.actv(x)
+        x = self.conv1(x)
+        if self.downsample:
+            x = F.avg_pool2d(x, 2)
+        if self.normalize:
+            x = self.norm2(x)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x):
+        x = self._shortcut(x) + self._residual(x)
+        return x /self.equal_var   # unit variance
+
+class AdaIN(nn.Module):
+    def __init__(self, style_dim, num_features):
+        super().__init__()
+        self.norm = nn.InstanceNorm2d(num_features, affine=False)
+        self.fc = nn.Linear(style_dim, num_features*2)
+
+    def forward(self, x, s):
+        h = self.fc(s)
+        h = h.view(h.size(0), h.size(1), 1, 1)
+        gamma, beta = torch.chunk(h, chunks=2, dim=1)
+        return (1 + gamma) * self.norm(x) + beta
+
+class ResUpBlk(nn.Module):
+    def __init__(self, dim_in, dim_out,actv=nn.LeakyReLU(0.2),normalize="in"):
+        super().__init__()
+        self.actv = actv
+        self.normalize = normalize
+        self.learned_sc = dim_in != dim_out
+        self.equal_var = math.sqrt(2)
+        self._build_weights(dim_in, dim_out)
+
+    def _build_weights(self, dim_in, dim_out):
+        self.conv1 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        self.conv2 = nn.Conv2d(dim_out, dim_out, 3, 1, 1)
+        if self.normalize.lower() == "in":
+            self.norm1 = nn.InstanceNorm2d(dim_in, affine=True)
+            self.norm2 = nn.InstanceNorm2d(dim_out, affine=True)
+        elif self.normalize.lower() == "bn":
+            self.norm1 = nn.BatchNorm2d(dim_in)
+            self.norm2 = nn.BatchNorm2d(dim_out)
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        x = F.interpolate(x, scale_factor=2, mode='nearest')
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        return x
+
+    def _residual(self, x):
+        x = self.norm1(x)
+        x = self.actv(x)
+        x = F.interpolate(x, scale_factor=2, mode='nearest')
+        x = self.conv1(x)
+        x = self.norm2(x)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x):
+        out = self._residual(x)
+        out = (out + self._shortcut(x)) / self.equal_var
+        return out
+
+class AdainResBlk(nn.Module):
+    def __init__(self, dim_in, dim_out, style_dim=512,
+                 actv=nn.LeakyReLU(0.2), upsample=False):
+        super().__init__()
+        self.actv = actv
+        self.upsample = upsample
+        self.learned_sc = dim_in != dim_out
+        self.equal_var = math.sqrt(2)
+        self._build_weights(dim_in, dim_out, style_dim)
+
+    def _build_weights(self, dim_in, dim_out, style_dim=64):
+        self.conv1 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        self.conv2 = nn.Conv2d(dim_out, dim_out, 3, 1, 1)
+        self.norm1 = AdaIN(style_dim, dim_in)
+        self.norm2 = AdaIN(style_dim, dim_out)
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        if self.upsample:
+            x = F.interpolate(x, scale_factor=2, mode='nearest')
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        return x
+
+    def _residual(self, x, s):
+        x = self.norm1(x, s)
+        x = self.actv(x)
+        if self.upsample:
+            x = F.interpolate(x, scale_factor=2, mode='nearest')
+        x = self.conv1(x)
+        x = self.norm2(x, s)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x, s):
+        out = self._residual(x, s)
+        out = (out + self._shortcut(x)) / self.equal_var
+        return out
+
+
+class Generator(nn.Module):
+    def __init__(
+                self,
+                **kwargs
+                ):
+        super().__init__()
+
+        id_dim      = kwargs["id_dim"]
+        k_size      = kwargs["g_kernel_size"]
+        res_num     = kwargs["res_num"]
+        in_channel  = kwargs["in_channel"]
+        up_mode     = kwargs["up_mode"]
+        norm        = kwargs["norm"]
+        norm        = norm.lower()
+        
+        aggregator  = kwargs["aggregator"]
+        res_mode    = kwargs["res_mode"]
+
+        padding_size= int((k_size -1)/2)
+        padding_type= 'reflect'
+        
+        
+        activation = nn.LeakyReLU(0.2)
+        # activation = nn.ReLU()
+
+        self.from_rgb = nn.Conv2d(3, in_channel, 1, 1, 0)
+        # self.first_layer = nn.Sequential(nn.Conv2d(3, 64, kernel_size=3, padding=1, bias=False),
+        #                         nn.BatchNorm2d(64), activation)
+        ### downsample
+        self.down1 = ResBlk(in_channel, in_channel, normalize=norm, downsample=True)# 256
+
+        self.down2 = ResBlk(in_channel, in_channel*2, normalize=norm, downsample=True)# 128
+
+        # self.sigma = ResBlk(in_channel*2,in_channel*2)
+                                
+        self.down3 = ResBlk(in_channel*2, in_channel*4,normalize=norm, downsample=True)# 64
+
+        self.down4 = ResBlk(in_channel*4, in_channel*8, normalize=norm, downsample=True)# 32
+
+        self.down5 = ResBlk(in_channel*8, in_channel*8, normalize=norm, downsample=True)# 16
+
+        # self.down6 = ResBlk(in_channel*8, in_channel*8, normalize=True, downsample=True)# 8
+        
+        
+        ### resnet blocks
+        BN = []
+        for i in range(res_num):
+            BN += [
+                AdainResBlk(in_channel*8, in_channel*8, style_dim=id_dim, upsample=False)]
+        self.BottleNeck = nn.Sequential(*BN)
+
+        # self.up6 = AdainResBlk(in_channel*8, in_channel*8, style_dim=id_dim, upsample=True) # 16
+        
+        self.up5 = AdainResBlk(in_channel*8, in_channel*8, style_dim=id_dim, upsample=True) # 32
+
+        self.up4 = AdainResBlk(in_channel*8, in_channel*4, style_dim=id_dim, upsample=True) # 64
+
+        self.up3 = AdainResBlk(in_channel*4, in_channel*2, style_dim=id_dim, upsample=True) # 128
+
+        self.maskhead_lr = nn.Sequential(
+                                nn.Conv2d(in_channel*2, in_channel, kernel_size=3, stride=1, padding=1, bias=False),
+                                nn.BatchNorm2d(in_channel), # 64
+                                activation,
+                                nn.Conv2d(in_channel, 1, kernel_size=3, stride=1, padding=1),
+                                nn.Sigmoid()
+        )
+        
+        # self.up2 = AdainResBlk(in_channel*2, in_channel, style_dim=id_dim, upsample=True)
+        # self.up2 = AdainResBlk(in_channel*2, in_channel, style_dim=id_dim, upsample=True)
+        self.up2 = ResUpBlk(in_channel*2, in_channel, normalize=norm)
+
+        # self.up1 = AdainResBlk(in_channel, in_channel, style_dim=id_dim, upsample=True)
+        self.up1 = ResUpBlk(in_channel, in_channel, normalize=norm)
+        # ResUpBlk(in_channel, in_channel, normalize="in") # 512
+        self.maskhead_hr = nn.Sequential(
+                                nn.Conv2d(in_channel, in_channel//8, kernel_size=3, stride=1, padding=1, bias=False),
+                                nn.BatchNorm2d(in_channel//8), # 64
+                                activation,
+                                nn.Conv2d(in_channel//8, 1, kernel_size=3, stride=1, padding=1),
+                                nn.Sigmoid()
+        )
+
+        if norm.lower() == "in":
+            norm_out = nn.InstanceNorm2d(in_channel, affine=True)
+        elif norm.lower() == "bn":
+            norm_out = nn.BatchNorm2d(in_channel)
+        
+        self.to_rgb = nn.Sequential(
+            norm_out,
+            activation,
+            nn.Conv2d(in_channel, 3, 3, 1, 1))
+
+        # self.last_layer = nn.Sequential(nn.ReflectionPad2d(3),
+        #             nn.Conv2d(64, 3, kernel_size=7, padding=0))
+
+
+    #     self.__weights_init__()
+
+    # def __weights_init__(self):
+    #     for layer in self.encoder:
+    #         if isinstance(layer,nn.Conv2d):
+    #             nn.init.xavier_uniform_(layer.weight)
+
+    #     for layer in self.encoder2:
+    #         if isinstance(layer,nn.Conv2d):
+    #             nn.init.xavier_uniform_(layer.weight)
+
+    def forward(self, img, id):
+        res = self.from_rgb(img)
+        res = self.down1(res)
+        skip = self.down2(res)
+        res = self.down3(skip)
+        res = self.down4(res)
+        res = self.down5(res)
+        
+        for i in range(len(self.BottleNeck)):
+            res = self.BottleNeck[i](res, id)
+        res = self.up5(res,id)
+        res = self.up4(res,id)
+        res = self.up3(res,id)
+        mask= self.maskhead_lr(res)
+        # res = (1-mask) * self.sigma(skip) + mask * res
+        res = (1-mask) * skip + mask * res
+        res = self.up2(res) #  + skip
+        res = self.up1(res)
+        mask_hr = self.maskhead_hr(res)
+        res = self.to_rgb(res)
+        res = (1-mask_hr)*img + mask_hr*res
+        return res, mask, mask_hr

components/Generator_ResSkip_config.py

@@ -0,0 +1,276 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+#############################################################
+# File: Generator_Invobn_config1.py
+# Created Date: Saturday February 26th 2022
+# Author: Chen Xuanhong
+# Email: chenxuanhongzju@outlook.com
+# Last Modified:  Tuesday, 29th March 2022 12:02:53 pm
+# Modified By: Chen Xuanhong
+# Copyright (c) 2022 Shanghai Jiao Tong University
+#############################################################
+
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+import math
+from components.LSTU import LSTU
+
+
+class ResBlk(nn.Module):
+    def __init__(self, dim_in, dim_out, actv=nn.LeakyReLU(0.2),
+                 normalize="in", downsample=False):
+        super().__init__()
+        self.actv = actv
+        self.normalize = normalize
+        self.downsample = downsample
+        self.learned_sc = dim_in != dim_out
+        self.equal_var  = math.sqrt(2)
+        self._build_weights(dim_in, dim_out)
+
+    def _build_weights(self, dim_in, dim_out):
+        self.conv1 = nn.Conv2d(dim_in, dim_in, 3, 1, 1)
+        self.conv2 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        if self.normalize.lower() == "in":
+            self.norm1 = nn.InstanceNorm2d(dim_in, affine=True)
+            self.norm2 = nn.InstanceNorm2d(dim_in, affine=True)
+        elif self.normalize.lower() == "bn":
+            self.norm1 = nn.BatchNorm2d(dim_in)
+            self.norm2 = nn.BatchNorm2d(dim_in)
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        if self.downsample:
+            x = F.avg_pool2d(x, 2)
+        return x
+
+    def _residual(self, x):
+        if self.normalize:
+            x = self.norm1(x)
+        x = self.actv(x)
+        x = self.conv1(x)
+        if self.downsample:
+            x = F.avg_pool2d(x, 2)
+        if self.normalize:
+            x = self.norm2(x)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x):
+        x = self._shortcut(x) + self._residual(x)
+        return x /self.equal_var   # unit variance
+
+class AdaIN(nn.Module):
+    def __init__(self, style_dim, num_features):
+        super().__init__()
+        self.norm = nn.InstanceNorm2d(num_features, affine=False)
+        self.fc = nn.Linear(style_dim, num_features*2)
+
+    def forward(self, x, s):
+        h = self.fc(s)
+        h = h.view(h.size(0), h.size(1), 1, 1)
+        gamma, beta = torch.chunk(h, chunks=2, dim=1)
+        return (1 + gamma) * self.norm(x) + beta
+
+class ResUpBlk(nn.Module):
+    def __init__(self, dim_in, dim_out,actv=nn.LeakyReLU(0.2),normalize="in"):
+        super().__init__()
+        self.actv = actv
+        self.normalize = normalize
+        self.learned_sc = dim_in != dim_out
+        self.equal_var = math.sqrt(2)
+        self._build_weights(dim_in, dim_out)
+
+    def _build_weights(self, dim_in, dim_out):
+        self.conv1 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        self.conv2 = nn.Conv2d(dim_out, dim_out, 3, 1, 1)
+        if self.normalize.lower() == "in":
+            self.norm1 = nn.InstanceNorm2d(dim_in, affine=True)
+            self.norm2 = nn.InstanceNorm2d(dim_out, affine=True)
+        elif self.normalize.lower() == "bn":
+            self.norm1 = nn.BatchNorm2d(dim_in)
+            self.norm2 = nn.BatchNorm2d(dim_out)
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        x = F.interpolate(x, scale_factor=2, mode='nearest')
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        return x
+
+    def _residual(self, x):
+        x = self.norm1(x)
+        x = self.actv(x)
+        x = F.interpolate(x, scale_factor=2, mode='nearest')
+        x = self.conv1(x)
+        x = self.norm2(x)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x):
+        out = self._residual(x)
+        out = (out + self._shortcut(x)) / self.equal_var
+        return out
+
+class AdainResBlk(nn.Module):
+    def __init__(self, dim_in, dim_out, style_dim=512,
+                 actv=nn.LeakyReLU(0.2), upsample=False):
+        super().__init__()
+        self.actv = actv
+        self.upsample = upsample
+        self.learned_sc = dim_in != dim_out
+        self.equal_var = math.sqrt(2)
+        self._build_weights(dim_in, dim_out, style_dim)
+
+    def _build_weights(self, dim_in, dim_out, style_dim=64):
+        self.conv1 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        self.conv2 = nn.Conv2d(dim_out, dim_out, 3, 1, 1)
+        self.norm1 = AdaIN(style_dim, dim_in)
+        self.norm2 = AdaIN(style_dim, dim_out)
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        if self.upsample:
+            x = F.interpolate(x, scale_factor=2, mode='nearest')
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        return x
+
+    def _residual(self, x, s):
+        x = self.norm1(x, s)
+        x = self.actv(x)
+        if self.upsample:
+            x = F.interpolate(x, scale_factor=2, mode='nearest')
+        x = self.conv1(x)
+        x = self.norm2(x, s)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x, s):
+        out = self._residual(x, s)
+        out = (out + self._shortcut(x)) / self.equal_var
+        return out
+
+
+class Generator(nn.Module):
+    def __init__(
+                self,
+                **kwargs
+                ):
+        super().__init__()
+
+        id_dim      = kwargs["id_dim"]
+        k_size      = kwargs["g_kernel_size"]
+        res_num     = kwargs["res_num"]
+        in_channel  = kwargs["in_channel"]
+        up_mode     = kwargs["up_mode"]
+        norm        = kwargs["norm"]
+        
+        aggregator  = kwargs["aggregator"]
+        res_mode    = kwargs["res_mode"]
+
+        padding_size= int((k_size -1)/2)
+        padding_type= 'reflect'
+        
+        activation = nn.LeakyReLU(0.2)
+        # activation = nn.ReLU()
+
+        self.from_rgb = nn.Conv2d(3, in_channel, 1, 1, 0)
+        # self.first_layer = nn.Sequential(nn.Conv2d(3, 64, kernel_size=3, padding=1, bias=False),
+        #                         nn.BatchNorm2d(64), activation)
+        ### downsample
+        self.down1 = ResBlk(in_channel, in_channel, normalize=norm, downsample=True)# 256
+
+        self.down2 = ResBlk(in_channel, in_channel*2, normalize=norm, downsample=True)# 128
+                                
+        self.down3 = ResBlk(in_channel*2, in_channel*4,normalize=norm, downsample=True)# 64
+
+        self.down4 = ResBlk(in_channel*4, in_channel*8, normalize=norm, downsample=True)# 32
+
+        self.down5 = ResBlk(in_channel*8, in_channel*8, normalize=norm, downsample=True)# 16
+
+        # self.down6 = ResBlk(in_channel*8, in_channel*8, normalize=True, downsample=True)# 8
+        
+        ### resnet blocks
+        BN = []
+        for i in range(res_num):
+            BN += [
+                AdainResBlk(in_channel*8, in_channel*8, style_dim=id_dim, upsample=False)]
+        self.BottleNeck = nn.Sequential(*BN)
+
+        # self.up6 = AdainResBlk(in_channel*8, in_channel*8, style_dim=id_dim, upsample=True) # 16
+        
+        self.up5 = AdainResBlk(in_channel*8, in_channel*8, style_dim=id_dim, upsample=True) # 32
+
+        self.up4 = AdainResBlk(in_channel*8, in_channel*4, style_dim=id_dim, upsample=True) # 64
+
+        self.up3 = AdainResBlk(in_channel*4, in_channel*2, style_dim=id_dim, upsample=True) # 128
+
+        self.lstu = LSTU(in_channel*2,norm)
+        
+        self.up2 = AdainResBlk(in_channel*2, in_channel, style_dim=id_dim, upsample=True)
+        # ResUpBlk(in_channel*2, in_channel, normalize="in") # 256
+
+        # self.lstu  = nn.Sequential(nn.Conv2d(in_channel*2, in_channel, kernel_size=3, stride=1, padding=1, bias=False),
+        #                         nn.BatchNorm2d(in_channel),
+        #                         activation,
+        #                         nn.Conv2d(in_channel, 1, kernel_size=3, stride=1, padding=1),
+        #                         nn.Sigmoid()
+        #                         )
+
+        self.up1 = AdainResBlk(in_channel, in_channel, style_dim=id_dim, upsample=True)
+        # ResUpBlk(in_channel, in_channel, normalize="in") # 512
+
+        if norm.lower() == "in":
+            norm_out = nn.InstanceNorm2d(in_channel, affine=True)
+        elif norm.lower() == "bn":
+            norm_out = nn.BatchNorm2d(in_channel)
+        
+        self.to_rgb = nn.Sequential(
+            norm_out,
+            activation,
+            nn.Conv2d(in_channel, 3, 3, 1, 1))
+
+        # self.last_layer = nn.Sequential(nn.ReflectionPad2d(3),
+        #             nn.Conv2d(64, 3, kernel_size=7, padding=0))
+
+
+    #     self.__weights_init__()
+
+    # def __weights_init__(self):
+    #     for layer in self.encoder:
+    #         if isinstance(layer,nn.Conv2d):
+    #             nn.init.xavier_uniform_(layer.weight)
+
+    #     for layer in self.encoder2:
+    #         if isinstance(layer,nn.Conv2d):
+    #             nn.init.xavier_uniform_(layer.weight)
+
+    def forward(self, img, id):
+        res = self.from_rgb(img)
+        res = self.down1(res)
+        skip = self.down2(res)
+        res = self.down3(skip)
+        res = self.down4(res)
+        res = self.down5(res)
+        # res = self.down6(res)
+        for i in range(len(self.BottleNeck)):
+            res = self.BottleNeck[i](res, id)
+        # res = self.up6(res,id)
+        res = self.up5(res,id)
+        res = self.up4(res,id)
+        res = self.up3(res,id)
+        res,mask = self.lstu(skip, res)
+        res = self.up2(res,id) #  + skip
+        res = self.up1(res,id)
+        res = self.to_rgb(res)
+        return res,mask

components/Generator_ResSkip_config1.py

@@ -0,0 +1,285 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+#############################################################
+# File: Generator_Invobn_config1.py
+# Created Date: Saturday February 26th 2022
+# Author: Chen Xuanhong
+# Email: chenxuanhongzju@outlook.com
+# Last Modified:  Tuesday, 29th March 2022 1:08:05 pm
+# Modified By: Chen Xuanhong
+# Copyright (c) 2022 Shanghai Jiao Tong University
+#############################################################
+
+import os
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+import math
+
+
+
+class ResBlk(nn.Module):
+    def __init__(self, dim_in, dim_out, actv=nn.LeakyReLU(0.2),
+                 normalize="in", downsample=False):
+        super().__init__()
+        self.actv = actv
+        self.normalize = normalize
+        self.downsample = downsample
+        self.learned_sc = dim_in != dim_out
+        self.equal_var  = math.sqrt(2)
+        self._build_weights(dim_in, dim_out)
+
+    def _build_weights(self, dim_in, dim_out):
+        self.conv1 = nn.Conv2d(dim_in, dim_in, 3, 1, 1)
+        self.conv2 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        if self.normalize.lower() == "in":
+            self.norm1 = nn.InstanceNorm2d(dim_in, affine=True)
+            self.norm2 = nn.InstanceNorm2d(dim_in, affine=True)
+        elif self.normalize.lower() == "bn":
+            self.norm1 = nn.BatchNorm2d(dim_in)
+            self.norm2 = nn.BatchNorm2d(dim_in)
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        if self.downsample:
+            x = F.avg_pool2d(x, 2)
+        return x
+
+    def _residual(self, x):
+        if self.normalize:
+            x = self.norm1(x)
+        x = self.actv(x)
+        x = self.conv1(x)
+        if self.downsample:
+            x = F.avg_pool2d(x, 2)
+        if self.normalize:
+            x = self.norm2(x)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x):
+        x = self._shortcut(x) + self._residual(x)
+        return x /self.equal_var   # unit variance
+
+class AdaIN(nn.Module):
+    def __init__(self, style_dim, num_features):
+        super().__init__()
+        self.norm = nn.InstanceNorm2d(num_features, affine=False)
+        self.fc = nn.Linear(style_dim, num_features*2)
+
+    def forward(self, x, s):
+        h = self.fc(s)
+        h = h.view(h.size(0), h.size(1), 1, 1)
+        gamma, beta = torch.chunk(h, chunks=2, dim=1)
+        return (1 + gamma) * self.norm(x) + beta
+
+class ResUpBlk(nn.Module):
+    def __init__(self, dim_in, dim_out,actv=nn.LeakyReLU(0.2),normalize="in"):
+        super().__init__()
+        self.actv = actv
+        self.normalize = normalize
+        self.learned_sc = dim_in != dim_out
+        self.equal_var = math.sqrt(2)
+        self._build_weights(dim_in, dim_out)
+
+    def _build_weights(self, dim_in, dim_out):
+        self.conv1 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        self.conv2 = nn.Conv2d(dim_out, dim_out, 3, 1, 1)
+        if self.normalize.lower() == "in":
+            self.norm1 = nn.InstanceNorm2d(dim_in, affine=True)
+            self.norm2 = nn.InstanceNorm2d(dim_out, affine=True)
+        elif self.normalize.lower() == "bn":
+            self.norm1 = nn.BatchNorm2d(dim_in)
+            self.norm2 = nn.BatchNorm2d(dim_out)
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        x = F.interpolate(x, scale_factor=2, mode='nearest')
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        return x
+
+    def _residual(self, x):
+        x = self.norm1(x)
+        x = self.actv(x)
+        x = F.interpolate(x, scale_factor=2, mode='nearest')
+        x = self.conv1(x)
+        x = self.norm2(x)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x):
+        out = self._residual(x)
+        out = (out + self._shortcut(x)) / self.equal_var
+        return out
+
+class AdainResBlk(nn.Module):
+    def __init__(self, dim_in, dim_out, style_dim=512,
+                 actv=nn.LeakyReLU(0.2), upsample=False):
+        super().__init__()
+        self.actv = actv
+        self.upsample = upsample
+        self.learned_sc = dim_in != dim_out
+        self.equal_var = math.sqrt(2)
+        self._build_weights(dim_in, dim_out, style_dim)
+
+    def _build_weights(self, dim_in, dim_out, style_dim=64):
+        self.conv1 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        self.conv2 = nn.Conv2d(dim_out, dim_out, 3, 1, 1)
+        self.norm1 = AdaIN(style_dim, dim_in)
+        self.norm2 = AdaIN(style_dim, dim_out)
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        if self.upsample:
+            x = F.interpolate(x, scale_factor=2, mode='nearest')
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        return x
+
+    def _residual(self, x, s):
+        x = self.norm1(x, s)
+        x = self.actv(x)
+        if self.upsample:
+            x = F.interpolate(x, scale_factor=2, mode='nearest')
+        x = self.conv1(x)
+        x = self.norm2(x, s)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x, s):
+        out = self._residual(x, s)
+        out = (out + self._shortcut(x)) / self.equal_var
+        return out
+
+
+class Generator(nn.Module):
+    def __init__(
+                self,
+                **kwargs
+                ):
+        super().__init__()
+
+        id_dim      = kwargs["id_dim"]
+        k_size      = kwargs["g_kernel_size"]
+        res_num     = kwargs["res_num"]
+        in_channel  = kwargs["in_channel"]
+        up_mode     = kwargs["up_mode"]
+        norm        = kwargs["norm"]
+
+        lstu_script = kwargs["lstu_script"]
+        lstu_class  = kwargs["lstu_class"]
+        
+        aggregator  = kwargs["aggregator"]
+        res_mode    = kwargs["res_mode"]
+
+        padding_size= int((k_size -1)/2)
+        padding_type= 'reflect'
+
+        script_name = "components." + lstu_script
+        package     = __import__(script_name, fromlist=True)
+        lstu_class  = getattr(package, lstu_class)
+        
+        
+        activation = nn.LeakyReLU(0.2)
+        # activation = nn.ReLU()
+
+        self.from_rgb = nn.Conv2d(3, in_channel, 1, 1, 0)
+        # self.first_layer = nn.Sequential(nn.Conv2d(3, 64, kernel_size=3, padding=1, bias=False),
+        #                         nn.BatchNorm2d(64), activation)
+        ### downsample
+        self.down1 = ResBlk(in_channel, in_channel, normalize=norm, downsample=True)# 256
+
+        self.down2 = ResBlk(in_channel, in_channel*2, normalize=norm, downsample=True)# 128
+                                
+        self.down3 = ResBlk(in_channel*2, in_channel*4,normalize=norm, downsample=True)# 64
+
+        self.down4 = ResBlk(in_channel*4, in_channel*8, normalize=norm, downsample=True)# 32
+
+        self.down5 = ResBlk(in_channel*8, in_channel*8, normalize=norm, downsample=True)# 16
+
+        # self.down6 = ResBlk(in_channel*8, in_channel*8, normalize=True, downsample=True)# 8
+        
+        ### resnet blocks
+        BN = []
+        for i in range(res_num):
+            BN += [
+                AdainResBlk(in_channel*8, in_channel*8, style_dim=id_dim, upsample=False)]
+        self.BottleNeck = nn.Sequential(*BN)
+
+        # self.up6 = AdainResBlk(in_channel*8, in_channel*8, style_dim=id_dim, upsample=True) # 16
+        
+        self.up5 = AdainResBlk(in_channel*8, in_channel*8, style_dim=id_dim, upsample=True) # 32
+
+        self.up4 = AdainResBlk(in_channel*8, in_channel*4, style_dim=id_dim, upsample=True) # 64
+
+        self.up3 = AdainResBlk(in_channel*4, in_channel*2, style_dim=id_dim, upsample=True) # 128
+
+        self.lstu = lstu_class(in_channel*2,norm)
+        
+        self.up2 = AdainResBlk(in_channel*2, in_channel, style_dim=id_dim, upsample=True)
+        # ResUpBlk(in_channel*2, in_channel, normalize="in") # 256
+
+        # self.lstu  = nn.Sequential(nn.Conv2d(in_channel*2, in_channel, kernel_size=3, stride=1, padding=1, bias=False),
+        #                         nn.BatchNorm2d(in_channel),
+        #                         activation,
+        #                         nn.Conv2d(in_channel, 1, kernel_size=3, stride=1, padding=1),
+        #                         nn.Sigmoid()
+        #                         )
+
+        self.up1 = AdainResBlk(in_channel, in_channel, style_dim=id_dim, upsample=True)
+        # ResUpBlk(in_channel, in_channel, normalize="in") # 512
+
+        if norm.lower() == "in":
+            norm_out = nn.InstanceNorm2d(in_channel, affine=True)
+        elif norm.lower() == "bn":
+            norm_out = nn.BatchNorm2d(in_channel)
+        
+        self.to_rgb = nn.Sequential(
+            norm_out,
+            activation,
+            nn.Conv2d(in_channel, 3, 3, 1, 1))
+
+        # self.last_layer = nn.Sequential(nn.ReflectionPad2d(3),
+        #             nn.Conv2d(64, 3, kernel_size=7, padding=0))
+
+
+    #     self.__weights_init__()
+
+    # def __weights_init__(self):
+    #     for layer in self.encoder:
+    #         if isinstance(layer,nn.Conv2d):
+    #             nn.init.xavier_uniform_(layer.weight)
+
+    #     for layer in self.encoder2:
+    #         if isinstance(layer,nn.Conv2d):
+    #             nn.init.xavier_uniform_(layer.weight)
+
+    def forward(self, img, id):
+        res = self.from_rgb(img)
+        res = self.down1(res)
+        skip = self.down2(res)
+        res = self.down3(skip)
+        res = self.down4(res)
+        res = self.down5(res)
+        # res = self.down6(res)
+        for i in range(len(self.BottleNeck)):
+            res = self.BottleNeck[i](res, id)
+        # res = self.up6(res,id)
+        res = self.up5(res,id)
+        res = self.up4(res,id)
+        res = self.up3(res,id)
+        res,mask = self.lstu(skip, res)
+        res = self.up2(res,id) #  + skip
+        res = self.up1(res,id)
+        res = self.to_rgb(res)
+        return res

components/Generator_Res_config3.py

@@ -0,0 +1,285 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+#############################################################
+# File: Generator_Invobn_config1.py
+# Created Date: Saturday February 26th 2022
+# Author: Chen Xuanhong
+# Email: chenxuanhongzju@outlook.com
+# Last Modified:  Wednesday, 30th March 2022 4:14:27 pm
+# Modified By: Chen Xuanhong
+# Copyright (c) 2022 Shanghai Jiao Tong University
+#############################################################
+
+import os
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+import math
+
+
+
+class ResBlk(nn.Module):
+    def __init__(self, dim_in, dim_out, actv=nn.LeakyReLU(0.2),
+                 normalize="in", downsample=False):
+        super().__init__()
+        self.actv = actv
+        self.normalize = normalize
+        self.downsample = downsample
+        self.learned_sc = dim_in != dim_out
+        self.equal_var  = math.sqrt(2)
+        self._build_weights(dim_in, dim_out)
+
+    def _build_weights(self, dim_in, dim_out):
+        self.conv1 = nn.Conv2d(dim_in, dim_in, 3, 1, 1)
+        self.conv2 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        if self.normalize.lower() == "in":
+            self.norm1 = nn.InstanceNorm2d(dim_in, affine=True)
+            self.norm2 = nn.InstanceNorm2d(dim_in, affine=True)
+        elif self.normalize.lower() == "bn":
+            self.norm1 = nn.BatchNorm2d(dim_in)
+            self.norm2 = nn.BatchNorm2d(dim_in)
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        if self.downsample:
+            x = F.avg_pool2d(x, 2)
+        return x
+
+    def _residual(self, x):
+        if self.normalize:
+            x = self.norm1(x)
+        x = self.actv(x)
+        x = self.conv1(x)
+        if self.downsample:
+            x = F.avg_pool2d(x, 2)
+        if self.normalize:
+            x = self.norm2(x)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x):
+        x = self._shortcut(x) + self._residual(x)
+        return x /self.equal_var   # unit variance
+
+class AdaIN(nn.Module):
+    def __init__(self, style_dim, num_features):
+        super().__init__()
+        self.norm = nn.InstanceNorm2d(num_features, affine=False)
+        self.fc = nn.Linear(style_dim, num_features*2)
+
+    def forward(self, x, s):
+        h = self.fc(s)
+        h = h.view(h.size(0), h.size(1), 1, 1)
+        gamma, beta = torch.chunk(h, chunks=2, dim=1)
+        return (1 + gamma) * self.norm(x) + beta
+
+class ResUpBlk(nn.Module):
+    def __init__(self, dim_in, dim_out,actv=nn.LeakyReLU(0.2),normalize="in"):
+        super().__init__()
+        self.actv = actv
+        self.normalize = normalize
+        self.learned_sc = dim_in != dim_out
+        self.equal_var = math.sqrt(2)
+        self._build_weights(dim_in, dim_out)
+
+    def _build_weights(self, dim_in, dim_out):
+        self.conv1 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        self.conv2 = nn.Conv2d(dim_out, dim_out, 3, 1, 1)
+        if self.normalize.lower() == "in":
+            self.norm1 = nn.InstanceNorm2d(dim_in, affine=True)
+            self.norm2 = nn.InstanceNorm2d(dim_out, affine=True)
+        elif self.normalize.lower() == "bn":
+            self.norm1 = nn.BatchNorm2d(dim_in)
+            self.norm2 = nn.BatchNorm2d(dim_out)
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        x = F.interpolate(x, scale_factor=2, mode='nearest')
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        return x
+
+    def _residual(self, x):
+        x = self.norm1(x)
+        x = self.actv(x)
+        x = F.interpolate(x, scale_factor=2, mode='nearest')
+        x = self.conv1(x)
+        x = self.norm2(x)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x):
+        out = self._residual(x)
+        out = (out + self._shortcut(x)) / self.equal_var
+        return out
+
+class AdainResBlk(nn.Module):
+    def __init__(self, dim_in, dim_out, style_dim=512,
+                 actv=nn.LeakyReLU(0.2), upsample=False):
+        super().__init__()
+        self.actv = actv
+        self.upsample = upsample
+        self.learned_sc = dim_in != dim_out
+        self.equal_var = math.sqrt(2)
+        self._build_weights(dim_in, dim_out, style_dim)
+
+    def _build_weights(self, dim_in, dim_out, style_dim=64):
+        self.conv1 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        self.conv2 = nn.Conv2d(dim_out, dim_out, 3, 1, 1)
+        self.norm1 = AdaIN(style_dim, dim_in)
+        self.norm2 = AdaIN(style_dim, dim_out)
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        if self.upsample:
+            x = F.interpolate(x, scale_factor=2, mode='nearest')
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        return x
+
+    def _residual(self, x, s):
+        x = self.norm1(x, s)
+        x = self.actv(x)
+        if self.upsample:
+            x = F.interpolate(x, scale_factor=2, mode='nearest')
+        x = self.conv1(x)
+        x = self.norm2(x, s)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x, s):
+        out = self._residual(x, s)
+        out = (out + self._shortcut(x)) / self.equal_var
+        return out
+
+
+class Generator(nn.Module):
+    def __init__(
+                self,
+                **kwargs
+                ):
+        super().__init__()
+
+        id_dim      = kwargs["id_dim"]
+        k_size      = kwargs["g_kernel_size"]
+        res_num     = kwargs["res_num"]
+        in_channel  = kwargs["in_channel"]
+        up_mode     = kwargs["up_mode"]
+        norm        = kwargs["norm"]
+
+        lstu_script = kwargs["lstu_script"]
+        lstu_class  = kwargs["lstu_class"]
+        
+        aggregator  = kwargs["aggregator"]
+        res_mode    = kwargs["res_mode"]
+
+        padding_size= int((k_size -1)/2)
+        padding_type= 'reflect'
+
+        script_name = "components." + lstu_script
+        package     = __import__(script_name, fromlist=True)
+        lstu_class  = getattr(package, lstu_class)
+        
+        
+        activation = nn.LeakyReLU(0.2)
+        # activation = nn.ReLU()
+
+        self.from_rgb = nn.Conv2d(3, in_channel, 1, 1, 0)
+        # self.first_layer = nn.Sequential(nn.Conv2d(3, 64, kernel_size=3, padding=1, bias=False),
+        #                         nn.BatchNorm2d(64), activation)
+        ### downsample
+        self.down1 = ResBlk(in_channel, in_channel, normalize=norm, downsample=True)# 256
+
+        self.down2 = ResBlk(in_channel, in_channel*2, normalize=norm, downsample=True)# 128
+                                
+        self.down3 = ResBlk(in_channel*2, in_channel*4,normalize=norm, downsample=True)# 64
+
+        self.down4 = ResBlk(in_channel*4, in_channel*8, normalize=norm, downsample=True)# 32
+
+        self.down5 = ResBlk(in_channel*8, in_channel*8, normalize=norm, downsample=True)# 16
+
+        # self.down6 = ResBlk(in_channel*8, in_channel*8, normalize=True, downsample=True)# 8
+        
+        ### resnet blocks
+        BN = []
+        for i in range(res_num):
+            BN += [
+                AdainResBlk(in_channel*8, in_channel*8, style_dim=id_dim, upsample=False)]
+        self.BottleNeck = nn.Sequential(*BN)
+
+        # self.up6 = AdainResBlk(in_channel*8, in_channel*8, style_dim=id_dim, upsample=True) # 16
+        
+        self.up5 = AdainResBlk(in_channel*8, in_channel*8, style_dim=id_dim, upsample=True) # 32
+
+        self.up4 = AdainResBlk(in_channel*8, in_channel*4, style_dim=id_dim, upsample=True) # 64
+
+        self.up3 = AdainResBlk(in_channel*4, in_channel*2, style_dim=id_dim, upsample=True) # 128
+
+        self.lstu = lstu_class(in_channel*2,norm)
+        
+        self.up2 = AdainResBlk(in_channel*2, in_channel, style_dim=id_dim, upsample=True)
+        # ResUpBlk(in_channel*2, in_channel, normalize="in") # 256
+
+        # self.lstu  = nn.Sequential(nn.Conv2d(in_channel*2, in_channel, kernel_size=3, stride=1, padding=1, bias=False),
+        #                         nn.BatchNorm2d(in_channel),
+        #                         activation,
+        #                         nn.Conv2d(in_channel, 1, kernel_size=3, stride=1, padding=1),
+        #                         nn.Sigmoid()
+        #                         )
+
+        self.up1 = AdainResBlk(in_channel, in_channel, style_dim=id_dim, upsample=True)
+        # ResUpBlk(in_channel, in_channel, normalize="in") # 512
+
+        if norm.lower() == "in":
+            norm_out = nn.InstanceNorm2d(in_channel, affine=True)
+        elif norm.lower() == "bn":
+            norm_out = nn.BatchNorm2d(in_channel)
+        
+        self.to_rgb = nn.Sequential(
+            norm_out,
+            activation,
+            nn.Conv2d(in_channel, 3, 3, 1, 1))
+
+        # self.last_layer = nn.Sequential(nn.ReflectionPad2d(3),
+        #             nn.Conv2d(64, 3, kernel_size=7, padding=0))
+
+
+    #     self.__weights_init__()
+
+    # def __weights_init__(self):
+    #     for layer in self.encoder:
+    #         if isinstance(layer,nn.Conv2d):
+    #             nn.init.xavier_uniform_(layer.weight)
+
+    #     for layer in self.encoder2:
+    #         if isinstance(layer,nn.Conv2d):
+    #             nn.init.xavier_uniform_(layer.weight)
+
+    def forward(self, img, id):
+        res = self.from_rgb(img)
+        res = self.down1(res)
+        skip = self.down2(res)
+        res = self.down3(skip)
+        res = self.down4(res)
+        res = self.down5(res)
+        # res = self.down6(res)
+        for i in range(len(self.BottleNeck)):
+            res = self.BottleNeck[i](res, id)
+        # res = self.up6(res,id)
+        res = self.up5(res,id)
+        res = self.up4(res,id)
+        res = self.up3(res,id)
+        # res,mask = self.lstu(skip, res)
+        res = self.up2(res,id) #  + skip
+        res = self.up1(res,id)
+        res = self.to_rgb(res)
+        return res

components/Generator_VGGStyle_maskhead_config.py

@@ -0,0 +1,204 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+#############################################################
+# File: Generator_Invobn_config1.py
+# Created Date: Saturday February 26th 2022
+# Author: Chen Xuanhong
+# Email: chenxuanhongzju@outlook.com
+# Last Modified:  Wednesday, 6th April 2022 12:55:51 am
+# Modified By: Chen Xuanhong
+# Copyright (c) 2022 Shanghai Jiao Tong University
+#############################################################
+
+import os
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+import math
+
+
+class AdaIN(nn.Module):
+    def __init__(self, style_dim, num_features):
+        super().__init__()
+        self.norm = nn.InstanceNorm2d(num_features, affine=False)
+        self.fc = nn.Linear(style_dim, num_features*2)
+
+    def forward(self, x, s):
+        h = self.fc(s)
+        h = h.view(h.size(0), h.size(1), 1, 1)
+        gamma, beta = torch.chunk(h, chunks=2, dim=1)
+        return (1 + gamma) * self.norm(x) + beta
+
+class AdainResBlk(nn.Module):
+    def __init__(self, dim_in, dim_out, style_dim=512,
+                 actv=nn.LeakyReLU(0.2), upsample=False):
+        super().__init__()
+        self.actv = actv
+        self.upsample = upsample
+        self.learned_sc = dim_in != dim_out
+        self.equal_var = math.sqrt(2)
+        self._build_weights(dim_in, dim_out, style_dim)
+
+    def _build_weights(self, dim_in, dim_out, style_dim=64):
+        self.conv1 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        self.conv2 = nn.Conv2d(dim_out, dim_out, 3, 1, 1)
+        self.norm1 = AdaIN(style_dim, dim_in)
+        self.norm2 = AdaIN(style_dim, dim_out)
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        if self.upsample:
+            x = F.interpolate(x, scale_factor=2, mode='nearest')
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        return x
+
+    def _residual(self, x, s):
+        x = self.norm1(x, s)
+        x = self.actv(x)
+        if self.upsample:
+            x = F.interpolate(x, scale_factor=2, mode='nearest')
+        x = self.conv1(x)
+        x = self.norm2(x, s)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x, s):
+        out = self._residual(x, s)
+        out = (out + self._shortcut(x)) / self.equal_var
+        return out
+
+class AdainUpBlock(nn.Module):
+    def __init__(self, dim_in, dim_out, style_dim=512,
+                 actv=nn.LeakyReLU(0.2)):
+        super().__init__()
+        self.actv = actv
+        self._build_weights(dim_in, dim_out, style_dim)
+
+    def _build_weights(self, dim_in, dim_out, style_dim=64):
+        self.conv = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        self.norm = AdaIN(style_dim, dim_out)
+
+    def forward(self, x, s):
+        x = F.interpolate(x, scale_factor=2, mode='nearest')
+        x = self.conv(x)
+        x = self.norm(x, s)
+        x = self.actv(x)
+        return x
+
+
+class Generator(nn.Module):
+    def __init__(
+                self,
+                **kwargs
+                ):
+        super().__init__()
+
+        id_dim      = kwargs["id_dim"]
+        k_size      = kwargs["g_kernel_size"]
+        res_num     = kwargs["res_num"]
+        in_channel  = kwargs["in_channel"]
+        up_mode     = kwargs["up_mode"]
+        norm        = kwargs["norm"]
+        
+        aggregator  = kwargs["aggregator"]
+        res_mode    = kwargs["res_mode"]
+
+        padding_size= int((k_size -1)/2)
+        padding_type= 'reflect'
+        
+        
+        activation = nn.LeakyReLU(0.2)
+        # activation = nn.ReLU()
+
+        self.from_rgb = nn.Conv2d(3, in_channel, 1, 1, 0)
+        # self.first_layer = nn.Sequential(nn.Conv2d(3, 64, kernel_size=3, padding=1, bias=False),
+        #                         nn.BatchNorm2d(64), activation)
+        ### downsample
+        self.down1 = nn.Sequential(nn.Conv2d(in_channel, in_channel, kernel_size=3, stride=2, padding=1, bias=False), # 256
+                                nn.BatchNorm2d(in_channel), activation)
+                                
+        self.down2 = nn.Sequential(nn.Conv2d(in_channel, in_channel*2, kernel_size=3, stride=2, padding=1, bias=False), # 128
+                                nn.BatchNorm2d(in_channel*2), activation)
+
+        self.down3 = nn.Sequential(nn.Conv2d(in_channel*2, in_channel*4, kernel_size=3, stride=2, padding=1, bias=False), # 64
+                                nn.BatchNorm2d(in_channel*4), activation)
+
+        self.down4 = nn.Sequential(nn.Conv2d(in_channel*4, in_channel*8, kernel_size=3, stride=2, padding=1, bias=False), # 32
+                                nn.BatchNorm2d(in_channel*8), activation)
+        
+        self.down5 = nn.Sequential(nn.Conv2d(in_channel*8, in_channel*8, kernel_size=3, stride=2, padding=1, bias=False), # 32
+                                nn.BatchNorm2d(in_channel*8), activation)
+
+        # self.down6 = ResBlk(in_channel*8, in_channel*8, normalize=True, downsample=True)# 8
+        self.maskhead = nn.Sequential(
+                                nn.UpsamplingNearest2d(scale_factor = 2),
+                                nn.Conv2d(in_channel*8, in_channel, kernel_size=3, stride=1, padding=1,bias=False),
+                                nn.BatchNorm2d(in_channel), # 32
+                                activation,
+                                nn.UpsamplingNearest2d(scale_factor = 2),
+                                nn.Conv2d(in_channel, in_channel//2, kernel_size=3, stride=1, padding=1, bias=False),
+                                nn.BatchNorm2d(in_channel//2), # 64
+                                activation,
+                                nn.UpsamplingNearest2d(scale_factor = 2),
+                                nn.Conv2d(in_channel//2, 1, kernel_size=3, stride=1, padding=1, bias=False),
+                                nn.Sigmoid()
+        )
+        
+        ### resnet blocks
+        BN = []
+        for i in range(res_num):
+            BN += [
+                AdainResBlk(in_channel*8, in_channel*8, style_dim=id_dim, upsample=False)]
+        self.BottleNeck = nn.Sequential(*BN)
+
+        # self.up6 = AdainResBlk(in_channel*8, in_channel*8, style_dim=id_dim, upsample=True) # 16
+        
+        self.up5 = AdainUpBlock(in_channel*8, in_channel*8, style_dim=id_dim) # 32
+
+        self.up4 = AdainUpBlock(in_channel*8, in_channel*4, style_dim=id_dim) # 64
+
+        self.up3 = AdainUpBlock(in_channel*4, in_channel*2, style_dim=id_dim) # 128
+        
+        self.up2 = AdainUpBlock(in_channel*2, in_channel, style_dim=id_dim)
+
+        self.up1 = AdainUpBlock(in_channel, in_channel, style_dim=id_dim)
+        # ResUpBlk(in_channel, in_channel, normalize="in") # 512
+
+
+        self.to_rgb = nn.Sequential(nn.ReflectionPad2d(1),
+                    nn.Conv2d(in_channel, 3, kernel_size=3, padding=0))
+
+
+    #     self.__weights_init__()
+
+    # def __weights_init__(self):
+    #     for layer in self.encoder:
+    #         if isinstance(layer,nn.Conv2d):
+    #             nn.init.xavier_uniform_(layer.weight)
+
+    #     for layer in self.encoder2:
+    #         if isinstance(layer,nn.Conv2d):
+    #             nn.init.xavier_uniform_(layer.weight)
+
+    def forward(self, img, id):
+        res = self.from_rgb(img)
+        res = self.down1(res)
+        skip = self.down2(res)
+        res = self.down3(skip)
+        res = self.down4(res)
+        res = self.down5(res)
+        mask= self.maskhead(res)
+        for i in range(len(self.BottleNeck)):
+            res = self.BottleNeck[i](res, id)
+        res = self.up5(res,id)
+        res = self.up4(res,id)
+        res = self.up3(res,id)
+        res = (1-mask) * skip + mask * res
+        res = self.up2(res,id) #  + skip
+        res = self.up1(res,id)
+        res = self.to_rgb(res)
+        return res, mask

components/Generator_featout_config.py

@@ -0,0 +1,298 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+#############################################################
+# File: Generator_Invobn_config1.py
+# Created Date: Saturday February 26th 2022
+# Author: Chen Xuanhong
+# Email: chenxuanhongzju@outlook.com
+# Last Modified:  Saturday, 2nd April 2022 1:27:23 pm
+# Modified By: Chen Xuanhong
+# Copyright (c) 2022 Shanghai Jiao Tong University
+#############################################################
+
+import os
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+import math
+
+
+
+class ResBlk(nn.Module):
+    def __init__(self, dim_in, dim_out, actv=nn.LeakyReLU(0.2),
+                 normalize="in", downsample=False):
+        super().__init__()
+        self.actv = actv
+        self.normalize = normalize
+        self.downsample = downsample
+        self.learned_sc = dim_in != dim_out
+        self.equal_var  = math.sqrt(2)
+        self._build_weights(dim_in, dim_out)
+
+    def _build_weights(self, dim_in, dim_out):
+        self.conv1 = nn.Conv2d(dim_in, dim_in, 3, 1, 1)
+        self.conv2 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        if self.normalize.lower() == "in":
+            self.norm1 = nn.InstanceNorm2d(dim_in, affine=True)
+            self.norm2 = nn.InstanceNorm2d(dim_in, affine=True)
+        elif self.normalize.lower() == "bn":
+            self.norm1 = nn.BatchNorm2d(dim_in)
+            self.norm2 = nn.BatchNorm2d(dim_in)
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        if self.downsample:
+            x = F.avg_pool2d(x, 2)
+        return x
+
+    def _residual(self, x):
+        if self.normalize:
+            x = self.norm1(x)
+        x = self.actv(x)
+        x = self.conv1(x)
+        if self.downsample:
+            x = F.avg_pool2d(x, 2)
+        if self.normalize:
+            x = self.norm2(x)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x):
+        x = self._shortcut(x) + self._residual(x)
+        return x /self.equal_var   # unit variance
+
+class AdaIN(nn.Module):
+    def __init__(self, style_dim, num_features):
+        super().__init__()
+        self.norm = nn.InstanceNorm2d(num_features, affine=False)
+        self.fc = nn.Linear(style_dim, num_features*2)
+
+    def forward(self, x, s):
+        h = self.fc(s)
+        h = h.view(h.size(0), h.size(1), 1, 1)
+        gamma, beta = torch.chunk(h, chunks=2, dim=1)
+        return (1 + gamma) * self.norm(x) + beta
+
+class ResUpBlk(nn.Module):
+    def __init__(self, dim_in, dim_out,actv=nn.LeakyReLU(0.2),normalize="in"):
+        super().__init__()
+        self.actv = actv
+        self.normalize = normalize
+        self.learned_sc = dim_in != dim_out
+        self.equal_var = math.sqrt(2)
+        self._build_weights(dim_in, dim_out)
+
+    def _build_weights(self, dim_in, dim_out):
+        self.conv1 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        self.conv2 = nn.Conv2d(dim_out, dim_out, 3, 1, 1)
+        if self.normalize.lower() == "in":
+            self.norm1 = nn.InstanceNorm2d(dim_in, affine=True)
+            self.norm2 = nn.InstanceNorm2d(dim_out, affine=True)
+        elif self.normalize.lower() == "bn":
+            self.norm1 = nn.BatchNorm2d(dim_in)
+            self.norm2 = nn.BatchNorm2d(dim_out)
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        x = F.interpolate(x, scale_factor=2, mode='nearest')
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        return x
+
+    def _residual(self, x):
+        x = self.norm1(x)
+        x = self.actv(x)
+        x = F.interpolate(x, scale_factor=2, mode='nearest')
+        x = self.conv1(x)
+        x = self.norm2(x)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x):
+        out = self._residual(x)
+        out = (out + self._shortcut(x)) / self.equal_var
+        return out
+
+class AdainResBlk(nn.Module):
+    def __init__(self, dim_in, dim_out, style_dim=512,
+                 actv=nn.LeakyReLU(0.2), upsample=False):
+        super().__init__()
+        self.actv = actv
+        self.upsample = upsample
+        self.learned_sc = dim_in != dim_out
+        self.equal_var = math.sqrt(2)
+        self._build_weights(dim_in, dim_out, style_dim)
+
+    def _build_weights(self, dim_in, dim_out, style_dim=64):
+        self.conv1 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        self.conv2 = nn.Conv2d(dim_out, dim_out, 3, 1, 1)
+        self.norm1 = AdaIN(style_dim, dim_in)
+        self.norm2 = AdaIN(style_dim, dim_out)
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        if self.upsample:
+            x = F.interpolate(x, scale_factor=2, mode='nearest')
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        return x
+
+    def _residual(self, x, s):
+        x = self.norm1(x, s)
+        x = self.actv(x)
+        if self.upsample:
+            x = F.interpolate(x, scale_factor=2, mode='nearest')
+        x = self.conv1(x)
+        x = self.norm2(x, s)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x, s):
+        out = self._residual(x, s)
+        out = (out + self._shortcut(x)) / self.equal_var
+        return out
+
+
+class Generator(nn.Module):
+    def __init__(
+                self,
+                **kwargs
+                ):
+        super().__init__()
+
+        id_dim      = kwargs["id_dim"]
+        k_size      = kwargs["g_kernel_size"]
+        res_num     = kwargs["res_num"]
+        in_channel  = kwargs["in_channel"]
+        up_mode     = kwargs["up_mode"]
+        norm        = kwargs["norm"]
+
+        lstu_script = kwargs["lstu_script"]
+        lstu_class  = kwargs["lstu_class"]
+        
+        aggregator  = kwargs["aggregator"]
+        res_mode    = kwargs["res_mode"]
+
+        padding_size= int((k_size -1)/2)
+        padding_type= 'reflect'
+
+        script_name = "components." + lstu_script
+        package     = __import__(script_name, fromlist=True)
+        lstu_class  = getattr(package, lstu_class)
+        
+        
+        activation = nn.LeakyReLU(0.2)
+        # activation = nn.ReLU()
+
+        self.from_rgb = nn.Conv2d(3, in_channel, 1, 1, 0)
+        # self.first_layer = nn.Sequential(nn.Conv2d(3, 64, kernel_size=3, padding=1, bias=False),
+        #                         nn.BatchNorm2d(64), activation)
+        ### downsample
+        self.down1 = ResBlk(in_channel, in_channel, normalize=norm, downsample=True)# 256
+
+        self.down2 = ResBlk(in_channel, in_channel*2, normalize=norm, downsample=True)# 128
+                                
+        self.down3 = ResBlk(in_channel*2, in_channel*4,normalize=norm, downsample=True)# 64
+
+        self.down4 = ResBlk(in_channel*4, in_channel*8, normalize=norm, downsample=True)# 32
+
+        self.down5 = ResBlk(in_channel*8, in_channel*8, normalize=norm, downsample=True)# 16
+
+        # self.down6 = ResBlk(in_channel*8, in_channel*8, normalize=True, downsample=True)# 8
+        self.maskhead = nn.Sequential(
+                                nn.ConvTranspose2d(in_channel*8, in_channel, kernel_size=4, stride=2, padding=1, bias=False),
+                                nn.BatchNorm2d(in_channel), # 32
+                                activation,
+                                nn.ConvTranspose2d(in_channel, in_channel//2, kernel_size=4, stride=2, padding=1, bias=False),
+                                nn.BatchNorm2d(in_channel), # 64
+                                activation,
+                                nn.ConvTranspose2d(in_channel//2, 1, kernel_size=4, stride=2, padding=1, bias=False),
+                                nn.BatchNorm2d(in_channel), # 128
+                                nn.Sigmoid()
+        )
+        
+        ### resnet blocks
+        BN = []
+        for i in range(res_num):
+            BN += [
+                AdainResBlk(in_channel*8, in_channel*8, style_dim=id_dim, upsample=False)]
+        self.BottleNeck = nn.Sequential(*BN)
+
+        # self.up6 = AdainResBlk(in_channel*8, in_channel*8, style_dim=id_dim, upsample=True) # 16
+        
+        self.up5 = AdainResBlk(in_channel*8, in_channel*8, style_dim=id_dim, upsample=True) # 32
+
+        self.up4 = AdainResBlk(in_channel*8, in_channel*4, style_dim=id_dim, upsample=True) # 64
+
+        self.up3 = AdainResBlk(in_channel*4, in_channel*2, style_dim=id_dim, upsample=True) # 128
+
+        self.lstu = lstu_class(in_channel*2,norm)
+        
+        self.up2 = AdainResBlk(in_channel*2, in_channel, style_dim=id_dim, upsample=True)
+        # ResUpBlk(in_channel*2, in_channel, normalize="in") # 256
+
+        # self.lstu  = nn.Sequential(nn.Conv2d(in_channel*2, in_channel, kernel_size=3, stride=1, padding=1, bias=False),
+        #                         nn.BatchNorm2d(in_channel),
+        #                         activation,
+        #                         nn.Conv2d(in_channel, 1, kernel_size=3, stride=1, padding=1),
+        #                         nn.Sigmoid()
+        #                         )
+
+        self.up1 = AdainResBlk(in_channel, in_channel, style_dim=id_dim, upsample=True)
+        # ResUpBlk(in_channel, in_channel, normalize="in") # 512
+
+        if norm.lower() == "in":
+            norm_out = nn.InstanceNorm2d(in_channel, affine=True)
+        elif norm.lower() == "bn":
+            norm_out = nn.BatchNorm2d(in_channel)
+        
+        self.to_rgb = nn.Sequential(
+            norm_out,
+            activation,
+            nn.Conv2d(in_channel, 3, 3, 1, 1))
+
+        # self.last_layer = nn.Sequential(nn.ReflectionPad2d(3),
+        #             nn.Conv2d(64, 3, kernel_size=7, padding=0))
+
+
+    #     self.__weights_init__()
+
+    # def __weights_init__(self):
+    #     for layer in self.encoder:
+    #         if isinstance(layer,nn.Conv2d):
+    #             nn.init.xavier_uniform_(layer.weight)
+
+    #     for layer in self.encoder2:
+    #         if isinstance(layer,nn.Conv2d):
+    #             nn.init.xavier_uniform_(layer.weight)
+
+    def forward(self, img, id, feat_out=False):
+        res = self.from_rgb(img)
+        res = self.down1(res)
+        skip = self.down2(res)
+        res = self.down3(skip)
+        res = self.down4(res)
+        res = self.down5(res)
+        if feat_out:
+            return res
+        # res = self.down6(res)
+        for i in range(len(self.BottleNeck)):
+            res = self.BottleNeck[i](res, id)
+        # res = self.up6(res,id)
+        res = self.up5(res,id)
+        res = self.up4(res,id)
+        res = self.up3(res,id)
+        res,mask = self.lstu(skip, res)
+        res = self.up2(res,id) #  + skip
+        res = self.up1(res,id)
+        res = self.to_rgb(res)
+        return res

components/Generator_involution_maskhead_config.py

@@ -0,0 +1,280 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+#############################################################
+# File: Generator_Invobn_config1.py
+# Created Date: Saturday February 26th 2022
+# Author: Chen Xuanhong
+# Email: chenxuanhongzju@outlook.com
+# Last Modified:  Sunday, 3rd April 2022 1:06:31 am
+# Modified By: Chen Xuanhong
+# Copyright (c) 2022 Shanghai Jiao Tong University
+#############################################################
+
+import os
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+import math
+
+
+
+class ResBlk(nn.Module):
+    def __init__(self, dim_in, dim_out, actv=nn.LeakyReLU(0.2),
+                 normalize="in", downsample=False):
+        super().__init__()
+        self.actv = actv
+        self.normalize = normalize
+        self.downsample = downsample
+        self.learned_sc = dim_in != dim_out
+        self.equal_var  = math.sqrt(2)
+        self._build_weights(dim_in, dim_out)
+
+    def _build_weights(self, dim_in, dim_out):
+        self.conv1 = nn.Conv2d(dim_in, dim_in, 3, 1, 1)
+        self.conv2 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        if self.normalize.lower() == "in":
+            self.norm1 = nn.InstanceNorm2d(dim_in, affine=True)
+            self.norm2 = nn.InstanceNorm2d(dim_in, affine=True)
+        elif self.normalize.lower() == "bn":
+            self.norm1 = nn.BatchNorm2d(dim_in)
+            self.norm2 = nn.BatchNorm2d(dim_in)
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        if self.downsample:
+            x = F.avg_pool2d(x, 2)
+        return x
+
+    def _residual(self, x):
+        if self.normalize:
+            x = self.norm1(x)
+        x = self.actv(x)
+        x = self.conv1(x)
+        if self.downsample:
+            x = F.avg_pool2d(x, 2)
+        if self.normalize:
+            x = self.norm2(x)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x):
+        x = self._shortcut(x) + self._residual(x)
+        return x /self.equal_var   # unit variance
+
+class AdaIN(nn.Module):
+    def __init__(self, style_dim, num_features):
+        super().__init__()
+        self.norm = nn.InstanceNorm2d(num_features, affine=False)
+        self.fc = nn.Linear(style_dim, num_features*2)
+
+    def forward(self, x, s):
+        h = self.fc(s)
+        h = h.view(h.size(0), h.size(1), 1, 1)
+        gamma, beta = torch.chunk(h, chunks=2, dim=1)
+        return (1 + gamma) * self.norm(x) + beta
+
+class ResUpBlk(nn.Module):
+    def __init__(self, dim_in, dim_out,actv=nn.LeakyReLU(0.2),normalize="in"):
+        super().__init__()
+        self.actv = actv
+        self.normalize = normalize
+        self.learned_sc = dim_in != dim_out
+        self.equal_var = math.sqrt(2)
+        self._build_weights(dim_in, dim_out)
+
+    def _build_weights(self, dim_in, dim_out):
+        self.conv1 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        self.conv2 = nn.Conv2d(dim_out, dim_out, 3, 1, 1)
+        if self.normalize.lower() == "in":
+            self.norm1 = nn.InstanceNorm2d(dim_in, affine=True)
+            self.norm2 = nn.InstanceNorm2d(dim_out, affine=True)
+        elif self.normalize.lower() == "bn":
+            self.norm1 = nn.BatchNorm2d(dim_in)
+            self.norm2 = nn.BatchNorm2d(dim_out)
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        x = F.interpolate(x, scale_factor=2, mode='nearest')
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        return x
+
+    def _residual(self, x):
+        x = self.norm1(x)
+        x = self.actv(x)
+        x = F.interpolate(x, scale_factor=2, mode='nearest')
+        x = self.conv1(x)
+        x = self.norm2(x)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x):
+        out = self._residual(x)
+        out = (out + self._shortcut(x)) / self.equal_var
+        return out
+
+class AdainResBlk(nn.Module):
+    def __init__(self, dim_in, dim_out, style_dim=512,
+                 actv=nn.LeakyReLU(0.2), upsample=False):
+        super().__init__()
+        self.actv = actv
+        self.upsample = upsample
+        self.learned_sc = dim_in != dim_out
+        self.equal_var = math.sqrt(2)
+        self._build_weights(dim_in, dim_out, style_dim)
+
+    def _build_weights(self, dim_in, dim_out, style_dim=64):
+        self.conv1 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        self.conv2 = nn.Conv2d(dim_out, dim_out, 3, 1, 1)
+        self.norm1 = AdaIN(style_dim, dim_in)
+        self.norm2 = AdaIN(style_dim, dim_out)
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        if self.upsample:
+            x = F.interpolate(x, scale_factor=2, mode='nearest')
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        return x
+
+    def _residual(self, x, s):
+        x = self.norm1(x, s)
+        x = self.actv(x)
+        if self.upsample:
+            x = F.interpolate(x, scale_factor=2, mode='nearest')
+        x = self.conv1(x)
+        x = self.norm2(x, s)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x, s):
+        out = self._residual(x, s)
+        out = (out + self._shortcut(x)) / self.equal_var
+        return out
+
+
+class Generator(nn.Module):
+    def __init__(
+                self,
+                **kwargs
+                ):
+        super().__init__()
+
+        id_dim      = kwargs["id_dim"]
+        k_size      = kwargs["g_kernel_size"]
+        res_num     = kwargs["res_num"]
+        in_channel  = kwargs["in_channel"]
+        up_mode     = kwargs["up_mode"]
+        norm        = kwargs["norm"]
+        
+        aggregator  = kwargs["aggregator"]
+        res_mode    = kwargs["res_mode"]
+
+        padding_size= int((k_size -1)/2)
+        padding_type= 'reflect'
+        
+        
+        activation = nn.LeakyReLU(0.2)
+        # activation = nn.ReLU()
+
+        self.from_rgb = nn.Conv2d(3, in_channel, 1, 1, 0)
+        # self.first_layer = nn.Sequential(nn.Conv2d(3, 64, kernel_size=3, padding=1, bias=False),
+        #                         nn.BatchNorm2d(64), activation)
+        ### downsample
+        self.down1 = ResBlk(in_channel, in_channel, normalize=norm, downsample=True)# 256
+
+        self.down2 = ResBlk(in_channel, in_channel*2, normalize=norm, downsample=True)# 128
+                                
+        self.down3 = ResBlk(in_channel*2, in_channel*4,normalize=norm, downsample=True)# 64
+
+        self.down4 = ResBlk(in_channel*4, in_channel*8, normalize=norm, downsample=True)# 32
+
+        self.down5 = ResBlk(in_channel*8, in_channel*8, normalize=norm, downsample=True)# 16
+
+        # self.down6 = ResBlk(in_channel*8, in_channel*8, normalize=True, downsample=True)# 8
+        self.maskhead = nn.Sequential(
+                                nn.UpsamplingNearest2d(scale_factor = 2),
+                                nn.Conv2d(in_channel*8, in_channel, kernel_size=3, stride=1, padding=1,bias=False),
+                                nn.BatchNorm2d(in_channel), # 32
+                                activation,
+                                nn.UpsamplingNearest2d(scale_factor = 2),
+                                nn.Conv2d(in_channel, in_channel//2, kernel_size=3, stride=1, padding=1, bias=False),
+                                nn.BatchNorm2d(in_channel//2), # 64
+                                activation,
+                                nn.UpsamplingNearest2d(scale_factor = 2),
+                                nn.Conv2d(in_channel//2, 1, kernel_size=3, stride=1, padding=1, bias=False),
+                                nn.Sigmoid()
+        )
+        
+        ### resnet blocks
+        BN = []
+        for i in range(res_num):
+            BN += [
+                AdainResBlk(in_channel*8, in_channel*8, style_dim=id_dim, upsample=False)]
+        self.BottleNeck = nn.Sequential(*BN)
+
+        # self.up6 = AdainResBlk(in_channel*8, in_channel*8, style_dim=id_dim, upsample=True) # 16
+        
+        self.up5 = AdainResBlk(in_channel*8, in_channel*8, style_dim=id_dim, upsample=True) # 32
+
+        self.up4 = AdainResBlk(in_channel*8, in_channel*4, style_dim=id_dim, upsample=True) # 64
+
+        self.up3 = AdainResBlk(in_channel*4, in_channel*2, style_dim=id_dim, upsample=True) # 128
+        
+        self.up2 = AdainResBlk(in_channel*2, in_channel, style_dim=id_dim, upsample=True)
+
+        self.up1 = AdainResBlk(in_channel, in_channel, style_dim=id_dim, upsample=True)
+        # ResUpBlk(in_channel, in_channel, normalize="in") # 512
+
+        if norm.lower() == "in":
+            norm_out = nn.InstanceNorm2d(in_channel, affine=True)
+        elif norm.lower() == "bn":
+            norm_out = nn.BatchNorm2d(in_channel)
+        
+        self.to_rgb = nn.Sequential(
+            norm_out,
+            activation,
+            nn.Conv2d(in_channel, 3, 3, 1, 1))
+
+        # self.last_layer = nn.Sequential(nn.ReflectionPad2d(3),
+        #             nn.Conv2d(64, 3, kernel_size=7, padding=0))
+
+
+    #     self.__weights_init__()
+
+    # def __weights_init__(self):
+    #     for layer in self.encoder:
+    #         if isinstance(layer,nn.Conv2d):
+    #             nn.init.xavier_uniform_(layer.weight)
+
+    #     for layer in self.encoder2:
+    #         if isinstance(layer,nn.Conv2d):
+    #             nn.init.xavier_uniform_(layer.weight)
+
+    def forward(self, img, id):
+        res = self.from_rgb(img)
+        res = self.down1(res)
+        skip = self.down2(res)
+        res = self.down3(skip)
+        res = self.down4(res)
+        res = self.down5(res)
+        mask= self.maskhead(res)
+        for i in range(len(self.BottleNeck)):
+            res = self.BottleNeck[i](res, id)
+        res = self.up5(res,id)
+        res = self.up4(res,id)
+        res = self.up3(res,id)
+        res = (1-mask) * skip + mask * res
+        res = self.up2(res,id) #  + skip
+        res = self.up1(res,id)
+        res = self.to_rgb(res)
+        return res, mask

components/Generator_maskhead_config.py

@@ -0,0 +1,280 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+#############################################################
+# File: Generator_Invobn_config1.py
+# Created Date: Saturday February 26th 2022
+# Author: Chen Xuanhong
+# Email: chenxuanhongzju@outlook.com
+# Last Modified:  Sunday, 3rd April 2022 1:06:31 am
+# Modified By: Chen Xuanhong
+# Copyright (c) 2022 Shanghai Jiao Tong University
+#############################################################
+
+import os
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+import math
+
+
+
+class ResBlk(nn.Module):
+    def __init__(self, dim_in, dim_out, actv=nn.LeakyReLU(0.2),
+                 normalize="in", downsample=False):
+        super().__init__()
+        self.actv = actv
+        self.normalize = normalize
+        self.downsample = downsample
+        self.learned_sc = dim_in != dim_out
+        self.equal_var  = math.sqrt(2)
+        self._build_weights(dim_in, dim_out)
+
+    def _build_weights(self, dim_in, dim_out):
+        self.conv1 = nn.Conv2d(dim_in, dim_in, 3, 1, 1)
+        self.conv2 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        if self.normalize.lower() == "in":
+            self.norm1 = nn.InstanceNorm2d(dim_in, affine=True)
+            self.norm2 = nn.InstanceNorm2d(dim_in, affine=True)
+        elif self.normalize.lower() == "bn":
+            self.norm1 = nn.BatchNorm2d(dim_in)
+            self.norm2 = nn.BatchNorm2d(dim_in)
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        if self.downsample:
+            x = F.avg_pool2d(x, 2)
+        return x
+
+    def _residual(self, x):
+        if self.normalize:
+            x = self.norm1(x)
+        x = self.actv(x)
+        x = self.conv1(x)
+        if self.downsample:
+            x = F.avg_pool2d(x, 2)
+        if self.normalize:
+            x = self.norm2(x)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x):
+        x = self._shortcut(x) + self._residual(x)
+        return x /self.equal_var   # unit variance
+
+class AdaIN(nn.Module):
+    def __init__(self, style_dim, num_features):
+        super().__init__()
+        self.norm = nn.InstanceNorm2d(num_features, affine=False)
+        self.fc = nn.Linear(style_dim, num_features*2)
+
+    def forward(self, x, s):
+        h = self.fc(s)
+        h = h.view(h.size(0), h.size(1), 1, 1)
+        gamma, beta = torch.chunk(h, chunks=2, dim=1)
+        return (1 + gamma) * self.norm(x) + beta
+
+class ResUpBlk(nn.Module):
+    def __init__(self, dim_in, dim_out,actv=nn.LeakyReLU(0.2),normalize="in"):
+        super().__init__()
+        self.actv = actv
+        self.normalize = normalize
+        self.learned_sc = dim_in != dim_out
+        self.equal_var = math.sqrt(2)
+        self._build_weights(dim_in, dim_out)
+
+    def _build_weights(self, dim_in, dim_out):
+        self.conv1 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        self.conv2 = nn.Conv2d(dim_out, dim_out, 3, 1, 1)
+        if self.normalize.lower() == "in":
+            self.norm1 = nn.InstanceNorm2d(dim_in, affine=True)
+            self.norm2 = nn.InstanceNorm2d(dim_out, affine=True)
+        elif self.normalize.lower() == "bn":
+            self.norm1 = nn.BatchNorm2d(dim_in)
+            self.norm2 = nn.BatchNorm2d(dim_out)
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        x = F.interpolate(x, scale_factor=2, mode='nearest')
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        return x
+
+    def _residual(self, x):
+        x = self.norm1(x)
+        x = self.actv(x)
+        x = F.interpolate(x, scale_factor=2, mode='nearest')
+        x = self.conv1(x)
+        x = self.norm2(x)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x):
+        out = self._residual(x)
+        out = (out + self._shortcut(x)) / self.equal_var
+        return out
+
+class AdainResBlk(nn.Module):
+    def __init__(self, dim_in, dim_out, style_dim=512,
+                 actv=nn.LeakyReLU(0.2), upsample=False):
+        super().__init__()
+        self.actv = actv
+        self.upsample = upsample
+        self.learned_sc = dim_in != dim_out
+        self.equal_var = math.sqrt(2)
+        self._build_weights(dim_in, dim_out, style_dim)
+
+    def _build_weights(self, dim_in, dim_out, style_dim=64):
+        self.conv1 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        self.conv2 = nn.Conv2d(dim_out, dim_out, 3, 1, 1)
+        self.norm1 = AdaIN(style_dim, dim_in)
+        self.norm2 = AdaIN(style_dim, dim_out)
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        if self.upsample:
+            x = F.interpolate(x, scale_factor=2, mode='nearest')
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        return x
+
+    def _residual(self, x, s):
+        x = self.norm1(x, s)
+        x = self.actv(x)
+        if self.upsample:
+            x = F.interpolate(x, scale_factor=2, mode='nearest')
+        x = self.conv1(x)
+        x = self.norm2(x, s)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x, s):
+        out = self._residual(x, s)
+        out = (out + self._shortcut(x)) / self.equal_var
+        return out
+
+
+class Generator(nn.Module):
+    def __init__(
+                self,
+                **kwargs
+                ):
+        super().__init__()
+
+        id_dim      = kwargs["id_dim"]
+        k_size      = kwargs["g_kernel_size"]
+        res_num     = kwargs["res_num"]
+        in_channel  = kwargs["in_channel"]
+        up_mode     = kwargs["up_mode"]
+        norm        = kwargs["norm"]
+        
+        aggregator  = kwargs["aggregator"]
+        res_mode    = kwargs["res_mode"]
+
+        padding_size= int((k_size -1)/2)
+        padding_type= 'reflect'
+        
+        
+        activation = nn.LeakyReLU(0.2)
+        # activation = nn.ReLU()
+
+        self.from_rgb = nn.Conv2d(3, in_channel, 1, 1, 0)
+        # self.first_layer = nn.Sequential(nn.Conv2d(3, 64, kernel_size=3, padding=1, bias=False),
+        #                         nn.BatchNorm2d(64), activation)
+        ### downsample
+        self.down1 = ResBlk(in_channel, in_channel, normalize=norm, downsample=True)# 256
+
+        self.down2 = ResBlk(in_channel, in_channel*2, normalize=norm, downsample=True)# 128
+                                
+        self.down3 = ResBlk(in_channel*2, in_channel*4,normalize=norm, downsample=True)# 64
+
+        self.down4 = ResBlk(in_channel*4, in_channel*8, normalize=norm, downsample=True)# 32
+
+        self.down5 = ResBlk(in_channel*8, in_channel*8, normalize=norm, downsample=True)# 16
+
+        # self.down6 = ResBlk(in_channel*8, in_channel*8, normalize=True, downsample=True)# 8
+        self.maskhead = nn.Sequential(
+                                nn.UpsamplingNearest2d(scale_factor = 2),
+                                nn.Conv2d(in_channel*8, in_channel, kernel_size=3, stride=1, padding=1,bias=False),
+                                nn.BatchNorm2d(in_channel), # 32
+                                activation,
+                                nn.UpsamplingNearest2d(scale_factor = 2),
+                                nn.Conv2d(in_channel, in_channel//2, kernel_size=3, stride=1, padding=1, bias=False),
+                                nn.BatchNorm2d(in_channel//2), # 64
+                                activation,
+                                nn.UpsamplingNearest2d(scale_factor = 2),
+                                nn.Conv2d(in_channel//2, 1, kernel_size=3, stride=1, padding=1, bias=False),
+                                nn.Sigmoid()
+        )
+        
+        ### resnet blocks
+        BN = []
+        for i in range(res_num):
+            BN += [
+                AdainResBlk(in_channel*8, in_channel*8, style_dim=id_dim, upsample=False)]
+        self.BottleNeck = nn.Sequential(*BN)
+
+        # self.up6 = AdainResBlk(in_channel*8, in_channel*8, style_dim=id_dim, upsample=True) # 16
+        
+        self.up5 = AdainResBlk(in_channel*8, in_channel*8, style_dim=id_dim, upsample=True) # 32
+
+        self.up4 = AdainResBlk(in_channel*8, in_channel*4, style_dim=id_dim, upsample=True) # 64
+
+        self.up3 = AdainResBlk(in_channel*4, in_channel*2, style_dim=id_dim, upsample=True) # 128
+        
+        self.up2 = AdainResBlk(in_channel*2, in_channel, style_dim=id_dim, upsample=True)
+
+        self.up1 = AdainResBlk(in_channel, in_channel, style_dim=id_dim, upsample=True)
+        # ResUpBlk(in_channel, in_channel, normalize="in") # 512
+
+        if norm.lower() == "in":
+            norm_out = nn.InstanceNorm2d(in_channel, affine=True)
+        elif norm.lower() == "bn":
+            norm_out = nn.BatchNorm2d(in_channel)
+        
+        self.to_rgb = nn.Sequential(
+            norm_out,
+            activation,
+            nn.Conv2d(in_channel, 3, 3, 1, 1))
+
+        # self.last_layer = nn.Sequential(nn.ReflectionPad2d(3),
+        #             nn.Conv2d(64, 3, kernel_size=7, padding=0))
+
+
+    #     self.__weights_init__()
+
+    # def __weights_init__(self):
+    #     for layer in self.encoder:
+    #         if isinstance(layer,nn.Conv2d):
+    #             nn.init.xavier_uniform_(layer.weight)
+
+    #     for layer in self.encoder2:
+    #         if isinstance(layer,nn.Conv2d):
+    #             nn.init.xavier_uniform_(layer.weight)
+
+    def forward(self, img, id):
+        res = self.from_rgb(img)
+        res = self.down1(res)
+        skip = self.down2(res)
+        res = self.down3(skip)
+        res = self.down4(res)
+        res = self.down5(res)
+        mask= self.maskhead(res)
+        for i in range(len(self.BottleNeck)):
+            res = self.BottleNeck[i](res, id)
+        res = self.up5(res,id)
+        res = self.up4(res,id)
+        res = self.up3(res,id)
+        res = (1-mask) * skip + mask * res
+        res = self.up2(res,id) #  + skip
+        res = self.up1(res,id)
+        res = self.to_rgb(res)
+        return res, mask

components/Generator_maskhead_config1.py

@@ -0,0 +1,279 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+#############################################################
+# File: Generator_Invobn_config1.py
+# Created Date: Saturday February 26th 2022
+# Author: Chen Xuanhong
+# Email: chenxuanhongzju@outlook.com
+# Last Modified:  Wednesday, 6th April 2022 8:38:50 pm
+# Modified By: Chen Xuanhong
+# Copyright (c) 2022 Shanghai Jiao Tong University
+#############################################################
+
+import os
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+import math
+
+
+
+class ResBlk(nn.Module):
+    def __init__(self, dim_in, dim_out, actv=nn.LeakyReLU(0.2),
+                 normalize="bn", downsample=False):
+        super().__init__()
+        self.actv = actv
+        self.normalize = normalize
+        self.downsample = downsample
+        self.learned_sc = dim_in != dim_out
+        self.equal_var  = math.sqrt(2)
+        self._build_weights(dim_in, dim_out)
+
+    def _build_weights(self, dim_in, dim_out):
+        self.conv1 = nn.Conv2d(dim_in, dim_in, 3, 1, 1)
+        self.conv2 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        if self.normalize.lower() == "in":
+            self.norm1 = nn.InstanceNorm2d(dim_in, affine=True)
+            self.norm2 = nn.InstanceNorm2d(dim_in, affine=True)
+        elif self.normalize.lower() == "bn":
+            self.norm1 = nn.BatchNorm2d(dim_in)
+            self.norm2 = nn.BatchNorm2d(dim_in)
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        if self.downsample:
+            x = F.avg_pool2d(x, 2)
+        return x
+
+    def _residual(self, x):
+        if self.normalize:
+            x = self.norm1(x)
+        x = self.actv(x)
+        x = self.conv1(x)
+        if self.downsample:
+            x = F.avg_pool2d(x, 2)
+        if self.normalize:
+            x = self.norm2(x)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x):
+        x = self._shortcut(x) + self._residual(x)
+        return x /self.equal_var   # unit variance
+
+class AdaIN(nn.Module):
+    def __init__(self, style_dim, num_features):
+        super().__init__()
+        self.norm = nn.InstanceNorm2d(num_features, affine=False)
+        self.fc = nn.Linear(style_dim, num_features*2)
+
+    def forward(self, x, s):
+        h = self.fc(s)
+        h = h.view(h.size(0), h.size(1), 1, 1)
+        gamma, beta = torch.chunk(h, chunks=2, dim=1)
+        return (1 + gamma) * self.norm(x) + beta
+
+class ResUpBlk(nn.Module):
+    def __init__(self, dim_in, dim_out,actv=nn.LeakyReLU(0.2),normalize="in"):
+        super().__init__()
+        self.actv = actv
+        self.normalize = normalize
+        self.learned_sc = dim_in != dim_out
+        self.equal_var = math.sqrt(2)
+        self._build_weights(dim_in, dim_out)
+
+    def _build_weights(self, dim_in, dim_out):
+        self.conv1 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        self.conv2 = nn.Conv2d(dim_out, dim_out, 3, 1, 1)
+        if self.normalize.lower() == "in":
+            self.norm1 = nn.InstanceNorm2d(dim_in, affine=True)
+            self.norm2 = nn.InstanceNorm2d(dim_out, affine=True)
+        elif self.normalize.lower() == "bn":
+            self.norm1 = nn.BatchNorm2d(dim_in)
+            self.norm2 = nn.BatchNorm2d(dim_out)
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        x = F.interpolate(x, scale_factor=2, mode='nearest')
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        return x
+
+    def _residual(self, x):
+        x = self.norm1(x)
+        x = self.actv(x)
+        x = F.interpolate(x, scale_factor=2, mode='nearest')
+        x = self.conv1(x)
+        x = self.norm2(x)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x):
+        out = self._residual(x)
+        out = (out + self._shortcut(x)) / self.equal_var
+        return out
+
+class AdainResBlk(nn.Module):
+    def __init__(self, dim_in, dim_out, style_dim=512,
+                 actv=nn.LeakyReLU(0.2), upsample=False):
+        super().__init__()
+        self.actv = actv
+        self.upsample = upsample
+        self.learned_sc = dim_in != dim_out
+        self.equal_var = math.sqrt(2)
+        self._build_weights(dim_in, dim_out, style_dim)
+
+    def _build_weights(self, dim_in, dim_out, style_dim=64):
+        self.conv1 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        self.conv2 = nn.Conv2d(dim_out, dim_out, 3, 1, 1)
+        self.norm1 = AdaIN(style_dim, dim_in)
+        self.norm2 = AdaIN(style_dim, dim_out)
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        if self.upsample:
+            x = F.interpolate(x, scale_factor=2, mode='nearest')
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        return x
+
+    def _residual(self, x, s):
+        x = self.norm1(x, s)
+        x = self.actv(x)
+        if self.upsample:
+            x = F.interpolate(x, scale_factor=2, mode='nearest')
+        x = self.conv1(x)
+        x = self.norm2(x, s)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x, s):
+        out = self._residual(x, s)
+        out = (out + self._shortcut(x)) / self.equal_var
+        return out
+
+
+class Generator(nn.Module):
+    def __init__(
+                self,
+                **kwargs
+                ):
+        super().__init__()
+
+        id_dim      = kwargs["id_dim"]
+        k_size      = kwargs["g_kernel_size"]
+        res_num     = kwargs["res_num"]
+        in_channel  = kwargs["in_channel"]
+        up_mode     = kwargs["up_mode"]
+        norm        = kwargs["norm"]
+        
+        aggregator  = kwargs["aggregator"]
+        res_mode    = kwargs["res_mode"]
+
+        padding_size= int((k_size -1)/2)
+        padding_type= 'reflect'
+        
+        
+        activation = nn.LeakyReLU(0.2)
+        # activation = nn.ReLU()
+
+        self.from_rgb = nn.Conv2d(3, in_channel, 1, 1, 0)
+        # self.first_layer = nn.Sequential(nn.Conv2d(3, 64, kernel_size=3, padding=1, bias=False),
+        #                         nn.BatchNorm2d(64), activation)
+        ### downsample
+        self.down1 = ResBlk(in_channel, in_channel, normalize=norm, downsample=True)# 256
+
+        self.down2 = ResBlk(in_channel, in_channel*2, normalize=norm, downsample=True)# 128
+
+        self.sigma = ResBlk(in_channel*2,in_channel*2)
+                                
+        self.down3 = ResBlk(in_channel*2, in_channel*4,normalize=norm, downsample=True)# 64
+
+        self.down4 = ResBlk(in_channel*4, in_channel*8, normalize=norm, downsample=True)# 32
+
+        self.down5 = ResBlk(in_channel*8, in_channel*8, normalize=norm, downsample=True)# 16
+
+        # self.down6 = ResBlk(in_channel*8, in_channel*8, normalize=True, downsample=True)# 8
+        
+        
+        ### resnet blocks
+        BN = []
+        for i in range(res_num):
+            BN += [
+                AdainResBlk(in_channel*8, in_channel*8, style_dim=id_dim, upsample=False)]
+        self.BottleNeck = nn.Sequential(*BN)
+
+        # self.up6 = AdainResBlk(in_channel*8, in_channel*8, style_dim=id_dim, upsample=True) # 16
+        
+        self.up5 = AdainResBlk(in_channel*8, in_channel*8, style_dim=id_dim, upsample=True) # 32
+
+        self.up4 = AdainResBlk(in_channel*8, in_channel*4, style_dim=id_dim, upsample=True) # 64
+
+        self.up3 = AdainResBlk(in_channel*4, in_channel*2, style_dim=id_dim, upsample=True) # 128
+
+        self.maskhead = nn.Sequential(
+                                nn.Conv2d(in_channel*2, in_channel, kernel_size=3, stride=1, padding=1, bias=False),
+                                nn.BatchNorm2d(in_channel), # 64
+                                activation,
+                                nn.Conv2d(in_channel, 1, kernel_size=3, stride=1, padding=1),
+                                nn.Sigmoid()
+        )
+        
+        self.up2 = AdainResBlk(in_channel*2, in_channel, style_dim=id_dim, upsample=True)
+
+        self.up1 = AdainResBlk(in_channel, in_channel, style_dim=id_dim, upsample=True)
+        # ResUpBlk(in_channel, in_channel, normalize="in") # 512
+
+        if norm.lower() == "in":
+            norm_out = nn.InstanceNorm2d(in_channel, affine=True)
+        elif norm.lower() == "bn":
+            norm_out = nn.BatchNorm2d(in_channel)
+        
+        self.to_rgb = nn.Sequential(
+            norm_out,
+            activation,
+            nn.Conv2d(in_channel, 3, 3, 1, 1))
+
+        # self.last_layer = nn.Sequential(nn.ReflectionPad2d(3),
+        #             nn.Conv2d(64, 3, kernel_size=7, padding=0))
+
+
+    #     self.__weights_init__()
+
+    # def __weights_init__(self):
+    #     for layer in self.encoder:
+    #         if isinstance(layer,nn.Conv2d):
+    #             nn.init.xavier_uniform_(layer.weight)
+
+    #     for layer in self.encoder2:
+    #         if isinstance(layer,nn.Conv2d):
+    #             nn.init.xavier_uniform_(layer.weight)
+
+    def forward(self, img, id):
+        res = self.from_rgb(img)
+        res = self.down1(res)
+        skip = self.down2(res)
+        res = self.down3(skip)
+        res = self.down4(res)
+        res = self.down5(res)
+        
+        for i in range(len(self.BottleNeck)):
+            res = self.BottleNeck[i](res, id)
+        res = self.up5(res,id)
+        res = self.up4(res,id)
+        res = self.up3(res,id)
+        mask= self.maskhead(res)
+        res = (1-mask) * self.sigma(skip) + mask * res
+        res = self.up2(res,id) #  + skip
+        res = self.up1(res,id)
+        res = self.to_rgb(res)
+        return res, mask

components/Generator_maskhead_config2.py

@@ -0,0 +1,283 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+#############################################################
+# File: Generator_Invobn_config1.py
+# Created Date: Saturday February 26th 2022
+# Author: Chen Xuanhong
+# Email: chenxuanhongzju@outlook.com
+# Last Modified:  Wednesday, 13th April 2022 3:12:53 pm
+# Modified By: Chen Xuanhong
+# Copyright (c) 2022 Shanghai Jiao Tong University
+#############################################################
+
+import os
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+import math
+
+
+
+class ResBlk(nn.Module):
+    def __init__(self, dim_in, dim_out, actv=nn.LeakyReLU(0.2),
+                 normalize="bn", downsample=False):
+        super().__init__()
+        self.actv = actv
+        self.normalize = normalize
+        self.downsample = downsample
+        self.learned_sc = dim_in != dim_out
+        self.equal_var  = math.sqrt(2)
+        self._build_weights(dim_in, dim_out)
+
+    def _build_weights(self, dim_in, dim_out):
+        self.conv1 = nn.Conv2d(dim_in, dim_in, 3, 1, 1)
+        self.conv2 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        if self.normalize.lower() == "in":
+            self.norm1 = nn.InstanceNorm2d(dim_in, affine=True)
+            self.norm2 = nn.InstanceNorm2d(dim_in, affine=True)
+        elif self.normalize.lower() == "bn":
+            self.norm1 = nn.BatchNorm2d(dim_in)
+            self.norm2 = nn.BatchNorm2d(dim_in)
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        if self.downsample:
+            x = F.avg_pool2d(x, 2)
+        return x
+
+    def _residual(self, x):
+        if self.normalize:
+            x = self.norm1(x)
+        x = self.actv(x)
+        x = self.conv1(x)
+        if self.downsample:
+            x = F.avg_pool2d(x, 2)
+        if self.normalize:
+            x = self.norm2(x)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x):
+        x = self._shortcut(x) + self._residual(x)
+        return x /self.equal_var   # unit variance
+
+class AdaIN(nn.Module):
+    def __init__(self, style_dim, num_features):
+        super().__init__()
+        self.norm = nn.InstanceNorm2d(num_features, affine=False)
+        self.fc = nn.Linear(style_dim, num_features*2)
+
+    def forward(self, x, s):
+        h = self.fc(s)
+        h = h.view(h.size(0), h.size(1), 1, 1)
+        gamma, beta = torch.chunk(h, chunks=2, dim=1)
+        return (1 + gamma) * self.norm(x) + beta
+
+class ResUpBlk(nn.Module):
+    def __init__(self, dim_in, dim_out,actv=nn.LeakyReLU(0.2),normalize="in"):
+        super().__init__()
+        self.actv = actv
+        self.normalize = normalize
+        self.learned_sc = dim_in != dim_out
+        self.equal_var = math.sqrt(2)
+        self._build_weights(dim_in, dim_out)
+
+    def _build_weights(self, dim_in, dim_out):
+        self.conv1 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        self.conv2 = nn.Conv2d(dim_out, dim_out, 3, 1, 1)
+        if self.normalize.lower() == "in":
+            self.norm1 = nn.InstanceNorm2d(dim_in, affine=True)
+            self.norm2 = nn.InstanceNorm2d(dim_out, affine=True)
+        elif self.normalize.lower() == "bn":
+            self.norm1 = nn.BatchNorm2d(dim_in)
+            self.norm2 = nn.BatchNorm2d(dim_out)
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        x = F.interpolate(x, scale_factor=2, mode='nearest')
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        return x
+
+    def _residual(self, x):
+        x = self.norm1(x)
+        x = self.actv(x)
+        x = F.interpolate(x, scale_factor=2, mode='nearest')
+        x = self.conv1(x)
+        x = self.norm2(x)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x):
+        out = self._residual(x)
+        out = (out + self._shortcut(x)) / self.equal_var
+        return out
+
+class AdainResBlk(nn.Module):
+    def __init__(self, dim_in, dim_out, style_dim=512,
+                 actv=nn.LeakyReLU(0.2), upsample=False):
+        super().__init__()
+        self.actv = actv
+        self.upsample = upsample
+        self.learned_sc = dim_in != dim_out
+        self.equal_var = math.sqrt(2)
+        self._build_weights(dim_in, dim_out, style_dim)
+
+    def _build_weights(self, dim_in, dim_out, style_dim=64):
+        self.conv1 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        self.conv2 = nn.Conv2d(dim_out, dim_out, 3, 1, 1)
+        self.norm1 = AdaIN(style_dim, dim_in)
+        self.norm2 = AdaIN(style_dim, dim_out)
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        if self.upsample:
+            x = F.interpolate(x, scale_factor=2, mode='nearest')
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        return x
+
+    def _residual(self, x, s):
+        x = self.norm1(x, s)
+        x = self.actv(x)
+        if self.upsample:
+            x = F.interpolate(x, scale_factor=2, mode='nearest')
+        x = self.conv1(x)
+        x = self.norm2(x, s)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x, s):
+        out = self._residual(x, s)
+        out = (out + self._shortcut(x)) / self.equal_var
+        return out
+
+
+class Generator(nn.Module):
+    def __init__(
+                self,
+                **kwargs
+                ):
+        super().__init__()
+
+        id_dim      = kwargs["id_dim"]
+        k_size      = kwargs["g_kernel_size"]
+        res_num     = kwargs["res_num"]
+        in_channel  = kwargs["in_channel"]
+        up_mode     = kwargs["up_mode"]
+        norm        = kwargs["norm"].lower()
+        
+        aggregator  = kwargs["aggregator"]
+        res_mode    = kwargs["res_mode"]
+
+        padding_size= int((k_size -1)/2)
+        padding_type= 'reflect'
+        
+        
+        activation = nn.LeakyReLU(0.2)
+        # activation = nn.ReLU()
+
+        self.from_rgb = nn.Conv2d(3, in_channel, 1, 1, 0)
+        # self.first_layer = nn.Sequential(nn.Conv2d(3, 64, kernel_size=3, padding=1, bias=False),
+        #                         nn.BatchNorm2d(64), activation)
+        ### downsample
+        self.down1 = ResBlk(in_channel, in_channel, normalize=norm, downsample=True)# 256
+
+        self.down2 = ResBlk(in_channel, in_channel*2, normalize=norm, downsample=True)# 128
+
+        # self.sigma = ResBlk(in_channel*2,in_channel*2)
+                                
+        self.down3 = ResBlk(in_channel*2, in_channel*4,normalize=norm, downsample=True)# 64
+
+        self.down4 = ResBlk(in_channel*4, in_channel*8, normalize=norm, downsample=True)# 32
+
+        self.down5 = ResBlk(in_channel*8, in_channel*8, normalize=norm, downsample=True)# 16
+
+        # self.down6 = ResBlk(in_channel*8, in_channel*8, normalize=True, downsample=True)# 8
+        
+        
+        ### resnet blocks
+        BN = []
+        for i in range(res_num):
+            BN += [
+                AdainResBlk(in_channel*8, in_channel*8, style_dim=id_dim, upsample=False)]
+        self.BottleNeck = nn.Sequential(*BN)
+
+        # self.up6 = AdainResBlk(in_channel*8, in_channel*8, style_dim=id_dim, upsample=True) # 16
+        
+        self.up5 = AdainResBlk(in_channel*8, in_channel*8, style_dim=id_dim, upsample=True) # 32
+
+        self.up4 = AdainResBlk(in_channel*8, in_channel*4, style_dim=id_dim, upsample=True) # 64
+
+        self.up3 = AdainResBlk(in_channel*4, in_channel*2, style_dim=id_dim, upsample=True) # 128
+
+        self.maskhead = nn.Sequential(
+                                nn.Conv2d(in_channel*2, in_channel, kernel_size=3, stride=1, padding=1, bias=False),
+                                nn.BatchNorm2d(in_channel), # 64
+                                activation,
+                                nn.Conv2d(in_channel, 1, kernel_size=3, stride=1, padding=1),
+                                nn.Sigmoid()
+        )
+        
+        # self.up2 = AdainResBlk(in_channel*2, in_channel, style_dim=id_dim, upsample=True)
+        # self.up2 = AdainResBlk(in_channel*2, in_channel, style_dim=id_dim, upsample=True)
+        self.up2 = ResUpBlk(in_channel*2, in_channel, normalize="bn")
+
+        # self.up1 = AdainResBlk(in_channel, in_channel, style_dim=id_dim, upsample=True)
+        self.up1 = ResUpBlk(in_channel, in_channel, normalize="bn")
+        # ResUpBlk(in_channel, in_channel, normalize="in") # 512
+
+        if norm.lower() == "in":
+            norm_out = nn.InstanceNorm2d(in_channel, affine=True)
+        elif norm.lower() == "bn":
+            norm_out = nn.BatchNorm2d(in_channel)
+        
+        self.to_rgb = nn.Sequential(
+            norm_out,
+            activation,
+            nn.Conv2d(in_channel, 3, 3, 1, 1))
+
+        # self.last_layer = nn.Sequential(nn.ReflectionPad2d(3),
+        #             nn.Conv2d(64, 3, kernel_size=7, padding=0))
+
+
+    #     self.__weights_init__()
+
+    # def __weights_init__(self):
+    #     for layer in self.encoder:
+    #         if isinstance(layer,nn.Conv2d):
+    #             nn.init.xavier_uniform_(layer.weight)
+
+    #     for layer in self.encoder2:
+    #         if isinstance(layer,nn.Conv2d):
+    #             nn.init.xavier_uniform_(layer.weight)
+
+    def forward(self, img, id):
+        res = self.from_rgb(img)
+        res = self.down1(res)
+        skip = self.down2(res)
+        res = self.down3(skip)
+        res = self.down4(res)
+        res = self.down5(res)
+        
+        for i in range(len(self.BottleNeck)):
+            res = self.BottleNeck[i](res, id)
+        res = self.up5(res,id)
+        res = self.up4(res,id)
+        res = self.up3(res,id)
+        mask= self.maskhead(res)
+        # res = (1-mask) * self.sigma(skip) + mask * res
+        res = (1-mask) * skip + mask * res
+        res = self.up2(res) #  + skip
+        res = self.up1(res)
+        res = self.to_rgb(res)
+        return res, mask

components/Generator_starganv2.py

@@ -0,0 +1,297 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+#############################################################
+# File: Generator_Invobn_config1.py
+# Created Date: Saturday February 26th 2022
+# Author: Chen Xuanhong
+# Email: chenxuanhongzju@outlook.com
+# Last Modified:  Wednesday, 13th April 2022 6:30:26 pm
+# Modified By: Chen Xuanhong
+# Copyright (c) 2022 Shanghai Jiao Tong University
+#############################################################
+
+import os
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+import math
+
+
+
+class ResBlk(nn.Module):
+    def __init__(self, dim_in, dim_out, actv=nn.LeakyReLU(0.2),
+                 normalize="in", downsample=False):
+        super().__init__()
+        self.actv = actv
+        self.normalize = normalize
+        self.downsample = downsample
+        self.learned_sc = dim_in != dim_out
+        self.equal_var  = math.sqrt(2)
+        self._build_weights(dim_in, dim_out)
+
+    def _build_weights(self, dim_in, dim_out):
+        self.conv1 = nn.Conv2d(dim_in, dim_in, 3, 1, 1)
+        self.conv2 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        if self.normalize.lower() == "in":
+            self.norm1 = nn.InstanceNorm2d(dim_in, affine=True)
+            self.norm2 = nn.InstanceNorm2d(dim_in, affine=True)
+        elif self.normalize.lower() == "bn":
+            self.norm1 = nn.BatchNorm2d(dim_in)
+            self.norm2 = nn.BatchNorm2d(dim_in)
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        if self.downsample:
+            x = F.avg_pool2d(x, 2)
+        return x
+
+    def _residual(self, x):
+        if self.normalize:
+            x = self.norm1(x)
+        x = self.actv(x)
+        x = self.conv1(x)
+        if self.downsample:
+            x = F.avg_pool2d(x, 2)
+        if self.normalize:
+            x = self.norm2(x)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x):
+        x = self._shortcut(x) + self._residual(x)
+        return x /self.equal_var   # unit variance
+
+class AdaIN(nn.Module):
+    def __init__(self, style_dim, num_features):
+        super().__init__()
+        self.norm = nn.InstanceNorm2d(num_features, affine=False)
+        self.fc = nn.Linear(style_dim, num_features*2)
+
+    def forward(self, x, s):
+        h = self.fc(s)
+        h = h.view(h.size(0), h.size(1), 1, 1)
+        gamma, beta = torch.chunk(h, chunks=2, dim=1)
+        return (1 + gamma) * self.norm(x) + beta
+
+class ResUpBlk(nn.Module):
+    def __init__(self, dim_in, dim_out,actv=nn.LeakyReLU(0.2),normalize="in"):
+        super().__init__()
+        self.actv = actv
+        self.normalize = normalize
+        self.learned_sc = dim_in != dim_out
+        self.equal_var = math.sqrt(2)
+        self._build_weights(dim_in, dim_out)
+
+    def _build_weights(self, dim_in, dim_out):
+        self.conv1 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        self.conv2 = nn.Conv2d(dim_out, dim_out, 3, 1, 1)
+        if self.normalize.lower() == "in":
+            self.norm1 = nn.InstanceNorm2d(dim_in, affine=True)
+            self.norm2 = nn.InstanceNorm2d(dim_out, affine=True)
+        elif self.normalize.lower() == "bn":
+            self.norm1 = nn.BatchNorm2d(dim_in)
+            self.norm2 = nn.BatchNorm2d(dim_out)
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        x = F.interpolate(x, scale_factor=2, mode='nearest')
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        return x
+
+    def _residual(self, x):
+        x = self.norm1(x)
+        x = self.actv(x)
+        x = F.interpolate(x, scale_factor=2, mode='nearest')
+        x = self.conv1(x)
+        x = self.norm2(x)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x):
+        out = self._residual(x)
+        out = (out + self._shortcut(x)) / self.equal_var
+        return out
+
+class AdainResBlk(nn.Module):
+    def __init__(self, dim_in, dim_out, style_dim=512,
+                 actv=nn.LeakyReLU(0.2), upsample=False):
+        super().__init__()
+        self.actv = actv
+        self.upsample = upsample
+        self.learned_sc = dim_in != dim_out
+        self.equal_var = math.sqrt(2)
+        self._build_weights(dim_in, dim_out, style_dim)
+
+    def _build_weights(self, dim_in, dim_out, style_dim=64):
+        self.conv1 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        self.conv2 = nn.Conv2d(dim_out, dim_out, 3, 1, 1)
+        self.norm1 = AdaIN(style_dim, dim_in)
+        self.norm2 = AdaIN(style_dim, dim_out)
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        if self.upsample:
+            x = F.interpolate(x, scale_factor=2, mode='nearest')
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        return x
+
+    def _residual(self, x, s):
+        x = self.norm1(x, s)
+        x = self.actv(x)
+        if self.upsample:
+            x = F.interpolate(x, scale_factor=2, mode='nearest')
+        x = self.conv1(x)
+        x = self.norm2(x, s)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x, s):
+        out = self._residual(x, s)
+        out = (out + self._shortcut(x)) / self.equal_var
+        return out
+
+
+class Generator(nn.Module):
+    def __init__(
+                self,
+                **kwargs
+                ):
+        super().__init__()
+
+        id_dim      = kwargs["id_dim"]
+        k_size      = kwargs["g_kernel_size"]
+        res_num     = kwargs["res_num"]
+        in_channel  = kwargs["in_channel"]
+        up_mode     = kwargs["up_mode"]
+        norm        = kwargs["norm"].lower()
+        
+        aggregator  = kwargs["aggregator"]
+        res_mode    = kwargs["res_mode"]
+
+        padding_size= int((k_size -1)/2)
+        padding_type= 'reflect'
+
+        if norm.lower() == "in":
+            norm_out = nn.InstanceNorm2d(in_channel, affine=True)
+            norm_mask= nn.InstanceNorm2d
+        elif norm.lower() == "bn":
+            norm_out = nn.BatchNorm2d(in_channel)
+            norm_mask = nn.BatchNorm2d
+        
+        
+        activation = nn.LeakyReLU(0.2)
+        # activation = nn.ReLU()
+
+        self.from_rgb = nn.Conv2d(3, in_channel, 1, 1, 0)
+        # self.first_layer = nn.Sequential(nn.Conv2d(3, 64, kernel_size=3, padding=1, bias=False),
+        #                         nn.BatchNorm2d(64), activation)
+        ### downsample
+        self.down1 = ResBlk(in_channel, in_channel, normalize=norm, downsample=True)# 256
+
+        self.down2 = ResBlk(in_channel, in_channel*2, normalize=norm, downsample=True)# 128
+                                
+        self.down3 = ResBlk(in_channel*2, in_channel*4,normalize=norm, downsample=True)# 64
+
+        self.down4 = ResBlk(in_channel*4, in_channel*8, normalize=norm, downsample=True)# 32
+
+        self.down5 = ResBlk(in_channel*8, in_channel*8, normalize=norm, downsample=True)# 16
+
+        # self.down6 = ResBlk(in_channel*8, in_channel*8, normalize=True, downsample=True)# 8
+        
+        
+        ### resnet blocks
+        BN = []
+        for i in range(res_num):
+            BN += [
+                AdainResBlk(in_channel*8, in_channel*8, style_dim=id_dim, upsample=False)]
+        self.BottleNeck = nn.Sequential(*BN)
+
+        # self.up6 = AdainResBlk(in_channel*8, in_channel*8, style_dim=id_dim, upsample=True) # 16
+        
+        self.up5 = AdainResBlk(in_channel*8, in_channel*8, style_dim=id_dim, upsample=True) # 32
+
+        self.up4 = AdainResBlk(in_channel*8, in_channel*4, style_dim=id_dim, upsample=True) # 64
+
+        self.up3 = AdainResBlk(in_channel*4, in_channel*2, style_dim=id_dim, upsample=True) # 128
+
+        # self.maskhead = nn.Sequential(
+        #                         nn.Conv2d(in_channel*2, in_channel//2, kernel_size=3, stride=1, padding=1, bias=False),
+        #                         norm_mask, # 64
+        #                         activation,
+        #                         nn.Conv2d(in_channel//2, 1, kernel_size=3, stride=1, padding=1),
+        #                         nn.Sigmoid())
+        self.maskhead = nn.Sequential(
+                                nn.UpsamplingNearest2d(scale_factor = 2),
+                                nn.Conv2d(in_channel*8, in_channel, kernel_size=3, stride=1, padding=1,bias=False),
+                                norm_mask(in_channel, affine=True), # 32
+                                activation,
+                                nn.UpsamplingNearest2d(scale_factor = 2),
+                                nn.Conv2d(in_channel, in_channel//4, kernel_size=3, stride=1, padding=1, bias=False),
+                                norm_mask(in_channel//4, affine=True), # 64
+                                activation,
+                                nn.UpsamplingNearest2d(scale_factor = 2),
+                                nn.Conv2d(in_channel//4, 1, kernel_size=3, stride=1, padding=1),
+                                nn.Sigmoid()
+        )
+        
+        # self.up2 = AdainResBlk(in_channel*2, in_channel, style_dim=id_dim, upsample=True)
+        # self.up2 = AdainResBlk(in_channel*2, in_channel, style_dim=id_dim, upsample=True)
+        self.up2 = ResUpBlk(in_channel*2, in_channel, normalize=norm)
+
+        # self.up1 = AdainResBlk(in_channel, in_channel, style_dim=id_dim, upsample=True)
+        self.up1 = ResUpBlk(in_channel, in_channel, normalize=norm)
+        # ResUpBlk(in_channel, in_channel, normalize="in") # 512
+
+        
+        
+        self.to_rgb = nn.Sequential(
+            norm_out,
+            activation,
+            nn.Conv2d(in_channel, 3, 3, 1, 1))
+
+        # self.last_layer = nn.Sequential(nn.ReflectionPad2d(3),
+        #             nn.Conv2d(64, 3, kernel_size=7, padding=0))
+
+
+    #     self.__weights_init__()
+
+    # def __weights_init__(self):
+    #     for layer in self.encoder:
+    #         if isinstance(layer,nn.Conv2d):
+    #             nn.init.xavier_uniform_(layer.weight)
+
+    #     for layer in self.encoder2:
+    #         if isinstance(layer,nn.Conv2d):
+    #             nn.init.xavier_uniform_(layer.weight)
+
+    def forward(self, img, id):
+        res = self.from_rgb(img)
+        res = self.down1(res)
+        skip = self.down2(res)
+        res = self.down3(skip)
+        res = self.down4(res)
+        res = self.down5(res)
+        mask= self.maskhead(res)
+        for i in range(len(self.BottleNeck)):
+            res = self.BottleNeck[i](res, id)
+        res = self.up5(res,id)
+        res = self.up4(res,id)
+        res = self.up3(res,id)
+        
+        # res = (1-mask) * self.sigma(skip) + mask * res
+        res = (1-mask) * skip + mask * res
+        res = self.up2(res) #  + skip
+        res = self.up1(res)
+        res = self.to_rgb(res)
+        return res, mask

components/LSTU.py

@@ -5,43 +5,115 @@
 # Created Date: Sunday January 16th 2022
 # Author: Chen Xuanhong
 # Email: chenxuanhongzju@outlook.com
-# Last Modified:  Sunday, 13th February 2022 2:03:21 am
+# Last Modified:  Monday, 28th March 2022 11:47:55 pm
 # Modified By: Chen Xuanhong
 # Copyright (c) 2022 Shanghai Jiao Tong University
 #############################################################
 
+import math
 import  torch
 from    torch import nn
 
+import torch.nn.functional as F
+
+# class LSTU(nn.Module):
+#     def __init__(
+#                 self,
+#                 in_channel,
+#                 out_channel,
+#                 latent_channel,
+#                 scale = 4
+#                 ):
+#         super().__init__()
+#         sig              = nn.Sigmoid()
+#         self.relu        = nn.ReLU(True)
+
+#         self.up_sample   = nn.Sequential(nn.Conv2d(latent_channel, out_channel/4, kernel_size=3, stride=1, padding=1, bias=False),
+#                                 nn.BatchNorm2d(out_channel/4),
+#                                 self.relu,
+#                                 nn.Conv2d(latent_channel/4, out_channel, kernel_size=3, stride=1, padding=1),
+#                                 )
+
+#         self.forget_gate = nn.Sequential(nn.Conv2d(in_channel, out_channel, kernel_size=3, padding=1, bias=False),
+#                                 nn.BatchNorm2d(out_channel), sig)
+        
+#         self.reset_gate  = nn.Sequential(nn.Conv2d(in_channel, out_channel, kernel_size=3, padding=1, bias=False),
+#                                 nn.BatchNorm2d(out_channel), sig)
+        
+#         self.conv11      = nn.Sequential(nn.Conv2d(out_channel, out_channel, kernel_size=1, bias=True))
+    
+#     def forward(self, encoder_in, bottleneck_in):
+#         h_hat_l_1   = self.up_sample(bottleneck_in)  # upsample and make `channel` identical to `out_channel`
+#         h_bar_l     = self.conv11(h_hat_l_1)
+#         f_l         = self.forget_gate(h_hat_l_1)
+#         r_l         = self.reset_gate (h_hat_l_1)
+#         h_hat_l     = (1-f_l)*h_bar_l + f_l* encoder_in
+#         x_hat_l     = r_l* self.relu(h_hat_l) + (1-r_l)* h_hat_l_1
+#         return x_hat_l
+
+
+class ResBlk(nn.Module):
+    def __init__(self, dim_in, dim_out, actv=nn.LeakyReLU(0.2),
+                 normalize="in", downsample=False):
+        super().__init__()
+        self.actv = actv
+        self.normalize = normalize
+        self.downsample = downsample
+        self.learned_sc = dim_in != dim_out
+        self.equal_var  = math.sqrt(2)
+        self._build_weights(dim_in, dim_out)
+
+    def _build_weights(self, dim_in, dim_out):
+        self.conv1 = nn.Conv2d(dim_in, dim_in, 3, 1, 1)
+        self.conv2 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        if self.normalize.lower() == "in":
+            self.norm1 = nn.InstanceNorm2d(dim_in, affine=True)
+            self.norm2 = nn.InstanceNorm2d(dim_in, affine=True)
+        elif self.normalize.lower() == "bn":
+            self.norm1 = nn.BatchNorm2d(dim_in)
+            self.norm2 = nn.BatchNorm2d(dim_in)
+            
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        if self.downsample:
+            x = F.avg_pool2d(x, 2)
+        return x
+
+    def _residual(self, x):
+        if self.normalize:
+            x = self.norm1(x)
+        x = self.actv(x)
+        x = self.conv1(x)
+        if self.downsample:
+            x = F.avg_pool2d(x, 2)
+        if self.normalize:
+            x = self.norm2(x)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x):
+        x = self._shortcut(x) + self._residual(x)
+        return x /self.equal_var   # unit variance
 
 class LSTU(nn.Module):
     def __init__(
                 self,
                 in_channel,
-                out_channel,
-                latent_channel,
-                scale = 4
+                norm
                 ):
         super().__init__()
-        sig              = nn.Sigmoid()
-        self.relu        = nn.ReLU(True)
+        self.sig   = nn.Sigmoid()
 
-        self.up_sample   = nn.Sequential(nn.ConvTranspose2d(latent_channel, out_channel, kernel_size=4, stride=scale, padding=0, bias=False),
-                                nn.BatchNorm2d(out_channel), sig)
-
-        self.forget_gate = nn.Sequential(nn.Conv2d(in_channel, out_channel, kernel_size=3, padding=1, bias=False),
-                                nn.BatchNorm2d(out_channel), sig)
-        
-        self.reset_gate  = nn.Sequential(nn.Conv2d(in_channel, out_channel, kernel_size=3, padding=1, bias=False),
-                                nn.BatchNorm2d(out_channel), sig)
-        
-        self.conv11      = nn.Sequential(nn.Conv2d(out_channel, out_channel, kernel_size=1, bias=True))
+        self.mask_head  = ResBlk(in_channel, 1, normalize=norm)
+        # self.forget_gate = ResBlk(in_channel,in_channel, normalize=norm)
     
-    def forward(self, encoder_in, bottleneck_in):
-        h_hat_l_1   = self.up_sample(bottleneck_in)  # upsample and make `channel` identical to `out_channel`
-        h_bar_l     = self.conv11(h_hat_l_1)
-        f_l         = self.forget_gate(h_hat_l_1)
-        r_l         = self.reset_gate (h_hat_l_1)
-        h_hat_l     = (1-f_l)*h_bar_l + f_l* encoder_in
-        x_hat_l     = r_l* self.relu(h_hat_l) + (1-r_l)* h_hat_l_1
-        return x_hat_l
+    def forward(self, encoder_in, decoder_in):
+        mask    = self.sig(self.mask_head(decoder_in))  # upsample and make `channel` identical to `out_channel`
+        # enc_feat= self.forget_gate(encoder_in)
+        out     = (1-mask)*encoder_in + mask * decoder_in
+        return out, mask

components/LSTU_Config.py

@@ -0,0 +1,124 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+#############################################################
+# File: Generator.py
+# Created Date: Sunday January 16th 2022
+# Author: Chen Xuanhong
+# Email: chenxuanhongzju@outlook.com
+# Last Modified:  Tuesday, 29th March 2022 12:20:26 pm
+# Modified By: Chen Xuanhong
+# Copyright (c) 2022 Shanghai Jiao Tong University
+#############################################################
+
+import math
+import  torch
+from    torch import nn
+
+import torch.nn.functional as F
+
+# class LSTU(nn.Module):
+#     def __init__(
+#                 self,
+#                 in_channel,
+#                 out_channel,
+#                 latent_channel,
+#                 scale = 4
+#                 ):
+#         super().__init__()
+#         sig              = nn.Sigmoid()
+#         self.relu        = nn.ReLU(True)
+
+#         self.up_sample   = nn.Sequential(nn.Conv2d(latent_channel, out_channel/4, kernel_size=3, stride=1, padding=1, bias=False),
+#                                 nn.BatchNorm2d(out_channel/4),
+#                                 self.relu,
+#                                 nn.Conv2d(latent_channel/4, out_channel, kernel_size=3, stride=1, padding=1),
+#                                 )
+
+#         self.forget_gate = nn.Sequential(nn.Conv2d(in_channel, out_channel, kernel_size=3, padding=1, bias=False),
+#                                 nn.BatchNorm2d(out_channel), sig)
+        
+#         self.reset_gate  = nn.Sequential(nn.Conv2d(in_channel, out_channel, kernel_size=3, padding=1, bias=False),
+#                                 nn.BatchNorm2d(out_channel), sig)
+        
+#         self.conv11      = nn.Sequential(nn.Conv2d(out_channel, out_channel, kernel_size=1, bias=True))
+    
+#     def forward(self, encoder_in, bottleneck_in):
+#         h_hat_l_1   = self.up_sample(bottleneck_in)  # upsample and make `channel` identical to `out_channel`
+#         h_bar_l     = self.conv11(h_hat_l_1)
+#         f_l         = self.forget_gate(h_hat_l_1)
+#         r_l         = self.reset_gate (h_hat_l_1)
+#         h_hat_l     = (1-f_l)*h_bar_l + f_l* encoder_in
+#         x_hat_l     = r_l* self.relu(h_hat_l) + (1-r_l)* h_hat_l_1
+#         return x_hat_l
+
+
+class ResBlk(nn.Module):
+    def __init__(self, dim_in, dim_out, actv=nn.LeakyReLU(0.2),
+                 normalize="in", downsample=False):
+        super().__init__()
+        self.actv = actv
+        self.normalize = normalize
+        self.downsample = downsample
+        self.learned_sc = dim_in != dim_out
+        self.equal_var  = math.sqrt(2)
+        self._build_weights(dim_in, dim_out)
+
+    def _build_weights(self, dim_in, dim_out):
+        self.conv1 = nn.Conv2d(dim_in, dim_in, 3, 1, 1)
+        self.conv2 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        if self.normalize.lower() == "in":
+            self.norm1 = nn.InstanceNorm2d(dim_in, affine=True)
+            self.norm2 = nn.InstanceNorm2d(dim_in, affine=True)
+        elif self.normalize.lower() == "bn":
+            self.norm1 = nn.BatchNorm2d(dim_in)
+            self.norm2 = nn.BatchNorm2d(dim_in)
+            
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        if self.downsample:
+            x = F.avg_pool2d(x, 2)
+        return x
+
+    def _residual(self, x):
+        if self.normalize:
+            x = self.norm1(x)
+        x = self.actv(x)
+        x = self.conv1(x)
+        if self.downsample:
+            x = F.avg_pool2d(x, 2)
+        if self.normalize:
+            x = self.norm2(x)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x):
+        x = self._shortcut(x) + self._residual(x)
+        return x /self.equal_var   # unit variance
+
+class LSTU(nn.Module):
+    def __init__(
+                self,
+                in_channel,
+                norm
+                ):
+        super().__init__()
+
+        # self.mask_head  = ResBlk(in_channel, 1, normalize=norm)
+        self.mask_head  = nn.Sequential(nn.Conv2d(in_channel, in_channel//2, kernel_size=3, stride=1, padding=1, bias=False),
+                                nn.BatchNorm2d(in_channel//2),
+                                nn.LeakyReLU(0.2),
+                                nn.Conv2d(in_channel//2, 1, kernel_size=3, stride=1, padding=1),
+                                nn.Sigmoid()
+                                )
+        # self.forget_gate = ResBlk(in_channel,in_channel, normalize=norm)
+    
+    def forward(self, encoder_in, decoder_in):
+        mask    = self.mask_head(decoder_in)  # upsample and make `channel` identical to `out_channel`
+        # enc_feat= self.forget_gate(encoder_in)
+        out     = (1-mask)*encoder_in + mask * decoder_in
+        return out, mask

components/ModulatedDWConv.py

@@ -0,0 +1,66 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+#############################################################
+# File: ModulatedDWConv.py
+# Created Date: Monday April 18th 2022
+# Author: Chen Xuanhong
+# Email: chenxuanhongzju@outlook.com
+# Last Modified:  Monday, 18th April 2022 10:33:48 am
+# Modified By: Chen Xuanhong
+# Modified from: https://github.com/bes-dev/MobileStyleGAN.pytorch
+# Copyright (c) 2022 Shanghai Jiao Tong University
+#############################################################
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class ModulatedDWConv2d(nn.Module):
+    def __init__(
+            self,
+            channels_in,
+            channels_out,
+            style_dim,
+            kernel_size,
+            demodulate=True
+    ):
+        super().__init__()
+        # create conv
+        self.weight_dw = nn.Parameter(
+            torch.randn(channels_in, 1, kernel_size, kernel_size)
+        )
+        self.weight_permute = nn.Parameter(
+            torch.randn(channels_out, channels_in, 1, 1)
+        )
+        # create modulation network
+        self.modulation = nn.Linear(style_dim, channels_in, bias=True)
+        self.modulation.bias.data.fill_(1.0)
+        # create demodulation parameters
+        self.demodulate = demodulate
+        if self.demodulate:
+            self.register_buffer("style_inv", torch.randn(1, 1, channels_in, 1, 1))
+        # some service staff
+        self.scale = 1.0 / math.sqrt(channels_in * kernel_size ** 2)
+        self.padding = kernel_size // 2
+
+    def forward(self, x, style):
+        modulation = self.get_modulation(style)
+        x = modulation * x
+        x = F.conv2d(x, self.weight_dw, padding=self.padding, groups=x.size(1))
+        x = F.conv2d(x, self.weight_permute)
+        if self.demodulate:
+            demodulation = self.get_demodulation(style)
+            x = demodulation * x
+        return x
+
+    def get_modulation(self, style):
+        style = self.modulation(style).view(style.size(0), -1, 1, 1)
+        modulation = self.scale * style
+        return modulation
+
+    def get_demodulation(self, style):
+        w = (self.weight_dw.transpose(0, 1) * self.weight_permute).unsqueeze(0)
+        norm = torch.rsqrt((self.scale * self.style_inv * w).pow(2).sum([2, 3, 4]) + 1e-8)
+        demodulation = norm
+        return demodulation.view(*demodulation.size(), 1, 1)

components/Nonstau_Discriminator.py

@@ -0,0 +1,96 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+#############################################################
+# File: Nonstau_Discriminator.py
+# Created Date: Monday March 28th 2022
+# Author: Chen Xuanhong
+# Email: chenxuanhongzju@outlook.com
+# Last Modified:  Monday, 28th March 2022 10:03:56 pm
+# Modified By: Chen Xuanhong
+# Copyright (c) 2022 Shanghai Jiao Tong University
+#############################################################
+import math
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class ResBlk(nn.Module):
+    def __init__(self, dim_in, dim_out, actv=nn.LeakyReLU(0.2),
+                 normalize="in", downsample=False):
+        super().__init__()
+        self.actv = actv
+        self.normalize = normalize
+        self.downsample = downsample
+        self.learned_sc = dim_in != dim_out
+        self.var        = math.sqrt(2)
+        self._build_weights(dim_in, dim_out)
+
+    def _build_weights(self, dim_in, dim_out):
+        self.conv1 = nn.Conv2d(dim_in, dim_in, 3, 1, 1)
+        self.conv2 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        if self.normalize.lower() == "in":
+            self.norm1 = nn.InstanceNorm2d(dim_in, affine=True)
+            self.norm2 = nn.InstanceNorm2d(dim_in, affine=True)
+        elif self.normalize.lower() == "bn":
+            self.norm1 = nn.BatchNorm2d(dim_in)
+            self.norm2 = nn.BatchNorm2d(dim_in)
+        elif self.normalize.lower() == "none":
+            self.normalize = False
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        if self.downsample:
+            x = F.avg_pool2d(x, 2)
+        return x
+
+    def _residual(self, x):
+        if self.normalize:
+            x = self.norm1(x)
+        x = self.actv(x)
+        x = self.conv1(x)
+        if self.downsample:
+            x = F.avg_pool2d(x, 2)
+        if self.normalize:
+            x = self.norm2(x)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x):
+        x = self._shortcut(x) + self._residual(x)
+        return x / self.var  # unit variance
+
+class Discriminator(torch.nn.Module):
+    def __init__(
+        self,
+        **kwargs
+    ):
+        super().__init__()
+        img_size     = kwargs["img_size"]
+        num_domains  = 1
+        max_conv_dim = kwargs["max_conv_dim"]
+        norm         = kwargs["norm"]
+        dim_in = 2**14 // img_size
+        blocks = []
+        blocks += [nn.Conv2d(3, dim_in, 3, 1, 1)]
+
+        repeat_num = int(np.log2(img_size)) - 2
+        for _ in range(repeat_num):
+            dim_out = min(dim_in*2, max_conv_dim)
+            blocks += [ResBlk(dim_in, dim_out, normalize=norm, downsample=True)]
+            dim_in = dim_out
+
+        blocks += [nn.LeakyReLU(0.2)]
+        blocks += [nn.Conv2d(dim_out, dim_out, 4, 1, 0)]
+        blocks += [nn.LeakyReLU(0.2)]
+        blocks += [nn.Conv2d(dim_out, num_domains, 1, 1, 0)]
+        self.main = nn.Sequential(*blocks)
+
+    def forward(self, x):
+        out = self.main(x)
+        out = out.view(out.size(0), -1)  # (batch, num_domains)
+        return out

components/Nonstau_Discriminator_FM.py

@@ -0,0 +1,107 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+#############################################################
+# File: Nonstau_Discriminator.py
+# Created Date: Monday March 28th 2022
+# Author: Chen Xuanhong
+# Email: chenxuanhongzju@outlook.com
+# Last Modified:  Wednesday, 13th April 2022 3:11:40 pm
+# Modified By: Chen Xuanhong
+# Copyright (c) 2022 Shanghai Jiao Tong University
+#############################################################
+import math
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class ResBlk(nn.Module):
+    def __init__(self, dim_in, dim_out, actv=nn.LeakyReLU(0.2),
+                 normalize="in", downsample=False):
+        super().__init__()
+        self.actv = actv
+        self.normalize = normalize
+        self.downsample = downsample
+        self.learned_sc = dim_in != dim_out
+        self.var        = math.sqrt(2)
+        self._build_weights(dim_in, dim_out)
+
+    def _build_weights(self, dim_in, dim_out):
+        self.conv1 = nn.Conv2d(dim_in, dim_in, 3, 1, 1)
+        self.conv2 = nn.Conv2d(dim_in, dim_out, 3, 1, 1)
+        if self.normalize.lower() == "in":
+            self.norm1 = nn.InstanceNorm2d(dim_in, affine=True)
+            self.norm2 = nn.InstanceNorm2d(dim_in, affine=True)
+        elif self.normalize.lower() == "bn":
+            self.norm1 = nn.BatchNorm2d(dim_in)
+            self.norm2 = nn.BatchNorm2d(dim_in)
+        elif self.normalize.lower() == "none":
+            self.normalize = False
+        if self.learned_sc:
+            self.conv1x1 = nn.Conv2d(dim_in, dim_out, 1, 1, 0, bias=False)
+
+    def _shortcut(self, x):
+        if self.learned_sc:
+            x = self.conv1x1(x)
+        if self.downsample:
+            x = F.avg_pool2d(x, 2)
+        return x
+
+    def _residual(self, x):
+        if self.normalize:
+            x = self.norm1(x)
+        x = self.actv(x)
+        x = self.conv1(x)
+        if self.downsample:
+            x = F.avg_pool2d(x, 2)
+        if self.normalize:
+            x = self.norm2(x)
+        x = self.actv(x)
+        x = self.conv2(x)
+        return x
+
+    def forward(self, x):
+        x = self._shortcut(x) + self._residual(x)
+        return x / self.var  # unit variance
+
+class Discriminator(torch.nn.Module):
+    def __init__(
+        self,
+        **kwargs
+    ):
+        super().__init__()
+        img_size     = kwargs["img_size"]
+        num_domains  = 1
+        max_conv_dim = kwargs["max_conv_dim"]
+        norm         = kwargs["norm"].lower()
+        dim_in = 2**14 // img_size
+        blocks = []
+        blocks += [nn.Conv2d(3, dim_in, 3, 1, 1)]
+
+        repeat_num = int(np.log2(img_size)) - 2
+        for _ in range(repeat_num-2):
+            dim_out = min(dim_in*2, max_conv_dim)
+            blocks += [ResBlk(dim_in, dim_out, normalize=norm, downsample=True)]
+            dim_in = dim_out
+        blocks1 = []
+        for _ in range(2): # 16
+            dim_out = min(dim_in*2, max_conv_dim)
+            blocks1 += [ResBlk(dim_in, dim_out, normalize=norm, downsample=True)]
+            dim_in = dim_out
+
+        blocks1 += [nn.LeakyReLU(0.2)]
+        blocks1 += [nn.Conv2d(dim_out, dim_out, 4, 1, 0)]
+        blocks1 += [nn.LeakyReLU(0.2)]
+        blocks1 += [nn.Conv2d(dim_out, num_domains, 1, 1, 0)]
+        self.main = nn.Sequential(*blocks)
+        self.tail = nn.Sequential(*blocks1)
+    
+    def get_feature(self,x):
+        mid = self.main(x)
+        return mid
+
+    def forward(self, x):
+        mid = self.main(x)
+        out = self.tail(mid)
+        out = out.view(out.size(0), -1)  # (batch, num_domains)
+        return out,mid