SnapOtter/packages/ai/python/dispatcher.py

"""
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 gc
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.

    A drain thread reads the pipe concurrently to prevent deadlock when
    scripts produce more than 64 KB of stdout (e.g. ONNX runtime logging).
    """
    import threading

    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)

    # Drain the pipe in a background thread so the pipe buffer never fills.
    captured_chunks = []

    def _drain():
        with os.fdopen(read_fd, "r") as f:
            while True:
                chunk = f.read(8192)
                if not chunk:
                    break
                captured_chunks.append(chunk)

    drain_thread = threading.Thread(target=_drain, daemon=True)
    drain_thread.start()

    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 (closes the pipe write end, unblocking the drain thread)
        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

    drain_thread.join(timeout=10)
    captured = "".join(captured_chunks)

    return captured.strip(), exit_code


# ── Main loop ───────────────────────────────────────────────────────


MAX_REQUESTS = int(os.environ.get("DISPATCHER_MAX_REQUESTS", "50"))


def _cleanup_after_request():
    """Free unreferenced objects and GPU memory after each request."""
    gc.collect()
    try:
        import torch
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
    except ImportError:
        pass


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}. Max requests: {MAX_REQUESTS}. Modules: {list(available_modules.keys())}", file=sys.stderr, flush=True)

    request_count = 0

    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()

        _cleanup_after_request()
        request_count += 1

        if request_count >= MAX_REQUESTS:
            print(f"[dispatcher] Reached max requests ({MAX_REQUESTS}), shutting down for restart",
                  file=sys.stderr, flush=True)
            break


if __name__ == "__main__":
    main()