| import os
|
| import cv2
|
| 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
|
| import torchvision
|
|
|
|
|
| from src.facerender.modules.keypoint_detector import HEEstimator, KPDetector
|
| from src.facerender.modules.mapping import MappingNet
|
| from src.facerender.modules.generator import OcclusionAwareGenerator, OcclusionAwareSPADEGenerator
|
| from src.facerender.modules.make_animation import make_animation
|
|
|
| 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 = True
|
| except:
|
| in_webui = False
|
|
|
| class AnimateFromCoeff():
|
|
|
| def __init__(self, sadtalker_path, device):
|
|
|
| with open(sadtalker_path['facerender_yaml']) as f:
|
| config = yaml.safe_load(f)
|
|
|
| generator = OcclusionAwareSPADEGenerator(**config['model_params']['generator_params'],
|
| **config['model_params']['common_params'])
|
| kp_extractor = KPDetector(**config['model_params']['kp_detector_params'],
|
| **config['model_params']['common_params'])
|
| he_estimator = HEEstimator(**config['model_params']['he_estimator_params'],
|
| **config['model_params']['common_params'])
|
| mapping = MappingNet(**config['model_params']['mapping_params'])
|
|
|
| generator.to(device)
|
| kp_extractor.to(device)
|
| he_estimator.to(device)
|
| mapping.to(device)
|
| for param in generator.parameters():
|
| param.requires_grad = False
|
| for param in kp_extractor.parameters():
|
| param.requires_grad = False
|
| for param in he_estimator.parameters():
|
| param.requires_grad = False
|
| for param in mapping.parameters():
|
| param.requires_grad = False
|
|
|
| if sadtalker_path is not None:
|
| if 'checkpoint' in sadtalker_path:
|
| self.load_cpk_facevid2vid_safetensor(sadtalker_path['checkpoint'], kp_detector=kp_extractor, generator=generator, he_estimator=None)
|
| else:
|
| self.load_cpk_facevid2vid(sadtalker_path['free_view_checkpoint'], kp_detector=kp_extractor, generator=generator, he_estimator=he_estimator)
|
| else:
|
| raise AttributeError("Checkpoint should be specified for video head pose estimator.")
|
|
|
| if sadtalker_path['mappingnet_checkpoint'] is not None:
|
| self.load_cpk_mapping(sadtalker_path['mappingnet_checkpoint'], mapping=mapping)
|
| else:
|
| raise AttributeError("Checkpoint should be specified for video head pose estimator.")
|
|
|
| self.kp_extractor = kp_extractor
|
| self.generator = generator
|
| self.he_estimator = he_estimator
|
| self.mapping = mapping
|
|
|
| self.kp_extractor.eval()
|
| self.generator.eval()
|
| self.he_estimator.eval()
|
| self.mapping.eval()
|
|
|
| self.device = device
|
|
|
| def load_cpk_facevid2vid_safetensor(self, checkpoint_path, generator=None,
|
| kp_detector=None, he_estimator=None,
|
| device="cpu"):
|
|
|
| checkpoint = safetensors.torch.load_file(checkpoint_path)
|
|
|
| if generator is not None:
|
| x_generator = {}
|
| for k,v in checkpoint.items():
|
| if 'generator' in k:
|
| x_generator[k.replace('generator.', '')] = v
|
| generator.load_state_dict(x_generator)
|
| if kp_detector is not None:
|
| x_generator = {}
|
| for k,v in checkpoint.items():
|
| if 'kp_extractor' in k:
|
| x_generator[k.replace('kp_extractor.', '')] = v
|
| kp_detector.load_state_dict(x_generator)
|
| if he_estimator is not None:
|
| x_generator = {}
|
| for k,v in checkpoint.items():
|
| if 'he_estimator' in k:
|
| x_generator[k.replace('he_estimator.', '')] = v
|
| he_estimator.load_state_dict(x_generator)
|
|
|
| return None
|
|
|
| def load_cpk_facevid2vid(self, checkpoint_path, generator=None, discriminator=None,
|
| kp_detector=None, he_estimator=None, optimizer_generator=None,
|
| optimizer_discriminator=None, optimizer_kp_detector=None,
|
| optimizer_he_estimator=None, device="cpu"):
|
| checkpoint = torch.load(checkpoint_path, map_location=torch.device(device))
|
| if generator is not None:
|
| generator.load_state_dict(checkpoint['generator'])
|
| if kp_detector is not None:
|
| kp_detector.load_state_dict(checkpoint['kp_detector'])
|
| if he_estimator is not None:
|
| he_estimator.load_state_dict(checkpoint['he_estimator'])
|
| if discriminator is not None:
|
| try:
|
| discriminator.load_state_dict(checkpoint['discriminator'])
|
| except:
|
| print ('No discriminator in the state-dict. Dicriminator will be randomly initialized')
|
| if optimizer_generator is not None:
|
| optimizer_generator.load_state_dict(checkpoint['optimizer_generator'])
|
| if optimizer_discriminator is not None:
|
| try:
|
| optimizer_discriminator.load_state_dict(checkpoint['optimizer_discriminator'])
|
| except RuntimeError as e:
|
| print ('No discriminator optimizer in the state-dict. Optimizer will be not initialized')
|
| if optimizer_kp_detector is not None:
|
| optimizer_kp_detector.load_state_dict(checkpoint['optimizer_kp_detector'])
|
| if optimizer_he_estimator is not None:
|
| optimizer_he_estimator.load_state_dict(checkpoint['optimizer_he_estimator'])
|
|
|
| return checkpoint['epoch']
|
|
|
| def load_cpk_mapping(self, checkpoint_path, mapping=None, discriminator=None,
|
| optimizer_mapping=None, optimizer_discriminator=None, device='cpu'):
|
| checkpoint = torch.load(checkpoint_path, map_location=torch.device(device))
|
| if mapping is not None:
|
| mapping.load_state_dict(checkpoint['mapping'])
|
| if discriminator is not None:
|
| discriminator.load_state_dict(checkpoint['discriminator'])
|
| if optimizer_mapping is not None:
|
| optimizer_mapping.load_state_dict(checkpoint['optimizer_mapping'])
|
| if optimizer_discriminator is not None:
|
| optimizer_discriminator.load_state_dict(checkpoint['optimizer_discriminator'])
|
|
|
| return checkpoint['epoch']
|
|
|
| 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)
|
| if 'yaw_c_seq' in x:
|
| yaw_c_seq = x['yaw_c_seq'].type(torch.FloatTensor)
|
| yaw_c_seq = x['yaw_c_seq'].to(self.device)
|
| else:
|
| yaw_c_seq = None
|
| if 'pitch_c_seq' in x:
|
| pitch_c_seq = x['pitch_c_seq'].type(torch.FloatTensor)
|
| pitch_c_seq = x['pitch_c_seq'].to(self.device)
|
| else:
|
| pitch_c_seq = None
|
| if 'roll_c_seq' in x:
|
| roll_c_seq = x['roll_c_seq'].type(torch.FloatTensor)
|
| roll_c_seq = x['roll_c_seq'].to(self.device)
|
| else:
|
| roll_c_seq = None
|
|
|
| frame_num = x['frame_num']
|
|
|
| predictions_video = make_animation(source_image, source_semantics, target_semantics,
|
| self.generator, self.kp_extractor, self.he_estimator, self.mapping,
|
| yaw_c_seq, pitch_c_seq, roll_c_seq, use_exp = True)
|
|
|
| predictions_video = predictions_video.reshape((-1,)+predictions_video.shape[2:])
|
| predictions_video = predictions_video[:frame_num]
|
|
|
| video = []
|
| for idx in range(predictions_video.shape[0]):
|
| 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)
|
|
|
|
|
| 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
|
|
|
| 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():
|
|
|
| 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
|
|
|
|
|
| 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
|
|
|
|
|
