@@ -1,13 +1,18 @@
""" MuseTalk Flask HTTP 服务 — 反向轮询架构的服务端部分.
部署在 RTX2060 本地,接收 gpu_worker.py 的推理请求,调用 MuseTalk 生成口型同步视频。
本文件修复了原 worker.py 的 8 个工程 bug,并新增 /cancel 端点。
#2000 关键修复:
- 集成真实 MuseTalk 推理(替换原有 stub 代码)
- 音频预处理:22050Hz MP3 → 16kHz mono 16bit WAV( MuseTalk 要求)
- 模型懒加载:首次推理时加载,后续复用,避免重复加载
- 视频帧循环使用 mirror indexing(乒乓模式),消除循环边界跳变
- bbox_shift 可通过请求参数配置
#1978 性能修复(v2 架构):
MuseTalk 原生支持长音频输入(内部循环视频帧),不需要我们先 loop 视频。
正确流程:原视频 + 全量音频 → MuseTalk 推理 → 输出时长=音频时长的无声画面
→ ffmpeg 快速 -c:v copy 替换音轨。推理时间不变(~14s),后处理几秒。
禁止在推理前用 ffmpeg 循环视频(会导致 MuseTalk 处理 2x+ 帧数,慢 16 倍)。
环境变量:
MUSE_PORT 监听端口,默认 7861
@@ -18,21 +23,27 @@
MUSE_DEFAULT_FPS 视频 fps 兜底值,默认 25.0
MUSE_TEMP_DIR 临时文件目录,默认 /tmp/musetalk_$$
MUSE_VIDEO_ENCODER 循环视频时的编码器(仅兜底):auto(默认)/h264_nvenc/libx264
MUSE_DIR MuseTalk 仓库路径,默认 /home/ying/projects/MuseTalk
MUSE_MODEL_DIR 模型目录(相对 MUSE_DIR),默认 models/musetalk
MUSE_USE_FLOAT16 使用 FP16 推理,默认 1(开启)
MUSE_BATCH_SIZE 推理批次大小,默认 8
接口:
GET /health 健康检查 + GPU 显存信息
POST /inference 推理请求(multipart: video + audio)
POST /inference 推理请求(multipart: video + audio, form: bbox_shift )
POST /cancel 终止当前推理任务
"""
from __future__ import annotations
import atexit
import copy
import logging
import os
import shutil
import signal
import subprocess
import sys
import threading
import time
from pathlib import Path
@@ -68,6 +79,16 @@ class Config:
video_encoder : str = _env ( " MUSE_VIDEO_ENCODER " , " auto " ) or " auto "
# 判定音视频时长差异的容差(秒)
duration_epsilon : float = 0.25
# MuseTalk 仓库路径
muse_dir : str = _env ( " MUSE_DIR " , " /home/ying/projects/MuseTalk " )
# 模型目录(相对 MUSE_DIR)
muse_model_dir : str = _env ( " MUSE_MODEL_DIR " , " models/musetalk " )
# 是否使用 FP16(节省显存,RTX2060 建议开启)
use_float16 : bool = _env ( " MUSE_USE_FLOAT16 " , " 1 " ) == " 1 "
# 推理批次大小(RTX2060 6G 显存建议 4-8)
batch_size : int = int ( _env ( " MUSE_BATCH_SIZE " , " 8 " ) )
# GFPGAN 人脸超分增强(提升生成人脸清晰度,+~170MB VRAM, +60ms/帧)
use_gfpgan : bool = _env ( " MUSE_USE_GFPGAN " , " 1 " ) == " 1 "
# ── 全局状态 ──────────────────────────────────────────────────────────
@@ -75,6 +96,13 @@ inference_lock = threading.Lock()
current_task : dict = { " task_id " : None , " process " : None , " start_time " : 0.0 }
shutdown_event = threading . Event ( )
# ── MuseTalk 模型懒加载 ─────────────────────────────────────────────
_muse_models = None
_muse_models_lock = threading . Lock ( )
_muse_models_loaded = False
_muse_load_error = None
# ── Flask App ─────────────────────────────────────────────────────────
app = Flask ( __name__ )
@@ -215,6 +243,29 @@ def _pick_video_encoder() -> str:
return " libx264 "
def _preprocess_audio ( input_path : Path , output_path : Path , target_sr : int = 16000 ) - > None :
""" 将输入音频转换为 MuseTalk 要求的格式:16kHz mono 16bit WAV.
MuseTalk 的 whisper audio2feature 要求 16kHz 采样率的单声道音频。
当前 TTS 输出为 22050Hz MP3,不转换会导致 mel 频谱错位、
音素特征提取错误,口型只跟能量不跟音素。
"""
cmd = [
" ffmpeg " , " -y " , " -v " , " warning " ,
" -i " , str ( input_path ) ,
" -ar " , str ( target_sr ) , # 重采样到 16kHz
" -ac " , " 1 " , # 单声道
" -sample_fmt " , " s16 " , # 16bit PCM
str ( output_path ) ,
]
_run_ffmpeg ( cmd , timeout = 60 )
if not output_path . exists ( ) or output_path . stat ( ) . st_size < 100 :
raise RuntimeError ( f " 音频预处理失败: { output_path } " )
logger . info ( " 音频预处理完成: %s → 16kHz mono WAV " , input_path . name )
def _mux_video_with_audio (
video_path : Path ,
audio_path : Path ,
@@ -340,129 +391,518 @@ def _run_ffmpeg(cmd: list, timeout: float = 120) -> subprocess.CompletedProcess:
raise RuntimeError ( f " ffmpeg 超时(> { timeout } s) " ) from exc
# ── MuseTalk 模型加载 ────────────────────────────────────────────────
def _load_musetalk_models ( ) :
""" 懒加载 MuseTalk 模型(全局单例,首次调用时加载).
加载 VAE、UNet、PositionalEncoder 三个核心组件。
加载到 GPU 后转为 FP16(如果配置开启)以节省显存。
RTX2060 6G 显存,FP16 大约需要 3-4GB。
"""
global _muse_models , _muse_models_loaded , _muse_load_error
if _muse_models_loaded :
return _muse_models
if _muse_load_error is not None :
raise _muse_load_error
with _muse_models_lock :
if _muse_models_loaded :
return _muse_models
try :
muse_dir = Path ( Config . muse_dir )
if not muse_dir . exists ( ) :
raise FileNotFoundError (
f " MuseTalk 目录不存在: { muse_dir } \n "
f " 请设置 MUSE_DIR 环境变量指向 MuseTalk 仓库路径 "
)
# 将 MuseTalk 加入 sys.path(只在首次加载时)
muse_str = str ( muse_dir )
if muse_str not in sys . path :
sys . path . insert ( 0 , muse_str )
import torch
from musetalk . utils . utils import load_all_model
device = torch . device ( " cuda:0 " if torch . cuda . is_available ( ) else " cpu " )
logger . info ( " MuseTalk 使用设备: %s " , device )
# 自动检测模型路径
# 优先检测 v1.5 模型,然后回退到 v1
v15_unet = muse_dir / " models " / " musetalkV15 " / " unet.pth "
v1_unet = muse_dir / " models " / " musetalk " / " pytorch_model.bin "
if v15_unet . exists ( ) :
unet_model_path = str ( v15_unet )
unet_config = str ( muse_dir / " models " / " musetalkV15 " / " musetalk.json " )
model_version = " v15 "
elif v1_unet . exists ( ) :
unet_model_path = str ( v1_unet )
unet_config = str ( muse_dir / " models " / " musetalk " / " config.json " )
model_version = " v1 "
else :
raise FileNotFoundError (
f " 未找到 MuseTalk 模型权重。 \n "
f " 检查路径: { v15_unet } 或 { v1_unet } \n "
f " 请确认模型已下载到 MuseTalk 仓库的 models/ 目录下 "
)
logger . info ( " 加载 MuseTalk %s 模型: %s " , model_version , unet_model_path )
vae , unet , pe = load_all_model (
unet_model_path = unet_model_path ,
vae_type = " sd-vae " ,
unet_config = unet_config ,
device = device ,
)
timesteps = torch . tensor ( [ 0 ] , device = device )
# FP16 转换(节省 ~50% 显存)
if Config . use_float16 :
pe = pe . half ( )
vae . vae = vae . vae . half ( )
unet . model = unet . model . half ( )
logger . info ( " 已启用 FP16 推理 " )
pe = pe . to ( device )
vae . vae = vae . vae . to ( device )
unet . model = unet . model . to ( device )
# 加载 AudioProcessor 和 face parsing
from musetalk . utils . audio_processor import AudioProcessor
from musetalk . utils . face_parsing import FaceParsing
audio_processor = AudioProcessor ( )
face_parsing = FaceParsing ( )
# 加载 GFPGAN 人脸超分模型(FP16,仅 ~170MB VRAM)
gfpgan_model = None
if Config . use_gfpgan :
try :
from gfpgan . archs . gfpganv1_clean_arch import GFPGANv1Clean
gfpgan_path = muse_dir / " models " / " GFPGAN " / " GFPGANv1.4.pth "
if gfpgan_path . exists ( ) :
logger . info ( " 加载 GFPGANv1.4 人脸超分模型: %s " , gfpgan_path )
gfpgan_ckpt = torch . load ( str ( gfpgan_path ) , map_location = " cpu " )
gfpgan_model = GFPGANv1Clean (
out_size = 512 , num_style_feat = 512 , channel_multiplier = 2 ,
decoder_load_path = None , fix_decoder = False , num_mlp = 8 ,
input_is_latent = True , different_w = True , narrow = 1 , sft_half = True ,
)
gfpgan_key = " params_ema " if " params_ema " in gfpgan_ckpt else " params "
gfpgan_model . load_state_dict ( gfpgan_ckpt [ gfpgan_key ] , strict = True )
gfpgan_model . eval ( )
if Config . use_float16 :
gfpgan_model = gfpgan_model . half ( )
gfpgan_model = gfpgan_model . to ( device )
del gfpgan_ckpt
logger . info ( " GFPGAN 加载完成 (FP16= %s ) " , Config . use_float16 )
else :
logger . warning ( " GFPGAN 模型不存在: %s ,跳过人脸增强 " , gfpgan_path )
except Exception as e :
logger . warning ( " GFPGAN 加载失败,跳过人脸增强: %s " , e )
gfpgan_model = None
else :
logger . info ( " GFPGAN 已禁用 (MUSE_USE_GFPGAN=0) " )
_muse_models = {
" vae " : vae ,
" unet " : unet ,
" pe " : pe ,
" timesteps " : timesteps ,
" audio_processor " : audio_processor ,
" face_parsing " : face_parsing ,
" gfpgan " : gfpgan_model ,
" device " : device ,
" model_version " : model_version ,
}
_muse_models_loaded = True
logger . info ( " MuseTalk 模型加载完成 (版本= %s , 设备= %s , fp16= %s ) " ,
model_version , device , Config . use_float16 )
# 打印显存使用情况
if torch . cuda . is_available ( ) :
allocated = torch . cuda . memory_allocated ( ) / 1024 * * 2
reserved = torch . cuda . memory_reserved ( ) / 1024 * * 2
logger . info ( " GPU 显存: 已分配 %.0f MB, 已预留 %.0f MB " , allocated , reserved )
return _muse_models
except Exception as exc :
_muse_load_error = exc
logger . error ( " MuseTalk 模型加载失败: %s " , exc )
raise
# ── MuseTalk 推理核心 ─────────────────────────────────────────────────
def _mirror_index ( size : int , index : int ) - > int :
""" 乒乓式循环索引,避免循环边界硬切跳变.
效果: 0→1→2→...→N→N-1→...→1→0→1→...
比简单的 index % s ize 在边界处更平滑。
"""
if size == 0 :
return 0
turn = index / / size
res = index % size
if turn % 2 == 0 :
return res
else :
return size - res - 1
def _run_inference (
video_path : Path ,
audio_path : Path ,
output_path : Path ,
bbox_shift : int = 0 ,
) - > None :
""" 执行 MuseTalk 推理(v2 架构:全量音频直传,不在推理前 loop 视频) .
""" 执行 MuseTalk 真实 推理.
#1978 性能修复核心 :
MuseTalk 原生支持长音频输入,内部会自动循环视频帧。
我们只需把【原视频】和【全量音频】传给 MuseTalk,
输出视频时长 = 音频时长(MuseTalk 自行处理帧循环)。
禁止在推理前用 ffmpeg 循环视频(会导致慢 16 倍)。
流程 :
1. 音频预处理:任意格式 → 16kHz mono 16bit WAV
2. 加载/复用 MuseTalk 模型(VAE + UNet + PE + Whisper)
3. 视频预处理:提取帧 → 人脸检测 → 获取 bbox → VAE 编码 latent
4. 音频特征提取:whisper 提取 audio features (50× 384 per chunk)
5. 批量推理:UNet 去噪 → VAE 解码 → 得到口型同步的人脸帧
6. 帧合成:将生成的人脸贴回原帧(使用 face parsing 做边缘融合)
7. 输出无声视频(后续由 _mux_video_with_audio 封装 TTS 音频)
实际部署时替换为 MuseTalk 真实推理逻辑。
此处为示例实现:提取帧 → 模拟 MuseTalk 产出音频时长的无声画面 → 快速封装。
Args:
video_path: 输入视频路径
audio_path: 输入音频路径(任意格式,会被预处理为 16kHz WAV)
output_path: 输出无声视频路径
bbox_shift: 口型区域垂直偏移量,默认 0,范围 [-5, 5]
"""
import cv2
import numpy as np
import torch
from tqdm import tqdm
from musetalk . utils . preprocessing import get_landmark_and_bbox as _orig_get_landmark_and_bbox
from musetalk . utils . blending import get_image
import tempfile as _tempfile , math as _math , shutil as _shutil
from einops import rearrange as _rearrange
# read_imgs: 读取视频帧(支持视频文件路径),返回 numpy BGR 帧列表
def read_imgs ( path ) :
import cv2 as _cv2
cap = _cv2 . VideoCapture ( str ( path ) )
frames = [ ]
while True :
ret , frame = cap . read ( )
if not ret :
break
frames . append ( frame )
cap . release ( )
return frames
# get_landmark_and_bbox 适配:旧版签名(img_list, upperbondrange=0),且 img_list 是文件路径列表
def get_landmark_and_bbox ( frames , vid_pts = 0 , bbox_shift = 0 ) :
import cv2 as _cv2
_tmpdir = _tempfile . mkdtemp ( prefix = " muse_frames_ " )
frame_paths = [ ]
for _i , _frm in enumerate ( frames ) :
_fp = f " { _tmpdir } / { _i : 08d } .png "
_cv2 . imwrite ( _fp , _frm )
frame_paths . append ( _fp )
coords_list , _ = _orig_get_landmark_and_bbox ( frame_paths , upperbondrange = bbox_shift )
_shutil . rmtree ( _tmpdir , ignore_errors = True )
_sentinel = object ( )
coords_list = [ c if c is not None else _sentinel for c in coords_list ]
return coords_list , _sentinel
# 加载模型(首次调用时加载,后续复用)
models = _load_musetalk_models ( )
vae = models [ " vae " ]
unet = models [ " unet " ]
pe = models [ " pe " ]
timesteps = models [ " timesteps " ]
audio_processor = models [ " audio_processor " ]
device = models [ " device " ]
model_version = models [ " model_version " ]
# 给旧版 AudioProcessor 动态添加 feature2chunks 方法
import types as _types
def _feature2chunks ( self , feature_array , fps = 25 , weight_dtype = None ,
batch_size = 8 , audio_padding_length_left = 2 ,
audio_padding_length_right = 2 ) :
import torch
sr = 16000
audio_fps = 50
chunk_len = 2 * ( audio_padding_length_left + audio_padding_length_right + 1 )
whisper_idx_multiplier = audio_fps / fps
num_frames = int ( _math . floor ( ( len ( feature_array ) / sr ) * fps ) )
actual_length = int ( _math . floor ( ( len ( feature_array ) / sr ) * audio_fps ) )
inputs = self . feature_extractor (
feature_array , return_tensors = " pt " , sampling_rate = sr
) . input_features . to ( device )
if weight_dtype is not None :
inputs = inputs . to ( dtype = weight_dtype )
global _whisper_enc_model
if " _whisper_enc_model " not in globals ( ) or _whisper_enc_model is None :
from transformers import WhisperModel
_wp = str ( Path ( Config . muse_dir ) / " models " / " whisper " )
_whisper_enc_model = WhisperModel . from_pretrained ( _wp ) . to ( device )
_whisper_enc_model . eval ( )
if Config . use_float16 :
_whisper_enc_model = _whisper_enc_model . half ( )
with torch . no_grad ( ) :
_af = _whisper_enc_model . encoder ( inputs , output_hidden_states = True ) . hidden_states
_af = torch . stack ( _af , dim = 2 )
_af = _af [ 0 , : actual_length , . . . ]
_pn = int ( _math . ceil ( whisper_idx_multiplier ) )
_af = torch . cat ( [
torch . zeros_like ( _af [ : _pn * audio_padding_length_left ] ) ,
_af ,
torch . zeros_like ( _af [ : _pn * 3 * audio_padding_length_right ] ) ,
] , dim = 0 )
_all = [ ]
for _fi in range ( num_frames ) :
_ai = int ( _math . floor ( _fi * whisper_idx_multiplier ) )
_clip = _af [ _ai : _ai + chunk_len ]
if _clip . shape [ 0 ] < chunk_len :
_pad = torch . zeros ( chunk_len - _clip . shape [ 0 ] , * _clip . shape [ 1 : ] ,
device = device , dtype = _clip . dtype )
_clip = torch . cat ( [ _clip , _pad ] , dim = 0 )
_all . append ( _clip )
_prompts = torch . stack ( _all , dim = 0 )
_prompts = _rearrange ( _prompts , " b c h w -> b (c h) w " )
return _prompts
audio_processor . feature2chunks = _types . MethodType ( _feature2chunks , audio_processor )
fps = _get_video_fps ( video_path )
audio_duration = _get_media_duration ( audio_path )
video_duration = _get_media_duration ( video_path )
logger . info (
" 推理开始: video= %.2f s, audio= %.2f s, fps= %.2f " ,
video_duration ,
audio_duration ,
fps ,
" MuseTalk 推理开始: video=%.2f s, audio= %.2f s, fps= %.1f , bbox_shift= %d " ,
video_duration , audio_duration , fps , bbox_shift ,
)
frames_dir = video_path . parent / " frames "
frames_dir . mkdir ( parents = True , exist_ok = True )
# ── Step 1: 音频预处理(关键修复:22050Hz MP3 → 16kHz mono WAV)──
audio_wav_path = video_path . parent / " audio_16k_mono.wav "
_preprocess_audio ( audio_path , audio_wav_path , target_sr = 16000 )
# 1. 从原 视频提取帧(仅原视频长度,不循环)
_run_ffmpeg (
[
" ffmpeg " ,
" -y " ,
" -i " ,
str ( video_path ) ,
" -r " ,
str ( fps ) ,
str ( frames_dir / " frame_ %05d .png " ) ,
] ,
timeout = 120 ,
# ── Step 2: 视频帧 提取 ──
input_frames = read_imgs ( str ( video_path ) )
total_frames = len ( input_frames )
if total_frames == 0 :
raise RuntimeError ( " 未能从视频中提取到任何帧 " )
logger . info ( " 提取到 %d 帧视频画面 " , total_frames )
# ── Step 3: 人脸检测 & bbox 计算 ──
coord_list , coord_placeholder = get_landmark_and_bbox (
input_frames , vid_pts = 0 , bbox_shift = bbox_shift
)
logger . info ( " 人脸检测完成,有效 bbox: %d / %d " , sum ( 1 for c in coord_list if c is not coord_placeholder ) , total_frames )
frame_files = sorted ( frames_dir . glob ( " *.png " ) )
if not frame_files :
raise RuntimeError ( " 未从视频中提取到帧 " )
# 使用 mirror indexing 循环帧和坐标(避免硬切跳变)
num_output_frames = int ( audio_duration * fps )
if num_output_frames < = 0 :
num_output_frames = total_frames
# 2. 模拟 MuseTalk 推理:输入原视频帧 + 全量音频,输出音频时长的无声画面。
# TODO: 替换为 MuseTalk 真实推理逻辑。
# MuseTalk 真实调用示例(伪代码):
# from musetalk import MuseTalkModel
# model = MuseTalkModel(...)
# silent_video = model.infer(video_path=video_path, audio_path=audio_path)
# # MuseTalk 内部会循环视频帧匹配音频长度,输出时长=音频时长
logger . warning ( " 使用示例推理逻辑,未实际调用 MuseTalk 模型 " )
# 示例:生成音频时长的无声画面(循环原视频帧到音频长度)
# 真实部署时 silent_video_path 应替换为 MuseTalk 输出的无声视频路径
silent_video_path = video_path . parent / " visual_silent.mp4 "
if audio_duration > video_duration + Config . duration_epsilon :
# 音频更长:循环视频帧到音频长度(仅用于示例,真实 MuseTalk 内部处理)
encoder = _pick_video_encoder ( )
preset = " p4 " if encoder == " h264_nvenc " else " veryfast "
logger . info (
" 示例:循环视频帧到音频长度 %.2f s(真实 MuseTalk 内部处理,无需此步骤) " ,
audio_duration ,
)
cmd = [
" ffmpeg " ,
" -y " ,
" -stream_loop " ,
" -1 " ,
" -i " ,
str ( video_path ) ,
" -an " ,
" -c:v " ,
encoder ,
" -preset " ,
preset ,
" -t " ,
f " { audio_duration : .3f } " ,
str ( silent_video_path ) ,
]
try :
_run_ffmpeg ( cmd , timeout = 300 )
except RuntimeError :
if encoder == " h264_nvenc " :
cmd [ cmd . index ( encoder ) ] = " libx264 "
cmd [ cmd . index ( preset ) + 1 ] = " veryfast "
_run_ffmpeg ( cmd , timeout = 300 )
else :
raise
else :
# 音频不长:直接生成无声视频(原视频长度)
_run_ffmpeg (
[
" ffmpeg " ,
" -y " ,
" -i " ,
str ( video_path ) ,
" -an " ,
" -c:v " ,
" libx264 " ,
" -preset " ,
" veryfast " ,
str ( silent_video_path ) ,
] ,
timeout = 300 ,
)
# 3. 快速封装:-map 取推理画面 + 驱动音频,-c:v copy 无损秒级封装
# MuseTalk 输出已匹配音频长度,此处无需循环,仅替换音轨
_mux_video_with_audio ( silent_video_path , audio_path , output_path )
if not output_path . exists ( ) or output_path . stat ( ) . st_size < 1024 :
raise RuntimeError ( " 推理产物不存在或过小 " )
logger . info (
" 推理完成: output= %.2f s (audio= %.2f s) " ,
_get_media_duration ( output_path ) ,
audio_duration ,
# ── Step 4: 音频特征提取 ──
# 使用 librosa 加载预处理后的 16kHz 音频
import librosa
audio_array , _ = librosa . load ( str ( audio_wav_path ) , sr = 16000 , mono = True )
whisper_features = audio_processor . feature2chunks (
feature_array = audio_array ,
fps = fps ,
weight_dtype = ( torch . float16 if Config . use_float16 else torch . float32 ) ,
batch_size = Config . batch_size ,
)
if isinstance ( whisper_features , torch . Tensor ) :
whisper_features = whisper_features . detach ( ) . cpu ( )
torch . cuda . empty_cache ( )
logger . info ( " 音频特征提取完成: %d 个 chunk " , len ( whisper_features ) )
# ── Step 5: 逐帧裁剪人脸并编码为 8ch latent (masked+ref) ──
face_parsing = models . get ( " face_parsing " , None )
input_latent_list = [ ]
valid_frame_indices = [ ] # 记录成功编码的帧索引(跳过无脸帧)
with torch . no_grad ( ) :
for idx , ( frame , bbox ) in enumerate ( zip ( input_frames , coord_list ) ) :
if bbox is coord_placeholder :
continue
x1 , y1 , x2 , y2 = bbox
# v1.5 额外扩展下边界(下巴区域),与 Step 7 保持一致
extra_y2 = 10 if model_version == " v15 " else 0
y2_eff = min ( y2 + extra_y2 , frame . shape [ 0 ] )
if y2_eff < = y1 or x2 < = x1 :
continue
# 裁剪人脸区域 → resize 256× 256
crop = frame [ y1 : y2_eff , x1 : x2 ]
if crop . size == 0 :
continue
crop_rgb = cv2 . cvtColor ( crop , cv2 . COLOR_BGR2RGB )
crop_resized = cv2 . resize ( crop_rgb , ( 256 , 256 ) , interpolation = cv2 . INTER_LANCZOS4 )
# 使用 VAE 的 get_latents_for_unet 得到 8 通道输入
# get_latents_for_unet 内部: preprocess(half_mask=True) encode + preprocess(half_mask=False) encode → cat → [1,8,32,32]
latents = vae . get_latents_for_unet ( crop_resized ) . detach ( ) . cpu ( )
input_latent_list . append ( latents )
# 保存此帧的实际bbox(含extra_y2)和原帧索引供 Step 7 使用
valid_frame_indices . append ( ( idx , x1 , y1 , x2 , y2_eff ) )
# 构建循环列表:正序+倒序,实现旧版的平滑首尾帧循环
frame_list_cycle = input_frames + input_frames [ : : - 1 ]
coord_cycle = [ ]
for _i , _x1 , _y1 , _x2 , _y2 in valid_frame_indices :
coord_cycle . append ( ( _x1 , _y1 , _x2 , _y2 ) )
coord_cycle = coord_cycle + coord_cycle [ : : - 1 ]
latent_cycle = input_latent_list + input_latent_list [ : : - 1 ]
valid_cycle = valid_frame_indices + [ ( i , x1 , y1 , x2 , y2 ) for ( i , x1 , y1 , x2 , y2 ) in reversed ( valid_frame_indices ) ]
torch . cuda . empty_cache ( )
logger . info ( " 人脸裁剪+VAE 编码完成: %d 个有效latent " , len ( input_latent_list ) )
# ── Step 6: 批量推理(仿旧版 datagen 循环)──
res_frame_list = [ ]
video_num = len ( whisper_features )
bs = min ( Config . batch_size , 2 ) # RTX2060 6G 限制batch=2防OOM
total_batches = ( video_num + bs - 1 ) / / bs
for bi in tqdm ( range ( total_batches ) , desc = " MuseTalk 推理 " ) :
whisper_batch = whisper_features [ bi * bs : ( bi + 1 ) * bs ]
if len ( whisper_batch ) == 0 :
break
# 对应 latent 索引(循环取 latent_cycle)
latent_batch_parts = [ ]
for j in range ( len ( whisper_batch ) ) :
global_idx = bi * bs + j
lat_idx = global_idx % len ( latent_cycle )
latent_batch_parts . append ( latent_cycle [ lat_idx ] )
# whisper_batch 是 [bs,50,384] tensor slice (feature2chunks 已返回 stacked tensor)
if isinstance ( whisper_batch , list ) :
whisper_batch_t = torch . stack ( whisper_batch ) . to ( device )
else :
whisper_batch_t = whisper_batch . to ( device )
latent_batch_t = torch . cat ( latent_batch_parts , dim = 0 ) . to ( device )
if Config . use_float16 :
latent_batch_t = latent_batch_t . to ( dtype = unet . model . dtype )
whisper_batch_t = whisper_batch_t . to ( dtype = unet . model . dtype )
audio_feature_batch = pe ( whisper_batch_t )
with torch . no_grad ( ) :
pred_latents = unet . model (
latent_batch_t ,
timesteps ,
encoder_hidden_states = audio_feature_batch ,
) . sample
recon_frames = vae . decode_latents ( pred_latents )
for rf in recon_frames :
res_frame_list . append ( rf )
del pred_latents , recon_frames , latent_batch_t , whisper_batch_t
if " audio_feature_batch " in dir ( ) :
try : del audio_feature_batch
except : pass
torch . cuda . empty_cache ( )
logger . info ( " 推理完成,生成 %d 帧 " , len ( res_frame_list ) )
# ── Step 7: 合成最终帧 → ffmpeg pipe 编码(零磁盘IO) ──
gfpgan_enhancer = models . get ( " gfpgan " )
silent_video_path = video_path . parent / " silent_output.mp4 "
frame_h , frame_w = frame_list_cycle [ 0 ] . shape [ : 2 ]
# 启动 ffmpeg: stdin 接收 raw BGR24 帧,直接编码 H.264(省去PNG落盘+回读)
_ff_cmd = [
" ffmpeg " , " -y " , " -v " , " warning " ,
" -f " , " rawvideo " , " -pix_fmt " , " bgr24 " ,
" -s " , f " { frame_w } x { frame_h } " , " -r " , str ( fps ) ,
" -i " , " - " ,
" -vcodec " , " libx264 " , " -preset " , " veryfast " ,
" -vf " , " format=yuv420p " , " -crf " , " 18 " ,
str ( silent_video_path ) ,
]
import subprocess as _sp
_ff_proc = _sp . Popen ( _ff_cmd , stdin = _sp . PIPE , stdout = _sp . DEVNULL , stderr = _sp . PIPE )
n_out = min ( len ( res_frame_list ) , num_output_frames )
try :
for i in tqdm ( range ( n_out ) , desc = " 合成帧 " ) :
cyc_i = i % len ( coord_cycle )
x1 , y1 , x2 , y2 = coord_cycle [ cyc_i ]
ori_frame = copy . deepcopy ( frame_list_cycle [ cyc_i ] )
res_frame = res_frame_list [ i ]
try :
res_frame_resized = cv2 . resize (
res_frame . astype ( np . uint8 ) , ( x2 - x1 , y2 - y1 ) ,
interpolation = cv2 . INTER_LANCZOS4
)
except Exception :
_ff_proc . stdin . write ( ori_frame . tobytes ( ) )
continue
# GFPGAN 人脸超分增强
if gfpgan_enhancer is not None :
try :
_fh , _fw = res_frame_resized . shape [ : 2 ]
_face_up = cv2 . resize ( res_frame_resized , ( 512 , 512 ) ,
interpolation = cv2 . INTER_LANCZOS4 )
_face_rgb = cv2 . cvtColor ( _face_up , cv2 . COLOR_BGR2RGB ) . astype ( np . float32 ) / 255.0
_face_t = torch . from_numpy ( _face_rgb . transpose ( 2 , 0 , 1 ) ) . unsqueeze ( 0 )
_face_t = ( ( _face_t - 0.5 ) / 0.5 ) . to ( device )
if Config . use_float16 :
_face_t = _face_t . half ( )
with torch . no_grad ( ) :
_out = gfpgan_enhancer ( _face_t , return_rgb = False , weight = 0.5 ) [ 0 ]
_out = _out . squeeze ( 0 ) . float ( ) . cpu ( ) . clamp_ ( - 1 , 1 )
_out = ( ( _out + 1 ) / 2 * 255 ) . numpy ( ) . transpose ( 1 , 2 , 0 )
_out_bgr = cv2 . cvtColor ( _out . astype ( np . uint8 ) , cv2 . COLOR_RGB2BGR )
res_frame_resized = cv2 . resize ( _out_bgr , ( _fw , _fh ) ,
interpolation = cv2 . INTER_LANCZOS4 )
del _face_t , _out , _out_bgr
except Exception :
pass
# face parsing 融合
try :
if face_parsing is not None :
combined = get_image ( ori_frame , res_frame_resized ,
[ x1 , y1 , x2 , y2 ] , fp = face_parsing )
else :
combined = get_image ( ori_frame , res_frame_resized , [ x1 , y1 , x2 , y2 ] )
except Exception :
combined = ori_frame . copy ( )
try : combined [ y1 : y2 , x1 : x2 ] = res_frame_resized
except Exception : combined = ori_frame
_ff_proc . stdin . write ( combined . tobytes ( ) )
_ff_proc . stdin . close ( )
_ff_ret = _ff_proc . wait ( timeout = 120 )
if _ff_ret != 0 :
_ff_err = _ff_proc . stderr . read ( ) . decode ( errors = " ignore " ) if _ff_proc . stderr else " "
raise RuntimeError ( f " ffmpeg编码失败(exit= { _ff_ret } ): { _ff_err [ - 300 : ] } " )
except Exception :
try : _ff_proc . kill ( )
except Exception : pass
raise
shutil . copy2 ( str ( silent_video_path ) , str ( output_path ) )
try :
torch . cuda . empty_cache ( )
if audio_wav_path . exists ( ) :
audio_wav_path . unlink ( )
if silent_video_path . exists ( ) and str ( silent_video_path ) != str ( output_path ) :
silent_video_path . unlink ( )
except Exception as e :
logger . warning ( " 清理中间文件失败: %s " , e )
logger . info ( " MuseTalk 推理完成: output= %s , duration= %.2f s " ,
output_path . name , _get_media_duration ( output_path ) )
# ── 路由 ──────────────────────────────────────────────────────────────
@@ -470,7 +910,7 @@ def _run_inference(
@app.route ( " /health " , methods = [ " GET " ] )
def health ( ) :
""" 健康检查 + GPU 显存信息. """
""" 健康检查 + GPU 显存信息 + MuseTalk 模型状态 . """
gpu_info = _get_gpu_info ( )
task_info = {
" task_id " : current_task [ " task_id " ] ,
@@ -482,6 +922,8 @@ def health():
" status " : " healthy " ,
" gpu " : gpu_info ,
" current_task " : task_info ,
" musetalk_loaded " : _muse_models_loaded ,
" musetalk_load_error " : str ( _muse_load_error ) if _muse_load_error else None ,
" timestamp " : time . time ( ) ,
}
)
@@ -491,7 +933,8 @@ def health():
def inference ( ) :
""" 推理请求:multipart form 包含 video 和 audio 文件.
#1978 v2: MuseTalk 直接处理全量音频,输出时长=音频时长,无需预处理循环。
可选 form 参数:
bbox_shift: 口型区域垂直偏移量,默认 0,范围 [-5, 5]
"""
# 并发控制:检查锁
if not inference_lock . acquire ( blocking = False ) :
@@ -510,6 +953,8 @@ def inference():
video_file = request . files [ " video " ]
audio_file = request . files [ " audio " ]
task_id = request . form . get ( " task_id " , f " task_ { int ( time . time ( ) ) } " )
bbox_shift = int ( request . form . get ( " bbox_shift " , " 0 " ) )
bbox_shift = max ( - 5 , min ( 5 , bbox_shift ) ) # 限制范围
# 文件大小检查
err = _check_file_size ( video_file , Config . video_max_mb , " 视频 " )
@@ -523,13 +968,14 @@ def inference():
task_dir = Path ( Config . temp_dir ) / task_id
task_dir . mkdir ( parents = True , exist_ok = True )
video_path = task_dir / " input.mp4 "
audio_path = task_dir / " input_audio.wav "
audio_path = task_dir / " input_audio.bin "
output_path = task_dir / " output.mp4 "
video_file . save ( str ( video_path ) )
audio_file . save ( str ( audio_path ) )
logger . info ( " 开始推理 task_id= %s , video= %s , audio= %s " , task_id , video_path . name , audio_path . name )
logger . info ( " 开始推理 task_id= %s , video= %s , audio= %s , bbox_shift= %d " ,
task_id , video_path . name , audio_path . name , bbox_shift )
# 更新当前任务信息
current_task [ " task_id " ] = task_id
@@ -541,8 +987,9 @@ def inference():
def inference_thread ( ) :
try :
_run_inference ( video_path , audio_path , output_path )
_run_inference ( video_path , audio_path , output_path , bbox_shift = bbox_shift )
except Exception as exc :
logger . exception ( " 推理异常: %s " , exc )
result_container [ " error " ] = str ( exc )
thread = threading . Thread ( target = inference_thread )
@@ -621,6 +1068,8 @@ def cancel():
def main ( ) :
import os as _os
_os . environ . setdefault ( " PYTORCH_CUDA_ALLOC_CONF " , " max_split_size_mb:128 " )
""" 启动 Flask 服务. """
# 创建临时目录
Path ( Config . temp_dir ) . mkdir ( parents = True , exist_ok = True )
@@ -633,11 +1082,14 @@ def main():
gpu_info [ " memory_used_mb " ] ,
gpu_info [ " memory_total_mb " ] ,
)
logger . info ( " MuseTalk 仓库路径: %s " , Config . muse_dir )
logger . info (
" 启动 MuseTalk Server: port= %d , timeout= %.0f s, max_concurrent= %d " ,
" 启动 MuseTalk Server: port= %d , timeout= %.0f s, max_concurrent= %d , fp16= %s , batch_size= %d " ,
Config . port ,
Config . inference_timeout ,
Config . max_concurrent ,
Config . use_float16 ,
Config . batch_size ,
)
app . run ( host = " 0.0.0.0 " , port = Config . port , threaded = True )