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). Each tool belongs to exactly one bundle — no partial functionality. When a bundle is installed, all its tools work fully. When it's not installed, those tools are locked entirely.

Feature Name	Python Packages	Models	Tools Fully Enabled	Est. Size
Background Removal	rembg, onnxruntime(-gpu), mediapipe	birefnet-general-lite, blaze_face, face_landmarker	remove-background, passport-photo	~700 MB - 1 GB
Face Detection	mediapipe	blaze_face, face_landmarker	blur-faces, red-eye-removal, smart-crop	~200-300 MB
Object Eraser & Colorize	onnxruntime(-gpu)	LaMa ONNX, DDColor ONNX, OpenCV colorize	erase-object, colorize	~600-800 MB
Upscale & Enhance	torch, torchvision, realesrgan, codeformer-pip (--no-deps), gfpgan, basicsr, lpips	RealESRGAN x4plus, GFPGANv1.3, CodeFormer (.pth), facexlib, SCUNet, NAFNet	upscale, enhance-faces, noise-removal	~4-5 GB
Photo Restoration	onnxruntime(-gpu), mediapipe	LaMa ONNX, DDColor ONNX, CodeFormer ONNX, blaze_face, face_landmarker, OpenCV colorize	restore-photo	~800 MB - 1 GB
OCR	paddlepaddle(-gpu), paddleocr	PP-OCRv5 (7 models), PaddleOCR-VL 1.5	ocr	~3-4 GB

Notes:

• passport-photo is in the Background Removal bundle because it primarily needs rembg; mediapipe (for face landmarks) is included in the same bundle so the tool works fully
• noise-removal is in the Upscale & Enhance bundle because its quality/maximum tiers need PyTorch; all 4 tiers (including OpenCV-based quick/balanced) are locked until the bundle is installed
• ocr is fully locked until the OCR bundle is installed, including the Tesseract-based fast tier — this keeps the UX clean even though Tesseract is pre-installed in the base image
• restore-photo is its own bundle because it needs models from multiple domains (inpainting, face enhancement, colorization); all stages work when installed
• Some packages appear in multiple bundles (e.g., mediapipe in Background Removal, Face Detection, and Photo Restoration; onnxruntime in Background Removal, Object Eraser, and Photo Restoration). The install script skips already-installed packages — pip handles this naturally
• Some models appear in multiple bundles (e.g., blaze_face in both Background Removal and Face Detection). The install script skips already-downloaded model files

Bundle Dependencies

Background Removal ───── standalone
Face Detection ────────── standalone
Object Eraser & Colorize ── standalone
Upscale & Enhance ─────── standalone
Photo Restoration ─────── standalone
OCR ───────────────────── standalone

All bundles are independently installable. Shared packages (mediapipe, onnxruntime) and shared models (blaze_face, LaMa, etc.) are silently skipped if already present from another bundle.

Single Venv Strategy

The current architecture uses a single venv at /opt/venv (set via PYTHON_VENV_PATH). The bridge (bridge.ts) constructs ${venvPath}/bin/python3 — it can only point to one interpreter. Having two venvs (base at /opt/venv, features at /data/ai/venv/) is fragile: C extensions and entry points reference their venv prefix, and PYTHONPATH hacks break in practice.

Solution: Use a single venv on the persistent volume at /data/ai/venv/.

• The Dockerfile creates /opt/venv with base packages (numpy, Pillow, opencv) as before
• The entrypoint script bootstraps /data/ai/venv/ on first run by copying /opt/venv into it (fast file copy, ~300 MB)
• PYTHON_VENV_PATH is set to /data/ai/venv/ so the bridge uses it
• Feature installs add packages to this same venv
• On container update, the entrypoint checks if base package versions changed and updates the venv accordingly (pip install from wheel cache)

This gives us one venv with all packages, living on a persistent volume, bootstrapped from the image's base packages.

Persistent Storage

All AI data lives under /data/ai/ on the existing Docker volume (no docker-compose changes):

/data/ai/
  venv/           # Single Python virtual environment (bootstrapped from /opt/venv, extended by feature installs)
  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", "mediapipe==0.10.21"],
        "arm64": ["onnxruntime==1.20.1", "rembg[cpu]==2.0.62", "mediapipe==0.10.18"]
      },
      "pipFlags": {},
      "models": [
        {
          "id": "birefnet-general-lite",
          "downloadFn": "rembg_session",
          "args": ["birefnet-general-lite"]
        },
        {
          "id": "blaze-face-short-range",
          "url": "https://storage.googleapis.com/mediapipe-models/face_detector/blaze_face_short_range/float16/latest/blaze_face_short_range.tflite",
          "path": "mediapipe/blaze_face_short_range.tflite",
          "minSize": 100000
        },
        {
          "id": "face-landmarker",
          "url": "https://storage.googleapis.com/mediapipe-models/face_landmarker/face_landmarker/float16/latest/face_landmarker.task",
          "path": "mediapipe/face_landmarker.task",
          "minSize": 5000000
        }
      ],
      "enablesTools": ["remove-background", "passport-photo"]
    },
    "upscale-enhance": {
      "name": "Upscale & Enhance",
      "description": "AI upscaling, face enhancement, and noise removal",
      "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 },
        { "id": "scunet", "url": "https://github.com/cszn/KAIR/releases/download/v1.0/scunet_color_real_psnr.pth", "path": "scunet/scunet_color_real_psnr.pth", "minSize": 4000000 },
        { "id": "nafnet", "url": "hf://mikestealth/nafnet-models/NAFNet-SIDD-width64.pth", "path": "nafnet/NAFNet-SIDD-width64.pth", "minSize": 67000000 }
      ],
      "enablesTools": ["upscale", "enhance-faces", "noise-removal"]
    }
  }
}

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.

