cat-shark
4746e0363f
为视频生成 VR 双眼字幕的单体实现：FastAPI 后端、调度器与全部节点 （提音/转写/翻译/ASS/抽帧/OCR/LLM 过滤）在单进程内运行。 - 节点协议（wov_sdk 数据模型）与分布式版保持一致，预留回退桥梁 - 工作流即数据：DAG 存于 workflows/*.json，模型/链路改动只改数据 - 调度器：拓扑顺序执行、断点续跑（产物重建）、任务暂停/继续 - 抽帧按帧间隔（select 按帧号精确取帧），VLM OCR 与 LLM 过滤使用 自适应线程池弹性并发，并打印数据处理速度进度日志 - 100% 行覆盖率（pytest --cov-fail-under=100）
275 lines
12 KiB
Python
Executable File
275 lines
12 KiB
Python
Executable File
"""faster-whisper ASR 节点。
|
||
|
||
单体版中作为进程内节点模块，由调度器直接调用。模型权重默认优先从本地
|
||
目录加载，避免从远端下载，仅在本地找不到模型时才回退到远端 large-v3。
|
||
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")
|
||
|
||
|
||
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-v3，
|
||
nodes/ 同级 model/ 允许部署时把权重随代码目录一起携带。
|
||
"""
|
||
monolith_root = Path(__file__).resolve().parent.parent
|
||
return [
|
||
monolith_root / "model" / "faster-whisper-large-v3",
|
||
monolith_root / "nodes" / "model" / "faster-whisper-large-v3",
|
||
]
|
||
|
||
|
||
def resolve_model_path(
|
||
params: dict,
|
||
env: dict | None = None,
|
||
candidates: list[Path] | None = None,
|
||
) -> str:
|
||
"""按 参数 > 环境变量 > 本地候选目录 > 远端 large-v3 的顺序解析模型路径。
|
||
|
||
本地优先是默认行为：只要候选目录存在且包含 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-v3"
|
||
|
||
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 "auto")
|
||
model = WhisperModel(
|
||
model_path,
|
||
device=device,
|
||
compute_type=compute_type,
|
||
)
|
||
# language 默认日语；vad_filter 默认开启（用户 2026-08 决定）：过滤静音
|
||
# 段以提速并减少无语音处幻觉；长静音时 VAD 压缩时间轴可能轻微错位，
|
||
# 如需极致对齐可在工作流参数中显式关闭。
|
||
# 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())
|
||
# 逐块转写并合并：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):
|
||
chunk_started = time.monotonic()
|
||
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=bool(request.params.get("vad_filter", True)),
|
||
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,
|
||
)
|
||
output_path = output_dir / "transcript.srt"
|
||
output_path.write_text("\n".join(lines), 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))
|
||
|
||
|