diff --git a/animation/SadTalker/checkpoints/.gitkeep b/animation/SadTalker/checkpoints/.gitkeep index 8b13789..e69de29 100644 --- a/animation/SadTalker/checkpoints/.gitkeep +++ b/animation/SadTalker/checkpoints/.gitkeep @@ -1 +0,0 @@ - diff --git a/animation/SadTalker/inference.py b/animation/SadTalker/inference.py index 5a405af..4cc0a00 100644 --- a/animation/SadTalker/inference.py +++ b/animation/SadTalker/inference.py @@ -3,11 +3,13 @@ import shutil import torch from time import strftime import os, sys, time +import platform from argparse import ArgumentParser from src.utils.preprocess import CropAndExtract from src.test_audio2coeff import Audio2Coeff from src.facerender.animate import AnimateFromCoeff +from src.facerender.pirender_animate import AnimateFromCoeff_PIRender from src.generate_batch import get_data from src.generate_facerender_batch import get_facerender_data from src.utils.init_path import init_path @@ -63,7 +65,12 @@ def main(args): audio_to_coeff = Audio2Coeff(sadtalker_paths, device) - animate_from_coeff = AnimateFromCoeff(sadtalker_paths, device) + if args.facerender == 'facevid2vid': + animate_from_coeff = AnimateFromCoeff(sadtalker_paths, device) + elif args.facerender == 'pirender': + animate_from_coeff = AnimateFromCoeff_PIRender(sadtalker_paths, device) + else: + raise(RuntimeError('Unknown model: {}'.format(args.facerender))) #crop image and extract 3dmm from image first_frame_dir = os.path.join(save_dir, 'first_frame_dir') @@ -108,8 +115,7 @@ def main(args): #coeff2video data = get_facerender_data(coeff_path, crop_pic_path, first_coeff_path, audio_path, batch_size, input_yaw_list, input_pitch_list, input_roll_list, - expression_scale=args.expression_scale, still_mode=args.still, preprocess=args.preprocess, size=args.size) - + expression_scale=args.expression_scale, still_mode=args.still, preprocess=args.preprocess, size=args.size, facemodel=args.facerender) result = animate_from_coeff.generate(data, save_dir, pic_path, crop_info, \ enhancer=args.enhancer, background_enhancer=args.background_enhancer, preprocess=args.preprocess, img_size=args.size) @@ -144,6 +150,7 @@ if __name__ == '__main__': parser.add_argument("--preprocess", default='crop', choices=['crop', 'extcrop', 'resize', 'full', 'extfull'], help="how to preprocess the images" ) parser.add_argument("--verbose",action="store_true", help="saving the intermedia output or not" ) parser.add_argument("--old_version",action="store_true", help="use the pth other than safetensor version" ) + parser.add_argument("--facerender", default='facevid2vid', choices=['pirender', 'facevid2vid'] ) # NB(Fakeyou): Added for fakeyou parser.add_argument("--result_file", default='./result.mp4', help="path to output file") @@ -167,6 +174,8 @@ if __name__ == '__main__': if torch.cuda.is_available() and not args.cpu: args.device = "cuda" + elif platform.system() == 'Darwin' and args.facerender == 'pirender': # macos + args.device = "mps" else: args.device = "cpu" diff --git a/animation/SadTalker/src/__pycache__/generate_batch.cpython-38.pyc b/animation/SadTalker/src/__pycache__/generate_batch.cpython-38.pyc index 5a80ca8..5c1a336 100644 Binary files a/animation/SadTalker/src/__pycache__/generate_batch.cpython-38.pyc and b/animation/SadTalker/src/__pycache__/generate_batch.cpython-38.pyc differ diff --git a/animation/SadTalker/src/__pycache__/generate_facerender_batch.cpython-38.pyc b/animation/SadTalker/src/__pycache__/generate_facerender_batch.cpython-38.pyc index 6f4e38c..ecefd25 100644 Binary files a/animation/SadTalker/src/__pycache__/generate_facerender_batch.cpython-38.pyc and b/animation/SadTalker/src/__pycache__/generate_facerender_batch.cpython-38.pyc differ diff --git a/animation/SadTalker/src/__pycache__/test_audio2coeff.cpython-38.pyc b/animation/SadTalker/src/__pycache__/test_audio2coeff.cpython-38.pyc index 2c27268..35d82b7 100644 Binary files a/animation/SadTalker/src/__pycache__/test_audio2coeff.cpython-38.pyc and b/animation/SadTalker/src/__pycache__/test_audio2coeff.cpython-38.pyc differ diff --git a/animation/SadTalker/src/audio2exp_models/__pycache__/audio2exp.cpython-38.pyc b/animation/SadTalker/src/audio2exp_models/__pycache__/audio2exp.cpython-38.pyc index e28abe8..cc248f9 100644 Binary files a/animation/SadTalker/src/audio2exp_models/__pycache__/audio2exp.cpython-38.pyc and b/animation/SadTalker/src/audio2exp_models/__pycache__/audio2exp.cpython-38.pyc differ diff --git a/animation/SadTalker/src/audio2exp_models/__pycache__/networks.cpython-38.pyc b/animation/SadTalker/src/audio2exp_models/__pycache__/networks.cpython-38.pyc index 623697a..6dad1ab 100644 Binary files a/animation/SadTalker/src/audio2exp_models/__pycache__/networks.cpython-38.pyc and b/animation/SadTalker/src/audio2exp_models/__pycache__/networks.cpython-38.pyc differ diff --git a/animation/SadTalker/src/audio2pose_models/__pycache__/audio2pose.cpython-38.pyc b/animation/SadTalker/src/audio2pose_models/__pycache__/audio2pose.cpython-38.pyc index 2be515b..2944ee2 100644 Binary files a/animation/SadTalker/src/audio2pose_models/__pycache__/audio2pose.cpython-38.pyc and b/animation/SadTalker/src/audio2pose_models/__pycache__/audio2pose.cpython-38.pyc differ diff --git a/animation/SadTalker/src/audio2pose_models/__pycache__/audio_encoder.cpython-38.pyc b/animation/SadTalker/src/audio2pose_models/__pycache__/audio_encoder.cpython-38.pyc index b6b644c..d22b3cb 100644 Binary files a/animation/SadTalker/src/audio2pose_models/__pycache__/audio_encoder.cpython-38.pyc and b/animation/SadTalker/src/audio2pose_models/__pycache__/audio_encoder.cpython-38.pyc differ diff --git a/animation/SadTalker/src/audio2pose_models/__pycache__/cvae.cpython-38.pyc b/animation/SadTalker/src/audio2pose_models/__pycache__/cvae.cpython-38.pyc index 90306be..9337815 100644 Binary files a/animation/SadTalker/src/audio2pose_models/__pycache__/cvae.cpython-38.pyc and b/animation/SadTalker/src/audio2pose_models/__pycache__/cvae.cpython-38.pyc differ diff --git a/animation/SadTalker/src/audio2pose_models/__pycache__/discriminator.cpython-38.pyc b/animation/SadTalker/src/audio2pose_models/__pycache__/discriminator.cpython-38.pyc index c993870..a8a098b 100644 Binary files a/animation/SadTalker/src/audio2pose_models/__pycache__/discriminator.cpython-38.pyc and b/animation/SadTalker/src/audio2pose_models/__pycache__/discriminator.cpython-38.pyc differ diff --git a/animation/SadTalker/src/audio2pose_models/__pycache__/networks.cpython-38.pyc b/animation/SadTalker/src/audio2pose_models/__pycache__/networks.cpython-38.pyc index 45ac321..8749bdf 100644 Binary files a/animation/SadTalker/src/audio2pose_models/__pycache__/networks.cpython-38.pyc and b/animation/SadTalker/src/audio2pose_models/__pycache__/networks.cpython-38.pyc differ diff --git a/animation/SadTalker/src/audio2pose_models/__pycache__/res_unet.cpython-38.pyc b/animation/SadTalker/src/audio2pose_models/__pycache__/res_unet.cpython-38.pyc index 177f47d..0599113 100644 Binary files a/animation/SadTalker/src/audio2pose_models/__pycache__/res_unet.cpython-38.pyc and b/animation/SadTalker/src/audio2pose_models/__pycache__/res_unet.cpython-38.pyc differ diff --git a/animation/SadTalker/src/config/facerender_pirender.yaml b/animation/SadTalker/src/config/facerender_pirender.yaml new file mode 100644 index 0000000..100e513 --- /dev/null +++ b/animation/SadTalker/src/config/facerender_pirender.yaml @@ -0,0 +1,81 @@ +# How often do you want to log the training stats. +# network_list: +# gen: gen_optimizer +# dis: dis_optimizer + +distributed: False +image_to_tensorboard: True +snapshot_save_iter: 40000 +snapshot_save_epoch: 20 +snapshot_save_start_iter: 20000 +snapshot_save_start_epoch: 10 +image_save_iter: 1000 +max_epoch: 200 +logging_iter: 100 +results_dir: ./eval_results + +gen_optimizer: + type: adam + lr: 0.0001 + adam_beta1: 0.5 + adam_beta2: 0.999 + lr_policy: + iteration_mode: True + type: step + step_size: 300000 + gamma: 0.2 + +trainer: + type: trainers.face_trainer::FaceTrainer + pretrain_warp_iteration: 200000 + loss_weight: + weight_perceptual_warp: 2.5 + weight_perceptual_final: 4 + vgg_param_warp: + network: vgg19 + layers: ['relu_1_1', 'relu_2_1', 'relu_3_1', 'relu_4_1', 'relu_5_1'] + use_style_loss: False + num_scales: 4 + vgg_param_final: + network: vgg19 + layers: ['relu_1_1', 'relu_2_1', 'relu_3_1', 'relu_4_1', 'relu_5_1'] + use_style_loss: True + num_scales: 4 + style_to_perceptual: 250 + init: + type: 'normal' + gain: 0.02 +gen: + type: generators.face_model::FaceGenerator + param: + mapping_net: + coeff_nc: 73 + descriptor_nc: 256 + layer: 3 + warpping_net: + encoder_layer: 5 + decoder_layer: 3 + base_nc: 32 + editing_net: + layer: 3 + num_res_blocks: 2 + base_nc: 64 + common: + image_nc: 3 + descriptor_nc: 256 + max_nc: 256 + use_spect: False + + +# Data options. +data: + type: data.vox_dataset::VoxDataset + path: ./dataset/vox_lmdb + resolution: 256 + semantic_radius: 13 + train: + batch_size: 5 + distributed: True + val: + batch_size: 8 + distributed: True \ No newline at end of file diff --git a/animation/SadTalker/src/face3d/__pycache__/extract_kp_videos_safe.cpython-38.pyc b/animation/SadTalker/src/face3d/__pycache__/extract_kp_videos_safe.cpython-38.pyc index 88808ea..703406e 100644 Binary files a/animation/SadTalker/src/face3d/__pycache__/extract_kp_videos_safe.cpython-38.pyc and b/animation/SadTalker/src/face3d/__pycache__/extract_kp_videos_safe.cpython-38.pyc differ diff --git a/animation/SadTalker/src/face3d/models/__pycache__/__init__.cpython-38.pyc b/animation/SadTalker/src/face3d/models/__pycache__/__init__.cpython-38.pyc index 0e378c1..168457e 100644 Binary files a/animation/SadTalker/src/face3d/models/__pycache__/__init__.cpython-38.pyc and b/animation/SadTalker/src/face3d/models/__pycache__/__init__.cpython-38.pyc differ diff --git a/animation/SadTalker/src/face3d/models/__pycache__/base_model.cpython-38.pyc b/animation/SadTalker/src/face3d/models/__pycache__/base_model.cpython-38.pyc index 9a9b589..bfde18f 100644 Binary files a/animation/SadTalker/src/face3d/models/__pycache__/base_model.cpython-38.pyc and b/animation/SadTalker/src/face3d/models/__pycache__/base_model.cpython-38.pyc differ diff --git a/animation/SadTalker/src/face3d/models/__pycache__/networks.cpython-38.pyc b/animation/SadTalker/src/face3d/models/__pycache__/networks.cpython-38.pyc index 88a7a4c..95314da 100644 Binary files a/animation/SadTalker/src/face3d/models/__pycache__/networks.cpython-38.pyc and b/animation/SadTalker/src/face3d/models/__pycache__/networks.cpython-38.pyc differ diff --git a/animation/SadTalker/src/face3d/models/arcface_torch/backbones/__pycache__/__init__.cpython-38.pyc b/animation/SadTalker/src/face3d/models/arcface_torch/backbones/__pycache__/__init__.cpython-38.pyc index ed1118e..268c595 100644 Binary files a/animation/SadTalker/src/face3d/models/arcface_torch/backbones/__pycache__/__init__.cpython-38.pyc and b/animation/SadTalker/src/face3d/models/arcface_torch/backbones/__pycache__/__init__.cpython-38.pyc differ diff --git a/animation/SadTalker/src/face3d/models/arcface_torch/backbones/__pycache__/iresnet.cpython-38.pyc b/animation/SadTalker/src/face3d/models/arcface_torch/backbones/__pycache__/iresnet.cpython-38.pyc index 7a12e0b..9f1252d 100644 Binary files a/animation/SadTalker/src/face3d/models/arcface_torch/backbones/__pycache__/iresnet.cpython-38.pyc and b/animation/SadTalker/src/face3d/models/arcface_torch/backbones/__pycache__/iresnet.cpython-38.pyc differ diff --git a/animation/SadTalker/src/face3d/models/arcface_torch/backbones/__pycache__/mobilefacenet.cpython-38.pyc b/animation/SadTalker/src/face3d/models/arcface_torch/backbones/__pycache__/mobilefacenet.cpython-38.pyc index 90d7b16..6bb49a2 100644 Binary files a/animation/SadTalker/src/face3d/models/arcface_torch/backbones/__pycache__/mobilefacenet.cpython-38.pyc and b/animation/SadTalker/src/face3d/models/arcface_torch/backbones/__pycache__/mobilefacenet.cpython-38.pyc differ diff --git a/animation/SadTalker/src/face3d/util/__pycache__/__init__.cpython-38.pyc b/animation/SadTalker/src/face3d/util/__pycache__/__init__.cpython-38.pyc index 0f84c19..e9807bb 100644 Binary files a/animation/SadTalker/src/face3d/util/__pycache__/__init__.cpython-38.pyc and b/animation/SadTalker/src/face3d/util/__pycache__/__init__.cpython-38.pyc differ diff --git a/animation/SadTalker/src/face3d/util/__pycache__/load_mats.cpython-38.pyc b/animation/SadTalker/src/face3d/util/__pycache__/load_mats.cpython-38.pyc index e4ae5d3..035313c 100644 Binary files a/animation/SadTalker/src/face3d/util/__pycache__/load_mats.cpython-38.pyc and b/animation/SadTalker/src/face3d/util/__pycache__/load_mats.cpython-38.pyc differ diff --git a/animation/SadTalker/src/face3d/util/__pycache__/my_awing_arch.cpython-38.pyc b/animation/SadTalker/src/face3d/util/__pycache__/my_awing_arch.cpython-38.pyc index c631217..a529ff0 100644 Binary files a/animation/SadTalker/src/face3d/util/__pycache__/my_awing_arch.cpython-38.pyc and b/animation/SadTalker/src/face3d/util/__pycache__/my_awing_arch.cpython-38.pyc differ diff --git a/animation/SadTalker/src/face3d/util/__pycache__/preprocess.cpython-38.pyc b/animation/SadTalker/src/face3d/util/__pycache__/preprocess.cpython-38.pyc index 78200df..0035393 100644 Binary files a/animation/SadTalker/src/face3d/util/__pycache__/preprocess.cpython-38.pyc and b/animation/SadTalker/src/face3d/util/__pycache__/preprocess.cpython-38.pyc differ diff --git a/animation/SadTalker/src/facerender/__pycache__/animate.cpython-38.pyc b/animation/SadTalker/src/facerender/__pycache__/animate.cpython-38.pyc index 4efa6fc..c907bd0 100644 Binary files a/animation/SadTalker/src/facerender/__pycache__/animate.cpython-38.pyc and b/animation/SadTalker/src/facerender/__pycache__/animate.cpython-38.pyc differ diff --git a/animation/SadTalker/src/facerender/modules/__pycache__/dense_motion.cpython-38.pyc b/animation/SadTalker/src/facerender/modules/__pycache__/dense_motion.cpython-38.pyc index c7c766a..ae4ec40 100644 Binary files a/animation/SadTalker/src/facerender/modules/__pycache__/dense_motion.cpython-38.pyc and b/animation/SadTalker/src/facerender/modules/__pycache__/dense_motion.cpython-38.pyc differ diff --git a/animation/SadTalker/src/facerender/modules/__pycache__/generator.cpython-38.pyc b/animation/SadTalker/src/facerender/modules/__pycache__/generator.cpython-38.pyc index 0741126..41c4546 100644 Binary files a/animation/SadTalker/src/facerender/modules/__pycache__/generator.cpython-38.pyc and b/animation/SadTalker/src/facerender/modules/__pycache__/generator.cpython-38.pyc differ diff --git a/animation/SadTalker/src/facerender/modules/__pycache__/keypoint_detector.cpython-38.pyc b/animation/SadTalker/src/facerender/modules/__pycache__/keypoint_detector.cpython-38.pyc index 7827245..2650cd1 100644 Binary files a/animation/SadTalker/src/facerender/modules/__pycache__/keypoint_detector.cpython-38.pyc and b/animation/SadTalker/src/facerender/modules/__pycache__/keypoint_detector.cpython-38.pyc differ diff --git a/animation/SadTalker/src/facerender/modules/__pycache__/make_animation.cpython-38.pyc b/animation/SadTalker/src/facerender/modules/__pycache__/make_animation.cpython-38.pyc index e703cf2..f6c2cb1 100644 Binary files a/animation/SadTalker/src/facerender/modules/__pycache__/make_animation.cpython-38.pyc and b/animation/SadTalker/src/facerender/modules/__pycache__/make_animation.cpython-38.pyc differ diff --git a/animation/SadTalker/src/facerender/modules/__pycache__/mapping.cpython-38.pyc b/animation/SadTalker/src/facerender/modules/__pycache__/mapping.cpython-38.pyc index ee400d0..4bbe2de 100644 Binary files a/animation/SadTalker/src/facerender/modules/__pycache__/mapping.cpython-38.pyc and b/animation/SadTalker/src/facerender/modules/__pycache__/mapping.cpython-38.pyc differ diff --git a/animation/SadTalker/src/facerender/modules/__pycache__/util.cpython-38.pyc b/animation/SadTalker/src/facerender/modules/__pycache__/util.cpython-38.pyc index e6a5622..68c299c 100644 Binary files a/animation/SadTalker/src/facerender/modules/__pycache__/util.cpython-38.pyc and b/animation/SadTalker/src/facerender/modules/__pycache__/util.cpython-38.pyc differ diff --git a/animation/SadTalker/src/facerender/pirender/base_function.py b/animation/SadTalker/src/facerender/pirender/base_function.py new file mode 100644 index 0000000..49fe4cf --- /dev/null +++ b/animation/SadTalker/src/facerender/pirender/base_function.py @@ -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 \ No newline at end of file diff --git a/animation/SadTalker/src/facerender/pirender/config.py b/animation/SadTalker/src/facerender/pirender/config.py new file mode 100644 index 0000000..3650582 --- /dev/null +++ b/animation/SadTalker/src/facerender/pirender/config.py @@ -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 \ No newline at end of file diff --git a/animation/SadTalker/src/facerender/pirender/face_model.py b/animation/SadTalker/src/facerender/pirender/face_model.py new file mode 100644 index 0000000..4a7840a --- /dev/null +++ b/animation/SadTalker/src/facerender/pirender/face_model.py @@ -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 \ No newline at end of file diff --git a/animation/SadTalker/src/facerender/pirender_animate.py b/animation/SadTalker/src/facerender/pirender_animate.py new file mode 100644 index 0000000..61a30ef --- /dev/null +++ b/animation/SadTalker/src/facerender/pirender_animate.py @@ -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 \ No newline at end of file diff --git a/animation/SadTalker/src/facerender/sync_batchnorm/__pycache__/__init__.cpython-38.pyc b/animation/SadTalker/src/facerender/sync_batchnorm/__pycache__/__init__.cpython-38.pyc index 52145db..e3b1d8b 100644 Binary files a/animation/SadTalker/src/facerender/sync_batchnorm/__pycache__/__init__.cpython-38.pyc and b/animation/SadTalker/src/facerender/sync_batchnorm/__pycache__/__init__.cpython-38.pyc differ diff --git a/animation/SadTalker/src/facerender/sync_batchnorm/__pycache__/batchnorm.cpython-38.pyc b/animation/SadTalker/src/facerender/sync_batchnorm/__pycache__/batchnorm.cpython-38.pyc index e7f0ba1..4fd223a 100644 Binary files a/animation/SadTalker/src/facerender/sync_batchnorm/__pycache__/batchnorm.cpython-38.pyc and b/animation/SadTalker/src/facerender/sync_batchnorm/__pycache__/batchnorm.cpython-38.pyc differ diff --git a/animation/SadTalker/src/facerender/sync_batchnorm/__pycache__/comm.cpython-38.pyc b/animation/SadTalker/src/facerender/sync_batchnorm/__pycache__/comm.cpython-38.pyc index 93ecaae..994f3c6 100644 Binary files a/animation/SadTalker/src/facerender/sync_batchnorm/__pycache__/comm.cpython-38.pyc and b/animation/SadTalker/src/facerender/sync_batchnorm/__pycache__/comm.cpython-38.pyc differ diff --git a/animation/SadTalker/src/facerender/sync_batchnorm/__pycache__/replicate.cpython-38.pyc b/animation/SadTalker/src/facerender/sync_batchnorm/__pycache__/replicate.cpython-38.pyc index 798aeb0..bd3fc98 100644 Binary files a/animation/SadTalker/src/facerender/sync_batchnorm/__pycache__/replicate.cpython-38.pyc and b/animation/SadTalker/src/facerender/sync_batchnorm/__pycache__/replicate.cpython-38.pyc differ diff --git a/animation/SadTalker/src/generate_facerender_batch.py b/animation/SadTalker/src/generate_facerender_batch.py index a62b6ed..680cc4a 100644 --- a/animation/SadTalker/src/generate_facerender_batch.py +++ b/animation/SadTalker/src/generate_facerender_batch.py @@ -7,7 +7,7 @@ import scipy.io as scio def get_facerender_data(coeff_path, pic_path, first_coeff_path, audio_path, batch_size, input_yaw_list=None, input_pitch_list=None, input_roll_list=None, - expression_scale=1.0, still_mode = False, preprocess='crop', size = 256): + expression_scale=1.0, still_mode = False, preprocess='crop', size = 256, facemodel='facevid2vid'): semantic_radius = 13 video_name = os.path.splitext(os.path.split(coeff_path)[-1])[0] @@ -27,7 +27,7 @@ def get_facerender_data(coeff_path, pic_path, first_coeff_path, audio_path, source_semantics_dict = scio.loadmat(first_coeff_path) generated_dict = scio.loadmat(coeff_path) - if 'full' not in preprocess.lower(): + if 'full' not in preprocess.lower() and facemodel != 'pirender': source_semantics = source_semantics_dict['coeff_3dmm'][:1,:70] #1 70 generated_3dmm = generated_dict['coeff_3dmm'][:,:70] @@ -43,7 +43,7 @@ def get_facerender_data(coeff_path, pic_path, first_coeff_path, audio_path, # target generated_3dmm[:, :64] = generated_3dmm[:, :64] * expression_scale - if 'full' in preprocess.lower(): + if 'full' in preprocess.lower() or facemodel == 'pirender': generated_3dmm = np.concatenate([generated_3dmm, np.repeat(source_semantics[:,70:], generated_3dmm.shape[0], axis=0)], axis=1) if still_mode: diff --git a/animation/SadTalker/src/gradio_demo.py b/animation/SadTalker/src/gradio_demo.py index 9c9ae05..b84c59f 100644 --- a/animation/SadTalker/src/gradio_demo.py +++ b/animation/SadTalker/src/gradio_demo.py @@ -1,5 +1,6 @@ import torch, uuid -import os, sys, shutil +import os, sys, shutil, platform +from src.facerender.pirender_animate import AnimateFromCoeff_PIRender from src.utils.preprocess import CropAndExtract from src.test_audio2coeff import Audio2Coeff from src.facerender.animate import AnimateFromCoeff @@ -20,8 +21,10 @@ class SadTalker(): def __init__(self, checkpoint_path='checkpoints', config_path='src/config', lazy_load=False): - if torch.cuda.is_available() : + if torch.cuda.is_available(): device = "cuda" + elif platform.system() == 'Darwin': # macos + device = "mps" else: device = "cpu" @@ -35,7 +38,9 @@ class SadTalker(): def test(self, source_image, driven_audio, preprocess='crop', still_mode=False, use_enhancer=False, batch_size=1, size=256, - pose_style = 0, exp_scale=1.0, + pose_style = 0, + facerender='facevid2vid', + exp_scale=1.0, use_ref_video = False, ref_video = None, ref_info = None, @@ -48,7 +53,14 @@ class SadTalker(): self.audio_to_coeff = Audio2Coeff(self.sadtalker_paths, self.device) self.preprocess_model = CropAndExtract(self.sadtalker_paths, self.device) - self.animate_from_coeff = AnimateFromCoeff(self.sadtalker_paths, self.device) + + if facerender == 'facevid2vid' and self.device != 'mps': + self.animate_from_coeff = AnimateFromCoeff(self.sadtalker_paths, self.device) + elif facerender == 'pirender' or self.device == 'mps': + self.animate_from_coeff = AnimateFromCoeff_PIRender(self.sadtalker_paths, self.device) + facerender = 'pirender' + else: + raise(RuntimeError('Unknown model: {}'.format(facerender))) time_tag = str(uuid.uuid4()) save_dir = os.path.join(result_dir, time_tag) @@ -131,11 +143,13 @@ class SadTalker(): if use_ref_video and ref_info == 'all': coeff_path = ref_video_coeff_path # self.audio_to_coeff.generate(batch, save_dir, pose_style, ref_pose_coeff_path) else: - batch = get_data(first_coeff_path, audio_path, self.device, ref_eyeblink_coeff_path=ref_eyeblink_coeff_path, still=still_mode, idlemode=use_idle_mode, length_of_audio=length_of_audio, use_blink=use_blink) # longer audio? + batch = get_data(first_coeff_path, audio_path, self.device, ref_eyeblink_coeff_path=ref_eyeblink_coeff_path, still=still_mode, \ + idlemode=use_idle_mode, length_of_audio=length_of_audio, use_blink=use_blink) # longer audio? coeff_path = self.audio_to_coeff.generate(batch, save_dir, pose_style, ref_pose_coeff_path) #coeff2video - data = get_facerender_data(coeff_path, crop_pic_path, first_coeff_path, audio_path, batch_size, still_mode=still_mode, preprocess=preprocess, size=size, expression_scale = exp_scale) + data = get_facerender_data(coeff_path, crop_pic_path, first_coeff_path, audio_path, batch_size, still_mode=still_mode, \ + preprocess=preprocess, size=size, expression_scale = exp_scale, facemodel=facerender) return_path = self.animate_from_coeff.generate(data, save_dir, pic_path, crop_info, enhancer='gfpgan' if use_enhancer else None, preprocess=preprocess, img_size=size) video_name = data['video_name'] print(f'The generated video is named {video_name} in {save_dir}') diff --git a/animation/SadTalker/src/utils/__pycache__/audio.cpython-38.pyc b/animation/SadTalker/src/utils/__pycache__/audio.cpython-38.pyc index fb84654..caa6ab4 100644 Binary files a/animation/SadTalker/src/utils/__pycache__/audio.cpython-38.pyc and b/animation/SadTalker/src/utils/__pycache__/audio.cpython-38.pyc differ diff --git a/animation/SadTalker/src/utils/__pycache__/croper.cpython-38.pyc b/animation/SadTalker/src/utils/__pycache__/croper.cpython-38.pyc index 71efb08..b9e6e4a 100644 Binary files a/animation/SadTalker/src/utils/__pycache__/croper.cpython-38.pyc and b/animation/SadTalker/src/utils/__pycache__/croper.cpython-38.pyc differ diff --git a/animation/SadTalker/src/utils/__pycache__/face_enhancer.cpython-38.pyc b/animation/SadTalker/src/utils/__pycache__/face_enhancer.cpython-38.pyc index 067fa24..06ea63c 100644 Binary files a/animation/SadTalker/src/utils/__pycache__/face_enhancer.cpython-38.pyc and b/animation/SadTalker/src/utils/__pycache__/face_enhancer.cpython-38.pyc differ diff --git a/animation/SadTalker/src/utils/__pycache__/hparams.cpython-38.pyc b/animation/SadTalker/src/utils/__pycache__/hparams.cpython-38.pyc index 7f056d2..2fe0088 100644 Binary files a/animation/SadTalker/src/utils/__pycache__/hparams.cpython-38.pyc and b/animation/SadTalker/src/utils/__pycache__/hparams.cpython-38.pyc differ diff --git a/animation/SadTalker/src/utils/__pycache__/init_path.cpython-38.pyc b/animation/SadTalker/src/utils/__pycache__/init_path.cpython-38.pyc index 36640ce..4d035aa 100644 Binary files a/animation/SadTalker/src/utils/__pycache__/init_path.cpython-38.pyc and b/animation/SadTalker/src/utils/__pycache__/init_path.cpython-38.pyc differ diff --git a/animation/SadTalker/src/utils/__pycache__/paste_pic.cpython-38.pyc b/animation/SadTalker/src/utils/__pycache__/paste_pic.cpython-38.pyc index 9cb21a3..ad7fc27 100644 Binary files a/animation/SadTalker/src/utils/__pycache__/paste_pic.cpython-38.pyc and b/animation/SadTalker/src/utils/__pycache__/paste_pic.cpython-38.pyc differ diff --git a/animation/SadTalker/src/utils/__pycache__/preprocess.cpython-38.pyc b/animation/SadTalker/src/utils/__pycache__/preprocess.cpython-38.pyc index a6975af..b56b55c 100644 Binary files a/animation/SadTalker/src/utils/__pycache__/preprocess.cpython-38.pyc and b/animation/SadTalker/src/utils/__pycache__/preprocess.cpython-38.pyc differ diff --git a/animation/SadTalker/src/utils/__pycache__/safetensor_helper.cpython-38.pyc b/animation/SadTalker/src/utils/__pycache__/safetensor_helper.cpython-38.pyc index 9061c81..e6f6ca5 100644 Binary files a/animation/SadTalker/src/utils/__pycache__/safetensor_helper.cpython-38.pyc and b/animation/SadTalker/src/utils/__pycache__/safetensor_helper.cpython-38.pyc differ diff --git a/animation/SadTalker/src/utils/__pycache__/videoio.cpython-38.pyc b/animation/SadTalker/src/utils/__pycache__/videoio.cpython-38.pyc index b124668..f9ad89e 100644 Binary files a/animation/SadTalker/src/utils/__pycache__/videoio.cpython-38.pyc and b/animation/SadTalker/src/utils/__pycache__/videoio.cpython-38.pyc differ diff --git a/animation/SadTalker/src/utils/init_path.py b/animation/SadTalker/src/utils/init_path.py index 5f38d11..18ca81e 100644 --- a/animation/SadTalker/src/utils/init_path.py +++ b/animation/SadTalker/src/utils/init_path.py @@ -35,6 +35,8 @@ def init_path(checkpoint_dir, config_dir, size=512, old_version=False, preproces sadtalker_paths['dir_of_BFM_fitting'] = os.path.join(config_dir) # , 'BFM_Fitting' sadtalker_paths['audio2pose_yaml_path'] = os.path.join(config_dir, 'auido2pose.yaml') sadtalker_paths['audio2exp_yaml_path'] = os.path.join(config_dir, 'auido2exp.yaml') + sadtalker_paths['pirender_yaml_path'] = os.path.join(config_dir, 'facerender_pirender.yaml') + sadtalker_paths['pirender_checkpoint'] = os.path.join(checkpoint_dir, 'epoch_00190_iteration_000400000_checkpoint.pt') sadtalker_paths['use_safetensor'] = use_safetensor # os.path.join(config_dir, 'auido2exp.yaml') if 'full' in preprocess: diff --git a/animation/SadTalker/src/utils/videoio.py b/animation/SadTalker/src/utils/videoio.py index 08bfbdd..d16ee66 100644 --- a/animation/SadTalker/src/utils/videoio.py +++ b/animation/SadTalker/src/utils/videoio.py @@ -19,7 +19,7 @@ def load_video_to_cv2(input_path): def save_video_with_watermark(video, audio, save_path, watermark=False): temp_file = str(uuid.uuid4())+'.mp4' - cmd = r'ffmpeg -y -hide_banner -loglevel error -i "%s" -i "%s" -vcodec copy "%s"' % (video, audio, temp_file) + cmd = r'ffmpeg -y -hide_banner -loglevel error -i "%s" -i "%s" -vcodec mpeg4 "%s"' % (video, audio, temp_file) os.system(cmd) if watermark is False: