94 lines
3.7 KiB
Python
94 lines
3.7 KiB
Python
"""faster-whisper ASR 节点入口。
|
|
|
|
使用 faster-whisper 加载 Whisper 模型，将音频转写为带时间轴的 SRT 文件。
|
|
模型、设备与计算类型均可通过参数或环境变量配置。
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
|
|
from wov_sdk.models import InvokeRequest, InvokeResponse, NodeManifest
|
|
from wov_sdk.server import run_node
|
|
|
|
|
|
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 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，保证健康检查等轻量路径不依赖重型依赖。
|
|
from faster_whisper import WhisperModel
|
|
|
|
# 参数优先于环境变量；模型路径缺省使用 faster-whisper 的 large-v3。
|
|
model_path = str(
|
|
request.params.get("model_path")
|
|
or os.getenv("WHISPER_MODEL_PATH", "large-v3")
|
|
)
|
|
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 过滤静音段以提升转写质量。
|
|
segments, _info = model.transcribe(
|
|
str(audio_path),
|
|
language=str(request.params.get("language", "ja")),
|
|
beam_size=int(request.params.get("beam_size", 1)),
|
|
vad_filter=bool(request.params.get("vad_filter", True)),
|
|
)
|
|
|
|
output_dir = Path(request.output_dir)
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
output_path = output_dir / "transcript.srt"
|
|
# 按 SRT 标准输出：序号、时间轴、文本和空行交替。
|
|
lines: list[str] = []
|
|
for index, segment in enumerate(segments, start=1):
|
|
lines.extend(
|
|
[
|
|
str(index),
|
|
f"{format_timestamp(segment.start)} --> {format_timestamp(segment.end)}",
|
|
segment.text.strip(),
|
|
"",
|
|
]
|
|
)
|
|
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))
|
|
|
|
|
|
def main() -> None:
|
|
"""加载节点清单并以本模块的 invoke 处理器启动服务。"""
|
|
manifest_path = Path(__file__).resolve().parent.parent / "node.manifest.json"
|
|
with open(manifest_path, "r", encoding="utf-8") as f:
|
|
manifest = NodeManifest.from_dict(json.load(f))
|
|
run_node(manifest, invoke)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|