docs: add design spec for on-demand AI feature downloads

Reduce Docker image from ~30GB to ~5-6GB by making AI features downloadable post-install. Users cherry-pick feature bundles (Background Removal, OCR, etc.) from the UI after pulling.
2026-08-03 07:46:42 +02:00 · 2026-04-17 11:31:35 +08:00
parent b76b9a1878
commit b9fbf9db67
1 changed files with 354 additions and 0 deletions
@@ -0,0 +1,354 @@
+# On-Demand AI Feature Downloads
+
+**Date:** 2026-04-17
+**Status:** Approved
+**Goal:** Reduce Docker image from ~30 GB to ~5-6 GB (amd64) / ~2-3 GB (arm64) by making AI features downloadable post-install.
+
+## Problem
+
+The Docker image bundles all Python ML packages (~8-10 GB) and model weights (~5-8 GB) regardless of whether users need AI features. Users who only want basic image tools (resize, crop, convert) must pull ~30 GB.
+
+## Design Decisions
+
+- **Single Docker image** — no lite/full variants
+- **Individual feature bundles** — users cherry-pick by feature name, not model name
+- **Admin-only downloads** — only admins can enable/disable AI features
+- **AI tools visible with badge** — uninstalled tools appear in grid with a download indicator
+- **Both tool-page and settings UI** — admins can download from the tool page or from a central management panel in settings
+
+## Architecture
+
+### Base Image Contents
+
+The base image includes everything needed for non-AI tools plus the prerequisites for AI feature installation:
+
+| Component | Rationale |
+|-----------|-----------|
+| Node.js 22 + pnpm + app source + frontend dist | Core application |
+| Sharp, imagemagick, tesseract-ocr, potrace, libheif, exiftool | Non-AI image processing |
+| caire binary | Content-aware resize |
+| Python 3 + pip + build-essential | Required for pip install at runtime |
+| numpy==1.26.4, Pillow, opencv-python-headless | Shared by all AI features, small (~300 MB) |
+| CUDA runtime (amd64 only, from nvidia/cuda base) | Required for GPU-accelerated AI |
+
+**Estimated size:** ~5-6 GB (amd64), ~2-3 GB (arm64)
+
+### Feature Bundles
+
+Six user-facing bundles, named by what they enable (not by model names):
+
+| Feature Name | Python Packages | Models | Tools Enabled | Est. Size |
+|---|---|---|---|---|
+| **Background Removal** | rembg, onnxruntime(-gpu) | birefnet-general-lite (default) | remove-background, passport-photo (partial) | ~500-700 MB |
+| **Face Detection** | mediapipe | blaze_face, face_landmarker | blur-faces, red-eye-removal, smart-crop, passport-photo (partial) | ~200-300 MB |
+| **Object Eraser & Colorize** | onnxruntime(-gpu) if not already installed | LaMa ONNX, DDColor ONNX, OpenCV colorize | erase-object, colorize, restore-photo (partial) | ~600-800 MB |
+| **Upscale & Face Enhance** | torch, torchvision, realesrgan, codeformer-pip (--no-deps), gfpgan, basicsr, lpips | RealESRGAN x4plus, GFPGANv1.3, CodeFormer (.pth + .onnx), facexlib models | upscale, enhance-faces, restore-photo (partial) | ~4-5 GB |
+| **OCR** | paddlepaddle(-gpu), paddleocr | PP-OCRv5 (7 models), PaddleOCR-VL 1.5 | ocr (balanced + best tiers) | ~3-4 GB |
+| **Advanced Noise Removal** | _(requires Upscale bundle for torch)_ | SCUNet, NAFNet | noise-removal (quality + maximum tiers) | ~100 MB |
+
+Notes:
+- `passport-photo` needs both Background Removal + Face Detection
+- `restore-photo` needs Object Eraser & Colorize + optionally Upscale & Face Enhance (for face restoration step)
+- `noise-removal` quick/balanced tiers work without any bundle (uses OpenCV)
+- `ocr` fast tier works without any bundle (uses Tesseract, pre-installed in base)
+- Advanced Noise Removal depends on the Upscale & Face Enhance bundle (shared PyTorch dependency)
+
+### Bundle Dependencies
+
+```
+Background Removal ─── standalone
+Face Detection ─────── standalone
+Object Eraser & Colorize ── standalone (uses onnxruntime from Background Removal if installed, otherwise installs it)
+Upscale & Face Enhance ─── standalone
+OCR ────────────────── standalone
+Advanced Noise Removal ─── depends on "Upscale & Face Enhance" (for PyTorch)
+```
+
+onnxruntime is needed by both Background Removal and Object Eraser & Colorize. The install script installs it with the first bundle that needs it, and skips it for subsequent bundles.
+
+### Persistent Storage
+
+All AI data lives under `/data/ai/` on the existing Docker volume (no docker-compose changes):
+
+```
+/data/ai/
+  venv/           # Python virtual environment with installed packages
+  models/         # Downloaded model weight files (same structure as /opt/models/)
+  pip-cache/      # Wheel cache for fast re-installs after updates
+  installed.json  # Tracks installed bundles, versions, timestamps
+```
+
+### Feature Manifest
+
+A `feature-manifest.json` file is baked into each Docker image at build time. It is the single source of truth for what each bundle installs:
+
+```json
+{
+  "manifestVersion": 1,
+  "imageVersion": "1.16.0",
+  "pythonVersion": "3.12",
+  "basePackages": ["numpy==1.26.4", "Pillow==11.1.0", "opencv-python-headless==4.10.0.84"],
+  "bundles": {
+    "background-removal": {
+      "name": "Background Removal",
+      "description": "Remove image backgrounds with AI",
+      "packages": {
+        "common": ["rembg==2.0.62"],
+        "amd64": ["onnxruntime-gpu==1.20.1"],
+        "arm64": ["onnxruntime==1.20.1", "rembg[cpu]==2.0.62"]
+      },
+      "pipFlags": {},
+      "models": [
+        {
+          "id": "birefnet-general-lite",
+          "name": "Default model",
+          "required": true,
+          "downloadFn": "rembg_session",
+          "args": ["birefnet-general-lite"]
+        }
+      ],
+      "optionalModels": [
+        {
+          "id": "u2net",
+          "name": "U2-Net (lightweight)",
+          "downloadFn": "rembg_session",
+          "args": ["u2net"]
+        },
+        {
+          "id": "birefnet-general",
+          "name": "BiRefNet General (high quality)",
+          "downloadFn": "rembg_session",
+          "args": ["birefnet-general"]
+        }
+      ],
+      "enablesTools": ["remove-background"],
+      "partialTools": ["passport-photo"]
+    },
+    "upscale-enhance": {
+      "name": "Upscale & Face Enhance",
+      "packages": {
+        "common": ["codeformer-pip==0.0.4", "lpips"],
+        "amd64": [
+          "torch torchvision --extra-index-url https://download.pytorch.org/whl/cu126",
+          "realesrgan==0.3.0 --extra-index-url https://download.pytorch.org/whl/cu126"
+        ],
+        "arm64": ["torch", "torchvision", "realesrgan==0.3.0"]
+      },
+      "pipFlags": {
+        "codeformer-pip==0.0.4": "--no-deps"
+      },
+      "postInstall": ["pip install numpy==1.26.4"],
+      "models": [
+        { "id": "realesrgan-x4plus", "url": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth", "path": "realesrgan/RealESRGAN_x4plus.pth", "minSize": 67000000 },
+        { "id": "gfpgan-v1.3", "url": "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth", "path": "gfpgan/GFPGANv1.3.pth", "minSize": 332000000 },
+        { "id": "codeformer-pth", "url": "https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth", "path": "codeformer/codeformer.pth", "minSize": 375000000 },
+        { "id": "codeformer-onnx", "url": "hf://facefusion/models-3.0.0/codeformer.onnx", "path": "codeformer/codeformer.onnx", "minSize": 377000000 },
+        { "id": "facexlib-detection", "url": "https://github.com/xinntao/facexlib/releases/download/v0.1.0/detection_Resnet50_Final.pth", "path": "gfpgan/facelib/detection_Resnet50_Final.pth", "minSize": 104000000 },
+        { "id": "facexlib-parsing", "url": "https://github.com/xinntao/facexlib/releases/download/v0.2.2/parsing_parsenet.pth", "path": "gfpgan/facelib/parsing_parsenet.pth", "minSize": 85000000 }
+      ],
+      "enablesTools": ["upscale", "enhance-faces"],
+      "partialTools": ["restore-photo"]
+    }
+  }
+}
+```
+
+### Install Script
+
+A Python script (`packages/ai/python/install_feature.py`) handles feature installation:
+
+1. Reads the feature manifest from the image
+2. Detects architecture (amd64/arm64) and GPU availability
+3. Creates or reuses the venv at `/data/ai/venv/`
+4. Runs pip install with the correct packages, flags, and index URLs per platform
+5. Handles the numpy version conflict (--no-deps for codeformer, re-pin numpy)
+6. Downloads model weights with retry logic (ported from `download_models.py`)
+7. Updates `/data/ai/installed.json` with bundle status
+8. Reports progress to stdout as JSON lines (consumed by the Node bridge)
+
+The script must be idempotent — running it twice for the same bundle is a no-op.
+
+### API Endpoints
+
+New routes under `/api/v1/admin/features/`:
+
+```
+GET  /api/v1/admin/features
+  Returns: list of all bundles with install status, sizes, enabled tools
+  Auth: any authenticated user (read-only)
+  Response: {
+    bundles: [{
+      id: "background-removal",
+      name: "Background Removal",
+      description: "Remove image backgrounds with AI",
+      status: "not_installed" | "installing" | "installed" | "error",
+      installedVersion: "1.15.3" | null,
+      estimatedSize: "500-700 MB",
+      enablesTools: ["remove-background"],
+      partialTools: ["passport-photo"],
+      progress: { percent: 45, stage: "Downloading models..." } | null,
+      error: "pip install failed: ..." | null,
+      dependencies: [] | ["upscale-enhance"]
+    }]
+  }
+
+POST /api/v1/admin/features/:bundleId/install
+  Starts background installation of a feature bundle.
+  Auth: admin only
+  Response: { jobId: "uuid" }
+  SSE progress at: GET /api/v1/jobs/:jobId/progress
+
+POST /api/v1/admin/features/:bundleId/uninstall
+  Removes a feature bundle (pip packages + models).
+  Auth: admin only
+  Response: { ok: true, freedSpace: "500 MB" }
+
+GET  /api/v1/admin/features/disk-usage
+  Returns total disk usage of /data/ai/.
+  Auth: admin only
+  Response: { totalBytes: 5368709120, byBundle: { "background-removal": 734003200, ... } }
+```
+
+### Background Job Mechanism
+
+Feature installation runs as a background child process (not inline with the HTTP request):
+
+1. `POST /admin/features/:bundleId/install` spawns the install script as a child process
+2. Progress is streamed via stderr JSON lines → captured by the Node process → pushed to SSE listeners
+3. The existing SSE infrastructure (`/api/v1/jobs/:jobId/progress`) is reused
+4. Job status is persisted to the `jobs` table for recovery on restart
+5. Only one install can run at a time (mutex). Concurrent install requests return 409 Conflict.
+
+### Python Sidecar Changes
+
+**dispatcher.py:**
+- On startup, read `/data/ai/installed.json` to know which features are available
+- Populate `available_modules` based on what's actually installed
+- When a script is requested for an uninstalled feature, return a structured error: `{"error": "feature_not_installed", "feature": "background-removal", "message": "Background Removal is not installed"}`
+- After a feature is installed, the dispatcher must be restarted (or sent a reload signal) to pick up new packages. The bridge handles this by killing and re-spawning the dispatcher.
+
+**Python scripts:**
+- Convert hard module-level imports in `colorize.py` and `restore.py` to lazy imports inside functions
+- All scripts should check for their feature's models and return a clear "not installed" error if missing
+- The `sys.path` must include `/data/ai/venv/lib/python3.X/site-packages/` (set by the dispatcher on startup based on installed.json)
+
+**Bridge (bridge.ts):**
+- Update `PYTHON_VENV_PATH` logic to prefer `/data/ai/venv/` when it exists
+- Add a `restartDispatcher()` function called after feature install completes
+- Handle the new `feature_not_installed` error type from the dispatcher
+
+### Model Path Resolution
+
+Currently models are at `/opt/models/`. With on-demand downloads, they'll be at `/data/ai/models/`. The resolution order:
+
+1. `/opt/models/<model>` (Docker-baked, for backwards compatibility if someone builds a full image)
+2. `/data/ai/models/<model>` (on-demand download location)
+3. `~/.cache/ashim/<model>` (local dev fallback)
+
+Environment variables (`U2NET_HOME`, etc.) are updated by the install script to point to `/data/ai/models/`.
+
+### Dockerfile Changes
+
+1. Remove all `pip install` commands for ML packages (lines 175-206)
+2. Remove `download_models.py` COPY and RUN (lines 219-231)
+3. Keep: Python 3 + pip + build-essential (do NOT purge build-essential)
+4. Keep: numpy, Pillow, opencv-python-headless install (lightweight shared deps)
+5. Add: COPY `feature-manifest.json` into the image
+6. Add: COPY `install_feature.py` into the image
+7. Update entrypoint to set up `/data/ai/` directory structure on first run
+8. Update env vars: `MODELS_PATH=/data/ai/models` as default, fallback to `/opt/models`
+
+### Frontend: Tool Page (Uninstalled State)
+
+When a user navigates to an AI tool that isn't installed:
+
+**For admins:**
+- Show a card replacing the normal upload area:
+  - Feature icon + name (e.g., "Background Removal")
+  - "This feature requires an additional download (~500-700 MB)"
+  - [Enable Feature] button
+  - After clicking: progress bar with stage text, estimated time
+  - On completion: page automatically transitions to the normal tool UI
+
+**For non-admins:**
+- Show: "This feature is not enabled. Ask your administrator to enable it in Settings."
+
+### Frontend: Tool Grid (Badge)
+
+AI tools in the grid show a small download icon overlay when not installed. When installed, the icon disappears and the tool looks like any other tool.
+
+Tools with partial dependencies (e.g., passport-photo needs 2 bundles) show the badge until ALL required bundles are installed.
+
+### Frontend: Settings Panel
+
+New "AI Features" section in the settings dialog (admin only):
+
+- List of all 6 feature bundles as cards
+- Each card shows: name, description, status (installed/not installed/installing), disk usage
+- Install/Uninstall buttons per bundle
+- "Install All" button at the top
+- Total AI disk usage summary at the bottom
+- Progress bar during installation
+- Dependency warnings (e.g., "Advanced Noise Removal requires Upscale & Face Enhance")
+
+### Container Update Flow
+
+When a user does `docker pull` + restart:
+
+1. **Pull:** Only app code layers changed → ~50-100 MB download
+2. **Startup:** Backend reads feature manifest from new image + installed.json from volume
+3. **Comparison:**
+   - If bundle package versions unchanged → no action, instant startup
+   - If a package version bumped → `pip install --upgrade` from wheel cache (seconds)
+   - If a model URL/version changed → re-download that model only
+   - If Python major version changed → rebuild venv from cached wheels (rare, ~2-5 min)
+4. **Dispatcher restart** if any packages changed
+
+This check runs at startup, not blocking the HTTP server. AI features show "Updating..." status until the check completes.
+
+### Error Handling
+
+| Scenario | Behavior |
+|---|---|
+| No internet during install | Error with clear message: "Could not download packages. Check your internet connection." |
+| Partial install (interrupted) | On next install attempt, detect incomplete state and resume/retry |
+| Disk full | Error with disk usage info: "Not enough disk space. Need ~500 MB, only 200 MB available." |
+| pip install failure | Error with the pip output. Bundle marked as "error" status, admin can retry. |
+| Model download failure | Retry 3 times with exponential backoff. On final failure, mark bundle as partially installed (packages OK, models missing). |
+| Container update breaks venv | Version manifest comparison detects mismatch, triggers venv rebuild from wheel cache |
+
+### Testing Strategy
+
+- **Unit tests:** Feature manifest parsing, version comparison logic, bundle dependency resolution
+- **Integration tests:** Install/uninstall API endpoints, status reporting, SSE progress
+- **E2E tests:** Admin enables a feature from settings, tool page transitions from "not installed" to working
+- **Docker build test:** Verify base image builds without ML packages, verify feature-manifest.json is present
+- **Install script test:** Run install script in a clean container, verify packages and models are correctly installed
+
+### Migration Path
+
+Since the new image is fundamentally different (no ML packages baked in), existing users upgrading from the full image will need to re-download their AI features. The Python ML packages are no longer in the system venv, so even if old model weights exist at `/opt/models/`, the features won't work without packages.
+
+The first-run experience for upgrading users:
+
+1. Detect this is an upgrade: no `/data/ai/installed.json` exists, but user data exists in `/data`
+2. Show a one-time banner in the UI: "We've reduced the image size from 30 GB to 5 GB! AI features are now downloaded on-demand. Visit Settings → AI Features to enable the ones you need."
+3. No automatic downloads — let the admin choose what to install
+4. Old model weights at `/opt/models/` are ignored (they won't exist in the new image anyway since that layer is removed)
+
+### Scope Boundaries
+
+**In scope:**
+- Dockerfile restructuring to remove ML packages and models
+- Feature manifest system
+- Install/uninstall API + background job
+- Python sidecar changes for dynamic feature detection
+- Frontend: tool page download prompt, grid badge, settings panel
+- Container update handling with version manifest
+
+**Out of scope (future work):**
+- Additional rembg model variants as sub-downloads within Background Removal
+- Automatic feature recommendations based on usage
+- Download from private/custom model registries
+- Bandwidth throttling for downloads
+- Multiple venv support (e.g., different Python versions)