diff --git a/.github/workflows/ai-bundles.yml b/.github/workflows/ai-bundles.yml new file mode 100644 index 00000000..932eadd9 --- /dev/null +++ b/.github/workflows/ai-bundles.yml @@ -0,0 +1,163 @@ +name: AI Bundles + +on: + workflow_dispatch: + inputs: + version: + description: "Version to build bundles for (e.g., 2.0.0)" + required: true + type: string + workflow_call: + inputs: + version: + type: string + required: true + +concurrency: + group: ai-bundles-${{ inputs.version }} + cancel-in-progress: false + +permissions: {} + +jobs: + build: + name: ${{ matrix.bundle }} (${{ matrix.arch }}) + strategy: + fail-fast: false + matrix: + include: + - bundle: background-removal + arch: amd64-gpu + runner: ubuntu-latest + - bundle: background-removal + arch: arm64-cpu + runner: ubuntu-24.04-arm + - bundle: face-detection + arch: amd64-gpu + runner: ubuntu-latest + - bundle: face-detection + arch: arm64-cpu + runner: ubuntu-24.04-arm + - bundle: object-eraser-colorize + arch: amd64-gpu + runner: ubuntu-latest + - bundle: object-eraser-colorize + arch: arm64-cpu + runner: ubuntu-24.04-arm + - bundle: upscale-enhance + arch: amd64-gpu + runner: ubuntu-latest + - bundle: upscale-enhance + arch: arm64-cpu + runner: ubuntu-24.04-arm + - bundle: photo-restoration + arch: amd64-gpu + runner: ubuntu-latest + - bundle: photo-restoration + arch: arm64-cpu + runner: ubuntu-24.04-arm + - bundle: ocr + arch: amd64-gpu + runner: ubuntu-latest + - bundle: ocr + arch: arm64-cpu + runner: ubuntu-24.04-arm + - bundle: transcription + arch: amd64-gpu + runner: ubuntu-latest + - bundle: transcription + arch: arm64-cpu + runner: ubuntu-24.04-arm + runs-on: ${{ matrix.runner }} + permissions: + contents: read + packages: read + steps: + - name: Free disk space + run: | + sudo rm -rf /usr/share/dotnet /usr/local/lib/android /opt/ghc \ + /usr/local/share/boost /opt/hostedtoolcache/CodeQL + sudo docker system prune -af + df -h / + + - name: Checkout + uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + ref: v${{ inputs.version }} + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@d7f5e7f509e45cec5c76c4d5afdd7de93d0b3df5 # v4.1.0 + + - name: Log in to GHCR + uses: docker/login-action@650006c6eb7dba73a995cc03b0b2d7f5ca915bee # v4.2.0 + with: + registry: ghcr.io + username: ${{ github.repository_owner }} + password: ${{ secrets.GHCR_TOKEN }} + + - name: Build bundle inside Docker image + env: + BUNDLE: ${{ matrix.bundle }} + ARCH: ${{ matrix.arch }} + VERSION: ${{ inputs.version }} + run: | + mkdir -p /tmp/bundles + docker run --rm --entrypoint bash \ + -v "$PWD/docker/build-bundle.sh:/build-bundle.sh:ro" \ + -v "/tmp/bundles:/output" \ + "ghcr.io/snapotter-hq/snapotter:${VERSION}" \ + /build-bundle.sh "${BUNDLE}" "${ARCH}" /output + + - name: Upload artifact + uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2 + with: + name: ${{ matrix.bundle }}-${{ matrix.arch }} + path: /tmp/bundles/ + retention-days: 1 + + publish: + name: Publish to HuggingFace + needs: build + runs-on: ubuntu-latest + permissions: + contents: read + steps: + - name: Download all artifacts + uses: actions/download-artifact@d3f86a106a0bac45b974a628896c90dbdf5c8093 # v4.3.0 + with: + path: /tmp/artifacts + + - name: Install huggingface-hub + run: pip install huggingface-hub + + - name: Organize and upload to HuggingFace + env: + HF_TOKEN: ${{ secrets.HF_TOKEN }} + VERSION: ${{ inputs.version }} + run: | + mkdir -p "/tmp/upload/v${VERSION}" + + find /tmp/artifacts -name '*.tar.gz' -exec cp {} "/tmp/upload/v${VERSION}/" \; + find /tmp/artifacts -name '*.sha256' -exec cp {} "/tmp/upload/v${VERSION}/" \; + + cd "/tmp/upload/v${VERSION}" + echo '{"version":"'"${VERSION}"'","archives":[' > manifest.json + first=true + for archive in *.tar.gz; do + [ -f "$archive" ] || continue + sha=$(sha256sum "$archive" | cut -d' ' -f1) + size=$(stat --format=%s "$archive") + if [ "$first" = true ]; then + first=false + else + echo ',' >> manifest.json + fi + printf '{"file":"%s","sha256":"%s","size":%s}' "$archive" "$sha" "$size" >> manifest.json + done + echo ']}' >> manifest.json + + huggingface-cli upload snapotter/feature-bundles \ + "/tmp/upload/v${VERSION}" "v${VERSION}" \ + --repo-type model --token "${HF_TOKEN}" + + echo "Published: https://huggingface.co/snapotter/feature-bundles/tree/main/v${VERSION}" diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 4bffde1a..f9aabc02 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -376,6 +376,15 @@ jobs: "snapotter-v${VERSION}-sbom.spdx.json" \ --clobber --repo snapotter-hq/SnapOtter + ai-bundles: + name: AI Bundles + needs: [release, docker] + if: needs.release.outputs.new_version + uses: ./.github/workflows/ai-bundles.yml + with: + version: ${{ needs.release.outputs.new_version }} + secrets: inherit + manifest: name: Create Multi-Arch Manifests needs: [release, docker, scan] diff --git a/apps/api/src/lib/feature-status.ts b/apps/api/src/lib/feature-status.ts index 912cc369..17567be9 100644 --- a/apps/api/src/lib/feature-status.ts +++ b/apps/api/src/lib/feature-status.ts @@ -1,3 +1,4 @@ +import { execFileSync } from "node:child_process"; import { randomUUID } from "node:crypto"; import { constants, @@ -378,6 +379,40 @@ export function recoverInterruptedInstalls(): void { } } + // 6. Delete staging-{bundleId}/ directories (incomplete extraction) + try { + const aiEntries = readdirSync(AI_DIR, { withFileTypes: true }); + for (const entry of aiEntries) { + if (entry.isDirectory() && entry.name.startsWith("staging-")) { + const stagingPath = join(AI_DIR, entry.name); + rmSync(stagingPath, { recursive: true, force: true }); + console.info(`[feature-status] Deleted orphaned ${entry.name}/`); + } + } + } catch { + // AI_DIR may not exist yet + } + + // 7. Clean up staging/ download directory (partial downloads, orphaned tars) + const downloadStaging = join(AI_DIR, "staging"); + if (existsSync(downloadStaging)) { + try { + const files = readdirSync(downloadStaging); + for (const file of files) { + const filePath = join(downloadStaging, file); + if (file.endsWith(".partial") || file.endsWith(".meta")) { + unlinkSync(filePath); + console.info(`[feature-status] Deleted stale download file: ${file}`); + } else if (file.endsWith(".tar.gz")) { + unlinkSync(filePath); + console.info(`[feature-status] Deleted orphaned archive: ${file}`); + } + } + } catch { + // best-effort + } + } + invalidateCache(); } @@ -598,6 +633,50 @@ export async function importBundleArchive( moveTreeRecursive(stagingModels, MODELS_DIR); } + // Move site-packages/* into venv site-packages + const stagingSitePackages = join(stagingDir, "site-packages"); + if (existsSync(stagingSitePackages)) { + const venvPath = process.env.PYTHON_VENV_PATH || join(AI_DIR, "venv"); + let sitePackagesDir = ""; + const libDir = join(venvPath, "lib"); + if (existsSync(libDir)) { + const pyDirs = readdirSync(libDir).filter((d) => d.startsWith("python")); + if (pyDirs.length > 0) { + sitePackagesDir = join(libDir, pyDirs[0], "site-packages"); + } + } + if (sitePackagesDir && existsSync(sitePackagesDir)) { + moveTreeRecursive(stagingSitePackages, sitePackagesDir); + } + } + + // Apply fixups (NCCL wheel) if present + const stagingFixups = join(stagingDir, "fixups"); + if (existsSync(stagingFixups)) { + const wheels = readdirSync(stagingFixups).filter((f) => f.endsWith(".whl")); + if (wheels.length > 0) { + const venvPython = `${process.env.PYTHON_VENV_PATH || join(AI_DIR, "venv")}/bin/python3`; + for (const wheel of wheels) { + try { + execFileSync( + venvPython, + [ + "-m", + "pip", + "install", + "--no-index", + `--find-links=${stagingFixups}`, + wheel.split("-")[0], + ], + { stdio: "ignore", timeout: 30_000 }, + ); + } catch { + // Non-fatal + } + } + } + } + markInstalled(descriptor.bundleId, descriptor.version, descriptor.models); return { diff --git a/docker/build-bundle.sh b/docker/build-bundle.sh new file mode 100755 index 00000000..417bc44e --- /dev/null +++ b/docker/build-bundle.sh @@ -0,0 +1,342 @@ +#!/usr/bin/env bash +set -euo pipefail + +# ────────────────────────────────────────────────────────────────────────────── +# build-bundle.sh -- Build a pre-built AI bundle tar.gz inside the Docker image +# +# Usage: build-bundle.sh +# +# bundleId - One of the bundle IDs in feature-manifest.json +# arch - Architecture variant: amd64-gpu or arm64-cpu +# outputDir - Directory for the output .tar.gz and .sha256 files +# +# Runs as root inside the SnapOtter Docker container. The venv at /opt/venv +# must be activated. Produces: +# /-.tar.gz +# /-.tar.gz.sha256 +# ────────────────────────────────────────────────────────────────────────────── + +export BUNDLE_ID="${1:?Usage: build-bundle.sh }" +export ARCH="${2:?Usage: build-bundle.sh }" +OUTPUT_DIR="${3:?Usage: build-bundle.sh }" + +MANIFEST="/app/docker/feature-manifest.json" +VENV_PATH="${PYTHON_VENV_PATH:-/opt/venv}" +export MODELS_DIR="/tmp/bundle-models" +export BUILD_DIR="/tmp/bundle-build" + +# When running with --entrypoint bash (bypassing entrypoint.sh), the venv +# at /data/ai/venv won't exist yet. Use /opt/venv directly -- it's the base +# venv baked into the Docker image, and that's exactly what we want as the +# starting point for building bundle deltas. +if [[ ! -f "${VENV_PATH}/bin/activate" && -f "/opt/venv/bin/activate" ]]; then + VENV_PATH="/opt/venv" +fi + +# Activate the venv so pip/python3 use it (not system Python) +if [[ -f "${VENV_PATH}/bin/activate" ]]; then + # shellcheck disable=SC1091 + source "${VENV_PATH}/bin/activate" +fi + +SITE_PACKAGES="$("${VENV_PATH}/bin/python3" -c 'import site; print(site.getsitepackages()[0])')" + +# Parse platform from arch: amd64-gpu -> amd64, arm64-cpu -> arm64 +export PLATFORM="${ARCH%%-*}" + +echo "=== Building bundle: ${BUNDLE_ID} arch=${ARCH} platform=${PLATFORM} ===" +echo "Venv: ${VENV_PATH}" +echo "Site-packages: ${SITE_PACKAGES}" + +# Validate manifest exists +if [[ ! -f "${MANIFEST}" ]]; then + echo "ERROR: Manifest not found at ${MANIFEST}" >&2 + exit 1 +fi + +# Validate bundle exists in manifest +python3 -c " +import json, sys +with open('${MANIFEST}') as f: + m = json.load(f) +if '${BUNDLE_ID}' not in m['bundles']: + print(f'ERROR: Bundle \"${BUNDLE_ID}\" not found in manifest', file=sys.stderr) + print(f'Available: {list(m[\"bundles\"].keys())}', file=sys.stderr) + sys.exit(1) +" + +# Clean previous build artifacts +rm -rf "${MODELS_DIR}" "${BUILD_DIR}" +mkdir -p "${MODELS_DIR}" "${BUILD_DIR}/site-packages" "${BUILD_DIR}/models" "${OUTPUT_DIR}" + +# ── Step 1: Record base site-packages ──────────────────────────────────── +echo "=== Recording base site-packages ===" +find "${SITE_PACKAGES}" -maxdepth 1 -mindepth 1 | sort > /tmp/base-packages.txt +echo " Base entries: $(wc -l < /tmp/base-packages.txt)" + +# ── Step 2: Install packages ───────────────────────────────────────────── +echo "=== Installing packages ===" +python3 << 'PYINSTALL' +import json, subprocess, sys, os + +with open("/app/docker/feature-manifest.json") as f: + manifest = json.load(f) + +bundle_id = os.environ["BUNDLE_ID"] +platform = os.environ["PLATFORM"] +bundle = manifest["bundles"][bundle_id] +pip_flags = bundle.get("pipFlags", {}) + +# Collect all packages: common + arch-specific +packages = list(bundle["packages"].get("common", [])) +packages.extend(bundle["packages"].get(platform, [])) + +print(f" Installing {len(packages)} package(s) for {bundle_id} ({platform})") + +for pkg_string in packages: + # Check if any pipFlags key matches the start of this package string + extra_flags = "" + for flag_key, flag_val in pip_flags.items(): + if pkg_string.startswith(flag_key): + extra_flags = flag_val + break + + # pkg_string may contain embedded flags (e.g. --index-url), so pass as-is + cmd = f"{sys.executable} -m pip install --no-cache-dir {extra_flags} {pkg_string}".strip() + print(f" > {cmd}", flush=True) + result = subprocess.run(cmd, shell=True) + if result.returncode != 0: + print(f"ERROR: pip install failed for: {pkg_string}", file=sys.stderr) + sys.exit(1) + +print(" Package installation complete") +PYINSTALL + +# ── Step 3: Post-install fixups ─────────────────────────────────────────── +echo "=== Running post-install fixups ===" +python3 << 'PYPOST' +import json, subprocess, sys, os + +with open("/app/docker/feature-manifest.json") as f: + manifest = json.load(f) + +bundle = manifest["bundles"][os.environ["BUNDLE_ID"]] +post_install = bundle.get("postInstall", []) + +if not post_install: + print(" No post-install fixups") + sys.exit(0) + +for pkg in post_install: + cmd = f"{sys.executable} -m pip install --no-cache-dir --force-reinstall {pkg}" + print(f" > {cmd}", flush=True) + result = subprocess.run(cmd, shell=True) + if result.returncode != 0: + print(f"ERROR: post-install fixup failed for: {pkg}", file=sys.stderr) + sys.exit(1) + +print(" Post-install fixups complete") +PYPOST + +# ── Step 4: Re-pin base packages ───────────────────────────────────────── +echo "=== Re-pinning base packages ===" +python3 << 'PYREPIN' +import json, subprocess, sys + +with open("/app/docker/feature-manifest.json") as f: + manifest = json.load(f) + +base_packages = manifest.get("basePackages", []) +if not base_packages: + print(" No base packages to re-pin") + sys.exit(0) + +pkgs = " ".join(base_packages) +cmd = f"{sys.executable} -m pip install --no-cache-dir --force-reinstall {pkgs}" +print(f" > {cmd}", flush=True) +result = subprocess.run(cmd, shell=True) +if result.returncode != 0: + print("ERROR: base package re-pin failed", file=sys.stderr) + sys.exit(1) + +print(" Base packages re-pinned") +PYREPIN + +# ── Step 5: Download models ────────────────────────────────────────────── +echo "=== Downloading models ===" +python3 << 'PYMODELS' +import json, os, sys, urllib.request, pathlib + +with open("/app/docker/feature-manifest.json") as f: + manifest = json.load(f) + +bundle = manifest["bundles"][os.environ["BUNDLE_ID"]] +models = bundle.get("models", []) +models_dir = os.environ["MODELS_DIR"] + +if not models: + print(" No models to download") + sys.exit(0) + +print(f" Downloading {len(models)} model(s)") + +for model in models: + model_id = model["id"] + download_fn = model.get("downloadFn") + url = model.get("url") + + if url: + # Direct URL download via urllib + dest = os.path.join(models_dir, model["path"]) + os.makedirs(os.path.dirname(dest), exist_ok=True) + print(f" [{model_id}] URL -> {model['path']}", flush=True) + urllib.request.urlretrieve(url, dest) + size = os.path.getsize(dest) + min_size = model.get("minSize", 0) + if min_size and size < min_size: + print(f" WARNING: {model_id} is {size:,} bytes, expected >= {min_size:,}", file=sys.stderr) + print(f" [{model_id}] Done ({size:,} bytes)") + + elif download_fn == "hf_snapshot": + from huggingface_hub import snapshot_download + args = model["args"] + repo_id = args[0] + local_dir = os.path.join(models_dir, args[1]) + + kwargs = {"repo_id": repo_id, "local_dir": local_dir} + + # Only download specific file if specified + if "file" in model: + kwargs["allow_patterns"] = [model["file"]] + + # Handle non-default repo types (e.g. "space") + if "repoType" in model: + kwargs["repo_type"] = model["repoType"] + + print(f" [{model_id}] HF snapshot: {repo_id} -> {args[1]}", flush=True) + snapshot_download(**kwargs) + print(f" [{model_id}] Done") + + elif download_fn == "rembg_session": + args = model["args"] + session_name = args[0] + print(f" [{model_id}] rembg session: {session_name}", flush=True) + from rembg.sessions import new_session + new_session(session_name) + print(f" [{model_id}] Done") + + else: + print(f" WARNING: Unknown download method for {model_id}", file=sys.stderr) + +print(" Model downloads complete") +PYMODELS + +# ── Step 6: Diff site-packages ──────────────────────────────────────────── +echo "=== Diffing site-packages ===" +find "${SITE_PACKAGES}" -maxdepth 1 -mindepth 1 | sort > /tmp/after-packages.txt +DELTA_DIRS="$(comm -13 /tmp/base-packages.txt /tmp/after-packages.txt)" +DELTA_COUNT="$(echo "${DELTA_DIRS}" | grep -c . || true)" +echo " New top-level dirs: ${DELTA_COUNT}" + +if [[ "${DELTA_COUNT}" -eq 0 ]]; then + echo "WARNING: No new site-packages detected. Bundle may be empty." >&2 +fi + +# Copy delta dirs to build staging +while IFS= read -r dir; do + [[ -z "${dir}" ]] && continue + cp -a "${dir}" "${BUILD_DIR}/site-packages/" +done <<< "${DELTA_DIRS}" + +# Copy models to build staging +if [[ -d "${MODELS_DIR}" ]] && [[ -n "$(ls -A "${MODELS_DIR}" 2>/dev/null)" ]]; then + cp -a "${MODELS_DIR}"/* "${BUILD_DIR}/models/" +fi + +# ── Step 7: Torch NCCL fixup ───────────────────────────────────────────── +echo "=== Checking for torch NCCL fixup ===" +python3 << 'PYNCCL' +import importlib.metadata, subprocess, os, sys + +try: + reqs = importlib.metadata.requires("torch") +except importlib.metadata.PackageNotFoundError: + print(" torch not installed, skipping NCCL check") + sys.exit(0) + +if reqs is None: + print(" No torch requirements found") + sys.exit(0) + +nccl_pkgs = [r.split(";")[0].strip() for r in reqs if "nccl" in r.lower()] +if not nccl_pkgs: + print(" No NCCL requirements found") + sys.exit(0) + +fixups_dir = os.path.join(os.environ["BUILD_DIR"], "fixups") +os.makedirs(fixups_dir, exist_ok=True) + +for pkg in nccl_pkgs: + print(f" Downloading NCCL wheel: {pkg}", flush=True) + result = subprocess.run( + [sys.executable, "-m", "pip", "download", "--no-cache-dir", "-d", fixups_dir, pkg] + ) + if result.returncode != 0: + print(f" WARNING: Failed to download NCCL wheel: {pkg}", file=sys.stderr) + +print(" NCCL fixup complete") +PYNCCL + +# ── Step 8: Write bundle.json ───────────────────────────────────────────── +echo "=== Writing bundle.json ===" +python3 << 'PYBUNDLE' +import json, os, sys + +with open("/app/docker/feature-manifest.json") as f: + manifest = json.load(f) + +bundle_id = os.environ["BUNDLE_ID"] +arch = os.environ["ARCH"] +bundle = manifest["bundles"][bundle_id] +model_ids = [m["id"] for m in bundle.get("models", [])] + +py_ver = f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}" + +bundle_meta = { + "bundleId": bundle_id, + "version": manifest["imageVersion"], + "arch": arch, + "imageVersion": manifest["imageVersion"], + "pythonVersion": py_ver, + "models": model_ids, +} + +out_path = os.path.join(os.environ["BUILD_DIR"], "bundle.json") +with open(out_path, "w") as f: + json.dump(bundle_meta, f, indent=2) + f.write("\n") + +print(f" {json.dumps(bundle_meta, indent=2)}") +PYBUNDLE + +# ── Step 9: Create archive ──────────────────────────────────────────────── +echo "=== Creating archive ===" +ARCHIVE_NAME="${BUNDLE_ID}-${ARCH}.tar.gz" +ARCHIVE_PATH="${OUTPUT_DIR}/${ARCHIVE_NAME}" + +tar -czf "${ARCHIVE_PATH}" -C "${BUILD_DIR}" . + +sha256sum "${ARCHIVE_PATH}" | awk '{print $1}' > "${ARCHIVE_PATH}.sha256" + +ARCHIVE_SIZE="$(stat -c%s "${ARCHIVE_PATH}" 2>/dev/null || stat -f%z "${ARCHIVE_PATH}")" +SHA256="$(cat "${ARCHIVE_PATH}.sha256")" + +echo " Archive: ${ARCHIVE_PATH}" +echo " Size: ${ARCHIVE_SIZE} bytes" +echo " SHA256: ${SHA256}" + +# ── Cleanup ─────────────────────────────────────────────────────────────── +echo "=== Cleaning up ===" +rm -rf "${MODELS_DIR}" "${BUILD_DIR}" /tmp/base-packages.txt /tmp/after-packages.txt + +echo "=== Done: ${BUNDLE_ID} ${ARCH} ===" diff --git a/docker/feature-manifest.json b/docker/feature-manifest.json index 57a10847..69c999f4 100644 --- a/docker/feature-manifest.json +++ b/docker/feature-manifest.json @@ -1,13 +1,28 @@ { - "manifestVersion": 1, + "manifestVersion": 2, "imageVersion": "2.0.0", "pythonVersion": "3.12", "basePackages": ["numpy==1.26.4", "Pillow==12.2.0", "opencv-python-headless==4.10.0.84"], + "bundleRepo": "snapotter/feature-bundles", "bundles": { "background-removal": { "name": "Background Removal", "description": "Remove image backgrounds with AI", "estimatedSize": "4-5 GB", + "archives": { + "amd64-gpu": { + "file": "v2.0.0/background-removal-amd64-gpu.tar.gz", + "sha256": "0000000000000000000000000000000000000000000000000000000000000000", + "compressedSize": 1, + "extractedSize": 1 + }, + "arm64-cpu": { + "file": "v2.0.0/background-removal-arm64-cpu.tar.gz", + "sha256": "0000000000000000000000000000000000000000000000000000000000000000", + "compressedSize": 1, + "extractedSize": 1 + } + }, "packages": { "common": ["rembg==2.0.62"], "amd64": ["onnxruntime-gpu==1.20.1", "mediapipe>=0.10.21"], @@ -61,6 +76,20 @@ "name": "Face Detection", "description": "Detect and blur faces, fix red-eye, smart crop", "estimatedSize": "200-300 MB", + "archives": { + "amd64-gpu": { + "file": "v2.0.0/face-detection-amd64-gpu.tar.gz", + "sha256": "0000000000000000000000000000000000000000000000000000000000000000", + "compressedSize": 1, + "extractedSize": 1 + }, + "arm64-cpu": { + "file": "v2.0.0/face-detection-arm64-cpu.tar.gz", + "sha256": "0000000000000000000000000000000000000000000000000000000000000000", + "compressedSize": 1, + "extractedSize": 1 + } + }, "packages": { "common": [], "amd64": ["mediapipe>=0.10.21"], @@ -88,6 +117,20 @@ "name": "Object Eraser & Colorize", "description": "Erase objects from photos and colorize B&W images", "estimatedSize": "1-2 GB", + "archives": { + "amd64-gpu": { + "file": "v2.0.0/object-eraser-colorize-amd64-gpu.tar.gz", + "sha256": "0000000000000000000000000000000000000000000000000000000000000000", + "compressedSize": 1, + "extractedSize": 1 + }, + "arm64-cpu": { + "file": "v2.0.0/object-eraser-colorize-arm64-cpu.tar.gz", + "sha256": "0000000000000000000000000000000000000000000000000000000000000000", + "compressedSize": 1, + "extractedSize": 1 + } + }, "packages": { "common": ["huggingface-hub"], "amd64": ["onnxruntime-gpu==1.20.1"], @@ -140,6 +183,20 @@ "name": "Upscale & Enhance", "description": "AI upscaling, face enhancement, and noise removal", "estimatedSize": "4-5 GB", + "archives": { + "amd64-gpu": { + "file": "v2.0.0/upscale-enhance-amd64-gpu.tar.gz", + "sha256": "0000000000000000000000000000000000000000000000000000000000000000", + "compressedSize": 1, + "extractedSize": 1 + }, + "arm64-cpu": { + "file": "v2.0.0/upscale-enhance-arm64-cpu.tar.gz", + "sha256": "0000000000000000000000000000000000000000000000000000000000000000", + "compressedSize": 1, + "extractedSize": 1 + } + }, "packages": { "common": ["codeformer-pip==0.0.4", "huggingface-hub", "einops", "setuptools<75"], "amd64": [ @@ -216,6 +273,20 @@ "name": "Photo Restoration", "description": "Restore old or damaged photos", "estimatedSize": "800 MB - 1 GB", + "archives": { + "amd64-gpu": { + "file": "v2.0.0/photo-restoration-amd64-gpu.tar.gz", + "sha256": "0000000000000000000000000000000000000000000000000000000000000000", + "compressedSize": 1, + "extractedSize": 1 + }, + "arm64-cpu": { + "file": "v2.0.0/photo-restoration-arm64-cpu.tar.gz", + "sha256": "0000000000000000000000000000000000000000000000000000000000000000", + "compressedSize": 1, + "extractedSize": 1 + } + }, "packages": { "common": ["codeformer-pip==0.0.4", "huggingface-hub", "setuptools<75"], "amd64": [ @@ -299,6 +370,20 @@ "name": "OCR", "description": "Extract text from images", "estimatedSize": "3-4 GB", + "archives": { + "amd64-gpu": { + "file": "v2.0.0/ocr-amd64-gpu.tar.gz", + "sha256": "0000000000000000000000000000000000000000000000000000000000000000", + "compressedSize": 1, + "extractedSize": 1 + }, + "arm64-cpu": { + "file": "v2.0.0/ocr-arm64-cpu.tar.gz", + "sha256": "0000000000000000000000000000000000000000000000000000000000000000", + "compressedSize": 1, + "extractedSize": 1 + } + }, "packages": { "common": ["huggingface-hub"], "amd64": [ @@ -357,6 +442,20 @@ "name": "Transcription", "description": "Speech to text for audio and video (subtitles)", "estimatedSize": "~600 MB", + "archives": { + "amd64-gpu": { + "file": "v2.0.0/transcription-amd64-gpu.tar.gz", + "sha256": "0000000000000000000000000000000000000000000000000000000000000000", + "compressedSize": 1, + "extractedSize": 1 + }, + "arm64-cpu": { + "file": "v2.0.0/transcription-arm64-cpu.tar.gz", + "sha256": "0000000000000000000000000000000000000000000000000000000000000000", + "compressedSize": 1, + "extractedSize": 1 + } + }, "packages": { "common": ["faster-whisper>=1.0.0"], "amd64": [], "arm64": [] }, "pipFlags": {}, "postInstall": [], diff --git a/packages/ai/python/install_feature.py b/packages/ai/python/install_feature.py index 633320f0..546791ad 100644 --- a/packages/ai/python/install_feature.py +++ b/packages/ai/python/install_feature.py @@ -1,4 +1,7 @@ -"""Install a feature bundle: pip packages + model downloads. +"""Pre-built AI bundle installer for SnapOtter. + +Downloads a pre-built tar.gz archive (or uses a local file), verifies its +SHA256 checksum, extracts site-packages and models, and writes installed.json. Invoked by the Node.js backend as a subprocess. @@ -9,21 +12,22 @@ Progress is reported via JSON lines on stderr (parsed by the Node bridge). Final result is a JSON object on stdout. """ -import concurrent.futures +import glob +import hashlib import json import os import platform import shutil import subprocess import sys +import tarfile import time import urllib.error import urllib.request from datetime import datetime, timezone -# ── Helpers ────────────────────────────────────────────────────────────── - +# -- Helpers -- def emit_progress(percent: int, stage: str) -> None: """Emit a progress update via stderr JSON line.""" @@ -38,521 +42,205 @@ def fail(message: str) -> None: sys.exit(1) +# -- Architecture detection -- + def detect_arch() -> str: - """Return 'arm64' or 'amd64' based on the host machine.""" + """Return 'amd64-gpu' or 'arm64-cpu' based on host + GPU.""" machine = platform.machine().lower() if machine in ("aarch64", "arm64"): - return "arm64" - return "amd64" + return "arm64-cpu" + return "amd64-gpu" -def has_nvidia_gpu() -> bool: - """Check whether an NVIDIA GPU is accessible at runtime.""" - try: - result = subprocess.run( - ["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"], - capture_output=True, text=True, timeout=5, +# -- Disk space -- + +def check_disk_space(path: str, needed_bytes: int) -> None: + """Fail if insufficient disk space.""" + usage = shutil.disk_usage(path) + if usage.free < needed_bytes: + free_gb = usage.free / (1024 ** 3) + need_gb = needed_bytes / (1024 ** 3) + fail( + f"Insufficient disk space: need {need_gb:.1f} GB, " + f"have {free_gb:.1f} GB free. " + f"Free up space and retry." ) - return result.returncode == 0 and len(result.stdout.strip()) > 0 - except (FileNotFoundError, subprocess.TimeoutExpired): - return False -def cpu_fallback_packages(packages: list[str]) -> list[str]: - """Replace GPU-only packages with their CPU equivalents. +# -- Venv site-packages discovery -- - Called on amd64 when no NVIDIA GPU is detected so that onnxruntime / - paddlepaddle don't crash with a CUDA segfault. - Also replaces CUDA-pinned torch/torchvision with CPU-only versions. - """ - replacements = { - "onnxruntime-gpu": "onnxruntime", - "paddlepaddle-gpu": "paddlepaddle", - } - result = [] - for pkg in packages: - # Handle multi-package CUDA torch entries like: - # "torch==2.7.0+cu126 torchvision==0.22.0+cu126 --index-url ..." - first_token = pkg.split()[0] if pkg.strip() else "" - if first_token.startswith("torch==") and "+cu" in first_token: - # Extract torch and torchvision versions, use CPU-only index - cpu_pkgs = [] - for token in pkg.split(): - if token.startswith("torch==") and "+cu" in token: - base_ver = token.split("+")[0] # "torch==2.6.0" - cpu_pkgs.append(base_ver) - elif token.startswith("torchvision==") and "+cu" in token: - base_ver = token.split("+")[0] # "torchvision==0.21.0" - cpu_pkgs.append(base_ver) - # Use CPU-only wheels (~200MB vs ~2.6GB with CUDA) - cpu_pkgs.append("--index-url") - cpu_pkgs.append("https://download.pytorch.org/whl/cpu") - # Join into a single string so pip_install processes them as one command - result.append(" ".join(cpu_pkgs)) - continue - - name = pkg.split("==")[0].split(">=")[0].split("[")[0].strip() - if name in replacements: - # Extract only the version spec, drop any inline flags - # (e.g. "--extra-index-url https://...cu126/" is GPU-specific) - tokens = pkg.split() - version_token = tokens[0][len(name):] # e.g. ">=3.2.1" - result.append(replacements[name] + version_token) - else: - result.append(pkg) - return result - - -def check_disk_space(path: str, min_bytes: int = 100 * 1024 * 1024) -> None: - """Exit with a clear error if free disk space is below min_bytes.""" - try: - usage = shutil.disk_usage(path) - if usage.free < min_bytes: - free_mb = usage.free / (1024 * 1024) - min_mb = min_bytes / (1024 * 1024) - fail( - f"Insufficient disk space: {free_mb:.0f} MB free, " - f"need at least {min_mb:.0f} MB" - ) - except OSError as e: - # If we can't check, warn but continue - sys.stderr.write(f"Warning: could not check disk space: {e}\n") - sys.stderr.flush() - - -# ── pip install ────────────────────────────────────────────────────────── - - -def _pip_error_hint(package: str, stderr: str) -> str: - """Return a user-friendly hint for known pip install failure patterns.""" - if "KeyError" in stderr and "__version__" in stderr: - return ( - "The 'basicsr' dependency failed to build due to a known " - "compatibility issue with newer setuptools versions. " - "Try running: pip install basicsr==1.4.2 --no-build-isolation " - "inside the container, then retry this installation." - ) - if "MemoryError" in stderr or "Cannot allocate memory" in stderr: - return ( - "Installation ran out of memory. " - "Increase the container's memory limit to at least 6 GB and retry." - ) - if "No space left on device" in stderr: - return ( - "Disk space exhausted during package installation. " - "Free up disk space or increase the container's disk size and retry." - ) +def get_site_packages_dir(venv_path: str) -> str: + """Find the site-packages directory inside a Python venv.""" + matches = glob.glob(os.path.join(venv_path, "lib", "python*", "site-packages")) + if matches: + return matches[0] return "" -def pip_install(package: str, extra_flags: list[str] | None = None) -> None: - """Run pip install for a single package spec. Raises on failure.""" - cmd = [sys.executable, "-m", "pip", "install", "--no-cache-dir"] - if extra_flags: - cmd.extend(extra_flags) +# -- SHA256 verification -- - # Package spec may include inline flags like - # "realesrgan==0.3.0 --extra-index-url https://..." - parts = package.split() - cmd.extend(parts) - - result = subprocess.run( - cmd, - capture_output=True, - text=True, - ) - if result.returncode != 0: - stderr = result.stderr.strip() - hint = _pip_error_hint(package, stderr) - if hint: - raise RuntimeError(f"pip install failed for '{package}': {hint}") - tail = stderr[-500:] if len(stderr) > 500 else stderr - raise RuntimeError(f"pip install failed for '{package}': {tail}") - - -def install_packages(bundle: dict, arch: str) -> None: - """Install pip packages for the bundle (common + arch-specific + post-install).""" - packages_section = bundle.get("packages", {}) - common_pkgs = packages_section.get("common", []) - arch_pkgs = packages_section.get(arch, []) - all_pkgs = common_pkgs + arch_pkgs - - # On amd64 without GPU, swap GPU packages for CPU equivalents to avoid - # segfaults from onnxruntime-gpu / paddlepaddle-gpu trying to init CUDA. - if arch == "amd64" and not has_nvidia_gpu(): - all_pkgs = cpu_fallback_packages(all_pkgs) - sys.stderr.write("No NVIDIA GPU detected — using CPU package variants\n") - sys.stderr.flush() - pip_flags = bundle.get("pipFlags", {}) - post_install = bundle.get("postInstall", []) - - total_pkgs = len(all_pkgs) + len(post_install) - if total_pkgs == 0: - return - - for i, pkg in enumerate(all_pkgs): - progress = int((i / total_pkgs) * 50) - # Extract display name(s) from package spec (may contain multiple - # packages and flags like "torch==2.6.0+cu126 torchvision==... --index-url ...") - tokens = [t for t in pkg.split() if not t.startswith("-") and "://" not in t] - pkg_name = ", ".join(t.split("==")[0].split(">=")[0].split("[")[0] for t in tokens) if tokens else pkg - emit_progress(progress, f"Installing {pkg_name}...") - - # Check for package-specific pip flags - extra = None - for flag_key, flag_val in pip_flags.items(): - if flag_key in pkg: - extra = flag_val.split() if isinstance(flag_val, str) else flag_val +def verify_sha256(filepath: str, expected: str) -> bool: + """Stream-hash a file and compare to expected hex digest.""" + h = hashlib.sha256() + with open(filepath, "rb") as f: + while True: + chunk = f.read(8192) + if not chunk: break - pip_install(pkg, extra) - - # Post-install fixups (e.g., re-pin numpy after codeformer drags in a newer one) - for j, pkg in enumerate(post_install): - progress = int(((len(all_pkgs) + j) / total_pkgs) * 50) - pkg_name = pkg.split("==")[0].split(">=")[0].strip() - emit_progress(progress, f"Post-install: {pkg_name}...") - pip_install(pkg) + h.update(chunk) + return h.hexdigest() == expected -def handle_nccl_conflict() -> None: - """Re-install torch's NCCL dependency if both torch and paddlepaddle-gpu coexist. +# -- Download with resume -- - PaddlePaddle ships its own NCCL, which can conflict with the version - that torch expects. Force-reinstalling torch's pinned nccl resolves this. +def download_with_resume( + url: str, + dest: str, + expected_size: int, + progress_start: int, + progress_end: int, +) -> None: + """Download a file with resume support via Range headers. + + Uses .partial and .meta sidecar files for crash recovery. """ - try: - from importlib.metadata import PackageNotFoundError, requires + partial_path = dest + ".partial" + meta_path = dest + ".meta" - # Only needed if both torch AND paddlepaddle-gpu are installed + # Check for existing partial download + bytes_downloaded = 0 + if os.path.exists(partial_path) and os.path.exists(meta_path): try: - requires("torch") - except PackageNotFoundError: - return - try: - requires("paddlepaddle-gpu") - except PackageNotFoundError: - return + with open(meta_path, "r") as f: + meta = json.load(f) + bytes_downloaded = meta.get("bytesDownloaded", 0) + if bytes_downloaded > 0: + actual_size = os.path.getsize(partial_path) + if actual_size != bytes_downloaded: + bytes_downloaded = 0 # Mismatch, restart + except (json.JSONDecodeError, OSError): + bytes_downloaded = 0 - # Find torch's NCCL requirement - reqs = requires("torch") or [] - nccl_reqs = [r.split(";")[0].strip() for r in reqs if "nccl" in r.lower()] - if nccl_reqs: - emit_progress(48, "Fixing NCCL conflict...") - subprocess.run( - [sys.executable, "-m", "pip", "install", nccl_reqs[0]], - capture_output=True, - text=True, - ) - except Exception: - # Non-fatal — if we can't fix it, the user may not even hit the conflict - pass - - -# ── Model downloads ────────────────────────────────────────────────────── - - -def urlretrieve_with_retry(url: str, dest: str, max_retries: int = 3) -> None: - """Download a URL to a local file with retry + exponential backoff.""" - for attempt in range(max_retries): - try: - req = urllib.request.Request( - url, headers={"User-Agent": "snapotter-installer/1.0"} - ) - with urllib.request.urlopen(req, timeout=300) as resp, open(dest, "wb") as f: - shutil.copyfileobj(resp, f) - return - except Exception as e: - if attempt < max_retries - 1: - time.sleep(10 * (2 ** attempt)) - else: - raise RuntimeError(f"Failed to download {url}: {e}") - - -def download_url_model(model: dict, models_dir: str) -> None: - """Download a model via direct URL with atomic rename.""" - rel_path = model["path"] - url = model["url"] - min_size = model.get("minSize", 0) - final_path = os.path.join(models_dir, rel_path) - tmp_path = final_path + ".downloading" - - # Idempotent: skip if already present and big enough - if os.path.exists(final_path): - if min_size <= 0 or os.path.getsize(final_path) >= min_size: - return - - os.makedirs(os.path.dirname(final_path), exist_ok=True) - - # Clean up orphaned partial download - if os.path.exists(tmp_path): - os.remove(tmp_path) - - urlretrieve_with_retry(url, tmp_path) - - # Verify size - actual_size = os.path.getsize(tmp_path) - if min_size > 0 and actual_size < min_size: - os.remove(tmp_path) - raise RuntimeError( - f"Model {model['id']} too small: {actual_size} bytes " - f"(expected >= {min_size})" - ) - - # Atomic rename - os.rename(tmp_path, final_path) - - -_matting_registered = False - - -def _register_birefnet_matting() -> None: - """Register the custom BiRefNet-matting ONNX session. - - This model is not built into rembg — it must be registered before - calling new_session("birefnet-matting"). The same registration is - done in remove_bg.py (runtime) and download_models.py (build-time). - """ - global _matting_registered - if _matting_registered: - return - _matting_registered = True - - import pooch - from rembg.sessions import sessions_class - from rembg.sessions.birefnet_general import BiRefNetSessionGeneral - - class BiRefNetMattingSession(BiRefNetSessionGeneral): - @classmethod - def download_models(cls, *args, **kwargs): - fname = f"{cls.name(*args, **kwargs)}.onnx" - pooch.retrieve( - "https://github.com/ZhengPeng7/BiRefNet/releases/download/v1/BiRefNet-matting-epoch_100.onnx", - None, - fname=fname, - path=cls.u2net_home(*args, **kwargs), - progressbar=True, - ) - return os.path.join(cls.u2net_home(*args, **kwargs), fname) - - @classmethod - def name(cls, *args, **kwargs): - return "birefnet-matting" - - sessions_class.append(BiRefNetMattingSession) - - -_hr_matting_registered = False - - -def _register_birefnet_hr_matting() -> None: - """Register the custom BiRefNet HR-matting ONNX session for 2048x2048 high-res matting. - - Like _register_birefnet_matting(), this model is not built into rembg and - must be registered before calling new_session("birefnet-hr-matting"). - """ - global _hr_matting_registered - if _hr_matting_registered: - return - _hr_matting_registered = True - - import numpy as np - import pooch - from PIL import Image - from rembg.sessions import sessions_class - from rembg.sessions.birefnet_general import BiRefNetSessionGeneral - - class BiRefNetHRMattingSession(BiRefNetSessionGeneral): - @classmethod - def download_models(cls, *args, **kwargs): - fname = f"{cls.name(*args, **kwargs)}.onnx" - pooch.retrieve( - "https://github.com/ZhengPeng7/BiRefNet/releases/download/v1/BiRefNet_HR-matting-epoch_135.onnx", - None, - fname=fname, - path=cls.u2net_home(*args, **kwargs), - progressbar=True, - ) - return os.path.join(cls.u2net_home(*args, **kwargs), fname) - - @classmethod - def name(cls, *args, **kwargs): - return "birefnet-hr-matting" - - def predict(self, img, *args, **kwargs): - ort_outs = self.inner_session.run( - None, - self.normalize( - img, (0.485, 0.456, 0.406), (0.229, 0.224, 0.225), (2048, 2048) - ), - ) - pred = ort_outs[0][:, 0, :, :] - ma = np.max(pred) - mi = np.min(pred) - denom = ma - mi - pred = (pred - mi) / denom if denom > 0 else pred * 0 - pred = np.squeeze(pred) - mask = Image.fromarray((pred * 255).astype("uint8"), mode="L") - mask = mask.resize(img.size, Image.LANCZOS) - return [mask] - - sessions_class.append(BiRefNetHRMattingSession) - - -def download_rembg_session(model: dict, models_dir: str) -> None: - """Download a rembg model by initializing a session.""" - args = model.get("args", []) - if not args: - raise RuntimeError(f"rembg_session model {model['id']} has no args") - - model_name = args[0] - - # Set U2NET_HOME so rembg stores models in our models_dir - u2net_dir = os.path.join(models_dir, "rembg") - os.makedirs(u2net_dir, exist_ok=True) - os.environ["U2NET_HOME"] = u2net_dir - - try: - from rembg import new_session - except ImportError: - raise RuntimeError( - f"rembg package not available for model '{model_name}' " - f"-- pip install may have failed in an earlier step" - ) - _register_birefnet_matting() - _register_birefnet_hr_matting() - try: - new_session(model_name) - except Exception as e: - raise RuntimeError( - f"Failed to download rembg model '{model_name}': {e}. " - f"This is usually caused by network issues (timeouts or rate limiting). " - f"Check your internet connection and retry." - ) - - -def download_hf_snapshot(model: dict, models_dir: str) -> None: - """Download a model via huggingface_hub.snapshot_download.""" - args = model.get("args", []) - if len(args) < 2: - raise RuntimeError( - f"hf_snapshot model {model['id']} needs [repo_id, local_subdir]" - ) - - repo_id = args[0] - local_subdir = args[1] - local_dir = os.path.join(models_dir, local_subdir) - repo_type = model.get("repoType", "model") - min_size = model.get("minSize", 0) - target_file = model.get("file") - - os.makedirs(local_dir, exist_ok=True) - - # Idempotent: if target file exists and meets minSize, skip - if target_file: - final_file = os.path.join(local_dir, target_file) - if os.path.exists(final_file): - if min_size <= 0 or os.path.getsize(final_file) >= min_size: - return - - from huggingface_hub import snapshot_download - - kwargs: dict = {"repo_id": repo_id, "local_dir": local_dir, "repo_type": repo_type} - if target_file: - kwargs["allow_patterns"] = [target_file] + if bytes_downloaded == 0 and os.path.exists(partial_path): + os.unlink(partial_path) max_retries = 3 for attempt in range(max_retries): try: - snapshot_download(**kwargs) - break + headers = {"User-Agent": "snapotter-installer/2.0"} + if bytes_downloaded > 0: + headers["Range"] = f"bytes={bytes_downloaded}-" + emit_progress( + progress_start, + f"Resuming download from {bytes_downloaded / (1024**3):.1f} GB...", + ) + + req = urllib.request.Request(url, headers=headers) + with urllib.request.urlopen(req, timeout=300) as resp: + mode = "ab" if bytes_downloaded > 0 else "wb" + with open(partial_path, mode) as f: + while True: + chunk = resp.read(65536) + if not chunk: + break + f.write(chunk) + bytes_downloaded += len(chunk) + + # Update progress + if expected_size > 0: + pct = bytes_downloaded / expected_size + progress = int( + progress_start + pct * (progress_end - progress_start) + ) + progress = min(progress, progress_end) + stage = f"Downloading... {bytes_downloaded / (1024**3):.1f} GB" + emit_progress(progress, stage) + + # Write meta periodically (every 10 MB) + if bytes_downloaded % (10 * 1024 * 1024) < 65536: + with open(meta_path, "w") as mf: + json.dump({"bytesDownloaded": bytes_downloaded}, mf) + + # Download complete + os.rename(partial_path, dest) + if os.path.exists(meta_path): + os.unlink(meta_path) + return + except Exception as e: + # Write meta for resume on next attempt + with open(meta_path, "w") as mf: + json.dump({"bytesDownloaded": bytes_downloaded}, mf) + if attempt < max_retries - 1: delay = 10 * (2 ** attempt) - sys.stderr.write( - f"HuggingFace download failed for {model.get('id', repo_id)} " - f"(attempt {attempt + 1}/{max_retries}), retrying in {delay}s: {e}\n" + emit_progress( + progress_start, + f"Download failed (attempt {attempt + 1}/{max_retries}), " + f"retrying in {delay}s: {e}", ) - sys.stderr.flush() time.sleep(delay) else: + # Clean up on final failure + for p in (partial_path, meta_path): + if os.path.exists(p): + os.unlink(p) raise RuntimeError( - f"Failed to download {model.get('id', repo_id)} from " - f"HuggingFace repo {repo_id} after {max_retries} attempts: {e}" + f"Failed to download after {max_retries} attempts: {e}" ) - # Verify file size if applicable - if target_file and min_size > 0: - final_file = os.path.join(local_dir, target_file) - if os.path.exists(final_file): - actual = os.path.getsize(final_file) - if actual < min_size: - raise RuntimeError( - f"Model {model['id']} file {target_file} too small: " - f"{actual} bytes (expected >= {min_size})" - ) - else: - raise RuntimeError( - f"Model {model['id']} file {target_file} not found after download" - ) + +# -- Safe tar extraction -- + +def safe_extract(tar_path: str, staging_dir: str) -> None: + """Extract a tar.gz with security guards.""" + os.makedirs(staging_dir, exist_ok=True) + with tarfile.open(tar_path, "r:gz") as tf: + for member in tf.getmembers(): + # Block symlinks, hardlinks, devices + if not member.isfile() and not member.isdir(): + raise RuntimeError(f"Blocked unsafe tar entry type: {member.name}") + # Block absolute paths and traversal + if member.name.startswith("/") or ".." in member.name.split("/"): + raise RuntimeError(f"Blocked unsafe tar path: {member.name}") + tf.extractall(staging_dir, filter="data") -def download_single_model(model: dict, models_dir: str) -> None: - """Dispatch to the correct download function for a single model entry.""" - download_fn = model.get("downloadFn") - if download_fn == "rembg_session": - download_rembg_session(model, models_dir) - elif download_fn == "hf_snapshot": - download_hf_snapshot(model, models_dir) - elif "url" in model and "path" in model: - download_url_model(model, models_dir) - else: - raise RuntimeError( - f"Model {model['id']} has no recognized download method" - ) +# -- File move -- + +def move_tree(src: str, dst: str) -> None: + """Recursively merge src into dst, overwriting existing files.""" + if os.path.isdir(src): + shutil.copytree(src, dst, dirs_exist_ok=True) + shutil.rmtree(src) -def download_models(models: list[dict], models_dir: str) -> list[str]: - """Download all models in parallel. Returns list of failed model IDs.""" - if not models: - return [] +# -- Fixups (NCCL wheel) -- - failed: list[str] = [] - total = len(models) - - def _download(idx: int, model: dict) -> tuple[str, Exception | None]: - model_id = model.get("id", f"model-{idx}") +def apply_fixups(staging_dir: str, venv_path: str) -> None: + """Install any wheels from fixups/ directory (local only, no network).""" + fixups_dir = os.path.join(staging_dir, "fixups") + if not os.path.isdir(fixups_dir): + return + wheels = [f for f in os.listdir(fixups_dir) if f.endswith(".whl")] + if not wheels: + return + python_path = os.path.join(venv_path, "bin", "python3") + if not os.path.exists(python_path): + return + for wheel in wheels: + pkg_name = wheel.split("-")[0] try: - download_single_model(model, models_dir) - return (model_id, None) - except Exception as e: - return (model_id, e) - - with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool: - futures = { - pool.submit(_download, i, m): i - for i, m in enumerate(models) - } - - completed = 0 - for future in concurrent.futures.as_completed(futures): - completed += 1 - progress = 50 + int((completed / total) * 50) - - model_id, error = future.result() - if error: - failed.append(model_id) - sys.stderr.write( - f"Error downloading {model_id}: {error}\n" - ) - sys.stderr.flush() - else: - emit_progress(progress, f"Downloaded {model_id}") - - return failed + subprocess.run( + [python_path, "-m", "pip", "install", "--no-index", + f"--find-links={fixups_dir}", pkg_name], + capture_output=True, text=True, timeout=60, + ) + except Exception: + pass # Non-fatal -# ── installed.json management ──────────────────────────────────────────── - +# -- installed.json management -- def read_installed(ai_dir: str) -> dict: """Read the current installed.json, returning empty structure if missing.""" @@ -576,15 +264,9 @@ def write_installed_atomic(ai_dir: str, data: dict) -> None: os.rename(tmp_path, path) -# ── Main ───────────────────────────────────────────────────────────────── - +# -- Main -- def main() -> None: - if sys.version_info >= (3, 14): - print(f"[WARN] Python {sys.version_info.major}.{sys.version_info.minor} detected. " - f"Some packages may not have pre-built wheels. Build from source may be attempted.", - file=sys.stderr, flush=True) - if len(sys.argv) < 4: fail( f"Usage: {sys.argv[0]} \n" @@ -594,71 +276,153 @@ def main() -> None: bundle_id = sys.argv[1] manifest_path = sys.argv[2] models_dir = sys.argv[3] - - # Derive AI dir (parent of models dir) ai_dir = os.path.dirname(models_dir) + staging_base = os.path.join(ai_dir, "staging") + venv_path = os.environ.get("PYTHON_VENV_PATH", os.path.join(ai_dir, "venv")) - # ── Load manifest ──────────────────────────────────────────────────── - + # -- Load manifest -- emit_progress(0, "Reading manifest...") try: with open(manifest_path, "r") as f: manifest = json.load(f) - except (OSError, json.JSONDecodeError) as e: - fail(f"Cannot read manifest at {manifest_path}: {e}") + except Exception as e: + fail(f"Failed to read manifest: {e}") bundles = manifest.get("bundles", {}) if bundle_id not in bundles: - fail(f"Bundle '{bundle_id}' not found in manifest") + fail(f"Unknown bundle: {bundle_id}") bundle = bundles[bundle_id] - version = manifest.get("imageVersion", "0.0.0") - - # ── Detect architecture ────────────────────────────────────────────── + archives = bundle.get("archives") + if not archives: + fail(f"Bundle '{bundle_id}' has no archives in manifest (v2 required)") + # -- Detect architecture -- arch = detect_arch() - emit_progress(1, f"Architecture: {arch}") + archive = archives.get(arch) + if not archive: + fail(f"No archive for architecture '{arch}' in bundle '{bundle_id}'") - # ── Disk space pre-check ───────────────────────────────────────────── + archive_file = archive["file"] + expected_sha256 = archive["sha256"] + compressed_size = archive.get("compressedSize", 0) + extracted_size = archive.get("extractedSize", 0) - check_disk_space(models_dir) + # -- Check for local file override (testing / offline) -- + local_path = os.environ.get("SNAPOTTER_BUNDLE_LOCAL_PATH") - # ── Install pip packages ───────────────────────────────────────────── + if local_path: + # Local mode: use the file directly, verify checksum + emit_progress(5, "Using local bundle archive...") + tar_path = local_path + + if not os.path.exists(tar_path): + fail(f"Local bundle file not found: {tar_path}") + + # Verify checksum + emit_progress(10, "Verifying checksum...") + if not verify_sha256(tar_path, expected_sha256): + fail( + f"SHA256 checksum mismatch for local file.\n" + f"Expected: {expected_sha256}\n" + f"This usually means the manifest and archive are out of sync." + ) + else: + # Remote mode: download from HuggingFace + bundle_repo = manifest.get("bundleRepo", "snapotter/feature-bundles") + url = f"https://huggingface.co/{bundle_repo}/resolve/main/{archive_file}" + + # Disk space check + needed = compressed_size + extracted_size + 500 * 1024 * 1024 # 500 MB buffer + if needed > 0: + check_disk_space(ai_dir, needed) + + # Download + os.makedirs(staging_base, exist_ok=True) + tar_path = os.path.join(staging_base, f"{bundle_id}-{arch}.tar.gz") + + emit_progress(2, f"Downloading {bundle.get('name', bundle_id)} bundle...") + + try: + download_with_resume(url, tar_path, compressed_size, 2, 85) + except RuntimeError as e: + fail( + f"{e}\n\n" + f"You can download the bundle manually from:\n" + f" {url}\n" + f"Then upload it via Settings > AI Features > Offline Import." + ) + + # Verify checksum + emit_progress(86, "Verifying integrity...") + if not verify_sha256(tar_path, expected_sha256): + # Delete and retry once from scratch + os.unlink(tar_path) + emit_progress(86, "Checksum mismatch, retrying download...") + try: + download_with_resume(url, tar_path, compressed_size, 2, 85) + except RuntimeError as e: + fail(str(e)) + + if not verify_sha256(tar_path, expected_sha256): + os.unlink(tar_path) + fail( + f"SHA256 checksum mismatch after re-download.\n" + f"Expected: {expected_sha256}\n" + f"The archive may be corrupted. Try again later." + ) + + # -- Extract to staging -- + staging_dir = os.path.join(ai_dir, f"staging-{bundle_id}") + emit_progress(88, "Extracting packages and models...") - emit_progress(2, "Installing packages...") try: - install_packages(bundle, arch) - except RuntimeError as e: - fail(str(e)) + if os.path.exists(staging_dir): + shutil.rmtree(staging_dir) + safe_extract(tar_path, staging_dir) + except Exception as e: + if os.path.exists(staging_dir): + shutil.rmtree(staging_dir, ignore_errors=True) + fail(f"Failed to extract archive: {e}") - emit_progress(50, "Packages installed") + # -- Read bundle.json from tar -- + bundle_json_path = os.path.join(staging_dir, "bundle.json") + if not os.path.exists(bundle_json_path): + shutil.rmtree(staging_dir, ignore_errors=True) + fail("Archive is missing bundle.json") - # ── NCCL conflict handling ─────────────────────────────────────────── + try: + with open(bundle_json_path, "r") as f: + bundle_meta = json.load(f) + except Exception as e: + shutil.rmtree(staging_dir, ignore_errors=True) + fail(f"Invalid bundle.json: {e}") - handle_nccl_conflict() + version = bundle_meta.get("version", manifest.get("imageVersion", "unknown")) + model_ids = bundle_meta.get("models", []) - # ── Download models ────────────────────────────────────────────────── + # -- Move site-packages -- + emit_progress(92, "Installing packages...") + site_packages_dir = get_site_packages_dir(venv_path) + staging_sp = os.path.join(staging_dir, "site-packages") - models = bundle.get("models", []) - model_ids = [m.get("id", f"model-{i}") for i, m in enumerate(models)] + if os.path.isdir(staging_sp) and site_packages_dir: + move_tree(staging_sp, site_packages_dir) - emit_progress(50, "Downloading models...") + # -- Move models -- + emit_progress(95, "Installing models...") + staging_models = os.path.join(staging_dir, "models") + if os.path.isdir(staging_models): + os.makedirs(models_dir, exist_ok=True) + move_tree(staging_models, models_dir) - os.makedirs(models_dir, exist_ok=True) - failed = download_models(models, models_dir) - - if failed: - fail( - f"Failed to download {len(failed)} model(s): {', '.join(failed)}. " - f"This is usually caused by network issues (timeouts, DNS, or rate limiting). " - f"Check your internet connection and retry the installation." - ) - - # ── Write installed.json ───────────────────────────────────────────── - - emit_progress(98, "Finalizing...") + # -- Apply fixups -- + emit_progress(97, "Finalizing...") + apply_fixups(staging_dir, venv_path) + # -- Write installed.json -- + emit_progress(98, "Recording installation...") installed = read_installed(ai_dir) installed["bundles"][bundle_id] = { "version": version, @@ -667,8 +431,14 @@ def main() -> None: } write_installed_atomic(ai_dir, installed) - # ── Report success ─────────────────────────────────────────────────── + # -- Cleanup -- + if os.path.exists(staging_dir): + shutil.rmtree(staging_dir, ignore_errors=True) + # Clean up downloaded tar (but not if local override) + if not local_path and os.path.exists(tar_path): + os.unlink(tar_path) + # -- Done -- emit_progress(100, "Complete") result = { diff --git a/tests/integration/feature-import.test.ts b/tests/integration/feature-import.test.ts index 2dd5e606..71c25116 100644 --- a/tests/integration/feature-import.test.ts +++ b/tests/integration/feature-import.test.ts @@ -143,6 +143,41 @@ async function buildSymlinkArchive(): Promise { return archivePath; } +async function createBundleTarWithSitePackages( + bundleId: string, + version: string, + modelFiles: Record, + sitePackageFiles: Record, +): Promise { + const tarDir = join(testRoot, `tar-src-${randomUUID()}`); + mkdirSync(tarDir, { recursive: true }); + + writeFileSync( + join(tarDir, "bundle.json"), + JSON.stringify({ bundleId, version, models: Object.keys(modelFiles) }), + ); + + const modelsSubdir = join(tarDir, "models"); + mkdirSync(modelsSubdir, { recursive: true }); + for (const [name, content] of Object.entries(modelFiles)) { + const modelPath = join(modelsSubdir, name); + mkdirSync(join(modelPath, ".."), { recursive: true }); + writeFileSync(modelPath, content); + } + + const spSubdir = join(tarDir, "site-packages"); + mkdirSync(spSubdir, { recursive: true }); + for (const [name, content] of Object.entries(sitePackageFiles)) { + const spPath = join(spSubdir, name); + mkdirSync(join(spPath, ".."), { recursive: true }); + writeFileSync(spPath, content); + } + + const tarPath = join(testRoot, `bundle-${randomUUID()}.tar.gz`); + await tar.create({ gzip: true, file: tarPath, cwd: tarDir }, ["."]); + return tarPath; +} + function resetState(): void { writeFileSync(installedPath, JSON.stringify({ bundles: {} }), "utf-8"); invalidateCache(); @@ -265,6 +300,30 @@ describe("importBundleArchive", () => { }); }); +describe("site-packages import", () => { + beforeEach(resetState); + + it("extracts site-packages into venv site-packages directory", async () => { + const venvSitePackages = join(aiDir, "venv", "lib", "python3.12", "site-packages"); + mkdirSync(venvSitePackages, { recursive: true }); + process.env.PYTHON_VENV_PATH = join(aiDir, "venv"); + + const tarPath = await createBundleTarWithSitePackages( + testBundleId, + testVersion, + { "mediapipe/face.tflite": "model-data" }, + { "fakepkg/__init__.py": "# fake package" }, + ); + + invalidateCache(); + const result = await importBundleArchive(createReadStream(tarPath)); + expect(result.bundleId).toBe(testBundleId); + expect(existsSync(join(venvSitePackages, "fakepkg", "__init__.py"))).toBe(true); + + delete process.env.PYTHON_VENV_PATH; + }); +}); + describe("POST /api/v1/admin/features/import", () => { let app: Awaited>["default"] extends ( ...args: infer _A diff --git a/tests/unit/features/feature-manifest.test.ts b/tests/unit/features/feature-manifest.test.ts index 0a36bc89..71cc523d 100644 --- a/tests/unit/features/feature-manifest.test.ts +++ b/tests/unit/features/feature-manifest.test.ts @@ -32,7 +32,7 @@ function archPackagesInclude(bundleId: string, arch: "amd64" | "arm64", pkg: str describe("Feature manifest structure", () => { it("manifest has valid version fields", () => { - expect(manifest.manifestVersion).toBe(1); + expect(manifest.manifestVersion).toBe(2); expect(manifest.pythonVersion).toBeDefined(); expect(manifest.basePackages).toBeInstanceOf(Array); }); @@ -211,3 +211,69 @@ describe("Feature manifest: enablesTools consistency with shared features", () = } }); }); + +describe("Manifest v2 archive fields", () => { + const ARCH_VARIANTS = ["amd64-gpu", "arm64-cpu"] as const; + + it("manifest version is 2", () => { + expect(manifest.manifestVersion).toBe(2); + }); + + it("has bundleRepo field", () => { + expect(manifest.bundleRepo).toBe("snapotter/feature-bundles"); + }); + + it("every bundle has archives with both arch variants", () => { + for (const [id, bundle] of Object.entries>(bundles)) { + const archives = bundle.archives as Record; + expect(archives, `${id} missing archives`).toBeDefined(); + for (const arch of ARCH_VARIANTS) { + expect(archives[arch], `${id} missing archives["${arch}"]`).toBeDefined(); + } + } + }); + + it("each archive entry has file, sha256, compressedSize, extractedSize", () => { + for (const [id, bundle] of Object.entries>(bundles)) { + const archives = bundle.archives as Record>; + for (const arch of ARCH_VARIANTS) { + const entry = archives[arch]; + expect(typeof entry.file, `${id}/${arch} file should be string`).toBe("string"); + expect(entry.file as string, `${id}/${arch} file should end with .tar.gz`).toMatch( + /\.tar\.gz$/, + ); + expect(typeof entry.sha256, `${id}/${arch} sha256 should be string`).toBe("string"); + expect(entry.sha256 as string, `${id}/${arch} sha256 should be 64 hex chars`).toMatch( + /^[0-9a-f]{64}$/, + ); + expect(typeof entry.compressedSize, `${id}/${arch} compressedSize should be number`).toBe( + "number", + ); + expect( + entry.compressedSize as number, + `${id}/${arch} compressedSize should be > 0`, + ).toBeGreaterThan(0); + expect(typeof entry.extractedSize, `${id}/${arch} extractedSize should be number`).toBe( + "number", + ); + expect( + entry.extractedSize as number, + `${id}/${arch} extractedSize should be > 0`, + ).toBeGreaterThan(0); + } + } + }); + + it("archive file paths include version prefix", () => { + const version = manifest.imageVersion; + for (const [id, bundle] of Object.entries>(bundles)) { + const archives = bundle.archives as Record>; + for (const arch of ARCH_VARIANTS) { + expect( + archives[arch].file as string, + `${id}/${arch} file should include v${version}/`, + ).toContain(`v${version}/`); + } + } + }); +}); diff --git a/tests/unit/features/feature-status.test.ts b/tests/unit/features/feature-status.test.ts index ad619a7c..b446d67d 100644 --- a/tests/unit/features/feature-status.test.ts +++ b/tests/unit/features/feature-status.test.ts @@ -438,6 +438,48 @@ describe("Crash recovery - recoverInterruptedInstalls", () => { mod.recoverInterruptedInstalls(); expect(mod.isFeatureInstalled("ocr")).toBe(false); }); + + it("deletes staging-{bundleId}/ directories", () => { + const staging = join(aiDir, "staging-ocr"); + mkdirSync(staging, { recursive: true }); + writeFileSync(join(staging, "somefile"), "data"); + mod.recoverInterruptedInstalls(); + expect(existsSync(staging)).toBe(false); + }); + + it("deletes multiple staging directories", () => { + mkdirSync(join(aiDir, "staging-ocr"), { recursive: true }); + mkdirSync(join(aiDir, "staging-upscale-enhance"), { recursive: true }); + mod.recoverInterruptedInstalls(); + expect(existsSync(join(aiDir, "staging-ocr"))).toBe(false); + expect(existsSync(join(aiDir, "staging-upscale-enhance"))).toBe(false); + }); + + it("does NOT delete non-staging directories", () => { + const venvDir = join(aiDir, "venv"); + mkdirSync(venvDir, { recursive: true }); + writeFileSync(join(venvDir, "file"), "data"); + mod.recoverInterruptedInstalls(); + expect(existsSync(venvDir)).toBe(true); + }); + + it("deletes stale download files in staging/", () => { + const staging = join(aiDir, "staging"); + mkdirSync(staging, { recursive: true }); + writeFileSync(join(staging, "bundle.tar.gz.partial"), "partial"); + writeFileSync(join(staging, "bundle.tar.gz.meta"), '{"bytesDownloaded":0}'); + mod.recoverInterruptedInstalls(); + expect(existsSync(join(staging, "bundle.tar.gz.partial"))).toBe(false); + expect(existsSync(join(staging, "bundle.tar.gz.meta"))).toBe(false); + }); + + it("deletes orphaned .tar.gz in staging when bundle not installed", () => { + const staging = join(aiDir, "staging"); + mkdirSync(staging, { recursive: true }); + writeFileSync(join(staging, "background-removal-amd64-gpu.tar.gz"), "tar-data"); + mod.recoverInterruptedInstalls(); + expect(existsSync(join(staging, "background-removal-amd64-gpu.tar.gz"))).toBe(false); + }); }); describe("Composite state - getFeatureStates", () => { diff --git a/tests/unit/features/install-feature-prebuilt.test.ts b/tests/unit/features/install-feature-prebuilt.test.ts new file mode 100644 index 00000000..9d1e905d --- /dev/null +++ b/tests/unit/features/install-feature-prebuilt.test.ts @@ -0,0 +1,153 @@ +import { execFileSync, spawnSync } from "node:child_process"; +import { createHash } from "node:crypto"; +import { existsSync, mkdirSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { afterEach, beforeEach, describe, expect, it } from "vitest"; + +const scriptPath = join(process.cwd(), "packages/ai/python/install_feature.py"); + +let tempDir: string; +let aiDir: string; +let modelsDir: string; +let venvDir: string; +let sitePackagesDir: string; +let manifestPath: string; + +beforeEach(() => { + tempDir = mkdtempSync(join(tmpdir(), "snapotter-install-test-")); + aiDir = join(tempDir, "ai"); + modelsDir = join(aiDir, "models"); + venvDir = join(aiDir, "venv"); + sitePackagesDir = join(venvDir, "lib", "python3.12", "site-packages"); + manifestPath = join(tempDir, "feature-manifest.json"); + + mkdirSync(sitePackagesDir, { recursive: true }); + mkdirSync(modelsDir, { recursive: true }); + mkdirSync(join(aiDir, "staging"), { recursive: true }); + writeFileSync(join(aiDir, "installed.json"), JSON.stringify({ bundles: {} })); +}); + +afterEach(() => { + rmSync(tempDir, { recursive: true, force: true }); +}); + +function createTestTar(bundleId: string): { tarPath: string; sha256: string } { + const buildDir = join(tempDir, "build"); + mkdirSync(join(buildDir, "models", "testmodel"), { recursive: true }); + mkdirSync(join(buildDir, "site-packages", "testpkg"), { recursive: true }); + writeFileSync(join(buildDir, "models", "testmodel", "weights.bin"), "model-weights"); + writeFileSync(join(buildDir, "site-packages", "testpkg", "__init__.py"), "# test"); + writeFileSync( + join(buildDir, "bundle.json"), + JSON.stringify({ + bundleId, + version: "1.0.0-test", + arch: "amd64-gpu", + imageVersion: "2.0.0", + pythonVersion: "3.12", + models: ["testmodel"], + }), + ); + + const tarPath = join(tempDir, `${bundleId}-test.tar.gz`); + execFileSync("tar", ["czf", tarPath, "-C", buildDir, "."]); + rmSync(buildDir, { recursive: true }); + + const hash = createHash("sha256").update(readFileSync(tarPath)).digest("hex"); + return { tarPath, sha256: hash }; +} + +function writeManifest(bundleId: string, tarPath: string, sha256: string) { + const size = readFileSync(tarPath).length; + const manifest = { + manifestVersion: 2, + imageVersion: "2.0.0", + pythonVersion: "3.12", + basePackages: [], + bundleRepo: "snapotter/feature-bundles", + bundles: { + [bundleId]: { + name: "Test Bundle", + archives: { + "amd64-gpu": { file: tarPath, sha256, compressedSize: size, extractedSize: size * 2 }, + "arm64-cpu": { file: tarPath, sha256, compressedSize: size, extractedSize: size * 2 }, + }, + models: [{ id: "testmodel", path: "testmodel/weights.bin", minSize: 0 }], + enablesTools: [], + }, + }, + }; + writeFileSync(manifestPath, JSON.stringify(manifest)); +} + +describe("install_feature.py prebuilt mode", () => { + it("extracts models and site-packages from a local tar", () => { + const { tarPath, sha256 } = createTestTar("face-detection"); + writeManifest("face-detection", tarPath, sha256); + + const result = spawnSync("python3", [scriptPath, "face-detection", manifestPath, modelsDir], { + env: { + ...process.env, + DATA_DIR: tempDir, + PYTHON_VENV_PATH: venvDir, + SNAPOTTER_BUNDLE_LOCAL_PATH: tarPath, + }, + timeout: 30_000, + }); + + expect(result.status, `stderr: ${result.stderr?.toString()}`).toBe(0); + expect(existsSync(join(modelsDir, "testmodel", "weights.bin"))).toBe(true); + expect(existsSync(join(sitePackagesDir, "testpkg", "__init__.py"))).toBe(true); + + const installed = JSON.parse(readFileSync(join(aiDir, "installed.json"), "utf-8")); + expect(installed.bundles["face-detection"]).toBeDefined(); + expect(installed.bundles["face-detection"].version).toBe("1.0.0-test"); + }); + + it("exits non-zero when checksum mismatches", () => { + const { tarPath } = createTestTar("face-detection"); + writeManifest("face-detection", tarPath, "badhash".padEnd(64, "0")); + + const result = spawnSync("python3", [scriptPath, "face-detection", manifestPath, modelsDir], { + env: { + ...process.env, + DATA_DIR: tempDir, + PYTHON_VENV_PATH: venvDir, + SNAPOTTER_BUNDLE_LOCAL_PATH: tarPath, + }, + timeout: 30_000, + }); + + expect(result.status).not.toBe(0); + }); + + it("writes progress JSON to stderr", () => { + const { tarPath, sha256 } = createTestTar("face-detection"); + writeManifest("face-detection", tarPath, sha256); + + const result = spawnSync("python3", [scriptPath, "face-detection", manifestPath, modelsDir], { + env: { + ...process.env, + DATA_DIR: tempDir, + PYTHON_VENV_PATH: venvDir, + SNAPOTTER_BUNDLE_LOCAL_PATH: tarPath, + }, + timeout: 30_000, + }); + + const stderr = result.stderr?.toString() ?? ""; + const progressLines = stderr.split("\n").filter((l) => { + try { + const p = JSON.parse(l); + return typeof p.progress === "number"; + } catch { + return false; + } + }); + expect(progressLines.length).toBeGreaterThan(0); + + const last = JSON.parse(progressLines[progressLines.length - 1]); + expect(last.progress).toBe(100); + }); +});