xiaoxia-saas/deploy/gpu_worker/musetalk_server.py

"""MuseTalk Flask HTTP 服务 — 反向轮询架构的服务端部分.

部署在 RTX2060 本地，接收 gpu_worker.py 的推理请求，调用 MuseTalk 生成口型同步视频。
本文件修复了原 worker.py 的 8 个工程 bug，并新增 /cancel 端点。

#1978 性能修复（v2 架构）：
  MuseTalk 原生支持长音频输入（内部循环视频帧），不需要我们先 loop 视频。
  正确流程：原视频 + 全量音频 → MuseTalk 推理 → 输出时长=音频时长的无声画面
  → ffmpeg 快速 -c:v copy 替换音轨。推理时间不变（~14s），后处理几秒。
  禁止在推理前用 ffmpeg 循环视频（会导致 MuseTalk 处理 2x+ 帧数，慢 16 倍）。

环境变量：
  MUSE_PORT            监听端口，默认 7861
  MUSE_MAX_CONCURRENT  最大并发推理数，默认 1（GPU 一次只能处理一个）
  MUSE_INFERENCE_TIMEOUT 推理超时秒数，默认 600
  MUSE_VIDEO_MAX_MB    视频上传大小限制 MB，默认 100
  MUSE_AUDIO_MAX_MB    音频上传大小限制 MB，默认 20
  MUSE_DEFAULT_FPS     视频 fps 兜底值，默认 25.0
  MUSE_TEMP_DIR        临时文件目录，默认 /tmp/musetalk_$$
  MUSE_VIDEO_ENCODER   循环视频时的编码器（仅兜底）：auto（默认）/h264_nvenc/libx264

接口：
  GET  /health          健康检查 + GPU 显存信息
  POST /inference       推理请求（multipart: video + audio）
  POST /cancel          终止当前推理任务
"""

from __future__ import annotations

import atexit
import logging
import os
import shutil
import signal
import subprocess
import threading
import time
from pathlib import Path
from typing import Optional

from flask import Flask, jsonify, request, send_file

# ── 日志 ──────────────────────────────────────────────────────────────
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
logger = logging.getLogger("musetalk-server")

# ── 配置 ──────────────────────────────────────────────────────────────


def _env(name: str, default: str = "") -> str:
    v = os.environ.get(name, default)
    return v.strip() if isinstance(v, str) else default


class Config:
    port: int = int(_env("MUSE_PORT", "7861"))
    max_concurrent: int = int(_env("MUSE_MAX_CONCURRENT", "1"))
    inference_timeout: float = float(_env("MUSE_INFERENCE_TIMEOUT", "600"))
    video_max_mb: int = int(_env("MUSE_VIDEO_MAX_MB", "100"))
    audio_max_mb: int = int(_env("MUSE_AUDIO_MAX_MB", "20"))
    default_fps: float = float(_env("MUSE_DEFAULT_FPS", "25.0"))
    temp_dir: str = _env("MUSE_TEMP_DIR", f"/tmp/musetalk_{os.getpid()}")
    # 循环视频时的编码器（仅当 MuseTalk 输出画面短于音频时的兜底）
    video_encoder: str = _env("MUSE_VIDEO_ENCODER", "auto") or "auto"
    # 判定音视频时长差异的容差（秒）
    duration_epsilon: float = 0.25


# ── 全局状态 ──────────────────────────────────────────────────────────
inference_lock = threading.Lock()
current_task: dict = {"task_id": None, "process": None, "start_time": 0.0}
shutdown_event = threading.Event()

# ── Flask App ─────────────────────────────────────────────────────────
app = Flask(__name__)


def _cleanup_temp_dir():
    """退出时清理临时目录."""
    if os.path.exists(Config.temp_dir):
        try:
            shutil.rmtree(Config.temp_dir)
            logger.info("已清理临时目录: %s", Config.temp_dir)
        except Exception as exc:
            logger.warning("清理临时目录失败: %s", exc)


atexit.register(_cleanup_temp_dir)


def _signal_handler(signum, frame):
    """优雅退出."""
    logger.info("收到信号 %s，准备退出...", signum)
    shutdown_event.set()
    if current_task["process"]:
        logger.info("终止正在进行的推理进程...")
        try:
            current_task["process"].terminate()
            current_task["process"].wait(timeout=5)
        except Exception:
            pass
    _cleanup_temp_dir()
    exit(0)


signal.signal(signal.SIGTERM, _signal_handler)
signal.signal(signal.SIGINT, _signal_handler)


# ── 工具函数 ──────────────────────────────────────────────────────────


