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
+4 -4
View File
@@ -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(