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)

.github/workflows/ci.yml

@@ -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 }}
+          build-args: |
-          tags: stirling-image:ci-${{ matrix.variant }}
+            VARIANT=${{ matrix.variant }}
-          cache-from: type=gha,scope=${{ matrix.variant }}
+            GPU=${{ matrix.gpu }}
-          cache-to: type=gha,mode=max,scope=${{ matrix.variant }}
+          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

@@ -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 }}
+          build-args: |
-          platforms: linux/amd64,linux/arm64
+            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-from: type=gha,scope=${{ matrix.tag }}
-          cache-to: type=gha,mode=max,scope=${{ matrix.variant }}
+          cache-to: type=gha,mode=max,scope=${{ matrix.tag }}

apps/api/src/index.ts

@@ -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

@@ -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 \
+    if [ "$GPU" = "true" ]; then \
-        Pillow numpy opencv-python-headless onnxruntime && \
+        /opt/venv/bin/pip install \
-    (/opt/venv/bin/pip install "rembg[cpu]" || echo "WARNING: rembg not installed") && \
+            Pillow numpy opencv-python-headless onnxruntime-gpu && \
-    (/opt/venv/bin/pip install realesrgan || echo "WARNING: realesrgan not installed") && \
+        (/opt/venv/bin/pip install rembg || echo "WARNING: rembg not installed") && \
-    (/opt/venv/bin/pip install paddlepaddle paddleocr || echo "WARNING: PaddleOCR not installed") && \
+        (/opt/venv/bin/pip install realesrgan \
-    (/opt/venv/bin/pip install mediapipe || echo "WARNING: mediapipe not installed") && \
+            --extra-index-url https://download.pytorch.org/whl/cu126 \
-    (/opt/venv/bin/pip install lama-cleaner || echo "WARNING: lama-cleaner not installed") \
+            || echo "WARNING: realesrgan not installed") && \
-; fi && rm -f /tmp/requirements.txt
+        (/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

@@ -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

@@ -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
+    # Signal readiness with GPU status
-    print(json.dumps({"ready": True}), file=sys.stderr, flush=True)
+    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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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,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";