""" 用于推理 inference/train_mimictalk_on_a_video.py 得到的person-specific模型 """ import os import torch import torch.nn as nn import torch.nn.functional as F # import librosa import random import time import numpy as np import importlib import tqdm import copy import cv2 # common utils from utils.commons.hparams import hparams, set_hparams from utils.commons.tensor_utils import move_to_cuda, convert_to_tensor from utils.commons.ckpt_utils import load_ckpt, get_last_checkpoint # 3DMM-related utils from deep_3drecon.deep_3drecon_models.bfm import ParametricFaceModel from data_util.face3d_helper import Face3DHelper from data_gen.utils.process_image.fit_3dmm_landmark import fit_3dmm_for_a_image from data_gen.utils.process_video.fit_3dmm_landmark import fit_3dmm_for_a_video from deep_3drecon.secc_renderer import SECC_Renderer from data_gen.eg3d.convert_to_eg3d_convention import get_eg3d_convention_camera_pose_intrinsic # Face Parsing from data_gen.utils.mp_feature_extractors.mp_segmenter import MediapipeSegmenter from data_gen.utils.process_video.extract_segment_imgs import inpaint_torso_job, extract_background # other inference utils from inference.infer_utils import mirror_index, load_img_to_512_hwc_array, load_img_to_normalized_512_bchw_tensor from inference.infer_utils import smooth_camera_sequence, smooth_features_xd from inference.edit_secc import blink_eye_for_secc, hold_eye_opened_for_secc from inference.real3d_infer import GeneFace2Infer class AdaptGeneFace2Infer(GeneFace2Infer): def __init__(self, audio2secc_dir, head_model_dir, torso_model_dir, device=None, **kwargs): if device is None: device = 'cuda' if torch.cuda.is_available() else 'cpu' self.device = device self.audio2secc_dir = audio2secc_dir self.head_model_dir = head_model_dir self.torso_model_dir = torso_model_dir self.audio2secc_model = self.load_audio2secc(audio2secc_dir) self.secc2video_model = self.load_secc2video(head_model_dir, torso_model_dir) self.audio2secc_model.to(device).eval() self.secc2video_model.to(device).eval() self.seg_model = MediapipeSegmenter() self.secc_renderer = SECC_Renderer(512) self.face3d_helper = Face3DHelper(use_gpu=True, keypoint_mode='lm68') self.mp_face3d_helper = Face3DHelper(use_gpu=True, keypoint_mode='mediapipe') # self.camera_selector = KNearestCameraSelector() def load_secc2video(self, head_model_dir, torso_model_dir): if torso_model_dir != '': config_dir = torso_model_dir if os.path.isdir(torso_model_dir) else os.path.dirname(torso_model_dir) set_hparams(f"{config_dir}/config.yaml", print_hparams=False) hparams['htbsr_head_threshold'] = 1.0 self.secc2video_hparams = copy.deepcopy(hparams) ckpt = get_last_checkpoint(torso_model_dir)[0] lora_args = ckpt.get("lora_args", None) from modules.real3d.secc_img2plane_torso import OSAvatarSECC_Img2plane_Torso model = OSAvatarSECC_Img2plane_Torso(self.secc2video_hparams, lora_args=lora_args) load_ckpt(model, f"{torso_model_dir}", model_name='model', strict=True) self.learnable_triplane = nn.Parameter(torch.zeros([1, 3, model.triplane_hid_dim*model.triplane_depth, 256, 256]).float().cuda(), requires_grad=True) load_ckpt(self.learnable_triplane, f"{torso_model_dir}", model_name='learnable_triplane', strict=True) model._last_cano_planes = self.learnable_triplane if head_model_dir != '': print("| Warning: Assigned --torso_ckpt which also contains head, but --head_ckpt is also assigned, skipping the --head_ckpt.") else: from modules.real3d.secc_img2plane_torso import OSAvatarSECC_Img2plane set_hparams(f"{head_model_dir}/config.yaml", print_hparams=False) ckpt = get_last_checkpoint(head_model_dir)[0] lora_args = ckpt.get("lora_args", None) self.secc2video_hparams = copy.deepcopy(hparams) model = OSAvatarSECC_Img2plane(self.secc2video_hparams, lora_args=lora_args) load_ckpt(model, f"{head_model_dir}", model_name='model', strict=True) self.learnable_triplane = nn.Parameter(torch.zeros([1, 3, model.triplane_hid_dim*model.triplane_depth, 256, 256]).float().cuda(), requires_grad=True) model._last_cano_planes = self.learnable_triplane load_ckpt(model._last_cano_planes, f"{head_model_dir}", model_name='learnable_triplane', strict=True) self.person_ds = ckpt['person_ds'] return model def prepare_batch_from_inp(self, inp): """ :param inp: {'audio_source_name': (str)} :return: a dict that contains the condition feature of NeRF """ sample = {} # Process Driving Motion if inp['drv_audio_name'][-4:] in ['.wav', '.mp3']: self.save_wav16k(inp['drv_audio_name']) if self.audio2secc_hparams['audio_type'] == 'hubert': hubert = self.get_hubert(self.wav16k_name) elif self.audio2secc_hparams['audio_type'] == 'mfcc': hubert = self.get_mfcc(self.wav16k_name) / 100 f0 = self.get_f0(self.wav16k_name) if f0.shape[0] > len(hubert): f0 = f0[:len(hubert)] else: num_to_pad = len(hubert) - len(f0) f0 = np.pad(f0, pad_width=((0,num_to_pad), (0,0))) t_x = hubert.shape[0] x_mask = torch.ones([1, t_x]).float() # mask for audio frames y_mask = torch.ones([1, t_x//2]).float() # mask for motion/image frames sample.update({ 'hubert': torch.from_numpy(hubert).float().unsqueeze(0).cuda(), 'f0': torch.from_numpy(f0).float().reshape([1,-1]).cuda(), 'x_mask': x_mask.cuda(), 'y_mask': y_mask.cuda(), }) sample['blink'] = torch.zeros([1, t_x, 1]).long().cuda() sample['audio'] = sample['hubert'] sample['eye_amp'] = torch.ones([1, 1]).cuda() * 1.0 elif inp['drv_audio_name'][-4:] in ['.mp4']: drv_motion_coeff_dict = fit_3dmm_for_a_video(inp['drv_audio_name'], save=False) drv_motion_coeff_dict = convert_to_tensor(drv_motion_coeff_dict) t_x = drv_motion_coeff_dict['exp'].shape[0] * 2 self.drv_motion_coeff_dict = drv_motion_coeff_dict elif inp['drv_audio_name'][-4:] in ['.npy']: drv_motion_coeff_dict = np.load(inp['drv_audio_name'], allow_pickle=True).tolist() drv_motion_coeff_dict = convert_to_tensor(drv_motion_coeff_dict) t_x = drv_motion_coeff_dict['exp'].shape[0] * 2 self.drv_motion_coeff_dict = drv_motion_coeff_dict # Face Parsing sample['ref_gt_img'] = self.person_ds['gt_img'].cuda() img = self.person_ds['gt_img'].reshape([3, 512, 512]).permute(1, 2, 0) img = (img + 1) * 127.5 img = np.ascontiguousarray(img.int().numpy()).astype(np.uint8) segmap = self.seg_model._cal_seg_map(img) sample['segmap'] = torch.tensor(segmap).float().unsqueeze(0).cuda() head_img = self.seg_model._seg_out_img_with_segmap(img, segmap, mode='head')[0] sample['ref_head_img'] = ((torch.tensor(head_img) - 127.5)/127.5).float().unsqueeze(0).permute(0, 3, 1,2).cuda() # [b,c,h,w] inpaint_torso_img, _, _, _ = inpaint_torso_job(img, segmap) sample['ref_torso_img'] = ((torch.tensor(inpaint_torso_img) - 127.5)/127.5).float().unsqueeze(0).permute(0, 3, 1,2).cuda() # [b,c,h,w] if inp['bg_image_name'] == '': bg_img = extract_background([img], [segmap], 'knn') else: bg_img = cv2.imread(inp['bg_image_name']) bg_img = cv2.cvtColor(bg_img, cv2.COLOR_BGR2RGB) bg_img = cv2.resize(bg_img, (512,512)) sample['bg_img'] = ((torch.tensor(bg_img) - 127.5)/127.5).float().unsqueeze(0).permute(0, 3, 1,2).cuda() # [b,c,h,w] # 3DMM, get identity code and camera pose image_name = f"data/raw/val_imgs/{self.person_ds['video_id']}_img.png" os.makedirs(os.path.dirname(image_name), exist_ok=True) cv2.imwrite(image_name, img[:,:,::-1]) coeff_dict = fit_3dmm_for_a_image(image_name, save=False) coeff_dict['id'] = self.person_ds['id'].reshape([1,80]).numpy() assert coeff_dict is not None src_id = torch.tensor(coeff_dict['id']).reshape([1,80]).cuda() src_exp = torch.tensor(coeff_dict['exp']).reshape([1,64]).cuda() src_euler = torch.tensor(coeff_dict['euler']).reshape([1,3]).cuda() src_trans = torch.tensor(coeff_dict['trans']).reshape([1,3]).cuda() sample['id'] = src_id.repeat([t_x//2,1]) # get the src_kp for torso model sample['src_kp'] = self.person_ds['src_kp'].cuda().reshape([1, 68, 3]).repeat([t_x//2,1,1])[..., :2] # [B, 68, 2] # get camera pose file random.seed(time.time()) if inp['drv_pose_name'] in ['nearest', 'topk']: camera_ret = get_eg3d_convention_camera_pose_intrinsic({'euler': torch.tensor(coeff_dict['euler']).reshape([1,3]), 'trans': torch.tensor(coeff_dict['trans']).reshape([1,3])}) c2w, intrinsics = camera_ret['c2w'], camera_ret['intrinsics'] camera = np.concatenate([c2w.reshape([1,16]), intrinsics.reshape([1,9])], axis=-1) coeff_names, distance_matrix = self.camera_selector.find_k_nearest(camera, k=100) coeff_names = coeff_names[0] # squeeze if inp['drv_pose_name'] == 'nearest': inp['drv_pose_name'] = coeff_names[0] else: inp['drv_pose_name'] = random.choice(coeff_names) # inp['drv_pose_name'] = coeff_names[0] elif inp['drv_pose_name'] == 'random': inp['drv_pose_name'] = self.camera_selector.random_select() else: inp['drv_pose_name'] = inp['drv_pose_name'] print(f"| To extract pose from {inp['drv_pose_name']}") # extract camera pose if inp['drv_pose_name'] == 'static': sample['euler'] = torch.tensor(coeff_dict['euler']).reshape([1,3]).cuda().repeat([t_x//2,1]) # default static pose sample['trans'] = torch.tensor(coeff_dict['trans']).reshape([1,3]).cuda().repeat([t_x//2,1]) else: # from file if inp['drv_pose_name'].endswith('.mp4'): # extract coeff from video drv_pose_coeff_dict = fit_3dmm_for_a_video(inp['drv_pose_name'], save=False) else: # load from npy drv_pose_coeff_dict = np.load(inp['drv_pose_name'], allow_pickle=True).tolist() print(f"| Extracted pose from {inp['drv_pose_name']}") eulers = convert_to_tensor(drv_pose_coeff_dict['euler']).reshape([-1,3]).cuda() trans = convert_to_tensor(drv_pose_coeff_dict['trans']).reshape([-1,3]).cuda() len_pose = len(eulers) index_lst = [mirror_index(i, len_pose) for i in range(t_x//2)] sample['euler'] = eulers[index_lst] sample['trans'] = trans[index_lst] # fix the z axis sample['trans'][:, -1] = sample['trans'][0:1, -1].repeat([sample['trans'].shape[0]]) # mapping to the init pose if inp.get("map_to_init_pose", 'False') == 'True': diff_euler = torch.tensor(coeff_dict['euler']).reshape([1,3]).cuda() - sample['euler'][0:1] sample['euler'] = sample['euler'] + diff_euler diff_trans = torch.tensor(coeff_dict['trans']).reshape([1,3]).cuda() - sample['trans'][0:1] sample['trans'] = sample['trans'] + diff_trans # prepare camera camera_ret = get_eg3d_convention_camera_pose_intrinsic({'euler':sample['euler'].cpu(), 'trans':sample['trans'].cpu()}) c2w, intrinsics = camera_ret['c2w'], camera_ret['intrinsics'] # smooth camera camera_smo_ksize = 7 camera = np.concatenate([c2w.reshape([-1,16]), intrinsics.reshape([-1,9])], axis=-1) camera = smooth_camera_sequence(camera, kernel_size=camera_smo_ksize) # [T, 25] camera = torch.tensor(camera).cuda().float() sample['camera'] = camera return sample @torch.no_grad() def forward_secc2video(self, batch, inp=None): num_frames = len(batch['drv_secc']) camera = batch['camera'] src_kps = batch['src_kp'] drv_kps = batch['drv_kp'] cano_secc_color = batch['cano_secc'] src_secc_color = batch['src_secc'] drv_secc_colors = batch['drv_secc'] ref_img_gt = batch['ref_gt_img'] ref_img_head = batch['ref_head_img'] ref_torso_img = batch['ref_torso_img'] bg_img = batch['bg_img'] segmap = batch['segmap'] # smooth torso drv_kp torso_smo_ksize = 7 drv_kps = smooth_features_xd(drv_kps.reshape([-1, 68*2]), kernel_size=torso_smo_ksize).reshape([-1, 68, 2]) # forward renderer img_raw_lst = [] img_lst = [] depth_img_lst = [] with torch.no_grad(): for i in tqdm.trange(num_frames, desc="MimicTalk is rendering frames"): kp_src = torch.cat([src_kps[i:i+1].reshape([1, 68, 2]), torch.zeros([1, 68,1]).to(src_kps.device)],dim=-1) kp_drv = torch.cat([drv_kps[i:i+1].reshape([1, 68, 2]), torch.zeros([1, 68,1]).to(drv_kps.device)],dim=-1) cond={'cond_cano': cano_secc_color,'cond_src': src_secc_color, 'cond_tgt': drv_secc_colors[i:i+1].cuda(), 'ref_torso_img': ref_torso_img, 'bg_img': bg_img, 'segmap': segmap, 'kp_s': kp_src, 'kp_d': kp_drv} ######################################################################################################## ### 相比real3d_infer只修改了这行👇，即cano_triplane来自cache里的learnable_triplane,而不是img预测的plane #### ######################################################################################################## gen_output = self.secc2video_model.forward(img=None, camera=camera[i:i+1], cond=cond, ret={}, cache_backbone=False, use_cached_backbone=True) img_lst.append(gen_output['image']) img_raw_lst.append(gen_output['image_raw']) depth_img_lst.append(gen_output['image_depth']) # save demo video depth_imgs = torch.cat(depth_img_lst) imgs = torch.cat(img_lst) imgs_raw = torch.cat(img_raw_lst) secc_img = torch.cat([torch.nn.functional.interpolate(drv_secc_colors[i:i+1], (512,512)) for i in range(num_frames)]) if inp['out_mode'] == 'concat_debug': secc_img = secc_img.cpu() secc_img = ((secc_img + 1) * 127.5).permute(0, 2, 3, 1).int().numpy() depth_img = F.interpolate(depth_imgs, (512,512)).cpu() depth_img = depth_img.repeat([1,3,1,1]) depth_img = (depth_img - depth_img.min()) / (depth_img.max() - depth_img.min()) depth_img = depth_img * 2 - 1 depth_img = depth_img.clamp(-1,1) secc_img = secc_img / 127.5 - 1 secc_img = torch.from_numpy(secc_img).permute(0, 3, 1, 2) imgs = torch.cat([ref_img_gt.repeat([imgs.shape[0],1,1,1]).cpu(), secc_img, F.interpolate(imgs_raw, (512,512)).cpu(), depth_img, imgs.cpu()], dim=-1) elif inp['out_mode'] == 'final': imgs = imgs.cpu() elif inp['out_mode'] == 'debug': raise NotImplementedError("to do: save separate videos") imgs = imgs.clamp(-1,1) import imageio import uuid debug_name = f'{uuid.uuid1()}.mp4' out_imgs = ((imgs.permute(0, 2, 3, 1) + 1)/2 * 255).int().cpu().numpy().astype(np.uint8) writer = imageio.get_writer(debug_name, fps=25, format='FFMPEG', codec='h264') for i in tqdm.trange(len(out_imgs), desc="Imageio is saving video"): writer.append_data(out_imgs[i]) writer.close() out_fname = 'infer_out/tmp/' + os.path.basename(inp['drv_pose_name'])[:-4] + '.mp4' if inp['out_name'] == '' else inp['out_name'] try: os.makedirs(os.path.dirname(out_fname), exist_ok=True) except: pass if inp['drv_audio_name'][-4:] in ['.wav', '.mp3']: # os.system(f"ffmpeg -i {debug_name} -i {inp['drv_audio_name']} -y -v quiet -shortest {out_fname}") cmd = f"/usr/bin/ffmpeg -i {debug_name} -i {self.wav16k_name} -y -r 25 -ar 16000 -c:v copy -c:a libmp3lame -pix_fmt yuv420p -b:v 2000k -strict experimental -shortest {out_fname}" os.system(cmd) os.system(f"rm {debug_name}") else: ret = os.system(f"ffmpeg -i {debug_name} -i {inp['drv_audio_name']} -map 0:v -map 1:a -y -v quiet -shortest {out_fname}") if ret != 0: # 没有成功从drv_audio_name里面提取到音频, 则直接输出无音频轨道的纯视频 os.system(f"mv {debug_name} {out_fname}") print(f"Saved at {out_fname}") return out_fname if __name__ == '__main__': import argparse, glob, tqdm parser = argparse.ArgumentParser() parser.add_argument("--a2m_ckpt", default='checkpoints/240112_icl_audio2secc_vox2_cmlr') # checkpoints/0727_audio2secc/audio2secc_withlm2d100_randomframe parser.add_argument("--head_ckpt", default='') # checkpoints/0729_th1kh/secc_img2plane checkpoints/0720_img2planes/secc_img2plane_two_stage parser.add_argument("--torso_ckpt", default='checkpoints_mimictalk/German_20s') parser.add_argument("--bg_img", default='') # data/raw/val_imgs/bg3.png parser.add_argument("--drv_aud", default='data/raw/examples/80_vs_60_10s.wav') parser.add_argument("--drv_pose", default='data/raw/examples/German_20s.mp4') # nearest | topk | random | static | vid_name parser.add_argument("--drv_style", default='data/raw/examples/angry.mp4') # nearest | topk | random | static | vid_name parser.add_argument("--blink_mode", default='period') # none | period parser.add_argument("--temperature", default=0.3, type=float) # nearest | random parser.add_argument("--denoising_steps", default=20, type=int) # nearest | random parser.add_argument("--cfg_scale", default=1.5, type=float) # nearest | random parser.add_argument("--out_name", default='') # nearest | random parser.add_argument("--out_mode", default='concat_debug') # concat_debug | debug | final parser.add_argument("--hold_eye_opened", default='False') # concat_debug | debug | final parser.add_argument("--map_to_init_pose", default='True') # concat_debug | debug | final parser.add_argument("--seed", default=None, type=int) # random seed, default None to use time.time() args = parser.parse_args() inp = { 'a2m_ckpt': args.a2m_ckpt, 'head_ckpt': args.head_ckpt, 'torso_ckpt': args.torso_ckpt, 'bg_image_name': args.bg_img, 'drv_audio_name': args.drv_aud, 'drv_pose_name': args.drv_pose, 'drv_talking_style_name': args.drv_style, 'blink_mode': args.blink_mode, 'temperature': args.temperature, 'denoising_steps': args.denoising_steps, 'cfg_scale': args.cfg_scale, 'out_name': args.out_name, 'out_mode': args.out_mode, 'map_to_init_pose': args.map_to_init_pose, 'hold_eye_opened': args.hold_eye_opened, 'seed': args.seed, } AdaptGeneFace2Infer.example_run(inp)