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

packages/ai/python/enhance_faces.py

@@ -151,7 +151,7 @@ def enhance_with_gfpgan(img_array, only_center_face):
     """Enhance faces using GFPGAN. Returns the enhanced image array."""
     import torch
     from gfpgan import GFPGANer
-    from gpu import gpu_available
+    from gpu import torch_gpu_available
 
     if not os.path.exists(GFPGAN_MODEL_PATH):
         raise FileNotFoundError(f"GFPGAN model not found: {GFPGAN_MODEL_PATH}")
@@ -162,7 +162,7 @@ def enhance_with_gfpgan(img_array, only_center_face):
     from offline_guard import prepare_gfpgan_helper_weights
     prepare_gfpgan_helper_weights(_MODELS_BASE)
 
-    use_gpu = gpu_available()
+    use_gpu = torch_gpu_available()
     device = torch.device("cuda" if use_gpu else "cpu")
 
     enhancer = GFPGANer(
@@ -197,9 +197,9 @@ def enhance_with_codeformer(img_array, fidelity_weight):
     import cv2
     import numpy as np
     import torch
-    from gpu import gpu_available
+    from gpu import torch_gpu_available
 
-    use_gpu = gpu_available()
+    use_gpu = torch_gpu_available()
 
     # codeformer-pip downloads four weights into a cwd-relative tree at import
     # time when they are missing; resolve the bundled ones first so only a

packages/ai/python/gpu.py

@@ -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."""
+    """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"):
+    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.
 

packages/ai/python/noise_removal.py

@@ -296,7 +296,7 @@ def denoise_quality(img_array, strength, detail, color_noise, model_path):
     Uses the Swin-Conv-UNet architecture trained on real-world noise.
     """
     import torch
-    from gpu import gpu_available
+    from gpu import torch_gpu_available
 
     emit_progress(15, "Loading SCUNet model")
 
@@ -313,7 +313,7 @@ def denoise_quality(img_array, strength, detail, color_noise, model_path):
         sys.path.insert(0, os.path.join(os.path.dirname(__file__), "models"))
         from scunet_arch import SCUNet
 
-        use_gpu = gpu_available()
+        use_gpu = torch_gpu_available()
         device = torch.device("cuda" if use_gpu else "cpu")
 
         model = SCUNet(in_nc=3, config=[4, 4, 4, 4, 4, 4, 4], dim=64)
@@ -352,7 +352,7 @@ def denoise_maximum(img_array, strength, detail, color_noise, model_path):
     state-of-the-art image restoration.
     """
     import torch
-    from gpu import gpu_available
+    from gpu import torch_gpu_available
 
     emit_progress(15, "Loading NAFNet model")
 
@@ -369,7 +369,7 @@ def denoise_maximum(img_array, strength, detail, color_noise, model_path):
         sys.path.insert(0, os.path.join(os.path.dirname(__file__), "models"))
         from nafnet_arch import NAFNet
 
-        use_gpu = gpu_available()
+        use_gpu = torch_gpu_available()
         device = torch.device("cuda" if use_gpu else "cpu")
 
         model = NAFNet(

packages/ai/python/restore.py

@@ -649,8 +649,8 @@ def main():
     colorize_strength = float(settings.get("colorizeStrength", 85)) / 100.0
 
     try:
-        from gpu import gpu_available
+        from gpu import torch_gpu_available
-        device = "cuda" if gpu_available() else "cpu"
+        device = "cuda" if torch_gpu_available() else "cpu"
 
         emit_progress(5, "Opening image")
         img_bgr = cv2.imread(input_path, cv2.IMREAD_COLOR)

packages/ai/python/tests/test_gpu_detection.py

@@ -11,6 +11,7 @@ GPU-less host and would wedge the shared AI dispatcher).
 import os
 import subprocess
 import sys
+import types
 
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
 import gpu  # noqa: E402
@@ -80,3 +81,59 @@ def test_gpu_available_never_probes_paddle_without_a_gpu(monkeypatch):
     monkeypatch.setattr(gpu, "_try_paddle_cuda_subprocess", spy)
     assert gpu.gpu_available() is False
     assert called["paddle"] is False
+
+
+# --- Per-framework detection (torch, ctranslate2) --------------------------
+#
+# Torch and CTranslate2 tools must gate on their OWN framework, not the general
+# gpu_available(), which can report True based on paddle or ONNX Runtime while
+# torch is a CPU-only build. Consuming the shared boolean would make those tools
+# route to CUDA on a device their framework cannot use.
+
+def test_torch_gpu_available_true_when_torch_can_use_cuda(monkeypatch):
+    monkeypatch.delenv("SNAPOTTER_GPU", raising=False)
+    monkeypatch.setattr(gpu, "_try_torch_cuda", lambda: True)
+    assert gpu.torch_gpu_available() is True
+
+
+def test_torch_gpu_available_false_when_override_disables_gpu(monkeypatch):
+    monkeypatch.setenv("SNAPOTTER_GPU", "0")
+    monkeypatch.setattr(gpu, "_try_torch_cuda", lambda: True)
+    assert gpu.torch_gpu_available() is False
+
+
+def test_torch_gpu_available_false_when_torch_is_cpu_only(monkeypatch):
+    # The crux: gpu_available() may be True via paddle or ONNX on a GPU box, but a
+    # CPU-only torch build must report no GPU so torch tools do not touch CUDA.
+    monkeypatch.delenv("SNAPOTTER_GPU", raising=False)
+    monkeypatch.setattr(gpu, "_try_torch_cuda", lambda: False)
+    assert gpu.torch_gpu_available() is False
+
+
+def test_ctranslate2_gpu_available_true_when_cuda_device_present(monkeypatch):
+    monkeypatch.delenv("SNAPOTTER_GPU", raising=False)
+    fake = types.SimpleNamespace(get_cuda_device_count=lambda: 1)
+    monkeypatch.setitem(sys.modules, "ctranslate2", fake)
+    assert gpu.ctranslate2_gpu_available() is True
+
+
+def test_ctranslate2_gpu_available_false_when_no_cuda_device(monkeypatch):
+    monkeypatch.delenv("SNAPOTTER_GPU", raising=False)
+    fake = types.SimpleNamespace(get_cuda_device_count=lambda: 0)
+    monkeypatch.setitem(sys.modules, "ctranslate2", fake)
+    assert gpu.ctranslate2_gpu_available() is False
+
+
+def test_ctranslate2_gpu_available_false_when_override_disables_gpu(monkeypatch):
+    monkeypatch.setenv("SNAPOTTER_GPU", "0")
+    fake = types.SimpleNamespace(get_cuda_device_count=lambda: 4)
+    monkeypatch.setitem(sys.modules, "ctranslate2", fake)
+    assert gpu.ctranslate2_gpu_available() is False
+
+
+def test_ctranslate2_gpu_available_false_when_not_installed(monkeypatch):
+    # A None entry in sys.modules makes `import ctranslate2` raise ImportError,
+    # which models the framework being absent (e.g. no transcription bundle).
+    monkeypatch.delenv("SNAPOTTER_GPU", raising=False)
+    monkeypatch.setitem(sys.modules, "ctranslate2", None)
+    assert gpu.ctranslate2_gpu_available() is False

packages/ai/python/transcribe.py

@@ -3,7 +3,7 @@ import sys
 import json
 import os
 
-from gpu import gpu_available
+from gpu import ctranslate2_gpu_available
 
 
 MODELS_PATH = os.environ.get(
@@ -42,7 +42,7 @@ def main():
         if not os.path.isdir(model_dir):
             ensure_download_allowed("Whisper transcription model (faster-whisper-small)")
 
-        if gpu_available():
+        if ctranslate2_gpu_available():
             device, compute_type = "cuda", "float16"
         else:
             device, compute_type = "cpu", "int8"

packages/ai/python/upscale.py

@@ -108,7 +108,7 @@ def main():
                 try:
                     from basicsr.archs.rrdbnet_arch import RRDBNet
                     from realesrgan import RealESRGANer
-                    from gpu import gpu_available
+                    from gpu import torch_gpu_available
                     import numpy as np
                     import torch
 
@@ -117,7 +117,7 @@ def main():
                             f"RealESRGAN model not found: {REALESRGAN_MODEL_PATH}"
                         )
 
-                    use_gpu = gpu_available()
+                    use_gpu = torch_gpu_available()
                     device = torch.device("cuda" if use_gpu else "cpu")
 
                     # RealESRGAN_x4plus is a 4x model internally