From fc9d5d94ea7bf4387970499b266a3bc54842e440 Mon Sep 17 00:00:00 2001
From: FizzleDorf <1fizzledorf@gmail.com>
Date: Thu, 12 Oct 2023 01:47:58 -0400
Subject: [PATCH] fixed up init

---
 .gitignore  | 1 +
 __init__.py | 8 +++++---
 ait_load.py | 4 ++--
 3 files changed, 8 insertions(+), 5 deletions(-)
 create mode 100644 .gitignore

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..c18dd8d
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1 @@
+__pycache__/
diff --git a/__init__.py b/__init__.py
index adbf285..65a45df 100644
--- a/__init__.py
+++ b/__init__.py
@@ -1,5 +1,7 @@
-from .ait_load import AITLoader
+from .ait_load import AIT_Unet_Loader
 
-NODE_DISPLAY_NAME_MAPPINGS = {
-    "AITLoader":"AITLoader",
+NODE_CLASS_MAPPINGS = {
+    "AIT_Unet_Loader":AIT_Unet_Loader,
 }
+
+print('\033[34mAIT Nodes: \033[92mLoaded\033[0m')
\ No newline at end of file
diff --git a/ait_load.py b/ait_load.py
index 1cd6f17..82ba077 100644
--- a/ait_load.py
+++ b/ait_load.py
@@ -88,7 +88,7 @@ class AITPatch:
                 self.ait_model = None
                 print("unloaded AIT")
 
-class AITLoader:
+class AIT_Unet_Loader:
     @classmethod
     def INPUT_TYPES(s):
         return {"required": { "model": ("MODEL", ),
@@ -97,7 +97,7 @@ class AITLoader:
     RETURN_TYPES = ("MODEL",)
     FUNCTION = "load_ait"
 
-    CATEGORY = "_for_testing"
+    CATEGORY = "loaders/AIT"
 
     def load_ait(self, model, ait_name):
         ait_path = folder_paths.get_full_path("ait", ait_name)