Merge pull request #1 from zer0thgear/feature/add-tavern-utils

Feature/add tavern utils

.gitignore

@@ -1 +1 @@
-/__pycache__
+**/__pycache__

__init__.py

@@ -7,10 +7,14 @@
 
 from .multiline_string import MultilineStringNode
 from .list_combine_node import ListCombineNode
+from .tavern_utils.read_tavern_card import ReadTavernCardNode
+from .tavern_utils.write_tavern_card import WriteTavernCardNode
 
 NODE_CLASS_MAPPINGS = {
     "List Combine Node (zer0)": ListCombineNode,
     "Multiline String Node (zer0)": MultilineStringNode,
+    "Tavern Card Info Node (zer0)": ReadTavernCardNode,
+    "Tavern Card Creation Node (zer0)": WriteTavernCardNode,
 }
 
 __all__ = ["NODE_CLASS_MAPPINGS"]

tavern_utils/read_tavern_card.py

@@ -0,0 +1,35 @@
+import base64
+import json
+import os
+
+from PIL import Image
+
+import folder_paths
+import node_helpers
+
+from .tavern_card_prototype import returnBasicOutputTypes, returnReadableBasicTupleNames, TavernCard
+
+class ReadTavernCardNode:
+    @classmethod
+    def INPUT_TYPES(self):
+        input_dir = folder_paths.get_input_directory()
+        files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f)) and f.endswith(".png")]
+        return {
+            "required": {
+                "image": (sorted(files), {"image_upload": True}),
+            },
+        }
+    RETURN_TYPES = returnBasicOutputTypes()
+    RETURN_NAMES = returnReadableBasicTupleNames()
+    OUTPUT_TOOLTIPS = ("Card info in the form of strings.",)
+    INPUT_IS_LIST = True
+    FUNCTION = "return_card_info"
+    CATEGORY = "utility"
+    DESCRIPTION = "Returns tavern card metadata."
+    def return_card_info(self, image):
+        image_path = folder_paths.get_annotated_filepath(image[0])
+        img = node_helpers.pillow(Image.open, image_path)
+        card_info = json.loads(base64.b64decode(img.text["chara"]))["data"]
+        parsed_card = TavernCard(card_info)
+        #prompt_json = json.loads(img.text["prompt"])
+        return parsed_card.return_basic_tuple() #(list(prompt_json.values()),)

tavern_utils/tavern_card_prototype.py

@@ -0,0 +1,66 @@
+class TavernCard:
+    def __init__(self, card_info=None):
+        self.name = card_info["name"] if card_info and "name" in card_info else ""
+        self.description = card_info["description"] if card_info and "description" in card_info else ""
+        self.personality = card_info["personality"] if card_info and "personality" in card_info else ""
+        self.scenario = card_info["scenario"] if card_info and "scenario" in card_info else ""
+        self.first_mes = card_info["first_mes"] if card_info and "first_mes" in card_info else ""
+        self.mes_example = card_info["mes_example"] if card_info and "mes_example" in card_info else ""
+        self.creator_notes = card_info["creator_notes"] if card_info and "creator_notes" in card_info else ""
+        self.system_prompt = card_info["system_prompt"] if card_info and "system_prompt" in card_info else ""
+        self.post_history_instructions = card_info["post_history_instructions"] if card_info and "post_history_instructions" in card_info else ""
+        self.alternate_greetings = card_info["alternate_greetings"] if card_info and "alternate_greetings" in card_info else []
+        self.character_book = card_info["character_book"] if card_info and "character_book" in card_info else {"name": "", "entries": []}
+        self.tags = card_info["tags"] if card_info and "tags" in card_info else ""
+        self.creator = card_info["creator"] if card_info and "creator" in card_info else ""
+        self.character_version = card_info["character_version"] if card_info and "character_version" in card_info else ""
+        self.extensions = card_info["extensions"] if card_info and "extensions" in card_info else ""
+
+    def __repr__(self):
+        return f"""name: {self.name},
+description: {self.description},
+personality: {self.personality},
+scenario: {self.scenario},
+first_mes: {self.first_mes},
+mes_example: {self.mes_example},
+creator_notes: {self.creator_notes},
+system_prompt: {self.system_prompt},
+post_history_instructions: {self.post_history_instructions},
+alternate_greetings: {self.alternate_greetings},
+character_book: {self.character_book},
+tags: {self.tags},
+creator: {self.creator},
+character_version: {self.character_version},
+extensions: {self.extensions}
+"""
+
+    def return_dict(self):
+        return {
+            "name": self.name,
+            "description": self.description,
+            "personality": self.personality,
+            "scenario": self.scenario,
+            "first_mes": self.first_mes,
+            "mes_example": self.mes_example,
+            "creator_notes": self.creator_notes,
+            "system_prompt": self.system_prompt,
+            "post_history_instructions": self.post_history_instructions,
+            "alternate_greetings": self.alternate_greetings,
+            "character_book": self.character_book,
+            "tags": self.tags,
+            "creator": self.creator,
+            "character_version": self.character_version,
+            "extensions": self.extensions
+        }
+
+    def return_basic_tuple(self):
+        return (self.name, self.description, self.personality, self.scenario, self.first_mes, self.mes_example, self.creator_notes, self.creator, self.character_version)
+
+    def return_full_tuple(self):
+        return (self.name, self.description, self.personality, self.scenario, self.first_mes, self.mes_example, self.creator_notes, self.system_prompt, self.post_history_instructions, self.alternate_greetings, self.character_book, self.tags, self.creator, self.character_version, self.extensions)
+    
+def returnBasicOutputTypes():
+    return ("STRING", "STRING", "STRING", "STRING", "STRING", "STRING", "STRING", "STRING", "STRING")
+
+def returnReadableBasicTupleNames():
+    return ("Card Name", "Card Description", "Personality", "Scenario", "First Message", "Message Examples", "Creator Notes", "Author", "Version")

