feat: unified Docker image with GPU auto-detection (#37)

Merge CPU, CUDA, and lite Docker images into a single unified image. One tag (latest) works on all platforms: amd64 (NVIDIA CUDA) and arm64 (CPU). GPU auto-detected at runtime. All ML models and packages baked in. Key changes: - Platform-conditional Dockerfile (nvidia/cuda on amd64, node on arm64) - tini as PID 1 for proper signal handling - Fix FILES_STORAGE_PATH data loss bug - Fix RealESRGAN upscaler (was broken, always fell back to Lanczos) - Fix PaddleOCR language codes and stdout corruption - Simplified CI/CD (single build, single tag) - Expanded model pre-download with verification - Shutdown timeout, improved health endpoint - Remove unused lama-cleaner
2026-08-03 07:46:42 +02:00 · 2026-04-10 13:21:06 +08:00
parent 7bc979f677
commit b0083e2b08
15 changed files with 374 additions and 310 deletions
@@ -1,12 +1,9 @@
 # syntax=docker/dockerfile:1
 # ============================================
-# Stirling Image - Production Dockerfile
-# Multi-stage build for single-container deployment
+# Stirling Image - Unified Production Dockerfile
+# Single image: GPU auto-detected on amd64, CPU on arm64
 # ============================================

-ARG VARIANT=full
-ARG GPU=false
-
 # ============================================
 # Stage 1: Build the frontend (Vite + React)
 # ============================================
@@ -39,25 +36,22 @@ RUN --mount=type=cache,id=turbo-cache,target=/app/.turbo \
    pnpm --filter @stirling-image/web build

 # ============================================
-# Stage 2: Base image selection
+# Stage 2: Platform-specific base images
 # ============================================
-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
+FROM node:22-bookworm AS base-linux-arm64
+FROM nvidia/cuda:12.6.3-runtime-ubuntu24.04 AS base-linux-amd64

 # ============================================
 # Stage 3: Production runtime
 # ============================================
-FROM production-base-${GPU} AS production
+ARG TARGETOS
+ARG TARGETARCH
+FROM base-${TARGETOS}-${TARGETARCH} AS production

-ARG VARIANT
-ARG GPU
+ARG TARGETARCH

-# Install Node.js when using CUDA base (node:22-bookworm already has it)
-RUN if [ "$GPU" = "true" ]; then \
+# Install Node.js on amd64 (CUDA base has no Node; arm64 base already has it)
+RUN if [ "$TARGETARCH" = "amd64" ]; then \
    apt-get update && apt-get install -y --no-install-recommends \
        curl ca-certificates gnupg && \
    mkdir -p /etc/apt/keyrings && \
@@ -71,67 +65,65 @@ RUN if [ "$GPU" = "true" ]; then \

 RUN corepack enable && corepack prepare pnpm@9.15.4 --activate

-# System dependencies shared by all variants
+# System dependencies (all platforms)
 RUN apt-get update && apt-get install -y --no-install-recommends \
+    tini \
    imagemagick \
    libraw-dev \
    potrace \
    curl \
    gosu \
    libheif-examples \
+    python3 python3-pip python3-venv python3-dev \
+    tesseract-ocr tesseract-ocr-eng tesseract-ocr-deu tesseract-ocr-fra tesseract-ocr-spa \
+    build-essential \
+    libgl1 libglib2.0-0 \
    && rm -rf /var/lib/apt/lists/*

-# Python/ML system dependencies (full variant only)
-RUN if [ "$VARIANT" = "full" ]; then \
-    apt-get update && apt-get install -y --no-install-recommends \
-        python3 python3-pip python3-venv python3-dev \
-        tesseract-ocr tesseract-ocr-eng tesseract-ocr-deu tesseract-ocr-fra tesseract-ocr-spa \
-        build-essential \
-        libgl1 libglib2.0-0 \
-    && rm -rf /var/lib/apt/lists/* \
+# Python venv - Layer 1: Base packages (rarely change, ~3 GB)
+RUN python3 -m venv /opt/venv && \
+    /opt/venv/bin/pip install --upgrade pip && \
+    /opt/venv/bin/pip install \
+        Pillow==11.1.0 \
+        numpy==1.26.4 \
+        opencv-python-headless==4.10.0.84
+
+# Platform-conditional ONNX runtime
+RUN if [ "$TARGETARCH" = "amd64" ]; then \
+    /opt/venv/bin/pip install onnxruntime-gpu==1.20.1 \
+; else \
+    /opt/venv/bin/pip install onnxruntime==1.20.1 \
 ; fi

-# 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 && \
-    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") && \
-        (/opt/venv/bin/pip install seam-carving || echo "WARNING: seam-carving 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") && \
-        (/opt/venv/bin/pip install seam-carving || echo "WARNING: seam-carving not installed") \
-    ; fi \
-; fi && rm -f /tmp/requirements.txt /tmp/requirements-gpu.txt
+# Python venv - Layer 2: Tool packages (change occasionally, ~2 GB)
+RUN if [ "$TARGETARCH" = "amd64" ]; then \
+    /opt/venv/bin/pip install rembg==2.0.62 && \
+    /opt/venv/bin/pip install realesrgan==0.3.0 \
+        --extra-index-url https://download.pytorch.org/whl/cu126 && \
+    /opt/venv/bin/pip install paddlepaddle-gpu==3.0.0 \
+        --extra-index-url https://www.paddlepaddle.org.cn/packages/stable/cu126/ && \
+    /opt/venv/bin/pip install paddleocr==2.9.1 \
+; else \
+    /opt/venv/bin/pip install "rembg[cpu]==2.0.62" && \
+    /opt/venv/bin/pip install realesrgan==0.3.0 && \
+    /opt/venv/bin/pip install paddlepaddle==3.0.0 paddleocr==2.9.1 \
+; fi

+# mediapipe 0.10.21 only has amd64 wheels; arm64 maxes out at 0.10.18
+RUN if [ "$TARGETARCH" = "amd64" ]; then \
+    /opt/venv/bin/pip install mediapipe==0.10.21 \
+; else \
+    /opt/venv/bin/pip install mediapipe==0.10.18 \
+; fi
+
+RUN /opt/venv/bin/pip install seam-carving==1.1.0
+
+# Pre-download and verify all ML models
+# Note: on amd64, paddlepaddle-gpu can't import without the CUDA driver (only
+# available at runtime). The download script gracefully skips PaddleOCR model
+# pre-download in this case; models download on first use at runtime instead.
 COPY docker/download_models.py /tmp/download_models.py
-RUN if [ "$VARIANT" = "full" ]; then \
-    /opt/venv/bin/python3 /tmp/download_models.py && \
-    /opt/venv/bin/python3 -c "\
-try: \
-    from paddleocr import PaddleOCR; \
-    print('Downloading PaddleOCR models...'); \
-    ocr = PaddleOCR(use_angle_cls=True, lang='en', show_log=False); \
-    print('PaddleOCR models ready'); \
-except: print('PaddleOCR model pre-download skipped') \
-" 2>/dev/null || echo "WARNING: Could not pre-download PaddleOCR models" \
-; fi && rm -f /tmp/download_models.py
+RUN /opt/venv/bin/python3 /tmp/download_models.py && rm -f /tmp/download_models.py

 WORKDIR /app

@@ -151,10 +143,8 @@ RUN --mount=type=cache,id=pnpm-store,target=/root/.local/share/pnpm/store/v3 \
    pnpm install --frozen-lockfile --prod

 # Remove build tools no longer needed in production
-RUN if [ "$VARIANT" = "full" ]; then \
-    apt-get purge -y --auto-remove build-essential python3-dev && \
-    rm -rf /var/lib/apt/lists/* \
-; fi
+RUN apt-get purge -y --auto-remove build-essential python3-dev && \
+    rm -rf /var/lib/apt/lists/*

 # Copy source code for API (tsx runs TS directly - no build step needed)
 COPY apps/api/src ./apps/api/src
@@ -170,9 +160,9 @@ COPY packages/ai/python ./packages/ai/python
 COPY --from=builder /app/apps/web/dist ./apps/web/dist

 # Create required directories
-RUN mkdir -p /data /tmp/workspace
+RUN mkdir -p /data /data/files /tmp/workspace

-# Environment defaults (matching PRD Section 16.1)
+# Environment defaults
 ENV PORT=1349 \
    NODE_ENV=production \
    AUTH_ENABLED=true \
@@ -181,6 +171,7 @@ ENV PORT=1349 \
    STORAGE_MODE=local \
    DB_PATH=/data/stirling.db \
    WORKSPACE_PATH=/tmp/workspace \
+    FILES_STORAGE_PATH=/data/files \
    PYTHON_VENV_PATH=/opt/venv \
    DEFAULT_THEME=light \
    DEFAULT_LOCALE=en \
@@ -191,13 +182,20 @@ ENV PORT=1349 \
    MAX_BATCH_SIZE=200 \
    CONCURRENT_JOBS=3 \
    MAX_MEGAPIXELS=100 \
-    RATE_LIMIT_PER_MIN=100 \
-    STIRLING_VARIANT=${VARIANT}
+    RATE_LIMIT_PER_MIN=100
+
+# NVIDIA Container Toolkit env vars (harmless on non-GPU systems)
+ENV NVIDIA_VISIBLE_DEVICES=all \
+    NVIDIA_DRIVER_CAPABILITIES=compute,utility
+
+# Suppress noisy ML library output in docker logs
+ENV PYTHONWARNINGS=ignore \
+    TF_CPP_MIN_LOG_LEVEL=3 \
+    PADDLE_PDX_DISABLE_MODEL_SOURCE_CHECK=True

 # Create non-root user for runtime
 RUN groupadd -r stirling && useradd -r -g stirling -d /app -s /sbin/nologin stirling
-RUN chown -R stirling:stirling /app /data /tmp/workspace && \
-    ([ -d /opt/venv ] && chown -R stirling:stirling /opt/venv || true)
+RUN chown -R stirling:stirling /app /data /tmp/workspace /opt/venv

 # Entrypoint fixes volume permissions then drops to stirling via gosu
 COPY docker/entrypoint.sh /usr/local/bin/entrypoint.sh
@@ -205,8 +203,9 @@ RUN chmod +x /usr/local/bin/entrypoint.sh

 EXPOSE 1349

-HEALTHCHECK --interval=30s --timeout=5s --start-period=15s --retries=3 \
+HEALTHCHECK --interval=30s --timeout=5s --start-period=60s --retries=3 \
    CMD curl -f http://localhost:1349/api/v1/health || exit 1

-ENTRYPOINT ["entrypoint.sh"]
+# tini as PID 1 for zombie reaping + signal forwarding
+ENTRYPOINT ["tini", "--", "entrypoint.sh"]
 CMD ["npx", "tsx", "apps/api/src/index.ts"]
@@ -1,17 +0,0 @@
-# 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]
@@ -1,32 +1,25 @@
-name: stirling-image
-
 services:
  stirling-image:
    build:
      context: ..
      dockerfile: docker/Dockerfile
+    image: stirlingimage/stirling-image:latest
    container_name: stirling-image
    ports:
      - "1349:1349"
+    volumes:
+      - stirling-data:/data
+      - stirling-workspace:/tmp/workspace
    environment:
      - AUTH_ENABLED=true
      - DEFAULT_USERNAME=admin
      - DEFAULT_PASSWORD=admin
-      - STORAGE_MODE=local
-      - FILE_MAX_AGE_HOURS=24
-      - CLEANUP_INTERVAL_MINUTES=30
-      - MAX_UPLOAD_SIZE_MB=100
-      - MAX_BATCH_SIZE=200
-      - CONCURRENT_JOBS=3
-      - MAX_MEGAPIXELS=100
-      - RATE_LIMIT_PER_MIN=100
-      - DEFAULT_THEME=light
-      - DEFAULT_LOCALE=en
-      - APP_NAME=Stirling Image
-    volumes:
-      - stirling-data:/data
-      - stirling-workspace:/tmp/workspace
    restart: unless-stopped
+    logging:
+      driver: json-file
+      options:
+        max-size: "10m"
+        max-file: "3"

 volumes:
  stirling-data:
@@ -1,7 +1,26 @@
-"""Pre-download all rembg models offered in the UI."""
-import sys
+"""Pre-download and verify all ML models for the Docker image.

-MODELS = [
+This script runs at Docker build time. Any failure exits non-zero,
+failing the build. No silent fallbacks.
+"""
+import os
+import sys
+import urllib.request
+
+# Force CPU mode during build - no GPU driver available at build time.
+# Must be set before any ML library import.
+os.environ["PADDLE_DEVICE"] = "cpu"
+os.environ["FLAGS_use_cuda"] = "0"
+os.environ["CUDA_VISIBLE_DEVICES"] = ""
+
+REALESRGAN_MODEL_DIR = "/opt/models/realesrgan"
+REALESRGAN_MODEL_URL = (
+    "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth"
+)
+REALESRGAN_MODEL_PATH = os.path.join(REALESRGAN_MODEL_DIR, "RealESRGAN_x4plus.pth")
+REALESRGAN_MIN_SIZE = 60_000_000  # ~67 MB
+
+REMBG_MODELS = [
    "u2net",
    "isnet-general-use",
    "bria-rmbg",
@@ -10,18 +29,116 @@ MODELS = [
    "birefnet-general",
 ]

-try:
-    from rembg import new_session
-except ImportError:
-    print("WARNING: rembg not installed, skipping model pre-download")
-    sys.exit(0)
+# PaddleOCR language codes (not ISO). German/French/Spanish use "latin" model.
+# Valid keys: ch, en, korean, japan, chinese_cht, ta, te, ka, latin, arabic, cyrillic, devanagari
+PADDLEOCR_LANGUAGES = ["en", "ch", "japan", "korean", "latin"]

-for model in MODELS:
-    print(f"Downloading {model}...")
-    try:
+
+def download_rembg_models():
+    """Download all rembg ONNX models."""
+    print("=== Downloading rembg models ===")
+    from rembg import new_session
+
+    for model in REMBG_MODELS:
+        print(f"  Downloading {model}...")
        new_session(model)
        print(f"  {model} ready")
-    except Exception as e:
-        print(f"  WARNING: {model} failed: {e}")
+    print(f"All {len(REMBG_MODELS)} rembg models downloaded.\n")

-print("Model pre-download complete")
+
+def download_realesrgan_model():
+    """Download RealESRGAN_x4plus.pth pretrained weights."""
+    print("=== Downloading RealESRGAN model ===")
+    os.makedirs(REALESRGAN_MODEL_DIR, exist_ok=True)
+    print(f"  Downloading from {REALESRGAN_MODEL_URL}...")
+    urllib.request.urlretrieve(REALESRGAN_MODEL_URL, REALESRGAN_MODEL_PATH)
+
+    size = os.path.getsize(REALESRGAN_MODEL_PATH)
+    assert size > REALESRGAN_MIN_SIZE, (
+        f"RealESRGAN model too small: {size} bytes (expected > {REALESRGAN_MIN_SIZE})"
+    )
+    print(f"  RealESRGAN_x4plus.pth downloaded ({size / 1_000_000:.1f} MB)\n")
+
+
+def download_paddleocr_models():
+    """Pre-download PaddleOCR models for all supported languages."""
+    print("=== Downloading PaddleOCR models ===")
+    try:
+        from paddleocr import PaddleOCR
+    except ImportError as e:
+        if "libcuda" in str(e):
+            # paddlepaddle-gpu can't import without CUDA driver at build time.
+            # Models will be downloaded on first use at runtime instead.
+            print(f"  Skipping PaddleOCR model pre-download (no CUDA driver at build time)")
+            print(f"  Models will download on first use at runtime.\n")
+            return
+        raise
+
+    for lang in PADDLEOCR_LANGUAGES:
+        print(f"  Downloading models for lang={lang}...")
+        PaddleOCR(lang=lang, use_gpu=False, show_log=False)
+        print(f"  {lang} ready")
+    print(f"All {len(PADDLEOCR_LANGUAGES)} PaddleOCR languages downloaded.\n")
+
+
+def verify_mediapipe():
+    """Verify MediaPipe face detection models are bundled in the wheel."""
+    print("=== Verifying MediaPipe models ===")
+    import mediapipe as mp
+
+    for selection in [0, 1]:
+        label = "short-range" if selection == 0 else "full-range"
+        print(f"  Verifying {label} model (selection={selection})...")
+        detector = mp.solutions.face_detection.FaceDetection(
+            model_selection=selection, min_detection_confidence=0.5
+        )
+        detector.close()
+        print(f"  {label} model OK")
+    print("MediaPipe models verified.\n")
+
+
+def smoke_test():
+    """Final verification that all ML libraries and models are loadable.
+
+    GPU-dependent libraries (paddlepaddle-gpu, torch CUDA) cannot be imported
+    at build time because the CUDA driver is only available at runtime. We verify
+    CPU-only imports and check that model files exist on disk.
+    """
+    print("=== Running smoke test ===")
+
+    # CPU-only imports that work on all platforms at build time
+    from PIL import Image
+    import cv2
+    import numpy
+    import seam_carving
+    from rembg import new_session
+    print("  CPU imports OK (Pillow, cv2, numpy, seam_carving, rembg)")
+
+    # MediaPipe is CPU-only, should always import
+    import mediapipe as mp
+    print("  MediaPipe import OK")
+
+    # RealESRGAN model file must exist
+    assert os.path.exists(REALESRGAN_MODEL_PATH), (
+        f"RealESRGAN model missing: {REALESRGAN_MODEL_PATH}"
+    )
+    assert os.path.getsize(REALESRGAN_MODEL_PATH) > REALESRGAN_MIN_SIZE, (
+        "RealESRGAN model file is too small"
+    )
+    print("  RealESRGAN model file verified")
+
+    print("Smoke test passed.\n")
+
+
+def main():
+    print("Pre-downloading all ML models...\n")
+    download_rembg_models()
+    download_realesrgan_model()
+    download_paddleocr_models()
+    verify_mediapipe()
+    smoke_test()
+    print("All models downloaded and verified.")
+
+
+if __name__ == "__main__":
+    main()