mirror of
https://github.com/snapotter-hq/SnapOtter.git
synced 2026-08-03 07:46:42 +02:00
fix(ai): gate AI tools on per-framework GPU detection, not a shared boolean (#445)
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
This commit is contained in:
@@ -25,11 +25,23 @@ def _nvidia_smi_gpu_name():
|
||||
return None
|
||||
|
||||
|
||||
def _override_disables_gpu():
|
||||
"""True if SNAPOTTER_GPU is explicitly set to a falsy value (0/false/no)."""
|
||||
override = os.environ.get("SNAPOTTER_GPU")
|
||||
return override is not None and override.lower() in ("0", "false", "no")
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=1)
|
||||
def gpu_available():
|
||||
"""Return True if a usable CUDA GPU is present at runtime."""
|
||||
override = os.environ.get("SNAPOTTER_GPU")
|
||||
if override is not None and override.lower() in ("0", "false", "no"):
|
||||
"""Return True if a usable CUDA GPU is present at runtime.
|
||||
|
||||
This is the general "can any framework use a GPU" check (torch, then ONNX
|
||||
Runtime, then paddle). Tools bound to a single framework should instead call
|
||||
the matching per-framework helper (torch_gpu_available,
|
||||
ctranslate2_gpu_available) so a GPU that only paddle or ONNX can use is not
|
||||
mistaken for a torch GPU.
|
||||
"""
|
||||
if _override_disables_gpu():
|
||||
return False
|
||||
|
||||
# Try torch first -- it probes the hardware directly.
|
||||
@@ -147,6 +159,35 @@ def _try_paddle_cuda_subprocess():
|
||||
return False
|
||||
|
||||
|
||||
def torch_gpu_available():
|
||||
"""True iff torch itself can use CUDA (honors the SNAPOTTER_GPU override).
|
||||
|
||||
Torch-based tools (upscale, denoise, face enhancement, restore) must gate on
|
||||
this rather than gpu_available(), which can report True based on paddle or
|
||||
ONNX Runtime while torch is a CPU-only build. Routing those tools to CUDA on a
|
||||
device torch cannot use would crash them.
|
||||
"""
|
||||
if _override_disables_gpu():
|
||||
return False
|
||||
return _try_torch_cuda()
|
||||
|
||||
|
||||
def ctranslate2_gpu_available():
|
||||
"""True iff CTranslate2 (faster-whisper's backend) can use CUDA.
|
||||
|
||||
Transcription runs on CTranslate2, not torch, so it cannot reuse torch's
|
||||
probe; torch may not even be installed in the transcription bundle. Honors the
|
||||
SNAPOTTER_GPU override and returns False when CTranslate2 is absent.
|
||||
"""
|
||||
if _override_disables_gpu():
|
||||
return False
|
||||
try:
|
||||
import ctranslate2
|
||||
return ctranslate2.get_cuda_device_count() > 0
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
def onnx_providers():
|
||||
"""Return (providers, device) tuple.
|
||||
|
||||
|
||||
Reference in New Issue
Block a user