xiaoxia-saas/deploy/gpu_worker/musetalk_server.py

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

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

环境变量：
  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）/h264_nvenc/libx264
  MUSE_ENABLE_VIDEO_LOOP 驱动音频比视频长时是否循环视频补齐，默认 1（开启）

接口：
  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()}")
    # 循环视频时编码器：auto 优先 h264_nvenc（RTX2060 支持），失败兜底 libx264
    video_encoder: str = _env("MUSE_VIDEO_ENCODER", "auto") or "auto"
    # 驱动音频比视频长时循环视频补齐画面
    enable_video_loop: bool = _env("MUSE_ENABLE_VIDEO_LOOP", "1") not in ("0", "false", "False", "")
    # 判定音视频时长差异的容差（秒），避免 ffprobe 微小误差触发无谓的循环/重编码
    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,
    enable_video_loop: Optional[bool] = None,
    timeout: float = 300,
) -> None:
    """把无声画面视频与驱动音频封装为最终结果.

    关键正确性要求：必须用 -map 0:v -map 1:a 显式指定取第一个输入（推理画面）的
    视频流和第二个输入（驱动音频 TTS）的音频流，禁止 ffmpeg 默认流选择行为
    （否则会把源视频自带音轨带进结果，口型与声音错位）。

    时长对齐：驱动音频比视频长时（TTS 15s vs 原视频 9s 很常见），用
    -stream_loop -1 循环视频画面到音频长度（NVENC 硬件重编码），-t 卡到音频时长；
    音频不超过视频时直接 -c:v copy 无损快封装，-shortest 以较短流为准。
    """
    video_duration = _get_media_duration(video_path)
    audio_duration = _get_media_duration(audio_path)

    loop_enabled = Config.enable_video_loop if enable_video_loop is None else enable_video_loop
    need_loop = bool(
        loop_enabled
        and audio_duration > 0
        and video_duration > 0
        and audio_duration > video_duration + Config.duration_epsilon
    )

    if need_loop:
        encoder = _pick_video_encoder()
        # preset 随编码器选择：h264_nvenc 用 p1-p7，libx264 用词形 preset
        preset = "p4" if encoder == "h264_nvenc" else "veryfast"
        logger.info(
            "音频(%.2fs)长于视频(%.2fs)，循环视频并以 %s(%s) 重编码至音频长度",
            audio_duration,
            video_duration,
            encoder,
            preset,
        )

        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:
            # NVENC 可能因驱动/占用失败，兜底 libx264 重试一次
            if encoder == "h264_nvenc":
                logger.warning("h264_nvenc 封装失败，回退 libx264 重试")
                _run_ffmpeg(build_cmd("libx264", "veryfast"), timeout=timeout)
            else:
                raise
    else:
        # 视频不短于音频：直接复制视频流，只把音频替换为驱动音频并转 AAC
        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,
    enable_video_loop: Optional[bool] = None,
) -> None:
    """执行 MuseTalk 推理（可被子线程和测试独立调用）.

    实际部署时替换为 MuseTalk 真实推理逻辑。
    此处为示例实现：提取帧 → 生成无声画面 → 用驱动音频封装。

    enable_video_loop: 驱动音频长于视频时是否循环视频；None 走全局配置。
    """
    fps = _get_video_fps(video_path)
    logger.info("视频 fps: %.2f", fps)

    frames_dir = video_path.parent / "frames"
    frames_dir.mkdir(parents=True, exist_ok=True)
    _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("未从视频中提取到帧")

    # TODO: 替换为 MuseTalk 实际推理逻辑。
    # MuseTalk 真实产物是「无声画面视频」，音轨必须在封装阶段用驱动音频替换。
    logger.warning("使用示例推理逻辑，未实际调用 MuseTalk 模型")

    # 示例：从源视频生成无声画面（-an 丢弃原音轨），模拟 MuseTalk 推理产物。
    # 真实部署时 silent_video_path 应替换为 MuseTalk 输出的无声视频路径。
    silent_video_path = video_path.parent / "visual_silent.mp4"
    _run_ffmpeg(
        [
            "ffmpeg",
            "-y",
            "-i",
            str(video_path),
            "-an",
            "-c:v",
            "libx264",
            "-preset",
            "veryfast",
            str(silent_video_path),
        ],
        timeout=300,
    )

    # 统一封装：显式 -map 取推理画面 + 驱动音频；音频更长时循环视频。
    _mux_video_with_audio(silent_video_path, audio_path, output_path, enable_video_loop=enable_video_loop)

    if not output_path.exists() or output_path.stat().st_size < 1024:
        raise RuntimeError("推理产物不存在或过小")


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


@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 文件."""
    # 并发控制：检查锁
    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())}")

        # 可选：本次任务是否在音频长于视频时循环视频（缺省走全局配置）
        loop_param = request.form.get("enable_video_loop")
        if loop_param is not None:
            task_enable_loop = loop_param.strip() not in ("0", "false", "False", "")
        else:
            task_enable_loop = None

        # 文件大小检查
        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()

        # 启动推理进程（用 subprocess 包装，便于超时终止）
        # 此处直接调用推理函数，实际可改为 subprocess 调用外部脚本
        current_task["process"] = "inference_thread"  # 标记为运行中

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

        def inference_thread():
            try:
                _run_inference(video_path, audio_path, output_path, enable_video_loop=task_enable_loop)
            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)

    # 打印配置
    logger.info("=" * 60)
    logger.info("MuseTalk Flask Server 启动")
    logger.info("  端口: %d", Config.port)
    logger.info("  最大并发: %d", Config.max_concurrent)
    logger.info("  推理超时: %.0fs", Config.inference_timeout)
    logger.info("  视频大小限制: %dMB", Config.video_max_mb)
    logger.info("  音频大小限制: %dMB", Config.audio_max_mb)
    logger.info("  默认 fps: %.1f", Config.default_fps)
    logger.info("  视频编码器: %s", Config.video_encoder)
    logger.info("  音频长于视频时循环视频: %s", Config.enable_video_loop)
    logger.info("=" * 60)

    # 检查 GPU
    gpu_info = _get_gpu_info()
    logger.info("GPU 信息: %s", gpu_info)

    # 启动 Flask（threaded=True 处理并发请求）
    app.run(host="0.0.0.0", port=Config.port, threaded=True)


if __name__ == "__main__":
    main()