mirror of
https://github.com/snapotter-hq/SnapOtter.git
synced 2026-08-03 07:46:42 +02:00
ashim-hq
e7eea34080
The torchvision compatibility shim for basicsr 1.4.2 was missing the parent-package binding and only proxied a single attribute, causing upscale and enhance-faces to fail at import time. The fix adds a __getattr__ proxy for all attributes, binds the shim to the parent package, and installs it in the dispatcher at startup for defense-in-depth. Also removes unused anyInstalling variable, redundant `as any` cast, and applies Biome formatting fixes across the codebase.
247 lines
8.3 KiB
Python
247 lines
8.3 KiB
Python
"""
|
|
Persistent Python sidecar dispatcher.
|
|
|
|
Runs as a long-lived process. Reads JSON requests from stdin (one per line),
|
|
dispatches to the appropriate AI handler, writes JSON responses to stdout.
|
|
Progress emissions continue via stderr (unchanged from the standalone scripts).
|
|
|
|
Request format: {"id": "uuid", "script": "remove_bg", "args": [...]}
|
|
Response format: {"id": "uuid", "stdout": "...", "exitCode": 0}
|
|
|
|
Pre-imports heavy libraries at startup to eliminate cold-start latency.
|
|
"""
|
|
import sys
|
|
import json
|
|
import io
|
|
import os
|
|
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):
|
|
"""Emit structured progress to stderr."""
|
|
print(json.dumps({"progress": percent, "stage": stage}), file=sys.stderr, flush=True)
|
|
|
|
|
|
# ── basicsr / torchvision compatibility shim ──────────────────────────
|
|
# basicsr 1.4.2 (pulled in by realesrgan) does:
|
|
# from torchvision.transforms.functional_tensor import rgb_to_grayscale
|
|
# but torchvision >= 0.17 removed the functional_tensor submodule,
|
|
# merging everything into torchvision.transforms.functional.
|
|
# We install a shim module ONCE here so every script in this process
|
|
# benefits, rather than relying on each script to patch individually.
|
|
try:
|
|
import torchvision.transforms.functional_tensor # noqa: F401
|
|
except (ImportError, ModuleNotFoundError):
|
|
try:
|
|
import types
|
|
import torchvision.transforms.functional as _F
|
|
import torchvision.transforms
|
|
|
|
_shim = types.ModuleType("torchvision.transforms.functional_tensor")
|
|
_shim.__getattr__ = lambda name: getattr(_F, name)
|
|
_shim.rgb_to_grayscale = _F.rgb_to_grayscale
|
|
sys.modules["torchvision.transforms.functional_tensor"] = _shim
|
|
torchvision.transforms.functional_tensor = _shim
|
|
print("[dispatcher] Installed torchvision.transforms.functional_tensor shim",
|
|
file=sys.stderr, flush=True)
|
|
except (ImportError, AttributeError):
|
|
# torchvision not installed yet — shim not needed until
|
|
# the upscale-enhance bundle is installed.
|
|
pass
|
|
except Exception:
|
|
# Catch-all so dispatcher startup is never blocked.
|
|
pass
|
|
|
|
# ── Pre-import heavy libraries ──────────────────────────────────────
|
|
# These imports are the main source of cold-start latency.
|
|
# By importing once at startup, subsequent requests skip the import cost.
|
|
|
|
available_modules = {}
|
|
|
|
|
|
def _try_import(name, import_fn):
|
|
try:
|
|
available_modules[name] = import_fn()
|
|
except ImportError as e:
|
|
print(f"[dispatcher] Module '{name}' not available: {e}", file=sys.stderr, flush=True)
|
|
|
|
|
|
_try_import("PIL", lambda: __import__("PIL"))
|
|
_try_import("mediapipe", lambda: __import__("mediapipe"))
|
|
_try_import("numpy", lambda: __import__("numpy"))
|
|
_try_import("gpu", lambda: __import__("gpu"))
|
|
|
|
# Heavy ML libraries - import but don't fail if unavailable
|
|
_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 ─────────────────────────────────────────────────
|
|
# Each handler sets sys.argv and calls the script's main() function,
|
|
# capturing stdout. The scripts remain unchanged.
|
|
|
|
|
|
def _run_script_main(script_name, args):
|
|
"""
|
|
Import and run a script's main() function, capturing its stdout output.
|
|
|
|
Since some scripts (like remove_bg.py) manipulate file descriptors directly
|
|
(os.dup2), we use a pipe at the fd level rather than StringIO.
|
|
"""
|
|
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
|
|
old_argv = sys.argv
|
|
|
|
# Create a pipe to capture stdout at the fd level
|
|
read_fd, write_fd = os.pipe()
|
|
|
|
# Save the real stdout fd
|
|
real_stdout_fd = os.dup(1)
|
|
|
|
# Redirect fd 1 to our pipe's write end
|
|
os.dup2(write_fd, 1)
|
|
os.close(write_fd)
|
|
|
|
# Also redirect sys.stdout to the same fd
|
|
old_sys_stdout = sys.stdout
|
|
sys.stdout = os.fdopen(1, "w", closefd=False)
|
|
|
|
exit_code = 0
|
|
try:
|
|
sys.argv = ["script.py"] + args
|
|
|
|
# Load and run the script
|
|
script_path = os.path.join(script_dir, script_name + ".py")
|
|
|
|
module_globals = {"__name__": "__main__", "__file__": script_path}
|
|
|
|
with open(script_path) as f:
|
|
code = compile(f.read(), script_path, "exec")
|
|
|
|
# Run the compiled script in its own namespace
|
|
exec(code, module_globals) # noqa: S102 - trusted internal scripts only
|
|
|
|
except SystemExit as e:
|
|
exit_code = e.code if isinstance(e.code, int) else 1
|
|
except Exception as e:
|
|
# Log full traceback to stderr for diagnostics
|
|
traceback.print_exc(file=sys.stderr)
|
|
# Write error to the captured stdout
|
|
sys.stdout.write(json.dumps({"success": False, "error": str(e)}) + "\n")
|
|
sys.stdout.flush()
|
|
exit_code = 1
|
|
finally:
|
|
# Flush before restoring
|
|
sys.stdout.flush()
|
|
|
|
# Restore stdout fd
|
|
os.dup2(real_stdout_fd, 1)
|
|
os.close(real_stdout_fd)
|
|
|
|
# Restore sys.stdout
|
|
sys.stdout = old_sys_stdout
|
|
|
|
# Restore sys.argv
|
|
sys.argv = old_argv
|
|
|
|
# Read captured output from the pipe
|
|
read_file = os.fdopen(read_fd, "r")
|
|
captured = read_file.read()
|
|
read_file.close()
|
|
|
|
return captured.strip(), exit_code
|
|
|
|
|
|
# ── Main loop ───────────────────────────────────────────────────────
|
|
|
|
|
|
def main():
|
|
# Signal readiness with GPU status
|
|
gpu = False
|
|
try:
|
|
from gpu import gpu_available
|
|
gpu = gpu_available()
|
|
except ImportError as e:
|
|
print(f"[dispatcher] GPU detection failed: {e}", file=sys.stderr, flush=True)
|
|
print(json.dumps({"ready": True, "gpu": gpu}), file=sys.stderr, flush=True)
|
|
print(f"[dispatcher] Ready. GPU: {gpu}. Modules: {list(available_modules.keys())}", file=sys.stderr, flush=True)
|
|
|
|
for line in sys.stdin:
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
|
|
try:
|
|
request = json.loads(line)
|
|
except json.JSONDecodeError:
|
|
continue
|
|
|
|
request_id = request.get("id", "unknown")
|
|
script_name = request.get("script", "")
|
|
args = request.get("args", [])
|
|
|
|
try:
|
|
stdout_output, exit_code = _run_script_main(script_name, args)
|
|
response = {
|
|
"id": request_id,
|
|
"stdout": stdout_output,
|
|
"exitCode": exit_code,
|
|
}
|
|
except Exception as e:
|
|
response = {
|
|
"id": request_id,
|
|
"stdout": json.dumps({"success": False, "error": str(e)}),
|
|
"exitCode": 1,
|
|
}
|
|
|
|
# Write response as a single JSON line to stdout
|
|
sys.stdout.write(json.dumps(response) + "\n")
|
|
sys.stdout.flush()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|