YinuoGuo27 commited on
Commit
bb12372
·
verified ·
1 Parent(s): cd1584f

Upload 2 files

Browse files
Files changed (2) hide show
  1. difpoint/.DS_Store +0 -0
  2. difpoint/inference.py +516 -0
difpoint/.DS_Store ADDED
Binary file (6.15 kB). View file
 
difpoint/inference.py ADDED
@@ -0,0 +1,516 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # -*- coding: UTF-8 -*-
2
+ '''
3
+ @File :inference.py
4
+ @Author :Chaolong Yang
5
+ @Date :2024/5/29 19:26
6
+ '''
7
+ import glob
8
+
9
+ import os
10
+ os.environ['HYDRA_FULL_ERROR']='1'
11
+ os.environ['CUDA_VISIBLE_DEVICES'] = '0'
12
+
13
+ import os
14
+ import time
15
+ import shutil
16
+ import uuid
17
+ import os
18
+ import cv2
19
+ import tyro
20
+ from difpoint.src.utils.crop import crop_image, parse_bbox_from_landmark, crop_image_by_bbox, paste_back, paste_back_pytorch
21
+ from difpoint.src.utils.utils import resize_to_limit, prepare_paste_back, get_rotation_matrix, calc_lip_close_ratio, \
22
+ calc_eye_close_ratio, transform_keypoint, concat_feat
23
+ from difpoint.src.utils import utils
24
+
25
+ import numpy as np
26
+ from tqdm import tqdm
27
+ import cv2
28
+ from rich.progress import track
29
+
30
+ from difpoint.croper import Croper
31
+ from PIL import Image
32
+ import time
33
+
34
+
35
+ import torch
36
+ import torch.nn.functional as F
37
+ from torch import nn
38
+ import imageio
39
+ from pydub import AudioSegment
40
+ from pykalman import KalmanFilter
41
+ import scipy
42
+ import matplotlib.pyplot as plt
43
+ import matplotlib
44
+ matplotlib.use('Agg')
45
+
46
+ from difpoint.dataset_process import audio
47
+ import os
48
+ import argparse
49
+ import pdb
50
+ import subprocess
51
+ import ffmpeg
52
+ import cv2
53
+ import time
54
+ import numpy as np
55
+ import os
56
+ import datetime
57
+ import platform
58
+ from omegaconf import OmegaConf
59
+ from difpoint.src.pipelines.faster_live_portrait_pipeline import FasterLivePortraitPipeline
60
+
61
+ FFMPEG = "ffmpeg"
62
+
63
+ def parse_audio_length(audio_length, sr, fps):
64
+ bit_per_frames = sr / fps
65
+ num_frames = int(audio_length / bit_per_frames)
66
+ audio_length = int(num_frames * bit_per_frames)
67
+ return audio_length, num_frames
68
+
69
+ def crop_pad_audio(wav, audio_length):
70
+ if len(wav) > audio_length:
71
+ wav = wav[:audio_length]
72
+ elif len(wav) < audio_length:
73
+ wav = np.pad(wav, [0, audio_length - len(wav)], mode='constant', constant_values=0)
74
+ return wav
75
+
76
+ class Conv2d(nn.Module):
77
+ def __init__(self, cin, cout, kernel_size, stride, padding, residual=False, use_act=True, *args, **kwargs):
78
+ super().__init__(*args, **kwargs)
79
+ self.conv_block = nn.Sequential(
80
+ nn.Conv2d(cin, cout, kernel_size, stride, padding),
81
+ nn.BatchNorm2d(cout)
82
+ )
83
+ self.act = nn.ReLU()
84
+ self.residual = residual
85
+ self.use_act = use_act
86
+
87
+ def forward(self, x):
88
+ out = self.conv_block(x)
89
+ if self.residual:
90
+ out += x
91
+
92
+ if self.use_act:
93
+ return self.act(out)
94
+ else:
95
+ return out
96
+
97
+ class AudioEncoder(nn.Module):
98
+ def __init__(self, wav2lip_checkpoint, device):
99
+ super(AudioEncoder, self).__init__()
100
+
101
+ self.audio_encoder = nn.Sequential(
102
+ Conv2d(1, 32, kernel_size=3, stride=1, padding=1),
103
+ Conv2d(32, 32, kernel_size=3, stride=1, padding=1, residual=True),
104
+ Conv2d(32, 32, kernel_size=3, stride=1, padding=1, residual=True),
105
+
106
+ Conv2d(32, 64, kernel_size=3, stride=(3, 1), padding=1),
107
+ Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True),
108
+ Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True),
109
+
110
+ Conv2d(64, 128, kernel_size=3, stride=3, padding=1),
111
+ Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True),
112
+ Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True),
113
+
114
+ Conv2d(128, 256, kernel_size=3, stride=(3, 2), padding=1),
115
+ Conv2d(256, 256, kernel_size=3, stride=1, padding=1, residual=True),
116
+
117
+ Conv2d(256, 512, kernel_size=3, stride=1, padding=0),
118
+ Conv2d(512, 512, kernel_size=1, stride=1, padding=0),)
119
+
120
+ #### load the pre-trained audio_encoder
121
+ wav2lip_state_dict = torch.load(wav2lip_checkpoint, map_location=torch.device(device))['state_dict']
122
+ state_dict = self.audio_encoder.state_dict()
123
+
124
+ for k,v in wav2lip_state_dict.items():
125
+ if 'audio_encoder' in k:
126
+ state_dict[k.replace('module.audio_encoder.', '')] = v
127
+ self.audio_encoder.load_state_dict(state_dict)
128
+
129
+ def forward(self, audio_sequences):
130
+ # audio_sequences = (B, T, 1, 80, 16)
131
+ B = audio_sequences.size(0)
132
+
133
+ audio_sequences = torch.cat([audio_sequences[:, i] for i in range(audio_sequences.size(1))], dim=0)
134
+
135
+ audio_embedding = self.audio_encoder(audio_sequences) # B, 512, 1, 1
136
+ dim = audio_embedding.shape[1]
137
+ audio_embedding = audio_embedding.reshape((B, -1, dim, 1, 1))
138
+
139
+ return audio_embedding.squeeze(-1).squeeze(-1) #B seq_len+1 512
140
+
141
+ def partial_fields(target_class, kwargs):
142
+ return target_class(**{k: v for k, v in kwargs.items() if hasattr(target_class, k)})
143
+
144
+ def dct2device(dct: dict, device):
145
+ for key in dct:
146
+ dct[key] = torch.tensor(dct[key]).to(device)
147
+ return dct
148
+
149
+ def save_video_with_watermark(video, audio, save_path, watermark=False):
150
+ temp_file = str(uuid.uuid4())+'.mp4'
151
+ cmd = r'ffmpeg -y -i "%s" -i "%s" -vcodec copy "%s"' % (video, audio, temp_file)
152
+ os.system(cmd)
153
+ shutil.move(temp_file, save_path)
154
+
155
+
156
+
157
+ class Inferencer(object):
158
+ def __init__(self):
159
+
160
+ st=time.time()
161
+ print('#'*25+'Start initialization'+'#'*25)
162
+ self.device = 'cuda'
163
+ from difpoint.model import get_model
164
+ self.point_diffusion = get_model()
165
+ ckpt = torch.load('/home/yinuo/Gradio-UI_copy/difpoint/outputs/2024.08.26_dim_70_frame_64_vox1_selected_d6.5_c8.5/2024-08-26--16-52-34/checkpoint-500000.pth')
166
+
167
+ self.point_diffusion.load_state_dict(ckpt['model'])
168
+ print('model', self.point_diffusion.children())
169
+ self.point_diffusion.eval()
170
+ self.point_diffusion.to(self.device)
171
+
172
+ lm_croper_checkpoint = os.path.join('difpoint/dataset_process/ckpts/', 'shape_predictor_68_face_landmarks.dat')
173
+ self.croper = Croper(lm_croper_checkpoint)
174
+
175
+ self.norm_info = dict(np.load(r'difpoint/datasets/norm_info_d6.5_c8.5_vox1_train.npz'))
176
+
177
+ wav2lip_checkpoint = 'difpoint/dataset_process/ckpts/wav2lip.pth'
178
+ self.wav2lip_model = AudioEncoder(wav2lip_checkpoint, 'cuda')
179
+ self.wav2lip_model.cuda()
180
+ self.wav2lip_model.eval()
181
+
182
+ # specify configs for inference
183
+ self.inf_cfg = OmegaConf.load("difpoint/configs/trt_mp_infer.yaml")
184
+ self.inf_cfg.infer_params.flag_pasteback = False
185
+
186
+ self.live_portrait_pipeline = FasterLivePortraitPipeline(cfg=self.inf_cfg, is_animal=False)
187
+ #ret = self.live_portrait_pipeline.prepare_source(source_image)
188
+
189
+ print('#'*25+f'End initialization, cost time {time.time()-st}'+'#'*25)
190
+
191
+ def _norm(self, data_dict):
192
+ for k in data_dict.keys():
193
+ if k in ['yaw', 'pitch', 'roll', 't', 'scale', 'c_lip', 'c_eye']:
194
+ v=data_dict[k]
195
+ data_dict[k] = (v - self.norm_info[k+'_mean'])/self.norm_info[k+'_std']
196
+ elif k in ['exp', 'kp']:
197
+ v=data_dict[k]
198
+ data_dict[k] = (v - self.norm_info[k+'_mean'].reshape(1,21,3))/self.norm_info[k+'_std'].reshape(1,21,3)
199
+ return data_dict
200
+
201
+ def _denorm(self, data_dict):
202
+ for k in data_dict.keys():
203
+ if k in ['yaw', 'pitch', 'roll', 't', 'scale', 'c_lip', 'c_eye']:
204
+ v=data_dict[k]
205
+ data_dict[k] = v * self.norm_info[k+'_std'] + self.norm_info[k+'_mean']
206
+ elif k in ['exp', 'kp']:
207
+ v=data_dict[k]
208
+ data_dict[k] = v * self.norm_info[k+'_std'] + self.norm_info[k+'_mean']
209
+ return data_dict
210
+
211
+
212
+ def output_to_dict(self, data):
213
+ output = {}
214
+
215
+ output['scale'] = data[:, 0]
216
+ output['yaw'] = data[:, 1, None]
217
+ output['pitch'] = data[:, 2, None]
218
+ output['roll'] = data[:, 3, None]
219
+ output['t'] = data[:, 4:7]
220
+ output['exp'] = data[:, 7:]
221
+
222
+ return output
223
+
224
+ def extract_mel_from_audio(self, audio_file_path):
225
+ syncnet_mel_step_size = 16
226
+ fps = 25
227
+ wav = audio.load_wav(audio_file_path, 16000)
228
+ wav_length, num_frames = parse_audio_length(len(wav), 16000, 25)
229
+ wav = crop_pad_audio(wav, wav_length)
230
+ orig_mel = audio.melspectrogram(wav).T
231
+ spec = orig_mel.copy()
232
+ indiv_mels = []
233
+
234
+ for i in tqdm(range(num_frames), 'mel:'):
235
+ start_frame_num = i - 2
236
+ start_idx = int(80. * (start_frame_num / float(fps)))
237
+ end_idx = start_idx + syncnet_mel_step_size
238
+ seq = list(range(start_idx, end_idx))
239
+ seq = [min(max(item, 0), orig_mel.shape[0] - 1) for item in seq]
240
+ m = spec[seq, :]
241
+ indiv_mels.append(m.T)
242
+ indiv_mels = np.asarray(indiv_mels) # T 80 16
243
+ return indiv_mels
244
+
245
+ def extract_wav2lip_from_audio(self, audio_file_path):
246
+ asd_mel = self.extract_mel_from_audio(audio_file_path)
247
+ asd_mel = torch.FloatTensor(asd_mel).cuda().unsqueeze(0).unsqueeze(2)
248
+ with torch.no_grad():
249
+ hidden = self.wav2lip_model(asd_mel)
250
+ return hidden[0].cpu().detach().numpy()
251
+
252
+ def headpose_pred_to_degree(self, pred):
253
+ device = pred.device
254
+ idx_tensor = [idx for idx in range(66)]
255
+ idx_tensor = torch.FloatTensor(idx_tensor).to(device)
256
+ pred = F.softmax(pred)
257
+ degree = torch.sum(pred * idx_tensor, 1) * 3 - 99
258
+ return degree
259
+
260
+ def calc_combined_eye_ratio(self, c_d_eyes_i, c_s_eyes):
261
+ c_s_eyes_tensor = torch.from_numpy(c_s_eyes).float().to(self.device)
262
+ c_d_eyes_i_tensor = c_d_eyes_i[0].reshape(1, 1).to(self.device)
263
+ # [c_s,eyes, c_d,eyes,i]
264
+ combined_eye_ratio_tensor = torch.cat([c_s_eyes_tensor, c_d_eyes_i_tensor], dim=1)
265
+ return combined_eye_ratio_tensor
266
+
267
+ def calc_combined_lip_ratio(self, c_d_lip_i, c_s_lip):
268
+ c_s_lip_tensor = torch.from_numpy(c_s_lip).float().to(self.device)
269
+ c_d_lip_i_tensor = c_d_lip_i[0].to(self.device).reshape(1, 1) # 1x1
270
+ # [c_s,lip, c_d,lip,i]
271
+ combined_lip_ratio_tensor = torch.cat([c_s_lip_tensor, c_d_lip_i_tensor], dim=1) # 1x2
272
+ return combined_lip_ratio_tensor
273
+
274
+ # 2024.06.26
275
+ @torch.no_grad()
276
+ def generate_with_audio_img(self, upload_audio_path, tts_audio_path, audio_type, image_path, smoothed_pitch, smoothed_yaw, smoothed_roll, smoothed_t, save_path='results'):
277
+ print(audio_type)
278
+ if audio_type == 'upload':
279
+ audio_path = upload_audio_path
280
+ elif audio_type == 'tts':
281
+ audio_path = tts_audio_path
282
+ save_path = os.path.join(save_path, "output.mp4")
283
+ image = [np.array(Image.open(image_path).convert('RGB'))]
284
+ if image[0].shape[0] != 256 or image[0].shape[1] != 256:
285
+ cropped_image, crop, quad = self.croper.crop(image, still=False, xsize=512)
286
+ input_image = cv2.resize(cropped_image[0], (256, 256))
287
+ else:
288
+ input_image = image[0]
289
+
290
+ I_s = torch.FloatTensor(input_image.transpose((2, 0, 1))).unsqueeze(0).cuda() / 255
291
+ pitch, yaw, roll, t, exp, scale, kp = self.live_portrait_pipeline.model_dict["motion_extractor"].predict(
292
+ I_s)
293
+ x_s_info = {
294
+ "pitch": pitch,
295
+ "yaw": yaw,
296
+ "roll": roll,
297
+ "t": t,
298
+ "exp": exp,
299
+ "scale": scale,
300
+ "kp": kp
301
+ }
302
+ x_c_s = kp.reshape(1, 21, -1)
303
+ R_s = get_rotation_matrix(x_s_info['pitch'], x_s_info['yaw'], x_s_info['roll'])
304
+ f_s = self.live_portrait_pipeline.model_dict["app_feat_extractor"].predict(I_s)
305
+ x_s = transform_keypoint(pitch, yaw, roll, t, exp, scale, kp)
306
+
307
+ flag_lip_zero = self.inf_cfg.infer_params.flag_normalize_lip
308
+
309
+ if flag_lip_zero:
310
+ # let lip-open scalar to be 0 at first
311
+ c_d_lip_before_animation = [0.]
312
+
313
+ lip_delta_before_animation = self.live_portrait_pipeline.model_dict['stitching_lip_retarget'].predict(
314
+ concat_feat(x_s, combined_lip_ratio_tensor_before_animation))
315
+
316
+ ######## process driving info ########
317
+ kp_info = {}
318
+ for k in x_s_info.keys():
319
+ kp_info[k] = x_s_info[k]
320
+ # kp_info['c_lip'] = c_s_lip
321
+ # kp_info['c_eye'] = c_s_eye
322
+
323
+ kp_info = self._norm(kp_info)
324
+
325
+ ori_kp = torch.cat([torch.zeros([1, 7]).to('cuda'), torch.Tensor(kp_info['kp'].reshape(1,63)).to('cuda')], -1).cuda()
326
+
327
+ input_x = np.concatenate([kp_info[k] for k in ['scale', 'yaw', 'pitch', 'roll', 't']], 1)
328
+ input_x = np.concatenate((input_x, kp_info['exp'].reshape(1, 63)), axis=1)
329
+ input_x = np.expand_dims(input_x, -1)
330
+ input_x = np.expand_dims(input_x, 0)
331
+ input_x = np.concatenate([input_x, input_x, input_x], -1)
332
+
333
+ aud_feat = self.extract_wav2lip_from_audio(audio_path)
334
+
335
+ outputs = [input_x]
336
+
337
+ st = time.time()
338
+ print('#' * 25 + 'Start Inference' + '#' * 25)
339
+ sample_frame = 64 # 32 aud_feat.shape[0]
340
+
341
+ for i in range(0, aud_feat.shape[0] - 1, sample_frame):
342
+ input_mel = torch.Tensor(aud_feat[i: i + sample_frame]).unsqueeze(0).cuda()
343
+ kp0 = torch.Tensor(outputs[-1])[:, -1].cuda()
344
+ pred_kp = self.point_diffusion.forward_sample(70, ref_kps=kp0, ori_kps=ori_kp, aud_feat=input_mel,
345
+ scheduler='ddim', num_inference_steps=50)
346
+ outputs.append(pred_kp.cpu().numpy())
347
+
348
+
349
+ outputs = np.mean(np.concatenate(outputs, 1)[0], -1)[1:, ]
350
+ output_dict = self.output_to_dict(outputs)
351
+ output_dict = self._denorm(output_dict)
352
+
353
+ num_frame = output_dict['yaw'].shape[0]
354
+ x_d_info = {}
355
+ for key in output_dict:
356
+ x_d_info[key] = torch.tensor(output_dict[key]).cuda()
357
+
358
+ # smooth
359
+ def smooth(sequence, n_dim_state=1):
360
+ kf = KalmanFilter(initial_state_mean=sequence[0],
361
+ transition_covariance=0.05 * np.eye(n_dim_state), # 较小的过程噪声
362
+ observation_covariance=0.001 * np.eye(n_dim_state)) # 可以增大观测噪声,减少敏感性
363
+ state_means, _ = kf.smooth(sequence)
364
+ return state_means
365
+
366
+ # scale_data = x_d_info['scale'].cpu().numpy()
367
+ yaw_data = x_d_info['yaw'].cpu().numpy()
368
+ pitch_data = x_d_info['pitch'].cpu().numpy()
369
+ roll_data = x_d_info['roll'].cpu().numpy()
370
+ t_data = x_d_info['t'].cpu().numpy()
371
+ exp_data = x_d_info['exp'].cpu().numpy()
372
+
373
+ smoothed_pitch = smooth(pitch_data, n_dim_state=1) * smoothed_pitch
374
+ smoothed_yaw = smooth(yaw_data, n_dim_state=1) * smoothed_yaw
375
+ smoothed_roll = smooth(roll_data, n_dim_state=1) * smoothed_roll
376
+ # smoothed_scale = smooth(scale_data, n_dim_state=1)
377
+ smoothed_t = smooth(t_data, n_dim_state=3) * smoothed_t
378
+ smoothed_exp = smooth(exp_data, n_dim_state=63)
379
+
380
+ # x_d_info['scale'] = torch.Tensor(smoothed_scale).cuda()
381
+ x_d_info['pitch'] = torch.Tensor(smoothed_pitch).cuda()
382
+ x_d_info['yaw'] = torch.Tensor(smoothed_yaw).cuda()
383
+ x_d_info['roll'] = torch.Tensor(smoothed_roll).cuda()
384
+ x_d_info['t'] = torch.Tensor(smoothed_t).cuda()
385
+ x_d_info['exp'] = torch.Tensor(smoothed_exp).cuda()
386
+
387
+
388
+
389
+ template_dct = {'motion': [], 'c_d_eyes_lst': [], 'c_d_lip_lst': []}
390
+ for i in track(range(num_frame), description='Making motion templates...', total=num_frame):
391
+ # collect s_d, R_d, δ_d and t_d for inference
392
+ x_d_i_info = x_d_info
393
+ R_d_i = get_rotation_matrix(x_d_i_info['pitch'][i], x_d_i_info['yaw'][i], x_d_i_info['roll'][i])
394
+
395
+ item_dct = {
396
+ 'scale': x_d_i_info['scale'][i].cpu().numpy().astype(np.float32),
397
+ 'R_d': R_d_i.astype(np.float32),
398
+ 'exp': x_d_i_info['exp'][i].reshape(1, 21, -1).cpu().numpy().astype(np.float32),
399
+ 't': x_d_i_info['t'][i].cpu().numpy().astype(np.float32),
400
+ }
401
+
402
+ template_dct['motion'].append(item_dct)
403
+ # template_dct['c_d_eyes_lst'].append(x_d_i_info['c_eye'][i])
404
+ # template_dct['c_d_lip_lst'].append(x_d_i_info['c_lip'][i])
405
+
406
+ I_p_lst = []
407
+ R_d_0, x_d_0_info = None, None
408
+
409
+ for i in track(range(num_frame), description='Animating...', total=num_frame):
410
+ x_d_i_info = template_dct['motion'][i]
411
+
412
+ for key in x_d_i_info:
413
+ x_d_i_info[key] = torch.tensor(x_d_i_info[key]).cuda()
414
+ for key in x_s_info:
415
+ x_s_info[key] = torch.tensor(x_s_info[key]).cuda()
416
+
417
+ R_d_i = x_d_i_info['R_d']
418
+
419
+ if i == 0:
420
+ R_d_0 = R_d_i
421
+ x_d_0_info = x_d_i_info
422
+
423
+
424
+ if self.inf_cfg.infer_params.flag_relative_motion:
425
+ R_new = (R_d_i.cpu().numpy() @ R_d_0.permute(0, 2, 1).cpu().numpy()) @ R_s
426
+ delta_new = x_s_info['exp'].reshape(1, 21, -1) + (x_d_i_info['exp'] - x_d_0_info['exp'])
427
+ scale_new = x_s_info['scale'] * (x_d_i_info['scale'] / x_d_0_info['scale'])
428
+ t_new = x_s_info['t'] + (x_d_i_info['t'] - x_d_0_info['t'])
429
+ else:
430
+ R_new = R_d_i
431
+ delta_new = x_d_i_info['exp']
432
+ scale_new = x_s_info['scale']
433
+ t_new = x_d_i_info['t']
434
+
435
+ t_new[..., 2] = 0 # zero tz
436
+ x_c_s = torch.tensor(x_c_s, dtype=torch.float32).cuda()
437
+ R_new = torch.tensor(R_new, dtype=torch.float32).cuda()
438
+ delta_new = torch.tensor(delta_new, dtype=torch.float32).cuda()
439
+ t_new = torch.tensor(t_new, dtype=torch.float32).cuda()
440
+ scale_new = torch.tensor(scale_new, dtype=torch.float32).cuda()
441
+ x_d_i_new = scale_new * (x_c_s @ R_new + delta_new) + t_new
442
+ x_d_i_new = x_d_i_new.cpu().numpy()
443
+
444
+ # Algorithm 1:
445
+ if not self.inf_cfg.infer_params.flag_stitching and not self.inf_cfg.infer_params.flag_eye_retargeting and not self.inf_cfg.infer_params.flag_lip_retargeting:
446
+ # without stitching or retargeting
447
+ if flag_lip_zero:
448
+ x_d_i_new += lip_delta_before_animation.reshape(-1, x_s.shape[1], 3)
449
+ else:
450
+ pass
451
+ elif self.inf_cfg.infer_params.flag_stitching and not self.inf_cfg.infer_params.flag_eye_retargeting and not self.inf_cfg.infer_params.flag_lip_retargeting:
452
+ # with stitching and without retargeting
453
+ if flag_lip_zero:
454
+ x_d_i_new = self.live_portrait_pipeline.stitching(x_s, x_d_i_new) + lip_delta_before_animation.reshape(
455
+ -1, x_s.shape[1], 3)
456
+ else:
457
+ x_d_i_new = self.live_portrait_pipeline.stitching(x_s, x_d_i_new)
458
+ else:
459
+ eyes_delta, lip_delta = None, None
460
+ if self.inf_cfg.infer_params.flag_eye_retargeting:
461
+ c_d_eyes_i = template_dct['c_d_eyes_lst'][i]
462
+ combined_eye_ratio_tensor = self.calc_combined_eye_ratio(c_d_eyes_i, c_s_eye)
463
+ # ∆_eyes,i = R_eyes(x_s; c_s,eyes, c_d,eyes,i)
464
+ eyes_delta = self.live_portrait_pipeline.retarget_eye(x_s, combined_eye_ratio_tensor)
465
+ if self.inf_cfg.infer_params.flag_lip_retargeting:
466
+ c_d_lip_i = template_dct['c_d_lip_lst'][i]
467
+ combined_lip_ratio_tensor = self.calc_combined_lip_ratio(c_d_lip_i, c_s_lip)
468
+ # ∆_lip,i = R_lip(x_s; c_s,lip, c_d,lip,i)
469
+ lip_delta = self.live_portrait_pipeline.retarget_lip(x_s, combined_lip_ratio_tensor)
470
+
471
+ if self.inf_cfg.infer_params.flag_relative_motion: # use x_s
472
+ x_d_i_new = x_s + \
473
+ (eyes_delta.reshape(-1, x_s.shape[1], 3) if eyes_delta is not None else 0) + \
474
+ (lip_delta.reshape(-1, x_s.shape[1], 3) if lip_delta is not None else 0)
475
+ else: # use x_d,i
476
+ x_d_i_new = x_d_i_new + \
477
+ (eyes_delta.reshape(-1, x_s.shape[1], 3) if eyes_delta is not None else 0) + \
478
+ (lip_delta.reshape(-1, x_s.shape[1], 3) if lip_delta is not None else 0)
479
+
480
+ if self.inf_cfg.infer_params.flag_stitching:
481
+ x_d_i_new = self.live_portrait_pipeline.stitching(x_s, x_d_i_new)
482
+
483
+ out = self.live_portrait_pipeline.model_dict["warping_spade"].predict(f_s, x_s, x_d_i_new).cpu().numpy().astype(np.uint8)
484
+ I_p_lst.append(out)
485
+
486
+ video_name = os.path.basename(save_path)
487
+ video_save_dir = os.path.dirname(save_path)
488
+ path = os.path.join(video_save_dir, video_name)
489
+
490
+ imageio.mimsave(path, I_p_lst, fps=float(25))
491
+
492
+ audio_name = audio_path.split('/')[-1]
493
+ new_audio_path = os.path.join(video_save_dir, audio_name)
494
+ start_time = 0
495
+ # cog will not keep the .mp3 filename
496
+ sound = AudioSegment.from_file(audio_path)
497
+ end_time = start_time + num_frame * 1 / 25 * 1000
498
+ word1 = sound.set_frame_rate(16000)
499
+ word = word1[start_time:end_time]
500
+ word.export(new_audio_path, format="wav")
501
+
502
+ save_video_with_watermark(path, new_audio_path, save_path, watermark=False)
503
+ print(f'The generated video is named {video_save_dir}/{video_name}')
504
+
505
+ print('#' * 25 + f'End Inference, cost time {time.time() - st}' + '#' * 25)
506
+ return save_path
507
+
508
+
509
+
510
+
511
+
512
+ import argparse
513
+ if __name__ == "__main__":
514
+ Infer = Inferencer()
515
+ Infer.generate_with_audio_img(None, 'difpoint/assets/test/test.wav', 'difpoint/assets/test/test2.jpg', 0.8, 0.8, 0.8, 0.8)
516
+