mirror of
https://github.com/snapotter-hq/SnapOtter.git
synced 2026-08-03 07:46:42 +02:00
AshimandClaude Opus 4.6
01d30cfb61
- Pin torch==2.6.0+cu126 and torchvision==0.21.0+cu126 in feature manifest to prevent NCCL symbol mismatch on CUDA 12.6 base images - Move lpips after torch in install order to prevent wrong version resolution from PyPI - Add einops to upscale-enhance common deps (required by SCUNet) - Update cpu_fallback_packages to handle multi-package CUDA torch entries on amd64 without GPU - Fix gpu.py ONNX CUDA detection: replace hardcoded .so path with cross-platform session smoke-test - Fix os.dup(1) crashes on Windows in upscale, enhance_faces, and noise_removal by wrapping in try/except with sys.stderr fallback - Guard top-level numpy/cv2 imports in colorize.py and restore.py with helpful error messages - Add weights_only=False fallback for torch.load in noise_removal - Fix integration tests to accept 501 for uninstalled AI features and 422 for missing system tools (exiftool, libheif) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
518 lines
18 KiB
Python
518 lines
18 KiB
Python
"""Install a feature bundle: pip packages + model downloads.
|
|
|
|
Invoked by the Node.js backend as a subprocess.
|
|
|
|
Usage:
|
|
python3 install_feature.py <bundleId> <manifestPath> <modelsDir>
|
|
|
|
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 json
|
|
import os
|
|
import platform
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
import urllib.error
|
|
import urllib.request
|
|
from datetime import datetime, timezone
|
|
|
|
|
|
# ── Helpers ──────────────────────────────────────────────────────────────
|
|
|
|
|
|
def emit_progress(percent: int, stage: str) -> None:
|
|
"""Emit a progress update via stderr JSON line."""
|
|
sys.stderr.write(json.dumps({"progress": percent, "stage": stage}) + "\n")
|
|
sys.stderr.flush()
|
|
|
|
|
|
def fail(message: str) -> None:
|
|
"""Print error to stderr and exit non-zero."""
|
|
sys.stderr.write(json.dumps({"error": message}) + "\n")
|
|
sys.stderr.flush()
|
|
sys.exit(1)
|
|
|
|
|
|
def detect_arch() -> str:
|
|
"""Return 'arm64' or 'amd64' based on the host machine."""
|
|
machine = platform.machine().lower()
|
|
if machine in ("aarch64", "arm64"):
|
|
return "arm64"
|
|
return "amd64"
|
|
|
|
|
|
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,
|
|
)
|
|
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.
|
|
|
|
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.6.0+cu126 torchvision==0.21.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, strip CUDA suffix
|
|
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)
|
|
# Drop --index-url and its argument (not needed for CPU torch)
|
|
result.extend(cpu_pkgs)
|
|
continue
|
|
|
|
name = pkg.split("==")[0].split(">=")[0].split("[")[0].strip()
|
|
if name in replacements:
|
|
version = pkg[len(name):] # e.g. "==1.20.1"
|
|
result.append(replacements[name] + version)
|
|
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_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"]
|
|
if extra_flags:
|
|
cmd.extend(extra_flags)
|
|
|
|
# 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:
|
|
raise RuntimeError(
|
|
f"pip install failed for '{package}': {result.stderr.strip()}"
|
|
)
|
|
|
|
|
|
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
|
|
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)
|
|
|
|
|
|
def handle_nccl_conflict() -> None:
|
|
"""Re-install torch's NCCL dependency if both torch and paddlepaddle-gpu coexist.
|
|
|
|
PaddlePaddle ships its own NCCL, which can conflict with the version
|
|
that torch expects. Force-reinstalling torch's pinned nccl resolves this.
|
|
"""
|
|
try:
|
|
from importlib.metadata import PackageNotFoundError, requires
|
|
|
|
# Only needed if both torch AND paddlepaddle-gpu are installed
|
|
try:
|
|
requires("torch")
|
|
except PackageNotFoundError:
|
|
return
|
|
try:
|
|
requires("paddlepaddle-gpu")
|
|
except PackageNotFoundError:
|
|
return
|
|
|
|
# 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": "ashim-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)
|
|
|
|
|
|
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
|
|
|
|
from rembg import new_session
|
|
new_session(model_name)
|
|
|
|
|
|
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]
|
|
|
|
snapshot_download(**kwargs)
|
|
|
|
# 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"
|
|
)
|
|
|
|
|
|
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"
|
|
)
|
|
|
|
|
|
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 []
|
|
|
|
failed: list[str] = []
|
|
total = len(models)
|
|
|
|
def _download(idx: int, model: dict) -> tuple[str, Exception | None]:
|
|
model_id = model.get("id", f"model-{idx}")
|
|
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=4) 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
|
|
|
|
|
|
# ── installed.json management ────────────────────────────────────────────
|
|
|
|
|
|
def read_installed(ai_dir: str) -> dict:
|
|
"""Read the current installed.json, returning empty structure if missing."""
|
|
path = os.path.join(ai_dir, "installed.json")
|
|
if not os.path.exists(path):
|
|
return {"bundles": {}}
|
|
try:
|
|
with open(path, "r") as f:
|
|
return json.load(f)
|
|
except (json.JSONDecodeError, OSError):
|
|
return {"bundles": {}}
|
|
|
|
|
|
def write_installed_atomic(ai_dir: str, data: dict) -> None:
|
|
"""Write installed.json atomically (write .tmp then rename)."""
|
|
path = os.path.join(ai_dir, "installed.json")
|
|
tmp_path = path + ".tmp"
|
|
with open(tmp_path, "w") as f:
|
|
json.dump(data, f, indent=2)
|
|
f.write("\n")
|
|
os.rename(tmp_path, path)
|
|
|
|
|
|
# ── Main ─────────────────────────────────────────────────────────────────
|
|
|
|
|
|
def main() -> None:
|
|
if len(sys.argv) < 4:
|
|
fail(
|
|
f"Usage: {sys.argv[0]} <bundleId> <manifestPath> <modelsDir>\n"
|
|
f"Got {len(sys.argv) - 1} argument(s)"
|
|
)
|
|
|
|
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)
|
|
|
|
# ── 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}")
|
|
|
|
bundles = manifest.get("bundles", {})
|
|
if bundle_id not in bundles:
|
|
fail(f"Bundle '{bundle_id}' not found in manifest")
|
|
|
|
bundle = bundles[bundle_id]
|
|
version = manifest.get("imageVersion", "0.0.0")
|
|
|
|
# ── Detect architecture ──────────────────────────────────────────────
|
|
|
|
arch = detect_arch()
|
|
emit_progress(1, f"Architecture: {arch}")
|
|
|
|
# ── Disk space pre-check ─────────────────────────────────────────────
|
|
|
|
check_disk_space(models_dir)
|
|
|
|
# ── Install pip packages ─────────────────────────────────────────────
|
|
|
|
emit_progress(2, "Installing packages...")
|
|
try:
|
|
install_packages(bundle, arch)
|
|
except RuntimeError as e:
|
|
fail(str(e))
|
|
|
|
emit_progress(50, "Packages installed")
|
|
|
|
# ── NCCL conflict handling ───────────────────────────────────────────
|
|
|
|
handle_nccl_conflict()
|
|
|
|
# ── Download models ──────────────────────────────────────────────────
|
|
|
|
models = bundle.get("models", [])
|
|
model_ids = [m.get("id", f"model-{i}") for i, m in enumerate(models)]
|
|
|
|
emit_progress(50, "Downloading models...")
|
|
|
|
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)}"
|
|
)
|
|
|
|
# ── Write installed.json ─────────────────────────────────────────────
|
|
|
|
emit_progress(98, "Finalizing...")
|
|
|
|
installed = read_installed(ai_dir)
|
|
installed["bundles"][bundle_id] = {
|
|
"version": version,
|
|
"installedAt": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
|
|
"models": model_ids,
|
|
}
|
|
write_installed_atomic(ai_dir, installed)
|
|
|
|
# ── Report success ───────────────────────────────────────────────────
|
|
|
|
emit_progress(100, "Complete")
|
|
|
|
result = {
|
|
"success": True,
|
|
"bundleId": bundle_id,
|
|
"version": version,
|
|
"models": model_ids,
|
|
}
|
|
print(json.dumps(result))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|