mirror of
https://github.com/snapotter-hq/SnapOtter.git
synced 2026-08-03 07:46:42 +02:00
feat: add feature gating to Python sidecar dispatcher
Check installed.json before exec()-ing AI scripts so that requests for uninstalled feature bundles return a structured error instead of crashing with an ImportError. Also sets U2NET_HOME to the bundled model directory when present.
This commit is contained in:
@@ -17,6 +17,33 @@ import os
|
|||||||
import traceback
|
import traceback
|
||||||
|
|
||||||
|
|
||||||
|
INSTALLED_PATH = os.path.join(os.environ.get("DATA_DIR", "/data"), "ai", "installed.json")
|
||||||
|
MODELS_DIR = os.path.join(os.environ.get("DATA_DIR", "/data"), "ai", "models")
|
||||||
|
|
||||||
|
TOOL_BUNDLE_MAP = {
|
||||||
|
"remove_bg": "background-removal",
|
||||||
|
"detect_faces": "face-detection",
|
||||||
|
"face_landmarks": "face-detection",
|
||||||
|
"red_eye_removal": "face-detection",
|
||||||
|
"inpaint": "object-eraser-colorize",
|
||||||
|
"colorize": "object-eraser-colorize",
|
||||||
|
"upscale": "upscale-enhance",
|
||||||
|
"enhance_faces": "upscale-enhance",
|
||||||
|
"noise_removal": "upscale-enhance",
|
||||||
|
"restore": "photo-restoration",
|
||||||
|
"ocr": "ocr",
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def _get_installed_bundles():
|
||||||
|
try:
|
||||||
|
with open(INSTALLED_PATH) as f:
|
||||||
|
data = json.load(f)
|
||||||
|
return set(data.get("bundles", {}).keys())
|
||||||
|
except (FileNotFoundError, json.JSONDecodeError):
|
||||||
|
return set()
|
||||||
|
|
||||||
|
|
||||||
def emit_progress(percent, stage):
|
def emit_progress(percent, stage):
|
||||||
"""Emit structured progress to stderr."""
|
"""Emit structured progress to stderr."""
|
||||||
print(json.dumps({"progress": percent, "stage": stage}), file=sys.stderr, flush=True)
|
print(json.dumps({"progress": percent, "stage": stage}), file=sys.stderr, flush=True)
|
||||||
@@ -44,6 +71,10 @@ _try_import("gpu", lambda: __import__("gpu"))
|
|||||||
# Heavy ML libraries - import but don't fail if unavailable
|
# Heavy ML libraries - import but don't fail if unavailable
|
||||||
_try_import("rembg", lambda: __import__("rembg"))
|
_try_import("rembg", lambda: __import__("rembg"))
|
||||||
|
|
||||||
|
# Point rembg at the bundled model directory if it exists
|
||||||
|
if os.path.isdir(MODELS_DIR):
|
||||||
|
os.environ.setdefault("U2NET_HOME", os.path.join(MODELS_DIR, "rembg"))
|
||||||
|
|
||||||
|
|
||||||
# ── Script handlers ─────────────────────────────────────────────────
|
# ── Script handlers ─────────────────────────────────────────────────
|
||||||
# Each handler sets sys.argv and calls the script's main() function,
|
# Each handler sets sys.argv and calls the script's main() function,
|
||||||
@@ -59,6 +90,18 @@ def _run_script_main(script_name, args):
|
|||||||
"""
|
"""
|
||||||
script_dir = os.path.dirname(os.path.abspath(__file__))
|
script_dir = os.path.dirname(os.path.abspath(__file__))
|
||||||
|
|
||||||
|
# ── Feature gate: reject scripts whose bundle is not installed ──
|
||||||
|
bundle_id = TOOL_BUNDLE_MAP.get(script_name)
|
||||||
|
if bundle_id:
|
||||||
|
installed = _get_installed_bundles()
|
||||||
|
if bundle_id not in installed:
|
||||||
|
return (json.dumps({
|
||||||
|
"success": False,
|
||||||
|
"error": "feature_not_installed",
|
||||||
|
"feature": bundle_id,
|
||||||
|
"message": f"Feature bundle '{bundle_id}' is not installed"
|
||||||
|
}), 1)
|
||||||
|
|
||||||
# Save original state
|
# Save original state
|
||||||
old_argv = sys.argv
|
old_argv = sys.argv
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user