diff --git a/src/character_select.py b/src/character_select.py index 490cd38..a51529b 100644 --- a/src/character_select.py +++ b/src/character_select.py @@ -1,50 +1,69 @@ import os import json +from .sn0w import ConfigReader class CharacterSelectNode: - # Get the directory of the current file - dir_path = os.path.dirname(os.path.realpath(__file__)) - # Check if the last part of the path is 'src' and adjust the path accordingly - if os.path.basename(dir_path) == "src": - dir_path = os.path.dirname(dir_path) + # Initialize class variables + character_dict = {} + final_character_dict = {} - # Construct the path to the 'characters.json' file - json_path = os.path.join(dir_path, 'characters.json') - # Load character data from the JSON file - with open(json_path, 'r') as file: - character_data = json.load(file) + @classmethod + def initialize(cls): + print("Initializing CharacterSelectNode...") - # Attempt to load 'custom_characters.json' and merge its contents - custom_json_path = os.path.join(dir_path, 'custom_characters.json') - if os.path.exists(custom_json_path): - with open(custom_json_path, 'r') as file: - custom_character_data = json.load(file) - if custom_character_data: # Check if the custom data is not empty + dir_path = os.path.dirname(os.path.realpath(__file__)) + if os.path.basename(dir_path) == "src": + dir_path = os.path.dirname(dir_path) + + cls.load_characters(dir_path) + + @classmethod + def load_characters(cls, dir_path): + json_path = os.path.join(dir_path, 'characters.json') + with open(json_path, 'r') as file: + character_data = json.load(file) + + custom_json_path = os.path.join(dir_path, 'custom_characters.json') + if os.path.exists(custom_json_path): + with open(custom_json_path, 'r') as file: + custom_character_data = json.load(file) for custom_character in custom_character_data: for character in character_data: if custom_character['name'] == character['name']: character['prompt'] += ", " + custom_character['prompt'] break - else: # If the custom character is not found, add it to character_data + else: character_data.append(custom_character) - # Create a dictionary mapping character names to their data - character_dict = {character['name']: character for character in character_data} + cls.character_dict = {character['name']: character for character in character_data} + + if ConfigReader.get_setting("sort_characters_by_series", False): + cls.final_character_dict = {name: cls.character_dict[name] for name in sorted(cls.character_dict, key=cls.extract_series_name)} + else: + cls.final_character_dict = {name: cls.character_dict[name] for name in sorted(cls.character_dict)} + + print("Final character dict initialized with:", len(cls.final_character_dict)) + + @staticmethod + def extract_series_name(character_name): + return character_name.split('(')[-1].split(')')[0].strip() @classmethod def INPUT_TYPES(cls): - character_names = ['None'] + sorted(cls.character_dict.keys()) + if not cls.final_character_dict: # Check if initialization is needed + cls.initialize() + character_names = ['None'] + list(cls.final_character_dict.keys()) return { "required": { "character": (character_names, ), "character_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.05, "round": 0.01}), "character_prompt": ("BOOLEAN", {"default": False}), - "xl": ("BOOLEAN", {"default": False},), + "xl": ("BOOLEAN", {"default": False}), }, } - RETURN_TYPES = ("STRING", "STRING", "BOOLEAN",) - RETURN_NAMES = ("CHARACTER STRING", "CHARACTER PROMPT", "XL",) + RETURN_TYPES = ("STRING", "STRING", "BOOLEAN") + RETURN_NAMES = ("CHARACTER STRING", "CHARACTER PROMPT", "XL") FUNCTION = "find_associated_string" CATEGORY = "sn0w" @@ -53,14 +72,11 @@ class CharacterSelectNode: if char_item: associated_string = char_item['associated_string'] prompt = char_item['prompt'] if character_prompt else "" - strength_part = f":{character_strength}" if character_strength != 1 else None + strength_part = f":{character_strength}" if character_strength != 1 else "" - if associated_string != "": - if strength_part: - return (f"({associated_string}{strength_part}), ", prompt, xl,) - else: - return (f"{associated_string}, ", prompt, xl,) + if associated_string: + return (f"{associated_string}{strength_part}, ", prompt, xl,) else: return ("", prompt, xl,) - return ("", "", "") + return ("", "", "") \ No newline at end of file