ComfyAssets-ComfyUI_PromptManager/prompt_manager.py

"""
PromptManager: Main custom node implementation that extends CLIPTextEncode
with persistent prompt storage and search capabilities.
"""

import time
from typing import Any, Dict, Optional, Tuple

try:
    from comfy.comfy_types import IO, ComfyNodeABC, InputTypeDict
except ImportError:
    # Fallback for older ComfyUI versions
    class ComfyNodeABC:
        pass

    class IO:
        STRING = "STRING"
        CLIP = "CLIP"
        CONDITIONING = "CONDITIONING"

    InputTypeDict = dict

try:
    from .prompt_manager_base import PromptManagerBase
except ImportError:
    import os
    import sys

    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
    from prompt_manager_base import PromptManagerBase


class PromptManager(PromptManagerBase, ComfyNodeABC):
    """
    A ComfyUI custom node that functions like CLIPTextEncode but adds:
    - Persistent storage of all prompts in SQLite database
    - Search and retrieval capabilities
    - Metadata management (categories, tags, ratings, notes)
    - Duplicate detection via SHA256 hashing
    """

    def __init__(self):
        super().__init__(logger_name="prompt_manager.node")

    @classmethod
    def INPUT_TYPES(cls) -> InputTypeDict:
        return {
            "required": {
                "text": (
                    IO.STRING,
                    {
                        "multiline": True,
                        "dynamicPrompts": True,
                        "tooltip": "The text prompt to be encoded and saved to database.",
                    },
                ),
                "clip": (
                    IO.CLIP,
                    {"tooltip": "The CLIP model used for encoding the text."},
                ),
            },
            "optional": {
                "category": (
                    IO.STRING,
                    {
                        "default": "",
                        "tooltip": "Optional category for organizing prompts (e.g., 'landscapes', 'portraits')",
                    },
                ),
                "tags": (
                    IO.STRING,
                    {
                        "default": "",
                        "tooltip": "Comma-separated tags for the prompt (e.g., 'anime, detailed, sunset')",
                    },
                ),
                "search_text": (
                    IO.STRING,
                    {
                        "default": "",
                        "tooltip": 'Search past prompts by text; add tag:name (or tag:"two words") to filter by tag',
                    },
                ),
                "prepend_text": (
                    IO.STRING,
                    {
                        "tooltip": "Text to prepend to the main prompt (connected STRING nodes will be added before the main text)"
                    },
                ),
                "append_text": (
                    IO.STRING,
                    {
                        "tooltip": "Text to append to the main prompt (connected STRING nodes will be added after the main text)"
                    },
                ),
            },
            "hidden": {
                # The queued graph and this node's id: used to tell the positive
                # prompt of a run from negative-prompt PromptManager nodes
                "prompt": "PROMPT",
                "unique_id": "UNIQUE_ID",
            },
        }

    RETURN_TYPES = (IO.CONDITIONING, IO.STRING)
    OUTPUT_TOOLTIPS = (
        "A conditioning containing the embedded text used to guide the diffusion model.",
        "The final combined text string (with prepend/append applied) that was encoded.",
    )
    FUNCTION = "encode_prompt"
    OUTPUT_NODE = True
    CATEGORY = "🫶 ComfyAssets/🧠 Prompts"
    DESCRIPTION = (
        "Encodes a text prompt using a CLIP model into an embedding that can be used to guide "
        "the diffusion model towards generating specific images. Additionally saves all prompts "
        "to a local SQLite database with optional metadata for search and retrieval."
    )

    def encode_prompt(
        self,
        clip,
        text: str,
        category: str = "",
        tags: str = "",
        search_text: str = "",
        prepend_text: str = "",
        append_text: str = "",
        prompt: Optional[Dict[str, Any]] = None,
        unique_id: Optional[str] = None,
    ) -> Tuple[Any]:
        """
        Encode the text prompt and save it to the database.

        Args:
            clip: The CLIP model for encoding
            text: The text prompt to encode
            category: Optional category for organization
            tags: Comma-separated tags
            search_text: Text to search for in past prompts
            prepend_text: Text to prepend to the main prompt
            append_text: Text to append to the main prompt

        Returns:
            Tuple containing the conditioning for the diffusion model and the final text string

        Raises:
            RuntimeError: If clip input is invalid
        """
        # Combine prepend, main text, and append text
        parts = []
        if prepend_text and prepend_text.strip():
            parts.append(prepend_text.strip())
        if text:
            parts.append(text)
        if append_text and append_text.strip():
            parts.append(append_text.strip())
        final_text = " ".join(parts)

        # Inject LoRA trigger words if integration is enabled
        final_text = self._inject_lora_trigger_words(final_text)

        # Use the combined text for encoding
        encoding_text = final_text

        # For database storage, save the original main text with metadata about prepend/append
        storage_text = text

        # Validate CLIP model
        if clip is None:
            error_msg = (
                "ERROR: clip input is invalid: None\n\n"
                "If the clip is from a checkpoint loader node your checkpoint does not "
                "contain a valid clip or text encoder model."
            )
            self.logger.error("CLIP validation failed: clip input is None")
            raise RuntimeError(error_msg)

        # Save prompt to database and set execution context for gallery tracking
        prompt_id = None
        extended_tags = []
        if storage_text and storage_text.strip():
            self.logger.debug(f"Processing prompt text: {storage_text[:100]}...")

            # Add prepend/append info to tags if they exist
            extended_tags = self._parse_tags(tags) or []
            if prepend_text and prepend_text.strip():
                extended_tags.append(f"prepend:{prepend_text.strip()[:50]}")
            if append_text and append_text.strip():
                extended_tags.append(f"append:{append_text.strip()[:50]}")

            try:
                prompt_id = self._track_prompt_execution(
                    text=storage_text.strip(),
                    encoding_text=encoding_text.strip(),
                    category=category.strip() if category else None,
                    tags=extended_tags if extended_tags else None,
                    additional_data={
                        "category": category.strip() if category else None,
                        "tags": extended_tags,
                        "prepend_text": (
                            prepend_text.strip() if prepend_text else None
                        ),
                        "append_text": append_text.strip() if append_text else None,
                        "final_text": encoding_text.strip(),
                    },
                    prompt_graph=prompt,
                    unique_id=unique_id,
                )
                self.logger.debug(f"Tracked prompt ID: {prompt_id}")

            except Exception as e:
                # Log error but don't fail the encoding
                self.logger.warning(f"Failed to save prompt to database: {e}")

        # Perform standard CLIP text encoding using the combined text
        self.logger.debug(
            f"Performing CLIP text encoding on combined text: {encoding_text[:100]}..."
        )
        tokens = clip.tokenize(encoding_text)
        conditioning = clip.encode_from_tokens_scheduled(tokens)

        # Register with ComfyUI integration for standard metadata compatibility
        node_id = f"promptmanager_{int(time.time() * 1000)}"
        self.comfyui_integration.register_prompt(
            node_id,
            encoding_text.strip(),
            {
                "category": category.strip() if category else None,
                "tags": extended_tags,
                "prompt_id": prompt_id,
                "prepend_text": prepend_text.strip() if prepend_text else None,
                "append_text": append_text.strip() if append_text else None,
            },
        )

        self.logger.info(f"CLIP encoding completed, text: {repr(encoding_text)[:80]}")
        return (conditioning, encoding_text)

    @classmethod
    def IS_CHANGED(
        cls,
        clip,
        text="",
        category="",
        tags="",
        search_text="",
        prepend_text="",
        append_text="",
        **kwargs,
    ):
        """
        ComfyUI method to determine if node needs re-execution.

        Returns a hash of input values that affect the conditioning output.
        This enables proper branch execution - only re-execute when inputs change.
        """
        import hashlib

        combined = f"{text}|{prepend_text}|{append_text}"
        return hashlib.sha256(combined.encode()).hexdigest()