Spaces:
Sleeping
Sleeping
Upload 2 files
Browse files- difpoint/.DS_Store +0 -0
- 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 |
+
|