From f25752b59d83441c29f76db3e5c2d9705c107021 Mon Sep 17 00:00:00 2001
From: Ivor Wanders <ivor@iwanders.net>
Date: Fri, 21 Jun 2024 20:10:22 -0400
Subject: [PATCH] Node to obtain json of the tokenizer vocabulary.

---
 __init__.py        |  7 +++++--
 save_string.py     |  5 +++++
 tokenizer_vocab.py | 35 +++++++++++++++++++++++++++++++++++
 3 files changed, 45 insertions(+), 2 deletions(-)
 create mode 100644 tokenizer_vocab.py

diff --git a/__init__.py b/__init__.py
index a57eacb..0afe487 100644
--- a/__init__.py
+++ b/__init__.py
@@ -1,11 +1,14 @@
 from .save_string import SaveString
+from .tokenizer_vocab import TokenizerVocab
 
 NODE_CLASS_MAPPINGS = {
-    "SaveString": SaveString
+    "SaveString": SaveString,
+    "TokenizerVocab": TokenizerVocab,
 }
 
 # A dictionary that contains the friendly/humanly readable titles for the nodes
 NODE_DISPLAY_NAME_MAPPINGS = {
-    "SaveString": "SaveString"
+    "SaveString": "IW SaveString",
+    "TokenizerVocab": "IW TokenizerVocab"
 }
 
diff --git a/save_string.py b/save_string.py
index 0972df7..ce1f75c 100644
--- a/save_string.py
+++ b/save_string.py
@@ -1,5 +1,6 @@
 import folder_paths
 import os
+import json
 
 class SaveString:
     def __init__(self):
@@ -30,6 +31,10 @@ class SaveString:
         full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir)
         file = f"{filename}_{counter:05}_{filename_suffix}.{extension}"
         path = os.path.join(full_output_folder, file)
+
+        if extension == "json":
+            string_field = json.dumps(json.loads(string_field), indent=4)
+
         with open(path, "w") as f:
             f.write(string_field)
         return tuple()
diff --git a/tokenizer_vocab.py b/tokenizer_vocab.py
new file mode 100644
index 0000000..e1a2627
--- /dev/null
+++ b/tokenizer_vocab.py
@@ -0,0 +1,35 @@
+import folder_paths
+from comfy import sdxl_clip
+import os
+import json
+
+class TokenizerVocab:
+    def __init__(self):
+        pass
+    
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "clip": ("CLIP",),
+            },
+            "optional": {
+            },
+        }
+
+    OUTPUT_NODE = False
+    RETURN_TYPES = ("STRING", )
+
+    FUNCTION = "execute"
+
+    CATEGORY = "iw"
+
+    def execute(self, clip):
+        tokenizer = clip.tokenizer
+        vocab = {}
+        # Is clip l and g different?
+        if isinstance(tokenizer, sdxl_clip.SDXLTokenizer):
+            clip_l_tokenizer = tokenizer.clip_l.tokenizer
+            vocab = clip_l_tokenizer.get_vocab()
+
+        return (json.dumps(vocab),)