feat(canvas): 图片转换插件 - 增加本地媒体转换节点与运行时
This commit is contained in:
QAyong
134
canvas-agent/python/depth_anything_runtime.py
Normal file
134
canvas-agent/python/depth_anything_runtime.py
Normal file
@@ -0,0 +1,134 @@
|
||||
"""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()
|
||||
132
canvas-agent/python/lineart_runtime.py
Normal file
132
canvas-agent/python/lineart_runtime.py
Normal file
@@ -0,0 +1,132 @@
|
||||
"""Small persistent worker for the local ControlNet Aux lineart detector.
|
||||
|
||||
This worker only extracts a lineart control image. It does not load Stable
|
||||
Diffusion or a ControlNet generation pipeline, so it stays suitable for a
|
||||
CPU-only notebook. The Node Local Runtime owns the process and sends one JSON
|
||||
request per line; stdout is reserved for that protocol.
|
||||
"""
|
||||
|
||||
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_LINEART_MODEL_ID", "lllyasviel/Annotators")
|
||||
MAX_INPUT_BYTES = 12 * 1024 * 1024
|
||||
DATA_URL_PATTERN = re.compile(r"^data:(image/(?:png|jpeg|webp));base64,([A-Za-z0-9+/=_-]+)$")
|
||||
DETECTOR: 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_detector() -> Any:
|
||||
global DETECTOR
|
||||
if DETECTOR is not None:
|
||||
return DETECTOR
|
||||
|
||||
try:
|
||||
from controlnet_aux import LineartDetector
|
||||
|
||||
# The detector downloads only the lineart preprocessor weights from the
|
||||
# configured Hugging Face cache. It is intentionally not a diffusion
|
||||
# pipeline and never loads Stable Diffusion or ControlNet weights.
|
||||
with contextlib.redirect_stdout(sys.stderr):
|
||||
DETECTOR = LineartDetector.from_pretrained(MODEL_ID)
|
||||
except Exception as error: # pragma: no cover - depends on local files
|
||||
raise RuntimeError("lineart_model_missing") from error
|
||||
return DETECTOR
|
||||
|
||||
|
||||
def image_from_data_url(value: Any) -> Image.Image:
|
||||
if not isinstance(value, str):
|
||||
raise ValueError("lineart_input_invalid")
|
||||
match = DATA_URL_PATTERN.fullmatch(value)
|
||||
if not match:
|
||||
raise ValueError("lineart_input_invalid")
|
||||
try:
|
||||
raw = base64.urlsafe_b64decode(match.group(2) + "=" * (-len(match.group(2)) % 4))
|
||||
except Exception as error:
|
||||
raise ValueError("lineart_input_invalid") from error
|
||||
if not raw or len(raw) > MAX_INPUT_BYTES:
|
||||
raise ValueError("lineart_input_invalid")
|
||||
try:
|
||||
image = Image.open(io.BytesIO(raw)).convert("RGB")
|
||||
image.load()
|
||||
return image
|
||||
except Exception as error:
|
||||
raise ValueError("lineart_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:
|
||||
lineart = load_detector()(image)
|
||||
if not isinstance(lineart, Image.Image):
|
||||
raise RuntimeError("lineart_inference_failed")
|
||||
return {
|
||||
"ok": True,
|
||||
"requestId": request.get("requestId"),
|
||||
"mimeType": "image/png",
|
||||
"width": lineart.width,
|
||||
"height": lineart.height,
|
||||
"modelId": MODEL_ID,
|
||||
"device": "cpu",
|
||||
"pngBase64": encode_png(lineart.convert("RGB")),
|
||||
}
|
||||
except RuntimeError:
|
||||
raise
|
||||
except Exception as error: # pragma: no cover - depends on local runtime
|
||||
raise RuntimeError("lineart_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:
|
||||
# PowerShell and a few Windows launchers may prepend a UTF-8 BOM to
|
||||
# the first piped line. The Node runtime sends plain UTF-8, but
|
||||
# accepting the marker keeps the worker protocol tolerant for manual
|
||||
# diagnostics as well.
|
||||
line = line.lstrip("\ufeff").strip()
|
||||
if not line:
|
||||
continue
|
||||
request_id: Any = None
|
||||
try:
|
||||
request = json.loads(line)
|
||||
if not isinstance(request, dict):
|
||||
raise ValueError("lineart_input_invalid")
|
||||
request_id = request.get("requestId")
|
||||
emit(handle(request))
|
||||
except ValueError as error:
|
||||
emit({"ok": False, "requestId": request_id, "code": str(error) or "lineart_input_invalid", "message": "线稿转换输入无效"})
|
||||
except RuntimeError as error:
|
||||
code = str(error) or "lineart_inference_failed"
|
||||
message = "本地 AI 线稿模型或依赖尚未安装" if code == "lineart_model_missing" else "本地 AI 线稿推理失败"
|
||||
emit({"ok": False, "requestId": request_id, "code": code, "message": message})
|
||||
except Exception:
|
||||
emit({"ok": False, "requestId": request_id, "code": "lineart_inference_failed", "message": "本地 AI 线稿推理失败"})
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
204
canvas-agent/python/pose_runtime.py
Normal file
204
canvas-agent/python/pose_runtime.py
Normal file
@@ -0,0 +1,204 @@
|
||||
"""Persistent CPU pose preprocessor for the local media conversion node.
|
||||
|
||||
The first notebook-friendly implementation uses the body-only OpenPose
|
||||
preprocessor from ControlNet Aux. It produces a colored skeleton map and a
|
||||
small list of normalized keypoints; it does not load a diffusion model.
|
||||
"""
|
||||
|
||||
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_POSE_MODEL_ID", "lllyasviel/Annotators")
|
||||
MAX_INPUT_BYTES = 12 * 1024 * 1024
|
||||
DETECT_RESOLUTION = max(256, min(1024, int(os.environ.get("CANVAS_POSE_DETECT_RESOLUTION", "512"))))
|
||||
DATA_URL_PATTERN = re.compile(r"^data:(image/(?:png|jpeg|webp));base64,([A-Za-z0-9+/=_-]+)$")
|
||||
DETECTOR: 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_detector() -> Any:
|
||||
global DETECTOR
|
||||
if DETECTOR is not None:
|
||||
return DETECTOR
|
||||
|
||||
try:
|
||||
import numpy as np
|
||||
from controlnet_aux.open_pose.body import Body
|
||||
from huggingface_hub import hf_hub_download
|
||||
|
||||
cache_dir = os.environ.get("HF_HUB_CACHE") or None
|
||||
with contextlib.redirect_stdout(sys.stderr):
|
||||
model_path = hf_hub_download(
|
||||
repo_id=MODEL_ID,
|
||||
filename="body_pose_model.pth",
|
||||
cache_dir=cache_dir,
|
||||
local_files_only=True,
|
||||
)
|
||||
DETECTOR = (Body(model_path), np)
|
||||
except Exception as error: # pragma: no cover - depends on local model files
|
||||
raise RuntimeError("pose_model_missing") from error
|
||||
return DETECTOR
|
||||
|
||||
|
||||
def image_from_data_url(value: Any) -> Image.Image:
|
||||
if not isinstance(value, str):
|
||||
raise ValueError("pose_input_invalid")
|
||||
match = DATA_URL_PATTERN.fullmatch(value)
|
||||
if not match:
|
||||
raise ValueError("pose_input_invalid")
|
||||
try:
|
||||
raw = base64.urlsafe_b64decode(match.group(2) + "=" * (-len(match.group(2)) % 4))
|
||||
except Exception as error:
|
||||
raise ValueError("pose_input_invalid") from error
|
||||
if not raw or len(raw) > MAX_INPUT_BYTES:
|
||||
raise ValueError("pose_input_invalid")
|
||||
try:
|
||||
image = Image.open(io.BytesIO(raw)).convert("RGB")
|
||||
image.load()
|
||||
return image
|
||||
except Exception as error:
|
||||
raise ValueError("pose_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 normalize_people(bodies: list[Any], width: int, height: int) -> list[dict[str, Any]]:
|
||||
people: list[dict[str, Any]] = []
|
||||
for body in bodies:
|
||||
keypoints = []
|
||||
visible = []
|
||||
for keypoint in body.keypoints:
|
||||
if keypoint is None:
|
||||
keypoints.append(None)
|
||||
else:
|
||||
normalized_x = max(0.0, min(1.0, float(keypoint.x) / float(width)))
|
||||
normalized_y = max(0.0, min(1.0, float(keypoint.y) / float(height)))
|
||||
keypoints.append({
|
||||
"x": normalized_x,
|
||||
"y": normalized_y,
|
||||
"score": float(keypoint.score),
|
||||
})
|
||||
visible.append((normalized_x, normalized_y, float(keypoint.score)))
|
||||
if visible:
|
||||
xs = [point[0] for point in visible]
|
||||
ys = [point[1] for point in visible]
|
||||
confidence = sum(point[2] for point in visible) / len(visible)
|
||||
box = {"x": min(xs), "y": min(ys), "width": max(xs) - min(xs), "height": max(ys) - min(ys)}
|
||||
else:
|
||||
confidence = 0.0
|
||||
box = {"x": 0.0, "y": 0.0, "width": 0.0, "height": 0.0}
|
||||
people.append({
|
||||
"keypoints": keypoints,
|
||||
"score": float(body.total_score),
|
||||
"confidence": confidence,
|
||||
"box": box,
|
||||
"parts": int(body.total_parts),
|
||||
})
|
||||
return people
|
||||
|
||||
|
||||
def handle(request: dict[str, Any]) -> dict[str, Any]:
|
||||
image = image_from_data_url(request.get("dataUrl"))
|
||||
try:
|
||||
import cv2
|
||||
import numpy as np
|
||||
from controlnet_aux.open_pose.util import draw_bodypose
|
||||
from controlnet_aux.util import HWC3, resize_image
|
||||
|
||||
detector, _ = load_detector()
|
||||
source = HWC3(np.asarray(image, dtype=np.uint8))
|
||||
detected = resize_image(source, DETECT_RESOLUTION)
|
||||
candidate, subset = detector(detected[:, :, ::-1].copy())
|
||||
bodies = detector.format_body_result(candidate, subset)
|
||||
if not bodies:
|
||||
raise RuntimeError("pose_no_person")
|
||||
|
||||
height, width = detected.shape[:2]
|
||||
people = normalize_people(bodies, width, height)
|
||||
canvas = np.zeros_like(detected, dtype=np.uint8)
|
||||
for body in bodies:
|
||||
normalized = [
|
||||
None if keypoint is None else type(keypoint)(
|
||||
x=float(keypoint.x) / float(width),
|
||||
y=float(keypoint.y) / float(height),
|
||||
score=keypoint.score,
|
||||
id=keypoint.id,
|
||||
)
|
||||
for keypoint in body.keypoints
|
||||
]
|
||||
draw_bodypose(canvas, normalized)
|
||||
|
||||
output = cv2.resize(canvas, (image.width, image.height), interpolation=cv2.INTER_LINEAR)
|
||||
output_image = Image.fromarray(output.astype(np.uint8), mode="RGB")
|
||||
return {
|
||||
"ok": True,
|
||||
"requestId": request.get("requestId"),
|
||||
"mimeType": "image/png",
|
||||
"width": image.width,
|
||||
"height": image.height,
|
||||
"modelId": MODEL_ID,
|
||||
"device": "cpu",
|
||||
"backend": "openpose-body",
|
||||
"personCount": len(bodies),
|
||||
"people": people,
|
||||
"pngBase64": encode_png(output_image),
|
||||
}
|
||||
except RuntimeError:
|
||||
raise
|
||||
except Exception as error: # pragma: no cover - depends on local runtime
|
||||
raise RuntimeError("pose_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.lstrip("\ufeff").strip()
|
||||
if not line:
|
||||
continue
|
||||
request_id: Any = None
|
||||
try:
|
||||
request = json.loads(line)
|
||||
if not isinstance(request, dict):
|
||||
raise ValueError("pose_input_invalid")
|
||||
request_id = request.get("requestId")
|
||||
emit(handle(request))
|
||||
except ValueError as error:
|
||||
emit({"ok": False, "requestId": request_id, "code": str(error) or "pose_input_invalid", "message": "姿态转换输入无效"})
|
||||
except RuntimeError as error:
|
||||
code = str(error) or "pose_inference_failed"
|
||||
if code == "pose_model_missing":
|
||||
message = "本地姿态模型或依赖尚未安装"
|
||||
elif code == "pose_no_person":
|
||||
message = "未检测到可用人物姿态"
|
||||
else:
|
||||
message = "本地姿态推理失败"
|
||||
emit({"ok": False, "requestId": request_id, "code": code, "message": message})
|
||||
except Exception:
|
||||
emit({"ok": False, "requestId": request_id, "code": "pose_inference_failed", "message": "本地姿态推理失败"})
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
1
canvas-agent/python/requirements-lineart.txt
Normal file
1
canvas-agent/python/requirements-lineart.txt
Normal file
@@ -0,0 +1 @@
|
||||
controlnet-aux==0.0.10
|
||||
Reference in New Issue
Block a user