import os # set CUDA_MODULE_LOADING=LAZY to speed up the serverless function os.environ["CUDA_MODULE_LOADING"] = "LAZY" # set SAFETENSORS_FAST_GPU=1 to speed up the serverless function os.environ["SAFETENSORS_FAST_GPU"] = "1" import cv2 import torch import time import imageio import numpy as np from tqdm import tqdm import moviepy.editor as mp import torch from audio import load_wav, melspectrogram from fete_model import FETE_model from preprocess_videos import face_detect, load_from_npz fps = 25 mel_idx_multiplier = 80.0 / fps mel_step_size = 16 batch_size = 64 if torch.cuda.is_available() else 4 device = "cuda" if torch.cuda.is_available() else "cpu" print("Using {} for inference.".format(device)) use_fp16 = True if torch.cuda.is_available() else False print("Using FP16 for inference.") if use_fp16 else None torch.backends.cudnn.benchmark = True if device == "cuda" else False def init_model(): checkpoint_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "checkpoints/obama-fp16.safetensors") model = FETE_model() if checkpoint_path.endswith(".pth") or checkpoint_path.endswith(".ckpt"): if device == "cuda": checkpoint = torch.load(checkpoint_path) else: checkpoint = torch.load(checkpoint_path, map_location=lambda storage, loc: storage) s = checkpoint["state_dict"] else: from safetensors import safe_open s = {} with safe_open(checkpoint_path, framework="pt", device=device) as f: for key in f.keys(): s[key] = f.get_tensor(key) new_s = {} for k, v in s.items(): new_s[k.replace("module.", "")] = v model.load_state_dict(new_s) model = model.to(device) model.eval() print("Model loaded") if use_fp16: for name, module in model.named_modules(): if ".query_conv" in name or ".key_conv" in name or ".value_conv" in name: # keep attention layers in full precision to avoid error module.to(torch.float) else: module.to(torch.half) print("Model converted to half precision to accelerate inference") return model def make_mask(image_size=256, border_size=32): mask_bar = np.linspace(1, 0, border_size).reshape(1, -1).repeat(image_size, axis=0) mask = np.zeros((image_size, image_size), dtype=np.float32) mask[-border_size:, :] += mask_bar.T[::-1] mask[:, :border_size] = mask_bar mask[:, -border_size:] = mask_bar[:, ::-1] mask[-border_size:, :][mask[-border_size:, :] < 0.6] = 0.6 mask = np.stack([mask] * 3, axis=-1).astype(np.float32) return mask face_mask = make_mask() def blend_images(foreground, background): # Blend the foreground and background images using the mask temp_mask = cv2.resize(face_mask, (foreground.shape[1], foreground.shape[0])) blended = cv2.multiply(foreground.astype(np.float32), temp_mask) blended += cv2.multiply(background.astype(np.float32), 1 - temp_mask) blended = np.clip(blended, 0, 255).astype(np.uint8) return blended def smooth_coord(last_coord, current_coord, factor=0.4): change = np.array(current_coord) - np.array(last_coord) change = change * factor return (np.array(last_coord) + np.array(change)).astype(int).tolist() def add_black(imgs): for i in range(len(imgs)): # print('x', imgs[i].shape) imgs[i] = cv2.vconcat( [np.zeros((100, imgs[i].shape[1], 3), dtype=np.uint8), imgs[i], np.zeros((20, imgs[i].shape[1], 3), dtype=np.uint8)] ) # imgs[i] = cv2.hconcat([np.zeros((imgs[i].shape[0], 100, 3), dtype=np.uint8), imgs[i], np.zeros((imgs[i].shape[0], 100, 3), dtype=np.uint8)])[:480+150,740-100:-740+100,:] # print('xx', imgs[i].shape) return imgs def remove_black(img): return img[100:-20] def resize_length(input_attributes, length): input_attributes = np.array(input_attributes) resized_attributes = [input_attributes[int(i_ * (input_attributes.shape[0] / length))] for i_ in range(length)] return np.array(resized_attributes).T def output_chunks(input_attributes): output_chunks = [] len_ = len(input_attributes[0]) i = 0 # print(mel.shape, pose.shape) # (80, 801) (3, 801) while 1: start_idx = int(i * mel_idx_multiplier) if start_idx + mel_step_size > len_: output_chunks.append(input_attributes[:, len_ - mel_step_size :]) break output_chunks.append(input_attributes[:, start_idx : start_idx + mel_step_size]) i += 1 return output_chunks def prepare_data(face_path, audio_path, pose, emotion, blink, img_size=256, pads=[0, 0, 0, 0]): if os.path.isfile(face_path) and face_path.split(".")[1] in ["jpg", "png", "jpeg"]: static = True full_frames = [cv2.imread(face_path)] else: static = False video_stream = cv2.VideoCapture(face_path) # print('Reading video frames...') full_frames = [] while 1: still_reading, frame = video_stream.read() if not still_reading: video_stream.release() break full_frames.append(frame) print("Number of frames available for inference: " + str(len(full_frames))) wav = load_wav(audio_path, 16000) mel = melspectrogram(wav) # take half len_ = mel.shape[1] # //2 mel = mel[:, :len_] # print('>>>', mel.shape) pose = resize_length(pose, len_) emotion = resize_length(emotion, len_) blink = resize_length(blink, len_) if np.isnan(mel.reshape(-1)).sum() > 0: raise ValueError("Mel contains nan! Using a TTS voice? Add a small epsilon noise to the wav file and try again") mel_chunks = output_chunks(mel) pose_chunks = output_chunks(pose) emotion_chunks = output_chunks(emotion) blink_chunks = output_chunks(blink) gen = datagen(face_path, full_frames, mel_chunks, pose_chunks, emotion_chunks, blink_chunks, static=static, img_size=img_size, pads=pads) steps = int(np.ceil(float(len(mel_chunks)) / batch_size)) return gen, steps def preprocess_batch(batch): return torch.FloatTensor(np.reshape(batch, [len(batch), 1, batch[0].shape[0], batch[0].shape[1]])).to(device) def datagen(face_path, frames, mels, poses, emotions, blinks, static=False, img_size=256, pads=[0, 0, 0, 0]): img_batch, mel_batch, pose_batch, emotion_batch, blink_batch, frame_batch, coords_batch = [], [], [], [], [], [], [] scale_factor = img_size // 128 # print("Length of mel chunks: {}".format(len(mel_chunks))) frames = frames[: len(mels)] frames = add_black(frames) try: video_name = os.path.basename(face_path).split(".")[0] coords = load_from_npz(video_name) face_det_results = [[image[y1:y2, x1:x2], (y1, y2, x1, x2)] for image, (x1, y1, x2, y2) in zip(frames, coords)] except Exception as e: print("No existing coords found, running face detection...", "Error: ", e) if not static: coords = face_detect(frames, pads) face_det_results = [[image[y1:y2, x1:x2], (y1, y2, x1, x2)] for image, (x1, y1, x2, y2) in zip(frames, coords)] else: coords = face_detect([frames[0]], pads) face_det_results = [[image[y1:y2, x1:x2], (y1, y2, x1, x2)] for image, (x1, y1, x2, y2) in zip(frames, coords)] face_det_results = face_det_results[: len(mels)] while len(frames) < len(mels): face_det_results = face_det_results + face_det_results[::-1] frames = frames + frames[::-1] else: face_det_results = face_det_results[: len(mels)] frames = frames[: len(mels)] for i in range(len(mels)): idx = 0 if static else i % len(frames) frame_to_save = frames[idx].copy() face, coords = face_det_results[idx].copy() face = cv2.resize(face, (img_size, img_size)) img_batch.append(face) mel_batch.append(mels[i]) pose_batch.append(poses[i]) emotion_batch.append(emotions[i]) blink_batch.append(blinks[i]) frame_batch.append(frame_to_save) coords_batch.append(coords) # print(m.shape, poses[i].shape) # (80, 16) (3, 16) if len(img_batch) >= batch_size: img_masked = np.asarray(img_batch).copy() img_masked[:, 16 * scale_factor : -16 * scale_factor, 16 * scale_factor : -16 * scale_factor] = 0.0 img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255.0 img_batch = torch.FloatTensor(np.transpose(img_batch, (0, 3, 1, 2))).to(device) mel_batch = preprocess_batch(mel_batch) pose_batch = preprocess_batch(pose_batch) emotion_batch = preprocess_batch(emotion_batch) blink_batch = preprocess_batch(blink_batch) if use_fp16: yield ( img_batch.half(), mel_batch.half(), pose_batch.half(), emotion_batch.half(), blink_batch.half(), ), frame_batch, coords_batch else: yield (img_batch, mel_batch, pose_batch, emotion_batch, blink_batch), frame_batch, coords_batch img_batch, mel_batch, pose_batch, emotion_batch, blink_batch, frame_batch, coords_batch = [], [], [], [], [], [], [] if len(img_batch) > 0: img_masked = np.asarray(img_batch).copy() img_masked[:, 16 * scale_factor : -16 * scale_factor, 16 * scale_factor : -16 * scale_factor] = 0.0 img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255.0 img_batch = torch.FloatTensor(np.transpose(img_batch, (0, 3, 1, 2))).to(device) mel_batch = preprocess_batch(mel_batch) pose_batch = preprocess_batch(pose_batch) emotion_batch = preprocess_batch(emotion_batch) blink_batch = preprocess_batch(blink_batch) if use_fp16: yield (img_batch.half(), mel_batch.half(), pose_batch.half(), emotion_batch.half(), blink_batch.half()), frame_batch, coords_batch else: yield (img_batch, mel_batch, pose_batch, emotion_batch, blink_batch), frame_batch, coords_batch def infenrece(model, face_path, audio_path, pose, emotion, blink, preview=False): timestamp = time.strftime("%Y-%m-%d-%H-%M-%S", time.gmtime(time.time())) gen, steps = prepare_data(face_path, audio_path, pose, emotion, blink) steps = 1 if preview else steps # duration = librosa.get_duration(filename=audio_path) if preview: outfile = "/tmp/{}.jpg".format(timestamp) else: outfile = "/tmp/{}.mp4".format(timestamp) tmp_video = "/tmp/temp_{}.mp4".format(timestamp) writer = imageio.get_writer(tmp_video, fps=fps, codec="libx264", quality=10, pixelformat="yuv420p", macro_block_size=1) if not preview else None # print('Generating frames...', outfile, steps) for inputs, frames, coords in tqdm(gen, total=steps): with torch.no_grad(): pred = model(*inputs) pred = pred.cpu().numpy().transpose(0, 2, 3, 1) * 255.0 for p, f, c in zip(pred, frames, coords): y1, y2, x1, x2 = c y1, y2, x1, x2 = int(y1), int(y2), int(x1), int(x2) y = round(y2 - y1) x = round(x2 - x1) # print(x, y, p.shape) p = cv2.resize(p.astype(np.uint8), (x, y)) try: f[y1 : y1 + y, x1 : x1 + x] = blend_images(f[y1 : y1 + y, x1 : x1 + x], p) except Exception as e: print(e) f[y1 : y1 + y, x1 : x1 + x] = p # out.write(f[100:-20]) f = remove_black(f) if preview: cv2.imwrite(outfile, f, [int(cv2.IMWRITE_JPEG_QUALITY), 95]) return outfile writer.append_data(cv2.cvtColor(f, cv2.COLOR_BGR2RGB)) writer.close() video_clip = mp.VideoFileClip(tmp_video) audio_clip = mp.AudioFileClip(audio_path) video_clip = video_clip.set_audio(audio_clip) video_clip.write_videofile(outfile, codec="libx264") print("Saved to {}".format(outfile) if os.path.exists(outfile) else "Failed to save {}".format(outfile)) try: os.remove(tmp_video) del video_clip del audio_clip del gen except: pass return outfile if __name__ == "__main__": model = init_model() from attributtes_utils import input_pose, input_emotion, input_blink pose = input_pose() emotion = input_emotion() blink = input_blink() audio_path = "./assets/sample.wav" face_path = "./assets/sample.mp4" infenrece(model, face_path, audio_path, pose, emotion, blink)