mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-18 00:33:15 +00:00
feature: new lightweight facerender
This commit is contained in:
Justin John
Binary file not shown.
Binary file not shown.
Binary file not shown.
BIN
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,368 @@
|
||||
import sys
|
||||
import math
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.nn import functional as F
|
||||
from torch.autograd import Function
|
||||
from torch.nn.utils.spectral_norm import spectral_norm as SpectralNorm
|
||||
|
||||
|
||||
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)
|
||||
|
||||
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 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 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
|
||||
|
||||
|
||||
def spectral_norm(module, use_spect=True):
|
||||
"""use spectral normal layer to stable the training process"""
|
||||
if use_spect:
|
||||
return SpectralNorm(module)
|
||||
else:
|
||||
return module
|
||||
|
||||
|
||||
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 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 FirstBlock2d(nn.Module):
|
||||
"""
|
||||
Downsampling block for use in encoder.
|
||||
"""
|
||||
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 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 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 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 FinalBlock2d(nn.Module):
|
||||
"""
|
||||
Define the output layer
|
||||
"""
|
||||
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
|
||||
@@ -0,0 +1,211 @@
|
||||
import collections
|
||||
import functools
|
||||
import os
|
||||
import re
|
||||
|
||||
import yaml
|
||||
|
||||
class AttrDict(dict):
|
||||
"""Dict as attribute trick."""
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super(AttrDict, self).__init__(*args, **kwargs)
|
||||
self.__dict__ = self
|
||||
for key, value in self.__dict__.items():
|
||||
if isinstance(value, dict):
|
||||
self.__dict__[key] = AttrDict(value)
|
||||
elif isinstance(value, (list, tuple)):
|
||||
if isinstance(value[0], dict):
|
||||
self.__dict__[key] = [AttrDict(item) for item in value]
|
||||
else:
|
||||
self.__dict__[key] = value
|
||||
|
||||
def yaml(self):
|
||||
"""Convert object to yaml dict and return."""
|
||||
yaml_dict = {}
|
||||
for key, value in self.__dict__.items():
|
||||
if isinstance(value, AttrDict):
|
||||
yaml_dict[key] = value.yaml()
|
||||
elif isinstance(value, list):
|
||||
if isinstance(value[0], AttrDict):
|
||||
new_l = []
|
||||
for item in value:
|
||||
new_l.append(item.yaml())
|
||||
yaml_dict[key] = new_l
|
||||
else:
|
||||
yaml_dict[key] = value
|
||||
else:
|
||||
yaml_dict[key] = value
|
||||
return yaml_dict
|
||||
|
||||
def __repr__(self):
|
||||
"""Print all variables."""
|
||||
ret_str = []
|
||||
for key, value in self.__dict__.items():
|
||||
if isinstance(value, AttrDict):
|
||||
ret_str.append('{}:'.format(key))
|
||||
child_ret_str = value.__repr__().split('\n')
|
||||
for item in child_ret_str:
|
||||
ret_str.append(' ' + item)
|
||||
elif isinstance(value, list):
|
||||
if isinstance(value[0], AttrDict):
|
||||
ret_str.append('{}:'.format(key))
|
||||
for item in value:
|
||||
# Treat as AttrDict above.
|
||||
child_ret_str = item.__repr__().split('\n')
|
||||
for item in child_ret_str:
|
||||
ret_str.append(' ' + item)
|
||||
else:
|
||||
ret_str.append('{}: {}'.format(key, value))
|
||||
else:
|
||||
ret_str.append('{}: {}'.format(key, value))
|
||||
return '\n'.join(ret_str)
|
||||
|
||||
|
||||
class Config(AttrDict):
|
||||
r"""Configuration class. This should include every human specifiable
|
||||
hyperparameter values for your training."""
|
||||
|
||||
def __init__(self, filename=None, args=None, verbose=False, is_train=True):
|
||||
super(Config, self).__init__()
|
||||
# Set default parameters.
|
||||
# Logging.
|
||||
|
||||
large_number = 1000000000
|
||||
self.snapshot_save_iter = large_number
|
||||
self.snapshot_save_epoch = large_number
|
||||
self.snapshot_save_start_iter = 0
|
||||
self.snapshot_save_start_epoch = 0
|
||||
self.image_save_iter = large_number
|
||||
self.eval_epoch = large_number
|
||||
self.start_eval_epoch = large_number
|
||||
self.eval_epoch = large_number
|
||||
self.max_epoch = large_number
|
||||
self.max_iter = large_number
|
||||
self.logging_iter = 100
|
||||
self.image_to_tensorboard=False
|
||||
self.which_iter = 0 # args.which_iter
|
||||
self.resume = False
|
||||
|
||||
self.checkpoints_dir = '/Users/shadowcun/Downloads/'
|
||||
self.name = 'face'
|
||||
self.phase = 'train' if is_train else 'test'
|
||||
|
||||
# Networks.
|
||||
self.gen = AttrDict(type='generators.dummy')
|
||||
self.dis = AttrDict(type='discriminators.dummy')
|
||||
|
||||
# Optimizers.
|
||||
self.gen_optimizer = AttrDict(type='adam',
|
||||
lr=0.0001,
|
||||
adam_beta1=0.0,
|
||||
adam_beta2=0.999,
|
||||
eps=1e-8,
|
||||
lr_policy=AttrDict(iteration_mode=False,
|
||||
type='step',
|
||||
step_size=large_number,
|
||||
gamma=1))
|
||||
self.dis_optimizer = AttrDict(type='adam',
|
||||
lr=0.0001,
|
||||
adam_beta1=0.0,
|
||||
adam_beta2=0.999,
|
||||
eps=1e-8,
|
||||
lr_policy=AttrDict(iteration_mode=False,
|
||||
type='step',
|
||||
step_size=large_number,
|
||||
gamma=1))
|
||||
# Data.
|
||||
self.data = AttrDict(name='dummy',
|
||||
type='datasets.images',
|
||||
num_workers=0)
|
||||
self.test_data = AttrDict(name='dummy',
|
||||
type='datasets.images',
|
||||
num_workers=0,
|
||||
test=AttrDict(is_lmdb=False,
|
||||
roots='',
|
||||
batch_size=1))
|
||||
self.trainer = AttrDict(
|
||||
model_average=False,
|
||||
model_average_beta=0.9999,
|
||||
model_average_start_iteration=1000,
|
||||
model_average_batch_norm_estimation_iteration=30,
|
||||
model_average_remove_sn=True,
|
||||
image_to_tensorboard=False,
|
||||
hparam_to_tensorboard=False,
|
||||
distributed_data_parallel='pytorch',
|
||||
delay_allreduce=True,
|
||||
gan_relativistic=False,
|
||||
gen_step=1,
|
||||
dis_step=1)
|
||||
|
||||
# # Cudnn.
|
||||
self.cudnn = AttrDict(deterministic=False,
|
||||
benchmark=True)
|
||||
|
||||
# Others.
|
||||
self.pretrained_weight = ''
|
||||
self.inference_args = AttrDict()
|
||||
|
||||
|
||||
# Update with given configurations.
|
||||
assert os.path.exists(filename), 'File {} not exist.'.format(filename)
|
||||
loader = yaml.SafeLoader
|
||||
loader.add_implicit_resolver(
|
||||
u'tag:yaml.org,2002:float',
|
||||
re.compile(u'''^(?:
|
||||
[-+]?(?:[0-9][0-9_]*)\\.[0-9_]*(?:[eE][-+]?[0-9]+)?
|
||||
|[-+]?(?:[0-9][0-9_]*)(?:[eE][-+]?[0-9]+)
|
||||
|\\.[0-9_]+(?:[eE][-+][0-9]+)?
|
||||
|[-+]?[0-9][0-9_]*(?::[0-5]?[0-9])+\\.[0-9_]*
|
||||
|[-+]?\\.(?:inf|Inf|INF)
|
||||
|\\.(?:nan|NaN|NAN))$''', re.X),
|
||||
list(u'-+0123456789.'))
|
||||
try:
|
||||
with open(filename, 'r') as f:
|
||||
cfg_dict = yaml.load(f, Loader=loader)
|
||||
except EnvironmentError:
|
||||
print('Please check the file with name of "%s"', filename)
|
||||
recursive_update(self, cfg_dict)
|
||||
|
||||
# Put common opts in both gen and dis.
|
||||
if 'common' in cfg_dict:
|
||||
self.common = AttrDict(**cfg_dict['common'])
|
||||
self.gen.common = self.common
|
||||
self.dis.common = self.common
|
||||
|
||||
|
||||
if verbose:
|
||||
print(' config '.center(80, '-'))
|
||||
print(self.__repr__())
|
||||
print(''.center(80, '-'))
|
||||
|
||||
|
||||
def rsetattr(obj, attr, val):
|
||||
"""Recursively find object and set value"""
|
||||
pre, _, post = attr.rpartition('.')
|
||||
return setattr(rgetattr(obj, pre) if pre else obj, post, val)
|
||||
|
||||
|
||||
def rgetattr(obj, attr, *args):
|
||||
"""Recursively find object and return value"""
|
||||
|
||||
def _getattr(obj, attr):
|
||||
r"""Get attribute."""
|
||||
return getattr(obj, attr, *args)
|
||||
|
||||
return functools.reduce(_getattr, [obj] + attr.split('.'))
|
||||
|
||||
|
||||
def recursive_update(d, u):
|
||||
"""Recursively update AttrDict d with AttrDict u"""
|
||||
for key, value in u.items():
|
||||
if isinstance(value, collections.abc.Mapping):
|
||||
d.__dict__[key] = recursive_update(d.get(key, AttrDict({})), value)
|
||||
elif isinstance(value, (list, tuple)):
|
||||
if isinstance(value[0], dict):
|
||||
d.__dict__[key] = [AttrDict(item) for item in value]
|
||||
else:
|
||||
d.__dict__[key] = value
|
||||
else:
|
||||
d.__dict__[key] = value
|
||||
return d
|
||||
@@ -0,0 +1,178 @@
|
||||
import functools
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from .base_function import LayerNorm2d, ADAINHourglass, FineEncoder, FineDecoder
|
||||
|
||||
def convert_flow_to_deformation(flow):
|
||||
r"""convert flow fields to deformations.
|
||||
|
||||
Args:
|
||||
flow (tensor): Flow field obtained by the model
|
||||
Returns:
|
||||
deformation (tensor): The deformation used for warpping
|
||||
"""
|
||||
b,c,h,w = flow.shape
|
||||
flow_norm = 2 * torch.cat([flow[:,:1,...]/(w-1),flow[:,1:,...]/(h-1)], 1)
|
||||
grid = make_coordinate_grid(flow)
|
||||
deformation = grid + flow_norm.permute(0,2,3,1)
|
||||
return deformation
|
||||
|
||||
def make_coordinate_grid(flow):
|
||||
r"""obtain coordinate grid with the same size as the flow filed.
|
||||
|
||||
Args:
|
||||
flow (tensor): Flow field obtained by the model
|
||||
Returns:
|
||||
grid (tensor): The grid with the same size as the input flow
|
||||
"""
|
||||
b,c,h,w = flow.shape
|
||||
|
||||
x = torch.arange(w).to(flow)
|
||||
y = torch.arange(h).to(flow)
|
||||
|
||||
x = (2 * (x / (w - 1)) - 1)
|
||||
y = (2 * (y / (h - 1)) - 1)
|
||||
|
||||
yy = y.view(-1, 1).repeat(1, w)
|
||||
xx = x.view(1, -1).repeat(h, 1)
|
||||
|
||||
meshed = torch.cat([xx.unsqueeze_(2), yy.unsqueeze_(2)], 2)
|
||||
meshed = meshed.expand(b, -1, -1, -1)
|
||||
return meshed
|
||||
|
||||
|
||||
def warp_image(source_image, deformation):
|
||||
r"""warp the input image according to the deformation
|
||||
|
||||
Args:
|
||||
source_image (tensor): source images to be warpped
|
||||
deformation (tensor): deformations used to warp the images; value in range (-1, 1)
|
||||
Returns:
|
||||
output (tensor): the warpped images
|
||||
"""
|
||||
_, h_old, w_old, _ = deformation.shape
|
||||
_, _, h, w = source_image.shape
|
||||
if h_old != h or w_old != w:
|
||||
deformation = deformation.permute(0, 3, 1, 2)
|
||||
deformation = torch.nn.functional.interpolate(deformation, size=(h, w), mode='bilinear')
|
||||
deformation = deformation.permute(0, 2, 3, 1)
|
||||
return torch.nn.functional.grid_sample(source_image, deformation)
|
||||
|
||||
|
||||
class FaceGenerator(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
mapping_net,
|
||||
warpping_net,
|
||||
editing_net,
|
||||
common
|
||||
):
|
||||
super(FaceGenerator, self).__init__()
|
||||
self.mapping_net = MappingNet(**mapping_net)
|
||||
self.warpping_net = WarpingNet(**warpping_net, **common)
|
||||
self.editing_net = EditingNet(**editing_net, **common)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_image,
|
||||
driving_source,
|
||||
stage=None
|
||||
):
|
||||
if stage == 'warp':
|
||||
descriptor = self.mapping_net(driving_source)
|
||||
output = self.warpping_net(input_image, descriptor)
|
||||
else:
|
||||
descriptor = self.mapping_net(driving_source)
|
||||
output = self.warpping_net(input_image, descriptor)
|
||||
output['fake_image'] = self.editing_net(input_image, output['warp_image'], descriptor)
|
||||
return output
|
||||
|
||||
class MappingNet(nn.Module):
|
||||
def __init__(self, coeff_nc, descriptor_nc, layer):
|
||||
super( MappingNet, self).__init__()
|
||||
|
||||
self.layer = layer
|
||||
nonlinearity = nn.LeakyReLU(0.1)
|
||||
|
||||
self.first = nn.Sequential(
|
||||
torch.nn.Conv1d(coeff_nc, descriptor_nc, kernel_size=7, padding=0, bias=True))
|
||||
|
||||
for i in range(layer):
|
||||
net = nn.Sequential(nonlinearity,
|
||||
torch.nn.Conv1d(descriptor_nc, descriptor_nc, kernel_size=3, padding=0, dilation=3))
|
||||
setattr(self, 'encoder' + str(i), net)
|
||||
|
||||
self.pooling = nn.AdaptiveAvgPool1d(1)
|
||||
self.output_nc = descriptor_nc
|
||||
|
||||
def forward(self, input_3dmm):
|
||||
out = self.first(input_3dmm)
|
||||
for i in range(self.layer):
|
||||
model = getattr(self, 'encoder' + str(i))
|
||||
out = model(out) + out[:,:,3:-3]
|
||||
out = self.pooling(out)
|
||||
return out
|
||||
|
||||
class WarpingNet(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
image_nc,
|
||||
descriptor_nc,
|
||||
base_nc,
|
||||
max_nc,
|
||||
encoder_layer,
|
||||
decoder_layer,
|
||||
use_spect
|
||||
):
|
||||
super( WarpingNet, self).__init__()
|
||||
|
||||
nonlinearity = nn.LeakyReLU(0.1)
|
||||
norm_layer = functools.partial(LayerNorm2d, affine=True)
|
||||
kwargs = {'nonlinearity':nonlinearity, 'use_spect':use_spect}
|
||||
|
||||
self.descriptor_nc = descriptor_nc
|
||||
self.hourglass = ADAINHourglass(image_nc, self.descriptor_nc, base_nc,
|
||||
max_nc, encoder_layer, decoder_layer, **kwargs)
|
||||
|
||||
self.flow_out = nn.Sequential(norm_layer(self.hourglass.output_nc),
|
||||
nonlinearity,
|
||||
nn.Conv2d(self.hourglass.output_nc, 2, kernel_size=7, stride=1, padding=3))
|
||||
|
||||
self.pool = nn.AdaptiveAvgPool2d(1)
|
||||
|
||||
def forward(self, input_image, descriptor):
|
||||
final_output={}
|
||||
output = self.hourglass(input_image, descriptor)
|
||||
final_output['flow_field'] = self.flow_out(output)
|
||||
|
||||
deformation = convert_flow_to_deformation(final_output['flow_field'])
|
||||
final_output['warp_image'] = warp_image(input_image, deformation)
|
||||
return final_output
|
||||
|
||||
|
||||
class EditingNet(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
image_nc,
|
||||
descriptor_nc,
|
||||
layer,
|
||||
base_nc,
|
||||
max_nc,
|
||||
num_res_blocks,
|
||||
use_spect):
|
||||
super(EditingNet, self).__init__()
|
||||
|
||||
nonlinearity = nn.LeakyReLU(0.1)
|
||||
norm_layer = functools.partial(LayerNorm2d, affine=True)
|
||||
kwargs = {'norm_layer':norm_layer, 'nonlinearity':nonlinearity, 'use_spect':use_spect}
|
||||
self.descriptor_nc = descriptor_nc
|
||||
|
||||
# encoder part
|
||||
self.encoder = FineEncoder(image_nc*2, base_nc, max_nc, layer, **kwargs)
|
||||
self.decoder = FineDecoder(image_nc, self.descriptor_nc, base_nc, max_nc, layer, num_res_blocks, **kwargs)
|
||||
|
||||
def forward(self, input_image, warp_image, descriptor):
|
||||
x = torch.cat([input_image, warp_image], 1)
|
||||
x = self.encoder(x)
|
||||
gen_image = self.decoder(x, descriptor)
|
||||
return gen_image
|
||||
@@ -0,0 +1,129 @@
|
||||
import os
|
||||
import cv2
|
||||
from tqdm import tqdm
|
||||
import yaml
|
||||
import numpy as np
|
||||
import warnings
|
||||
from skimage import img_as_ubyte
|
||||
import safetensors
|
||||
import safetensors.torch
|
||||
warnings.filterwarnings('ignore')
|
||||
|
||||
|
||||
import imageio
|
||||
import torch
|
||||
|
||||
from src.facerender.pirender.config import Config
|
||||
from src.facerender.pirender.face_model import FaceGenerator
|
||||
|
||||
from pydub import AudioSegment
|
||||
from src.utils.face_enhancer import enhancer_generator_with_len, enhancer_list
|
||||
from src.utils.paste_pic import paste_pic
|
||||
from src.utils.videoio import save_video_with_watermark
|
||||
|
||||
try:
|
||||
import webui # in webui
|
||||
in_webui = True
|
||||
except:
|
||||
in_webui = False
|
||||
|
||||
class AnimateFromCoeff_PIRender():
|
||||
|
||||
def __init__(self, sadtalker_path, device):
|
||||
|
||||
opt = Config(sadtalker_path['pirender_yaml_path'], None, is_train=False)
|
||||
opt.device = device
|
||||
self.net_G_ema = FaceGenerator(**opt.gen.param).to(opt.device)
|
||||
checkpoint_path = sadtalker_path['pirender_checkpoint']
|
||||
checkpoint = torch.load(checkpoint_path, map_location=lambda storage, loc: storage)
|
||||
self.net_G_ema.load_state_dict(checkpoint['net_G_ema'], strict=False)
|
||||
print('load [net_G] and [net_G_ema] from {}'.format(checkpoint_path))
|
||||
self.net_G = self.net_G_ema.eval()
|
||||
self.device = device
|
||||
|
||||
|
||||
def generate(self, x, video_save_dir, pic_path, crop_info, enhancer=None, background_enhancer=None, preprocess='crop', img_size=256):
|
||||
|
||||
source_image=x['source_image'].type(torch.FloatTensor)
|
||||
source_semantics=x['source_semantics'].type(torch.FloatTensor)
|
||||
target_semantics=x['target_semantics_list'].type(torch.FloatTensor)
|
||||
source_image=source_image.to(self.device)
|
||||
source_semantics=source_semantics.to(self.device)
|
||||
target_semantics=target_semantics.to(self.device)
|
||||
frame_num = x['frame_num']
|
||||
|
||||
with torch.no_grad():
|
||||
predictions_video = []
|
||||
for i in tqdm(range(target_semantics.shape[1]), 'FaceRender:'):
|
||||
predictions_video.append(self.net_G(source_image, target_semantics[:, i])['fake_image'])
|
||||
|
||||
predictions_video = torch.stack(predictions_video, dim=1)
|
||||
predictions_video = predictions_video.reshape((-1,)+predictions_video.shape[2:])
|
||||
|
||||
video = []
|
||||
for idx in range(len(predictions_video)):
|
||||
image = predictions_video[idx]
|
||||
image = np.transpose(image.data.cpu().numpy(), [1, 2, 0]).astype(np.float32)
|
||||
video.append(image)
|
||||
result = img_as_ubyte(video)
|
||||
|
||||
### the generated video is 256x256, so we keep the aspect ratio,
|
||||
original_size = crop_info[0]
|
||||
if original_size:
|
||||
result = [ cv2.resize(result_i,(img_size, int(img_size * original_size[1]/original_size[0]) )) for result_i in result ]
|
||||
|
||||
video_name = x['video_name'] + '.mp4'
|
||||
path = os.path.join(video_save_dir, 'temp_'+video_name)
|
||||
|
||||
imageio.mimsave(path, result, fps=float(25))
|
||||
|
||||
av_path = os.path.join(video_save_dir, video_name)
|
||||
return_path = av_path
|
||||
|
||||
audio_path = x['audio_path']
|
||||
audio_name = os.path.splitext(os.path.split(audio_path)[-1])[0]
|
||||
new_audio_path = os.path.join(video_save_dir, audio_name+'.wav')
|
||||
start_time = 0
|
||||
# cog will not keep the .mp3 filename
|
||||
sound = AudioSegment.from_file(audio_path)
|
||||
frames = frame_num
|
||||
end_time = start_time + frames*1/25*1000
|
||||
word1=sound.set_frame_rate(16000)
|
||||
word = word1[start_time:end_time]
|
||||
word.export(new_audio_path, format="wav")
|
||||
|
||||
save_video_with_watermark(path, new_audio_path, av_path, watermark= False)
|
||||
print(f'The generated video is named {video_save_dir}/{video_name}')
|
||||
|
||||
if 'full' in preprocess.lower():
|
||||
# only add watermark to the full image.
|
||||
video_name_full = x['video_name'] + '_full.mp4'
|
||||
full_video_path = os.path.join(video_save_dir, video_name_full)
|
||||
return_path = full_video_path
|
||||
paste_pic(path, pic_path, crop_info, new_audio_path, full_video_path, extended_crop= True if 'ext' in preprocess.lower() else False)
|
||||
print(f'The generated video is named {video_save_dir}/{video_name_full}')
|
||||
else:
|
||||
full_video_path = av_path
|
||||
|
||||
#### paste back then enhancers
|
||||
if enhancer:
|
||||
video_name_enhancer = x['video_name'] + '_enhanced.mp4'
|
||||
enhanced_path = os.path.join(video_save_dir, 'temp_'+video_name_enhancer)
|
||||
av_path_enhancer = os.path.join(video_save_dir, video_name_enhancer)
|
||||
return_path = av_path_enhancer
|
||||
|
||||
try:
|
||||
enhanced_images_gen_with_len = enhancer_generator_with_len(full_video_path, method=enhancer, bg_upsampler=background_enhancer)
|
||||
imageio.mimsave(enhanced_path, enhanced_images_gen_with_len, fps=float(25))
|
||||
except:
|
||||
enhanced_images_gen_with_len = enhancer_list(full_video_path, method=enhancer, bg_upsampler=background_enhancer)
|
||||
imageio.mimsave(enhanced_path, enhanced_images_gen_with_len, fps=float(25))
|
||||
|
||||
save_video_with_watermark(enhanced_path, new_audio_path, av_path_enhancer, watermark= False)
|
||||
print(f'The generated video is named {video_save_dir}/{video_name_enhancer}')
|
||||
os.remove(enhanced_path)
|
||||
|
||||
os.remove(path)
|
||||
os.remove(new_audio_path)
|
||||
|
||||
return return_path
|
||||
Binary file not shown.
BIN
Binary file not shown.
Binary file not shown.
BIN
Binary file not shown.
Reference in New Issue
Block a user