mirror of
https://github.com/TMElyralab/MuseTalk.git
synced 2026-02-04 09:29:20 +08:00
feat: real-time infer (#286)
* feat: realtime infer * cchore: infer script
This commit is contained in:
Zhizhou Zhong
GitHub
@@ -1,8 +1,9 @@
|
||||
import os
|
||||
import cv2
|
||||
import math
|
||||
import copy
|
||||
import glob
|
||||
import torch
|
||||
import glob
|
||||
import shutil
|
||||
import pickle
|
||||
import argparse
|
||||
@@ -17,18 +18,16 @@ from musetalk.utils.audio_processor import AudioProcessor
|
||||
from musetalk.utils.utils import get_file_type, get_video_fps, datagen, load_all_model
|
||||
from musetalk.utils.preprocessing import get_landmark_and_bbox, read_imgs, coord_placeholder
|
||||
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def main(args):
|
||||
# Configure ffmpeg path
|
||||
if args.ffmpeg_path not in os.getenv('PATH'):
|
||||
print("Adding ffmpeg to PATH")
|
||||
os.environ["PATH"] = f"{args.ffmpeg_path}:{os.environ['PATH']}"
|
||||
|
||||
|
||||
# Set computing device
|
||||
device = torch.device(f"cuda:{args.gpu_id}" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
|
||||
# Load model weights
|
||||
vae, unet, pe = load_all_model(
|
||||
unet_model_path=args.unet_model_path,
|
||||
@@ -37,164 +36,229 @@ def main(args):
|
||||
device=device
|
||||
)
|
||||
timesteps = torch.tensor([0], device=device)
|
||||
|
||||
|
||||
if args.use_float16 is True:
|
||||
|
||||
# Convert models to half precision if float16 is enabled
|
||||
if args.use_float16:
|
||||
pe = pe.half()
|
||||
vae.vae = vae.vae.half()
|
||||
unet.model = unet.model.half()
|
||||
|
||||
# Move models to specified device
|
||||
pe = pe.to(device)
|
||||
vae.vae = vae.vae.to(device)
|
||||
unet.model = unet.model.to(device)
|
||||
|
||||
# Initialize audio processor and Whisper model
|
||||
# Initialize audio processor and Whisper model
|
||||
audio_processor = AudioProcessor(feature_extractor_path=args.whisper_dir)
|
||||
weight_dtype = unet.model.dtype
|
||||
whisper = WhisperModel.from_pretrained(args.whisper_dir)
|
||||
whisper = whisper.to(device=device, dtype=weight_dtype).eval()
|
||||
whisper.requires_grad_(False)
|
||||
|
||||
# Initialize face parser
|
||||
fp = FaceParsing()
|
||||
|
||||
inference_config = OmegaConf.load(args.inference_config)
|
||||
print(inference_config)
|
||||
for task_id in inference_config:
|
||||
video_path = inference_config[task_id]["video_path"]
|
||||
audio_path = inference_config[task_id]["audio_path"]
|
||||
bbox_shift = inference_config[task_id].get("bbox_shift", args.bbox_shift)
|
||||
|
||||
input_basename = os.path.basename(video_path).split('.')[0]
|
||||
audio_basename = os.path.basename(audio_path).split('.')[0]
|
||||
output_basename = f"{input_basename}_{audio_basename}"
|
||||
result_img_save_path = os.path.join(args.result_dir, output_basename) # related to video & audio inputs
|
||||
crop_coord_save_path = os.path.join(result_img_save_path, input_basename+".pkl") # only related to video input
|
||||
os.makedirs(result_img_save_path,exist_ok =True)
|
||||
|
||||
if args.output_vid_name is None:
|
||||
output_vid_name = os.path.join(args.result_dir, output_basename+".mp4")
|
||||
else:
|
||||
output_vid_name = os.path.join(args.result_dir, args.output_vid_name)
|
||||
############################################## extract frames from source video ##############################################
|
||||
if get_file_type(video_path)=="video":
|
||||
save_dir_full = os.path.join(args.result_dir, input_basename)
|
||||
os.makedirs(save_dir_full,exist_ok = True)
|
||||
cmd = f"ffmpeg -v fatal -i {video_path} -start_number 0 {save_dir_full}/%08d.png"
|
||||
os.system(cmd)
|
||||
input_img_list = sorted(glob.glob(os.path.join(save_dir_full, '*.[jpJP][pnPN]*[gG]')))
|
||||
fps = get_video_fps(video_path)
|
||||
elif get_file_type(video_path)=="image":
|
||||
input_img_list = [video_path, ]
|
||||
fps = args.fps
|
||||
elif os.path.isdir(video_path): # input img folder
|
||||
input_img_list = glob.glob(os.path.join(video_path, '*.[jpJP][pnPN]*[gG]'))
|
||||
input_img_list = sorted(input_img_list, key=lambda x: int(os.path.splitext(os.path.basename(x))[0]))
|
||||
fps = args.fps
|
||||
else:
|
||||
raise ValueError(f"{video_path} should be a video file, an image file or a directory of images")
|
||||
|
||||
############################################## extract audio feature ##############################################
|
||||
# Extract audio features
|
||||
whisper_input_features, librosa_length = audio_processor.get_audio_feature(audio_path)
|
||||
whisper_chunks = audio_processor.get_whisper_chunk(
|
||||
whisper_input_features,
|
||||
device,
|
||||
weight_dtype,
|
||||
whisper,
|
||||
librosa_length,
|
||||
fps=fps,
|
||||
audio_padding_length_left=args.audio_padding_length_left,
|
||||
audio_padding_length_right=args.audio_padding_length_right,
|
||||
# Initialize face parser with configurable parameters based on version
|
||||
if args.version == "v15":
|
||||
fp = FaceParsing(
|
||||
left_cheek_width=args.left_cheek_width,
|
||||
right_cheek_width=args.right_cheek_width
|
||||
)
|
||||
|
||||
############################################## preprocess input image ##############################################
|
||||
if os.path.exists(crop_coord_save_path) and args.use_saved_coord:
|
||||
print("using extracted coordinates")
|
||||
with open(crop_coord_save_path,'rb') as f:
|
||||
coord_list = pickle.load(f)
|
||||
frame_list = read_imgs(input_img_list)
|
||||
else:
|
||||
print("extracting landmarks...time consuming")
|
||||
coord_list, frame_list = get_landmark_and_bbox(input_img_list, bbox_shift)
|
||||
with open(crop_coord_save_path, 'wb') as f:
|
||||
pickle.dump(coord_list, f)
|
||||
|
||||
i = 0
|
||||
input_latent_list = []
|
||||
for bbox, frame in zip(coord_list, frame_list):
|
||||
if bbox == coord_placeholder:
|
||||
continue
|
||||
x1, y1, x2, y2 = bbox
|
||||
crop_frame = frame[y1:y2, x1:x2]
|
||||
crop_frame = cv2.resize(crop_frame,(256,256),interpolation = cv2.INTER_LANCZOS4)
|
||||
latents = vae.get_latents_for_unet(crop_frame)
|
||||
input_latent_list.append(latents)
|
||||
else: # v1
|
||||
fp = FaceParsing()
|
||||
|
||||
# to smooth the first and the last frame
|
||||
frame_list_cycle = frame_list + frame_list[::-1]
|
||||
coord_list_cycle = coord_list + coord_list[::-1]
|
||||
input_latent_list_cycle = input_latent_list + input_latent_list[::-1]
|
||||
############################################## inference batch by batch ##############################################
|
||||
print("start inference")
|
||||
video_num = len(whisper_chunks)
|
||||
batch_size = args.batch_size
|
||||
gen = datagen(whisper_chunks,input_latent_list_cycle,batch_size)
|
||||
res_frame_list = []
|
||||
for i, (whisper_batch,latent_batch) in enumerate(tqdm(gen,total=int(np.ceil(float(video_num)/batch_size)))):
|
||||
audio_feature_batch = pe(whisper_batch)
|
||||
latent_batch = latent_batch.to(dtype=unet.model.dtype)
|
||||
# Load inference configuration
|
||||
inference_config = OmegaConf.load(args.inference_config)
|
||||
print("Loaded inference config:", inference_config)
|
||||
|
||||
# Process each task
|
||||
for task_id in inference_config:
|
||||
try:
|
||||
# Get task configuration
|
||||
video_path = inference_config[task_id]["video_path"]
|
||||
audio_path = inference_config[task_id]["audio_path"]
|
||||
if "result_name" in inference_config[task_id]:
|
||||
args.output_vid_name = inference_config[task_id]["result_name"]
|
||||
|
||||
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_list.append(res_frame)
|
||||
|
||||
############################################## pad to full image ##############################################
|
||||
print("pad talking image to original video")
|
||||
for i, res_frame in enumerate(tqdm(res_frame_list)):
|
||||
bbox = coord_list_cycle[i%(len(coord_list_cycle))]
|
||||
ori_frame = copy.deepcopy(frame_list_cycle[i%(len(frame_list_cycle))])
|
||||
x1, y1, x2, y2 = bbox
|
||||
try:
|
||||
res_frame = cv2.resize(res_frame.astype(np.uint8),(x2-x1,y2-y1))
|
||||
except:
|
||||
continue
|
||||
# Set bbox_shift based on version
|
||||
if args.version == "v15":
|
||||
bbox_shift = 0 # v15 uses fixed bbox_shift
|
||||
else:
|
||||
bbox_shift = inference_config[task_id].get("bbox_shift", args.bbox_shift) # v1 uses config or default
|
||||
|
||||
# Merge results
|
||||
combine_frame = get_image(ori_frame, res_frame, [x1, y1, x2, y2], fp=fp)
|
||||
cv2.imwrite(f"{result_img_save_path}/{str(i).zfill(8)}.png",combine_frame)
|
||||
# Set output paths
|
||||
input_basename = os.path.basename(video_path).split('.')[0]
|
||||
audio_basename = os.path.basename(audio_path).split('.')[0]
|
||||
output_basename = f"{input_basename}_{audio_basename}"
|
||||
|
||||
# Create temporary directories
|
||||
temp_dir = os.path.join(args.result_dir, f"{args.version}")
|
||||
os.makedirs(temp_dir, exist_ok=True)
|
||||
|
||||
# Set result save paths
|
||||
result_img_save_path = os.path.join(temp_dir, output_basename)
|
||||
crop_coord_save_path = os.path.join(args.result_dir, "../", input_basename+".pkl")
|
||||
os.makedirs(result_img_save_path, exist_ok=True)
|
||||
|
||||
# Set output video paths
|
||||
if args.output_vid_name is None:
|
||||
output_vid_name = os.path.join(temp_dir, output_basename + ".mp4")
|
||||
else:
|
||||
output_vid_name = os.path.join(temp_dir, args.output_vid_name)
|
||||
output_vid_name_concat = os.path.join(temp_dir, output_basename + "_concat.mp4")
|
||||
|
||||
# Extract frames from source video
|
||||
if get_file_type(video_path) == "video":
|
||||
save_dir_full = os.path.join(temp_dir, input_basename)
|
||||
os.makedirs(save_dir_full, exist_ok=True)
|
||||
cmd = f"ffmpeg -v fatal -i {video_path} -start_number 0 {save_dir_full}/%08d.png"
|
||||
os.system(cmd)
|
||||
input_img_list = sorted(glob.glob(os.path.join(save_dir_full, '*.[jpJP][pnPN]*[gG]')))
|
||||
fps = get_video_fps(video_path)
|
||||
elif get_file_type(video_path) == "image":
|
||||
input_img_list = [video_path]
|
||||
fps = args.fps
|
||||
elif os.path.isdir(video_path):
|
||||
input_img_list = glob.glob(os.path.join(video_path, '*.[jpJP][pnPN]*[gG]'))
|
||||
input_img_list = sorted(input_img_list, key=lambda x: int(os.path.splitext(os.path.basename(x))[0]))
|
||||
fps = args.fps
|
||||
else:
|
||||
raise ValueError(f"{video_path} should be a video file, an image file or a directory of images")
|
||||
|
||||
cmd_img2video = f"ffmpeg -y -v warning -r {fps} -f image2 -i {result_img_save_path}/%08d.png -vcodec libx264 -vf format=rgb24,scale=out_color_matrix=bt709,format=yuv420p -crf 18 temp.mp4"
|
||||
print(cmd_img2video)
|
||||
os.system(cmd_img2video)
|
||||
# Extract audio features
|
||||
whisper_input_features, librosa_length = audio_processor.get_audio_feature(audio_path)
|
||||
whisper_chunks = audio_processor.get_whisper_chunk(
|
||||
whisper_input_features,
|
||||
device,
|
||||
weight_dtype,
|
||||
whisper,
|
||||
librosa_length,
|
||||
fps=fps,
|
||||
audio_padding_length_left=args.audio_padding_length_left,
|
||||
audio_padding_length_right=args.audio_padding_length_right,
|
||||
)
|
||||
|
||||
# Preprocess input images
|
||||
if os.path.exists(crop_coord_save_path) and args.use_saved_coord:
|
||||
print("Using saved coordinates")
|
||||
with open(crop_coord_save_path, 'rb') as f:
|
||||
coord_list = pickle.load(f)
|
||||
frame_list = read_imgs(input_img_list)
|
||||
else:
|
||||
print("Extracting landmarks... time-consuming operation")
|
||||
coord_list, frame_list = get_landmark_and_bbox(input_img_list, bbox_shift)
|
||||
with open(crop_coord_save_path, 'wb') as f:
|
||||
pickle.dump(coord_list, f)
|
||||
|
||||
print(f"Number of frames: {len(frame_list)}")
|
||||
|
||||
# Process each frame
|
||||
input_latent_list = []
|
||||
for bbox, frame in zip(coord_list, frame_list):
|
||||
if bbox == coord_placeholder:
|
||||
continue
|
||||
x1, y1, x2, y2 = bbox
|
||||
if args.version == "v15":
|
||||
y2 = y2 + args.extra_margin
|
||||
y2 = min(y2, frame.shape[0])
|
||||
crop_frame = frame[y1:y2, x1:x2]
|
||||
crop_frame = cv2.resize(crop_frame, (256,256), interpolation=cv2.INTER_LANCZOS4)
|
||||
latents = vae.get_latents_for_unet(crop_frame)
|
||||
input_latent_list.append(latents)
|
||||
|
||||
cmd_combine_audio = f"ffmpeg -y -v warning -i {audio_path} -i temp.mp4 {output_vid_name}"
|
||||
print(cmd_combine_audio)
|
||||
os.system(cmd_combine_audio)
|
||||
|
||||
os.remove("temp.mp4")
|
||||
shutil.rmtree(result_img_save_path)
|
||||
print(f"result is save to {output_vid_name}")
|
||||
# Smooth first and last frames
|
||||
frame_list_cycle = frame_list + frame_list[::-1]
|
||||
coord_list_cycle = coord_list + coord_list[::-1]
|
||||
input_latent_list_cycle = input_latent_list + input_latent_list[::-1]
|
||||
|
||||
# Batch inference
|
||||
print("Starting inference")
|
||||
video_num = len(whisper_chunks)
|
||||
batch_size = args.batch_size
|
||||
gen = datagen(
|
||||
whisper_chunks=whisper_chunks,
|
||||
vae_encode_latents=input_latent_list_cycle,
|
||||
batch_size=batch_size,
|
||||
delay_frame=0,
|
||||
device=device,
|
||||
)
|
||||
|
||||
res_frame_list = []
|
||||
total = int(np.ceil(float(video_num) / batch_size))
|
||||
|
||||
# Execute inference
|
||||
for i, (whisper_batch, latent_batch) in enumerate(tqdm(gen, total=total)):
|
||||
audio_feature_batch = pe(whisper_batch)
|
||||
latent_batch = latent_batch.to(dtype=unet.model.dtype)
|
||||
|
||||
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_list.append(res_frame)
|
||||
|
||||
# Pad generated images to original video size
|
||||
print("Padding generated images to original video size")
|
||||
for i, res_frame in enumerate(tqdm(res_frame_list)):
|
||||
bbox = coord_list_cycle[i%(len(coord_list_cycle))]
|
||||
ori_frame = copy.deepcopy(frame_list_cycle[i%(len(frame_list_cycle))])
|
||||
x1, y1, x2, y2 = bbox
|
||||
if args.version == "v15":
|
||||
y2 = y2 + args.extra_margin
|
||||
y2 = min(y2, frame.shape[0])
|
||||
try:
|
||||
res_frame = cv2.resize(res_frame.astype(np.uint8), (x2-x1, y2-y1))
|
||||
except:
|
||||
continue
|
||||
|
||||
# Merge results with version-specific parameters
|
||||
if args.version == "v15":
|
||||
combine_frame = get_image(ori_frame, res_frame, [x1, y1, x2, y2], mode=args.parsing_mode, fp=fp)
|
||||
else:
|
||||
combine_frame = get_image(ori_frame, res_frame, [x1, y1, x2, y2], fp=fp)
|
||||
cv2.imwrite(f"{result_img_save_path}/{str(i).zfill(8)}.png", combine_frame)
|
||||
|
||||
# Save prediction results
|
||||
temp_vid_path = f"{temp_dir}/temp_{input_basename}_{audio_basename}.mp4"
|
||||
cmd_img2video = f"ffmpeg -y -v warning -r {fps} -f image2 -i {result_img_save_path}/%08d.png -vcodec libx264 -vf format=yuv420p -crf 18 {temp_vid_path}"
|
||||
print("Video generation command:", cmd_img2video)
|
||||
os.system(cmd_img2video)
|
||||
|
||||
cmd_combine_audio = f"ffmpeg -y -v warning -i {audio_path} -i {temp_vid_path} {output_vid_name}"
|
||||
print("Audio combination command:", cmd_combine_audio)
|
||||
os.system(cmd_combine_audio)
|
||||
|
||||
# Clean up temporary files
|
||||
shutil.rmtree(result_img_save_path)
|
||||
os.remove(temp_vid_path)
|
||||
|
||||
shutil.rmtree(save_dir_full)
|
||||
if not args.saved_coord:
|
||||
os.remove(crop_coord_save_path)
|
||||
|
||||
print(f"Results saved to {output_vid_name}")
|
||||
except Exception as e:
|
||||
print("Error occurred during processing:", e)
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--ffmpeg_path", type=str, default="./ffmpeg-4.4-amd64-static/", help="Path to ffmpeg executable")
|
||||
parser.add_argument("--inference_config", type=str, default="configs/inference/test_img.yaml")
|
||||
parser.add_argument("--bbox_shift", type=int, default=0)
|
||||
parser.add_argument("--result_dir", default='./results', help="path to output")
|
||||
parser.add_argument("--gpu_id", type=int, default=0, help="GPU ID to use")
|
||||
parser.add_argument("--batch_size", type=int, default=8)
|
||||
parser.add_argument("--output_vid_name", type=str, default=None)
|
||||
parser.add_argument("--use_saved_coord",
|
||||
action="store_true",
|
||||
help='use saved coordinate to save time')
|
||||
parser.add_argument("--use_float16",
|
||||
action="store_true",
|
||||
help="Whether use float16 to speed up inference",
|
||||
)
|
||||
parser.add_argument("--fps", type=int, default=25, help="Video frames per second")
|
||||
parser.add_argument("--unet_model_path", type=str, default="./models/musetalk/pytorch_model.bin", help="Path to UNet model weights")
|
||||
parser.add_argument("--vae_type", type=str, default="sd-vae", help="Type of VAE model")
|
||||
parser.add_argument("--unet_config", type=str, default="./models/musetalk/config.json", help="Path to UNet configuration file")
|
||||
parser.add_argument("--unet_model_path", type=str, default="./models/musetalkV15/unet.pth", help="Path to UNet model weights")
|
||||
parser.add_argument("--whisper_dir", type=str, default="./models/whisper", help="Directory containing Whisper model")
|
||||
parser.add_argument("--inference_config", type=str, default="configs/inference/test_img.yaml", help="Path to inference configuration file")
|
||||
parser.add_argument("--bbox_shift", type=int, default=0, help="Bounding box shift value")
|
||||
parser.add_argument("--result_dir", default='./results', help="Directory for output results")
|
||||
parser.add_argument("--extra_margin", type=int, default=10, help="Extra margin for face cropping")
|
||||
parser.add_argument("--fps", type=int, default=25, help="Video frames per second")
|
||||
parser.add_argument("--audio_padding_length_left", type=int, default=2, help="Left padding length for audio")
|
||||
parser.add_argument("--audio_padding_length_right", type=int, default=2, help="Right padding length for audio")
|
||||
parser.add_argument("--batch_size", type=int, default=8, help="Batch size for inference")
|
||||
parser.add_argument("--output_vid_name", type=str, default=None, help="Name of output video file")
|
||||
parser.add_argument("--use_saved_coord", action="store_true", help='Use saved coordinates to save time')
|
||||
parser.add_argument("--saved_coord", action="store_true", help='Save coordinates for future use')
|
||||
parser.add_argument("--use_float16", action="store_true", help="Use float16 for faster inference")
|
||||
parser.add_argument("--parsing_mode", default='jaw', help="Face blending parsing mode")
|
||||
parser.add_argument("--left_cheek_width", type=int, default=90, help="Width of left cheek region")
|
||||
parser.add_argument("--right_cheek_width", type=int, default=90, help="Width of right cheek region")
|
||||
parser.add_argument("--version", type=str, default="v15", choices=["v1", "v15"], help="Model version to use")
|
||||
args = parser.parse_args()
|
||||
main(args)
|
||||
|
||||
@@ -1,252 +0,0 @@
|
||||
import os
|
||||
import cv2
|
||||
import math
|
||||
import copy
|
||||
import torch
|
||||
import glob
|
||||
import shutil
|
||||
import pickle
|
||||
import argparse
|
||||
import subprocess
|
||||
import numpy as np
|
||||
from tqdm import tqdm
|
||||
from omegaconf import OmegaConf
|
||||
from transformers import WhisperModel
|
||||
|
||||
from musetalk.utils.blending import get_image
|
||||
from musetalk.utils.face_parsing import FaceParsing
|
||||
from musetalk.utils.audio_processor import AudioProcessor
|
||||
from musetalk.utils.utils import get_file_type, get_video_fps, datagen, load_all_model
|
||||
from musetalk.utils.preprocessing import get_landmark_and_bbox, read_imgs, coord_placeholder
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def main(args):
|
||||
# Configure ffmpeg path
|
||||
if args.ffmpeg_path not in os.getenv('PATH'):
|
||||
print("Adding ffmpeg to PATH")
|
||||
os.environ["PATH"] = f"{args.ffmpeg_path}:{os.environ['PATH']}"
|
||||
|
||||
# Set computing device
|
||||
device = torch.device(f"cuda:{args.gpu_id}" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
# Load model weights
|
||||
vae, unet, pe = load_all_model(
|
||||
unet_model_path=args.unet_model_path,
|
||||
vae_type=args.vae_type,
|
||||
unet_config=args.unet_config,
|
||||
device=device
|
||||
)
|
||||
timesteps = torch.tensor([0], device=device)
|
||||
|
||||
# Convert models to half precision if float16 is enabled
|
||||
if args.use_float16:
|
||||
pe = pe.half()
|
||||
vae.vae = vae.vae.half()
|
||||
unet.model = unet.model.half()
|
||||
|
||||
# Move models to specified device
|
||||
pe = pe.to(device)
|
||||
vae.vae = vae.vae.to(device)
|
||||
unet.model = unet.model.to(device)
|
||||
|
||||
# Initialize audio processor and Whisper model
|
||||
audio_processor = AudioProcessor(feature_extractor_path=args.whisper_dir)
|
||||
weight_dtype = unet.model.dtype
|
||||
whisper = WhisperModel.from_pretrained(args.whisper_dir)
|
||||
whisper = whisper.to(device=device, dtype=weight_dtype).eval()
|
||||
whisper.requires_grad_(False)
|
||||
|
||||
# Initialize face parser
|
||||
fp = FaceParsing(left_cheek_width=args.left_cheek_width, right_cheek_width=args.right_cheek_width)
|
||||
|
||||
# Load inference configuration
|
||||
inference_config = OmegaConf.load(args.inference_config)
|
||||
print("Loaded inference config:", inference_config)
|
||||
|
||||
# Process each task
|
||||
for task_id in inference_config:
|
||||
try:
|
||||
# Get task configuration
|
||||
video_path = inference_config[task_id]["video_path"]
|
||||
audio_path = inference_config[task_id]["audio_path"]
|
||||
if "result_name" in inference_config[task_id]:
|
||||
args.output_vid_name = inference_config[task_id]["result_name"]
|
||||
bbox_shift = args.bbox_shift
|
||||
# Set output paths
|
||||
input_basename = os.path.basename(video_path).split('.')[0]
|
||||
audio_basename = os.path.basename(audio_path).split('.')[0]
|
||||
output_basename = f"{input_basename}_{audio_basename}"
|
||||
|
||||
# Create temporary directories
|
||||
temp_dir = os.path.join(args.result_dir, "frames_result")
|
||||
os.makedirs(temp_dir, exist_ok=True)
|
||||
|
||||
# Set result save paths
|
||||
result_img_save_path = os.path.join(temp_dir, output_basename) # related to video & audio inputs
|
||||
crop_coord_save_path = os.path.join(args.result_dir, "../", input_basename+".pkl") # only related to video input
|
||||
os.makedirs(result_img_save_path, exist_ok=True)
|
||||
# Set output video paths
|
||||
if args.output_vid_name is None:
|
||||
output_vid_name = os.path.join(temp_dir, output_basename + ".mp4")
|
||||
else:
|
||||
output_vid_name = os.path.join(temp_dir, args.output_vid_name)
|
||||
output_vid_name_concat = os.path.join(temp_dir, output_basename + "_concat.mp4")
|
||||
|
||||
# Skip if output file already exists
|
||||
if os.path.exists(output_vid_name):
|
||||
print(f"{output_vid_name} already exists, skipping!")
|
||||
continue
|
||||
|
||||
# Extract frames from source video
|
||||
if get_file_type(video_path) == "video":
|
||||
save_dir_full = os.path.join(temp_dir, input_basename)
|
||||
os.makedirs(save_dir_full, exist_ok=True)
|
||||
cmd = f"ffmpeg -v fatal -i {video_path} -start_number 0 {save_dir_full}/%08d.png"
|
||||
os.system(cmd)
|
||||
input_img_list = sorted(glob.glob(os.path.join(save_dir_full, '*.[jpJP][pnPN]*[gG]')))
|
||||
fps = get_video_fps(video_path)
|
||||
elif get_file_type(video_path) == "image":
|
||||
input_img_list = [video_path]
|
||||
fps = args.fps
|
||||
elif os.path.isdir(video_path):
|
||||
input_img_list = glob.glob(os.path.join(video_path, '*.[jpJP][pnPN]*[gG]'))
|
||||
input_img_list = sorted(input_img_list, key=lambda x: int(os.path.splitext(os.path.basename(x))[0]))
|
||||
fps = args.fps
|
||||
else:
|
||||
raise ValueError(f"{video_path} should be a video file, an image file or a directory of images")
|
||||
|
||||
# Extract audio features
|
||||
whisper_input_features, librosa_length = audio_processor.get_audio_feature(audio_path)
|
||||
whisper_chunks = audio_processor.get_whisper_chunk(
|
||||
whisper_input_features,
|
||||
device,
|
||||
weight_dtype,
|
||||
whisper,
|
||||
librosa_length,
|
||||
fps=fps,
|
||||
audio_padding_length_left=args.audio_padding_length_left,
|
||||
audio_padding_length_right=args.audio_padding_length_right,
|
||||
)
|
||||
|
||||
# Preprocess input images
|
||||
if os.path.exists(crop_coord_save_path) and args.use_saved_coord:
|
||||
print("Using saved coordinates")
|
||||
with open(crop_coord_save_path, 'rb') as f:
|
||||
coord_list = pickle.load(f)
|
||||
frame_list = read_imgs(input_img_list)
|
||||
else:
|
||||
print("Extracting landmarks... time-consuming operation")
|
||||
coord_list, frame_list = get_landmark_and_bbox(input_img_list, bbox_shift)
|
||||
with open(crop_coord_save_path, 'wb') as f:
|
||||
pickle.dump(coord_list, f)
|
||||
|
||||
print(f"Number of frames: {len(frame_list)}")
|
||||
|
||||
# Process each frame
|
||||
input_latent_list = []
|
||||
for bbox, frame in zip(coord_list, frame_list):
|
||||
if bbox == coord_placeholder:
|
||||
continue
|
||||
x1, y1, x2, y2 = bbox
|
||||
y2 = y2 + args.extra_margin
|
||||
y2 = min(y2, frame.shape[0])
|
||||
crop_frame = frame[y1:y2, x1:x2]
|
||||
crop_frame = cv2.resize(crop_frame, (256,256), interpolation=cv2.INTER_LANCZOS4)
|
||||
latents = vae.get_latents_for_unet(crop_frame)
|
||||
input_latent_list.append(latents)
|
||||
|
||||
# Smooth first and last frames
|
||||
frame_list_cycle = frame_list + frame_list[::-1]
|
||||
coord_list_cycle = coord_list + coord_list[::-1]
|
||||
input_latent_list_cycle = input_latent_list + input_latent_list[::-1]
|
||||
|
||||
# Batch inference
|
||||
print("Starting inference")
|
||||
video_num = len(whisper_chunks)
|
||||
batch_size = args.batch_size
|
||||
gen = datagen(
|
||||
whisper_chunks=whisper_chunks,
|
||||
vae_encode_latents=input_latent_list_cycle,
|
||||
batch_size=batch_size,
|
||||
delay_frame=0,
|
||||
device=device,
|
||||
)
|
||||
|
||||
res_frame_list = []
|
||||
total = int(np.ceil(float(video_num) / batch_size))
|
||||
|
||||
# Execute inference
|
||||
for i, (whisper_batch, latent_batch) in enumerate(tqdm(gen, total=total)):
|
||||
audio_feature_batch = pe(whisper_batch)
|
||||
latent_batch = latent_batch.to(dtype=unet.model.dtype)
|
||||
|
||||
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_list.append(res_frame)
|
||||
|
||||
# Pad generated images to original video size
|
||||
print("Padding generated images to original video size")
|
||||
for i, res_frame in enumerate(tqdm(res_frame_list)):
|
||||
bbox = coord_list_cycle[i%(len(coord_list_cycle))]
|
||||
ori_frame = copy.deepcopy(frame_list_cycle[i%(len(frame_list_cycle))])
|
||||
x1, y1, x2, y2 = bbox
|
||||
y2 = y2 + args.extra_margin
|
||||
y2 = min(y2, frame.shape[0])
|
||||
try:
|
||||
res_frame = cv2.resize(res_frame.astype(np.uint8), (x2-x1, y2-y1))
|
||||
except:
|
||||
continue
|
||||
|
||||
# Merge results
|
||||
combine_frame = get_image(ori_frame, res_frame, [x1, y1, x2, y2], mode=args.parsing_mode, fp=fp)
|
||||
cv2.imwrite(f"{result_img_save_path}/{str(i).zfill(8)}.png", combine_frame)
|
||||
|
||||
# Save prediction results
|
||||
temp_vid_path = f"{temp_dir}/temp_{input_basename}_{audio_basename}.mp4"
|
||||
cmd_img2video = f"ffmpeg -y -v warning -r {fps} -f image2 -i {result_img_save_path}/%08d.png -vcodec libx264 -vf format=yuv420p -crf 18 {temp_vid_path}"
|
||||
print("Video generation command:", cmd_img2video)
|
||||
os.system(cmd_img2video)
|
||||
|
||||
cmd_combine_audio = f"ffmpeg -y -v warning -i {audio_path} -i {temp_vid_path} {output_vid_name}"
|
||||
print("Audio combination command:", cmd_combine_audio)
|
||||
os.system(cmd_combine_audio)
|
||||
|
||||
# Clean up temporary files
|
||||
shutil.rmtree(result_img_save_path)
|
||||
os.remove(temp_vid_path)
|
||||
|
||||
shutil.rmtree(save_dir_full)
|
||||
if not args.saved_coord:
|
||||
os.remove(crop_coord_save_path)
|
||||
|
||||
print(f"Results saved to {output_vid_name}")
|
||||
except Exception as e:
|
||||
print("Error occurred during processing:", e)
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--ffmpeg_path", type=str, default="./ffmpeg-4.4-amd64-static/", help="Path to ffmpeg executable")
|
||||
parser.add_argument("--gpu_id", type=int, default=0, help="GPU ID to use")
|
||||
parser.add_argument("--vae_type", type=str, default="sd-vae", help="Type of VAE model")
|
||||
parser.add_argument("--unet_config", type=str, default="./models/musetalk/config.json", help="Path to UNet configuration file")
|
||||
parser.add_argument("--unet_model_path", type=str, default="./models/musetalkV15/unet.pth", help="Path to UNet model weights")
|
||||
parser.add_argument("--whisper_dir", type=str, default="./models/whisper", help="Directory containing Whisper model")
|
||||
parser.add_argument("--inference_config", type=str, default="configs/inference/test_img.yaml", help="Path to inference configuration file")
|
||||
parser.add_argument("--bbox_shift", type=int, default=0, help="Bounding box shift value")
|
||||
parser.add_argument("--result_dir", default='./results', help="Directory for output results")
|
||||
parser.add_argument("--extra_margin", type=int, default=10, help="Extra margin for face cropping")
|
||||
parser.add_argument("--fps", type=int, default=25, help="Video frames per second")
|
||||
parser.add_argument("--audio_padding_length_left", type=int, default=2, help="Left padding length for audio")
|
||||
parser.add_argument("--audio_padding_length_right", type=int, default=2, help="Right padding length for audio")
|
||||
parser.add_argument("--batch_size", type=int, default=8, help="Batch size for inference")
|
||||
parser.add_argument("--output_vid_name", type=str, default=None, help="Name of output video file")
|
||||
parser.add_argument("--use_saved_coord", action="store_true", help='Use saved coordinates to save time')
|
||||
parser.add_argument("--saved_coord", action="store_true", help='Save coordinates for future use')
|
||||
parser.add_argument("--use_float16", action="store_true", help="Use float16 for faster inference")
|
||||
parser.add_argument("--parsing_mode", default='jaw', help="Face blending parsing mode")
|
||||
parser.add_argument("--left_cheek_width", type=int, default=90, help="Width of left cheek region")
|
||||
parser.add_argument("--right_cheek_width", type=int, default=90, help="Width of right cheek region")
|
||||
args = parser.parse_args()
|
||||
main(args)
|
||||
@@ -10,26 +10,22 @@ 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
|
||||
from transformers import WhisperModel
|
||||
|
||||
from musetalk.utils.face_parsing import FaceParsing
|
||||
from musetalk.utils.utils import datagen
|
||||
from musetalk.utils.preprocessing import get_landmark_and_bbox, read_imgs
|
||||
from musetalk.utils.blending import get_image_prepare_material, get_image_blending
|
||||
from musetalk.utils.utils import load_all_model
|
||||
from musetalk.utils.audio_processor import AudioProcessor
|
||||
|
||||
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)
|
||||
pe = pe.half()
|
||||
vae.vae = vae.vae.half()
|
||||
unet.model = unet.model.half()
|
||||
|
||||
def video2imgs(vid_path, save_path, ext = '.png',cut_frame = 10000000):
|
||||
def video2imgs(vid_path, save_path, ext='.png', cut_frame=10000000):
|
||||
cap = cv2.VideoCapture(vid_path)
|
||||
count = 0
|
||||
while True:
|
||||
@@ -42,35 +38,43 @@ def video2imgs(vid_path, save_path, ext = '.png',cut_frame = 10000000):
|
||||
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()
|
||||
|
||||
@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"
|
||||
# 根据版本设置不同的基础路径
|
||||
if args.version == "v15":
|
||||
self.base_path = f"./results/{args.version}/avatars/{avatar_id}"
|
||||
else: # v1
|
||||
self.base_path = f"./results/avatars/{avatar_id}"
|
||||
|
||||
self.avatar_path = self.base_path
|
||||
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.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.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
|
||||
"avatar_id": avatar_id,
|
||||
"video_path": video_path,
|
||||
"bbox_shift": bbox_shift,
|
||||
"version": args.version
|
||||
}
|
||||
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):
|
||||
@@ -80,7 +84,7 @@ class Avatar:
|
||||
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])
|
||||
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)
|
||||
@@ -98,16 +102,16 @@ class Avatar:
|
||||
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])
|
||||
osmakedirs([self.avatar_path, self.full_imgs_path, self.video_out_path, self.mask_out_path])
|
||||
self.prepare_material()
|
||||
else:
|
||||
else:
|
||||
if not os.path.exists(self.avatar_path):
|
||||
print(f"{self.avatar_id} does not exist, you should set preparation to True")
|
||||
sys.exit()
|
||||
|
||||
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":
|
||||
@@ -115,11 +119,11 @@ class Avatar:
|
||||
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])
|
||||
osmakedirs([self.avatar_path, self.full_imgs_path, self.video_out_path, self.mask_out_path])
|
||||
self.prepare_material()
|
||||
else:
|
||||
sys.exit()
|
||||
else:
|
||||
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)
|
||||
@@ -131,36 +135,40 @@ class Avatar:
|
||||
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')
|
||||
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"]
|
||||
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)
|
||||
# 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
|
||||
if args.version == "v15":
|
||||
y2 = y2 + args.extra_margin
|
||||
y2 = min(y2, frame.shape[0])
|
||||
coord_list[idx] = [x1, y1, x2, y2] # 更新coord_list中的bbox
|
||||
crop_frame = frame[y1:y2, x1:x2]
|
||||
resized_crop_frame = cv2.resize(crop_frame,(256,256),interpolation = cv2.INTER_LANCZOS4)
|
||||
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)
|
||||
|
||||
@@ -170,112 +178,116 @@ class Avatar:
|
||||
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)
|
||||
for i, frame in enumerate(tqdm(self.frame_list_cycle)):
|
||||
cv2.imwrite(f"{self.full_imgs_path}/{str(i).zfill(8)}.png", frame)
|
||||
|
||||
x1, y1, x2, y2 = self.coord_list_cycle[i]
|
||||
if args.version == "v15":
|
||||
mode = args.parsing_mode
|
||||
else:
|
||||
mode = "raw"
|
||||
mask, crop_box = get_image_prepare_material(frame, [x1, y1, x2, y2], fp=fp, mode=mode)
|
||||
|
||||
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,
|
||||
skip_save_images):
|
||||
|
||||
torch.save(self.input_latent_list_cycle, os.path.join(self.latents_out_path))
|
||||
|
||||
def process_frames(self, res_frame_queue, video_len, skip_save_images):
|
||||
print(video_len)
|
||||
while True:
|
||||
if self.idx>=video_len-1:
|
||||
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))])
|
||||
|
||||
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))
|
||||
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)
|
||||
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_blending(ori_frame,res_frame,bbox,mask,mask_crop_box)
|
||||
|
||||
if skip_save_images is False:
|
||||
cv2.imwrite(f"{self.avatar_path}/tmp/{str(self.idx).zfill(8)}.png",combine_frame)
|
||||
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,
|
||||
skip_save_images):
|
||||
os.makedirs(self.avatar_path+'/tmp',exist_ok =True)
|
||||
def inference(self, audio_path, out_vid_name, fps, skip_save_images):
|
||||
os.makedirs(self.avatar_path + '/tmp', exist_ok=True)
|
||||
print("start inference")
|
||||
############################################## extract audio feature ##############################################
|
||||
start_time = time.time()
|
||||
whisper_feature = audio_processor.audio2feat(audio_path)
|
||||
whisper_chunks = audio_processor.feature2chunks(feature_array=whisper_feature,fps=fps)
|
||||
# Extract audio features
|
||||
whisper_input_features, librosa_length = audio_processor.get_audio_feature(audio_path, weight_dtype=weight_dtype)
|
||||
whisper_chunks = audio_processor.get_whisper_chunk(
|
||||
whisper_input_features,
|
||||
device,
|
||||
weight_dtype,
|
||||
whisper,
|
||||
librosa_length,
|
||||
fps=fps,
|
||||
audio_padding_length_left=args.audio_padding_length_left,
|
||||
audio_padding_length_right=args.audio_padding_length_right,
|
||||
)
|
||||
print(f"processing audio:{audio_path} costs {(time.time() - start_time) * 1000}ms")
|
||||
############################################## inference batch by batch ##############################################
|
||||
video_num = len(whisper_chunks)
|
||||
video_num = len(whisper_chunks)
|
||||
res_frame_queue = queue.Queue()
|
||||
self.idx = 0
|
||||
# # Create a sub-thread and start it
|
||||
# Create a sub-thread and start it
|
||||
process_thread = threading.Thread(target=self.process_frames, args=(res_frame_queue, video_num, skip_save_images))
|
||||
process_thread.start()
|
||||
|
||||
gen = datagen(whisper_chunks,
|
||||
self.input_latent_list_cycle,
|
||||
self.batch_size)
|
||||
self.input_latent_list_cycle,
|
||||
self.batch_size)
|
||||
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)))):
|
||||
audio_feature_batch = torch.from_numpy(whisper_batch)
|
||||
audio_feature_batch = audio_feature_batch.to(device=unet.device,
|
||||
dtype=unet.model.dtype)
|
||||
audio_feature_batch = pe(audio_feature_batch)
|
||||
latent_batch = latent_batch.to(dtype=unet.model.dtype)
|
||||
|
||||
pred_latents = unet.model(latent_batch,
|
||||
timesteps,
|
||||
encoder_hidden_states=audio_feature_batch).sample
|
||||
for i, (whisper_batch, latent_batch) in enumerate(tqdm(gen, total=int(np.ceil(float(video_num) / self.batch_size)))):
|
||||
audio_feature_batch = pe(whisper_batch.to(device))
|
||||
latent_batch = latent_batch.to(device=device, dtype=unet.model.dtype)
|
||||
|
||||
pred_latents = unet.model(latent_batch,
|
||||
timesteps,
|
||||
encoder_hidden_states=audio_feature_batch).sample
|
||||
pred_latents = pred_latents.to(device=device, dtype=vae.vae.dtype)
|
||||
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 args.skip_save_images is True:
|
||||
print('Total process time of {} frames without saving images = {}s'.format(
|
||||
video_num,
|
||||
time.time()-start_time))
|
||||
video_num,
|
||||
time.time() - start_time))
|
||||
else:
|
||||
print('Total process time of {} frames including saving images = {}s'.format(
|
||||
video_num,
|
||||
time.time()-start_time))
|
||||
video_num,
|
||||
time.time() - start_time))
|
||||
|
||||
if out_vid_name is not None and args.skip_save_images is False:
|
||||
if out_vid_name is not None and args.skip_save_images is False:
|
||||
# optional
|
||||
cmd_img2video = f"ffmpeg -y -v warning -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"
|
||||
cmd_img2video = f"ffmpeg -y -v warning -r {fps} -f image2 -i {self.avatar_path}/tmp/%08d.png -vcodec libx264 -vf 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
|
||||
output_vid = os.path.join(self.video_out_path, out_vid_name + ".mp4") # on
|
||||
cmd_combine_audio = f"ffmpeg -y -v warning -i {audio_path} -i {self.avatar_path}/temp.mp4 {output_vid}"
|
||||
print(cmd_combine_audio)
|
||||
os.system(cmd_combine_audio)
|
||||
@@ -284,52 +296,95 @@ class Avatar:
|
||||
shutil.rmtree(f"{self.avatar_path}/tmp")
|
||||
print(f"result is save to {output_vid}")
|
||||
print("\n")
|
||||
|
||||
|
||||
|
||||
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,
|
||||
)
|
||||
parser.add_argument("--version", type=str, default="v15", choices=["v1", "v15"], help="Version of MuseTalk: v1 or v15")
|
||||
parser.add_argument("--ffmpeg_path", type=str, default="./ffmpeg-4.4-amd64-static/", help="Path to ffmpeg executable")
|
||||
parser.add_argument("--gpu_id", type=int, default=0, help="GPU ID to use")
|
||||
parser.add_argument("--vae_type", type=str, default="sd-vae", help="Type of VAE model")
|
||||
parser.add_argument("--unet_config", type=str, default="./models/musetalk/musetalk.json", help="Path to UNet configuration file")
|
||||
parser.add_argument("--unet_model_path", type=str, default="./models/musetalk/pytorch_model.bin", help="Path to UNet model weights")
|
||||
parser.add_argument("--whisper_dir", type=str, default="./models/whisper", help="Directory containing Whisper model")
|
||||
parser.add_argument("--inference_config", type=str, default="configs/inference/realtime.yaml")
|
||||
parser.add_argument("--bbox_shift", type=int, default=0, help="Bounding box shift value")
|
||||
parser.add_argument("--result_dir", default='./results', help="Directory for output results")
|
||||
parser.add_argument("--extra_margin", type=int, default=10, help="Extra margin for face cropping")
|
||||
parser.add_argument("--fps", type=int, default=25, help="Video frames per second")
|
||||
parser.add_argument("--audio_padding_length_left", type=int, default=2, help="Left padding length for audio")
|
||||
parser.add_argument("--audio_padding_length_right", type=int, default=2, help="Right padding length for audio")
|
||||
parser.add_argument("--batch_size", type=int, default=25, help="Batch size for inference")
|
||||
parser.add_argument("--output_vid_name", type=str, default=None, help="Name of output video file")
|
||||
parser.add_argument("--use_saved_coord", action="store_true", help='Use saved coordinates to save time')
|
||||
parser.add_argument("--saved_coord", action="store_true", help='Save coordinates for future use')
|
||||
parser.add_argument("--parsing_mode", default='jaw', help="Face blending parsing mode")
|
||||
parser.add_argument("--left_cheek_width", type=int, default=90, help="Width of left cheek region")
|
||||
parser.add_argument("--right_cheek_width", type=int, default=90, help="Width of right cheek region")
|
||||
parser.add_argument("--skip_save_images",
|
||||
action="store_true",
|
||||
help="Whether skip saving images for better generation speed calculation",
|
||||
)
|
||||
action="store_true",
|
||||
help="Whether skip saving images for better generation speed calculation",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
|
||||
# Set computing device
|
||||
device = torch.device(f"cuda:{args.gpu_id}" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
# Load model weights
|
||||
vae, unet, pe = load_all_model(
|
||||
unet_model_path=args.unet_model_path,
|
||||
vae_type=args.vae_type,
|
||||
unet_config=args.unet_config,
|
||||
device=device
|
||||
)
|
||||
timesteps = torch.tensor([0], device=device)
|
||||
|
||||
pe = pe.half().to(device)
|
||||
vae.vae = vae.vae.half().to(device)
|
||||
unet.model = unet.model.half().to(device)
|
||||
|
||||
# Initialize audio processor and Whisper model
|
||||
audio_processor = AudioProcessor(feature_extractor_path=args.whisper_dir)
|
||||
weight_dtype = unet.model.dtype
|
||||
whisper = WhisperModel.from_pretrained(args.whisper_dir)
|
||||
whisper = whisper.to(device=device, dtype=weight_dtype).eval()
|
||||
whisper.requires_grad_(False)
|
||||
|
||||
# Initialize face parser with configurable parameters based on version
|
||||
if args.version == "v15":
|
||||
fp = FaceParsing(
|
||||
left_cheek_width=args.left_cheek_width,
|
||||
right_cheek_width=args.right_cheek_width
|
||||
)
|
||||
else: # v1
|
||||
fp = FaceParsing()
|
||||
|
||||
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"]
|
||||
if args.version == "v15":
|
||||
bbox_shift = 0
|
||||
else:
|
||||
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)
|
||||
|
||||
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,
|
||||
args.skip_save_images)
|
||||
print("Inferring using:", audio_path)
|
||||
avatar.inference(audio_path,
|
||||
audio_num,
|
||||
args.fps,
|
||||
args.skip_save_images)
|
||||
|
||||
Reference in New Issue
Block a user