470 lines
16 KiB
JavaScript
470 lines
16 KiB
JavaScript
import { app } from "../../scripts/app.js";
|
||
import { api } from "../../scripts/api.js";
|
||
import { compileFlowPrompt, NODE_TYPES, normalizeChannel } from "./compiler.mjs";
|
||
import { nextAvailablePairedPosition, uniquePublisherNode } from "./node_actions.mjs";
|
||
import { parsePipelineConfig } from "./pipeline_config.mjs";
|
||
import { PIPELINE_NODE_MIN_WIDTH, setupPipelineEditor } from "./pipeline_editor.js";
|
||
|
||
const FLOW_TYPES = new Set(Object.values(NODE_TYPES));
|
||
const LEGACY_PUBLISHERS = new Set([NODE_TYPES.legacyStage, NODE_TYPES.legacyRoute]);
|
||
const LEGACY_CONSUMERS = new Set([NODE_TYPES.legacyStage, NODE_TYPES.legacyRoute]);
|
||
const COLORS = { ok: "#56b889", warning: "#d9a441", error: "#df6666", idle: "#7b8492" };
|
||
const STATUS_PRIORITY = { idle: 0, ok: 1, warning: 2, error: 3 };
|
||
const INACTIVE_MODES = new Set([2, 4]);
|
||
const INTERNAL_INPUTS = ["source", "pipeline_result", "selected_value"];
|
||
const INTERNAL_WIDGETS = [
|
||
"compile_error",
|
||
"pipeline_config",
|
||
"stage_name",
|
||
"selected_name",
|
||
"__stage_internal",
|
||
"__flow_generated",
|
||
];
|
||
const LABELS = {
|
||
channel: "结果名称",
|
||
value: "数据",
|
||
fallback: "找不到时使用",
|
||
input_channel: "起点结果",
|
||
output_channel: "最终发布为",
|
||
enabled: "启用阶段",
|
||
processed: "处理结果",
|
||
route: "选择方案",
|
||
option_1: "方案 1 结果",
|
||
option_2: "方案 2 结果",
|
||
option_3: "方案 3 结果",
|
||
condition: "条件",
|
||
on_true: "为真时",
|
||
on_false: "为假时",
|
||
};
|
||
const OUTPUT_LABELS = {
|
||
[NODE_TYPES.publish]: "数据",
|
||
[NODE_TYPES.get]: "数据",
|
||
[NODE_TYPES.pipeline]: "流程结果",
|
||
[NODE_TYPES.legacyStage]: "阶段结果",
|
||
[NODE_TYPES.legacyRoute]: "选择结果",
|
||
[NODE_TYPES.legacyIf]: "选择结果",
|
||
};
|
||
let refreshPending = false;
|
||
|
||
function ensureStyles() {
|
||
const id = "flow-branch-styles";
|
||
if (document.getElementById(id)) return;
|
||
const link = document.createElement("link");
|
||
link.id = id;
|
||
link.rel = "stylesheet";
|
||
link.href = new URL("./flow_branch.css", import.meta.url).href;
|
||
document.head.append(link);
|
||
}
|
||
|
||
function nodeType(node) {
|
||
return node?.comfyClass || node?.type || "";
|
||
}
|
||
|
||
function widget(node, name) {
|
||
return node.widgets?.find((item) => item.name === name);
|
||
}
|
||
|
||
function input(node, name) {
|
||
return node.inputs?.find((item) => item.name === name);
|
||
}
|
||
|
||
function isConnected(slot) {
|
||
return slot?.link !== null && slot?.link !== undefined;
|
||
}
|
||
|
||
function channelFor(node, direction) {
|
||
const type = nodeType(node);
|
||
const name = direction === "out"
|
||
? (type === NODE_TYPES.publish ? "channel" : "output_channel")
|
||
: (type === NODE_TYPES.get ? "channel" : "input_channel");
|
||
return normalizeChannel(widget(node, name)?.value);
|
||
}
|
||
|
||
function removeInputByName(node, name) {
|
||
const index = node.inputs?.findIndex((item) => item.name === name) ?? -1;
|
||
if (index >= 0) node.removeInput(index);
|
||
}
|
||
|
||
function hideWidget(item) {
|
||
if (!item) return;
|
||
item.hidden = true;
|
||
item.computeSize = () => [0, -4];
|
||
item.options = { ...(item.options || {}), serialize: true };
|
||
}
|
||
|
||
function hideInternalControls(node) {
|
||
for (const name of INTERNAL_INPUTS) removeInputByName(node, name);
|
||
for (const name of INTERNAL_WIDGETS) hideWidget(widget(node, name));
|
||
}
|
||
|
||
function applyChineseLabels(node) {
|
||
for (const item of node.widgets || []) {
|
||
if (LABELS[item.name]) item.label = LABELS[item.name];
|
||
}
|
||
for (const item of node.inputs || []) {
|
||
if (LABELS[item.name]) item.label = LABELS[item.name];
|
||
}
|
||
const outputLabel = OUTPUT_LABELS[nodeType(node)];
|
||
if (outputLabel && node.outputs?.[0]) node.outputs[0].label = outputLabel;
|
||
}
|
||
|
||
function setStatus(node, level, text, detail = text) {
|
||
node.__flowBranchStatus = { level, text, detail };
|
||
node.boxcolor = COLORS[level] || COLORS.idle;
|
||
}
|
||
|
||
function setHigherPriorityStatus(node, level, text, detail = text) {
|
||
const current = node.__flowBranchStatus;
|
||
if (!current || STATUS_PRIORITY[level] >= STATUS_PRIORITY[current.level]) {
|
||
setStatus(node, level, text, detail);
|
||
}
|
||
}
|
||
|
||
function addPublisher(publishers, channel, node, key = channel) {
|
||
if (!channel) return;
|
||
const matches = publishers.get(channel) || [];
|
||
matches.push({ node, key });
|
||
publishers.set(channel, matches);
|
||
}
|
||
|
||
function publisherHasValue(node) {
|
||
return nodeType(node) !== NODE_TYPES.publish || isConnected(input(node, "value"));
|
||
}
|
||
|
||
function registerPublishers(nodes, requireValue = true) {
|
||
const publishers = new Map();
|
||
for (const node of nodes) {
|
||
const type = nodeType(node);
|
||
if (type === NODE_TYPES.publish) {
|
||
const channel = channelFor(node, "out");
|
||
if (!requireValue || publisherHasValue(node)) addPublisher(publishers, channel, node);
|
||
} else if (LEGACY_PUBLISHERS.has(type)) {
|
||
addPublisher(publishers, channelFor(node, "out"), node);
|
||
} else if (type === NODE_TYPES.pipeline) {
|
||
const config = parsePipelineConfig(widget(node, "pipeline_config")?.value);
|
||
for (const stage of config.stages) addPublisher(publishers, normalizeChannel(stage.name), node, stage.id);
|
||
addPublisher(publishers, channelFor(node, "out"), node, "final");
|
||
}
|
||
}
|
||
return publishers;
|
||
}
|
||
|
||
function registerLivePublishers(nodes) {
|
||
return registerPublishers(nodes, true);
|
||
}
|
||
|
||
function graphPublishers(graph) {
|
||
return registerPublishers(graph?._nodes || [], false);
|
||
}
|
||
|
||
function createReaderNode(sourceNode, channel) {
|
||
const normalized = normalizeChannel(channel);
|
||
const graph = sourceNode?.graph || app.graph;
|
||
const reader = globalThis.LiteGraph?.createNode?.(NODE_TYPES.get);
|
||
if (!normalized || !graph || !reader) return null;
|
||
|
||
reader.pos = nextAvailablePairedPosition(sourceNode, reader, graph._nodes || []);
|
||
graph.add(reader);
|
||
const channelWidget = widget(reader, "channel");
|
||
if (channelWidget) {
|
||
channelWidget.value = normalized;
|
||
channelWidget.callback?.(normalized);
|
||
}
|
||
if (app.canvas?.graph === graph) app.canvas.selectNode?.(reader, false);
|
||
graph.setDirtyCanvas?.(true, true);
|
||
scheduleDiagnostics();
|
||
return reader;
|
||
}
|
||
|
||
function publisherForReader(node) {
|
||
const channel = channelFor(node, "in");
|
||
if (!channel) return null;
|
||
return uniquePublisherNode(graphPublishers(node.graph || app.graph).get(channel));
|
||
}
|
||
|
||
function jumpToNode(target) {
|
||
if (!target || app.canvas?.graph !== target.graph) return;
|
||
app.canvas.centerOnNode?.(target);
|
||
app.canvas.selectNode?.(target, false);
|
||
app.canvas.setDirty?.(true, true);
|
||
}
|
||
|
||
function pipelineReaderMenu(node) {
|
||
const config = parsePipelineConfig(widget(node, "pipeline_config")?.value);
|
||
const options = [];
|
||
for (const [index, stage] of config.stages.entries()) {
|
||
const channel = normalizeChannel(stage.name);
|
||
if (!channel) continue;
|
||
options.push({
|
||
content: `阶段 ${index + 1}：${channel}`,
|
||
callback: () => createReaderNode(node, channel),
|
||
});
|
||
}
|
||
const finalChannel = channelFor(node, "out");
|
||
if (finalChannel) {
|
||
if (options.length) options.push(null);
|
||
options.push({
|
||
content: `最终结果：${finalChannel}`,
|
||
callback: () => createReaderNode(node, finalChannel),
|
||
});
|
||
}
|
||
return options;
|
||
}
|
||
|
||
function flowNodeMenuItems(node) {
|
||
const type = nodeType(node);
|
||
if (type === NODE_TYPES.publish) {
|
||
const channel = channelFor(node, "out");
|
||
return channel ? [{
|
||
content: `创建配对读取：${channel}`,
|
||
callback: () => createReaderNode(node, channel),
|
||
}] : [];
|
||
}
|
||
if (type === NODE_TYPES.pipeline) {
|
||
const options = pipelineReaderMenu(node);
|
||
return options.length ? [{
|
||
content: "创建结果读取节点",
|
||
has_submenu: true,
|
||
submenu: { options },
|
||
}] : [];
|
||
}
|
||
if (type === NODE_TYPES.get) {
|
||
const publisher = publisherForReader(node);
|
||
return publisher ? [{
|
||
content: "跳转到发送位置",
|
||
callback: () => jumpToNode(publisher),
|
||
}] : [];
|
||
}
|
||
return [];
|
||
}
|
||
|
||
function diagnosePublishers(nodes, publishers) {
|
||
for (const node of nodes) {
|
||
const type = nodeType(node);
|
||
if (type === NODE_TYPES.publish) {
|
||
const channel = channelFor(node, "out");
|
||
if (!channel) setStatus(node, "warning", "结果名称为空");
|
||
else if (!publisherHasValue(node)) setStatus(node, "warning", `未接数据 ${channel}`);
|
||
else {
|
||
const matches = publishers.get(channel) || [];
|
||
setStatus(
|
||
node,
|
||
matches.length > 1 ? "error" : "ok",
|
||
matches.length > 1 ? `名称冲突 ${channel}` : `发布 ${channel}`,
|
||
);
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
function diagnosePipeline(node, publishers) {
|
||
const config = parsePipelineConfig(widget(node, "pipeline_config")?.value);
|
||
const start = channelFor(node, "in");
|
||
const finalName = channelFor(node, "out");
|
||
const names = config.stages.map((stage) => normalizeChannel(stage.name));
|
||
const duplicate = names.find((name, index) => name && names.indexOf(name) !== index);
|
||
if (!start) setStatus(node, "error", "起点名称为空");
|
||
else if (!finalName) setStatus(node, "error", "最终名称为空");
|
||
else if (names.some((name) => !name)) setStatus(node, "error", "阶段名称为空");
|
||
else if (duplicate) setStatus(node, "error", `阶段重名 ${duplicate}`);
|
||
else {
|
||
const ownNames = new Set([...names, finalName]);
|
||
const conflict = [...ownNames].find((name) => (publishers.get(name) || []).length > 1);
|
||
const sources = (publishers.get(start) || []).filter((item) => item.node !== node);
|
||
if (conflict) setStatus(node, "error", `名称冲突 ${conflict}`);
|
||
else if (sources.length === 0) setStatus(node, "warning", `缺少起点 ${start}`);
|
||
else if (sources.length > 1) setStatus(node, "error", `起点冲突 ${start}`);
|
||
else if (config.stages.length === 0) setStatus(node, "warning", "尚未添加阶段");
|
||
else setStatus(node, "ok", `${config.stages.length} 个阶段`);
|
||
}
|
||
}
|
||
|
||
function diagnoseConsumer(node, publishers) {
|
||
const type = nodeType(node);
|
||
if (type !== NODE_TYPES.get && !LEGACY_CONSUMERS.has(type)) return;
|
||
const channel = channelFor(node, "in");
|
||
const matches = (publishers.get(channel) || []).filter((item) => item.node !== node);
|
||
if (!channel) setHigherPriorityStatus(node, "warning", "读取名称为空");
|
||
else if (matches.length === 1) setHigherPriorityStatus(node, "ok", `读取 ${channel}`);
|
||
else if (matches.length > 1) setHigherPriorityStatus(node, "error", `名称冲突 ${channel}`);
|
||
else setHigherPriorityStatus(node, "warning", `未找到 ${channel}`);
|
||
}
|
||
|
||
function refreshDiagnostics() {
|
||
refreshPending = false;
|
||
const allNodes = (app.graph?._nodes || []).filter((node) => FLOW_TYPES.has(nodeType(node)));
|
||
const nodes = allNodes.filter((node) => !INACTIVE_MODES.has(node.mode));
|
||
for (const node of allNodes) setStatus(node, "idle", "已停用");
|
||
for (const node of nodes) setStatus(node, "idle", "待检查");
|
||
const publishers = registerLivePublishers(nodes);
|
||
diagnosePublishers(nodes, publishers);
|
||
for (const node of nodes) {
|
||
if (nodeType(node) === NODE_TYPES.pipeline) diagnosePipeline(node, publishers);
|
||
diagnoseConsumer(node, publishers);
|
||
if (nodeType(node) === NODE_TYPES.legacyIf) setStatus(node, "idle", "旧版节点");
|
||
}
|
||
app.graph?.setDirtyCanvas(true, true);
|
||
}
|
||
|
||
function scheduleDiagnostics() {
|
||
if (refreshPending) return;
|
||
refreshPending = true;
|
||
requestAnimationFrame(refreshDiagnostics);
|
||
}
|
||
|
||
function wrapWidget(widgetItem, node) {
|
||
if (!widgetItem || widgetItem.__flowBranchWrapped) return;
|
||
widgetItem.__flowBranchWrapped = true;
|
||
const original = widgetItem.callback;
|
||
widgetItem.callback = function () {
|
||
const result = original?.apply(this, arguments);
|
||
if (nodeType(node) === NODE_TYPES.pipeline && widgetItem.name === "input_channel") {
|
||
node.__flowPipelineEditor?.render();
|
||
}
|
||
scheduleDiagnostics();
|
||
return result;
|
||
};
|
||
}
|
||
|
||
function fitStatusText(ctx, text, maxWidth) {
|
||
if (maxWidth <= 0) return "";
|
||
if (ctx.measureText(text).width <= maxWidth) return text;
|
||
let shortened = text;
|
||
while (shortened.length > 1 && ctx.measureText(`${shortened}...`).width > maxWidth) {
|
||
shortened = shortened.slice(0, -1);
|
||
}
|
||
return shortened.length ? `${shortened}...` : "";
|
||
}
|
||
|
||
function setupNode(node, reload = false) {
|
||
hideInternalControls(node);
|
||
applyChineseLabels(node);
|
||
const minWidth = nodeType(node) === NODE_TYPES.pipeline ? PIPELINE_NODE_MIN_WIDTH : 250;
|
||
node.setSize([Math.max(node.size[0], minWidth), node.size[1]]);
|
||
|
||
if (node.__flowBranchSetup) {
|
||
if (reload) node.__flowPipelineEditor?.reload();
|
||
scheduleDiagnostics();
|
||
return;
|
||
}
|
||
node.__flowBranchSetup = true;
|
||
|
||
for (const name of ["channel", "input_channel", "output_channel"]) {
|
||
wrapWidget(widget(node, name), node);
|
||
}
|
||
|
||
if (nodeType(node) === NODE_TYPES.pipeline) {
|
||
setupPipelineEditor(node, app, scheduleDiagnostics);
|
||
}
|
||
|
||
if (nodeType(node) === NODE_TYPES.get) {
|
||
const originalDoubleClick = node.onDblClick;
|
||
node.onDblClick = function () {
|
||
const publisher = publisherForReader(this);
|
||
if (publisher) {
|
||
jumpToNode(publisher);
|
||
return;
|
||
}
|
||
return originalDoubleClick?.apply(this, arguments);
|
||
};
|
||
}
|
||
|
||
const originalConnectionsChange = node.onConnectionsChange;
|
||
node.onConnectionsChange = function () {
|
||
const result = originalConnectionsChange?.apply(this, arguments);
|
||
this.__flowPipelineEditor?.refreshConnections();
|
||
scheduleDiagnostics();
|
||
return result;
|
||
};
|
||
|
||
const originalSerialize = node.onSerialize;
|
||
node.onSerialize = function () {
|
||
this.__flowPipelineEditor?.persist();
|
||
return originalSerialize?.apply(this, arguments);
|
||
};
|
||
|
||
const originalRemoved = node.onRemoved;
|
||
node.onRemoved = function () {
|
||
this.__flowPipelineEditor?.destroy();
|
||
const result = originalRemoved?.apply(this, arguments);
|
||
scheduleDiagnostics();
|
||
return result;
|
||
};
|
||
|
||
const originalModeChange = node.onModeChange;
|
||
node.onModeChange = function () {
|
||
const result = originalModeChange?.apply(this, arguments);
|
||
scheduleDiagnostics();
|
||
return result;
|
||
};
|
||
|
||
const originalDraw = node.onDrawForeground;
|
||
node.onDrawForeground = function (ctx) {
|
||
originalDraw?.apply(this, arguments);
|
||
this.__flowPipelineEditor?.updateCanvasPositions();
|
||
const status = this.__flowBranchStatus;
|
||
if (!status || this.flags?.collapsed) return;
|
||
ctx.save();
|
||
ctx.font = "11px sans-serif";
|
||
ctx.textAlign = "right";
|
||
ctx.textBaseline = "middle";
|
||
ctx.fillStyle = COLORS[status.level] || COLORS.idle;
|
||
const title = String(this.getTitle?.() || this.title || "");
|
||
const titleWidth = ctx.measureText(title).width;
|
||
const maxWidth = Math.max(0, this.size[0] - titleWidth - 48);
|
||
const text = fitStatusText(ctx, status.text, maxWidth);
|
||
if (text) ctx.fillText(text, this.size[0] - 10, -15);
|
||
ctx.restore();
|
||
};
|
||
scheduleDiagnostics();
|
||
}
|
||
|
||
function installPromptCompiler() {
|
||
if (api.__flowBranchCompilerInstalled) return;
|
||
api.__flowBranchCompilerInstalled = true;
|
||
const originalQueuePrompt = api.queuePrompt;
|
||
api.queuePrompt = async function () {
|
||
try {
|
||
const diagnostics = compileFlowPrompt(arguments[1]);
|
||
const warnings = diagnostics.filter((item) => item.level !== "ok");
|
||
if (warnings.length) console.warn("[FlowBranch]", warnings);
|
||
} catch (error) {
|
||
console.error("[FlowBranch] 编译流程失败", error);
|
||
}
|
||
return originalQueuePrompt.apply(this, arguments);
|
||
};
|
||
}
|
||
|
||
app.registerExtension({
|
||
name: "Comfy.FlowBranch",
|
||
init() {
|
||
ensureStyles();
|
||
},
|
||
setup() {
|
||
installPromptCompiler();
|
||
},
|
||
getNodeMenuItems(node) {
|
||
const items = flowNodeMenuItems(node);
|
||
return items.length ? [null, ...items] : [];
|
||
},
|
||
async beforeRegisterNodeDef(nodeType, nodeData) {
|
||
if (!FLOW_TYPES.has(nodeData.name)) return;
|
||
const originalCreated = nodeType.prototype.onNodeCreated;
|
||
nodeType.prototype.onNodeCreated = function () {
|
||
originalCreated?.apply(this, arguments);
|
||
setupNode(this);
|
||
};
|
||
if (nodeData.name === NODE_TYPES.pipeline) {
|
||
const originalConfigure = nodeType.prototype.configure;
|
||
nodeType.prototype.configure = function () {
|
||
const result = originalConfigure?.apply(this, arguments);
|
||
setupNode(this, true);
|
||
return result;
|
||
};
|
||
}
|
||
},
|
||
loadedGraphNode(node) {
|
||
if (FLOW_TYPES.has(nodeType(node))) setupNode(node, true);
|
||
},
|
||
});
|