inference/real3d_infer.py

import os
import sys
sys.path.append('./')
import torch
import torch.nn.functional as F
import torchshow as ts
import librosa
import random
import time
import numpy as np
import importlib
import tqdm
import copy
import cv2
import math

# 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
from data_gen.utils.process_image.extract_lm2d import extract_lms_mediapipe_job

# 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


def read_first_frame_from_a_video(vid_name):
 frames = []
 cap = cv2.VideoCapture(vid_name)
 ret, frame_bgr = cap.read()
 frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
 return frame_rgb

def analyze_weights_img(gen_output):
 img_raw = gen_output['image_raw']
 mask_005_to_03 = torch.bitwise_and(gen_output['weights_img']>0.05, gen_output['weights_img']<0.3).repeat([1,3,1,1])
 mask_005_to_05 = torch.bitwise_and(gen_output['weights_img']>0.05, gen_output['weights_img']<0.5).repeat([1,3,1,1])
 mask_005_to_07 = torch.bitwise_and(gen_output['weights_img']>0.05, gen_output['weights_img']<0.7).repeat([1,3,1,1])
 mask_005_to_09 = torch.bitwise_and(gen_output['weights_img']>0.05, gen_output['weights_img']<0.9).repeat([1,3,1,1])
 mask_005_to_10 = torch.bitwise_and(gen_output['weights_img']>0.05, gen_output['weights_img']<1.0).repeat([1,3,1,1])

 img_raw_005_to_03 = img_raw.clone()
 img_raw_005_to_03[~mask_005_to_03] = -1
 img_raw_005_to_05 = img_raw.clone()
 img_raw_005_to_05[~mask_005_to_05] = -1
 img_raw_005_to_07 = img_raw.clone()
 img_raw_005_to_07[~mask_005_to_07] = -1
 img_raw_005_to_09 = img_raw.clone()
 img_raw_005_to_09[~mask_005_to_09] = -1
 img_raw_005_to_10 = img_raw.clone()
 img_raw_005_to_10[~mask_005_to_10] = -1
 ts.save([img_raw_005_to_03[0], img_raw_005_to_05[0], img_raw_005_to_07[0], img_raw_005_to_09[0], img_raw_005_to_10[0]])

def cal_face_area_percent(img_name):
 img = cv2.resize(cv2.imread(img_name)[:,:,::-1], (512,512))
 lm478 = extract_lms_mediapipe_job(img) / 512
 min_x = lm478[:,0].min()
 max_x = lm478[:,0].max()
 min_y = lm478[:,1].min()
 max_y = lm478[:,1].max()
 area = (max_x - min_x) * (max_y - min_y)
 return area

def crop_img_on_face_area_percent(img_name, out_name='temp/cropped_src_img.png', min_face_area_percent=0.2):
 try:
 os.makedirs(os.path.dirname(out_name), exist_ok=True)
 except: pass
 face_area_percent = cal_face_area_percent(img_name)
 if face_area_percent >= min_face_area_percent:
 print(f"face area percent {face_area_percent} larger than threshold {min_face_area_percent}, directly use the input image...")
 cmd = f"cp {img_name} {out_name}"
 os.system(cmd)
 return out_name
 else:
 print(f"face area percent {face_area_percent} smaller than threshold {min_face_area_percent}, crop the input image...")
 img = cv2.resize(cv2.imread(img_name)[:,:,::-1], (512,512))
 lm478 = extract_lms_mediapipe_job(img).astype(int)
 min_x = lm478[:,0].min()
 max_x = lm478[:,0].max()
 min_y = lm478[:,1].min()
 max_y = lm478[:,1].max()
 face_area = (max_x - min_x) * (max_y - min_y)
 target_total_area = face_area / min_face_area_percent
 target_hw = int(target_total_area**0.5)
 center_x, center_y = (min_x+max_x)/2, (min_y+max_y)/2
 shrink_pixels = 2 * max(-(center_x - target_hw/2), center_x + target_hw/2 - 512, -(center_y - target_hw/2), center_y + target_hw/2-512)
 shrink_pixels = max(0, shrink_pixels)
 hw = math.floor(target_hw - shrink_pixels)
 new_min_x = int(center_x - hw/2)
 new_max_x = int(center_x + hw/2)
 new_min_y = int(center_y - hw/2)
 new_max_y = int(center_y + hw/2)

 img = img[new_min_y:new_max_y, new_min_x:new_max_x]
 img = cv2.resize(img, (512, 512))
 cv2.imwrite(out_name, img[:,:,::-1])
 return out_name