def _get_gpu_info() -> dict:
    """获取 GPU 显存信息（通过 nvidia-smi）."""
    try:
        out = subprocess.check_output(
            [
                "nvidia-smi",
                "--query-gpu=name,memory.total,memory.used,memory.free",
                "--format=csv,noheader,nounits",
            ],
            stderr=subprocess.DEVNULL,
            timeout=5,
        )
        parts = out.decode().strip().split(",")
        if len(parts) >= 4:
            return {
                "gpu_name": parts[0].strip(),
                "memory_total_mb": int(parts[1].strip()),
                "memory_used_mb": int(parts[2].strip()),
                "memory_free_mb": int(parts[3].strip()),
            }
    except Exception as exc:
        logger.warning("nvidia-smi 失败: %s", exc)
    return {"gpu_name": "unknown", "memory_total_mb": 0, "memory_used_mb": 0, "memory_free_mb": 0}


def _get_video_fps(video_path: Path) -> float:
    """用 ffprobe 读视频帧率，失败或为 0 时返回 default_fps."""
    try:
        out = subprocess.check_output(
            [
                "ffprobe",
                "-v",
                "error",
                "-select_streams",
                "v:0",
                "-show_entries",
                "stream=r_frame_rate",
                "-of",
                "default=noprint_wrappers=1:nokey=1",
                str(video_path),
            ],
            stderr=subprocess.DEVNULL,
            timeout=10,
        )
        fps_str = out.decode().strip()
        if "/" in fps_str:
            num, den = fps_str.split("/")
            fps = float(num) / float(den) if float(den) != 0 else 0.0
        else:
            fps = float(fps_str) if fps_str else 0.0
        return fps if fps > 0 else Config.default_fps
    except Exception as exc:
        logger.warning("ffprobe 读 fps 失败: %s，使用默认 %.1f", exc, Config.default_fps)
        return Config.default_fps


def _get_media_duration(path: Path) -> float:
    """用 ffprobe 读媒体时长（秒），失败返回 0.0."""
    try:
        out = subprocess.check_output(
            [
                "ffprobe",
                "-v",
                "error",
                "-show_entries",
                "format=duration",
                "-of",
                "default=noprint_wrappers=1:nokey=1",
                str(path),
            ],
            stderr=subprocess.DEVNULL,
            timeout=10,
        )
        duration = float(out.decode().strip())
        return duration if duration > 0 else 0.0
    except Exception as exc:
        logger.warning("ffprobe 读时长失败 %s: %s", path, exc)
        return 0.0


def _pick_video_encoder() -> str:
    """选择视频编码器：配置指定则用指定值；auto 时探测 NVENC 是否可用，不可用回退 libx264."""
    configured = Config.video_encoder.strip()
    if configured in ("h264_nvenc", "libx264"):
        return configured
    # auto：探测本机 ffmpeg 是否编译了 h264_nvenc
    try:
        result = subprocess.run(
            ["ffmpeg", "-hide_banner", "-encoders"],
            stdout=subprocess.PIPE,
            stderr=subprocess.DEVNULL,
            timeout=10,
            check=False,
        )
        if b"h264_nvenc" in result.stdout:
            return "h264_nvenc"
    except Exception as exc:
        logger.warning("探测 ffmpeg 编码器失败，回退 libx264: %s", exc)
    return "libx264"


def _mux_video_with_audio(
    video_path: Path,
    audio_path: Path,
    output_path: Path,
    timeout: float = 300,
) -> None:
    """把无声画面视频与驱动音频封装为最终结果.

    #1978 v2 架构：MuseTalk 已处理全量音频，输出视频时长=音频时长。
    此处仅做快速封装：-map 0:v:0 -map 1:a:0 强制取画面+驱动音频，
    -c:v copy 无损秒级封装（不重编码），-shortest 以较短流为准。

    仅当 MuseTalk 输出画面短于音频时（极端兜底），才启用 -stream_loop + NVENC
    循环视频到音频长度。正常情况下走 copy 快速路径。
    """
    video_duration = _get_media_duration(video_path)
    audio_duration = _get_media_duration(audio_path)

    # 判断是否需要兜底循环（正常情况下 MuseTalk 输出已 >= 音频时长）
    need_loop_fallback = bool(
        audio_duration > 0 and video_duration > 0 and video_duration < audio_duration - Config.duration_epsilon
    )

    if need_loop_fallback:
        # 兜底：MuseTalk 输出画面不足，循环补齐
        encoder = _pick_video_encoder()
        preset = "p4" if encoder == "h264_nvenc" else "veryfast"
        logger.warning(
            "MuseTalk 输出(%.2fs)短于音频(%.2fs)，兜底循环视频以 %s 重编码",
            video_duration,
            audio_duration,
            encoder,
        )

        def build_cmd(enc: str, pre: str) -> list:
            return [
                "ffmpeg",
                "-y",
                "-stream_loop",
                "-1",
                "-i",
                str(video_path),
                "-i",
                str(audio_path),
                "-map",
                "0:v:0",
                "-map",
                "1:a:0",
                "-c:v",
                enc,
                "-preset",
                pre,
                "-c:a",
                "aac",
                "-b:a",
                "128k",
                "-t",
                f"{audio_duration:.3f}",
                str(output_path),
            ]

        try:
            _run_ffmpeg(build_cmd(encoder, preset), timeout=timeout)
        except RuntimeError:
            if encoder == "h264_nvenc":
                logger.warning("h264_nvenc 兜底失败，回退 libx264 重试")
                _run_ffmpeg(build_cmd("libx264", "veryfast"), timeout=timeout)
            else:
                raise
    else:
        # 正常快速路径：-c:v copy 无损封装，仅替换音轨为驱动音频
        cmd = [
            "ffmpeg",
            "-y",
            "-i",
            str(video_path),
            "-i",
            str(audio_path),
            "-map",
            "0:v:0",
            "-map",
            "1:a:0",
            "-c:v",
            "copy",
            "-c:a",
            "aac",
            "-b:a",
            "128k",
            "-shortest",
            str(output_path),
        ]
        _run_ffmpeg(cmd, timeout=timeout)


