diff --git a/.gitignore b/.gitignore index 1d7901e..eeb8a6e 100644 --- a/.gitignore +++ b/.gitignore @@ -1 +1 @@ -/__pycache__ \ No newline at end of file +**/__pycache__ diff --git a/__init__.py b/__init__.py index 5a7d35a..acf126c 100644 --- a/__init__.py +++ b/__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"] \ No newline at end of file diff --git a/tavern_utils/__init__.py b/tavern_utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tavern_utils/read_tavern_card.py b/tavern_utils/read_tavern_card.py new file mode 100644 index 0000000..c79293a --- /dev/null +++ b/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()),) \ No newline at end of file diff --git a/tavern_utils/tavern_card_prototype.py b/tavern_utils/tavern_card_prototype.py new file mode 100644 index 0000000..dff031e --- /dev/null +++ b/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") \ No newline at end of file diff --git a/tavern_utils/write_tavern_card.py b/tavern_utils/write_tavern_card.py new file mode 100644 index 0000000..078b1af --- /dev/null +++ b/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}]}} \ No newline at end of file