From d027ff121c904f5d21b5d9fd8607fcdb2b166ec3 Mon Sep 17 00:00:00 2001
From: pythongosssss <125205205+pythongosssss@users.noreply.github.com>
Date: Sun, 2 Apr 2023 15:33:34 +0100
Subject: [PATCH] Snap to grid

---
 web/extensions/core/snapToGrid.js | 86 +++++++++++++++++++++++++++++++
 web/scripts/app.js                |  8 ++-
 2 files changed, 93 insertions(+), 1 deletion(-)
 create mode 100644 web/extensions/core/snapToGrid.js

diff --git a/web/extensions/core/snapToGrid.js b/web/extensions/core/snapToGrid.js
new file mode 100644
index 00000000..80e836a0
--- /dev/null
+++ b/web/extensions/core/snapToGrid.js
@@ -0,0 +1,86 @@
+import { app } from "/scripts/app.js";
+
+// Shift + drag/resize to snap to grid 
+
+app.registerExtension({
+	name: "Comfy.SnapToGrid",
+	init() {
+		// Add setting to control grid size
+		app.ui.settings.addSetting({
+			id: "Comfy.SnapToGrid.GridSize",
+			name: "Grid Size",
+			type: "number",
+			attrs: {
+				min: 1,
+				max: 500,
+			},
+			tooltip:
+				"When dragging and resizing nodes while holding shift they will be aligned to the grid, this controls the size of that grid.",
+			defaultValue: LiteGraph.CANVAS_GRID_SIZE,
+			onChange(value) {
+				LiteGraph.CANVAS_GRID_SIZE = +value;
+			},
+		});
+
+		// After moving a node, if the shift key is down align it to grid
+		const onNodeMoved = app.canvas.onNodeMoved;
+		app.canvas.onNodeMoved = function (node) {
+			const r = onNodeMoved?.apply(this, arguments);
+
+			if (app.shiftDown) {
+				node.alignToGrid();
+			}
+
+			return r;
+		};
+
+      // When a node is added, add a resize handler to it so we can fix align the size with the grid
+		const onNodeAdded = app.graph.onNodeAdded;
+		app.graph.onNodeAdded = function (node) {
+			const onResize = node.onResize;
+			node.onResize = function () {
+            if(app.shiftDown) {
+               const w = LiteGraph.CANVAS_GRID_SIZE * Math.round(node.size[0] / LiteGraph.CANVAS_GRID_SIZE);
+               const h = LiteGraph.CANVAS_GRID_SIZE * Math.round(node.size[1] / LiteGraph.CANVAS_GRID_SIZE);
+               node.size[0] = w;
+               node.size[1] = h;
+            }
+				return onResize?.apply(this, arguments);
+			};
+			return onNodeAdded?.apply(this, arguments);
+		};
+
+		// Draw a preview of where the node will go if holding shift
+		const origDrawNode = LGraphCanvas.prototype.drawNode;
+		LGraphCanvas.prototype.drawNode = function (node, ctx) {
+			if (app.shiftDown && node === this.node_dragged) {
+				const x = LiteGraph.CANVAS_GRID_SIZE * Math.round(node.pos[0] / LiteGraph.CANVAS_GRID_SIZE);
+				const y = LiteGraph.CANVAS_GRID_SIZE * Math.round(node.pos[1] / LiteGraph.CANVAS_GRID_SIZE);
+
+				const shiftX = x - node.pos[0];
+				let shiftY = y - node.pos[1];
+
+				let w, h;
+				if (node.flags.collapsed) {
+					w = node._collapsed_width;
+					h = LiteGraph.NODE_TITLE_HEIGHT;
+					shiftY -= LiteGraph.NODE_TITLE_HEIGHT;
+				} else {
+					w = node.size[0];
+					h = node.size[1];
+					let titleMode = node.constructor.title_mode;
+					if (titleMode !== LiteGraph.TRANSPARENT_TITLE && titleMode !== LiteGraph.NO_TITLE) {
+						h += LiteGraph.NODE_TITLE_HEIGHT;
+						shiftY -= LiteGraph.NODE_TITLE_HEIGHT;
+					}
+				}
+				const f = ctx.fillStyle;
+				ctx.fillStyle = "rgba(100, 100, 100, 0.5)";
+				ctx.fillRect(shiftX, shiftY, w, h);
+				ctx.fillStyle = f;
+			}
+
+			return origDrawNode.apply(this, arguments);
+		};
+	},
+});
diff --git a/web/scripts/app.js b/web/scripts/app.js
index 5af6d5fc..6f8ac067 100644
--- a/web/scripts/app.js
+++ b/web/scripts/app.js
@@ -18,6 +18,7 @@ class ComfyApp {
 		this.ui = new ComfyUI(this);
 		this.extensions = [];
 		this.nodeOutputs = {};
+		this.shiftDown = false;
 	}
 
 	/**
@@ -538,7 +539,7 @@ class ComfyApp {
 				color = "#0f0";
 			} else if (self.dragOverNode && node.id === self.dragOverNode.id) {
 				color = "dodgerblue";
-			}
+			} 
 
 			if (color) {
 				const shape = node._shape || node.constructor.shape || LiteGraph.ROUND_SHAPE;
@@ -637,11 +638,16 @@ class ComfyApp {
 
 	#addKeyboardHandler() {
 		window.addEventListener("keydown", (e) => {
+			this.shiftDown = e.shiftKey;
+
 			// Queue prompt using ctrl or command + enter
 			if ((e.ctrlKey || e.metaKey) && (e.key === "Enter" || e.keyCode === 13 || e.keyCode === 10)) {
 				this.queuePrompt(e.shiftKey ? -1 : 0);
 			}
 		});
+		window.addEventListener("keyup", (e) => {
+			this.shiftDown = e.shiftKey;
+		});
 	}
 
 	/**