From 3b062d0dfc4345bc175f7f1f072885679e9812a5 Mon Sep 17 00:00:00 2001
From: Haoming <hmstudy02@gmail.com>
Date: Wed, 19 Nov 2025 15:30:33 +0800
Subject: [PATCH] rewrite

---
 LICENSE             |  2 +-
 README.md           |  5 +----
 js/functions.js     | 39 ++++++++++++++++++++++++++-------------
 js/node_beautify.js | 35 +++++++++++++++++++----------------
 pyproject.toml      |  2 +-
 5 files changed, 48 insertions(+), 35 deletions(-)

diff --git a/LICENSE b/LICENSE
index 91b74be..f727380 100644
--- a/LICENSE
+++ b/LICENSE
@@ -1,6 +1,6 @@
 MIT License
 
-Copyright (c) 2024 Haoming
+Copyright (c) 2025 Haoming
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal
diff --git a/README.md b/README.md
index f29f912..95e9209 100644
--- a/README.md
+++ b/README.md
@@ -2,7 +2,4 @@
 This is an Extension for [ComfyUI](https://github.com/comfyanonymous/ComfyUI), which helps formatting the workflow.
 
 ## Feature
-This will add a button, **Beautify**, to the menu that rounds the **size** and **position** of each **node** and **group** in a workflow using a specified value
-
-## Config
-You can edit the **Beautify Grid Size** field in the settings to change the value used for rounding
+This will add a button, **Beautify**, to the menu that rounds the **size** and **position** of each **node** and **group** in a workflow using the set value
diff --git a/js/functions.js b/js/functions.js
index e851fa9..80411e4 100644
--- a/js/functions.js
+++ b/js/functions.js
@@ -1,23 +1,36 @@
 import { app } from "../../scripts/app.js";
 
-function round(value, grid_size) {
-    return value - value % grid_size;
+/** @param {number} value @param {number} size @returns {number} */
+function round(value, size) {
+    return Math.round(value / size) * size;
 }
 
-export function Process(graph) {
-    const GRID_SIZE = app.ui.settings.getSettingValue("beautify.GridSize", 10);
+export function process_graph() {
+    const GRID_SIZE = app.extensionManager.setting.get("Beautify.GridSize");
 
-    graph["nodes"].forEach((node) => {
-        ["pos", "size"].forEach((prop) => {
-            node[prop][0] = round(node[prop][0], GRID_SIZE);
-            node[prop][1] = round(node[prop][1], GRID_SIZE);
-        });
+    app.graph._nodes?.forEach((node) => {
+        const [x, y] = node.pos;
+        const [w, h] = node.size;
+
+        node.pos[0] = round(x, GRID_SIZE);
+        node.pos[1] = round(y, GRID_SIZE);
+
+        node.size[0] = round(w, GRID_SIZE);
+        node.size[1] = round(h, GRID_SIZE);
     });
 
-    graph["groups"].forEach((group) => {
-        for (let i = 0; i < 4; i++)
-            group["bounding"][i] = round(group["bounding"][i], GRID_SIZE);
+    app.graph._groups?.forEach((group) => {
+        const [x, y] = group._pos;
+        const [w, h] = group._size;
+
+        group.pos[0] = round(x, GRID_SIZE);
+        group.pos[1] = round(y, GRID_SIZE);
+
+        group.size[0] = round(w, GRID_SIZE);
+        group.size[1] = round(h, GRID_SIZE);
+
+        group.recomputeInsideNodes();
     });
 
-    return graph;
+    app.graph.change();
 }
diff --git a/js/node_beautify.js b/js/node_beautify.js
index 16fad31..8382920 100644
--- a/js/node_beautify.js
+++ b/js/node_beautify.js
@@ -1,12 +1,5 @@
 import { app } from "../../scripts/app.js";
-import { Process } from "./functions.js";
-
-app.ui.settings.addSetting({
-	id: "beautify.GridSize",
-	name: "Beautify Grid Size",
-	type: "number",
-	defaultValue: 10,
-});
+import { process_graph } from "./functions.js";
 
 function legacy() {
 	const menu = document.querySelector(".comfy-menu");
@@ -14,38 +7,48 @@ function legacy() {
 	const beautifyButton = document.createElement("button");
 	beautifyButton.textContent = "Beautify";
 	beautifyButton.addEventListener("click", () => {
-		app.loadGraphData(Process(app.graph.serialize()));
+		process_graph();
 	});
 
 	const refreshButton = document.getElementById("comfy-refresh-button");
 	menu.insertBefore(beautifyButton, refreshButton);
+	document.addEventListener("keydown", (e) => {
+		if (e.altKey && e.shiftKey && e.code === "KeyF")
+			beautifyButton.click();
+	});
 }
 
 async function frontend() {
 	const btn = new (await import("../../scripts/ui/components/button.js")).ComfyButton({
 		icon: "creation",
 		action: () => {
-			app.loadGraphData(Process(app.graph.serialize()));
+			process_graph();
 		},
 		tooltip: "Beautify Workflow",
 		content: "Beautify",
-		classList: "comfyui-button comfyui-menu-mobile-collapse"
+		classList: "comfyui-button comfyui-menu-mobile-collapse",
 	}).element;
 
-	app.menu?.actionsGroup.element.after(btn);
+	app.menu.actionsGroup.element.after(btn);
 }
 
 app.registerExtension({
 	name: "Comfy.NodeBeautify",
 	async setup() {
-
 		try {
 			await frontend();
 		} catch {
-			// No Frontend?
+			// No Frontend
 		}
 
 		legacy();
-
-	}
+	},
+	settings: [
+		{
+			id: "Beautify.GridSize",
+			name: "Grid Size to Round To",
+			type: "number",
+			defaultValue: 10,
+		},
+	]
 });
diff --git a/pyproject.toml b/pyproject.toml
index 5f8d246..71b7d5e 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,7 +1,7 @@
 [project]
 name = "comfyui-node-beautify"
 description = "Add a button that formats the workflow graph"
-version = "1.1.0"
+version = "1.2.0"
 license = { text = "MIT License" }
 dependencies = []