Files
ComfyAssets-ComfyUI_PromptM…/prompt_manager.py
T
Vito 55fa2c4870 fix(runs): count cached PromptManager nodes fed by a runtime text link on every job (#347) (#423)
* fix(usage): count cached re-runs of runtime-linked prompt text (#347)

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.

* fix(usage): never replay a node twice in one job; dedupe receipts once complete (#347)

A repeated execution_cached id or message for the same job now counts the
node once. Receipts replace equal older ones only when their job finishes
(or the next job starts), so a list still recording [A] of [A, B] no longer
drops an older [A] receipt. Node-level tests check the receipt each
execution records.

* test(comfyui): count #347's cached-run cases in the execution suite pin (#347)

The pin listed 19 execution-text-link cases per project; #347 adds five node
kinds and eight one-renderer cases, so the suite lists 32 per project and 118
in total with #332's image-attribution case from main.
2026-10-28 16:37:24 +00:00

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."""