Skip to content

Commit f6a2039

Browse files
authored
Extend core snapToGrid to LiteGraph Groups. (comfyanonymous#3393)
Extends the core Comfy.SnapToGrid behavior for nodes to apply to LiteGraph's LGraphGroup with the same behavior. Also, pulls out redundant rounding code into util function.
1 parent 16a493a commit f6a2039

File tree

1 file changed

+89
-7
lines changed

1 file changed

+89
-7
lines changed

web/extensions/core/snapToGrid.js

+89-7
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,13 @@ import { app } from "../../scripts/app.js";
22

33
// Shift + drag/resize to snap to grid
44

5+
/** Rounds a Vector2 in-place to the current CANVAS_GRID_SIZE. */
6+
function roundVectorToGrid(vec) {
7+
vec[0] = LiteGraph.CANVAS_GRID_SIZE * Math.round(vec[0] / LiteGraph.CANVAS_GRID_SIZE);
8+
vec[1] = LiteGraph.CANVAS_GRID_SIZE * Math.round(vec[1] / LiteGraph.CANVAS_GRID_SIZE);
9+
return vec;
10+
}
11+
512
app.registerExtension({
613
name: "Comfy.SnapToGrid",
714
init() {
@@ -43,10 +50,7 @@ app.registerExtension({
4350
const onResize = node.onResize;
4451
node.onResize = function () {
4552
if (app.shiftDown) {
46-
const w = LiteGraph.CANVAS_GRID_SIZE * Math.round(node.size[0] / LiteGraph.CANVAS_GRID_SIZE);
47-
const h = LiteGraph.CANVAS_GRID_SIZE * Math.round(node.size[1] / LiteGraph.CANVAS_GRID_SIZE);
48-
node.size[0] = w;
49-
node.size[1] = h;
53+
roundVectorToGrid(node.size);
5054
}
5155
return onResize?.apply(this, arguments);
5256
};
@@ -57,9 +61,7 @@ app.registerExtension({
5761
const origDrawNode = LGraphCanvas.prototype.drawNode;
5862
LGraphCanvas.prototype.drawNode = function (node, ctx) {
5963
if (app.shiftDown && this.node_dragged && node.id in this.selected_nodes) {
60-
const x = LiteGraph.CANVAS_GRID_SIZE * Math.round(node.pos[0] / LiteGraph.CANVAS_GRID_SIZE);
61-
const y = LiteGraph.CANVAS_GRID_SIZE * Math.round(node.pos[1] / LiteGraph.CANVAS_GRID_SIZE);
62-
64+
const [x, y] = roundVectorToGrid([...node.pos]);
6365
const shiftX = x - node.pos[0];
6466
let shiftY = y - node.pos[1];
6567

@@ -85,5 +87,85 @@ app.registerExtension({
8587

8688
return origDrawNode.apply(this, arguments);
8789
};
90+
91+
92+
93+
/**
94+
* The currently moving, selected group only. Set after the `selected_group` has actually started
95+
* moving.
96+
*/
97+
let selectedAndMovingGroup = null;
98+
99+
/**
100+
* Handles moving a group; tracking when a group has been moved (to show the ghost in `drawGroups`
101+
* below) as well as handle the last move call from LiteGraph's `processMouseUp`.
102+
*/
103+
const groupMove = LGraphGroup.prototype.move;
104+
LGraphGroup.prototype.move = function(deltax, deltay, ignore_nodes) {
105+
const v = groupMove.apply(this, arguments);
106+
// When we've started moving, set `selectedAndMovingGroup` as LiteGraph sets `selected_group`
107+
// too eagerly and we don't want to behave like we're moving until we get a delta.
108+
if (!selectedAndMovingGroup && app.canvas.selected_group === this && (deltax || deltay)) {
109+
selectedAndMovingGroup = this;
110+
}
111+
112+
// LiteGraph will call group.move both on mouse-move as well as mouse-up though we only want
113+
// to snap on a mouse-up which we can determine by checking if `app.canvas.last_mouse_dragging`
114+
// has been set to `false`. Essentially, this check here is the equivilant to calling an
115+
// `LGraphGroup.prototype.onNodeMoved` if it had existed.
116+
if (app.canvas.last_mouse_dragging === false && app.shiftDown) {
117+
// After moving a group (while app.shiftDown), snap all the child nodes and, finally,
118+
// align the group itself.
119+
this.recomputeInsideNodes();
120+
for (const node of this._nodes) {
121+
node.alignToGrid();
122+
}
123+
LGraphNode.prototype.alignToGrid.apply(this);
124+
}
125+
return v;
126+
};
127+
128+
/**
129+
* Handles drawing a group when, snapping the size when one is actively being resized tracking and/or
130+
* drawing a ghost box when one is actively being moved. This mimics the node snapping behavior for
131+
* both.
132+
*/
133+
const drawGroups = LGraphCanvas.prototype.drawGroups;
134+
LGraphCanvas.prototype.drawGroups = function (canvas, ctx) {
135+
if (this.selected_group && app.shiftDown) {
136+
if (this.selected_group_resizing) {
137+
roundVectorToGrid(this.selected_group.size);
138+
} else if (selectedAndMovingGroup) {
139+
const [x, y] = roundVectorToGrid([...selectedAndMovingGroup.pos]);
140+
const f = ctx.fillStyle;
141+
const s = ctx.strokeStyle;
142+
ctx.fillStyle = "rgba(100, 100, 100, 0.33)";
143+
ctx.strokeStyle = "rgba(100, 100, 100, 0.66)";
144+
ctx.rect(x, y, ...selectedAndMovingGroup.size);
145+
ctx.fill();
146+
ctx.stroke();
147+
ctx.fillStyle = f;
148+
ctx.strokeStyle = s;
149+
}
150+
} else if (!this.selected_group) {
151+
selectedAndMovingGroup = null;
152+
}
153+
return drawGroups.apply(this, arguments);
154+
};
155+
156+
157+
/** Handles adding a group in a snapping-enabled state. */
158+
const onGroupAdd = LGraphCanvas.onGroupAdd;
159+
LGraphCanvas.onGroupAdd = function() {
160+
const v = onGroupAdd.apply(app.canvas, arguments);
161+
if (app.shiftDown) {
162+
const lastGroup = app.graph._groups[app.graph._groups.length - 1];
163+
if (lastGroup) {
164+
roundVectorToGrid(lastGroup.pos);
165+
roundVectorToGrid(lastGroup.size);
166+
}
167+
}
168+
return v;
169+
};
88170
},
89171
});

0 commit comments

Comments
 (0)