abdul168/Free-View_Expressive_Talking_Head_Video_Editing

1import os2 3# set CUDA_MODULE_LOADING=LAZY to speed up the serverless function4os.environ["CUDA_MODULE_LOADING"] = "LAZY"5# set SAFETENSORS_FAST_GPU=1 to speed up the serverless function6os.environ["SAFETENSORS_FAST_GPU"] = "1"7import cv28import torch9import time10import imageio11import numpy as np12from tqdm import tqdm13import moviepy.editor as mp14import torch15 16from audio import load_wav, melspectrogram17from fete_model import FETE_model18from preprocess_videos import face_detect, load_from_npz19 20fps = 2521mel_idx_multiplier = 80.0 / fps22 23mel_step_size = 1624batch_size = 64 if torch.cuda.is_available() else 425device = "cuda" if torch.cuda.is_available() else "cpu"26print("Using {} for inference.".format(device))27use_fp16 = True if torch.cuda.is_available() else False28print("Using FP16 for inference.") if use_fp16 else None29torch.backends.cudnn.benchmark = True if device == "cuda" else False30 31 32def init_model():33    checkpoint_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "checkpoints/obama-fp16.safetensors")34    model = FETE_model()35    if checkpoint_path.endswith(".pth") or checkpoint_path.endswith(".ckpt"):36        if device == "cuda":37            checkpoint = torch.load(checkpoint_path)38        else:39            checkpoint = torch.load(checkpoint_path, map_location=lambda storage, loc: storage)40        s = checkpoint["state_dict"]41    else:42        from safetensors import safe_open43 44        s = {}45        with safe_open(checkpoint_path, framework="pt", device=device) as f:46            for key in f.keys():47                s[key] = f.get_tensor(key)48    new_s = {}49    for k, v in s.items():50        new_s[k.replace("module.", "")] = v51    model.load_state_dict(new_s)52 53    model = model.to(device)54    model.eval()55    print("Model loaded")56    if use_fp16:57        for name, module in model.named_modules():58            if ".query_conv" in name or ".key_conv" in name or ".value_conv" in name:59                # keep attention layers in full precision to avoid error60                module.to(torch.float)61            else:62                module.to(torch.half)63        print("Model converted to half precision to accelerate inference")64    return model65 66 67def make_mask(image_size=256, border_size=32):68    mask_bar = np.linspace(1, 0, border_size).reshape(1, -1).repeat(image_size, axis=0)69    mask = np.zeros((image_size, image_size), dtype=np.float32)70    mask[-border_size:, :] += mask_bar.T[::-1]71    mask[:, :border_size] = mask_bar72    mask[:, -border_size:] = mask_bar[:, ::-1]73    mask[-border_size:, :][mask[-border_size:, :] < 0.6] = 0.674    mask = np.stack([mask] * 3, axis=-1).astype(np.float32)75    return mask76 77 78face_mask = make_mask()79 80 81def blend_images(foreground, background):82    # Blend the foreground and background images using the mask83    temp_mask = cv2.resize(face_mask, (foreground.shape[1], foreground.shape[0]))84    blended = cv2.multiply(foreground.astype(np.float32), temp_mask)85    blended += cv2.multiply(background.astype(np.float32), 1 - temp_mask)86    blended = np.clip(blended, 0, 255).astype(np.uint8)87    return blended88 89 90def smooth_coord(last_coord, current_coord, factor=0.4):91    change = np.array(current_coord) - np.array(last_coord)92    change = change * factor93    return (np.array(last_coord) + np.array(change)).astype(int).tolist()94 95 96def add_black(imgs):97    for i in range(len(imgs)):98        # print('x', imgs[i].shape)99        imgs[i] = cv2.vconcat(100            [np.zeros((100, imgs[i].shape[1], 3), dtype=np.uint8), imgs[i], np.zeros((20, imgs[i].shape[1], 3), dtype=np.uint8)]101        )102        # 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,:]103 104        # print('xx', imgs[i].shape)105    return imgs106 107 108def remove_black(img):109    return img[100:-20]110 111 112def resize_length(input_attributes, length):113    input_attributes = np.array(input_attributes)114    resized_attributes = [input_attributes[int(i_ * (input_attributes.shape[0] / length))] for i_ in range(length)]115    return np.array(resized_attributes).T116 117 118def output_chunks(input_attributes):119    output_chunks = []120    len_ = len(input_attributes[0])121 122    i = 0123    # print(mel.shape, pose.shape)124    # (80, 801) (3, 801)125    while 1:126        start_idx = int(i * mel_idx_multiplier)127        if start_idx + mel_step_size > len_:128            output_chunks.append(input_attributes[:, len_ - mel_step_size :])129            break130        output_chunks.append(input_attributes[:, start_idx : start_idx + mel_step_size])131        i += 1132    return output_chunks133 134 135def prepare_data(face_path, audio_path, pose, emotion, blink, img_size=256, pads=[0, 0, 0, 0]):136    if os.path.isfile(face_path) and face_path.split(".")[1] in ["jpg", "png", "jpeg"]:137        static = True138        full_frames = [cv2.imread(face_path)]139    else:140        static = False141        video_stream = cv2.VideoCapture(face_path)142 143        # print('Reading video frames...')144        full_frames = []145        while 1:146            still_reading, frame = video_stream.read()147            if not still_reading:148                video_stream.release()149                break150            full_frames.append(frame)151    print("Number of frames available for inference: " + str(len(full_frames)))152 153    wav = load_wav(audio_path, 16000)154    mel = melspectrogram(wav)155    # take half156    len_ = mel.shape[1]  #  //2157    mel = mel[:, :len_]158    # print('>>>', mel.shape)159 160    pose = resize_length(pose, len_)161    emotion = resize_length(emotion, len_)162    blink = resize_length(blink, len_)163 164    if np.isnan(mel.reshape(-1)).sum() > 0:165        raise ValueError("Mel contains nan! Using a TTS voice? Add a small epsilon noise to the wav file and try again")166 167    mel_chunks = output_chunks(mel)168    pose_chunks = output_chunks(pose)169    emotion_chunks = output_chunks(emotion)170    blink_chunks = output_chunks(blink)171 172    gen = datagen(face_path, full_frames, mel_chunks, pose_chunks, emotion_chunks, blink_chunks, static=static, img_size=img_size, pads=pads)173    steps = int(np.ceil(float(len(mel_chunks)) / batch_size))174 175    return gen, steps176 177 178def preprocess_batch(batch):179    return torch.FloatTensor(np.reshape(batch, [len(batch), 1, batch[0].shape[0], batch[0].shape[1]])).to(device)180 181 182def datagen(face_path, frames, mels, poses, emotions, blinks, static=False, img_size=256, pads=[0, 0, 0, 0]):183    img_batch, mel_batch, pose_batch, emotion_batch, blink_batch, frame_batch, coords_batch = [], [], [], [], [], [], []184    scale_factor = img_size // 128185 186    # print("Length of mel chunks: {}".format(len(mel_chunks)))187    frames = frames[: len(mels)]188    frames = add_black(frames)189    try:190        video_name = os.path.basename(face_path).split(".")[0]191        coords = load_from_npz(video_name)192        face_det_results = [[image[y1:y2, x1:x2], (y1, y2, x1, x2)] for image, (x1, y1, x2, y2) in zip(frames, coords)]193 194    except Exception as e:195        print("No existing coords found, running face detection...", "Error: ", e)196        if not static:197            coords = face_detect(frames, pads)198            face_det_results = [[image[y1:y2, x1:x2], (y1, y2, x1, x2)] for image, (x1, y1, x2, y2) in zip(frames, coords)]199        else:200            coords = face_detect([frames[0]], pads)201            face_det_results = [[image[y1:y2, x1:x2], (y1, y2, x1, x2)] for image, (x1, y1, x2, y2) in zip(frames, coords)]202 203    face_det_results = face_det_results[: len(mels)]204 205    while len(frames) < len(mels):206        face_det_results = face_det_results + face_det_results[::-1]207        frames = frames + frames[::-1]208    else:209        face_det_results = face_det_results[: len(mels)]210        frames = frames[: len(mels)]211 212    for i in range(len(mels)):213        idx = 0 if static else i % len(frames)214        frame_to_save = frames[idx].copy()215        face, coords = face_det_results[idx].copy()216        face = cv2.resize(face, (img_size, img_size))217 218        img_batch.append(face)219        mel_batch.append(mels[i])220        pose_batch.append(poses[i])221        emotion_batch.append(emotions[i])222        blink_batch.append(blinks[i])223        frame_batch.append(frame_to_save)224        coords_batch.append(coords)225 226        # print(m.shape, poses[i].shape)227        # (80, 16) (3, 16)228        if len(img_batch) >= batch_size:229            img_masked = np.asarray(img_batch).copy()230 231            img_masked[:, 16 * scale_factor : -16 * scale_factor, 16 * scale_factor : -16 * scale_factor] = 0.0232 233            img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255.0234            img_batch = torch.FloatTensor(np.transpose(img_batch, (0, 3, 1, 2))).to(device)235 236            mel_batch = preprocess_batch(mel_batch)237            pose_batch = preprocess_batch(pose_batch)238            emotion_batch = preprocess_batch(emotion_batch)239            blink_batch = preprocess_batch(blink_batch)240 241            if use_fp16:242                yield (243                    img_batch.half(),244                    mel_batch.half(),245                    pose_batch.half(),246                    emotion_batch.half(),247                    blink_batch.half(),248                ), frame_batch, coords_batch249            else:250                yield (img_batch, mel_batch, pose_batch, emotion_batch, blink_batch), frame_batch, coords_batch251            img_batch, mel_batch, pose_batch, emotion_batch, blink_batch, frame_batch, coords_batch = [], [], [], [], [], [], []252 253    if len(img_batch) > 0:254        img_masked = np.asarray(img_batch).copy()255 256        img_masked[:, 16 * scale_factor : -16 * scale_factor, 16 * scale_factor : -16 * scale_factor] = 0.0257 258        img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255.0259        img_batch = torch.FloatTensor(np.transpose(img_batch, (0, 3, 1, 2))).to(device)260 261        mel_batch = preprocess_batch(mel_batch)262        pose_batch = preprocess_batch(pose_batch)263        emotion_batch = preprocess_batch(emotion_batch)264        blink_batch = preprocess_batch(blink_batch)265 266        if use_fp16:267            yield (img_batch.half(), mel_batch.half(), pose_batch.half(), emotion_batch.half(), blink_batch.half()), frame_batch, coords_batch268        else:269            yield (img_batch, mel_batch, pose_batch, emotion_batch, blink_batch), frame_batch, coords_batch270 271 272def infenrece(model, face_path, audio_path, pose, emotion, blink, preview=False):273    timestamp = time.strftime("%Y-%m-%d-%H-%M-%S", time.gmtime(time.time()))274    gen, steps = prepare_data(face_path, audio_path, pose, emotion, blink)275    steps = 1 if preview else steps276    # duration = librosa.get_duration(filename=audio_path)277 278    if preview:279        outfile = "/tmp/{}.jpg".format(timestamp)280    else:281        outfile = "/tmp/{}.mp4".format(timestamp)282        tmp_video = "/tmp/temp_{}.mp4".format(timestamp)283        writer = (284            imageio.get_writer(tmp_video, fps=fps, codec="libx264", quality=10, pixelformat="yuv420p", macro_block_size=1)285            if not preview286            else None287        )288    # print('Generating frames...', outfile, steps)289    for inputs, frames, coords in tqdm(gen, total=steps):290        with torch.no_grad():291            pred = model(*inputs)292 293        pred = pred.cpu().numpy().transpose(0, 2, 3, 1) * 255.0294 295        for p, f, c in zip(pred, frames, coords):296            y1, y2, x1, x2 = c297            y1, y2, x1, x2 = int(y1), int(y2), int(x1), int(x2)298            y = round(y2 - y1)299            x = round(x2 - x1)300            p = cv2.resize(p.astype(np.uint8), (x, y))301 302            try:303                f[y1 : y1 + y, x1 : x1 + x] = blend_images(f[y1 : y1 + y, x1 : x1 + x], p)304            except Exception as e:305                print(e)306                f[y1 : y1 + y, x1 : x1 + x] = p307            f = remove_black(f)308            if preview:309                cv2.imwrite(outfile, f, [int(cv2.IMWRITE_JPEG_QUALITY), 95])310                return outfile311            writer.append_data(cv2.cvtColor(f, cv2.COLOR_BGR2RGB))312    writer.close()313    video_clip = mp.VideoFileClip(tmp_video)314    audio_clip = mp.AudioFileClip(audio_path)315    video_clip = video_clip.set_audio(audio_clip)316    video_clip.write_videofile(outfile, codec="libx264")317 318    print("Saved to {}".format(outfile) if os.path.exists(outfile) else "Failed to save {}".format(outfile))319    try:320        os.remove(tmp_video)321        del video_clip322        del audio_clip323        del gen324    except:325        pass326    return outfile327 328 329if __name__ == "__main__":330    model = init_model()331 332    from attributtes_utils import input_pose, input_emotion, input_blink333 334    pose = input_pose()335    emotion = input_emotion()336    blink = input_blink()337    audio_path = "./assets/sample.wav"338    face_path = "./assets/sample.mp4"339 340    infenrece(model, face_path, audio_path, pose, emotion, blink)341