import { app } from "../../../scripts/app.js";
// Based on NodeAligner by Tenney95 GitHub:https://github.com/Tenney95/ComfyUI-NodeAligner
const kayAlignBottomSvg = ``;
const kayAlignCenterHorizontallySvg = ``;
const kayAlignCenterVerticallySvg = ``;
const kayAlignLeftSvg = ``;
const kayAlignRightSvg = ``;
const kayAlignTopSvg = ``;
const kayEqualWidthSvg = ``;
const kayEqualHeightSvg = ``;
const kayHorizontalDistributionSvg = ``;
const kayVerticalDistributionSvg = ``;
let stylesInjected = false;
// 工具栏是 document.body 上的 fixed 浮窗，不插进 ComfyUI 的 DOM。
// 前端的顶栏由 Vue 托管，往里塞节点会在重渲染时被清掉；挂到官方遗留插槽
// (app.menu.element) 又会被关进「运行」按钮那张卡片里跟着它一起动。
// 所以干脆完全浮动：拖到哪就待在哪，不跟 ComfyUI 的布局有任何牵扯。
const KayNodeAlignmentManager = {
isInitialized: false,
toolbarContainer: null,
dragState: { isDragging: false, offsetX: 0, offsetY: 0 },
hasShownTooltip: false,
isVisible: true,
position: { leftPercentage: 50, topPercentage: 5 },
async init() {
if (this.isInitialized) return;
this.canvas = await this.waitForCanvas();
this.injectStyles();
this.loadPosition();
this.setupToolbar();
this.restorePosition();
this.bindCanvasEvents();
this.bindKeyboardShortcuts();
window.addEventListener('resize', () => this.updatePosition());
const displayMode = app.ui.settings.getSettingValue("KayTool.NodeAlignDisplayMode");
this.updateDisplayMode(displayMode);
// 放在最后：初始化中途失败时不要留下「已初始化」的假象，挡住下次重试。
this.isInitialized = true;
},
// 只等画布，并且带超时。等不到也照常初始化，最多是快捷键的点击刷新失效，
// 绝不像以前那样卡在 requestAnimationFrame 里空转到天荒地老。
waitForCanvas(timeout = 10000) {
return new Promise(resolve => {
const deadline = performance.now() + timeout;
const check = () => {
const canvas = app.canvasEl || document.querySelector('#graph-canvas') || app.canvas?.canvas || null;
if (canvas || performance.now() > deadline) resolve(canvas);
else requestAnimationFrame(check);
};
check();
});
},
injectStyles() {
if (stylesInjected) return;
const bgColor = app.ui.settings.getSettingValue("KayTool.NodeAlignBackgroundColor");
const iconColor = app.ui.settings.getSettingValue("KayTool.NodeAlignIconColor");
const dividerColor = app.ui.settings.getSettingValue("KayTool.NodeAlignDividerColor");
document.head.insertAdjacentHTML('beforeend', ``);
stylesInjected = true;
},
loadPosition() {
let saved = {};
try {
saved = JSON.parse(localStorage.getItem('KayNodeAlignToolbarPosition')) || {};
} catch (e) {
saved = {};
}
this.position = {
leftPercentage: saved.leftPercentage ?? 50,
topPercentage: saved.topPercentage ?? 5
};
},
savePosition() {
try {
localStorage.setItem('KayNodeAlignToolbarPosition', JSON.stringify(this.position));
} catch (e) { /* 隐私模式下 localStorage 可能不可写，忽略即可 */ }
},
setupToolbar() {
const bgColor = app.ui.settings.getSettingValue("KayTool.NodeAlignBackgroundColor");
const opacity = app.ui.settings.getSettingValue("KayTool.NodeAlignBackgroundOpacity") / 100;
const iconColor = app.ui.settings.getSettingValue("KayTool.NodeAlignIconColor");
this.toolbarContainer = document.createElement('div');
this.toolbarContainer.id = 'kay-node-alignment-toolbar';
this.toolbarContainer.classList.add('floating');
if (opacity > 0 && /^[0-9A-Fa-f]{6}$/.test(bgColor)) {
this.toolbarContainer.style.background = `rgba(${parseInt(bgColor.substr(0, 2), 16)}, ${parseInt(bgColor.substr(2, 2), 16)}, ${parseInt(bgColor.substr(4, 2), 16)}, ${opacity})`;
}
document.body.appendChild(this.toolbarContainer);
this.getAlignmentButtons().forEach(btn => {
const el = document.createElement(btn.type === 'divider' ? 'div' : 'button');
el.className = btn.type === 'divider' ? 'kay-toolbar-divider' : 'kay-align-button';
if (btn.type !== 'divider') {
el.id = btn.id;
el.innerHTML = btn.svg.replace(/fill="#666666"/g, `fill="#${iconColor}"`);
el.addEventListener('click', e => btn.action.call(this, e));
el.addEventListener('mouseover', () => {
const baseColor = app.ui.settings.getSettingValue("KayTool.NodeAlignIconBackgroundColor");
el.style.backgroundColor = this.adjustColor(baseColor, 65);
});
el.addEventListener('mouseout', () => {
const baseColor = app.ui.settings.getSettingValue("KayTool.NodeAlignIconBackgroundColor");
el.style.backgroundColor = `#${baseColor}`;
});
el.addEventListener('mousedown', () => el.style.transform = 'scale(0.95)');
el.addEventListener('mouseup', () => el.style.transform = '');
el.addEventListener('mouseleave', () => el.style.transform = '');
} else {
el.addEventListener('mousedown', e => this.onDragStart(e));
}
this.toolbarContainer.appendChild(el);
});
document.addEventListener('mousemove', (e) => this.onDragging(e));
document.addEventListener('mouseup', (e) => this.onDragEnd(e));
document.addEventListener('selectstart', e => this.dragState.isDragging && e.preventDefault());
this.addTooltip();
},
adjustColor(hex, amount) {
let color = hex.replace("#", "");
if (color.length === 3) color = color.split('').map(c => c + c).join('');
const r = parseInt(color.substr(0, 2), 16);
const g = parseInt(color.substr(2, 2), 16);
const b = parseInt(color.substr(4, 2), 16);
const brightness = Math.round(0.299 * r + 0.587 * g + 0.114 * b);
const adjust = brightness > 127 ? -amount : amount;
return `#${[r, g, b].map(c => Math.max(0, Math.min(255, c + adjust)).toString(16).padStart(2, '0')).join('')}`;
},
getAlignmentButtons() {
return [
// 开头这根就是工具栏自己的拖拽手柄，和「运行」按钮那张卡片互不相干。
{ type: 'divider' },
{ id: 'kay-align-left', svg: kayAlignLeftSvg, action: this.alignLeft },
{ id: 'kay-align-center-vertically', svg: kayAlignCenterVerticallySvg, action: this.alignCenterVertically },
{ id: 'kay-align-right', svg: kayAlignRightSvg, action: this.alignRight },
{ type: 'divider' },
{ id: 'kay-align-top', svg: kayAlignTopSvg, action: this.alignTop },
{ id: 'kay-align-center-horizontally', svg: kayAlignCenterHorizontallySvg, action: this.alignCenterHorizontally },
{ id: 'kay-align-bottom', svg: kayAlignBottomSvg, action: this.alignBottom },
{ type: 'divider' },
{ id: 'kay-equal-width', svg: kayEqualWidthSvg, action: this.equalWidth },
{ id: 'kay-equal-height', svg: kayEqualHeightSvg, action: this.equalHeight },
{ type: 'divider' },
{ id: 'kay-horizontal-distribution', svg: kayHorizontalDistributionSvg, action: this.horizontalDistribution },
{ id: 'kay-vertical-distribution', svg: kayVerticalDistributionSvg, action: this.verticalDistribution }
];
},
addTooltip() {
const tooltip = document.createElement('div');
tooltip.id = 'kay-align-tooltip';
tooltip.textContent = 'Node alignment toolbar';
this.toolbarContainer.appendChild(tooltip);
if (!this.hasShownTooltip) {
this.toolbarContainer.addEventListener('mouseenter', () => {
if (!this.hasShownTooltip) {
setTimeout(() => {
tooltip.style.display = 'block';
this.hasShownTooltip = true;
}, 1000);
}
});
this.toolbarContainer.addEventListener('mouseleave', () => tooltip.remove());
}
},
show() {
this.isVisible = true;
if (this.toolbarContainer) {
this.toolbarContainer.style.display = 'flex';
this.updatePosition();
}
},
hide() {
this.isVisible = false;
if (this.toolbarContainer) {
this.toolbarContainer.style.display = 'none';
}
},
updatePosition() {
if (!this.toolbarContainer || !this.isVisible) return;
const { windowRect, toolbarRect } = this.getRect();
let left = (this.position.leftPercentage / 100) * windowRect.width - toolbarRect.width / 2;
let top = (this.position.topPercentage / 100) * windowRect.height;
left = Math.max(0, Math.min(left, windowRect.width - toolbarRect.width));
top = Math.max(0, Math.min(top, windowRect.height - toolbarRect.height));
this.toolbarContainer.style.left = `${left}px`;
this.toolbarContainer.style.top = `${top}px`;
},
getRect() {
return {
windowRect: { width: window.innerWidth, height: window.innerHeight },
toolbarRect: this.toolbarContainer.getBoundingClientRect()
};
},
onDragStart(e) {
e.preventDefault();
const { toolbarRect } = this.getRect();
this.dragState = {
isDragging: true,
offsetX: e.clientX - toolbarRect.left,
offsetY: e.clientY - toolbarRect.top
};
},
onDragging(e) {
if (!this.dragState.isDragging) return;
const { windowRect, toolbarRect } = this.getRect();
let left = e.clientX - this.dragState.offsetX;
let top = e.clientY - this.dragState.offsetY;
left = Math.max(0, Math.min(left, windowRect.width - toolbarRect.width));
top = Math.max(0, Math.min(top, windowRect.height - toolbarRect.height));
this.toolbarContainer.style.left = `${left}px`;
this.toolbarContainer.style.top = `${top}px`;
},
onDragEnd() {
if (!this.dragState.isDragging) return;
this.dragState = { isDragging: false, offsetX: 0, offsetY: 0 };
const { windowRect, toolbarRect } = this.getRect();
this.position.leftPercentage = ((toolbarRect.left + toolbarRect.width / 2) / windowRect.width) * 100;
this.position.topPercentage = (toolbarRect.top / windowRect.height) * 100;
this.savePosition();
},
restorePosition() {
this.updatePosition();
},
getSelectedNodes() {
const nodes = Object.values(app.canvas.selected_nodes || {});
return nodes;
},
alignLeft() {
const nodes = this.getSelectedNodes();
const minX = Math.min(...nodes.map(n => n.pos[0]));
nodes.forEach(n => n.pos[0] = minX);
this.redraw();
},
alignRight() {
const nodes = this.getSelectedNodes();
const maxX = Math.max(...nodes.map(n => n.pos[0] + n.size[0]));
nodes.forEach(n => n.pos[0] = maxX - n.size[0]);
this.redraw();
},
alignTop() {
const nodes = this.getSelectedNodes();
const minY = Math.min(...nodes.map(n => n.pos[1]));
nodes.forEach(n => n.pos[1] = minY);
this.redraw();
},
alignBottom() {
const nodes = this.getSelectedNodes();
const maxY = Math.max(...nodes.map(n => n.pos[1] + n.size[1]));
nodes.forEach(n => n.pos[1] = maxY - n.size[1]);
this.redraw();
},
alignCenterHorizontally() {
const nodes = this.getSelectedNodes();
const minY = Math.min(...nodes.map(n => n.pos[1]));
const maxY = Math.max(...nodes.map(n => n.pos[1] + n.size[1]));
const centerY = (minY + maxY) / 2;
nodes.forEach(n => n.pos[1] = centerY - n.size[1] / 2);
this.redraw();
},
alignCenterVertically() {
const nodes = this.getSelectedNodes();
const minX = Math.min(...nodes.map(n => n.pos[0]));
const maxX = Math.max(...nodes.map(n => n.pos[0] + n.size[0]));
const centerX = (minX + maxX) / 2;
nodes.forEach(n => n.pos[0] = centerX - n.size[0] / 2);
this.redraw();
},
// 统一尺寸时取选中节点里的最大值。以前取的是「第一个被选中的节点」，
// 谁先被点到就以谁为准，看起来像是随机缩水；取最大值也不会把节点压到
// 它自己的最小尺寸以下。
resizeSelected(axis) {
const nodes = this.getSelectedNodes();
if (!nodes.length) return;
const target = Math.max(...nodes.map(n => n.size[axis]));
nodes.forEach(n => {
const size = [n.size[0], n.size[1]];
size[axis] = target;
if (typeof n.setSize === 'function') n.setSize(size);
else n.size = size;
});
this.redraw();
},
equalWidth() {
this.resizeSelected(0);
},
equalHeight() {
this.resizeSelected(1);
},
horizontalDistribution() {
const nodes = this.getSelectedNodes();
const axis = 0;
if (nodes.length > 1) {
nodes.sort((a, b) => a.pos[axis] - b.pos[axis]);
const min = Math.min(...nodes.map(node => node.pos[axis]));
const max = Math.max(...nodes.map(node => node.pos[axis] + node.size[axis]));
const totalSize = nodes.reduce((sum, node) => sum + node.size[axis], 0);
const spacing = (max - min - totalSize) / (nodes.length - 1);
let current = min;
nodes.forEach(node => {
node.pos[axis] = current;
current += node.size[axis] + spacing;
});
this.redraw();
}
},
verticalDistribution() {
const nodes = this.getSelectedNodes();
if (nodes.length > 1) {
const axis = 1;
const otherAxis = 0;
const tolerance = 100;
const minSpacing = 20;
const columns = [];
nodes.forEach(node => {
let foundColumn = null;
for (let column of columns) {
const columnX = column[0].pos[otherAxis];
if (Math.abs(columnX - node.pos[otherAxis]) <= tolerance) {
foundColumn = column;
break;
}
}
if (foundColumn) {
foundColumn.push(node);
} else {
columns.push([node]);
}
});
const columnHeights = columns.map(column => {
const minY = Math.min(...column.map(node => node.pos[axis]));
const maxY = Math.max(...column.map(node => node.pos[axis] + node.size[axis]));
return maxY - minY;
});
const maxColumnHeight = Math.max(...columnHeights);
columns.forEach(column => {
if (column.length > 1) {
column.sort((a, b) => a.pos[axis] - b.pos[axis]);
const minY = Math.min(...column.map(node => node.pos[axis]));
const totalSize = column.reduce((sum, node) => sum + node.size[axis], 0);
let spacing = (maxColumnHeight - totalSize) / (column.length - 1);
spacing = Math.max(spacing, minSpacing);
let currentY = column[0].pos[axis];
column.forEach((node, idx) => {
if (idx === 0) {
currentY += node.size[axis] + spacing;
} else {
node.pos[axis] = currentY;
currentY += node.size[axis] + spacing;
}
node.pos[otherAxis] = column[0].pos[otherAxis];
});
}
});
this.redraw();
}
},
redraw() {
if (app.canvas) app.canvas.setDirty(true, true);
},
updateDisplayMode(mode) {
const effectiveMode = mode || "permanent";
if (effectiveMode === "disabled") {
this.hide();
} else if (effectiveMode === "permanent") {
this.show();
} else if (effectiveMode === "on-select") {
const selectedNodes = this.getSelectedNodes();
if (selectedNodes.length >= 2) this.show();
else this.hide();
}
},
bindCanvasEvents() {
if (!this.canvas) return;
this.canvas.addEventListener('click', () => {
const currentMode = app.ui.settings.getSettingValue("KayTool.NodeAlignDisplayMode");
this.updateDisplayMode(currentMode);
});
},
bindKeyboardShortcuts() {
document.addEventListener('keydown', (e) => {
// 只在 Shift 键按下且没有其他修饰键（如 Ctrl、Alt）时检查 WASD
if (!e.shiftKey || e.ctrlKey || e.altKey || e.metaKey) return;
const enableShortcuts = app.ui.settings.getSettingValue("KayTool.EnableAlignmentShortcuts");
if (!enableShortcuts || this.getSelectedNodes().length < 2) return;
switch (e.key.toLowerCase()) {
case 'w':
e.preventDefault();
this.alignTop();
break;
case 'a':
e.preventDefault();
this.alignLeft();
break;
case 's':
e.preventDefault();
this.alignBottom();
break;
case 'd':
e.preventDefault();
this.alignRight();
break;
}
});
}
};
function initializeKayNodeAlignment() {
if (!KayNodeAlignmentManager.isInitialized) {
KayNodeAlignmentManager.init();
} else {
const displayMode = app.ui.settings.getSettingValue("KayTool.NodeAlignDisplayMode");
KayNodeAlignmentManager.updateDisplayMode(displayMode);
}
}
app.registerExtension({
name: "KayTool.NodeAlign",
init() {
initializeKayNodeAlignment();
window.KayNodeAlignmentManager = KayNodeAlignmentManager;
window.initializeKayNodeAlignment = initializeKayNodeAlignment;
}
});