mirror of
https://github.com/snapotter-hq/SnapOtter.git
synced 2026-08-03 07:46:42 +02:00
SnapOtterandGitHub
36dde9ad87
gpu_available() answers "can ANY framework use a GPU" (torch, then ONNX, then
paddle). But torch tools consumed that shared boolean directly as
device = torch.device("cuda" if gpu_available() else "cpu"). On a GPU host where
gpu_available() is True via paddle or ONNX while torch is a CPU-only build, those
tools would route to a CUDA torch cannot use and crash. Transcription had the
mirror problem: it runs on CTranslate2 (not torch), so on a transcription-only
GPU box gpu_available() returned False and Whisper ran on CPU despite a GPU.
Add per-framework helpers to gpu.py:
- torch_gpu_available(): torch.cuda.is_available(), honoring SNAPOTTER_GPU.
- ctranslate2_gpu_available(): ctranslate2.get_cuda_device_count() > 0.
Point each tool at the helper for its own framework: upscale, noise_removal,
enhance_faces and restore use torch_gpu_available(); transcribe uses
ctranslate2_gpu_available(). ocr.py keeps gpu_available() (paddle-aware) and the
dispatcher keeps it for its startup GPU-status line. The SNAPOTTER_GPU override
check is factored into a shared _override_disables_gpu() helper.
TDD: 7 new tests in tests/test_gpu_detection.py cover both helpers (override,
CPU-only, absent framework), including the crux that torch_gpu_available() stays
False on a CPU-only torch build even when a GPU exists for another framework.
Claude-Session: https://claude.ai/code/session_01NfaRxjek8ex5nawvx3mVMf
98 lines
2.9 KiB
Python
98 lines
2.9 KiB
Python
"""Speech-to-text transcription using faster-whisper (CTranslate2)."""
|
|
import sys
|
|
import json
|
|
import os
|
|
|
|
from gpu import ctranslate2_gpu_available
|
|
|
|
|
|
MODELS_PATH = os.environ.get(
|
|
"MODELS_PATH",
|
|
os.path.join(os.environ.get("DATA_DIR", "/data"), "ai", "models"),
|
|
)
|
|
|
|
|
|
def emit_progress(percent, stage):
|
|
"""Emit structured progress to stderr for bridge.ts to capture."""
|
|
print(json.dumps({"progress": percent, "stage": stage}), file=sys.stderr, flush=True)
|
|
|
|
|
|
def main():
|
|
input_path = sys.argv[1]
|
|
settings = json.loads(sys.argv[2]) if len(sys.argv) > 2 else {}
|
|
|
|
language = settings.get("language", "auto")
|
|
# "task" is accepted for forward-compat but only "transcribe" is used today
|
|
_task = settings.get("task", "transcribe")
|
|
|
|
try:
|
|
emit_progress(5, "Loading model")
|
|
|
|
# Lazy import -- faster_whisper is only available when the
|
|
# transcription bundle is installed; keeping it lazy lets
|
|
# py_compile succeed without the dependency.
|
|
from faster_whisper import WhisperModel
|
|
|
|
model_dir = os.path.join(MODELS_PATH, "faster-whisper-small")
|
|
|
|
# When the bundled model dir is absent, faster-whisper treats the
|
|
# argument as a Hugging Face repo id and downloads it; strict offline
|
|
# mode blocks that fallback with a clear error.
|
|
from offline_guard import downloads_allowed, ensure_download_allowed
|
|
if not os.path.isdir(model_dir):
|
|
ensure_download_allowed("Whisper transcription model (faster-whisper-small)")
|
|
|
|
if ctranslate2_gpu_available():
|
|
device, compute_type = "cuda", "float16"
|
|
else:
|
|
device, compute_type = "cpu", "int8"
|
|
|
|
model = WhisperModel(
|
|
model_dir,
|
|
device=device,
|
|
compute_type=compute_type,
|
|
local_files_only=not downloads_allowed(),
|
|
)
|
|
|
|
emit_progress(20, "Transcribing")
|
|
|
|
lang_arg = None if language == "auto" else language
|
|
segments_iter, info = model.transcribe(
|
|
input_path,
|
|
language=lang_arg,
|
|
vad_filter=True,
|
|
)
|
|
|
|
detected_language = info.language if info else (language if language != "auto" else "en")
|
|
|
|
segments = []
|
|
batch_count = 0
|
|
for seg in segments_iter:
|
|
segments.append({
|
|
"start": round(seg.start, 3),
|
|
"end": round(seg.end, 3),
|
|
"text": seg.text.strip(),
|
|
})
|
|
batch_count += 1
|
|
if batch_count % 5 == 0:
|
|
emit_progress(min(20 + batch_count, 90), "Transcribing")
|
|
|
|
emit_progress(95, "Done")
|
|
|
|
full_text = " ".join(s["text"] for s in segments)
|
|
|
|
print(json.dumps({
|
|
"success": True,
|
|
"language": detected_language,
|
|
"segments": segments,
|
|
"text": full_text,
|
|
}))
|
|
|
|
except Exception as e:
|
|
print(json.dumps({"error": str(e)}))
|
|
sys.exit(1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|