Spaces:
Sleeping
Sleeping
| 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 as face_enhancer | |
| from src.utils.face_enhancer import enhancer_with_len as face_enhancer | |
| from src.utils.paste_pic import paste_pic | |
| from src.utils.videoio import save_video_with_watermark | |
| import time | |
| class AnimateFromCoeff(): | |
| def __init__(self, sadtalker_path, device, bf16=False): | |
| 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: # use safe tensor | |
| 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.") | |
| def optimize_cpu_model(model): | |
| return model | |
| self.kp_extractor = optimize_cpu_model(kp_extractor) | |
| self.generator = optimize_cpu_model(generator) | |
| self.he_estimator = optimize_cpu_model(he_estimator) | |
| self.mapping = optimize_cpu_model(mapping) | |
| self.kp_extractor.eval() | |
| self.generator.eval() | |
| self.he_estimator.eval() | |
| self.mapping.eval() | |
| self.bf16 = bf16 | |
| if self.bf16: | |
| import intel_extension_for_pytorch as ipex | |
| self.generator = ipex.optimize(self.generator, dtype=torch.bfloat16) | |
| 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, rank=0, p_num=1, bf16=False): | |
| 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'] | |
| with torch.autocast(device_type="cuda" if "cuda" in str(self.device) else "cpu", dtype=torch.float16): | |
| 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, rank=rank, p_num=p_num, bf16=bf16) | |
| if predictions_video is None and rank != 0: | |
| if enhancer is None: | |
| exit(0) | |
| while not os.path.exists("workspace/rendering_video.mp4"): | |
| time.sleep(0.5) | |
| with open("workspace/rendering_video.mp4", 'r') as f: | |
| full_video_path = f.readline() | |
| else: | |
| predictions_video = predictions_video.reshape((-1,)+predictions_video.shape[2:]) | |
| predictions_video = predictions_video[:frame_num] | |
| # Vectorized tensor to numpy conversion | |
| predictions_video_np = predictions_video.data.cpu().numpy() | |
| predictions_video_np = np.transpose(predictions_video_np, [0, 2, 3, 1]).astype(np.float32) | |
| video = list(predictions_video_np) | |
| 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 | |
| os.makedirs("workspace", exist_ok=True) | |
| with open("workspace/rendering_video.mp4", 'w') as f: | |
| f.write(full_video_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 | |
| enhanced_images = face_enhancer(full_video_path, method=enhancer, bg_upsampler=background_enhancer, | |
| rank=rank, p_num=p_num, bf16=bf16) | |
| imageio.mimsave(enhanced_path, enhanced_images, 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 | |