import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.modules.batchnorm import BatchNorm2d
from torch.nn.utils.spectral_norm import spectral_norm as SpectralNorm
from videoretalking.models.ffc import FFC
from basicsr.archs.arch_util import default_init_weights
class Conv2d(nn.Module):
def __init__(self, cin, cout, kernel_size, stride, padding, residual=False, *args, **kwargs):
super().__init__(*args, **kwargs)
self.conv_block = nn.Sequential(
nn.Conv2d(cin, cout, kernel_size, stride, padding),
nn.BatchNorm2d(cout)
)
self.act = nn.ReLU()
self.residual = residual
def forward(self, x):
out = self.conv_block(x)
if self.residual:
out += x
return self.act(out)
class ResBlock(nn.Module):
def __init__(self, in_channels, out_channels, mode='down'):
super(ResBlock, self).__init__()
self.conv1 = nn.Conv2d(in_channels, in_channels, 3, 1, 1)
self.conv2 = nn.Conv2d(in_channels, out_channels, 3, 1, 1)
self.skip = nn.Conv2d(in_channels, out_channels, 1, bias=False)
if mode == 'down':
self.scale_factor = 0.5
elif mode == 'up':
self.scale_factor = 2
def forward(self, x):
out = F.leaky_relu_(self.conv1(x), negative_slope=0.2)
# upsample/downsample
out = F.interpolate(out, scale_factor=self.scale_factor, mode='bilinear', align_corners=False)
out = F.leaky_relu_(self.conv2(out), negative_slope=0.2)
# skip
x = F.interpolate(x, scale_factor=self.scale_factor, mode='bilinear', align_corners=False)
skip = self.skip(x)
out = out + skip
return out
class LayerNorm2d(nn.Module):
def __init__(self, n_out, affine=True):
super(LayerNorm2d, self).__init__()
self.n_out = n_out
self.affine = affine
if self.affine:
self.weight = nn.Parameter(torch.ones(n_out, 1, 1))
self.bias = nn.Parameter(torch.zeros(n_out, 1, 1))
def forward(self, x):
normalized_shape = x.size()[1:]
if self.affine:
return F.layer_norm(x, normalized_shape, \
self.weight.expand(normalized_shape),
self.bias.expand(normalized_shape))
else:
return F.layer_norm(x, normalized_shape)
def spectral_norm(module, use_spect=True):
if use_spect:
return SpectralNorm(module)
else:
return module
class FirstBlock2d(nn.Module):
def __init__(self, input_nc, output_nc, norm_layer=nn.BatchNorm2d, nonlinearity=nn.LeakyReLU(), use_spect=False):
super(FirstBlock2d, self).__init__()
kwargs = {'kernel_size': 7, 'stride': 1, 'padding': 3}
conv = spectral_norm(nn.Conv2d(input_nc, output_nc, **kwargs), use_spect)
if type(norm_layer) == type(None):
self.model = nn.Sequential(conv, nonlinearity)
else:
self.model = nn.Sequential(conv, norm_layer(output_nc), nonlinearity)
def forward(self, x):
out = self.model(x)
return out
class DownBlock2d(nn.Module):
def __init__(self, input_nc, output_nc, norm_layer=nn.BatchNorm2d, nonlinearity=nn.LeakyReLU(), use_spect=False):
super(DownBlock2d, self).__init__()
kwargs = {'kernel_size': 3, 'stride': 1, 'padding': 1}
conv = spectral_norm(nn.Conv2d(input_nc, output_nc, **kwargs), use_spect)
pool = nn.AvgPool2d(kernel_size=(2, 2))
if type(norm_layer) == type(None):
self.model = nn.Sequential(conv, nonlinearity, pool)
else:
self.model = nn.Sequential(conv, norm_layer(output_nc), nonlinearity, pool)
def forward(self, x):
out = self.model(x)
return out
class UpBlock2d(nn.Module):
def __init__(self, input_nc, output_nc, norm_layer=nn.BatchNorm2d, nonlinearity=nn.LeakyReLU(), use_spect=False):
super(UpBlock2d, self).__init__()
kwargs = {'kernel_size': 3, 'stride': 1, 'padding': 1}
conv = spectral_norm(nn.Conv2d(input_nc, output_nc, **kwargs), use_spect)
if type(norm_layer) == type(None):
self.model = nn.Sequential(conv, nonlinearity)
else:
self.model = nn.Sequential(conv, norm_layer(output_nc), nonlinearity)
def forward(self, x):
out = self.model(F.interpolate(x, scale_factor=2))
return out
class ADAIN(nn.Module):
def __init__(self, norm_nc, feature_nc):
super().__init__()
self.param_free_norm = nn.InstanceNorm2d(norm_nc, affine=False)
nhidden = 128
use_bias=True
self.mlp_shared = nn.Sequential(
nn.Linear(feature_nc, nhidden, bias=use_bias),
nn.ReLU()
)
self.mlp_gamma = nn.Linear(nhidden, norm_nc, bias=use_bias)
self.mlp_beta = nn.Linear(nhidden, norm_nc, bias=use_bias)
def forward(self, x, feature):
# Part 1. generate parameter-free normalized activations
normalized = self.param_free_norm(x)
# Part 2. produce scaling and bias conditioned on feature
feature = feature.view(feature.size(0), -1)
actv = self.mlp_shared(feature)
gamma = self.mlp_gamma(actv)
beta = self.mlp_beta(actv)
# apply scale and bias
gamma = gamma.view(*gamma.size()[:2], 1,1)
beta = beta.view(*beta.size()[:2], 1,1)
out = normalized * (1 + gamma) + beta
return out
class FineADAINResBlock2d(nn.Module):
"""
Define an Residual block for different types
"""
def __init__(self, input_nc, feature_nc, norm_layer=nn.BatchNorm2d, nonlinearity=nn.LeakyReLU(), use_spect=False):
super(FineADAINResBlock2d, self).__init__()
kwargs = {'kernel_size': 3, 'stride': 1, 'padding': 1}
self.conv1 = spectral_norm(nn.Conv2d(input_nc, input_nc, **kwargs), use_spect)
self.conv2 = spectral_norm(nn.Conv2d(input_nc, input_nc, **kwargs), use_spect)
self.norm1 = ADAIN(input_nc, feature_nc)
self.norm2 = ADAIN(input_nc, feature_nc)
self.actvn = nonlinearity
def forward(self, x, z):
dx = self.actvn(self.norm1(self.conv1(x), z))
dx = self.norm2(self.conv2(x), z)
out = dx + x
return out
class FineADAINResBlocks(nn.Module):
def __init__(self, num_block, input_nc, feature_nc, norm_layer=nn.BatchNorm2d, nonlinearity=nn.LeakyReLU(), use_spect=False):
super(FineADAINResBlocks, self).__init__()
self.num_block = num_block
for i in range(num_block):
model = FineADAINResBlock2d(input_nc, feature_nc, norm_layer, nonlinearity, use_spect)
setattr(self, 'res'+str(i), model)
def forward(self, x, z):
for i in range(self.num_block):
model = getattr(self, 'res'+str(i))
x = model(x, z)
return x
class ADAINEncoderBlock(nn.Module):
def __init__(self, input_nc, output_nc, feature_nc, nonlinearity=nn.LeakyReLU(), use_spect=False):
super(ADAINEncoderBlock, self).__init__()
kwargs_down = {'kernel_size': 4, 'stride': 2, 'padding': 1}
kwargs_fine = {'kernel_size': 3, 'stride': 1, 'padding': 1}
self.conv_0 = spectral_norm(nn.Conv2d(input_nc, output_nc, **kwargs_down), use_spect)
self.conv_1 = spectral_norm(nn.Conv2d(output_nc, output_nc, **kwargs_fine), use_spect)
self.norm_0 = ADAIN(input_nc, feature_nc)
self.norm_1 = ADAIN(output_nc, feature_nc)
self.actvn = nonlinearity
def forward(self, x, z):
x = self.conv_0(self.actvn(self.norm_0(x, z)))
x = self.conv_1(self.actvn(self.norm_1(x, z)))
return x
class ADAINDecoderBlock(nn.Module):
def __init__(self, input_nc, output_nc, hidden_nc, feature_nc, use_transpose=True, nonlinearity=nn.LeakyReLU(), use_spect=False):
super(ADAINDecoderBlock, self).__init__()
# Attributes
self.actvn = nonlinearity
hidden_nc = min(input_nc, output_nc) if hidden_nc is None else hidden_nc
kwargs_fine = {'kernel_size':3, 'stride':1, 'padding':1}
if use_transpose:
kwargs_up = {'kernel_size':3, 'stride':2, 'padding':1, 'output_padding':1}
else:
kwargs_up = {'kernel_size':3, 'stride':1, 'padding':1}
# create conv layers
self.conv_0 = spectral_norm(nn.Conv2d(input_nc, hidden_nc, **kwargs_fine), use_spect)
if use_transpose:
self.conv_1 = spectral_norm(nn.ConvTranspose2d(hidden_nc, output_nc, **kwargs_up), use_spect)
self.conv_s = spectral_norm(nn.ConvTranspose2d(input_nc, output_nc, **kwargs_up), use_spect)
else:
self.conv_1 = nn.Sequential(spectral_norm(nn.Conv2d(hidden_nc, output_nc, **kwargs_up), use_spect),
nn.Upsample(scale_factor=2))
self.conv_s = nn.Sequential(spectral_norm(nn.Conv2d(input_nc, output_nc, **kwargs_up), use_spect),
nn.Upsample(scale_factor=2))
# define normalization layers
self.norm_0 = ADAIN(input_nc, feature_nc)
self.norm_1 = ADAIN(hidden_nc, feature_nc)
self.norm_s = ADAIN(input_nc, feature_nc)
def forward(self, x, z):
x_s = self.shortcut(x, z)
dx = self.conv_0(self.actvn(self.norm_0(x, z)))
dx = self.conv_1(self.actvn(self.norm_1(dx, z)))
out = x_s + dx
return out
def shortcut(self, x, z):
x_s = self.conv_s(self.actvn(self.norm_s(x, z)))
return x_s
class FineEncoder(nn.Module):
"""docstring for Encoder"""
def __init__(self, image_nc, ngf, img_f, layers, norm_layer=nn.BatchNorm2d, nonlinearity=nn.LeakyReLU(), use_spect=False):
super(FineEncoder, self).__init__()
self.layers = layers
self.first = FirstBlock2d(image_nc, ngf, norm_layer, nonlinearity, use_spect)
for i in range(layers):
in_channels = min(ngf*(2**i), img_f)
out_channels = min(ngf*(2**(i+1)), img_f)
model = DownBlock2d(in_channels, out_channels, norm_layer, nonlinearity, use_spect)
setattr(self, 'down' + str(i), model)
self.output_nc = out_channels
def forward(self, x):
x = self.first(x)
out=[x]
for i in range(self.layers):
model = getattr(self, 'down'+str(i))
x = model(x)
out.append(x)
return out
class FineDecoder(nn.Module):
"""docstring for FineDecoder"""
def __init__(self, image_nc, feature_nc, ngf, img_f, layers, num_block, norm_layer=nn.BatchNorm2d, nonlinearity=nn.LeakyReLU(), use_spect=False):
super(FineDecoder, self).__init__()
self.layers = layers
for i in range(layers)[::-1]:
in_channels = min(ngf*(2**(i+1)), img_f)
out_channels = min(ngf*(2**i), img_f)
up = UpBlock2d(in_channels, out_channels, norm_layer, nonlinearity, use_spect)
res = FineADAINResBlocks(num_block, in_channels, feature_nc, norm_layer, nonlinearity, use_spect)
jump = Jump(out_channels, norm_layer, nonlinearity, use_spect)
setattr(self, 'up' + str(i), up)
setattr(self, 'res' + str(i), res)
setattr(self, 'jump' + str(i), jump)
self.final = FinalBlock2d(out_channels, image_nc, use_spect, 'tanh')
self.output_nc = out_channels
def forward(self, x, z):
out = x.pop()
for i in range(self.layers)[::-1]:
res_model = getattr(self, 'res' + str(i))
up_model = getattr(self, 'up' + str(i))
jump_model = getattr(self, 'jump' + str(i))
out = res_model(out, z)
out = up_model(out)
out = jump_model(x.pop()) + out
out_image = self.final(out)
return out_image
class ADAINEncoder(nn.Module):
def __init__(self, image_nc, pose_nc, ngf, img_f, layers, nonlinearity=nn.LeakyReLU(), use_spect=False):
super(ADAINEncoder, self).__init__()
self.layers = layers
self.input_layer = nn.Conv2d(image_nc, ngf, kernel_size=7, stride=1, padding=3)
for i in range(layers):
in_channels = min(ngf * (2**i), img_f)
out_channels = min(ngf *(2**(i+1)), img_f)
model = ADAINEncoderBlock(in_channels, out_channels, pose_nc, nonlinearity, use_spect)
setattr(self, 'encoder' + str(i), model)
self.output_nc = out_channels
def forward(self, x, z):
out = self.input_layer(x)
out_list = [out]
for i in range(self.layers):
model = getattr(self, 'encoder' + str(i))
out = model(out, z)
out_list.append(out)
return out_list
class ADAINDecoder(nn.Module):
"""docstring for ADAINDecoder"""
def __init__(self, pose_nc, ngf, img_f, encoder_layers, decoder_layers, skip_connect=True,
nonlinearity=nn.LeakyReLU(), use_spect=False):
super(ADAINDecoder, self).__init__()
self.encoder_layers = encoder_layers
self.decoder_layers = decoder_layers
self.skip_connect = skip_connect
use_transpose = True
for i in range(encoder_layers-decoder_layers, encoder_layers)[::-1]:
in_channels = min(ngf * (2**(i+1)), img_f)
in_channels = in_channels*2 if i != (encoder_layers-1) and self.skip_connect else in_channels
out_channels = min(ngf * (2**i), img_f)
model = ADAINDecoderBlock(in_channels, out_channels, out_channels, pose_nc, use_transpose, nonlinearity, use_spect)
setattr(self, 'decoder' + str(i), model)
self.output_nc = out_channels*2 if self.skip_connect else out_channels
def forward(self, x, z):
out = x.pop() if self.skip_connect else x
for i in range(self.encoder_layers-self.decoder_layers, self.encoder_layers)[::-1]:
model = getattr(self, 'decoder' + str(i))
out = model(out, z)
out = torch.cat([out, x.pop()], 1) if self.skip_connect else out
return out
class ADAINHourglass(nn.Module):
def __init__(self, image_nc, pose_nc, ngf, img_f, encoder_layers, decoder_layers, nonlinearity, use_spect):
super(ADAINHourglass, self).__init__()
self.encoder = ADAINEncoder(image_nc, pose_nc, ngf, img_f, encoder_layers, nonlinearity, use_spect)
self.decoder = ADAINDecoder(pose_nc, ngf, img_f, encoder_layers, decoder_layers, True, nonlinearity, use_spect)
self.output_nc = self.decoder.output_nc
def forward(self, x, z):
return self.decoder(self.encoder(x, z), z)
class FineADAINLama(nn.Module):
def __init__(self, input_nc, feature_nc, norm_layer=nn.BatchNorm2d, nonlinearity=nn.LeakyReLU(), use_spect=False):
super(FineADAINLama, self).__init__()
kwargs = {'kernel_size': 3, 'stride': 1, 'padding': 1}
self.actvn = nonlinearity
ratio_gin = 0.75
ratio_gout = 0.75
self.ffc = FFC(input_nc, input_nc, 3,
ratio_gin, ratio_gout, 1, 1, 1,
1, False, False, padding_type='reflect')
global_channels = int(input_nc * ratio_gout)
self.bn_l = ADAIN(input_nc - global_channels, feature_nc)
self.bn_g = ADAIN(global_channels, feature_nc)
def forward(self, x, z):
x_l, x_g = self.ffc(x)
x_l = self.actvn(self.bn_l(x_l,z))
x_g = self.actvn(self.bn_g(x_g,z))
return x_l, x_g
class FFCResnetBlock(nn.Module):
def __init__(self, dim, feature_dim, padding_type='reflect', norm_layer=BatchNorm2d, activation_layer=nn.ReLU, dilation=1,
spatial_transform_kwargs=None, inline=False, **conv_kwargs):
super().__init__()
self.conv1 = FineADAINLama(dim, feature_dim, **conv_kwargs)
self.conv2 = FineADAINLama(dim, feature_dim, **conv_kwargs)
self.inline = True
def forward(self, x, z):
if self.inline:
x_l, x_g = x[:, :-self.conv1.ffc.global_in_num], x[:, -self.conv1.ffc.global_in_num:]
else:
x_l, x_g = x if type(x) is tuple else (x, 0)
id_l, id_g = x_l, x_g
x_l, x_g = self.conv1((x_l, x_g), z)
x_l, x_g = self.conv2((x_l, x_g), z)
x_l, x_g = id_l + x_l, id_g + x_g
out = x_l, x_g
if self.inline:
out = torch.cat(out, dim=1)
return out
class FFCADAINResBlocks(nn.Module):
def __init__(self, num_block, input_nc, feature_nc, norm_layer=nn.BatchNorm2d, nonlinearity=nn.LeakyReLU(), use_spect=False):
super(FFCADAINResBlocks, self).__init__()
self.num_block = num_block
for i in range(num_block):
model = FFCResnetBlock(input_nc, feature_nc, norm_layer, nonlinearity, use_spect)
setattr(self, 'res'+str(i), model)
def forward(self, x, z):
for i in range(self.num_block):
model = getattr(self, 'res'+str(i))
x = model(x, z)
return x
class Jump(nn.Module):
def __init__(self, input_nc, norm_layer=nn.BatchNorm2d, nonlinearity=nn.LeakyReLU(), use_spect=False):
super(Jump, self).__init__()
kwargs = {'kernel_size': 3, 'stride': 1, 'padding': 1}
conv = spectral_norm(nn.Conv2d(input_nc, input_nc, **kwargs), use_spect)
if type(norm_layer) == type(None):
self.model = nn.Sequential(conv, nonlinearity)
else:
self.model = nn.Sequential(conv, norm_layer(input_nc), nonlinearity)
def forward(self, x):
out = self.model(x)
return out
class FinalBlock2d(nn.Module):
def __init__(self, input_nc, output_nc, use_spect=False, tanh_or_sigmoid='tanh'):
super(FinalBlock2d, self).__init__()
kwargs = {'kernel_size': 7, 'stride': 1, 'padding':3}
conv = spectral_norm(nn.Conv2d(input_nc, output_nc, **kwargs), use_spect)
if tanh_or_sigmoid == 'sigmoid':
out_nonlinearity = nn.Sigmoid()
else:
out_nonlinearity = nn.Tanh()
self.model = nn.Sequential(conv, out_nonlinearity)
def forward(self, x):
out = self.model(x)
return out
class ModulatedConv2d(nn.Module):
def __init__(self,
in_channels,
out_channels,
kernel_size,
num_style_feat,
demodulate=True,
sample_mode=None,
eps=1e-8):
super(ModulatedConv2d, self).__init__()
self.in_channels = in_channels
self.out_channels = out_channels
self.kernel_size = kernel_size
self.demodulate = demodulate
self.sample_mode = sample_mode
self.eps = eps
# modulation inside each modulated conv
self.modulation = nn.Linear(num_style_feat, in_channels, bias=True)
# initialization
default_init_weights(self.modulation, scale=1, bias_fill=1, a=0, mode='fan_in', nonlinearity='linear')
self.weight = nn.Parameter(
torch.randn(1, out_channels, in_channels, kernel_size, kernel_size) /
math.sqrt(in_channels * kernel_size**2))
self.padding = kernel_size // 2
def forward(self, x, style):
b, c, h, w = x.shape
style = self.modulation(style).view(b, 1, c, 1, 1)
weight = self.weight * style
if self.demodulate:
demod = torch.rsqrt(weight.pow(2).sum([2, 3, 4]) + self.eps)
weight = weight * demod.view(b, self.out_channels, 1, 1, 1)
weight = weight.view(b * self.out_channels, c, self.kernel_size, self.kernel_size)
# upsample or downsample if necessary
if self.sample_mode == 'upsample':
x = F.interpolate(x, scale_factor=2, mode='bilinear', align_corners=False)
elif self.sample_mode == 'downsample':
x = F.interpolate(x, scale_factor=0.5, mode='bilinear', align_corners=False)
b, c, h, w = x.shape
x = x.view(1, b * c, h, w)
out = F.conv2d(x, weight, padding=self.padding, groups=b)
out = out.view(b, self.out_channels, *out.shape[2:4])
return out
def __repr__(self):
return (f'{self.__class__.__name__}(in_channels={self.in_channels}, out_channels={self.out_channels}, '
f'kernel_size={self.kernel_size}, demodulate={self.demodulate}, sample_mode={self.sample_mode})')
class StyleConv(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size, num_style_feat, demodulate=True, sample_mode=None):
super(StyleConv, self).__init__()
self.modulated_conv = ModulatedConv2d(
in_channels, out_channels, kernel_size, num_style_feat, demodulate=demodulate, sample_mode=sample_mode)
self.weight = nn.Parameter(torch.zeros(1)) # for noise injection
self.bias = nn.Parameter(torch.zeros(1, out_channels, 1, 1))
self.activate = nn.LeakyReLU(negative_slope=0.2, inplace=True)
def forward(self, x, style, noise=None):
# modulate
out = self.modulated_conv(x, style) * 2**0.5 # for conversion
# noise injection
if noise is None:
b, _, h, w = out.shape
noise = out.new_empty(b, 1, h, w).normal_()
out = out + self.weight * noise
# add bias
out = out + self.bias
# activation
out = self.activate(out)
return out
class ToRGB(nn.Module):
def __init__(self, in_channels, num_style_feat, upsample=True):
super(ToRGB, self).__init__()
self.upsample = upsample
self.modulated_conv = ModulatedConv2d(
in_channels, 3, kernel_size=1, num_style_feat=num_style_feat, demodulate=False, sample_mode=None)
self.bias = nn.Parameter(torch.zeros(1, 3, 1, 1))
def forward(self, x, style, skip=None):
out = self.modulated_conv(x, style)
out = out + self.bias
if skip is not None:
if self.upsample:
skip = F.interpolate(skip, scale_factor=2, mode='bilinear', align_corners=False)
out = out + skip
return out