初始化
This commit is contained in:
@@ -0,0 +1,469 @@
|
||||
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);
|
||||
},
|
||||
});
|
||||
Reference in New Issue
Block a user