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:
@@ -26,7 +26,12 @@ def main():
|
||||
try:
|
||||
from basicsr.archs.rrdbnet_arch import RRDBNet
|
||||
from realesrgan import RealESRGANer
|
||||
from gpu import gpu_available
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
use_gpu = gpu_available()
|
||||
device = torch.device("cuda" if use_gpu else "cpu")
|
||||
|
||||
model = RRDBNet(
|
||||
num_in_ch=3,
|
||||
@@ -40,7 +45,8 @@ def main():
|
||||
scale=scale,
|
||||
model_path=None,
|
||||
model=model,
|
||||
half=False,
|
||||
half=use_gpu,
|
||||
device=device,
|
||||
)
|
||||
emit_progress(20, "Model ready")
|
||||
img_array = np.array(img.convert("RGB"))
|
||||
|
||||
Reference in New Issue
Block a user