tavern_utils/write_tavern_card.py

@@ -0,0 +1,63 @@
+import base64
+import json
+import os
+
+import numpy as np
+from PIL import Image
+from PIL.PngImagePlugin import PngInfo
+
+import folder_paths
+
+from .tavern_card_prototype import TavernCard
+
+class WriteTavernCardNode:
+    def __init__(self):
+        self.output_dir = folder_paths.get_output_directory()
+        self.type = "output"
+        self.compress_level = 4
+
+    @classmethod
+    def INPUT_TYPES(self):
+        return {
+            "required": {
+                "image": ("IMAGE", {"tooltip": "The image to be displayed on the tavern card."}),
+                "filename": ("STRING", {"default": "tavern_card", "tooltip": "The filename to save the tavern card as."}),
+                "card_name": ("STRING", {"default": "",}),
+                "card_description": ("STRING", {"default": "", "multiline": True}),
+                "personality": ("STRING", {"default": "",}),
+                "scenario": ("STRING", {"default": "",}),
+                "first_message": ("STRING", {"default": "", "multiline": True}),
+                "message_examples": ("STRING", {"default": "", "multiline": True}),
+                "creator_notes": ("STRING", {"default": "", "multiline": True}),
+                "author": ("STRING", {"default": "",}),
+                "version": ("STRING", {"default": "",}),
+            }
+        }
+    RETURN_TYPES = ()
+    FUNCTION = "create_tavern_card"
+    CATEGORY = "utility"
+    OUTPUT_NODE = True
+    DESCRIPTION = "Saves a created tavern card as a PNG."
+    def create_tavern_card(self, image, filename, card_name, card_description, personality, scenario, first_message, message_examples, creator_notes, author, version):
+        full_output_folder, filename, _, subfolder, _ = folder_paths.get_save_image_path(filename, self.output_dir, image[0].shape[1], image[0].shape[0])
+        i = 255. * image[0].cpu().numpy()
+        img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
+
+        card_info = {
+            "name": card_name,
+            "description": card_description,
+            "personality": personality,
+            "scenario": scenario,
+            "first_mes": first_message,
+            "mes_example": message_examples,
+            "creator_notes": creator_notes,
+            "creator": author,
+            "character_version": version,
+        }
+
+        card = TavernCard(card_info)
+        metadata = PngInfo()
+        metadata.add_text("chara", base64.b64encode(json.dumps({"data": card.return_dict()}).encode("ascii")))
+        file = f"{filename}.png"
+        img.save(os.path.join(full_output_folder, file), pnginfo=metadata, compress_level=self.compress_level)
+        return {"ui": {"images": [{"filename": file, "subfolder": subfolder, "type": self.type}]}}