def _check_file_size(file, max_mb: int, label: str) -> Optional[str]:
    """检查文件大小，超限返回错误信息，否则返回 None."""
    file.seek(0, 2)
    size = file.tell()
    file.seek(0)
    max_bytes = max_mb * 1024 * 1024
    if size > max_bytes:
        return f"{label} 文件大小 {size / (1024*1024):.1f}MB 超过限制 {max_mb}MB"
    if size == 0:
        return f"{label} 文件为空"
    return None


def _run_ffmpeg(cmd: list, timeout: float = 120) -> subprocess.CompletedProcess:
    """运行 ffmpeg 命令，检查返回码和超时."""
    try:
        result = subprocess.run(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            timeout=timeout,
            check=True,
        )
        return result
    except subprocess.CalledProcessError as exc:
        stderr = exc.stderr.decode(errors="ignore") if exc.stderr else ""
        raise RuntimeError(f"ffmpeg 失败 (code={exc.returncode}): {stderr[:500]}") from exc
    except subprocess.TimeoutExpired as exc:
        raise RuntimeError(f"ffmpeg 超时（>{timeout}s）") from exc


def _run_inference(
    video_path: Path,
    audio_path: Path,
    output_path: Path,
) -> None:
    """执行 MuseTalk 推理（v2 架构：全量音频直传，不在推理前 loop 视频）.

    #1978 性能修复核心：
      MuseTalk 原生支持长音频输入，内部会自动循环视频帧。
      我们只需把【原视频】和【全量音频】传给 MuseTalk，
      输出视频时长 = 音频时长（MuseTalk 自行处理帧循环）。
      禁止在推理前用 ffmpeg 循环视频（会导致慢 16 倍）。

    实际部署时替换为 MuseTalk 真实推理逻辑。
    此处为示例实现：提取帧 → 模拟 MuseTalk 产出音频时长的无声画面 → 快速封装。
    """
    fps = _get_video_fps(video_path)
    audio_duration = _get_media_duration(audio_path)
    video_duration = _get_media_duration(video_path)
    logger.info(
        "推理开始: video=%.2fs, audio=%.2fs, fps=%.2f",
        video_duration,
        audio_duration,
        fps,
    )

    frames_dir = video_path.parent / "frames"
    frames_dir.mkdir(parents=True, exist_ok=True)

    # 1. 从原视频提取帧（仅原视频长度，不循环）
    _run_ffmpeg(
        [
            "ffmpeg",
            "-y",
            "-i",
            str(video_path),
            "-r",
            str(fps),
            str(frames_dir / "frame_%05d.png"),
        ],
        timeout=120,
    )

    frame_files = sorted(frames_dir.glob("*.png"))
    if not frame_files:
        raise RuntimeError("未从视频中提取到帧")

    # 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(
            "示例：循环视频帧到音频长度 %.2fs（真实 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=%.2fs (audio=%.2fs)",
        _get_media_duration(output_path),
        audio_duration,
    )


# ── 路由 ──────────────────────────────────────────────────────────────


@app.route("/health", methods=["GET"])
def health():
    """健康检查 + GPU 显存信息."""
    gpu_info = _get_gpu_info()
    task_info = {
        "task_id": current_task["task_id"],
        "running": current_task["process"] is not None,
        "elapsed_seconds": time.time() - current_task["start_time"] if current_task["start_time"] else 0.0,
    }
    return jsonify(
        {
            "status": "healthy",
            "gpu": gpu_info,
            "current_task": task_info,
            "timestamp": time.time(),
        }
    )


