From 27dfcd8cea332238d3e721ca4e0b55ea788913ef Mon Sep 17 00:00:00 2001
From: nkchocoai <nkai3359@gmail.com>
Date: Mon, 22 Jan 2024 16:04:23 +0900
Subject: [PATCH] Add PromptUtilitiesLoadPreset node

---
 __init__.py             |  4 +++-
 presets/animagineXL.csv |  3 +++
 {nodes => py}/node.py   | 26 ++++++++++++++++++++++++--
 py/preset.py            | 38 ++++++++++++++++++++++++++++++++++++++
 4 files changed, 68 insertions(+), 3 deletions(-)
 create mode 100644 presets/animagineXL.csv
 rename {nodes => py}/node.py (72%)
 create mode 100644 py/preset.py

diff --git a/__init__.py b/__init__.py
index 4730262..9f2f1c5 100644
--- a/__init__.py
+++ b/__init__.py
@@ -1,13 +1,15 @@
-from .nodes.node import *
+from .py.node import *
 
 NODE_CLASS_MAPPINGS = {
     "PromptUtilitiesFormatString": PromptUtilitiesFormatString,
     "PromptUtilitiesJoinStringList": PromptUtilitiesJoinStringList,
+    "PromptUtilitiesLoadPreset": PromptUtilitiesLoadPreset,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
     "PromptUtilitiesFormatString": "Format String",
     "PromptUtilitiesJoinStringList": "Join String List",
+    "PromptUtilitiesLoadPreset": "Load Preset",
 }
 
 WEB_DIRECTORY = "./js"
diff --git a/presets/animagineXL.csv b/presets/animagineXL.csv
new file mode 100644
index 0000000..24d7a7f
--- /dev/null
+++ b/presets/animagineXL.csv
@@ -0,0 +1,3 @@
+name,prompt
+"neg","lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry, artist name"
+"pos postfix","masterpiece, best quality"
\ No newline at end of file
diff --git a/nodes/node.py b/py/node.py
similarity index 72%
rename from nodes/node.py
rename to py/node.py
index 689df74..a3669a1 100644
--- a/nodes/node.py
+++ b/py/node.py
@@ -1,4 +1,6 @@
+import folder_paths
 
+from .preset import PresetManager
 
 class PromptUtilitiesFormatString:
     @classmethod
@@ -28,7 +30,7 @@ class PromptUtilitiesFormatString:
         for i in range(1, len(kwargs) + 1):
             result = result.replace(f'[{i}]',kwargs[f"arg{i}"]) 
         return (result,)
-    
+
 
 class PromptUtilitiesJoinStringList:
     @classmethod
@@ -50,7 +52,27 @@ class PromptUtilitiesJoinStringList:
     CATEGORY = "PromptUtilities"
 
     def join(self, separator, **kwargs):
+        PresetManager.load_presets()
         print(kwargs)
         print(kwargs.values())
         result = separator.join(kwargs.values())
-        return (result,)
\ No newline at end of file
+        return (result,)
+
+
+class PromptUtilitiesLoadPreset:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": { 
+                "preset": (list(PresetManager.get_presets().keys()), ),
+            }
+        }
+            
+    RETURN_TYPES = ("STRING",)
+    FUNCTION = "load_preset"
+
+    CATEGORY = "PromptUtilities"
+
+    def load_preset(self, preset):
+        prompt = PresetManager.get_prompt(preset)
+        return (prompt,)
\ No newline at end of file
diff --git a/py/preset.py b/py/preset.py
new file mode 100644
index 0000000..7a04cee
--- /dev/null
+++ b/py/preset.py
@@ -0,0 +1,38 @@
+import csv
+import os
+
+import folder_paths
+
+class PresetManager:
+    _presets = None
+
+    custom_nodes_dir = folder_paths.get_folder_paths("custom_nodes")[0]
+    presets_dir = os.path.join(custom_nodes_dir, "ComfyUI-PromptUtilities", "presets")
+
+    @classmethod
+    def get_presets(cls):
+        if cls._presets is None:
+            cls.load_presets()
+        return cls._presets
+
+    @classmethod
+    def get_prompt(cls, key):
+        presets = cls.get_presets()
+        print(key, presets, presets[key])
+        return presets[key]
+
+    @classmethod
+    def load_presets(cls):
+        cls._presets = dict()
+        preset_filename_list = cls.get_preset_filename_list()
+        for preset_filename in preset_filename_list:
+            with open(os.path.join(cls.presets_dir, preset_filename),"r") as f:
+                reader = csv.DictReader(f)
+                for row in reader:
+                    cls._presets[f"{preset_filename[:-4]} : {row['name']}"] = row['prompt']
+
+    @classmethod
+    def get_preset_filename_list(cls):
+        files, _ = folder_paths.recursive_search(cls.presets_dir, excluded_dir_names=[".git"])
+
+        return folder_paths.filter_files_extensions(files, ['.csv'])