SnapOtter/docs/superpowers/specs/2026-04-17-on-demand-ai-features-design.md

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:

{
  "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)