class GeneFace2Infer:
 def __init__(self, audio2secc_dir, head_model_dir, torso_model_dir, device=None, inp=None):
 if device is None:
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 self.device = device
 self.audio2secc_model = self.load_audio2secc(audio2secc_dir)
 self.secc2video_model = self.load_secc2video(head_model_dir, torso_model_dir, inp)
 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')

 def load_audio2secc(self, audio2secc_dir):
 config_name = f"{audio2secc_dir}/config.yaml" if not audio2secc_dir.endswith(".ckpt") else f"{os.path.dirname(audio2secc_dir)}/config.yaml"
 set_hparams(f"{config_name}", print_hparams=False)
 self.audio2secc_dir = audio2secc_dir
 self.audio2secc_hparams = copy.deepcopy(hparams)
 from modules.audio2motion.vae import VAEModel, PitchContourVAEModel
 if self.audio2secc_hparams['audio_type'] == 'hubert':
 audio_in_dim = 1024
 elif self.audio2secc_hparams['audio_type'] == 'mfcc':
 audio_in_dim = 13

 if 'icl' in hparams['task_cls']:
 self.use_icl_audio2motion = True
 model = InContextAudio2MotionModel(hparams['icl_model_type'], hparams=self.audio2secc_hparams)
 else:
 self.use_icl_audio2motion = False
 if hparams.get("use_pitch", False) is True:
 model = PitchContourVAEModel(hparams, in_out_dim=64, audio_in_dim=audio_in_dim)
 else:
 model = VAEModel(in_out_dim=64, audio_in_dim=audio_in_dim)
 load_ckpt(model, f"{audio2secc_dir}", model_name='model', strict=True)
 return model

 def load_secc2video(self, head_model_dir, torso_model_dir, inp):
 if inp is None:
 inp = {}
 self.head_model_dir = head_model_dir
 self.torso_model_dir = torso_model_dir
 if torso_model_dir != '':
 if torso_model_dir.endswith(".ckpt"):
 set_hparams(f"{os.path.dirname(torso_model_dir)}/config.yaml", print_hparams=False)
 else:
 set_hparams(f"{torso_model_dir}/config.yaml", print_hparams=False)
 if inp.get('head_torso_threshold', None) is not None:
 hparams['htbsr_head_threshold'] = inp['head_torso_threshold']
 self.secc2video_hparams = copy.deepcopy(hparams)
 from modules.real3d.secc_img2plane_torso import OSAvatarSECC_Img2plane_Torso
 model = OSAvatarSECC_Img2plane_Torso()
 load_ckpt(model, f"{torso_model_dir}", model_name='model', strict=True)
 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
 if head_model_dir.endswith(".ckpt"):
 set_hparams(f"{os.path.dirname(head_model_dir)}/config.yaml", print_hparams=False)
 else:
 set_hparams(f"{head_model_dir}/config.yaml", print_hparams=False)
 if inp.get('head_torso_threshold', None) is not None:
 hparams['htbsr_head_threshold'] = inp['head_torso_threshold']
 self.secc2video_hparams = copy.deepcopy(hparams)
 model = OSAvatarSECC_Img2plane()
 load_ckpt(model, f"{head_model_dir}", model_name='model', strict=True)
 return model

 def infer_once(self, inp):
 self.inp = inp
 samples = self.prepare_batch_from_inp(inp)
 seed = inp['seed'] if inp['seed'] is not None else int(time.time())
 random.seed(seed)
 torch.manual_seed(seed)
 np.random.seed(seed)
 out_name = self.forward_system(samples, inp)
 return out_name
 
 def prepare_batch_from_inp(self, inp):
 """
 :param inp: {'audio_source_name': (str)}
 :return: a dict that contains the condition feature of NeRF
 """
 tmp_img_name = 'infer_out/tmp/cropped_src_img.png'
 crop_img_on_face_area_percent(inp['src_image_name'], tmp_img_name, min_face_area_percent=inp['min_face_area_percent'])
 inp['src_image_name'] = tmp_img_name

 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
 sample['mouth_amp'] = torch.ones([1, 1]).cuda() * inp['mouth_amp']
 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
 image_name = inp['src_image_name']
 if image_name.endswith(".mp4"):
 img = read_first_frame_from_a_video(image_name)
 image_name = inp['src_image_name'] = image_name[:-4] + '.png'
 cv2.imwrite(image_name, cv2.cvtColor(img, cv2.COLOR_RGB2BGR))
 sample['ref_gt_img'] = load_img_to_normalized_512_bchw_tensor(image_name).cuda()
 img = load_img_to_512_hwc_array(image_name)
 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]
 ts.save(sample['ref_head_img'])
 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
 coeff_dict = fit_3dmm_for_a_image(image_name, save=False)
 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
 src_kp = self.face3d_helper.reconstruct_lm2d(src_id, src_exp, src_euler, src_trans) # [1, 68, 2]
 src_kp = (src_kp-0.5) / 0.5 # rescale to -1~1
 sample['src_kp'] = torch.clamp(src_kp, -1, 1).repeat([t_x//2,1,1])

 # get camera pose file
 # random.seed(time.time())
 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
 print(inp)
 if inp.get("map_to_init_pose", 'True') in ['True', 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 get_hubert(self, wav16k_name):
 from data_gen.utils.process_audio.extract_hubert import get_hubert_from_16k_wav
 hubert = get_hubert_from_16k_wav(wav16k_name).detach().numpy()
 len_mel = hubert.shape[0]
 x_multiply = 8
 if len_mel % x_multiply == 0:
 num_to_pad = 0
 else:
 num_to_pad = x_multiply - len_mel % x_multiply
 hubert = np.pad(hubert, pad_width=((0,num_to_pad), (0,0)))
 return hubert

 def get_mfcc(self, wav16k_name):
 from utils.audio import librosa_wav2mfcc
 hparams['fft_size'] = 1200
 hparams['win_size'] = 1200
 hparams['hop_size'] = 480
 hparams['audio_num_mel_bins'] = 80
 hparams['fmin'] = 80
 hparams['fmax'] = 12000
 hparams['audio_sample_rate'] = 24000
 mfcc = librosa_wav2mfcc(wav16k_name,
 fft_size=hparams['fft_size'],
 hop_size=hparams['hop_size'],
 win_length=hparams['win_size'],
 num_mels=hparams['audio_num_mel_bins'],
 fmin=hparams['fmin'],
 fmax=hparams['fmax'],
 sample_rate=hparams['audio_sample_rate'],
 center=True)
 mfcc = np.array(mfcc).reshape([-1, 13])
 len_mel = mfcc.shape[0]
 x_multiply = 8
 if len_mel % x_multiply == 0:
 num_to_pad = 0
 else:
 num_to_pad = x_multiply - len_mel % x_multiply
 mfcc = np.pad(mfcc, pad_width=((0,num_to_pad), (0,0)))
 return mfcc

 @torch.no_grad()
 def forward_audio2secc(self, batch, inp=None):
 if inp['drv_audio_name'][-4:] in ['.wav', '.mp3']:
 # audio-to-exp
 ret = {}
 pred = self.audio2secc_model.forward(batch, ret=ret,train=False, temperature=inp['temperature'],)
 print("| audio-to-motion finished")
 if pred.shape[-1] == 144:
 id = ret['pred'][0][:,:80]
 exp = ret['pred'][0][:,80:]
 else:
 id = batch['id']
 exp = ret['pred'][0]
 if len(id) < len(exp): # happens when use ICL
 id = torch.cat([id, id[0].unsqueeze(0).repeat([len(exp)-len(id),1])])
 batch['id'] = id
 batch['exp'] = exp
 else:
 drv_motion_coeff_dict = self.drv_motion_coeff_dict
 batch['exp'] = torch.FloatTensor(drv_motion_coeff_dict['exp']).cuda()

 batch = self.get_driving_motion(batch['id'], batch['exp'], batch['euler'], batch['trans'], batch, inp)
 if self.use_icl_audio2motion:
 self.audio2secc_model.empty_context()
 return batch

 @torch.no_grad()
 def get_driving_motion(self, id, exp, euler, trans, batch, inp):
 zero_eulers = torch.zeros([id.shape[0], 3]).to(id.device)
 zero_trans = torch.zeros([id.shape[0], 3]).to(exp.device)
 # render the secc given the id,exp
 with torch.no_grad():
 chunk_size = 50
 drv_secc_color_lst = []
 num_iters = len(id)//chunk_size if len(id)%chunk_size == 0 else len(id)//chunk_size+1
 for i in tqdm.trange(num_iters, desc="rendering drv secc"):
 torch.cuda.empty_cache()
 face_mask, drv_secc_color = self.secc_renderer(id[i*chunk_size:(i+1)*chunk_size], exp[i*chunk_size:(i+1)*chunk_size], zero_eulers[i*chunk_size:(i+1)*chunk_size], zero_trans[i*chunk_size:(i+1)*chunk_size])
 drv_secc_color_lst.append(drv_secc_color.cpu())
 drv_secc_colors = torch.cat(drv_secc_color_lst, dim=0)
 _, src_secc_color = self.secc_renderer(id[0:1], exp[0:1], zero_eulers[0:1], zero_trans[0:1])
 _, cano_secc_color = self.secc_renderer(id[0:1], exp[0:1]*0, zero_eulers[0:1], zero_trans[0:1])
 batch['drv_secc'] = drv_secc_colors.cuda()
 batch['src_secc'] = src_secc_color.cuda()
 batch['cano_secc'] = cano_secc_color.cuda()
 
 # blinking secc
 if inp['blink_mode'] == 'period':
 period = 5 # second

 for i in tqdm.trange(len(drv_secc_colors),desc="blinking secc"):
 if i % (25*period) == 0:
 blink_dur_frames = random.randint(8, 12)
 for offset in range(blink_dur_frames):
 j = offset + i
 if j >= len(drv_secc_colors)-1: break
 def blink_percent_fn(t, T):
 return -4/T**2 * t**2 + 4/T * t
 blink_percent = blink_percent_fn(offset, blink_dur_frames)
 secc = batch['drv_secc'][j]
 out_secc = blink_eye_for_secc(secc, blink_percent)
 out_secc = out_secc.cuda()
 batch['drv_secc'][j] = out_secc

 # get the drv_kp for torso model, using the transformed trajectory
 drv_kp = self.face3d_helper.reconstruct_lm2d(id, exp, euler, trans) # [T, 68, 2]

 drv_kp = (drv_kp-0.5) / 0.5 # rescale to -1~1
 batch['drv_kp'] = torch.clamp(drv_kp, -1, 1)
 return batch

 @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
 if inp['low_memory_usage']:
 # save memory, when one image is rendered, write it into video
 import imageio
 debug_name = 'demo.mp4'
 writer = imageio.get_writer(debug_name, fps=25, format='FFMPEG', codec='h264')
 
 with torch.no_grad():
 for i in tqdm.trange(num_frames, desc="Real3D-Portrait 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}
 if i == 0:
 gen_output = self.secc2video_model.forward(img=ref_img_head, camera=camera[i:i+1], cond=cond, ret={}, cache_backbone=True, use_cached_backbone=False)
 else:
 gen_output = self.secc2video_model.forward(img=ref_img_head, camera=camera[i:i+1], cond=cond, ret={}, cache_backbone=False, use_cached_backbone=True)
 img = ((gen_output['image']+1)/2 * 255.).permute(0, 2, 3, 1)[0].int().cpu().numpy().astype(np.uint8)
 writer.append_data(img)
 writer.close()
 else:
 img_raw_lst = []
 img_lst = []
 depth_img_lst = []
 with torch.no_grad():
 for i in tqdm.trange(num_frames, desc="Real3D-Portrait 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}
 if i == 0:
 gen_output = self.secc2video_model.forward(img=ref_img_head, camera=camera[i:i+1], cond=cond, ret={}, cache_backbone=True, use_cached_backbone=False)
 else:
 gen_output = self.secc2video_model.forward(img=ref_img_head, 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
 debug_name = 'demo.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()
 
 # add audio track
 out_fname = 'infer_out/tmp/' + os.path.basename(inp['src_image_name'])[:-4] + '_' + 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 {self.wav16k_name} -y -v quiet -shortest {out_fname}")
 os.system(f"rm {debug_name}")
 os.system(f"rm {self.wav16k_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
 
 @torch.no_grad()
 def forward_system(self, batch, inp):
 self.forward_audio2secc(batch, inp)
 out_fname = self.forward_secc2video(batch, inp)
 return out_fname

 @classmethod
 def example_run(cls, inp=None):
 inp_tmp = {
 'drv_audio_name': 'data/raw/val_wavs/zozo.wav',
 'src_image_name': 'data/raw/val_imgs/Macron.png'
 }
 if inp is not None:
 inp_tmp.update(inp)
 inp = inp_tmp

 infer_instance = cls(inp['a2m_ckpt'], inp['head_ckpt'], inp['torso_ckpt'], inp=inp)
 infer_instance.infer_once(inp)

 ##############
 # IO-related
 ##############
 def save_wav16k(self, audio_name):
 supported_types = ('.wav', '.mp3', '.mp4', '.avi')
 assert audio_name.endswith(supported_types), f"Now we only support {','.join(supported_types)} as audio source!"
 wav16k_name = audio_name[:-4] + '_16k.wav'
 self.wav16k_name = wav16k_name
 extract_wav_cmd = f"ffmpeg -i {audio_name} -f wav -ar 16000 -v quiet -y {wav16k_name} -y"
 os.system(extract_wav_cmd)
 print(f"Extracted wav file (16khz) from {audio_name} to {wav16k_name}.")

 def get_f0(self, wav16k_name):
 from data_gen.utils.process_audio.extract_mel_f0 import extract_mel_from_fname, extract_f0_from_wav_and_mel
 wav, mel = extract_mel_from_fname(self.wav16k_name)
 f0, f0_coarse = extract_f0_from_wav_and_mel(wav, mel)
 f0 = f0.reshape([-1,1])
 return f0

if __name__ == '__main__':
 import argparse, glob, tqdm
 parser = argparse.ArgumentParser()
 parser.add_argument("--a2m_ckpt", default='checkpoints/240210_real3dportrait_orig/audio2secc_vae', type=str) 
 parser.add_argument("--head_ckpt", default='', type=str)
 parser.add_argument("--torso_ckpt", default='checkpoints/240210_real3dportrait_orig/secc2plane_torso_orig', type=str) 
 parser.add_argument("--src_img", default='data/raw/examples/Macron.png', type=str) # data/raw/examples/Macron.png
 parser.add_argument("--bg_img", default='', type=str) # data/raw/examples/bg.png
 parser.add_argument("--drv_aud", default='data/raw/examples/Obama_5s.wav', type=str) # data/raw/examples/Obama_5s.wav
 parser.add_argument("--drv_pose", default='data/raw/examples/May_5s.mp4', type=str) # data/raw/examples/May_5s.mp4
 parser.add_argument("--blink_mode", default='period', type=str) # none | period
 parser.add_argument("--temperature", default=0.2, type=float) # sampling temperature in audio2motion, higher -> more diverse, less accurate
 parser.add_argument("--mouth_amp", default=0.45, type=float) # scale of predicted mouth, enabled in audio-driven
 parser.add_argument("--head_torso_threshold", default=None, type=float, help="0.1~1.0, turn up this value if the hair is translucent")
 parser.add_argument("--out_name", default='') # output filename
 parser.add_argument("--out_mode", default='concat_debug') # final: only output talking head video; concat_debug: talking head with internel features 
 parser.add_argument("--map_to_init_pose", default='True') # whether to map the pose of first frame to source image
 parser.add_argument("--seed", default=None, type=int) # random seed, default None to use time.time()
 parser.add_argument("--min_face_area_percent", default=0.2, type=float) # scale of predicted mouth, enabled in audio-driven
 parser.add_argument("--low_memory_usage", action='store_true', help='write img to video upon generated, leads to slower fps, but use less memory')

 args = parser.parse_args()

 inp = {
 'a2m_ckpt': args.a2m_ckpt,
 'head_ckpt': args.head_ckpt,
 'torso_ckpt': args.torso_ckpt,
 'src_image_name': args.src_img,
 'bg_image_name': args.bg_img,
 'drv_audio_name': args.drv_aud,
 'drv_pose_name': args.drv_pose,
 'blink_mode': args.blink_mode,
 'temperature': args.temperature,
 'mouth_amp': args.mouth_amp,
 'out_name': args.out_name,
 'out_mode': args.out_mode,
 'map_to_init_pose': args.map_to_init_pose,
 'head_torso_threshold': args.head_torso_threshold,
 'seed': args.seed,
 'min_face_area_percent': args.min_face_area_percent,
 'low_memory_usage': args.low_memory_usage,
 }

 GeneFace2Infer.example_run(inp)