palette-lab/LipsyncDocker
0
1# Prediction interface for Cog ⚙️2# https://github.com/replicate/cog/blob/main/docs/python.md3 4import os5import sys6import argparse7import subprocess8import numpy as np9from tqdm import tqdm10from PIL import Image11from scipy.io import loadmat12import torch13import cv214from cog import BasePredictor, Input, Path15 16sys.path.insert(0, "third_part")17sys.path.insert(0, "third_part/GPEN")18sys.path.insert(0, "third_part/GFPGAN")19 20# 3dmm extraction21from third_part.face3d.util.preprocess import align_img22from third_part.face3d.util.load_mats import load_lm3d23from third_part.face3d.extract_kp_videos import KeypointExtractor24 25# face enhancement26from third_part.GPEN.gpen_face_enhancer import FaceEnhancement27from third_part.GFPGAN.gfpgan import GFPGANer28 29# expression control30from third_part.ganimation_replicate.model.ganimation import GANimationModel31 32from utils import audio33from utils.ffhq_preprocess import Croper34from utils.alignment_stit import crop_faces, calc_alignment_coefficients, paste_image35from utils.inference_utils import (36 Laplacian_Pyramid_Blending_with_mask,37 face_detect,38 load_model,39 options,40 split_coeff,41 trans_image,42 transform_semantic,43 find_crop_norm_ratio,44 load_face3d_net,45 exp_aus_dict,46)47 48 49class Predictor(BasePredictor):50 def setup(self) -> None:51 """Load the model into memory to make running multiple predictions efficient"""52 self.enhancer = FaceEnhancement(53 base_dir="checkpoints",54 size=512,55 model="GPEN-BFR-512",56 use_sr=False,57 sr_model="rrdb_realesrnet_psnr",58 channel_multiplier=2,59 narrow=1,60 device="cuda",61 )62 self.restorer = GFPGANer(63 model_path="checkpoints/GFPGANv1.3.pth",64 upscale=1,65 arch="clean",66 channel_multiplier=2,67 bg_upsampler=None,68 )69 self.croper = Croper("checkpoints/shape_predictor_68_face_landmarks.dat")70 self.kp_extractor = KeypointExtractor()71 72 face3d_net_path = "checkpoints/face3d_pretrain_epoch_20.pth"73 74 self.net_recon = load_face3d_net(face3d_net_path, "cuda")75 self.lm3d_std = load_lm3d("checkpoints/BFM")76 77 def predict(78 self,79 face: Path = Input(description="Input video file of a talking-head."),80 input_audio: Path = Input(description="Input audio file."),81 ) -> Path:82 """Run a single prediction on the model"""83 device = "cuda"84 args = argparse.Namespace(85 DNet_path="checkpoints/DNet.pt",86 LNet_path="checkpoints/LNet.pth",87 ENet_path="checkpoints/ENet.pth",88 face3d_net_path="checkpoints/face3d_pretrain_epoch_20.pth",89 face=str(face),90 audio=str(input_audio),91 exp_img="neutral",92 outfile=None,93 fps=25,94 pads=[0, 20, 0, 0],95 face_det_batch_size=4,96 LNet_batch_size=16,97 img_size=384,98 crop=[0, -1, 0, -1],99 box=[-1, -1, -1, -1],100 nosmooth=False,101 static=False,102 up_face="original",103 one_shot=False,104 without_rl1=False,105 tmp_dir="temp",106 re_preprocess=False,107 )108 109 base_name = args.face.split("/")[-1]110 111 if args.face.split(".")[1] in ["jpg", "png", "jpeg"]:112 full_frames = [cv2.imread(args.face)]113 args.static = True114 fps = args.fps115 else:116 video_stream = cv2.VideoCapture(args.face)117 fps = video_stream.get(cv2.CAP_PROP_FPS)118 full_frames = []119 while True:120 still_reading, frame = video_stream.read()121 if not still_reading:122 video_stream.release()123 break124 y1, y2, x1, x2 = args.crop125 if x2 == -1:126 x2 = frame.shape[1]127 if y2 == -1:128 y2 = frame.shape[0]129 frame = frame[y1:y2, x1:x2]130 full_frames.append(frame)131 132 full_frames_RGB = [133 cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) for frame in full_frames134 ]135 full_frames_RGB, crop, quad = self.croper.crop(full_frames_RGB, xsize=512)136 137 clx, cly, crx, cry = crop138 lx, ly, rx, ry = quad139 lx, ly, rx, ry = int(lx), int(ly), int(rx), int(ry)140 oy1, oy2, ox1, ox2 = (141 cly + ly,142 min(cly + ry, full_frames[0].shape[0]),143 clx + lx,144 min(clx + rx, full_frames[0].shape[1]),145 )146 # original_size = (ox2 - ox1, oy2 - oy1)147 frames_pil = [148 Image.fromarray(cv2.resize(frame, (256, 256))) for frame in full_frames_RGB149 ]150 151 # get the landmark according to the detected face.152 if (153 not os.path.isfile("temp/" + base_name + "_landmarks.txt")154 or args.re_preprocess155 ):156 print("[Step 1] Landmarks Extraction in Video.")157 lm = self.kp_extractor.extract_keypoint(158 frames_pil, "./temp/" + base_name + "_landmarks.txt"159 )160 else:161 print("[Step 1] Using saved landmarks.")162 lm = np.loadtxt("temp/" + base_name + "_landmarks.txt").astype(np.float32)163 lm = lm.reshape([len(full_frames), -1, 2])164 165 if (166 not os.path.isfile("temp/" + base_name + "_coeffs.npy")167 or args.exp_img is not None168 or args.re_preprocess169 ):170 video_coeffs = []171 for idx in tqdm(172 range(len(frames_pil)), desc="[Step 2] 3DMM Extraction In Video:"173 ):174 frame = frames_pil[idx]175 W, H = frame.size176 lm_idx = lm[idx].reshape([-1, 2])177 if np.mean(lm_idx) == -1:178 lm_idx = (self.lm3d_std[:, :2] + 1) / 2.0179 lm_idx = np.concatenate([lm_idx[:, :1] * W, lm_idx[:, 1:2] * H], 1)180 else:181 lm_idx[:, -1] = H - 1 - lm_idx[:, -1]182 183 trans_params, im_idx, lm_idx, _ = align_img(184 frame, lm_idx, self.lm3d_std185 )186 trans_params = np.array(187 [float(item) for item in np.hsplit(trans_params, 5)]188 ).astype(np.float32)189 im_idx_tensor = (190 torch.tensor(np.array(im_idx) / 255.0, dtype=torch.float32)191 .permute(2, 0, 1)192 .to(device)193 .unsqueeze(0)194 )195 with torch.no_grad():196 coeffs = split_coeff(self.net_recon(im_idx_tensor))197 198 pred_coeff = {key: coeffs[key].cpu().numpy() for key in coeffs}199 pred_coeff = np.concatenate(200 [201 pred_coeff["id"],202 pred_coeff["exp"],203 pred_coeff["tex"],204 pred_coeff["angle"],205 pred_coeff["gamma"],206 pred_coeff["trans"],207 trans_params[None],208 ],209 1,210 )211 video_coeffs.append(pred_coeff)212 semantic_npy = np.array(video_coeffs)[:, 0]213 np.save("temp/" + base_name + "_coeffs.npy", semantic_npy)214 else:215 print("[Step 2] Using saved coeffs.")216 semantic_npy = np.load("temp/" + base_name + "_coeffs.npy").astype(217 np.float32218 )219 220 # generate the 3dmm coeff from a single image221 if args.exp_img == "smile":222 expression = torch.tensor(223 loadmat("checkpoints/expression.mat")["expression_mouth"]224 )[0]225 else:226 print("using expression center")227 expression = torch.tensor(228 loadmat("checkpoints/expression.mat")["expression_center"]229 )[0]230 231 # load DNet, model(LNet and ENet)232 D_Net, model = load_model(args, device)233 234 if (235 not os.path.isfile("temp/" + base_name + "_stablized.npy")236 or args.re_preprocess237 ):238 imgs = []239 for idx in tqdm(240 range(len(frames_pil)),241 desc="[Step 3] Stabilize the expression In Video:",242 ):243 if args.one_shot:244 source_img = trans_image(frames_pil[0]).unsqueeze(0).to(device)245 semantic_source_numpy = semantic_npy[0:1]246 else:247 source_img = trans_image(frames_pil[idx]).unsqueeze(0).to(device)248 semantic_source_numpy = semantic_npy[idx : idx + 1]249 ratio = find_crop_norm_ratio(semantic_source_numpy, semantic_npy)250 coeff = (251 transform_semantic(semantic_npy, idx, ratio).unsqueeze(0).to(device)252 )253 254 # hacking the new expression255 coeff[:, :64, :] = expression[None, :64, None].to(device)256 with torch.no_grad():257 output = D_Net(source_img, coeff)258 img_stablized = np.uint8(259 (260 output["fake_image"]261 .squeeze(0)262 .permute(1, 2, 0)263 .cpu()264 .clamp_(-1, 1)265 .numpy()266 + 1267 )268 / 2.0269 * 255270 )271 imgs.append(cv2.cvtColor(img_stablized, cv2.COLOR_RGB2BGR))272 np.save("temp/" + base_name + "_stablized.npy", imgs)273 del D_Net274 else:275 print("[Step 3] Using saved stabilized video.")276 imgs = np.load("temp/" + base_name + "_stablized.npy")277 torch.cuda.empty_cache()278 279 if not args.audio.endswith(".wav"):280 command = "ffmpeg -loglevel error -y -i {} -strict -2 {}".format(281 args.audio, "temp/{}/temp.wav".format(args.tmp_dir)282 )283 subprocess.call(command, shell=True)284 args.audio = "temp/{}/temp.wav".format(args.tmp_dir)285 wav = audio.load_wav(args.audio, 16000)286 mel = audio.melspectrogram(wav)287 if np.isnan(mel.reshape(-1)).sum() > 0:288 raise ValueError(289 "Mel contains nan! Using a TTS voice? Add a small epsilon noise to the wav file and try again"290 )291 292 mel_step_size, mel_idx_multiplier, i, mel_chunks = 16, 80.0 / fps, 0, []293 while True:294 start_idx = int(i * mel_idx_multiplier)295 if start_idx + mel_step_size > len(mel[0]):296 mel_chunks.append(mel[:, len(mel[0]) - mel_step_size :])297 break298 mel_chunks.append(mel[:, start_idx : start_idx + mel_step_size])299 i += 1300 301 print("[Step 4] Load audio; Length of mel chunks: {}".format(len(mel_chunks)))302 imgs = imgs[: len(mel_chunks)]303 full_frames = full_frames[: len(mel_chunks)]304 lm = lm[: len(mel_chunks)]305 306 imgs_enhanced = []307 for idx in tqdm(range(len(imgs)), desc="[Step 5] Reference Enhancement"):308 img = imgs[idx]309 pred, _, _ = self.enhancer.process(310 img, img, face_enhance=True, possion_blending=False311 )312 imgs_enhanced.append(pred)313 gen = datagen(314 imgs_enhanced.copy(), mel_chunks, full_frames, args, (oy1, oy2, ox1, ox2)315 )316 317 frame_h, frame_w = full_frames[0].shape[:-1]318 out = cv2.VideoWriter(319 "temp/{}/result.mp4".format(args.tmp_dir),320 cv2.VideoWriter_fourcc(*"mp4v"),321 fps,322 (frame_w, frame_h),323 )324 325 if args.up_face != "original":326 instance = GANimationModel()327 instance.initialize()328 instance.setup()329 330 # kp_extractor = KeypointExtractor()331 for i, (332 img_batch,333 mel_batch,334 frames,335 coords,336 img_original,337 f_frames,338 ) in enumerate(339 tqdm(340 gen,341 desc="[Step 6] Lip Synthesis:",342 total=int(np.ceil(float(len(mel_chunks)) / args.LNet_batch_size)),343 )344 ):345 img_batch = torch.FloatTensor(np.transpose(img_batch, (0, 3, 1, 2))).to(346 device347 )348 mel_batch = torch.FloatTensor(np.transpose(mel_batch, (0, 3, 1, 2))).to(349 device350 )351 img_original = (352 torch.FloatTensor(np.transpose(img_original, (0, 3, 1, 2))).to(device)353 / 255.0354 ) # BGR -> RGB355 356 with torch.no_grad():357 incomplete, reference = torch.split(img_batch, 3, dim=1)358 pred, low_res = model(mel_batch, img_batch, reference)359 pred = torch.clamp(pred, 0, 1)360 361 if args.up_face in ["sad", "angry", "surprise"]:362 tar_aus = exp_aus_dict[args.up_face]363 else:364 pass365 366 if args.up_face == "original":367 cur_gen_faces = img_original368 else:369 test_batch = {370 "src_img": torch.nn.functional.interpolate(371 (img_original * 2 - 1), size=(128, 128), mode="bilinear"372 ),373 "tar_aus": tar_aus.repeat(len(incomplete), 1),374 }375 instance.feed_batch(test_batch)376 instance.forward()377 cur_gen_faces = torch.nn.functional.interpolate(378 instance.fake_img / 2.0 + 0.5, size=(384, 384), mode="bilinear"379 )380 381 if args.without_rl1 is not False:382 incomplete, reference = torch.split(img_batch, 3, dim=1)383 mask = torch.where(384 incomplete == 0,385 torch.ones_like(incomplete),386 torch.zeros_like(incomplete),387 )388 pred = pred * mask + cur_gen_faces * (1 - mask)389 390 pred = pred.cpu().numpy().transpose(0, 2, 3, 1) * 255.0391 392 torch.cuda.empty_cache()393 for p, f, xf, c in zip(pred, frames, f_frames, coords):394 y1, y2, x1, x2 = c395 p = cv2.resize(p.astype(np.uint8), (x2 - x1, y2 - y1))396 397 ff = xf.copy()398 ff[y1:y2, x1:x2] = p399 400 # month region enhancement by GFPGAN401 cropped_faces, restored_faces, restored_img = self.restorer.enhance(402 ff, has_aligned=False, only_center_face=True, paste_back=True403 )404 # 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12,405 mm = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 255, 255, 255, 0, 0, 0, 0, 0, 0]406 mouse_mask = np.zeros_like(restored_img)407 tmp_mask = self.enhancer.faceparser.process(408 restored_img[y1:y2, x1:x2], mm409 )[0]410 mouse_mask[y1:y2, x1:x2] = (411 cv2.resize(tmp_mask, (x2 - x1, y2 - y1))[:, :, np.newaxis] / 255.0412 )413 414 height, width = ff.shape[:2]415 restored_img, ff, full_mask = [416 cv2.resize(x, (512, 512))417 for x in (restored_img, ff, np.float32(mouse_mask))418 ]419 img = Laplacian_Pyramid_Blending_with_mask(420 restored_img, ff, full_mask[:, :, 0], 10421 )422 pp = np.uint8(cv2.resize(np.clip(img, 0, 255), (width, height)))423 424 pp, orig_faces, enhanced_faces = self.enhancer.process(425 pp, xf, bbox=c, face_enhance=False, possion_blending=True426 )427 out.write(pp)428 out.release()429 430 output_file = "/tmp/output.mp4"431 command = "ffmpeg -loglevel error -y -i {} -i {} -strict -2 -q:v 1 {}".format(432 args.audio, "temp/{}/result.mp4".format(args.tmp_dir), output_file433 )434 subprocess.call(command, shell=True)435 436 return Path(output_file)437 438 439# frames:256x256, full_frames: original size440def datagen(frames, mels, full_frames, args, cox):441 img_batch, mel_batch, frame_batch, coords_batch, ref_batch, full_frame_batch = (442 [],443 [],444 [],445 [],446 [],447 [],448 )449 base_name = args.face.split("/")[-1]450 refs = []451 image_size = 256452 453 # original frames454 kp_extractor = KeypointExtractor()455 fr_pil = [Image.fromarray(frame) for frame in frames]456 lms = kp_extractor.extract_keypoint(457 fr_pil, "temp/" + base_name + "x12_landmarks.txt"458 )459 frames_pil = [460 (lm, frame) for frame, lm in zip(fr_pil, lms)461 ] # frames is the croped version of modified face462 crops, orig_images, quads = crop_faces(463 image_size, frames_pil, scale=1.0, use_fa=True464 )465 inverse_transforms = [466 calc_alignment_coefficients(467 quad + 0.5,468 [[0, 0], [0, image_size], [image_size, image_size], [image_size, 0]],469 )470 for quad in quads471 ]472 del kp_extractor.detector473 474 oy1, oy2, ox1, ox2 = cox475 face_det_results = face_detect(full_frames, args, jaw_correction=True)476 477 for inverse_transform, crop, full_frame, face_det in zip(478 inverse_transforms, crops, full_frames, face_det_results479 ):480 imc_pil = paste_image(481 inverse_transform,482 crop,483 Image.fromarray(484 cv2.resize(485 full_frame[int(oy1) : int(oy2), int(ox1) : int(ox2)], (256, 256)486 )487 ),488 )489 490 ff = full_frame.copy()491 ff[int(oy1) : int(oy2), int(ox1) : int(ox2)] = cv2.resize(492 np.array(imc_pil.convert("RGB")), (ox2 - ox1, oy2 - oy1)493 )494 oface, coords = face_det495 y1, y2, x1, x2 = coords496 refs.append(ff[y1:y2, x1:x2])497 498 for i, m in enumerate(mels):499 idx = 0 if args.static else i % len(frames)500 frame_to_save = frames[idx].copy()501 face = refs[idx]502 oface, coords = face_det_results[idx].copy()503 504 face = cv2.resize(face, (args.img_size, args.img_size))505 oface = cv2.resize(oface, (args.img_size, args.img_size))506 507 img_batch.append(oface)508 ref_batch.append(face)509 mel_batch.append(m)510 coords_batch.append(coords)511 frame_batch.append(frame_to_save)512 full_frame_batch.append(full_frames[idx].copy())513 514 if len(img_batch) >= args.LNet_batch_size:515 img_batch, mel_batch, ref_batch = (516 np.asarray(img_batch),517 np.asarray(mel_batch),518 np.asarray(ref_batch),519 )520 img_masked = img_batch.copy()521 img_original = img_batch.copy()522 img_masked[:, args.img_size // 2 :] = 0523 img_batch = np.concatenate((img_masked, ref_batch), axis=3) / 255.0524 mel_batch = np.reshape(525 mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1]526 )527 528 yield img_batch, mel_batch, frame_batch, coords_batch, img_original, full_frame_batch529 (530 img_batch,531 mel_batch,532 frame_batch,533 coords_batch,534 img_original,535 full_frame_batch,536 ref_batch,537 ) = ([], [], [], [], [], [], [])538 539 if len(img_batch) > 0:540 img_batch, mel_batch, ref_batch = (541 np.asarray(img_batch),542 np.asarray(mel_batch),543 np.asarray(ref_batch),544 )545 img_masked = img_batch.copy()546 img_original = img_batch.copy()547 img_masked[:, args.img_size // 2 :] = 0548 img_batch = np.concatenate((img_masked, ref_batch), axis=3) / 255.0549 mel_batch = np.reshape(550 mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1]551 )552 yield img_batch, mel_batch, frame_batch, coords_batch, img_original, full_frame_batch553