diff --git a/deploy/gpu_worker/musetalk_server.py b/deploy/gpu_worker/musetalk_server.py index 31213a69a..ef6b2b072 100644 --- a/deploy/gpu_worker/musetalk_server.py +++ b/deploy/gpu_worker/musetalk_server.py @@ -27,10 +27,13 @@ MUSE_MODEL_DIR 模型目录（相对 MUSE_DIR），默认 models/musetalk MUSE_USE_FLOAT16 使用 FP16 推理，默认 1（开启） MUSE_BATCH_SIZE 推理批次大小，默认 8 + MUSE_FACE_ENHANCE 启用 GFPGAN 面部增强后处理，默认 1（开启） + MUSE_ENHANCE_REGION 增强区域：mouth（仅嘴部，快）或 face（全脸，更自然），默认 mouth + MUSE_ENHANCE_UPSCALE 面部增强上采样倍数，默认 1（不放大） 接口： - GET /health 健康检查 + GPU 显存信息 - POST /inference 推理请求（multipart: video + audio, form: bbox_shift） + GET /health 健康检查 + GPU 显存信息 + MuseTalk/GFPGAN 模型状态 + POST /inference 推理请求（multipart: video + audio, form: bbox_shift, face_enhance） POST /cancel 终止当前推理任务 """ @@ -87,6 +90,14 @@ class Config: use_float16: bool = _env("MUSE_USE_FLOAT16", "1") == "1" # 推理批次大小（RTX2060 6G 显存建议 4-8） batch_size: int = int(_env("MUSE_BATCH_SIZE", "8")) + # 面部增强（GFPGAN）：启用开关、区域、上采样倍数 + face_enhance: bool = _env("MUSE_FACE_ENHANCE", "1") == "1" + enhance_region: str = _env("MUSE_ENHANCE_REGION", "mouth") or "mouth" + enhance_upscale: int = max(1, min(4, int(_env("MUSE_ENHANCE_UPSCALE", "1")))) + # GFPGAN 模型权重路径（相对 MUSE_DIR 或绝对路径） + gfpgan_model_path: str = _env( + "MUSE_GFPGAN_MODEL", "models/GFPGAN/GFPGANv1.4.pth" + ) # ── 全局状态 ────────────────────────────────────────────────────────── @@ -100,6 +111,12 @@ _muse_models_lock = threading.Lock() _muse_models_loaded = False _muse_load_error = None +# ── GFPGAN 面部增强模型懒加载 ─────────────────────────────────────── +_gfpgan = None +_gfpgan_lock = threading.Lock() +_gfpgan_loaded = False +_gfpgan_load_error = None + # ── Flask App ───────────────────────────────────────────────────────── app = Flask(__name__) @@ -247,21 +264,38 @@ def _preprocess_audio(input_path: Path, output_path: Path, target_sr: int = 1600 MuseTalk 的 whisper audio2feature 要求 16kHz 采样率的单声道音频。 当前 TTS 输出为 22050Hz MP3，不转换会导致 mel 频谱错位、 音素特征提取错误，口型只跟能量不跟音素。 + + 使用 soxr 高质量重采样器（如果可用），避免线性插值带来的频谱失真。 """ + # 尝试使用 soxr 高质量重采样；不可用时回退默认 resampler cmd = [ "ffmpeg", "-y", "-v", "warning", "-i", str(input_path), + "-af", "aresample=resampler=soxr:precision=28", # 高质量重采样 "-ar", str(target_sr), # 重采样到 16kHz "-ac", "1", # 单声道 "-sample_fmt", "s16", # 16bit PCM str(output_path), ] - _run_ffmpeg(cmd, timeout=60) + try: + _run_ffmpeg(cmd, timeout=60) + except RuntimeError: + # soxr 不可用时回退到默认重采样器 + logger.warning("soxr 重采样器不可用，回退默认重采样器") + cmd_fallback = [ + "ffmpeg", "-y", "-v", "warning", + "-i", str(input_path), + "-ar", str(target_sr), + "-ac", "1", + "-sample_fmt", "s16", + str(output_path), + ] + _run_ffmpeg(cmd_fallback, 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) + logger.info("音频预处理完成: %s → 16kHz mono WAV（高质量重采样）", input_path.name) def _mux_video_with_audio( @@ -505,6 +539,193 @@ def _load_musetalk_models(): raise +# ── GFPGAN 面部增强 ─────────────────────────────────────────────────── + + +def _load_gfpgan(): + """懒加载 GFPGAN 面部增强模型（全局单例，首次调用时加载）. + + GFPGAN 用于提升 MuseTalk 输出的面部清晰度和纹理细节， + 特别改善嘴部区域模糊问题。 + 优先从 MuseTalk 自带的 models/GFPGAN 目录加载权重， + 也支持通过 MUSE_GFPGAN_MODEL 环境变量指定绝对路径。 + """ + global _gfpgan, _gfpgan_loaded, _gfpgan_load_error + + if _gfpgan_loaded: + return _gfpgan + if _gfpgan_load_error is not None: + raise _gfpgan_load_error + + with _gfpgan_lock: + if _gfpgan_loaded: + return _gfpgan + + try: + import torch + + # 解析模型权重路径 + model_path = Config.gfpgan_model_path + if not os.path.isabs(model_path): + model_path = os.path.join(Config.muse_dir, model_path) + + if not os.path.isfile(model_path): + raise FileNotFoundError( + f"GFPGAN 模型权重不存在: {model_path}\n" + f"请下载到 MuseTalk 仓库的 models/GFPGAN/ 目录下，\n" + f"或通过 MUSE_GFPGAN_MODEL 环境变量指定路径。\n" + f"下载地址: https://github.com/TencentARC/GFPGAN/releases" + ) + + # 尝试多种导入方式 + gfpgan_model = None + device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + + # 方式 1: 从 gfpgan 官方包加载 + try: + from gfpgan import GFPGANer + gfpgan_model = GFPGANer( + model_path=model_path, + upscale=Config.enhance_upscale, + arch="clean", + channel_multiplier=2, + device=device, + ) + logger.info("GFPGAN 加载成功（gfpgan 官方包）") + except ImportError: + pass + + # 方式 2: 从 MuseTalk 自带的 facelib/gfpgan 加载 + if gfpgan_model is None: + muse_str = str(Path(Config.muse_dir)) + if muse_str not in sys.path: + sys.path.insert(0, muse_str) + try: + from gfpgan import GFPGANer + gfpgan_model = GFPGANer( + model_path=model_path, + upscale=Config.enhance_upscale, + arch="clean", + channel_multiplier=2, + device=device, + ) + logger.info("GFPGAN 加载成功（MuseTalk 环境）") + except ImportError: + pass + + if gfpgan_model is None: + raise ImportError( + "无法导入 GFPGAN。请先安装: pip install gfpgan\n" + "并确保模型权重文件存在于指定路径。" + ) + + _gfpgan = { + "model": gfpgan_model, + "device": device, + } + _gfpgan_loaded = True + logger.info("GFPGAN 面部增强模型加载完成 (device=%s)", device) + + return _gfpgan + + except Exception as exc: + _gfpgan_load_error = exc + logger.error("GFPGAN 模型加载失败: %s", exc) + raise + + +def _enhance_face(frame, bbox, enhance_region="mouth", padding_ratio=0.3): + """对帧中的人脸/嘴部区域应用 GFPGAN 增强. + + Args: + frame: BGR numpy 数组 + bbox: [x1, y1, x2, y2] 嘴部/面部区域 + enhance_region: 'mouth' 仅增强嘴部区域（更快）, 'face' 增强全脸（更自然） + padding_ratio: 区域扩展比例，用于提供足够的上下文给 GFPGAN + + Returns: + 增强后的帧 (BGR numpy 数组) + """ + import cv2 + import numpy as np + + models = _load_gfpgan() + gfpgan_model = models["model"] + + x1, y1, x2, y2 = bbox + h, w = frame.shape[:2] + + if enhance_region == "mouth": + # 嘴部区域：扩展 bbox 提供足够上下文 + pad_x = int((x2 - x1) * padding_ratio) + pad_y = int((y2 - y1) * padding_ratio * 1.5) # 嘴部上下多留一些空间 + crop_x1 = max(0, x1 - pad_x) + crop_y1 = max(0, y1 - pad_y) + crop_x2 = min(w, x2 + pad_x) + crop_y2 = min(h, y2 + pad_y) + else: + # 全脸区域：使用更大的 padding + face_w = x2 - x1 + face_h = y2 - y1 + pad_x = int(face_w * 0.5) + pad_y = int(face_h * 0.6) + crop_x1 = max(0, x1 - pad_x) + crop_y1 = max(0, y1 - pad_y) + crop_x2 = min(w, x2 + pad_x) + crop_y2 = min(h, y2 + pad_y) + + # 裁剪区域 + crop = frame[crop_y1:crop_y2, crop_x1:crop_x2].copy() + if crop.size == 0: + return frame + + # 确保最小尺寸（GFPGAN 需要至少 64x64） + crop_h, crop_w = crop.shape[:2] + if crop_h < 64 or crop_w < 64: + return frame + + try: + # GFPGAN 增强 + _, _, enhanced_crop = gfpgan_model.enhance( + crop, + has_aligned=False, + only_center_face=False, + paste_back=True, + ) + + # 确保尺寸一致（upscale 可能导致尺寸变化） + if enhanced_crop.shape[:2] != crop.shape[:2]: + enhanced_crop = cv2.resize(enhanced_crop, (crop_w, crop_h)) + + # 创建高斯模糊遮罩用于平滑融合 + mask_h, mask_w = enhanced_crop.shape[:2] + # 使用较大的模糊核以获得更柔和的过渡 + ksize = max(mask_h, mask_w) // 3 + if ksize % 2 == 0: + ksize += 1 + ksize = max(ksize, 11) # 最小核大小 + mask = np.zeros((mask_h, mask_w), dtype=np.float32) + # 中心区域为 1 + margin_y = int(mask_h * 0.15) + margin_x = int(mask_w * 0.15) + mask[margin_y:mask_h-margin_y, margin_x:mask_w-margin_x] = 1.0 + mask = cv2.GaussianBlur(mask, (ksize, ksize), 0) + + # 融合：增强区域与原图平滑过渡 + mask_3ch = mask[:, :, np.newaxis] + blended_crop = (enhanced_crop.astype(np.float32) * mask_3ch + + crop.astype(np.float32) * (1 - mask_3ch)) + blended_crop = np.clip(blended_crop, 0, 255).astype(np.uint8) + + # 放回原帧 + frame[crop_y1:crop_y2, crop_x1:crop_x2] = blended_crop + + except Exception as e: + logger.warning("GFPGAN 面部增强失败，使用原帧: %s", e) + + return frame + + # ── MuseTalk 推理核心 ───────────────────────────────────────────────── @@ -529,23 +750,27 @@ def _run_inference( audio_path: Path, output_path: Path, bbox_shift: int = 0, + face_enhance: bool = None, ) -> None: """执行 MuseTalk 真实推理. 流程： - 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 音频） + 1. 音频预处理：任意格式 → 16kHz mono 16bit WAV（高质量 soxr 重采样） + 2. 帧率适配：高帧率视频（>30fps）降至 25fps 推理，输出时恢复原始帧率 + 3. 加载/复用 MuseTalk 模型（VAE + UNet + PE + Whisper） + 4. 视频预处理：提取帧 → 人脸检测 → 获取 bbox → VAE 编码 latent + 5. 音频特征提取：whisper 提取 audio features (50×384 per chunk) + 6. 批量推理：UNet 去噪 → VAE 解码 → 得到口型同步的人脸帧 + 7. 帧合成：将生成的人脸贴回原帧（使用 face parsing + 高斯模糊边缘融合） + 8. GFPGAN 面部增强后处理（可选）：提升嘴部/面部区域的清晰度和纹理 + 9. 输出无声视频（后续由 _mux_video_with_audio 封装 TTS 音频） Args: video_path: 输入视频路径 audio_path: 输入音频路径（任意格式，会被预处理为 16kHz WAV） output_path: 输出无声视频路径 bbox_shift: 口型区域垂直偏移量，默认 0，范围 [-5, 5] + face_enhance: 是否启用 GFPGAN 面部增强；None 时使用 Config.face_enhance 配置 """ import cv2 import numpy as np @@ -556,6 +781,10 @@ def _run_inference( from musetalk.utils.utils import datagen from musetalk.utils.blending import get_image + # 是否启用面部增强 + if face_enhance is None: + face_enhance = Config.face_enhance + # 加载模型（首次调用时加载，后续复用） models = _load_musetalk_models() vae = models["vae"] @@ -569,9 +798,43 @@ def _run_inference( fps = _get_video_fps(video_path) audio_duration = _get_media_duration(audio_path) video_duration = _get_media_duration(video_path) + + # ── 帧率适配：MuseTalk 训练在 25fps，高帧率输入需降采样 ── + original_fps = fps + inference_fps = 25.0 # MuseTalk 标准帧率 + need_fps_conversion = fps > 30 # 超过 30fps 的输入需要处理 + + if need_fps_conversion: + logger.info( + "检测到高帧率输入 (%.1f fps)，将先降至 %.1f fps 推理，推理后恢复至原始帧率", + fps, inference_fps, + ) + # 使用 ffmpeg 将高帧率视频降至 25fps + downsampled_video = video_path.parent / "video_25fps.mp4" + _run_ffmpeg( + [ + "ffmpeg", "-y", "-v", "warning", + "-i", str(video_path), + "-r", str(int(inference_fps)), + "-vsync", "cfr", + "-c:v", "libx264", + "-preset", "fast", + "-crf", "18", + "-an", # 去除音频，后续单独处理 + str(downsampled_video), + ], + timeout=300, + ) + work_video_path = downsampled_video + fps = inference_fps + logger.info("视频降采样至 25fps 完成") + else: + work_video_path = video_path + logger.info( - "MuseTalk 推理开始: video=%.2fs, audio=%.2fs, fps=%.1f, bbox_shift=%d", - video_duration, audio_duration, fps, bbox_shift, + "MuseTalk 推理开始: video=%.2fs, audio=%.2fs, original_fps=%.1f, " + "inference_fps=%.1f, bbox_shift=%d, face_enhance=%s", + video_duration, audio_duration, original_fps, fps, bbox_shift, face_enhance, ) # ── Step 1: 音频预处理（关键修复：22050Hz MP3 → 16kHz mono WAV）── @@ -579,17 +842,18 @@ def _run_inference( _preprocess_audio(audio_path, audio_wav_path, target_sr=16000) # ── Step 2: 视频帧提取 ── - input_frames = read_imgs(str(video_path)) + input_frames = read_imgs(str(work_video_path)) total_frames = len(input_frames) if total_frames == 0: raise RuntimeError("未能从视频中提取到任何帧") - logger.info("提取到 %d 帧视频画面", total_frames) + logger.info("提取到 %d 帧视频画面 (工作帧率 %.1f fps)", total_frames, fps) # ── 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) + valid_bboxes = sum(1 for c in coord_list if c is not coord_placeholder) + logger.info("人脸检测完成，有效 bbox: %d/%d", valid_bboxes, total_frames) # 使用 mirror indexing 循环帧和坐标（避免硬切跳变） num_output_frames = int(audio_duration * fps) @@ -597,7 +861,6 @@ def _run_inference( num_output_frames = total_frames # ── Step 4: 音频特征提取 ── - # 使用 librosa 加载预处理后的 16kHz 音频 import librosa audio_array, _ = librosa.load(str(audio_wav_path), sr=16000, mono=True) whisper_features = audio_processor.feature2chunks( @@ -670,6 +933,15 @@ def _run_inference( output_frames_dir = video_path.parent / "output_frames" output_frames_dir.mkdir(parents=True, exist_ok=True) + # 预加载 GFPGAN（如果启用面部增强） + if face_enhance: + try: + _load_gfpgan() + logger.info("GFPGAN 面部增强已启用，区域: %s", Config.enhance_region) + except Exception as e: + logger.warning("GFPGAN 加载失败，跳过面部增强: %s", e) + face_enhance = False + for i, res_frame in enumerate(tqdm(result_frames[:num_output_frames], desc="合成帧")): orig_idx = _mirror_index(total_frames, i) ori_frame = copy.deepcopy(input_frames[orig_idx]) @@ -681,8 +953,11 @@ def _run_inference( else: x1, y1, x2, y2 = bbox if model_version == "v15": - # v1.5 额外扩展下边界（下巴区域） - y2 = min(y2 + 10, ori_frame.shape[0]) + # v1.5 额外扩展下边界（下巴区域），扩大融合范围 + y2 = min(y2 + 15, ori_frame.shape[0]) + x1 = max(0, x1 - 2) + x2 = min(ori_frame.shape[1], x2 + 2) + try: res_frame_resized = cv2.resize( res_frame.astype(np.uint8), (x2 - x1, y2 - y1) @@ -701,23 +976,53 @@ def _run_inference( [x1, y1, x2, y2], ) + # ── P0: GFPGAN 面部增强后处理 ── + if face_enhance: + try: + combined = _enhance_face( + combined, + [x1, y1, x2, y2], + enhance_region=Config.enhance_region, + ) + except Exception as e: + logger.warning("第 %d 帧面部增强失败，使用未增强结果: %s", i, e) + cv2.imwrite(str(output_frames_dir / f"{i:08d}.png"), combined) - # 帧序列 → 无声视频 + # ── Step 8: 帧序列 → 无声视频 ── silent_video_path = video_path.parent / "silent_output.mp4" - _run_ffmpeg( - [ - "ffmpeg", "-y", "-v", "warning", - "-r", str(fps), - "-f", "image2", - "-i", str(output_frames_dir / "%08d.png"), - "-vcodec", "libx264", - "-vf", "format=yuv420p", - "-crf", "18", - str(silent_video_path), - ], - timeout=300, - ) + + if need_fps_conversion: + # 需要恢复原始帧率：使用 minterpolate 做运动补偿插帧 + logger.info("将推理结果从 25fps 恢复至 %.1f fps（运动补偿插帧）", original_fps) + _run_ffmpeg( + [ + "ffmpeg", "-y", "-v", "warning", + "-r", str(int(inference_fps)), + "-f", "image2", + "-i", str(output_frames_dir / "%08d.png"), + "-vf", f"minterpolate=fps={int(original_fps)}:mi_mode=mci:mc_mode=aobmc:me_mode=bidir:vsbmc=1,format=yuv420p", + "-c:v", "libx264", + "-preset", "medium", + "-crf", "18", + str(silent_video_path), + ], + timeout=600, + ) + else: + _run_ffmpeg( + [ + "ffmpeg", "-y", "-v", "warning", + "-r", str(int(original_fps)), + "-f", "image2", + "-i", str(output_frames_dir / "%08d.png"), + "-vcodec", "libx264", + "-vf", "format=yuv420p", + "-crf", "18", + str(silent_video_path), + ], + timeout=300, + ) # 将无声视频复制到输出路径 shutil.copy2(str(silent_video_path), str(output_path)) @@ -729,11 +1034,17 @@ def _run_inference( audio_wav_path.unlink() if silent_video_path.exists() and str(silent_video_path) != str(output_path): silent_video_path.unlink() + if need_fps_conversion: + downsampled_video = video_path.parent / "video_25fps.mp4" + if downsampled_video.exists(): + downsampled_video.unlink() except Exception as e: logger.warning("清理中间文件失败: %s", e) - logger.info("MuseTalk 推理完成: output=%s, duration=%.2fs", - output_path.name, _get_media_duration(output_path)) + logger.info( + "MuseTalk 推理完成: output=%s, duration=%.2fs, fps=%.1f, face_enhance=%s", + output_path.name, _get_media_duration(output_path), original_fps, face_enhance, + ) # ── 路由 ────────────────────────────────────────────────────────────── @@ -741,7 +1052,7 @@ def _run_inference( @app.route("/health", methods=["GET"]) def health(): - """健康检查 + GPU 显存信息 + MuseTalk 模型状态.""" + """健康检查 + GPU 显存信息 + MuseTalk/GFPGAN 模型状态.""" gpu_info = _get_gpu_info() task_info = { "task_id": current_task["task_id"], @@ -755,6 +1066,9 @@ def health(): "current_task": task_info, "musetalk_loaded": _muse_models_loaded, "musetalk_load_error": str(_muse_load_error) if _muse_load_error else None, + "gfpgan_loaded": _gfpgan_loaded, + "gfpgan_load_error": str(_gfpgan_load_error) if _gfpgan_load_error else None, + "face_enhance_enabled": Config.face_enhance, "timestamp": time.time(), } ) @@ -766,6 +1080,7 @@ def inference(): 可选 form 参数： bbox_shift: 口型区域垂直偏移量，默认 0，范围 [-5, 5] + face_enhance: 是否启用 GFPGAN 面部增强后处理，"0" 或 "1"，默认跟随配置 """ # 并发控制：检查锁 if not inference_lock.acquire(blocking=False): @@ -787,6 +1102,13 @@ def inference(): bbox_shift = int(request.form.get("bbox_shift", "0")) bbox_shift = max(-5, min(5, bbox_shift)) # 限制范围 + # 解析 face_enhance 参数：支持请求级覆盖 + face_enhance_param = request.form.get("face_enhance", None) + if face_enhance_param is not None: + face_enhance = face_enhance_param.strip() in ("1", "true", "yes", "on") + else: + face_enhance = None # 跟随 Config 默认配置 + # 文件大小检查 err = _check_file_size(video_file, Config.video_max_mb, "视频") if err: @@ -805,8 +1127,8 @@ def inference(): video_file.save(str(video_path)) audio_file.save(str(audio_path)) - logger.info("开始推理 task_id=%s, video=%s, audio=%s, bbox_shift=%d", - task_id, video_path.name, audio_path.name, bbox_shift) + logger.info("开始推理 task_id=%s, video=%s, audio=%s, bbox_shift=%d, face_enhance=%s", + task_id, video_path.name, audio_path.name, bbox_shift, face_enhance) # 更新当前任务信息 current_task["task_id"] = task_id @@ -818,7 +1140,11 @@ def inference(): def inference_thread(): try: - _run_inference(video_path, audio_path, output_path, bbox_shift=bbox_shift) + _run_inference( + video_path, audio_path, output_path, + bbox_shift=bbox_shift, + face_enhance=face_enhance, + ) except Exception as exc: logger.exception("推理异常: %s", exc) result_container["error"] = str(exc)