Vito Sansevero
1b67cb7ccb
A PromptManager-family node whose text comes from a runtime link (PreviewAny, an LLM, a list node) counted only the job that executed it; later jobs served from ComfyUI's output cache counted nothing. Each execution now records a receipt (digests of its STRING outputs and hashes of the text it stores) per run key and queue item. On execution_cached the receipt whose digests equal the cache entry's text is replayed once for the job, with the node's role and what the queue hook already saw recomputed from the job's own graph. The executing node and the replay count through one shared own_counts. Unmatched cache entries count nothing, as before; a node that executes after its replay in the same job skips its own count. No schema change, no backfill.
316 lines
11 KiB
Python
316 lines
11 KiB
Python
"""
|
|
PromptManager: Main custom node implementation that extends CLIPTextEncode
|
|
with persistent prompt storage and search capabilities.
|
|
"""
|
|
|
|
import time
|
|
from typing import Any, Dict, Optional, Tuple
|
|
|
|
if __package__:
|
|
from .utils.execution_context_capture import record_runtime_context
|
|
else:
|
|
from utils.execution_context_capture import record_runtime_context
|
|
|
|
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}")
|
|
self._mark_untracked(prompt, unique_id)
|
|
else:
|
|
self._skip_empty_prompt(prompt, unique_id)
|
|
|
|
# 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]}")
|
|
result = (conditioning, encoding_text)
|
|
record_runtime_context(
|
|
type(self).__name__, text, prepend_text, append_text, encoding_text
|
|
)
|
|
self._record_cached_item(unique_id, (encoding_text,), storage_text)
|
|
return result
|
|
|
|
@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 or
|
|
the saved prompt's metadata. This enables proper branch execution -
|
|
only re-execute when inputs change.
|
|
"""
|
|
import hashlib
|
|
|
|
combined = f"{text}|{prepend_text}|{append_text}|{category}|{tags}"
|
|
return hashlib.sha256(combined.encode()).hexdigest()
|
|
|
|
|
|
class PromptManagerFocused(PromptManager):
|
|
"""The editor-first UI uses the existing encoding and storage contract."""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> InputTypeDict:
|
|
inputs = super().INPUT_TYPES()
|
|
for name in ("prepend_text", "append_text"):
|
|
kind, options = inputs["optional"][name]
|
|
inputs["optional"][name] = (
|
|
kind,
|
|
{**options, "default": "", "forceInput": True},
|
|
)
|
|
return inputs
|
|
|
|
|
|
class PromptManagerSplit(PromptManagerFocused):
|
|
"""The persistent library layout shares Focused encoding and storage."""
|