77 lines
2.6 KiB
Python
77 lines
2.6 KiB
Python
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:
|
|
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:
|
|
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:
|
|
from faster_whisper import WhisperModel
|
|
|
|
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"))
|
|
compute_type = str(request.params.get("compute_type") or "auto")
|
|
model = WhisperModel(
|
|
model_path,
|
|
device=device,
|
|
compute_type=compute_type,
|
|
)
|
|
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"
|
|
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
|
|
return InvokeResponse(status="failed", error=str(exc))
|
|
|
|
|
|
def main() -> None:
|
|
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()
|