External/SnapOtter
1
0
Fork 0
mirror of https://github.com/snapotter-hq/SnapOtter.git synced 2026-08-03 07:46:42 +02:00
Code Issues Packages Projects Releases Wiki Activity

feat: add GPU/CUDA acceleration support (:cuda Docker tag)

Browse Source 
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:
Siddharth Kumar Sah
2026-04-05 19:12:45 +08:00
parent d0c69d6a46
commit 29a382e9e0
13 changed files with 182 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/ci.yml
+17 -6
View File
@@ -80,11 +80,20 @@ jobs:
- run: pnpm build
docker:
name: Docker Build Test (${{ matrix.variant }})
name: Docker Build Test (${{ matrix.tag }})
runs-on: ubuntu-latest
strategy:
matrix:
variant: [full, lite]
include:
- tag: full
variant: full
gpu: "false"
- tag: lite
variant: lite
gpu: "false"
- tag: cuda
variant: full
gpu: "true"
steps:
- uses: actions/checkout@v4
@@ -95,7 +104,9 @@ jobs:
context: .
file: docker/Dockerfile
push: false
build-args: VARIANT=${{ matrix.variant }}
tags: stirling-image:ci-${{ matrix.variant }}
cache-from: type=gha,scope=${{ matrix.variant }}
cache-to: type=gha,mode=max,scope=${{ matrix.variant }}
build-args: |
VARIANT=${{ matrix.variant }}
GPU=${{ matrix.gpu }}
tags: stirling-image:ci-${{ matrix.tag }}
cache-from: type=gha,scope=${{ matrix.tag }}
cache-to: type=gha,mode=max,scope=${{ matrix.tag }}
.github/workflows/release.yml
+21 -9
View File
@@ -49,18 +49,28 @@ jobs:
fi
docker:
name: Docker (${{ matrix.variant }})
name: Docker (${{ matrix.tag }})
needs: release
if: needs.release.outputs.new_version != ''
runs-on: ubuntu-latest
strategy:
matrix:
variant: [full, lite]
include:
- variant: full
- tag: full
variant: full
gpu: "false"
suffix: ""
- variant: lite
platforms: "linux/amd64,linux/arm64"
- tag: lite
variant: lite
gpu: "false"
suffix: "-lite"
platforms: "linux/amd64,linux/arm64"
- tag: cuda
variant: full
gpu: "true"
suffix: "-cuda"
platforms: "linux/amd64"
steps:
- name: Checkout release tag
uses: actions/checkout@v4
@@ -97,7 +107,7 @@ jobs:
type=semver,pattern={{version}}${{ matrix.suffix }},value=v${{ needs.release.outputs.new_version }}
type=semver,pattern={{major}}.{{minor}}${{ matrix.suffix }},value=v${{ needs.release.outputs.new_version }}
type=semver,pattern={{major}}${{ matrix.suffix }},value=v${{ needs.release.outputs.new_version }}
type=raw,value=${{ matrix.variant == 'full' && 'latest' || 'lite' }}
type=raw,value=${{ matrix.tag == 'full' && 'latest' || matrix.tag }}
- name: Build and push
uses: docker/build-push-action@v6
@@ -105,9 +115,11 @@ jobs:
context: .
file: docker/Dockerfile
push: true
build-args: VARIANT=${{ matrix.variant }}
platforms: linux/amd64,linux/arm64
build-args: |
VARIANT=${{ matrix.variant }}
GPU=${{ matrix.gpu }}
platforms: ${{ matrix.platforms }}
tags: ${{ steps.meta.outputs.tags }}
labels: ${{ steps.meta.outputs.labels }}
cache-from: type=gha,scope=${{ matrix.variant }}
cache-to: type=gha,mode=max,scope=${{ matrix.variant }}
cache-from: type=gha,scope=${{ matrix.tag }}
cache-to: type=gha,mode=max,scope=${{ matrix.tag }}
apps/api/src/index.ts
+2 -1
View File
@@ -1,5 +1,6 @@
import cors from "@fastify/cors";
import rateLimit from "@fastify/rate-limit";
import { isGpuAvailable } from "@stirling-image/ai";
import { APP_VERSION } from "@stirling-image/shared";
import Fastify from "fastify";
import { env } from "./config.js";
@@ -132,7 +133,7 @@ app.get("/api/v1/admin/health", async (request, reply) => {
storage: { mode: env.STORAGE_MODE, available: "N/A" },
database: dbOk ? "ok" : "error",
queue: { active: 0, pending: 0 },
ai: {},
ai: { gpu: isGpuAvailable() },
};
});
docker/Dockerfile
+50 -11
View File
@@ -5,6 +5,7 @@
# ============================================
ARG VARIANT=full
ARG GPU=false
# ============================================
# Stage 1: Build the frontend (Vite + React)
@@ -38,11 +39,35 @@ RUN --mount=type=cache,id=turbo-cache,target=/app/.turbo \
pnpm --filter @stirling-image/web build
# ============================================
# Stage 2: Production runtime
# Stage 2: Base image selection
# ============================================
FROM node:22-bookworm AS production
FROM node:22-bookworm AS base-cpu
FROM nvidia/cuda:12.6.3-runtime-ubuntu24.04 AS base-gpu
# Select base: GPU=false -> base-cpu, GPU=true -> base-gpu
FROM base-cpu AS production-base-false
FROM base-gpu AS production-base-true
# ============================================
# Stage 3: Production runtime
# ============================================
FROM production-base-${GPU} AS production
ARG VARIANT
ARG GPU
# Install Node.js when using CUDA base (node:22-bookworm already has it)
RUN if [ "$GPU" = "true" ]; then \
apt-get update && apt-get install -y --no-install-recommends \
curl ca-certificates gnupg && \
mkdir -p /etc/apt/keyrings && \
curl -fsSL https://deb.nodesource.com/gpgkey/nodesource-repo.gpg.key | \
gpg --dearmor -o /etc/apt/keyrings/nodesource.gpg && \
echo "deb [signed-by=/etc/apt/keyrings/nodesource.gpg] https://deb.nodesource.com/node_22.x nodistro main" > \
/etc/apt/sources.list.d/nodesource.list && \
apt-get update && apt-get install -y nodejs && \
rm -rf /var/lib/apt/lists/* \
; fi
RUN corepack enable && corepack prepare pnpm@9.15.4 --activate
@@ -68,17 +93,30 @@ RUN if [ "$VARIANT" = "full" ]; then \
# Python venv + ML packages + model weights (full variant only)
COPY packages/ai/python/requirements.txt /tmp/requirements.txt
COPY packages/ai/python/requirements-gpu.txt /tmp/requirements-gpu.txt
RUN if [ "$VARIANT" = "full" ]; then \
python3 -m venv /opt/venv && \
/opt/venv/bin/pip install --upgrade pip && \
/opt/venv/bin/pip install \
Pillow numpy opencv-python-headless onnxruntime && \
(/opt/venv/bin/pip install "rembg[cpu]" || echo "WARNING: rembg not installed") && \
(/opt/venv/bin/pip install realesrgan || echo "WARNING: realesrgan not installed") && \
(/opt/venv/bin/pip install paddlepaddle paddleocr || echo "WARNING: PaddleOCR not installed") && \
(/opt/venv/bin/pip install mediapipe || echo "WARNING: mediapipe not installed") && \
(/opt/venv/bin/pip install lama-cleaner || echo "WARNING: lama-cleaner not installed") \
; fi && rm -f /tmp/requirements.txt
if [ "$GPU" = "true" ]; then \
/opt/venv/bin/pip install \
Pillow numpy opencv-python-headless onnxruntime-gpu && \
(/opt/venv/bin/pip install rembg || echo "WARNING: rembg not installed") && \
(/opt/venv/bin/pip install realesrgan \
--extra-index-url https://download.pytorch.org/whl/cu126 \
|| echo "WARNING: realesrgan not installed") && \
(/opt/venv/bin/pip install paddlepaddle-gpu paddleocr || echo "WARNING: PaddleOCR not installed") && \
(/opt/venv/bin/pip install mediapipe || echo "WARNING: mediapipe not installed") && \
(/opt/venv/bin/pip install lama-cleaner || echo "WARNING: lama-cleaner not installed") \
; else \
/opt/venv/bin/pip install \
Pillow numpy opencv-python-headless onnxruntime && \
(/opt/venv/bin/pip install "rembg[cpu]" || echo "WARNING: rembg not installed") && \
(/opt/venv/bin/pip install realesrgan || echo "WARNING: realesrgan not installed") && \
(/opt/venv/bin/pip install paddlepaddle paddleocr || echo "WARNING: PaddleOCR not installed") && \
(/opt/venv/bin/pip install mediapipe || echo "WARNING: mediapipe not installed") && \
(/opt/venv/bin/pip install lama-cleaner || echo "WARNING: lama-cleaner not installed") \
; fi \
; fi && rm -f /tmp/requirements.txt /tmp/requirements-gpu.txt
COPY docker/download_models.py /tmp/download_models.py
RUN if [ "$VARIANT" = "full" ]; then \
@@ -152,7 +190,8 @@ ENV PORT=1349 \
CONCURRENT_JOBS=3 \
MAX_MEGAPIXELS=100 \
RATE_LIMIT_PER_MIN=100 \
STIRLING_VARIANT=${VARIANT}
STIRLING_VARIANT=${VARIANT} \
STIRLING_GPU=${GPU}
# Create non-root user for runtime
RUN groupadd -r stirling && useradd -r -g stirling -d /app -s /sbin/nologin stirling
docker/docker-compose.gpu.yml
+17
View File
@@ -0,0 +1,17 @@
# GPU override - use with:
# docker compose -f docker/docker-compose.yml -f docker/docker-compose.gpu.yml up
services:
stirling-image:
build:
context: ..
dockerfile: docker/Dockerfile
args:
GPU: "true"
image: stirlingimage/stirling-image:cuda
deploy:
resources:
reservations:
devices:
- driver: nvidia
count: 1
capabilities: [gpu]
packages/ai/python/dispatcher.py
+9 -2
View File
@@ -39,6 +39,7 @@ def _try_import(name, import_fn):
_try_import("PIL", lambda: __import__("PIL"))
_try_import("cv2", lambda: __import__("cv2"))
_try_import("numpy", lambda: __import__("numpy"))
_try_import("gpu", lambda: __import__("gpu"))
# Heavy ML libraries - import but don't fail if unavailable
_try_import("rembg", lambda: __import__("rembg"))
@@ -123,8 +124,14 @@ def _run_script_main(script_name, args):
def main():
# Signal readiness
print(json.dumps({"ready": True}), file=sys.stderr, flush=True)
# Signal readiness with GPU status
gpu = False
try:
from gpu import gpu_available
gpu = gpu_available()
except ImportError:
pass
print(json.dumps({"ready": True, "gpu": gpu}), file=sys.stderr, flush=True)
for line in sys.stdin:
line = line.strip()
packages/ai/python/gpu.py
+34
View File
@@ -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"]
packages/ai/python/ocr.py
+2 -1
View File
@@ -34,9 +34,10 @@ def run_paddleocr(input_path, language):
"""Run PaddleOCR."""
os.environ["PADDLE_PDX_DISABLE_MODEL_SOURCE_CHECK"] = "True"
from paddleocr import PaddleOCR
from gpu import gpu_available
emit_progress(20, "Loading")
ocr = PaddleOCR(lang=language)
ocr = PaddleOCR(lang=language, use_gpu=gpu_available())
emit_progress(30, "Scanning")
result = ocr.ocr(input_path)
emit_progress(70, "Extracting text")
packages/ai/python/remove_bg.py
+2 -1
View File
@@ -24,11 +24,12 @@ def main():
try:
from rembg import remove, new_session
from gpu import onnx_providers
import io
emit_progress(10, "Loading model")
session = new_session(model)
session = new_session(model, providers=onnx_providers())
emit_progress(25, "Model loaded")
packages/ai/python/requirements-gpu.txt
+10
View File
@@ -0,0 +1,10 @@
rembg==2.0.62
realesrgan==0.3.0
lama-cleaner==1.2.5
paddleocr==2.9.1
paddlepaddle-gpu==3.0.0
mediapipe==0.10.21
onnxruntime-gpu==1.20.1
numpy==1.26.4
Pillow==11.1.0
opencv-python-headless==4.10.0.84
packages/ai/python/upscale.py
+7 -1
View File
@@ -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"))
packages/ai/src/bridge.ts
+10
View File
@@ -54,6 +54,8 @@ interface PendingRequest {
let dispatcher: ChildProcess | null = null;
let dispatcherReady = false;
let dispatcherFailed = false;
// biome-ignore lint/style/useConst: reassigned on dispatcher readiness signal
let dispatcherGpuAvailable = false;
const pendingRequests = new Map<string, PendingRequest>();
let stdoutBuffer = "";
@@ -82,6 +84,7 @@ function startDispatcher(): ChildProcess | null {
// Readiness signal
if (parsed.ready === true) {
dispatcherReady = true;
dispatcherGpuAvailable = parsed.gpu === true;
continue;
}
@@ -221,6 +224,13 @@ function dispatcherRun(
});
}
/**
* Whether the Python dispatcher detected a CUDA GPU at startup.
*/
export function isGpuAvailable(): boolean {
return dispatcherGpuAvailable;
}
/**
* Shut down the persistent dispatcher process.
*/
packages/ai/src/index.ts
+1 -1
View File
@@ -1,5 +1,5 @@
export { removeBackground } from "./background-removal.js";
export { shutdownDispatcher } from "./bridge.js";
export { isGpuAvailable, shutdownDispatcher } from "./bridge.js";
export { blurFaces } from "./face-detection.js";
export { inpaint } from "./inpainting.js";
export { extractText } from "./ocr.js";