328 lines
13 KiB
TypeScript
328 lines
13 KiB
TypeScript
import fs from "node:fs";
|
|
import path from "node:path";
|
|
import { fileURLToPath } from "node:url";
|
|
import { spawn, type ChildProcessWithoutNullStreams } from "node:child_process";
|
|
|
|
import type { Request, RequestHandler, Response } from "express";
|
|
|
|
import type { LocalRuntimeModule } from "../local-runtime.js";
|
|
|
|
const DEFAULT_MODEL_ID = "depth-anything/Depth-Anything-V2-Small-hf";
|
|
const MAX_INPUT_BYTES = 12 * 1024 * 1024;
|
|
const MAX_PROTOCOL_LINE_BYTES = 48 * 1024 * 1024;
|
|
const DATA_URL_PATTERN = /^data:(image\/(?:png|jpeg|webp));base64,([A-Za-z0-9+/=_-]+)$/;
|
|
|
|
export type DepthRuntimePaths = {
|
|
pythonPath: string;
|
|
scriptPath: string;
|
|
hfHome: string;
|
|
modelId: string;
|
|
configured: boolean;
|
|
modelCached: boolean;
|
|
reason?: "python_missing" | "worker_missing";
|
|
};
|
|
|
|
export type DepthRuntimeStatus = {
|
|
configured: boolean;
|
|
modelCached: boolean;
|
|
loaded: boolean;
|
|
device: "cpu";
|
|
modelId: string;
|
|
reason?: string;
|
|
};
|
|
|
|
type DepthWorkerSuccess = {
|
|
ok: true;
|
|
requestId: string;
|
|
mimeType: "image/png";
|
|
width: number;
|
|
height: number;
|
|
modelId: string;
|
|
device: "cpu";
|
|
pngBase64: string;
|
|
};
|
|
|
|
export class DepthRuntimeError extends Error {
|
|
constructor(readonly code: string, message: string, readonly statusCode = 503) {
|
|
super(message);
|
|
this.name = "DepthRuntimeError";
|
|
}
|
|
}
|
|
|
|
export function createDepthEstimationHttpModule(): LocalRuntimeModule {
|
|
const worker = new DepthEstimationWorker();
|
|
return {
|
|
descriptor: {
|
|
id: "depth-estimation",
|
|
displayName: "Depth Anything V2 本机深度",
|
|
apiVersion: 1,
|
|
scopes: ["depth:status", "depth:run"],
|
|
},
|
|
routes: [
|
|
{
|
|
method: "GET",
|
|
path: "/depth-estimation/status",
|
|
scope: "depth:status",
|
|
handler: asyncRoute(async (_request, response) => {
|
|
const status = worker.status();
|
|
response.json({
|
|
ok: true,
|
|
module: "depth-estimation",
|
|
apiVersion: 1,
|
|
ready: status.configured && status.modelCached,
|
|
...status,
|
|
});
|
|
}),
|
|
},
|
|
{
|
|
method: "POST",
|
|
path: "/depth-estimation/run",
|
|
scope: "depth:run",
|
|
handler: asyncRoute(async (request, response) => {
|
|
const body = parseBody(request);
|
|
const dataUrl = parseDepthInput(body);
|
|
const result = await worker.run({ dataUrl });
|
|
response.json({
|
|
ok: true,
|
|
module: "depth-estimation",
|
|
apiVersion: 1,
|
|
result,
|
|
});
|
|
}),
|
|
},
|
|
],
|
|
dispose: () => worker.dispose(),
|
|
publicHealth: () => {
|
|
const status = worker.status();
|
|
return { depthEstimation: status.configured && status.modelCached ? "available" : "unavailable" };
|
|
},
|
|
};
|
|
}
|
|
|
|
export function resolveDepthRuntimePaths(env: NodeJS.ProcessEnv = process.env): DepthRuntimePaths {
|
|
const moduleRoot = path.resolve(path.dirname(fileURLToPath(import.meta.url)), "../..");
|
|
const configuredProjectRoot = env.CANVAS_PROJECT_ROOT?.trim();
|
|
const projectRoots = [
|
|
configuredProjectRoot ? path.resolve(configuredProjectRoot) : undefined,
|
|
path.resolve(moduleRoot, ".."),
|
|
path.resolve(process.cwd()),
|
|
].filter((value): value is string => Boolean(value));
|
|
const projectRoot = configuredProjectRoot
|
|
? path.resolve(configuredProjectRoot)
|
|
: projectRoots.find((value) => fs.existsSync(path.join(value, ".local"))) || projectRoots[0]!;
|
|
const scriptCandidates = [
|
|
path.join(moduleRoot, "python", "depth_anything_runtime.py"),
|
|
path.join(projectRoot, "canvas-agent", "python", "depth_anything_runtime.py"),
|
|
];
|
|
const scriptPath = scriptCandidates.find((value) => fs.existsSync(value)) || scriptCandidates[0]!;
|
|
const pythonCandidates = env.CANVAS_DEPTH_PYTHON?.trim()
|
|
? [path.resolve(env.CANVAS_DEPTH_PYTHON.trim())]
|
|
: process.platform === "win32"
|
|
? [path.join(projectRoot, ".local", "depth-anything-v2", "venv", "Scripts", "python.exe")]
|
|
: [path.join(projectRoot, ".local", "depth-anything-v2", "venv", "bin", "python")];
|
|
const pythonPath = pythonCandidates[0]!;
|
|
const hfHome = path.resolve(env.CANVAS_DEPTH_HF_HOME?.trim() || path.join(projectRoot, ".local", "cache", "huggingface"));
|
|
const modelId = env.CANVAS_DEPTH_MODEL_ID?.trim() || DEFAULT_MODEL_ID;
|
|
const modelCacheDir = path.join(hfHome, "hub", `models--${modelId.replaceAll("/", "--")}`);
|
|
const workerExists = fs.existsSync(scriptPath);
|
|
const pythonExists = fs.existsSync(pythonPath);
|
|
return {
|
|
pythonPath,
|
|
scriptPath,
|
|
hfHome,
|
|
modelId,
|
|
configured: workerExists && pythonExists,
|
|
modelCached: workerExists && pythonExists && fs.existsSync(modelCacheDir),
|
|
...(pythonExists ? {} : { reason: "python_missing" as const }),
|
|
...(!workerExists ? { reason: "worker_missing" as const } : {}),
|
|
};
|
|
}
|
|
|
|
export class DepthEstimationWorker {
|
|
private paths = resolveDepthRuntimePaths();
|
|
private child?: ChildProcessWithoutNullStreams;
|
|
private stdoutBuffer = "";
|
|
private stderrBuffer = "";
|
|
private requestCounter = 0;
|
|
private loaded = false;
|
|
private readonly pending = new Map<string, { resolve: (value: DepthWorkerSuccess) => void; reject: (error: unknown) => void }>();
|
|
|
|
status(): DepthRuntimeStatus {
|
|
const paths = this.refreshPaths();
|
|
return {
|
|
configured: paths.configured,
|
|
modelCached: paths.modelCached,
|
|
loaded: this.loaded,
|
|
device: "cpu",
|
|
modelId: paths.modelId,
|
|
...(paths.reason ? { reason: paths.reason } : {}),
|
|
};
|
|
}
|
|
|
|
async run(input: { dataUrl: string }) {
|
|
validateDepthDataUrl(input.dataUrl);
|
|
const paths = this.refreshPaths();
|
|
if (!paths.configured) {
|
|
throw new DepthRuntimeError(
|
|
"depth_runtime_unavailable",
|
|
paths.reason === "worker_missing" ? "本机深度运行文件不存在" : "未找到本机深度 Python 环境",
|
|
);
|
|
}
|
|
const child = this.ensureChild();
|
|
const requestId = `depth-${Date.now().toString(36)}-${(++this.requestCounter).toString(36)}`;
|
|
const response = new Promise<DepthWorkerSuccess>((resolve, reject) => {
|
|
this.pending.set(requestId, { resolve, reject });
|
|
});
|
|
try {
|
|
child.stdin.write(JSON.stringify({ requestId, dataUrl: input.dataUrl }) + "\n");
|
|
} catch (error) {
|
|
this.pending.delete(requestId);
|
|
throw new DepthRuntimeError("depth_runtime_unavailable", "本机深度进程无法写入");
|
|
}
|
|
return response;
|
|
}
|
|
|
|
dispose() {
|
|
const child = this.child;
|
|
this.child = undefined;
|
|
this.loaded = false;
|
|
this.rejectPending(new DepthRuntimeError("depth_runtime_unavailable", "本机深度进程已停止"));
|
|
if (child && !child.killed) child.kill();
|
|
}
|
|
|
|
private ensureChild() {
|
|
if (this.child && !this.child.killed) return this.child;
|
|
let child: ChildProcessWithoutNullStreams;
|
|
try {
|
|
child = spawn(this.paths.pythonPath, [this.paths.scriptPath], {
|
|
cwd: path.dirname(this.paths.scriptPath),
|
|
env: {
|
|
...process.env,
|
|
HF_HOME: this.paths.hfHome,
|
|
HF_HUB_CACHE: path.join(this.paths.hfHome, "hub"),
|
|
HF_HUB_OFFLINE: "1",
|
|
TRANSFORMERS_OFFLINE: "1",
|
|
HF_HUB_DISABLE_TELEMETRY: "1",
|
|
HF_HUB_DISABLE_PROGRESS_BARS: "1",
|
|
TRANSFORMERS_VERBOSITY: "error",
|
|
CANVAS_DEPTH_MODEL_ID: this.paths.modelId,
|
|
},
|
|
stdio: ["pipe", "pipe", "pipe"],
|
|
windowsHide: true,
|
|
});
|
|
} catch (error) {
|
|
throw new DepthRuntimeError("depth_runtime_unavailable", "本机深度进程无法启动");
|
|
}
|
|
this.child = child;
|
|
child.stdout.setEncoding("utf8");
|
|
child.stderr.setEncoding("utf8");
|
|
child.stdout.on("data", (chunk: string) => this.consumeStdout(chunk));
|
|
child.stderr.on("data", (chunk: string) => {
|
|
this.stderrBuffer = `${this.stderrBuffer}${chunk}`.slice(-4_000);
|
|
});
|
|
child.on("error", () => this.handleChildExit(child));
|
|
child.on("exit", () => this.handleChildExit(child));
|
|
return child;
|
|
}
|
|
|
|
private refreshPaths() {
|
|
if (!this.child) this.paths = resolveDepthRuntimePaths();
|
|
return this.paths;
|
|
}
|
|
|
|
private consumeStdout(chunk: string) {
|
|
this.stdoutBuffer += chunk;
|
|
if (Buffer.byteLength(this.stdoutBuffer, "utf8") > MAX_PROTOCOL_LINE_BYTES) {
|
|
this.child?.kill();
|
|
this.handleChildExit(this.child);
|
|
return;
|
|
}
|
|
while (true) {
|
|
const newline = this.stdoutBuffer.indexOf("\n");
|
|
if (newline < 0) return;
|
|
const line = this.stdoutBuffer.slice(0, newline).trim();
|
|
this.stdoutBuffer = this.stdoutBuffer.slice(newline + 1);
|
|
if (!line) continue;
|
|
let value: unknown;
|
|
try { value = JSON.parse(line); } catch { continue; }
|
|
if (!isRecord(value) || typeof value.requestId !== "string") continue;
|
|
const pending = this.pending.get(value.requestId);
|
|
if (!pending) continue;
|
|
this.pending.delete(value.requestId);
|
|
if (value.ok === true) {
|
|
if (value.mimeType !== "image/png" || !isPositiveInteger(value.width) || !isPositiveInteger(value.height) || typeof value.modelId !== "string" || value.device !== "cpu" || typeof value.pngBase64 !== "string" || !value.pngBase64) {
|
|
pending.reject(new DepthRuntimeError("depth_response_invalid", "本机深度响应无效", 500));
|
|
continue;
|
|
}
|
|
this.loaded = true;
|
|
pending.resolve(value as unknown as DepthWorkerSuccess);
|
|
continue;
|
|
}
|
|
const code = typeof value.code === "string" ? value.code : "depth_inference_failed";
|
|
const statusCode = code === "depth_model_missing" ? 503 : code === "depth_input_invalid" ? 400 : 500;
|
|
pending.reject(new DepthRuntimeError(code, typeof value.message === "string" ? value.message : "本机深度推理失败", statusCode));
|
|
}
|
|
}
|
|
|
|
private handleChildExit(child?: ChildProcessWithoutNullStreams) {
|
|
if (!child || this.child !== child) return;
|
|
this.child = undefined;
|
|
this.loaded = false;
|
|
this.stderrBuffer = "";
|
|
this.stdoutBuffer = "";
|
|
this.rejectPending(new DepthRuntimeError("depth_runtime_unavailable", "本机深度进程已退出"));
|
|
}
|
|
|
|
private rejectPending(error: DepthRuntimeError) {
|
|
for (const { reject } of this.pending.values()) reject(error);
|
|
this.pending.clear();
|
|
}
|
|
}
|
|
|
|
export function validateDepthDataUrl(value: string) {
|
|
const match = DATA_URL_PATTERN.exec(value);
|
|
if (!match) throw new DepthRuntimeError("depth_input_invalid", "深度转换只支持 PNG、JPEG 或 WebP 图片", 400);
|
|
const bytes = Buffer.from(match[2]!, "base64");
|
|
if (!bytes.length || bytes.length > MAX_INPUT_BYTES) throw new DepthRuntimeError("depth_input_invalid", "深度转换图片不能超过 12MB", 400);
|
|
}
|
|
|
|
function parseDepthInput(body: Record<string, unknown>) {
|
|
const keys = Object.keys(body).sort();
|
|
if (keys.length !== 3 || keys[0] !== "dataUrl" || keys[1] !== "operation" || keys[2] !== "schemaVersion" || body.schemaVersion !== 1 || body.operation !== "depth" || typeof body.dataUrl !== "string") {
|
|
throw new DepthRuntimeError("depth_input_invalid", "深度转换请求字段无效", 400);
|
|
}
|
|
validateDepthDataUrl(body.dataUrl);
|
|
return body.dataUrl;
|
|
}
|
|
|
|
function parseBody(request: Request): Record<string, unknown> {
|
|
if (!Buffer.isBuffer(request.body) || request.body.length === 0) return {};
|
|
try {
|
|
const value = JSON.parse(request.body.toString("utf8")) as unknown;
|
|
return isRecord(value) ? value : {};
|
|
} catch {
|
|
return {};
|
|
}
|
|
}
|
|
|
|
function asyncRoute(action: (request: Request, response: Response) => Promise<void>): RequestHandler {
|
|
return (request, response, next) => {
|
|
void action(request, response).catch((error) => {
|
|
if (response.headersSent) return next(error);
|
|
if (error instanceof DepthRuntimeError) {
|
|
response.status(error.statusCode).json({ ok: false, code: error.code, message: error.message });
|
|
return;
|
|
}
|
|
response.status(503).json({ ok: false, code: "depth_runtime_unavailable", message: "本机深度运行时不可用" });
|
|
});
|
|
};
|
|
}
|
|
|
|
function isRecord(value: unknown): value is Record<string, unknown> {
|
|
return Boolean(value && typeof value === "object" && !Array.isArray(value));
|
|
}
|
|
|
|
function isPositiveInteger(value: unknown): value is number {
|
|
return typeof value === "number" && Number.isInteger(value) && value > 0;
|
|
}
|