HalfikCode-ComfyUI-Prompt-DB/nodes.py

"""
ComfyUI custom nodes for Prompt Database.
"""

import json as _json
import re as _re
import folder_paths


def _clean_name(name):
    """Strip the 📖 LoRA marker suffix added by the JS frontend."""
    return _re.sub(r'\s*📖$', '', name).strip()
import comfy.sd
import comfy.utils
from . import db


class PromptDBManager:
    """
    Management node — provides a UI button to open the prompt database manager.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {},
            "hidden": {
                "unique_id": "UNIQUE_ID",
            },
        }

    RETURN_TYPES = ()
    FUNCTION = "manage"
    CATEGORY = "Prompt Database"
    OUTPUT_NODE = True

    def manage(self, unique_id=None):
        return {}


class PromptSelector:
    """
    Selector node — pick a category and prompt.
    Outputs:
      - prompt_text (STRING): the selected prompt text (no LoRA tags)
      - db_selection (DB_SELECTION): the category name (for PromptCombiner)
      - lora_stack (LORA_STACK): LoRAs attached to the selected prompt
    """

    @classmethod
    def INPUT_TYPES(cls):
        categories = db.get_categories()
        if not categories:
            categories = ["(no categories)"]

        all_prompts = db.get_prompts()
        prompt_names = list({p["name"] for p in all_prompts})
        if not prompt_names:
            prompt_names = ["(no prompts)"]

        return {
            "required": {
                "category": (categories, {"default": categories[0]}),
                "prompt_name": (prompt_names, {"default": prompt_names[0]}),
            },
        }

    RETURN_TYPES = ("STRING", "DB_SELECTION", "LORA_STACK")
    RETURN_NAMES = ("prompt_text", "db_selection", "lora_stack")
    FUNCTION = "select_prompt"
    CATEGORY = "Prompt Database"

    def select_prompt(self, category, prompt_name):
        prompt_name = _clean_name(prompt_name)
        prompt = db.get_prompt_by_name(prompt_name, category)
        if not prompt:
            prompt = db.get_prompt_by_name(prompt_name)
        if not prompt:
            return ("", category, [])

        text = prompt["text"]
        loras = prompt.get("loras", [])
        # Normalize lora entries
        lora_stack = []
        for l in loras:
            lora_stack.append({
                "name": l.get("name", ""),
                "strength": float(l.get("strength", 1.0)),
            })

        return (text, category, lora_stack)

    @classmethod
    def IS_CHANGED(cls, category, prompt_name):
        prompt_name = _clean_name(prompt_name)
        prompt = db.get_prompt_by_name(prompt_name, category)
        if prompt:
            return prompt["text"] + str(prompt.get("loras", []))
        return ""


class PromptCombiner:
    """
    Combiner node — accepts multiple DB_SELECTION inputs (category names),
    shows a prompt dropdown per category inside the node,
    and combines selected prompts into one STRING + merged LORA_STACK.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "separator": (
                    ["comma_space", "newline", "space", "none"],
                    {"default": "comma_space"},
                ),
                "prompt_selections": (
                    "STRING",
                    {"default": "{}", "multiline": False},
                ),
            },
            "optional": {
                "selection_1": ("DB_SELECTION", {"forceInput": True}),
            },
        }

    RETURN_TYPES = ("STRING", "LORA_STACK")
    RETURN_NAMES = ("combined_prompt", "lora_stack")
    FUNCTION = "combine"
    CATEGORY = "Prompt Database"
    OUTPUT_NODE = True

    def combine(self, separator="comma_space", prompt_selections="{}", **kwargs):
        sep_map = {
            "comma_space": ", ",
            "newline": "\n",
            "space": " ",
            "none": "",
        }
        sep = sep_map.get(separator, ", ")

        try:
            selections = _json.loads(prompt_selections)
        except _json.JSONDecodeError:
            selections = {}

        texts = []
        merged_loras = []
        seen_loras = set()

        selection_keys = [k for k in kwargs.keys() if k.startswith("selection_")]
        selection_keys.sort(key=lambda k: int(k.split("_")[1]))

        for key in selection_keys:
            category = kwargs[key]
            if not category:
                continue
            prompt_name = _clean_name(selections.get(key, ""))
            if not prompt_name:
                continue
            prompt = db.get_prompt_by_name(prompt_name, category)
            if prompt:
                texts.append(prompt["text"])
                # Collect LoRAs (deduplicate by name)
                for l in prompt.get("loras", []):
                    lname = l.get("name", "")
                    if lname and lname not in seen_loras:
                        seen_loras.add(lname)
                        merged_loras.append({
                            "name": lname,
                            "strength": float(l.get("strength", 1.0)),
                        })

        combined = sep.join(texts)
        return (combined, merged_loras)

    @classmethod
    def IS_CHANGED(cls, separator, prompt_selections="{}", **kwargs):
        return prompt_selections + str(kwargs)


class ApplyLoraStack:
    """
    Applies LoRAs from a LORA_STACK to a model and clip.
    The lora_config hidden widget stores per-LoRA overrides
    (enable/disable + custom strength).
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model": ("MODEL",),
                "clip": ("CLIP",),
                "lora_config": (
                    "STRING",
                    {"default": "{}", "multiline": False},
                ),
            },
            "optional": {
                "lora_stack": ("LORA_STACK",),
            },
        }

    RETURN_TYPES = ("MODEL", "CLIP")
    RETURN_NAMES = ("model", "clip")
    FUNCTION = "apply_loras"
    CATEGORY = "Prompt Database"

    def apply_loras(self, model, clip, lora_config="{}", lora_stack=None):
        if not lora_stack:
            return (model, clip)

        # Parse overrides
        try:
            config = _json.loads(lora_config)
        except _json.JSONDecodeError:
            config = {}

        for lora_entry in lora_stack:
            name = lora_entry.get("name", "")
            default_strength = lora_entry.get("strength", 1.0)
            if not name:
                continue

            # Apply per-LoRA overrides from config
            override = config.get(name, {})
            if not override.get("enabled", True):
                print(f"[PromptDB] LoRA skipped (disabled): {name}")
                continue

            strength = float(override.get("strength", default_strength))

            # Load LoRA file
            lora_path = folder_paths.get_full_path("loras", name)
            if not lora_path:
                print(f"[PromptDB] LoRA file not found: {name}")
                continue

            try:
                lora = comfy.utils.load_torch_file(lora_path, safe_load=True)
                model, clip = comfy.sd.load_lora_for_models(
                    model, clip, lora, strength, strength
                )
                print(f"[PromptDB] LoRA applied: {name} (strength={strength:.2f})")
            except Exception as e:
                print(f"[PromptDB] Error loading LoRA {name}: {e}")

        return (model, clip)

    @classmethod
    def IS_CHANGED(cls, model, clip, lora_config="{}", lora_stack=None):
        return str(lora_stack) + lora_config