import numpy as np
import torch
import torch.nn.functional as F
from torch import nn
from torch.nn import Linear, Conv2d, BatchNorm2d, PReLU, Sequential, Module
from models.encoders.helpers import get_blocks, Flatten, bottleneck_IR, bottleneck_IR_SE
from models.stylegan2.model import EqualLinear
class GradualStyleBlock(Module):
def __init__(self, in_c, out_c, spatial, max_pooling=False):
super(GradualStyleBlock, self).__init__()
self.out_c = out_c
self.spatial = spatial
self.max_pooling = max_pooling
num_pools = int(np.log2(spatial))
modules = []
modules += [Conv2d(in_c, out_c, kernel_size=3, stride=2, padding=1),
nn.LeakyReLU()]
for i in range(num_pools - 1):
modules += [
Conv2d(out_c, out_c, kernel_size=3, stride=2, padding=1),
nn.LeakyReLU()
]
self.convs = nn.Sequential(*modules)
self.linear = EqualLinear(out_c, out_c, lr_mul=1)
def forward(self, x):
x = self.convs(x)
# To make E accept more general H*W images, we add global average pooling to
# resize all features to 1*1*512 before mapping to latent codes
if self.max_pooling:
x = F.adaptive_max_pool2d(x, 1) ##### modified
else:
x = F.adaptive_avg_pool2d(x, 1) ##### modified
x = x.view(-1, self.out_c)
x = self.linear(x)
return x
class AdaptiveInstanceNorm(nn.Module):
def __init__(self, fin, style_dim=512):
super().__init__()
self.norm = nn.InstanceNorm2d(fin, affine=False)
self.style = nn.Linear(style_dim, fin * 2)
self.style.bias.data[:fin] = 1
self.style.bias.data[fin:] = 0
def forward(self, input, style):
style = self.style(style).unsqueeze(2).unsqueeze(3)
gamma, beta = style.chunk(2, 1)
out = self.norm(input)
out = gamma * out + beta
return out
class FusionLayer(Module): ##### modified
def __init__(self, inchannel, outchannel, use_skip_torgb=True, use_att=0):
super(FusionLayer, self).__init__()
self.transform = nn.Sequential(nn.Conv2d(inchannel, outchannel, kernel_size=3, stride=1, padding=1),
nn.LeakyReLU())
self.fusion_out = nn.Conv2d(outchannel*2, outchannel, kernel_size=3, stride=1, padding=1)
self.fusion_out.weight.data *= 0.01
self.fusion_out.weight[:,0:outchannel,1,1].data += torch.eye(outchannel)
self.use_skip_torgb = use_skip_torgb
if use_skip_torgb:
self.fusion_skip = nn.Conv2d(3+outchannel, 3, kernel_size=3, stride=1, padding=1)
self.fusion_skip.weight.data *= 0.01
self.fusion_skip.weight[:,0:3,1,1].data += torch.eye(3)
self.use_att = use_att
if use_att:
modules = []
modules.append(nn.Linear(512, outchannel))
for _ in range(use_att):
modules.append(nn.LeakyReLU(negative_slope=0.2, inplace=True))
modules.append(nn.Linear(outchannel, outchannel))
modules.append(nn.LeakyReLU(negative_slope=0.2, inplace=True))
self.linear = Sequential(*modules)
self.norm = AdaptiveInstanceNorm(outchannel*2, outchannel)
self.conv = nn.Conv2d(outchannel*2, 1, 3, 1, 1, bias=True)
def forward(self, feat, out, skip, editing_w=None):
x = self.transform(feat)
# similar to VToonify, use editing vector as condition
# fuse encoder feature and decoder feature with a predicted attention mask m_E
# if self.use_att = False, just fuse them with a simple conv layer
if self.use_att and editing_w is not None:
label = self.linear(editing_w)
m_E = (F.relu(self.conv(self.norm(torch.cat([out, abs(out-x)], dim=1), label)))).tanh()
x = x * m_E
out = self.fusion_out(torch.cat((out, x), dim=1))
if self.use_skip_torgb:
skip = self.fusion_skip(torch.cat((skip, x), dim=1))
return out, skip
class ResnetBlock(nn.Module):
def __init__(self, dim):
super(ResnetBlock, self).__init__()
self.conv_block = nn.Sequential(Conv2d(dim, dim, 3, 1, 1),
nn.LeakyReLU(),
Conv2d(dim, dim, 3, 1, 1))
self.relu = nn.LeakyReLU()
def forward(self, x):
out = x + self.conv_block(x)
return self.relu(out)
# trainable light-weight translation network T
# for sketch/mask-to-face translation,
# we add a trainable T to map y to an intermediate domain where E can more easily extract features.
class ResnetGenerator(nn.Module):
def __init__(self, in_channel=19, res_num=2):
super(ResnetGenerator, self).__init__()
modules = []
modules.append(Conv2d(in_channel, 16, 3, 2, 1))
modules.append(nn.LeakyReLU())
modules.append(Conv2d(16, 16, 3, 2, 1))
modules.append(nn.LeakyReLU())
for _ in range(res_num):
modules.append(ResnetBlock(16))
for _ in range(2):
modules.append(nn.ConvTranspose2d(16, 16, 3, 2, 1, output_padding=1))
modules.append(nn.LeakyReLU())
modules.append(Conv2d(16, 64, 3, 1, 1, bias=False))
modules.append(BatchNorm2d(64))
modules.append(PReLU(64))
self.model = Sequential(*modules)
def forward(self, input):
return self.model(input)
class GradualStyleEncoder(Module):
def __init__(self, num_layers, mode='ir', opts=None):
super(GradualStyleEncoder, self).__init__()
assert num_layers in [50, 100, 152], 'num_layers should be 50,100, or 152'
assert mode in ['ir', 'ir_se'], 'mode should be ir or ir_se'
blocks = get_blocks(num_layers)
if mode == 'ir':
unit_module = bottleneck_IR
elif mode == 'ir_se':
unit_module = bottleneck_IR_SE
# for sketch/mask-to-face translation, add a new network T
if opts.input_nc != 3:
self.input_label_layer = ResnetGenerator(opts.input_nc, opts.res_num)
self.input_layer = Sequential(Conv2d(3, 64, (3, 3), 1, 1, bias=False),
BatchNorm2d(64),
PReLU(64))
modules = []
for block in blocks:
for bottleneck in block:
modules.append(unit_module(bottleneck.in_channel,
bottleneck.depth,
bottleneck.stride))
self.body = Sequential(*modules)
self.styles = nn.ModuleList()
self.style_count = opts.n_styles
self.coarse_ind = 3
self.middle_ind = 7
for i in range(self.style_count):
if i < self.coarse_ind:
style = GradualStyleBlock(512, 512, 16, 'max_pooling' in opts and opts.max_pooling)
elif i < self.middle_ind:
style = GradualStyleBlock(512, 512, 32, 'max_pooling' in opts and opts.max_pooling)
else:
style = GradualStyleBlock(512, 512, 64, 'max_pooling' in opts and opts.max_pooling)
self.styles.append(style)
self.latlayer1 = nn.Conv2d(256, 512, kernel_size=1, stride=1, padding=0)
self.latlayer2 = nn.Conv2d(128, 512, kernel_size=1, stride=1, padding=0)
# we concatenate pSp features in the middle layers and
# add a convolution layer to map the concatenated features to the first-layer input feature f of G.
self.featlayer = nn.Conv2d(768, 512, kernel_size=1, stride=1, padding=0) ##### modified
self.skiplayer = nn.Conv2d(768, 3, kernel_size=1, stride=1, padding=0) ##### modified
# skip connection
if 'use_skip' in opts and opts.use_skip: ##### modified
self.fusion = nn.ModuleList()
channels = [[256,512], [256,512], [256,512], [256,512], [128,512], [64,256], [64,128]]
# opts.skip_max_layer: how many layers are skipped to the decoder
for inc, outc in channels[:max(1, min(7, opts.skip_max_layer))]: # from 4 to 256
self.fusion.append(FusionLayer(inc, outc, opts.use_skip_torgb, opts.use_att))
def _upsample_add(self, x, y):
'''Upsample and add two feature maps.
Args:
x: (Variable) top feature map to be upsampled.
y: (Variable) lateral feature map.
Returns:
(Variable) added feature map.
Note in PyTorch, when input size is odd, the upsampled feature map
with `F.upsample(..., scale_factor=2, mode='nearest')`
maybe not equal to the lateral feature map size.
e.g.
original input size: [N,_,15,15] ->
conv2d feature map size: [N,_,8,8] ->
upsampled feature map size: [N,_,16,16]
So we choose bilinear upsample which supports arbitrary output sizes.
'''
_, _, H, W = y.size()
return F.interpolate(x, size=(H, W), mode='bilinear', align_corners=True) + y
# return_feat: return f
# return_full: return f and the skipped encoder features
# return [out, feats]
# out is the style latent code w+
# feats[0] is f for the 1st conv layer, feats[1] is f for the 1st torgb layer
# feats[2-8] is the skipped encoder features
def forward(self, x, return_feat=False, return_full=False): ##### modified
if x.shape[1] != 3:
x = self.input_label_layer(x)
else:
x = self.input_layer(x)
c256 = x ##### modified
latents = []
modulelist = list(self.body._modules.values())
for i, l in enumerate(modulelist):
x = l(x)
if i == 2: ##### modified
c128 = x
elif i == 6:
c1 = x
elif i == 10: ##### modified
c21 = x ##### modified
elif i == 15: ##### modified
c22 = x ##### modified
elif i == 20:
c2 = x
elif i == 23:
c3 = x
for j in range(self.coarse_ind):
latents.append(self.styles[j](c3))
p2 = self._upsample_add(c3, self.latlayer1(c2))
for j in range(self.coarse_ind, self.middle_ind):
latents.append(self.styles[j](p2))
p1 = self._upsample_add(p2, self.latlayer2(c1))
for j in range(self.middle_ind, self.style_count):
latents.append(self.styles[j](p1))
out = torch.stack(latents, dim=1)
if not return_feat:
return out
feats = [self.featlayer(torch.cat((c21, c22, c2), dim=1)), self.skiplayer(torch.cat((c21, c22, c2), dim=1))]
if return_full: ##### modified
feats += [c2, c2, c22, c21, c1, c128, c256]
return out, feats
# only compute the first-layer feature f
# E_F in the paper
def get_feat(self, x): ##### modified
# for sketch/mask-to-face translation
# use a trainable light-weight translation network T
if x.shape[1] != 3:
x = self.input_label_layer(x)
else:
x = self.input_layer(x)
latents = []
modulelist = list(self.body._modules.values())
for i, l in enumerate(modulelist):
x = l(x)
if i == 10: ##### modified
c21 = x ##### modified
elif i == 15: ##### modified
c22 = x ##### modified
elif i == 20:
c2 = x
break
return self.featlayer(torch.cat((c21, c22, c2), dim=1))
class BackboneEncoderUsingLastLayerIntoW(Module):
def __init__(self, num_layers, mode='ir', opts=None):
super(BackboneEncoderUsingLastLayerIntoW, self).__init__()
print('Using BackboneEncoderUsingLastLayerIntoW')
assert num_layers in [50, 100, 152], 'num_layers should be 50,100, or 152'
assert mode in ['ir', 'ir_se'], 'mode should be ir or ir_se'
blocks = get_blocks(num_layers)
if mode == 'ir':
unit_module = bottleneck_IR
elif mode == 'ir_se':
unit_module = bottleneck_IR_SE
self.input_layer = Sequential(Conv2d(opts.input_nc, 64, (3, 3), 1, 1, bias=False),
BatchNorm2d(64),
PReLU(64))
self.output_pool = torch.nn.AdaptiveAvgPool2d((1, 1))
self.linear = EqualLinear(512, 512, lr_mul=1)
modules = []
for block in blocks:
for bottleneck in block:
modules.append(unit_module(bottleneck.in_channel,
bottleneck.depth,
bottleneck.stride))
self.body = Sequential(*modules)
def forward(self, x):
x = self.input_layer(x)
x = self.body(x)
x = self.output_pool(x)
x = x.view(-1, 512)
x = self.linear(x)
return x
class BackboneEncoderUsingLastLayerIntoWPlus(Module):
def __init__(self, num_layers, mode='ir', opts=None):
super(BackboneEncoderUsingLastLayerIntoWPlus, self).__init__()
print('Using BackboneEncoderUsingLastLayerIntoWPlus')
assert num_layers in [50, 100, 152], 'num_layers should be 50,100, or 152'
assert mode in ['ir', 'ir_se'], 'mode should be ir or ir_se'
blocks = get_blocks(num_layers)
if mode == 'ir':
unit_module = bottleneck_IR
elif mode == 'ir_se':
unit_module = bottleneck_IR_SE
self.n_styles = opts.n_styles
self.input_layer = Sequential(Conv2d(opts.input_nc, 64, (3, 3), 1, 1, bias=False),
BatchNorm2d(64),
PReLU(64))
self.output_layer_2 = Sequential(BatchNorm2d(512),
torch.nn.AdaptiveAvgPool2d((7, 7)),
Flatten(),
Linear(512 * 7 * 7, 512))
self.linear = EqualLinear(512, 512 * self.n_styles, lr_mul=1)
modules = []
for block in blocks:
for bottleneck in block:
modules.append(unit_module(bottleneck.in_channel,
bottleneck.depth,
bottleneck.stride))
self.body = Sequential(*modules)
def forward(self, x):
x = self.input_layer(x)
x = self.body(x)
x = self.output_layer_2(x)
x = self.linear(x)
x = x.view(-1, self.n_styles, 512)
return x