Uninstall and Shared Package Strategy

Bundles share Python packages (e.g., onnxruntime in Background Removal, Object Eraser, and Photo Restoration). Naively pip-uninstalling a bundle's packages could break other installed bundles.

v1 approach (simple): Uninstall removes model files and updates installed.json. Orphaned pip packages stay in the venv — they use disk but don't cause issues. A "Clean up" button in the AI Features settings panel rebuilds the venv from scratch: creates a fresh venv, installs only packages needed by currently-installed bundles, removes the old venv.

Future improvement: Reference counting — track which bundles need which packages, only remove packages exclusively owned by the target bundle.

Tool Route Registration for Uninstalled Features

Currently registerToolRoutes() either registers a route or doesn't (disabled tools get 404). For uninstalled AI features, we need routes that return a structured error instead of 404.

Solution: Register ALL tool routes always, add a pre-processing guard.

In tool-factory.ts, after settings validation (around line 198) and before calling config.process(), check feature installation status. The response follows the existing { error, details } shape used by formatZodErrors (apps/api/src/lib/errors.ts) and consumed by parseApiError (apps/web/src/lib/api.ts):

if (isAiTool(config.toolId) && !isFeatureInstalled(config.toolId)) {
  const bundle = getBundleForTool(config.toolId);
  return reply.status(501).send({
    error: "Feature not installed",
    code: "FEATURE_NOT_INSTALLED",
    feature: bundle.id,
    featureName: bundle.name,
    estimatedSize: bundle.estimatedSize,
  });
}

This also applies to:

• restore-photo.ts (uses its own route handler, not the factory)
• batch.ts — the getToolConfig(toolId) call at line 35 is the gating point. Add a feature-installed check alongside the existing 404 check.
• pipeline.ts — the pre-validation loops (lines 143-172 for execute, lines 441-462 for batch) already validate all tool IDs before processing starts. Extend to also check feature installation.

Frontend error detection: Extend parseApiError in apps/web/src/lib/api.ts to detect the FEATURE_NOT_INSTALLED code and return structured data (bundle id, name, size) instead of a plain error string. This enables use-tool-processor.ts and use-pipeline-processor.ts (both already use parseApiError) to trigger the install prompt rather than showing a generic error.

The global Fastify error handler in apps/api/src/index.ts (lines 41-51) provides a safety net — any unhandled Python import errors will produce structured JSON rather than crashing.

API Endpoints

New routes — read endpoint is public (no /admin/ prefix), mutation endpoints are admin-only:

GET  /api/v1/features
  Returns: list of all bundles with install status, sizes, enabled tools
  Auth: any authenticated user (read-only, needed by frontend for badges/tool page state)
  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"],
      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

All testing runs against Docker containers using the existing playwright.docker.config.ts and tests/e2e-docker/ infrastructure:

• Unit tests: Feature manifest parsing, version comparison logic, bundle dependency resolution (Vitest, excluded from e2e via vitest.config.ts)
• Integration tests: Install/uninstall API endpoints, status reporting, SSE progress (Vitest integration suite)
• E2e-docker tests: Add to tests/e2e-docker/ alongside existing fixes-verification.spec.ts:
  • Verify uninstalled AI tool returns 501 with FEATURE_NOT_INSTALLED code
  • Admin enables a feature from settings, tool page transitions from "not installed" to working
  • Non-admin sees "not enabled" message on uninstalled tool page
  • Feature install/uninstall round-trip
• Docker build test: Verify base image builds without ML packages, verify feature-manifest.json is present (CI, SKIP_MODEL_DOWNLOADS already exists)

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)

Frontend: Feature Status Propagation

The frontend needs to know which tools are installed for three purposes: tool grid badges, tool page state, and settings panel.

Features store (apps/web/src/stores/features-store.ts):

• Zustand store fetched on app load (like settings-store.ts)
• Calls GET /api/v1/features to get bundle statuses
• Provides a derived mapping: toolInstallStatus: Record<string, "installed" | "not_installed" | "installing"> (each tool maps to exactly one bundle, no partial states)
• Provides isToolInstalled(toolId): boolean and getBundleForTool(toolId): BundleInfo | null helpers
• Refreshes on install/uninstall completion

Tool grid integration:

• ToolCard checks isToolInstalled(tool.id) from the features store
• If not installed: show a download icon badge (similar to existing "Experimental" badge)
• The tool remains clickable (not disabled) — clicking navigates to the tool page where the install prompt appears
• PYTHON_SIDECAR_TOOLS constant is used to determine which tools are AI tools (only AI tools can be "not installed")

Tool page integration:

• ToolPage component checks feature status after the tool lookup
• If the user is admin and feature not installed: render FeatureInstallPrompt component instead of the normal tool UI
• If the user is non-admin and feature not installed: render "This feature is not enabled. Contact your administrator."
• The install prompt shows feature name, description, estimated size, and an "Enable" button
• After clicking "Enable": show progress bar with SSE-streamed progress, auto-transition to normal tool UI on completion

Development and Testing

All development and testing is done via Docker containers — the same environment users run. Build the image locally and run it with:

docker run -d --name ashim -p 1349:1349 -v ashim-data:/data ghcr.io/ashim-hq/ashim:latest

Auth can be disabled for development by passing -e AUTH_ENABLED=false.

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)