From 578e38e58bd9dd998062c3859fa9145f4fed1c41 Mon Sep 17 00:00:00 2001
From: space-nuko <24979496+space-nuko@users.noreply.github.com>
Date: Fri, 5 May 2023 07:26:45 -0500
Subject: [PATCH] Integrate align to grid extension

---
 src/lib/ComfyGraphCanvas.ts     | 47 ++++++++++++++++++++++++++++++++-
 src/lib/components/ComfyApp.ts  |  2 ++
 src/lib/nodes/ComfyGraphNode.ts | 16 ++++++++++-
 src/lib/stores/uiState.ts       |  2 +-
 4 files changed, 64 insertions(+), 3 deletions(-)

diff --git a/src/lib/ComfyGraphCanvas.ts b/src/lib/ComfyGraphCanvas.ts
index 3a464bf..24279be 100644
--- a/src/lib/ComfyGraphCanvas.ts
+++ b/src/lib/ComfyGraphCanvas.ts
@@ -1,7 +1,8 @@
-import { BuiltInSlotShape, LGraph, LGraphCanvas, LGraphNode, LiteGraph, NodeMode, type MouseEventExt, type Vector2, type Vector4 } from "@litegraph-ts/core";
+import { BuiltInSlotShape, LGraph, LGraphCanvas, LGraphNode, LiteGraph, NodeMode, type MouseEventExt, type Vector2, type Vector4, TitleMode } from "@litegraph-ts/core";
 import type ComfyApp from "./components/ComfyApp";
 import queueState from "./stores/queueState";
 import { get } from "svelte/store";
+import uiState from "./stores/uiState";
 
 export type SerializedGraphCanvasState = {
     offset: Vector2,
@@ -101,7 +102,39 @@ export default class ComfyGraphCanvas extends LGraphCanvas {
         }
     }
 
+    private alignToGrid(node: LGraphNode, ctx: CanvasRenderingContext2D) {
+        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.titleMode
+            if (titleMode !== TitleMode.TRANSPARENT_TITLE && titleMode !== TitleMode.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;
+    }
+
     override drawNode(node: LGraphNode, ctx: CanvasRenderingContext2D): void {
+        if ((window as any)?.app?.shiftDown && this.node_dragged && node.id in this.selected_nodes) {
+            this.alignToGrid(node, ctx)
+        }
+
+        // Fade out inactive nodes
         var editor_alpha = this.editor_alpha;
         if (node.mode === NodeMode.NEVER) { // never
             this.editor_alpha = 0.4;
@@ -196,4 +229,16 @@ export default class ComfyGraphCanvas extends LGraphCanvas {
 
         return res;
     }
+
+    override onNodeMoved(node: LGraphNode) {
+        if (super.onNodeMoved)
+            super.onNodeMoved(node);
+
+        if ((window as any)?.app?.shiftDown) {
+            // Ensure all selected nodes are realigned
+            for (const id in this.selected_nodes) {
+                this.selected_nodes[id].alignToGrid();
+            }
+        }
+    }
 }
diff --git a/src/lib/components/ComfyApp.ts b/src/lib/components/ComfyApp.ts
index 132c74c..b406252 100644
--- a/src/lib/components/ComfyApp.ts
+++ b/src/lib/components/ComfyApp.ts
@@ -24,10 +24,12 @@ import { toast } from '@zerodevx/svelte-toast'
 import ComfyGraph from "$lib/ComfyGraph";
 import { ComfyBackendNode } from "$lib/nodes/ComfyBackendNode";
 import { get } from "svelte/store";
+import uiState from "$lib/stores/uiState";
 
 export const COMFYBOX_SERIAL_VERSION = 1;
 
 LiteGraph.catch_exceptions = false;
+LiteGraph.CANVAS_GRID_SIZE = 32;
 
 if (typeof window !== "undefined") {
     // Load default visibility
diff --git a/src/lib/nodes/ComfyGraphNode.ts b/src/lib/nodes/ComfyGraphNode.ts
index 17eb405..b9d1ca6 100644
--- a/src/lib/nodes/ComfyGraphNode.ts
+++ b/src/lib/nodes/ComfyGraphNode.ts
@@ -1,10 +1,12 @@
 import type { ComfyInputConfig } from "$lib/IComfyInputSlot";
 import type { SerializedPrompt } from "$lib/components/ComfyApp";
 import type ComfyWidget from "$lib/components/widgets/ComfyWidget";
-import { LGraph, LGraphNode, LiteGraph, type SerializedLGraphNode } from "@litegraph-ts/core";
+import { LGraph, LGraphNode, LiteGraph, type SerializedLGraphNode, type Vector2 } from "@litegraph-ts/core";
 import type { SvelteComponentDev } from "svelte/internal";
 import type { ComfyWidgetNode } from "./ComfyWidgetNodes";
 import type IComfyInputSlot from "$lib/IComfyInputSlot";
+import uiState from "$lib/stores/uiState";
+import { get } from "svelte/store";
 
 export type DefaultWidgetSpec = {
     defaultWidgetNode: new (name?: string) => ComfyWidgetNode,
@@ -27,6 +29,18 @@ export default class ComfyGraphNode extends LGraphNode {
     receiveOutput(output: any) {
     }
 
+    override onResize(size: Vector2) {
+        if ((window as any)?.app?.shiftDown) {
+            const w = LiteGraph.CANVAS_GRID_SIZE * Math.round(this.size[0] / LiteGraph.CANVAS_GRID_SIZE);
+            const h = LiteGraph.CANVAS_GRID_SIZE * Math.round(this.size[1] / LiteGraph.CANVAS_GRID_SIZE);
+            this.size[0] = w;
+            this.size[1] = h;
+        }
+
+        if (super.onResize)
+            super.onResize(size)
+    }
+
     override onSerialize(o: SerializedLGraphNode) {
         for (let index = 0; index < this.inputs.length; index++) {
             const input = this.inputs[index]
diff --git a/src/lib/stores/uiState.ts b/src/lib/stores/uiState.ts
index e9422be..456982a 100644
--- a/src/lib/stores/uiState.ts
+++ b/src/lib/stores/uiState.ts
@@ -19,7 +19,7 @@ const store: WritableUIStateStore = writable(
         graphLocked: false,
         nodesLocked: false,
         autoAddUI: true,
-        uiEditMode: "disabled",
+        uiEditMode: "disabled"
     })
 
 const uiStateStore: WritableUIStateStore =