From eb85896b731aa9d51c523eaf6e1cd6580da76d16 Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Thu, 26 Jan 2023 01:26:28 -0500
Subject: [PATCH] Add code for categories.

---
 main.py            | 3 +++
 webshit/index.html | 2 +-
 2 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/main.py b/main.py
index 20193f82..3e25114d 100644
--- a/main.py
+++ b/main.py
@@ -291,6 +291,9 @@ class PromptServer(BaseHTTPRequestHandler):
                 info['output'] = obj_class.RETURN_TYPES
                 info['name'] = x #TODO
                 info['description'] = ''
+                info['category'] = 'sd'
+                if hasattr(obj_class, 'CATEGORY'):
+                    info['category'] = obj_class.CATEGORY
                 out[x] = info
             self.wfile.write(json.dumps(out).encode('utf-8'))
         elif self.path[1:] in os.listdir(self.server.server_dir):
diff --git a/webshit/index.html b/webshit/index.html
index 809517a4..31c36708 100644
--- a/webshit/index.html
+++ b/webshit/index.html
@@ -237,7 +237,7 @@ function onObjectInfo(json) {
         MyNode.__json_data = json[key]
 
         LiteGraph.registerNodeType(key, MyNode);
-        MyNode.category = "sd"; //TODO: proper categories
+        MyNode.category = json[key]['category'];
     };
 
     afterLoadGraph();