diff --git a/README.md b/README.md
index 258f069..deb3d89 100644
--- a/README.md
+++ b/README.md
@@ -11,7 +11,7 @@ Chao Zhan,
Wenjiang Zhou
(*Equal Contribution, †Corresponding Author, benbinwu@tencent.com)
-**[github](https://github.com/TMElyralab/MuseTalk)** **[huggingface](https://huggingface.co/TMElyralab/MuseTalk)** **[gradio](https://huggingface.co/spaces/TMElyralab/MuseTalk)** **Project (comming soon)** **Technical report (comming soon)**
+**[github](https://github.com/TMElyralab/MuseTalk)** **[huggingface](https://huggingface.co/TMElyralab/MuseTalk)** **[space](https://huggingface.co/spaces/TMElyralab/MuseTalk)** **Project (comming soon)** **Technical report (comming soon)**
We introduce `MuseTalk`, a **real-time high quality** lip-syncing model (30fps+ on an NVIDIA Tesla V100). MuseTalk can be applied with input videos, e.g., generated by [MuseV](https://github.com/TMElyralab/MuseV), as a complete virtual human solution.
@@ -28,12 +28,13 @@ We introduce `MuseTalk`, a **real-time high quality** lip-syncing model (30fps+
# News
- [04/02/2024] Release MuseTalk project and pretrained models.
- [04/16/2024] Release Gradio [demo](https://huggingface.co/spaces/TMElyralab/MuseTalk) on HuggingFace Spaces (thanks to HF team for their community grant)
+- [04/17/2024] :mega: We release a pipeline that utilizes MuseTalk for real-time inference.
## Model

MuseTalk was trained in latent spaces, where the images were encoded by a freezed VAE. The audio was encoded by a freezed `whisper-tiny` model. The architecture of the generation network was borrowed from the UNet of the `stable-diffusion-v1-4`, where the audio embeddings were fused to the image embeddings by cross-attention.
-Note that although we use a very similar architecture as Stable Diffusion, MuseTalk is distinct in that it is `Not` a diffusion model. Instead, MuseTalk operates by inpainting in the latent space with `a single step`.
+Note that although we use a very similar architecture as Stable Diffusion, MuseTalk is distinct in that it is **NOT** a diffusion model. Instead, MuseTalk operates by inpainting in the latent space with a single step.
## Cases
### MuseV + MuseTalk make human photos alive!
@@ -162,7 +163,7 @@ Note that although we use a very similar architecture as Stable Diffusion, MuseT
# TODO:
- [x] trained models and inference codes.
- [x] Huggingface Gradio [demo](https://huggingface.co/spaces/TMElyralab/MuseTalk).
-- [ ] codes for real-time inference.
+- [x] codes for real-time inference.
- [ ] technical report.
- [ ] training codes.
- [ ] a better model (may take longer).
@@ -262,9 +263,30 @@ python -m scripts.inference --inference_config configs/inference/test.yaml --bbo
As a complete solution to virtual human generation, you are suggested to first apply [MuseV](https://github.com/TMElyralab/MuseV) to generate a video (text-to-video, image-to-video or pose-to-video) by referring [this](https://github.com/TMElyralab/MuseV?tab=readme-ov-file#text2video). Frame interpolation is suggested to increase frame rate. Then, you can use `MuseTalk` to generate a lip-sync video by referring [this](https://github.com/TMElyralab/MuseTalk?tab=readme-ov-file#inference).
-# Note
+#### :new: Real-time inference
-If you want to launch online video chats, you are suggested to generate videos using MuseV and apply necessary pre-processing such as face detection and face parsing in advance. During online chatting, only UNet and the VAE decoder are involved, which makes MuseTalk real-time.
+Here, we provide the inference script. This script first applies necessary pre-processing such as face detection, face parsing and VAE encode in advance. During inference, only UNet and the VAE decoder are involved, which makes MuseTalk real-time.
+```
+python -m scripts.realtime_inference --inference_config configs/inference/realtime.yaml
+```
+configs/inference/realtime.yaml is the path to the real-time inference configuration file, including `preparation`, `video_path` , `bbox_shift` and `audio_clips`.
+
+1. Set `preparation` to `True` in `realtime.yaml` to prepare the materials for a new `avatar`. (If the `bbox_shift` has changed, you also need to re-prepare the materials.)
+1. After that, the `avatar` will use an audio clip selected from `audio_clips` to generate video.
+ ```
+ Inferring using: data/audio/yongen.wav
+ ```
+1. While MuseTalk is inferring, sub-threads can simultaneously stream the results to the users. The generation process can achieve up to 50fps on an NVIDIA Tesla V100.
+ ```
+ 2%|██▍ | 3/141 [00:00<00:32, 4.30it/s] # inference process
+ Generating the 6-th frame with FPS: 48.58 # playing process
+ Generating the 7-th frame with FPS: 48.74
+ Generating the 8-th frame with FPS: 49.17
+ 3%|███▎ | 4/141 [00:00<00:32, 4.21it/s]
+ ```
+1. Set `preparation` to `False` and run this script if you want to genrate more videos using the same avatar.
+
+If you want to generate multiple videos using the same avatar/video, you can also use this script to **SIGNIFICANTLY** expedite the generation process.
# Acknowledgement
diff --git a/configs/inference/realtime.yaml b/configs/inference/realtime.yaml
new file mode 100644
index 0000000..d4092ac
--- /dev/null
+++ b/configs/inference/realtime.yaml
@@ -0,0 +1,10 @@
+avator_1:
+ preparation: False
+ bbox_shift: 5
+ video_path: "data/video/sun.mp4"
+ audio_clips:
+ audio_0: "data/audio/yongen.wav"
+ audio_1: "data/audio/sun.wav"
+
+
+
diff --git a/musetalk/utils/blending.py b/musetalk/utils/blending.py
index 0a9ec3e..d69e435 100644
--- a/musetalk/utils/blending.py
+++ b/musetalk/utils/blending.py
@@ -57,3 +57,44 @@ def get_image(image,face,face_box,upper_boundary_ratio = 0.5,expand=1.2):
body.paste(face_large, crop_box[:2], mask_image)
body = np.array(body)
return body[:,:,::-1]
+
+def get_image_prepare_material(image,face_box,upper_boundary_ratio = 0.5,expand=1.2):
+ body = Image.fromarray(image[:,:,::-1])
+
+ x, y, x1, y1 = face_box
+ #print(x1-x,y1-y)
+ crop_box, s = get_crop_box(face_box, expand)
+ x_s, y_s, x_e, y_e = crop_box
+
+ face_large = body.crop(crop_box)
+ ori_shape = face_large.size
+
+ mask_image = face_seg(face_large)
+ mask_small = mask_image.crop((x-x_s, y-y_s, x1-x_s, y1-y_s))
+ mask_image = Image.new('L', ori_shape, 0)
+ mask_image.paste(mask_small, (x-x_s, y-y_s, x1-x_s, y1-y_s))
+
+ # keep upper_boundary_ratio of talking area
+ width, height = mask_image.size
+ top_boundary = int(height * upper_boundary_ratio)
+ modified_mask_image = Image.new('L', ori_shape, 0)
+ modified_mask_image.paste(mask_image.crop((0, top_boundary, width, height)), (0, top_boundary))
+
+ blur_kernel_size = int(0.1 * ori_shape[0] // 2 * 2) + 1
+ mask_array = cv2.GaussianBlur(np.array(modified_mask_image), (blur_kernel_size, blur_kernel_size), 0)
+ return mask_array,crop_box
+
+def get_image_blending(image,face,face_box,mask_array,crop_box):
+ body = Image.fromarray(image[:,:,::-1])
+ face = Image.fromarray(face[:,:,::-1])
+
+ x, y, x1, y1 = face_box
+ x_s, y_s, x_e, y_e = crop_box
+ face_large = body.crop(crop_box)
+
+ mask_image = Image.fromarray(mask_array)
+ mask_image = mask_image.convert("L")
+ face_large.paste(face, (x-x_s, y-y_s, x1-x_s, y1-y_s))
+ body.paste(face_large, crop_box[:2], mask_image)
+ body = np.array(body)
+ return body[:,:,::-1]
\ No newline at end of file
diff --git a/scripts/realtime_inference.py b/scripts/realtime_inference.py
new file mode 100644
index 0000000..07874c9
--- /dev/null
+++ b/scripts/realtime_inference.py
@@ -0,0 +1,295 @@
+import argparse
+import os
+from omegaconf import OmegaConf
+import numpy as np
+import cv2
+import torch
+import glob
+import pickle
+import sys
+from tqdm import tqdm
+import copy
+import json
+from musetalk.utils.utils import get_file_type,get_video_fps,datagen
+from musetalk.utils.preprocessing import get_landmark_and_bbox,read_imgs,coord_placeholder
+from musetalk.utils.blending import get_image,get_image_prepare_material,get_image_blending
+from musetalk.utils.utils import load_all_model
+import shutil
+
+import threading
+import queue
+
+import time
+
+# load model weights
+audio_processor,vae,unet,pe = load_all_model()
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+timesteps = torch.tensor([0], device=device)
+
+
+def video2imgs(vid_path, save_path, ext = '.png',cut_frame = 10000000):
+ cap = cv2.VideoCapture(vid_path)
+ count = 0
+ while True:
+ if count > cut_frame:
+ break
+ ret, frame = cap.read()
+ if ret:
+ cv2.imwrite(f"{save_path}/{count:08d}.png", frame)
+ count += 1
+ else:
+ break
+
+def osmakedirs(path_list):
+ for path in path_list:
+ os.makedirs(path) if not os.path.exists(path) else None
+
+
+@torch.no_grad()
+class Avatar:
+ def __init__(self, avatar_id, video_path, bbox_shift, batch_size, preparation):
+ self.avatar_id = avatar_id
+ self.video_path = video_path
+ self.bbox_shift = bbox_shift
+ self.avatar_path = f"./results/avatars/{avatar_id}"
+ self.full_imgs_path = f"{self.avatar_path}/full_imgs"
+ self.coords_path = f"{self.avatar_path}/coords.pkl"
+ self.latents_out_path= f"{self.avatar_path}/latents.pt"
+ self.video_out_path = f"{self.avatar_path}/vid_output/"
+ self.mask_out_path =f"{self.avatar_path}/mask"
+ self.mask_coords_path =f"{self.avatar_path}/mask_coords.pkl"
+ self.avatar_info_path = f"{self.avatar_path}/avator_info.json"
+ self.avatar_info = {
+ "avatar_id":avatar_id,
+ "video_path":video_path,
+ "bbox_shift":bbox_shift
+ }
+ self.preparation = preparation
+ self.batch_size = batch_size
+ self.idx = 0
+ self.init()
+
+ def init(self):
+ if self.preparation:
+ if os.path.exists(self.avatar_path):
+ response = input(f"{self.avatar_id} exists, Do you want to re-create it ? (y/n)")
+ if response.lower() == "y":
+ shutil.rmtree(self.avatar_path)
+ print("*********************************")
+ print(f" creating avator: {self.avatar_id}")
+ print("*********************************")
+ osmakedirs([self.avatar_path,self.full_imgs_path,self.video_out_path,self.mask_out_path])
+ self.prepare_material()
+ else:
+ self.input_latent_list_cycle = torch.load(self.latents_out_path)
+ with open(self.coords_path, 'rb') as f:
+ self.coord_list_cycle = pickle.load(f)
+ input_img_list = glob.glob(os.path.join(self.full_imgs_path, '*.[jpJP][pnPN]*[gG]'))
+ input_img_list = sorted(input_img_list, key=lambda x: int(os.path.splitext(os.path.basename(x))[0]))
+ self.frame_list_cycle = read_imgs(input_img_list)
+ with open(self.mask_coords_path, 'rb') as f:
+ self.mask_coords_list_cycle = pickle.load(f)
+ input_mask_list = glob.glob(os.path.join(self.mask_out_path, '*.[jpJP][pnPN]*[gG]'))
+ input_mask_list = sorted(input_mask_list, key=lambda x: int(os.path.splitext(os.path.basename(x))[0]))
+ self.mask_list_cycle = read_imgs(input_mask_list)
+ else:
+ print("*********************************")
+ print(f" creating avator: {self.avatar_id}")
+ print("*********************************")
+ osmakedirs([self.avatar_path,self.full_imgs_path,self.video_out_path,self.mask_out_path])
+ self.prepare_material()
+ else:
+ with open(self.avatar_info_path, "r") as f:
+ avatar_info = json.load(f)
+
+ if avatar_info['bbox_shift'] != self.avatar_info['bbox_shift']:
+ response = input(f" 【bbox_shift】 is changed, you need to re-create it ! (c/continue)")
+ if response.lower() == "c":
+ shutil.rmtree(self.avatar_path)
+ print("*********************************")
+ print(f" creating avator: {self.avatar_id}")
+ print("*********************************")
+ osmakedirs([self.avatar_path,self.full_imgs_path,self.video_out_path,self.mask_out_path])
+ self.prepare_material()
+ else:
+ sys.exit()
+ else:
+ self.input_latent_list_cycle = torch.load(self.latents_out_path)
+ with open(self.coords_path, 'rb') as f:
+ self.coord_list_cycle = pickle.load(f)
+ input_img_list = glob.glob(os.path.join(self.full_imgs_path, '*.[jpJP][pnPN]*[gG]'))
+ input_img_list = sorted(input_img_list, key=lambda x: int(os.path.splitext(os.path.basename(x))[0]))
+ self.frame_list_cycle = read_imgs(input_img_list)
+ with open(self.mask_coords_path, 'rb') as f:
+ self.mask_coords_list_cycle = pickle.load(f)
+ input_mask_list = glob.glob(os.path.join(self.mask_out_path, '*.[jpJP][pnPN]*[gG]'))
+ input_mask_list = sorted(input_mask_list, key=lambda x: int(os.path.splitext(os.path.basename(x))[0]))
+ self.mask_list_cycle = read_imgs(input_mask_list)
+
+ def prepare_material(self):
+ print("preparing data materials ... ...")
+ with open(self.avatar_info_path, "w") as f:
+ json.dump(self.avatar_info, f)
+
+ if os.path.isfile(self.video_path):
+ video2imgs(self.video_path, self.full_imgs_path, ext = 'png')
+ else:
+ print(f"copy files in {self.video_path}")
+ files = os.listdir(self.video_path)
+ files.sort()
+ files = [file for file in files if file.split(".")[-1]=="png"]
+ for filename in files:
+ shutil.copyfile(f"{self.video_path}/{filename}", f"{self.full_imgs_path}/{filename}")
+ input_img_list = sorted(glob.glob(os.path.join(self.full_imgs_path, '*.[jpJP][pnPN]*[gG]')))
+
+ print("extracting landmarks...")
+ coord_list, frame_list = get_landmark_and_bbox(input_img_list, self.bbox_shift)
+ input_latent_list = []
+ idx = -1
+ # maker if the bbox is not sufficient
+ coord_placeholder = (0.0,0.0,0.0,0.0)
+ for bbox, frame in zip(coord_list, frame_list):
+ idx = idx + 1
+ if bbox == coord_placeholder:
+ continue
+ x1, y1, x2, y2 = bbox
+ crop_frame = frame[y1:y2, x1:x2]
+ resized_crop_frame = cv2.resize(crop_frame,(256,256),interpolation = cv2.INTER_LANCZOS4)
+ latents = vae.get_latents_for_unet(resized_crop_frame)
+ input_latent_list.append(latents)
+
+ self.frame_list_cycle = frame_list + frame_list[::-1]
+ self.coord_list_cycle = coord_list + coord_list[::-1]
+ self.input_latent_list_cycle = input_latent_list + input_latent_list[::-1]
+ self.mask_coords_list_cycle = []
+ self.mask_list_cycle = []
+
+ for i,frame in enumerate(tqdm(self.frame_list_cycle)):
+ cv2.imwrite(f"{self.full_imgs_path}/{str(i).zfill(8)}.png",frame)
+
+ face_box = self.coord_list_cycle[i]
+ mask,crop_box = get_image_prepare_material(frame,face_box)
+ cv2.imwrite(f"{self.mask_out_path}/{str(i).zfill(8)}.png",mask)
+ self.mask_coords_list_cycle += [crop_box]
+ self.mask_list_cycle.append(mask)
+
+ with open(self.mask_coords_path, 'wb') as f:
+ pickle.dump(self.mask_coords_list_cycle, f)
+
+ with open(self.coords_path, 'wb') as f:
+ pickle.dump(self.coord_list_cycle, f)
+
+ torch.save(self.input_latent_list_cycle, os.path.join(self.latents_out_path))
+ #
+
+ def process_frames(self, res_frame_queue,video_len):
+ print(video_len)
+ while True:
+ if self.idx>=video_len-1:
+ break
+ try:
+ start = time.time()
+ res_frame = res_frame_queue.get(block=True, timeout=1)
+ except queue.Empty:
+ continue
+
+ bbox = self.coord_list_cycle[self.idx%(len(self.coord_list_cycle))]
+ ori_frame = copy.deepcopy(self.frame_list_cycle[self.idx%(len(self.frame_list_cycle))])
+ x1, y1, x2, y2 = bbox
+ try:
+ res_frame = cv2.resize(res_frame.astype(np.uint8),(x2-x1,y2-y1))
+ except:
+ continue
+ mask = self.mask_list_cycle[self.idx%(len(self.mask_list_cycle))]
+ mask_crop_box = self.mask_coords_list_cycle[self.idx%(len(self.mask_coords_list_cycle))]
+ #combine_frame = get_image(ori_frame,res_frame,bbox)
+ combine_frame = get_image_blending(ori_frame,res_frame,bbox,mask,mask_crop_box)
+
+ fps = 1/(time.time()-start)
+ print(f"Generating the {self.idx}-th frame with FPS: {fps:.2f}")
+ cv2.imwrite(f"{self.avatar_path}/tmp/{str(self.idx).zfill(8)}.png",combine_frame)
+ self.idx = self.idx + 1
+
+ def inference(self, audio_path, out_vid_name, fps):
+ os.makedirs(self.avatar_path+'/tmp',exist_ok =True)
+ ############################################## extract audio feature ##############################################
+ whisper_feature = audio_processor.audio2feat(audio_path)
+ whisper_chunks = audio_processor.feature2chunks(feature_array=whisper_feature,fps=fps)
+ ############################################## inference batch by batch ##############################################
+ video_num = len(whisper_chunks)
+ print("start inference")
+ res_frame_queue = queue.Queue()
+ self.idx = 0
+ # # Create a sub-thread and start it
+ process_thread = threading.Thread(target=self.process_frames, args=(res_frame_queue,video_num))
+ process_thread.start()
+ start_time = time.time()
+ gen = datagen(whisper_chunks,self.input_latent_list_cycle, self.batch_size)
+ print(f"processing audio:{audio_path} costs {(time.time() - start_time) * 1000}ms")
+ start_time = time.time()
+ res_frame_list = []
+
+ for i, (whisper_batch,latent_batch) in enumerate(tqdm(gen,total=int(np.ceil(float(video_num)/self.batch_size)))):
+ start_time = time.time()
+ tensor_list = [torch.FloatTensor(arr) for arr in whisper_batch]
+ audio_feature_batch = torch.stack(tensor_list).to(unet.device) # torch, B, 5*N,384
+ audio_feature_batch = pe(audio_feature_batch)
+
+ pred_latents = unet.model(latent_batch, timesteps, encoder_hidden_states=audio_feature_batch).sample
+ recon = vae.decode_latents(pred_latents)
+ for res_frame in recon:
+ res_frame_queue.put(res_frame)
+ # Close the queue and sub-thread after all tasks are completed
+ process_thread.join()
+
+ if out_vid_name is not None:
+ # optional
+ cmd_img2video = f"ffmpeg -y -v fatal -r {fps} -f image2 -i {self.avatar_path}/tmp/%08d.png -vcodec libx264 -vf format=rgb24,scale=out_color_matrix=bt709,format=yuv420p -crf 18 {self.avatar_path}/temp.mp4"
+ print(cmd_img2video)
+ os.system(cmd_img2video)
+
+ output_vid = os.path.join(self.video_out_path, out_vid_name+".mp4") # on
+ cmd_combine_audio = f"ffmpeg -y -v fatal -i {audio_path} -i {self.avatar_path}/temp.mp4 {output_vid}"
+ print(cmd_combine_audio)
+ os.system(cmd_combine_audio)
+
+ os.remove(f"{self.avatar_path}/temp.mp4")
+ shutil.rmtree(f"{self.avatar_path}/tmp")
+ print(f"result is save to {output_vid}")
+
+
+
+
+
+if __name__ == "__main__":
+ '''
+ This script is used to simulate online chatting and applies necessary pre-processing such as face detection and face parsing in advance. During online chatting, only UNet and the VAE decoder are involved, which makes MuseTalk real-time.
+ '''
+
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--inference_config", type=str, default="configs/inference/realtime.yaml")
+ parser.add_argument("--fps", type=int, default=25)
+ parser.add_argument("--batch_size", type=int, default=4)
+
+ args = parser.parse_args()
+
+ inference_config = OmegaConf.load(args.inference_config)
+ print(inference_config)
+
+
+ for avatar_id in inference_config:
+ data_preparation = inference_config[avatar_id]["preparation"]
+ video_path = inference_config[avatar_id]["video_path"]
+ bbox_shift = inference_config[avatar_id]["bbox_shift"]
+ avatar = Avatar(
+ avatar_id = avatar_id,
+ video_path = video_path,
+ bbox_shift = bbox_shift,
+ batch_size = args.batch_size,
+ preparation= data_preparation)
+
+ audio_clips = inference_config[avatar_id]["audio_clips"]
+ for audio_num, audio_path in audio_clips.items():
+ print("Inferring using:",audio_path)
+ avatar.inference(audio_path, audio_num, args.fps)
+
\ No newline at end of file