From 6fbe9f9e29a0e5dac1917d363bcd2f3358b18f71 Mon Sep 17 00:00:00 2001
From: daxcay <daxtoncaylor@gmail.com>
Date: Sat, 6 Jul 2024 16:43:24 +0530
Subject: [PATCH] Update

---
 classes/DataSet_TextFilesSave.py |  2 +-
 web/util.js                      | 47 +++++++++++++++++++++++++++++---
 2 files changed, 44 insertions(+), 5 deletions(-)

diff --git a/classes/DataSet_TextFilesSave.py b/classes/DataSet_TextFilesSave.py
index 80f6e29..e83f183 100644
--- a/classes/DataSet_TextFilesSave.py
+++ b/classes/DataSet_TextFilesSave.py
@@ -41,7 +41,7 @@ class DataSet_TextFilesSave:
                 "TextFileNames": ("STRING",{"forceInput": True}),
                 "TextFileContents": ("STRING",{"forceInput": True}),
                 "destination": ("STRING", {"default": "directory path"}),
-                "save_mode": (['Merge','Overwrite','SaveNew','MergeAndSaveNew'],),
+                "save_mode": (['Overwrite','Merge','SaveNew','MergeAndSaveNew'],),
             },
         }
 
diff --git a/web/util.js b/web/util.js
index da092ee..ad4dc26 100644
--- a/web/util.js
+++ b/web/util.js
@@ -1,12 +1,51 @@
 import { app } from "../../scripts/app.js";
 
+let groups = {
+    "loaders": ['DataSet_TextFilesLoad', 'DataSet_TextFilesLoadFromList', 'DataSet_LoadImage', 'DataSet_PathSelector'],
+    "processors": ['DataSet_Visualizer', 'DataSet_FindAndReplace', 'DataSet_ConceptManager', 'DataSet_TriggerWords', 'DataSet_OpenAIChatImageBatch', 'DataSet_OpenAIChatImage', 'DataSet_OpenAIChat'],
+    "handlers": ['DataSet_TextFilesSave', 'DataSet_SaveImage', 'DataSet_CopyFiles'],
+}
+
+let groupsNodesColor = {
+    "loaders": "#823329",
+    "processors": "#1a472a",
+    "handlers": "#084c61"
+}
+
+function getColorByName(name) {
+    for (let group in groups) {
+        if (groups[group].includes(name)) {
+            return groupsNodesColor[group];
+        }
+    }
+    return null;
+}
+
+function isInGroup(name, group) {
+    if (groups[group]) {
+        return groups[group].includes(name);
+    }
+    return false;
+}
+
+function onNodeAdded(node, app) {
+    let color = getColorByName(node.title)
+    if (color) {
+        node.bgcolor = color
+        node.color = "#232c33"
+    }
+}
+
 let DataSet = {
     name: "ComfyUI-Dataset",
     async beforeRegisterNodeDef(nodeType, nodeData, app) {
-        if(nodeData.name.indexOf('DATASET') == 0) {
-            nodeData.input.required['seed'] = ["INT", {"default": 0, "min": 0, "max": 18446744073709552000}]
+        if (isInGroup(nodeData.name, "loaders")) {
+            nodeData.input.required['seed'] = ["INT", { "default": 0, "min": 0, "max": 18446744073709552000 }]
         }
-	},
+    },
+    loadedGraphNode(node, app) {
+        onNodeAdded(node, app)
+    },
 }
 
-// app.registerExtension(DataSet);
\ No newline at end of file
+app.registerExtension(DataSet);
\ No newline at end of file