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:
SnapOtter
2026-07-06 18:39:01 +08:00
committed by GitHub
parent bc59114dcb
commit 36dde9ad87
7 changed files with 115 additions and 17 deletions
+44 -3
View File
@@ -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.