腾讯开始数字人mousetalk 部署笔记
目录
训练:
自动下载s3fd人脸检测器:
依赖项安装:
模型下载
hf.exe找不到
模型下载脚本,需要用notepad 把格式改为windows格式,
使用效果笔记
测试脚本,infer.py
是2d数字人。
开源地址:
https://github.com/TMElyralab/MuseTalk/blob/main/download_weights.bat
训练:
提供了训练代码:
train.py
自动下载s3fd人脸检测器:
s3fd-619a316812.pth
依赖项安装:
pip install --no-cache-dir -U openmim
mim install mmengine
mim install "mmcv==2.0.1"
mim install "mmdet==3.1.0"
mim install "mmpose==1.1.0"
mmcv安装的2.2.0,mmdet和mmpose版本依赖性报错,把报错注释了也可以跑。
mmcv安装手册:
https://blog.csdn.net/jacke121/article/details/152331246
模型下载
windows download_weights.bat
hf.exe找不到
everything 搜索文件
hf.exe
把锁在目录 添加到系统环境变量中,
模型下载脚本,需要用notepad 把格式改为windows格式,
@echo off
setlocal:: Set the checkpoints directory
set CheckpointsDir=models:: Create necessary directories
mkdir %CheckpointsDir%\musetalk
mkdir %CheckpointsDir%\musetalkV15
mkdir %CheckpointsDir%\syncnet
mkdir %CheckpointsDir%\dwpose
mkdir %CheckpointsDir%\face-parse-bisent
mkdir %CheckpointsDir%\sd-vae-ft-mse
mkdir %CheckpointsDir%\whisper:: Install required packages
pip install -U "huggingface_hub[hf_xet]":: 请将下面的"你的访问令牌"替换为你的真实Hugging Face Token(以hf_开头的那串字符)
huggingface-cli login --token "hf_Akxxxxxxxxxxxxxxxxx" --add-to-git-credential:: 登录成功后,如果需要,再设置镜像端点用于后续的模型下载
set HF_ENDPOINT=https://hf-mirror.com:: Download MuseTalk weights
hf download TMElyralab/MuseTalk --local-dir %CheckpointsDir%:: Download SD VAE weights
hf download stabilityai/sd-vae-ft-mse --local-dir %CheckpointsDir%\sd-vae --include "config.json" "diffusion_pytorch_model.bin":: Download Whisper weights
hf download openai/whisper-tiny --local-dir %CheckpointsDir%\whisper --include "config.json" "pytorch_model.bin" "preprocessor_config.json":: Download DWPose weights
hf download yzd-v/DWPose --local-dir %CheckpointsDir%\dwpose --include "dw-ll_ucoco_384.pth":: Download SyncNet weights
hf download ByteDance/LatentSync --local-dir %CheckpointsDir%\syncnet --include "latentsync_syncnet.pt":: Download face-parse-bisent weights
hf download ManyOtherFunctions/face-parse-bisent --local-dir %CheckpointsDir%\face-parse-bisent --include "79999_iter.pth" "resnet18-5c106cde.pth"echo All weights have been downloaded successfully!
endlocal
使用效果笔记
8g显存,推理挺慢的,一个视频需要好几分钟。
测试脚本,infer.py
注意,要放到根目录。
import os
import cv2
import math
import copy
import torch
import glob
import shutil
import pickle
import argparse
import numpy as np
import subprocess
from tqdm import tqdm
from omegaconf import OmegaConf
from transformers import WhisperModel
import sysfrom 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_placeholderdef fast_check_ffmpeg():try:subprocess.run(["ffmpeg", "-version"], capture_output=True, check=True)return Trueexcept:return False@torch.no_grad()
def main(args):# Configure ffmpeg pathif not fast_check_ffmpeg():print("Adding ffmpeg to PATH")# Choose path separator based on operating systempath_separator = ';' if sys.platform == 'win32' else ':'# os.environ["PATH"] = f"{args.ffmpeg_path}{path_separator}{os.environ['PATH']}"if not fast_check_ffmpeg():print("Warning: Unable to find ffmpeg, please ensure ffmpeg is properly installed")# Set computing devicedevice = torch.device(f"cuda:{args.gpu_id}" if torch.cuda.is_available() else "cpu")# Load model weightsvae, 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 enabledif args.use_float16:pe = pe.half()vae.vae = vae.vae.half()unet.model = unet.model.half()# Move models to specified devicepe = pe.to(device)vae.vae = vae.vae.to(device)unet.model = unet.model.to(device)# Initialize audio processor and Whisper modelaudio_processor = AudioProcessor(feature_extractor_path=args.whisper_dir)weight_dtype = unet.model.dtypewhisper = 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 versionif args.version == "v15":fp = FaceParsing(left_cheek_width=args.left_cheek_width,right_cheek_width=args.right_cheek_width)else: # v1fp = FaceParsing()# Load inference configurationinference_config = OmegaConf.load(args.inference_config)print("Loaded inference config:", inference_config)# Process each taskfor task_id in inference_config:try:# Get task configurationvideo_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"]# Set bbox_shift based on versionif args.version == "v15":bbox_shift = 0 # v15 uses fixed bbox_shiftelse:bbox_shift = inference_config[task_id].get("bbox_shift", args.bbox_shift) # v1 uses config or default# Set output pathsinput_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 directoriestemp_dir = os.path.join(args.result_dir, f"{args.version}")os.makedirs(temp_dir, exist_ok=True)# Set result save pathsresult_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 pathsif 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 videoif 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.fpselif 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.fpselse:raise ValueError(f"{video_path} should be a video file, an image file or a directory of images")# Extract audio featureswhisper_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 imagesif 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 frameinput_latent_list = []for bbox, frame in zip(coord_list, frame_list):if bbox == coord_placeholder:continuex1, y1, x2, y2 = bboxif args.version == "v15":y2 = y2 + args.extra_marginy2 = 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 framesframe_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 inferenceprint("Starting inference")video_num = len(whisper_chunks)batch_size = args.batch_sizegen = 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 inferencefor 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).samplerecon = vae.decode_latents(pred_latents)for res_frame in recon:res_frame_list.append(res_frame)# Pad generated images to original video sizeprint("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 = bboxif args.version == "v15":y2 = y2 + args.extra_marginy2 = 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 parametersif 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 resultstemp_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 filesshutil.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_config", type=str, default="./models\musetalkV15\musetalk.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.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/test', 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)