337 lines
17 KiB
Python
Executable File
337 lines
17 KiB
Python
Executable File
"""faster-whisper ASR 节点。
|
||
|
||
单体版中作为进程内节点模块，由调度器直接调用。模型权重默认优先从本地
|
||
目录加载，避免从远端下载，仅在本地找不到模型时才回退到远端 large-v2。
|
||
CUDA 动态库通过 ctypes 在进程内预加载，替代分布式版的 LD_LIBRARY_PATH 注入。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import ctypes
|
||
import os
|
||
import subprocess
|
||
import sysconfig
|
||
import time
|
||
import wave
|
||
from pathlib import Path
|
||
|
||
from nodes.ffmpeg import _ffmpeg_bin
|
||
from wov_app.logging import get_logger
|
||
from wov_sdk.models import InvokeRequest, InvokeResponse
|
||
|
||
# 转写进度日志：输出到主进程控制台，长音频分块时可见每块进度。
|
||
logger = get_logger("whisper")
|
||
|
||
# 暂停信号文件名：位于 run 根目录（<storage>/runs/<run_id>/paused.flag），
|
||
# 与 subtitle-ocr 节点约定一致；批量暂停时由暂停接口写入，分块间检查即中止。
|
||
PAUSE_FLAG = "paused.flag"
|
||
|
||
def _is_windows() -> bool:
|
||
"""判断当前是否为 Windows，供测试单独注入覆盖。"""
|
||
return os.name == "nt"
|
||
|
||
|
||
def _load_cuda_libraries() -> None:
|
||
"""在进程内预加载 pip 安装的 NVIDIA 动态库。
|
||
|
||
分布式版通过给节点子进程注入 LD_LIBRARY_PATH（Windows 为 PATH）解决；
|
||
单体版没有进程边界，必须在导入 faster-whisper 前加载 nvidia 轮子自带
|
||
的 .so/.dll，否则 ctranslate2 初始化 CUDA 时找不到 libcublas.so.12。
|
||
"""
|
||
# site-packages 目录，nvidia 各包的动态库位于其下。
|
||
site_packages = Path(sysconfig.get_paths()["purelib"])
|
||
for vendor in ("cublas", "cudnn", "cuda_nvrtc"):
|
||
# Windows 轮子把 dll 放在 bin/，Linux 放在 lib/。
|
||
for subdir in ("bin", "lib"):
|
||
lib_dir = site_packages / "nvidia" / vendor / subdir
|
||
if not lib_dir.is_dir():
|
||
continue
|
||
if _is_windows():
|
||
# Windows 通过 DLL 搜索目录注册，等价于进程内 PATH 注入。
|
||
os.add_dll_directory(str(lib_dir))
|
||
else:
|
||
for so_file in sorted(lib_dir.glob("*.so*")):
|
||
try:
|
||
ctypes.CDLL(str(so_file))
|
||
except OSError:
|
||
# 个别依赖缺失（如 libcudart）时跳过，交由 ctranslate2 报错。
|
||
continue
|
||
|
||
def _local_model_candidates() -> list[Path]:
|
||
"""返回本地模型候选目录：单体根目录 model/ 优先，其次 nodes/ 同级 model/。
|
||
|
||
单体根目录 model/ 对应仓库根下的 model/faster-whisper-large-v2，
|
||
nodes/ 同级 model/ 允许部署时把权重随代码目录一起携带。
|
||
"""
|
||
monolith_root = Path(__file__).resolve().parent.parent
|
||
return [
|
||
monolith_root / "model" / "faster-whisper-large-v2",
|
||
monolith_root / "nodes" / "model" / "faster-whisper-large-v2",
|
||
]
|
||
|
||
|
||
def resolve_model_path(
|
||
params: dict,
|
||
env: dict | None = None,
|
||
candidates: list[Path] | None = None,
|
||
) -> str:
|
||
"""按 参数 > 环境变量 > 本地候选目录 > 远端 large-v2 的顺序解析模型路径。
|
||
|
||
本地优先是默认行为：只要候选目录存在且包含 model.bin 就使用本地权重，
|
||
避免从 Hugging Face 下载；远端下载仅在全部本地候选缺失时作为兜底。
|
||
参数/环境变量传入的是裸模型名（不含路径分隔符）时，会先在本地模型
|
||
目录（model/）下按名解析，方便工作流直接引用下载好的模型。
|
||
candidates 参数供测试注入临时目录，默认使用 _local_model_candidates()。
|
||
"""
|
||
env = env if env is not None else os.environ
|
||
candidates = candidates if candidates is not None else _local_model_candidates()
|
||
explicit = params.get("model_path") or env.get("WHISPER_MODEL_PATH")
|
||
if explicit:
|
||
explicit_str = str(explicit)
|
||
# 裸模型名按 <模型目录>/<名称> 在本地解析，例如
|
||
# "whisper-large-v3-translate-zh-v0.1-lt-ct2"。
|
||
if "/" not in explicit_str and "\\" not in explicit_str:
|
||
named = candidates[0].parent / explicit_str
|
||
if (named / "model.bin").is_file():
|
||
return str(named)
|
||
return explicit_str
|
||
for candidate in candidates:
|
||
# model.bin 是 CTranslate2 权重的必需文件，存在才认为模型完整。
|
||
if candidate.is_dir() and (candidate / "model.bin").is_file():
|
||
return str(candidate)
|
||
return "large-v2"
|
||
|
||
def format_timestamp(seconds: float) -> str:
|
||
"""把秒数格式化为 SRT 时间戳，例如 01:00:00,500。"""
|
||
# 先换算成毫秒再逐级拆分为时/分/秒/毫秒，避免浮点误差。
|
||
total_ms = int(seconds * 1000)
|
||
hours, remainder = divmod(total_ms, 3600000)
|
||
minutes, remainder = divmod(remainder, 60000)
|
||
secs, millis = divmod(remainder, 1000)
|
||
return f"{hours:02d}:{minutes:02d}:{secs:02d},{millis:03d}"
|
||
|
||
|
||
def _split_audio(
|
||
audio_path: Path,
|
||
output_dir: Path,
|
||
chunk_seconds: int,
|
||
ffmpeg_bin: str | None,
|
||
) -> list[Path]:
|
||
"""用 ffmpeg 把音频切成 chunk_seconds 秒一块的 wav，返回块路径列表。
|
||
|
||
分块是应用层工程策略（内存有界、失败粒度小），与模型 30s 窗口无关；
|
||
whisper 训练与推理都按 30s 窗口解码，任意块大小都适用。以下情况回退
|
||
为整段单次转写：chunk_seconds <= 0、找不到 ffmpeg、切块失败、音频本身
|
||
不足一块（ffmpeg 产出单块）。
|
||
"""
|
||
if chunk_seconds <= 0 or not ffmpeg_bin:
|
||
return [audio_path]
|
||
chunk_dir = output_dir / "chunks"
|
||
chunk_dir.mkdir(parents=True, exist_ok=True)
|
||
pattern = str(chunk_dir / "chunk_%03d.wav")
|
||
# 音频已是 16kHz 单声道 WAV，流拷贝切块即可，无需重编码。
|
||
result = subprocess.run(
|
||
[
|
||
ffmpeg_bin,
|
||
"-y",
|
||
"-i",
|
||
str(audio_path),
|
||
"-f",
|
||
"segment",
|
||
"-segment_time",
|
||
str(chunk_seconds),
|
||
"-c",
|
||
"copy",
|
||
pattern,
|
||
],
|
||
capture_output=True,
|
||
text=True,
|
||
)
|
||
if result.returncode != 0:
|
||
# 切块失败（输入损坏等）回退整段，不让转写流程中断。
|
||
return [audio_path]
|
||
chunks = sorted(chunk_dir.glob("chunk_*.wav"))
|
||
return chunks or [audio_path]
|
||
|
||
|
||
def _wav_duration_seconds(path: Path, fallback: float) -> float:
|
||
"""从 WAV 头精确读取时长；文件非法/非 WAV 时用 fallback 兜底。
|
||
|
||
分块偏移必须用每块的实际时长累积，而不是 块序号×块长 的假设值——
|
||
ffmpeg 切出的块实际时长并不精确等于块长（如 60.05s），假设值会随
|
||
块数累积漂移，造成字幕时间轴逐渐错位。
|
||
"""
|
||
try:
|
||
with wave.open(str(path), "rb") as wav:
|
||
rate = wav.getframerate()
|
||
return wav.getnframes() / rate if rate else fallback
|
||
except (wave.Error, EOFError, OSError):
|
||
# 文件损坏/非 WAV（如 mp4 直传）时用 fallback 兜底。
|
||
return fallback
|
||
|
||
|
||
def _append_srt_lines(lines: list[str], segments, offset: float, start_index: int) -> int:
|
||
"""把一段转写结果按 SRT 格式追加到 lines，时间加上 offset 偏移。
|
||
|
||
分块合并时每块 offset 为前面所有块的实际时长累积；单次调用 offset=0。
|
||
每写出一条字幕就打印其编号与完整视频角度的时间范围，便于对照对齐。
|
||
返回本段新增的条数，用于全局序号递增。
|
||
"""
|
||
count = 0
|
||
for segment in segments:
|
||
# 完整视频角度的时间 = 模型预测时间 + 累积偏移。
|
||
start_time = segment.start + offset
|
||
end_time = segment.end + offset
|
||
lines.extend(
|
||
[
|
||
str(start_index + count),
|
||
f"{format_timestamp(start_time)} --> {format_timestamp(end_time)}",
|
||
segment.text.strip(),
|
||
"",
|
||
]
|
||
)
|
||
logger.info(
|
||
"分段 #%d: %s --> %s",
|
||
start_index + count,
|
||
format_timestamp(start_time),
|
||
format_timestamp(end_time),
|
||
)
|
||
count += 1
|
||
return count
|
||
|
||
def invoke(request: InvokeRequest) -> InvokeResponse:
|
||
"""转写音频并生成 SRT 字幕，产物为 transcript.srt。"""
|
||
audio_uri = request.inputs.get("audio_uri")
|
||
if not audio_uri:
|
||
return InvokeResponse(status="failed", error="audio_uri is required")
|
||
|
||
# 文件不存在时提前失败，避免进入耗时的模型加载流程。
|
||
audio_path = Path(audio_uri)
|
||
if not audio_path.is_file():
|
||
return InvokeResponse(status="failed", error="audio file not found")
|
||
|
||
try:
|
||
# 延迟导入 faster-whisper，保证节点注册与调度等轻量路径不依赖重型依赖；
|
||
# 导入前先预加载 NVIDIA 动态库，否则 ctranslate2 找不到 libcublas。
|
||
_load_cuda_libraries()
|
||
from faster_whisper import WhisperModel
|
||
# 模型路径默认本地优先：参数 > 环境变量 > 工作区本地目录 > 远端兜底。
|
||
model_path = resolve_model_path(request.params)
|
||
device = str(request.params.get("device") or os.getenv("WHISPER_DEVICE", "auto"))
|
||
# auto 让 faster-whisper 根据硬件自动选择 float16/int8 等计算类型。
|
||
compute_type = str(request.params.get("compute_type") or "float16")
|
||
model = WhisperModel(
|
||
model_path,
|
||
device=device,
|
||
compute_type=compute_type,
|
||
)
|
||
# language 默认日语；vad_filter 默认开启（用户 2026-08 决定）：过滤静音
|
||
# 段以提速并减少无语音处幻觉；长静音时 VAD 压缩时间轴可能轻微错位，
|
||
# 如需极致对齐可在工作流参数中显式关闭。
|
||
# decode_full=true（默认 false）：VAD 对呻吟/轻语/BGM 混叠声学切段会
|
||
# 把真话当非语音剔除（实测 savr-1054 全片仅召回 115 条），开启后强制
|
||
# 无 VAD 整段解码（vad_filter=False 且跳过自动 VAD 分析）以召回弱语音，
|
||
# 代价是无语音段会产生长时套话幻觉，由下方日语幻觉清洗兜底移除。
|
||
# 另：decode_full 救回的弱语音中混有纯语气词碎片（あ/ん？等），由
|
||
# short_moan_max_chars（默认 3，0=关闭）参数控制短呻吟整条删除。
|
||
# task 默认 transcribe，中文直出模型可传 translate 直接翻译为目标语言。
|
||
# condition_on_previous_text 默认 False：长音频下开启会导致重复/漂移，
|
||
# 关闭后每个 30s 窗口独立解码，是 faster-whisper 官方建议的长音频方案。
|
||
output_dir = Path(request.output_dir)
|
||
output_dir.mkdir(parents=True, exist_ok=True)
|
||
# 分块转写：默认每 1 分钟一块（chunk_seconds=60），切块失败自动回退整段。
|
||
chunk_seconds = int(request.params.get("chunk_seconds", 60))
|
||
chunks = _split_audio(audio_path, output_dir, chunk_seconds, _ffmpeg_bin())
|
||
# decode_full 时强制无 VAD：不传 vad_filter/vad_parameters，也不跑自动 VAD
|
||
# 分析（分析结果对呻吟/轻语类音频无效，只会把整块切碎/剔除真话）。
|
||
decode_full = bool(request.params.get("decode_full", False))
|
||
vad_parameters = request.params.get("vad_parameters")
|
||
vad_filter = bool(request.params.get("vad_filter", True))
|
||
if not decode_full and vad_filter and not vad_parameters and os.getenv("WOV_AUTO_VAD", "1") == "1":
|
||
try:
|
||
from nodes.vad_profiler import vad_parameters_for_audio
|
||
|
||
vad_parameters = vad_parameters_for_audio(
|
||
audio_path,
|
||
sample_rate=int(request.params.get("sample_rate", 16000)),
|
||
whisper_invoke=None, # 片段网格需模型；生产默认仅信号分析（快速）
|
||
run_dir=Path(request.output_dir),
|
||
chunk_seconds=chunk_seconds,
|
||
)
|
||
logger.info("自动 VAD 参数: %s", vad_parameters)
|
||
except Exception as exc: # noqa: BLE001
|
||
# 分析失败不影响转写：回退默认 bentenVAD，仅记录。
|
||
logger.warning("自动 VAD 分析失败，回退默认: %s", exc)
|
||
vad_parameters = None
|
||
# 逐块转写并合并：offset 用每块实际时长累积（WAV 头精确），SRT 序号连续。
|
||
logger.info("转写开始: %d 个分块", len(chunks))
|
||
lines: list[str] = []
|
||
offset = 0.0
|
||
# SRT 序号从 1 开始，跨块连续递增。
|
||
srt_number = 1
|
||
transcribe_started = time.monotonic()
|
||
for chunk_index, chunk in enumerate(chunks, start=1):
|
||
# 暂停检查：批量暂停时在 run 根目录写 paused.flag，whisper 在分块
|
||
# 之间检查该信号（默认 60s 一块，暂停粒度不超过一块），检测到即抛
|
||
# 异常，由 invoke 转 failed、调度器保持任务 PAUSED，继续时整个节点
|
||
# 重新转写（whisper 没有节点级断点存档，产物只在结束时一次性写出）。
|
||
if (Path(request.output_dir).parent.parent / PAUSE_FLAG).exists():
|
||
raise RuntimeError(f"whisper 被暂停（run {request.run_id}）")
|
||
chunk_started = time.monotonic()
|
||
# decode_full=true 时 vad_filter 传 False：跳过 VAD 剔除弱语音段。
|
||
segments, _info = model.transcribe(
|
||
str(chunk),
|
||
language=str(request.params.get("language", "ja")),
|
||
task=str(request.params.get("task", "transcribe")),
|
||
beam_size=int(request.params.get("beam_size", 1)),
|
||
vad_filter=False if decode_full else vad_filter,
|
||
vad_parameters=None if decode_full else vad_parameters,
|
||
condition_on_previous_text=bool(
|
||
request.params.get("condition_on_previous_text", False)
|
||
),
|
||
)
|
||
# 进度日志：块序号/总数、单块耗时、实时倍率（块音频时长/墙钟耗时）
|
||
# 与转写累计耗时，直观反映数据处理速度。
|
||
chunk_elapsed = time.monotonic() - chunk_started
|
||
srt_number += _append_srt_lines(lines, segments, offset, srt_number)
|
||
# 偏移按本块实际时长推进，避免假设块长导致的累积漂移。
|
||
offset += _wav_duration_seconds(chunk, chunk_seconds)
|
||
logger.info(
|
||
"分块 %d/%d 完成 offset=%.2fs 耗时 %.1fs (%.2fx 实时, 累计 %.1fs)",
|
||
chunk_index, len(chunks), offset, chunk_elapsed,
|
||
chunk_seconds / chunk_elapsed if chunk_elapsed > 0 else 0.0,
|
||
time.monotonic() - transcribe_started,
|
||
)
|
||
# decode_full（无 VAD）副作用：无语音/音乐/呻吟段会产生长时寒暄套话幻觉
|
||
# （おやすみなさい/ご視聴ありがとうございました 等），在此**连带时间戳整条
|
||
# 剔除**（序号/时间轴/文本全删、剩余重编号），不留下 '-' 占位污染下游
|
||
# （占位会渲染进 ASS 成减号、翻译/过滤都要额外处理）；短时（≤15s）相同词
|
||
# 可能是剧情真实道晚安，保留。见 nodes/subtitle_cleanup.py。
|
||
# 其次（2026-09 用户决策）：decode_full 也会把呻吟/BGM 混叠的弱语音整段
|
||
# 救回，其中混有大量**纯语气词碎片**（あ…/ん?/はぁ…/あ!あ!/んふふ 等），
|
||
# 这类噪声影响字幕观感；在此按'全部字符∈纯呻吟集合 且 有效假名≤max_chars'
|
||
# 判据**整条删除**（remove_short_moan_entries），真实短对话（そこ/やばい/
|
||
# ねえ/やだ）天然不命中。仅 decode_full 生效，demo 等 VAD 链路不受影响；
|
||
# 参数 short_moan_max_chars 可调（默认 3，设 0 关闭）。
|
||
if decode_full:
|
||
from nodes.subtitle_cleanup import (
|
||
clean_japanese_hallucinations,
|
||
remove_short_moan_entries,
|
||
)
|
||
|
||
body = clean_japanese_hallucinations("\n".join(lines))
|
||
body = remove_short_moan_entries(
|
||
body,
|
||
max_chars=int(request.params.get("short_moan_max_chars", 3)),
|
||
)
|
||
else:
|
||
body = "\n".join(lines)
|
||
output_path = output_dir / "transcript.srt"
|
||
output_path.write_text(body, encoding="utf-8")
|
||
return InvokeResponse(status="completed", outputs={"srt_uri": str(output_path)})
|
||
except Exception as exc: # noqa: BLE001
|
||
# 模型加载或转写异常统一转换为 failed 响应，不让调度线程崩溃。
|
||
return InvokeResponse(status="failed", error=str(exc))
|
||
|
||
|