@app.route("/inference", methods=["POST"])
def inference():
    """推理请求：multipart form 包含 video 和 audio 文件.

    #1978 v2：MuseTalk 直接处理全量音频，输出时长=音频时长，无需预处理循环。
    """
    # 并发控制：检查锁
    if not inference_lock.acquire(blocking=False):
        return jsonify({"error": "GPU 正在处理其他任务，请稍后重试", "status": "busy"}), 503

    task_id = None
    video_path = None
    audio_path = None
    output_path = None

    try:
        # 解析参数
        if "video" not in request.files or "audio" not in request.files:
            return jsonify({"error": "缺少 video 或 audio 文件"}), 400

        video_file = request.files["video"]
        audio_file = request.files["audio"]
        task_id = request.form.get("task_id", f"task_{int(time.time())}")

        # 文件大小检查
        err = _check_file_size(video_file, Config.video_max_mb, "视频")
        if err:
            return jsonify({"error": err}), 413
        err = _check_file_size(audio_file, Config.audio_max_mb, "音频")
        if err:
            return jsonify({"error": err}), 413

        # 保存到临时目录
        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"
        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)

        # 更新当前任务信息
        current_task["task_id"] = task_id
        current_task["start_time"] = time.time()
        current_task["process"] = "inference_thread"  # 标记为运行中

        # 在线程中运行推理（支持超时）
        result_container = {"error": None}

        def inference_thread():
            try:
                _run_inference(video_path, audio_path, output_path)
            except Exception as exc:
                result_container["error"] = str(exc)

        thread = threading.Thread(target=inference_thread)
        thread.start()
        thread.join(timeout=Config.inference_timeout)

        if thread.is_alive():
            # 超时，终止
            logger.error("推理超时 (>%ds)，终止任务 %s", Config.inference_timeout, task_id)
            return jsonify({"error": f"推理超时（>{Config.inference_timeout}s）", "task_id": task_id}), 504

        if result_container["error"]:
            logger.error("推理失败 task_id=%s: %s", task_id, result_container["error"])
            return jsonify({"error": result_container["error"], "task_id": task_id}), 500

        # 返回结果文件
        logger.info("推理完成 task_id=%s, output=%s", task_id, output_path)
        return send_file(str(output_path), mimetype="video/mp4", as_attachment=True, download_name=f"{task_id}.mp4")

    except Exception as exc:
        logger.exception("推理异常: %s", exc)
        return jsonify({"error": str(exc)}), 500

    finally:
        # 释放锁，清理当前任务信息
        inference_lock.release()
        current_task["task_id"] = None
        current_task["process"] = None
        current_task["start_time"] = 0.0

        # 清理临时文件
        if video_path and video_path.parent.exists():
            try:
                shutil.rmtree(video_path.parent)
                logger.info("已清理临时目录: %s", video_path.parent)
            except Exception as exc:
                logger.warning("清理临时目录失败: %s", exc)


@app.route("/cancel", methods=["POST"])
def cancel():
    """终止当前正在进行的推理任务."""
    if current_task["task_id"] is None:
        return jsonify({"message": "当前无正在运行的任务"})

    task_id = current_task["task_id"]
    logger.info("收到取消请求，终止任务 %s", task_id)

    # 终止推理进程（如果是 subprocess）
    if current_task["process"] and current_task["process"] != "inference_thread":
        try:
            current_task["process"].terminate()
            current_task["process"].wait(timeout=5)
            logger.info("已终止推理进程")
        except Exception as exc:
            logger.warning("终止进程失败: %s", exc)

    # 清理临时文件
    task_dir = Path(Config.temp_dir) / task_id
    if task_dir.exists():
        try:
            shutil.rmtree(task_dir)
            logger.info("已清理临时目录: %s", task_dir)
        except Exception as exc:
            logger.warning("清理临时目录失败: %s", exc)

    # 重置当前任务
    current_task["task_id"] = None
    current_task["process"] = None
    current_task["start_time"] = 0.0

    return jsonify({"message": f"已取消任务 {task_id}"})


# ── 主入口 ────────────────────────────────────────────────────────────


def main():
    """启动 Flask 服务."""
    # 创建临时目录
    Path(Config.temp_dir).mkdir(parents=True, exist_ok=True)
    logger.info("临时目录: %s", Config.temp_dir)

    gpu_info = _get_gpu_info()
    logger.info(
        "GPU: %s (显存 %dMB / %dMB)",
        gpu_info["gpu_name"],
        gpu_info["memory_used_mb"],
        gpu_info["memory_total_mb"],
    )
    logger.info(
        "启动 MuseTalk Server: port=%d, timeout=%.0fs, max_concurrent=%d",
        Config.port,
        Config.inference_timeout,
        Config.max_concurrent,
    )

    app.run(host="0.0.0.0", port=Config.port, threaded=True)


if __name__ == "__main__":
    main()