TKRLAB-ComfyUI_Prompt_List_JSON/ComfyUI_Prompt_JSON.py

import json
import os
import folder_paths
from collections import OrderedDict
import random

# Define a default boolean parameter for input
BOOLEAN = ("BOOLEAN", {"default": False})

class ComfyUI_Prompt_JSON:
    def __init__(self):
        # Define the folder for storing prompts
        self.prompt_list_folder = os.path.join(folder_paths.base_path, "Prompt")
        # Ensure the folder exists
        if not os.path.exists(self.prompt_list_folder):
            os.makedirs(self.prompt_list_folder, exist_ok=True)

    def load_json(self, file_path):
        """Load data from the specified JSON file into an OrderedDict."""
        if not os.path.exists(file_path):
            return OrderedDict()
        with open(file_path, 'r', encoding='utf-8') as file:
            return json.load(file, object_pairs_hook=OrderedDict) or OrderedDict()

    def save_json(self, file_path, data):
        """Save data to the specified JSON file."""
        with open(file_path, 'w', encoding='utf-8') as file:
            json.dump(data, file, ensure_ascii=False, indent=4)

    @classmethod
    def INPUT_TYPES(cls):
        """
        Define the input types for the node interface.
        - 'Prompt List': A dropdown to select existing prompts or create a new one.
        - 'New List Name': A string input for creating a new list (used when creating a new list).
        - 'Random': A boolean to enable/disable random prompt selection.
        - 'Overwrite': A boolean to allow overwriting an existing prompt.
        - 'Console Log': A boolean to enable logging to the console.
        - 'Prompt Name': A string for naming the current prompt.
        - 'Positive Prompt': Multiline text for the positive part of the prompt.
        - 'Negative Prompt': Multiline text for the negative part of the prompt.
        """
        instance = cls()
        # Include existing JSON files along with the option to create a "New List"
        prompt_files = ["New List"] + [f for f in os.listdir(instance.prompt_list_folder) if f.endswith('.json')]
        return {
            "required": {
                "Prompt List": (prompt_files, ),
                "New List Name": ("STRING", {"multiline": False}),
                "Random": BOOLEAN,
                "Overwrite": BOOLEAN,
                "Console Log": BOOLEAN,
                "Prompt Name": ("STRING", {"multiline": False}),
                "Positive Prompt": ("STRING", {"multiline": True}),
                "Negative Prompt": ("STRING", {"multiline": True}),
            }
        }

    # Define return types for the process function
    RETURN_TYPES = ("STRING", "STRING", "STRING")
    RETURN_NAMES = ("positive", "negative", "full list")
    FUNCTION = "process"
    CATEGORY = "prompt"

    def process(self, **kwargs):
        """
        Process prompts:
        - Load the selected prompt list.
        - Optionally select a random prompt from the list.
        - Return the positive and negative prompts, as well as the full list.
        """
        prompt_list_name = kwargs.get("Prompt List", "").strip()
        new_list_name = kwargs.get("New List Name", "").strip()
        random_selection = kwargs.get("Random", False)
        overwrite = kwargs.get("Overwrite", False)
        console_log = kwargs.get("Console Log", False)
        prompt_name = kwargs.get("Prompt Name", "").strip()
        positive_prompt = kwargs.get("Positive Prompt", "").strip()
        negative_prompt = kwargs.get("Negative Prompt", "").strip()

        # Handle "New List" option: Create a new list if selected
        if prompt_list_name == "New List":
            if not new_list_name:
                raise ValueError("'New List Name' is required when creating a new list.")
            prompt_list_name = new_list_name + ".json"

        # Construct the file path for the selected or new prompt list
        prompt_list_path = os.path.join(self.prompt_list_folder, prompt_list_name)
        data = self.load_json(prompt_list_path)

        # Handle random selection of a prompt from the list
        if random_selection:
            if not data:
                raise ValueError("The selected prompt list is empty or invalid.")
            # Choose a random prompt name and fetch its details
            prompt_name = random.choice(list(data.keys()))
            positive_prompt = data[prompt_name]["positive"]
            negative_prompt = data[prompt_name]["negative"]
        else:
            # Validate required fields for manual input
            if not prompt_name:
                raise ValueError("'Prompt Name' is required when Random is not enabled.")
            if prompt_name in data:
                if overwrite:
                    # Overwrite existing positive/negative prompts if allowed
                    if positive_prompt:
                        data[prompt_name]["positive"] = positive_prompt
                    if negative_prompt or not data[prompt_name]["negative"]:
                        data[prompt_name]["negative"] = negative_prompt
                else:
                    # Load existing prompt details
                    positive_prompt = data[prompt_name]["positive"]
                    negative_prompt = data[prompt_name]["negative"]
            else:
                # Add a new prompt to the list
                if not positive_prompt:
                    raise ValueError("'Positive Prompt' is required when adding a new prompt.")
                data[prompt_name] = OrderedDict([
                    ("positive", positive_prompt),
                    ("negative", negative_prompt),
                ])

        # Optionally log the prompt details to the console
        if console_log:
            import shutil
            terminal_width = shutil.get_terminal_size((80, 20)).columns

            # Function to format multiline text with a background color
            def format_with_bg_multiline(text, color_code):
                lines = [text[i:i + terminal_width] for i in range(0, len(text), terminal_width)]
                return "\n".join(f"\033[1;{color_code}m{line.ljust(terminal_width)}\033[0;0m" for line in lines)

            # Log formatted details
            print(format_with_bg_multiline(f"{prompt_name} from {prompt_list_name}", "40;37"))
            print(format_with_bg_multiline(f"POSITIVE: {positive_prompt}", "40;32"))
            print(format_with_bg_multiline(f"NEGATIVE: {negative_prompt}", "40;31"))

        # Save the updated list back to the file
        self.save_json(prompt_list_path, data)

        # Return the processed positive prompt, negative prompt, and full list as JSON
        full_list = json.dumps(data, ensure_ascii=False, indent=4)
        return positive_prompt, negative_prompt, full_list

    @classmethod
    def IS_CHANGED(cls, *args, **kwargs):
        """Check for changes in the node (stub implementation)."""
        return float("NaN")

# Node mappings for the interface
NODE_CLASS_MAPPINGS = {
    "ComfyUI_Prompt_JSON": ComfyUI_Prompt_JSON
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "ComfyUI_Prompt_JSON": "Prompt JSON"
}