135 lines
4.6 KiB
Python
135 lines
4.6 KiB
Python
"""Small persistent worker for the local Depth Anything V2 Small model.
|
|
|
|
The Node Local Runtime owns the process and sends one JSON request per line.
|
|
Keeping the pipeline alive avoids loading the model for every canvas conversion.
|
|
The worker is deliberately offline-only: the model must already exist in the
|
|
configured Hugging Face cache.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
import contextlib
|
|
import io
|
|
import json
|
|
import os
|
|
import re
|
|
import sys
|
|
from typing import Any
|
|
|
|
from PIL import Image
|
|
|
|
|
|
MODEL_ID = os.environ.get("CANVAS_DEPTH_MODEL_ID", "depth-anything/Depth-Anything-V2-Small-hf")
|
|
MAX_INPUT_BYTES = 12 * 1024 * 1024
|
|
DATA_URL_PATTERN = re.compile(r"^data:(image/(?:png|jpeg|webp));base64,([A-Za-z0-9+/=_-]+)$")
|
|
PIPELINE: Any = None
|
|
|
|
|
|
def emit(value: dict[str, Any]) -> None:
|
|
sys.stdout.write(json.dumps(value, ensure_ascii=False, separators=(",", ":")) + "\n")
|
|
sys.stdout.flush()
|
|
|
|
|
|
def load_pipeline() -> Any:
|
|
global PIPELINE
|
|
if PIPELINE is not None:
|
|
return PIPELINE
|
|
|
|
try:
|
|
from transformers import pipeline
|
|
from transformers.utils import logging as transformers_logging
|
|
|
|
transformers_logging.set_verbosity_error()
|
|
# Transformers/tqdm can write progress text while loading weights.
|
|
# Keep stdout reserved for the JSON-lines protocol.
|
|
with contextlib.redirect_stdout(sys.stderr):
|
|
PIPELINE = pipeline(
|
|
"depth-estimation",
|
|
model=MODEL_ID,
|
|
device=-1,
|
|
local_files_only=True,
|
|
)
|
|
except Exception as error: # pragma: no cover - depends on local model files
|
|
raise RuntimeError("depth_model_missing") from error
|
|
return PIPELINE
|
|
|
|
|
|
def image_from_data_url(value: Any) -> Image.Image:
|
|
if not isinstance(value, str):
|
|
raise ValueError("depth_input_invalid")
|
|
match = DATA_URL_PATTERN.fullmatch(value)
|
|
if not match:
|
|
raise ValueError("depth_input_invalid")
|
|
try:
|
|
raw = base64.urlsafe_b64decode(match.group(2) + "=" * (-len(match.group(2)) % 4))
|
|
except Exception as error:
|
|
raise ValueError("depth_input_invalid") from error
|
|
if not raw or len(raw) > MAX_INPUT_BYTES:
|
|
raise ValueError("depth_input_invalid")
|
|
try:
|
|
image = Image.open(io.BytesIO(raw)).convert("RGB")
|
|
image.load()
|
|
return image
|
|
except Exception as error:
|
|
raise ValueError("depth_input_invalid") from error
|
|
|
|
|
|
def encode_png(image: Image.Image) -> str:
|
|
output = io.BytesIO()
|
|
image.save(output, format="PNG", optimize=True)
|
|
return base64.b64encode(output.getvalue()).decode("ascii")
|
|
|
|
|
|
def handle(request: dict[str, Any]) -> dict[str, Any]:
|
|
image = image_from_data_url(request.get("dataUrl"))
|
|
try:
|
|
result = load_pipeline()(image)
|
|
depth = result.get("depth") if isinstance(result, dict) else None
|
|
if not isinstance(depth, Image.Image):
|
|
raise RuntimeError("depth_inference_failed")
|
|
return {
|
|
"ok": True,
|
|
"requestId": request.get("requestId"),
|
|
"mimeType": "image/png",
|
|
"width": depth.width,
|
|
"height": depth.height,
|
|
"modelId": MODEL_ID,
|
|
"device": "cpu",
|
|
"pngBase64": encode_png(depth),
|
|
}
|
|
except RuntimeError:
|
|
raise
|
|
except Exception as error: # pragma: no cover - depends on local runtime
|
|
raise RuntimeError("depth_inference_failed") from error
|
|
|
|
|
|
def main() -> None:
|
|
os.environ.setdefault("HF_HUB_DISABLE_TELEMETRY", "1")
|
|
os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")
|
|
os.environ.setdefault("HF_HUB_OFFLINE", "1")
|
|
os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")
|
|
|
|
for line in sys.stdin:
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
request_id: Any = None
|
|
try:
|
|
request = json.loads(line)
|
|
if not isinstance(request, dict):
|
|
raise ValueError("depth_input_invalid")
|
|
request_id = request.get("requestId")
|
|
emit(handle(request))
|
|
except ValueError as error:
|
|
emit({"ok": False, "requestId": request_id, "code": str(error) or "depth_input_invalid", "message": "深度转换输入无效"})
|
|
except RuntimeError as error:
|
|
code = str(error) or "depth_inference_failed"
|
|
emit({"ok": False, "requestId": request_id, "code": code, "message": "本机深度模型不可用" if code == "depth_model_missing" else "本机深度推理失败"})
|
|
except Exception:
|
|
emit({"ok": False, "requestId": request_id, "code": "depth_inference_failed", "message": "本机深度推理失败"})
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|