diff --git a/animation/SadTalker/src/generate_batch.py b/animation/SadTalker/src/generate_batch.py new file mode 100644 index 0000000..95f2152 --- /dev/null +++ b/animation/SadTalker/src/generate_batch.py @@ -0,0 +1,120 @@ +import os + +from tqdm import tqdm +import torch +import numpy as np +import random +import scipy.io as scio +import src.utils.audio as audio + +def crop_pad_audio(wav, audio_length): + if len(wav) > audio_length: + wav = wav[:audio_length] + elif len(wav) < audio_length: + wav = np.pad(wav, [0, audio_length - len(wav)], mode='constant', constant_values=0) + return wav + +def parse_audio_length(audio_length, sr, fps): + bit_per_frames = sr / fps + + num_frames = int(audio_length / bit_per_frames) + audio_length = int(num_frames * bit_per_frames) + + return audio_length, num_frames + +def generate_blink_seq(num_frames): + ratio = np.zeros((num_frames,1)) + frame_id = 0 + while frame_id in range(num_frames): + start = 80 + if frame_id+start+9<=num_frames - 1: + ratio[frame_id+start:frame_id+start+9, 0] = [0.5,0.6,0.7,0.9,1, 0.9, 0.7,0.6,0.5] + frame_id = frame_id+start+9 + else: + break + return ratio + +def generate_blink_seq_randomly(num_frames): + ratio = np.zeros((num_frames,1)) + if num_frames<=20: + return ratio + frame_id = 0 + while frame_id in range(num_frames): + start = random.choice(range(min(10,num_frames), min(int(num_frames/2), 70))) + if frame_id+start+5<=num_frames - 1: + ratio[frame_id+start:frame_id+start+5, 0] = [0.5, 0.9, 1.0, 0.9, 0.5] + frame_id = frame_id+start+5 + else: + break + return ratio + +def get_data(first_coeff_path, audio_path, device, ref_eyeblink_coeff_path, still=False, idlemode=False, length_of_audio=False, use_blink=True): + + syncnet_mel_step_size = 16 + fps = 25 + + pic_name = os.path.splitext(os.path.split(first_coeff_path)[-1])[0] + audio_name = os.path.splitext(os.path.split(audio_path)[-1])[0] + + + if idlemode: + num_frames = int(length_of_audio * 25) + indiv_mels = np.zeros((num_frames, 80, 16)) + else: + wav = audio.load_wav(audio_path, 16000) + wav_length, num_frames = parse_audio_length(len(wav), 16000, 25) + wav = crop_pad_audio(wav, wav_length) + orig_mel = audio.melspectrogram(wav).T + spec = orig_mel.copy() # nframes 80 + indiv_mels = [] + + for i in tqdm(range(num_frames), 'mel:'): + start_frame_num = i-2 + start_idx = int(80. * (start_frame_num / float(fps))) + end_idx = start_idx + syncnet_mel_step_size + seq = list(range(start_idx, end_idx)) + seq = [ min(max(item, 0), orig_mel.shape[0]-1) for item in seq ] + m = spec[seq, :] + indiv_mels.append(m.T) + indiv_mels = np.asarray(indiv_mels) # T 80 16 + + ratio = generate_blink_seq_randomly(num_frames) # T + source_semantics_path = first_coeff_path + source_semantics_dict = scio.loadmat(source_semantics_path) + ref_coeff = source_semantics_dict['coeff_3dmm'][:1,:70] #1 70 + ref_coeff = np.repeat(ref_coeff, num_frames, axis=0) + + if ref_eyeblink_coeff_path is not None: + ratio[:num_frames] = 0 + refeyeblink_coeff_dict = scio.loadmat(ref_eyeblink_coeff_path) + refeyeblink_coeff = refeyeblink_coeff_dict['coeff_3dmm'][:,:64] + refeyeblink_num_frames = refeyeblink_coeff.shape[0] + if refeyeblink_num_frames frame_num: + new_degree_list = new_degree_list[:frame_num] + elif len(new_degree_list) < frame_num: + for _ in range(frame_num-len(new_degree_list)): + new_degree_list.append(new_degree_list[-1]) + print(len(new_degree_list)) + print(frame_num) + + remainder = frame_num%batch_size + if remainder!=0: + for _ in range(batch_size-remainder): + new_degree_list.append(new_degree_list[-1]) + new_degree_np = np.array(new_degree_list).reshape(batch_size, -1) + return new_degree_np + diff --git a/animation/SadTalker/src/gradio_demo.py b/animation/SadTalker/src/gradio_demo.py new file mode 100644 index 0000000..9c9ae05 --- /dev/null +++ b/animation/SadTalker/src/gradio_demo.py @@ -0,0 +1,155 @@ +import torch, uuid +import os, sys, shutil +from src.utils.preprocess import CropAndExtract +from src.test_audio2coeff import Audio2Coeff +from src.facerender.animate import AnimateFromCoeff +from src.generate_batch import get_data +from src.generate_facerender_batch import get_facerender_data + +from src.utils.init_path import init_path + +from pydub import AudioSegment + + +def mp3_to_wav(mp3_filename,wav_filename,frame_rate): + mp3_file = AudioSegment.from_file(file=mp3_filename) + mp3_file.set_frame_rate(frame_rate).export(wav_filename,format="wav") + + +class SadTalker(): + + def __init__(self, checkpoint_path='checkpoints', config_path='src/config', lazy_load=False): + + if torch.cuda.is_available() : + device = "cuda" + else: + device = "cpu" + + self.device = device + + os.environ['TORCH_HOME']= checkpoint_path + + self.checkpoint_path = checkpoint_path + self.config_path = config_path + + + 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, + use_ref_video = False, + ref_video = None, + ref_info = None, + use_idle_mode = False, + length_of_audio = 0, use_blink=True, + result_dir='./results/'): + + self.sadtalker_paths = init_path(self.checkpoint_path, self.config_path, size, False, preprocess) + print(self.sadtalker_paths) + + 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) + + time_tag = str(uuid.uuid4()) + save_dir = os.path.join(result_dir, time_tag) + os.makedirs(save_dir, exist_ok=True) + + input_dir = os.path.join(save_dir, 'input') + os.makedirs(input_dir, exist_ok=True) + + print(source_image) + pic_path = os.path.join(input_dir, os.path.basename(source_image)) + shutil.move(source_image, input_dir) + + if driven_audio is not None and os.path.isfile(driven_audio): + audio_path = os.path.join(input_dir, os.path.basename(driven_audio)) + + #### mp3 to wav + if '.mp3' in audio_path: + mp3_to_wav(driven_audio, audio_path.replace('.mp3', '.wav'), 16000) + audio_path = audio_path.replace('.mp3', '.wav') + else: + shutil.move(driven_audio, input_dir) + + elif use_idle_mode: + audio_path = os.path.join(input_dir, 'idlemode_'+str(length_of_audio)+'.wav') ## generate audio from this new audio_path + from pydub import AudioSegment + one_sec_segment = AudioSegment.silent(duration=1000*length_of_audio) #duration in milliseconds + one_sec_segment.export(audio_path, format="wav") + else: + print(use_ref_video, ref_info) + assert use_ref_video == True and ref_info == 'all' + + if use_ref_video and ref_info == 'all': # full ref mode + ref_video_videoname = os.path.basename(ref_video) + audio_path = os.path.join(save_dir, ref_video_videoname+'.wav') + print('new audiopath:',audio_path) + # if ref_video contains audio, set the audio from ref_video. + cmd = r"ffmpeg -y -hide_banner -loglevel error -i %s %s"%(ref_video, audio_path) + os.system(cmd) + + os.makedirs(save_dir, exist_ok=True) + + #crop image and extract 3dmm from image + first_frame_dir = os.path.join(save_dir, 'first_frame_dir') + os.makedirs(first_frame_dir, exist_ok=True) + first_coeff_path, crop_pic_path, crop_info = self.preprocess_model.generate(pic_path, first_frame_dir, preprocess, True, size) + + if first_coeff_path is None: + raise AttributeError("No face is detected") + + if use_ref_video: + print('using ref video for genreation') + ref_video_videoname = os.path.splitext(os.path.split(ref_video)[-1])[0] + ref_video_frame_dir = os.path.join(save_dir, ref_video_videoname) + os.makedirs(ref_video_frame_dir, exist_ok=True) + print('3DMM Extraction for the reference video providing pose') + ref_video_coeff_path, _, _ = self.preprocess_model.generate(ref_video, ref_video_frame_dir, preprocess, source_image_flag=False) + else: + ref_video_coeff_path = None + + if use_ref_video: + if ref_info == 'pose': + ref_pose_coeff_path = ref_video_coeff_path + ref_eyeblink_coeff_path = None + elif ref_info == 'blink': + ref_pose_coeff_path = None + ref_eyeblink_coeff_path = ref_video_coeff_path + elif ref_info == 'pose+blink': + ref_pose_coeff_path = ref_video_coeff_path + ref_eyeblink_coeff_path = ref_video_coeff_path + elif ref_info == 'all': + ref_pose_coeff_path = None + ref_eyeblink_coeff_path = None + else: + raise('error in refinfo') + else: + ref_pose_coeff_path = None + ref_eyeblink_coeff_path = None + + #audio2ceoff + 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? + 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) + 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}') + + del self.preprocess_model + del self.audio_to_coeff + del self.animate_from_coeff + + if torch.cuda.is_available(): + torch.cuda.empty_cache() + torch.cuda.synchronize() + + import gc; gc.collect() + + return return_path + + \ No newline at end of file diff --git a/animation/SadTalker/src/test_audio2coeff.py b/animation/SadTalker/src/test_audio2coeff.py new file mode 100644 index 0000000..bbf19f4 --- /dev/null +++ b/animation/SadTalker/src/test_audio2coeff.py @@ -0,0 +1,123 @@ +import os +import torch +import numpy as np +from scipy.io import savemat, loadmat +from yacs.config import CfgNode as CN +from scipy.signal import savgol_filter + +import safetensors +import safetensors.torch + +from src.audio2pose_models.audio2pose import Audio2Pose +from src.audio2exp_models.networks import SimpleWrapperV2 +from src.audio2exp_models.audio2exp import Audio2Exp +from src.utils.safetensor_helper import load_x_from_safetensor + +def load_cpk(checkpoint_path, model=None, optimizer=None, device="cpu"): + checkpoint = torch.load(checkpoint_path, map_location=torch.device(device)) + if model is not None: + model.load_state_dict(checkpoint['model']) + if optimizer is not None: + optimizer.load_state_dict(checkpoint['optimizer']) + + return checkpoint['epoch'] + +class Audio2Coeff(): + + def __init__(self, sadtalker_path, device): + #load config + fcfg_pose = open(sadtalker_path['audio2pose_yaml_path']) + cfg_pose = CN.load_cfg(fcfg_pose) + cfg_pose.freeze() + fcfg_exp = open(sadtalker_path['audio2exp_yaml_path']) + cfg_exp = CN.load_cfg(fcfg_exp) + cfg_exp.freeze() + + # load audio2pose_model + self.audio2pose_model = Audio2Pose(cfg_pose, None, device=device) + self.audio2pose_model = self.audio2pose_model.to(device) + self.audio2pose_model.eval() + for param in self.audio2pose_model.parameters(): + param.requires_grad = False + + try: + if sadtalker_path['use_safetensor']: + checkpoints = safetensors.torch.load_file(sadtalker_path['checkpoint']) + self.audio2pose_model.load_state_dict(load_x_from_safetensor(checkpoints, 'audio2pose')) + else: + load_cpk(sadtalker_path['audio2pose_checkpoint'], model=self.audio2pose_model, device=device) + except: + raise Exception("Failed in loading audio2pose_checkpoint") + + # load audio2exp_model + netG = SimpleWrapperV2() + netG = netG.to(device) + for param in netG.parameters(): + netG.requires_grad = False + netG.eval() + try: + if sadtalker_path['use_safetensor']: + checkpoints = safetensors.torch.load_file(sadtalker_path['checkpoint']) + netG.load_state_dict(load_x_from_safetensor(checkpoints, 'audio2exp')) + else: + load_cpk(sadtalker_path['audio2exp_checkpoint'], model=netG, device=device) + except: + raise Exception("Failed in loading audio2exp_checkpoint") + self.audio2exp_model = Audio2Exp(netG, cfg_exp, device=device, prepare_training_loss=False) + self.audio2exp_model = self.audio2exp_model.to(device) + for param in self.audio2exp_model.parameters(): + param.requires_grad = False + self.audio2exp_model.eval() + + self.device = device + + def generate(self, batch, coeff_save_dir, pose_style, ref_pose_coeff_path=None): + + with torch.no_grad(): + #test + results_dict_exp= self.audio2exp_model.test(batch) + exp_pred = results_dict_exp['exp_coeff_pred'] #bs T 64 + + #for class_id in range(1): + #class_id = 0#(i+10)%45 + #class_id = random.randint(0,46) #46 styles can be selected + batch['class'] = torch.LongTensor([pose_style]).to(self.device) + results_dict_pose = self.audio2pose_model.test(batch) + pose_pred = results_dict_pose['pose_pred'] #bs T 6 + + pose_len = pose_pred.shape[1] + if pose_len<13: + pose_len = int((pose_len-1)/2)*2+1 + pose_pred = torch.Tensor(savgol_filter(np.array(pose_pred.cpu()), pose_len, 2, axis=1)).to(self.device) + else: + pose_pred = torch.Tensor(savgol_filter(np.array(pose_pred.cpu()), 13, 2, axis=1)).to(self.device) + + coeffs_pred = torch.cat((exp_pred, pose_pred), dim=-1) #bs T 70 + + coeffs_pred_numpy = coeffs_pred[0].clone().detach().cpu().numpy() + + if ref_pose_coeff_path is not None: + coeffs_pred_numpy = self.using_refpose(coeffs_pred_numpy, ref_pose_coeff_path) + + savemat(os.path.join(coeff_save_dir, '%s##%s.mat'%(batch['pic_name'], batch['audio_name'])), + {'coeff_3dmm': coeffs_pred_numpy}) + + return os.path.join(coeff_save_dir, '%s##%s.mat'%(batch['pic_name'], batch['audio_name'])) + + def using_refpose(self, coeffs_pred_numpy, ref_pose_coeff_path): + num_frames = coeffs_pred_numpy.shape[0] + refpose_coeff_dict = loadmat(ref_pose_coeff_path) + refpose_coeff = refpose_coeff_dict['coeff_3dmm'][:,64:70] + refpose_num_frames = refpose_coeff.shape[0] + if refpose_num_frames