mirror of
https://github.com/snapotter-hq/SnapOtter.git
synced 2026-08-03 07:46:42 +02:00
feat: add GPU/CUDA acceleration support (:cuda Docker tag)
Add a :cuda Docker image tag that auto-detects NVIDIA GPU at runtime and falls back gracefully to CPU. Same pattern as Immich. - New gpu.py shared utility for cached CUDA detection - Background removal (rembg): pass CUDAExecutionProvider to ONNX Runtime - Upscaling (Real-ESRGAN): use CUDA device + FP16 when GPU available - OCR (PaddleOCR): enable use_gpu when CUDA detected - Dispatcher reports GPU status at startup via readiness signal - Admin health endpoint exposes GPU availability - Dockerfile uses ARG GPU=false with conditional NVIDIA CUDA base image - docker-compose.gpu.yml override for GPU users - CI/CD workflows build and publish :cuda tag (amd64 only) Three tags: :latest (CPU), :lite (no AI), :cuda (GPU with CPU fallback)
This commit is contained in:
@@ -0,0 +1,34 @@
|
||||
"""Runtime GPU/CUDA detection utility."""
|
||||
import functools
|
||||
import os
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=1)
|
||||
def gpu_available():
|
||||
"""Return True if a usable CUDA GPU is present at runtime."""
|
||||
override = os.environ.get("STIRLING_GPU")
|
||||
if override is not None:
|
||||
return override.lower() in ("1", "true", "yes")
|
||||
|
||||
try:
|
||||
import onnxruntime
|
||||
if "CUDAExecutionProvider" in onnxruntime.get_available_providers():
|
||||
return True
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
try:
|
||||
import torch
|
||||
if torch.cuda.is_available():
|
||||
return True
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def onnx_providers():
|
||||
"""Return ONNX Runtime execution providers in priority order."""
|
||||
if gpu_available():
|
||||
return ["CUDAExecutionProvider", "CPUExecutionProvider"]
|
||||
return ["CPUExecutionProvider"]
|
||||
Reference in New Issue
Block a user