92 lines
2.9 KiB
Python
92 lines
2.9 KiB
Python
"""
|
|
RunPod Serverless worker: Speech-to-Text (Whisper via faster-whisper).
|
|
|
|
Chosen over Parakeet/NeMo because it has NO heavy dependencies (no NeMo, no
|
|
torch — just ctranslate2), so it builds cleanly, AND it's multilingual (99
|
|
languages incl. Hindi/regional + English) — a better fit for a mixed crew.
|
|
|
|
Lets a worker speak instead of typing annotation text (image + video editors).
|
|
|
|
Matches the app's transcribe contract:
|
|
{"input": {
|
|
"audio": "<base64 audio>", # required (wav/mp3/m4a; any sample rate — resampled internally)
|
|
"language": "en", # optional (e.g. "hi", "en"); omit = auto-detect
|
|
"timestamps": false # optional — return per-segment start/end
|
|
}}
|
|
Returns: {"transcript": "...", "language": "en", "segments"?: [{text,start_sec,end_sec}]}
|
|
(or {"error": "..."}).
|
|
|
|
Tuning (env): WHISPER_MODEL (default large-v3; "medium"/"small" are faster),
|
|
WHISPER_COMPUTE (default float16; "int8_float16" for less VRAM).
|
|
"""
|
|
|
|
import base64
|
|
import os
|
|
import tempfile
|
|
|
|
_VOLUME = "/runpod-volume"
|
|
if os.path.isdir(_VOLUME):
|
|
os.environ.setdefault("HF_HOME", os.path.join(_VOLUME, "huggingface"))
|
|
|
|
import runpod # noqa: E402
|
|
from faster_whisper import WhisperModel # noqa: E402
|
|
|
|
MODEL_SIZE = os.environ.get("WHISPER_MODEL", "large-v3")
|
|
DEVICE = os.environ.get("WHISPER_DEVICE", "cuda")
|
|
COMPUTE = os.environ.get("WHISPER_COMPUTE", "float16")
|
|
|
|
_model = None
|
|
|
|
|
|
def _get_model():
|
|
global _model
|
|
if _model is None:
|
|
_model = WhisperModel(
|
|
MODEL_SIZE,
|
|
device=DEVICE,
|
|
compute_type=COMPUTE,
|
|
download_root=os.environ.get("HF_HOME"),
|
|
)
|
|
return _model
|
|
|
|
|
|
def handler(event):
|
|
inp = (event or {}).get("input") or {}
|
|
audio_b64 = inp.get("audio")
|
|
if not audio_b64:
|
|
return {"error": "Missing 'audio' (base64)."}
|
|
|
|
language = inp.get("language") or None
|
|
want_ts = bool(inp.get("timestamps", False))
|
|
|
|
path = None
|
|
try:
|
|
data = audio_b64.split(",", 1)[1] if "," in audio_b64 else audio_b64
|
|
with tempfile.NamedTemporaryFile(suffix=".audio", delete=False) as f:
|
|
f.write(base64.b64decode(data))
|
|
path = f.name
|
|
|
|
segments, info = _get_model().transcribe(path, language=language, vad_filter=True)
|
|
|
|
parts = []
|
|
segs = []
|
|
for s in segments:
|
|
parts.append(s.text)
|
|
if want_ts:
|
|
segs.append(
|
|
{"text": s.text.strip(), "start_sec": round(s.start, 2), "end_sec": round(s.end, 2)}
|
|
)
|
|
|
|
out = {"transcript": "".join(parts).strip(), "language": info.language}
|
|
if want_ts:
|
|
out["segments"] = segs
|
|
return out
|
|
except Exception as exc:
|
|
return {"error": f"{type(exc).__name__}: {exc}"}
|
|
finally:
|
|
if path and os.path.exists(path):
|
|
os.unlink(path)
|
|
|
|
|
|
runpod.serverless.start({"handler": handler})
|