Vito Sansevero
2774ed921b
The api.queuePrompt wrapper introduced in PR #91 only forwarded 2 of 3 arguments, dropping the options object that carries partialExecutionTargets. This caused the server to execute all output nodes instead of just the targeted branch. Also removes debug console.log and print statements added during development. Fixes #94
504 lines
18 KiB
Python
504 lines
18 KiB
Python
"""
|
|
PromptManager: Main custom node implementation that extends CLIPTextEncode
|
|
with persistent prompt storage and search capabilities.
|
|
"""
|
|
|
|
import datetime
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import time
|
|
import webbrowser
|
|
from typing import Any, Dict, List, Optional, Tuple
|
|
|
|
# Import logging system
|
|
try:
|
|
from .utils.logging_config import get_logger
|
|
except ImportError:
|
|
import sys
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
from utils.logging_config import get_logger
|
|
|
|
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 .database.operations import PromptDatabase
|
|
from .utils.comfyui_integration import get_comfyui_integration
|
|
from .utils.image_monitor import get_image_monitor
|
|
from .utils.prompt_tracker import PromptExecutionContext, get_prompt_tracker
|
|
except ImportError:
|
|
# For direct imports when not in a package
|
|
import os
|
|
import sys
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
from database.operations import PromptDatabase
|
|
from utils.comfyui_integration import get_comfyui_integration
|
|
from utils.image_monitor import get_image_monitor
|
|
from utils.prompt_tracker import PromptExecutionContext, get_prompt_tracker
|
|
|
|
|
|
class PromptManager(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):
|
|
self.logger = get_logger("prompt_manager.node")
|
|
self.logger.debug("Initializing PromptManager node")
|
|
|
|
self.db = PromptDatabase()
|
|
# Use singleton getters to ensure only one tracker/monitor exists
|
|
self.prompt_tracker = get_prompt_tracker(self.db)
|
|
self.image_monitor = get_image_monitor(self.db, self.prompt_tracker)
|
|
self.comfyui_integration = get_comfyui_integration()
|
|
|
|
# Start image monitoring automatically
|
|
self._start_gallery_system()
|
|
self.logger.debug("PromptManager node initialization completed")
|
|
|
|
@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 for past prompts containing this text",
|
|
},
|
|
),
|
|
"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)"
|
|
},
|
|
),
|
|
},
|
|
}
|
|
|
|
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"
|
|
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 = "",
|
|
) -> 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
|
|
final_text = ""
|
|
if prepend_text and prepend_text.strip():
|
|
final_text += prepend_text.strip() + " "
|
|
final_text += text if text else ""
|
|
if append_text and append_text.strip():
|
|
final_text += " " + append_text.strip()
|
|
|
|
# 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
|
|
|
|
# Search functionality is now handled by the JavaScript UI
|
|
# The search parameters are still available for backend processing if needed
|
|
|
|
# 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
|
|
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._save_prompt_to_database(
|
|
text=storage_text.strip(), # Always strip whitespace
|
|
category=category.strip() if category else None,
|
|
tags=extended_tags if extended_tags else None,
|
|
)
|
|
|
|
# Set current prompt for image tracking
|
|
if prompt_id:
|
|
execution_id = self.prompt_tracker.set_current_prompt(
|
|
prompt_text=encoding_text.strip(), # Use final combined text for tracking
|
|
additional_data={
|
|
"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,
|
|
"final_text": encoding_text.strip(), # Store final combined text
|
|
},
|
|
)
|
|
self.logger.debug(
|
|
f"Set execution context: {execution_id} for 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}")
|
|
# Already logged above, no need for additional print
|
|
|
|
# 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)}" # Unique node ID
|
|
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)
|
|
|
|
def _save_prompt_to_database(
|
|
self, text: str, category: Optional[str] = None, tags: Optional[list] = None
|
|
) -> Optional[int]:
|
|
"""
|
|
Save the prompt to the SQLite database.
|
|
|
|
Args:
|
|
text: The prompt text
|
|
category: Optional category
|
|
tags: List of tags
|
|
|
|
Returns:
|
|
The prompt ID if saved successfully, None otherwise
|
|
"""
|
|
try:
|
|
# Generate hash for duplicate detection
|
|
prompt_hash = self._generate_hash(text)
|
|
self.logger.debug(f"Generated hash for prompt: {prompt_hash[:16]}...")
|
|
|
|
# Check if prompt already exists
|
|
existing = self.db.get_prompt_by_hash(prompt_hash)
|
|
if existing:
|
|
self.logger.info(
|
|
f"Found existing prompt with ID {existing['id']}, updating metadata"
|
|
)
|
|
# Update metadata if this is a duplicate with new info
|
|
if any([category, tags]):
|
|
self.db.update_prompt_metadata(
|
|
prompt_id=existing["id"], category=category, tags=tags
|
|
)
|
|
self.logger.debug("Updated metadata for existing prompt")
|
|
return existing["id"]
|
|
|
|
# Save new prompt
|
|
self.logger.debug(
|
|
f"Saving new prompt with category: {category}, tags: {tags}"
|
|
)
|
|
prompt_id = self.db.save_prompt(
|
|
text=text, category=category, tags=tags, prompt_hash=prompt_hash
|
|
)
|
|
|
|
if prompt_id:
|
|
self.logger.debug(f"Successfully saved new prompt with ID: {prompt_id}")
|
|
else:
|
|
self.logger.warning("Failed to save prompt - no ID returned")
|
|
|
|
return prompt_id
|
|
|
|
except Exception as e:
|
|
self.logger.error(f"Error saving prompt to database: {e}")
|
|
# Already logged above, no need for additional print
|
|
return None
|
|
|
|
def _generate_hash(self, text: str) -> str:
|
|
"""
|
|
Generate SHA256 hash for the prompt text.
|
|
|
|
Args:
|
|
text: The prompt text to hash
|
|
|
|
Returns:
|
|
Hexadecimal string representation of the SHA256 hash
|
|
"""
|
|
# Normalize text for consistent hashing (strip whitespace, normalize case)
|
|
normalized_text = text.strip().lower()
|
|
return hashlib.sha256(normalized_text.encode("utf-8")).hexdigest()
|
|
|
|
def _parse_tags(self, tags_string: str) -> Optional[list]:
|
|
"""
|
|
Parse comma-separated tags string into a list.
|
|
|
|
Args:
|
|
tags_string: Comma-separated string of tags
|
|
|
|
Returns:
|
|
List of parsed tags, or None if no valid tags found
|
|
"""
|
|
if not tags_string or not tags_string.strip():
|
|
return None
|
|
|
|
tags = [tag.strip() for tag in tags_string.split(",") if tag.strip()]
|
|
return tags if tags else None
|
|
|
|
def _search_prompts(self, search_text: str = "") -> List[Dict[str, Any]]:
|
|
"""
|
|
Search for past prompts by text content.
|
|
|
|
Args:
|
|
search_text: Text to search for in prompt database
|
|
|
|
Returns:
|
|
List of matching prompt dictionaries with metadata
|
|
"""
|
|
try:
|
|
if not search_text or not search_text.strip():
|
|
return []
|
|
|
|
results = self.db.search_prompts(
|
|
text=search_text.strip(),
|
|
category=None,
|
|
tags=None,
|
|
rating_min=None,
|
|
limit=50,
|
|
)
|
|
|
|
return results
|
|
|
|
except Exception as e:
|
|
self.logger.error(f"Error searching prompts: {e}")
|
|
return []
|
|
|
|
def _open_web_interface(self):
|
|
"""
|
|
Open the web interface in the default browser.
|
|
|
|
Attempts to locate and open the web interface HTML file.
|
|
Logs warnings if the interface is not properly configured.
|
|
"""
|
|
try:
|
|
# Look for a web interface directory
|
|
current_dir = os.path.dirname(os.path.abspath(__file__))
|
|
web_dir = os.path.join(current_dir, "web_interface")
|
|
|
|
if os.path.exists(web_dir):
|
|
# If web interface exists, try to start it
|
|
index_path = os.path.join(web_dir, "index.html")
|
|
if os.path.exists(index_path):
|
|
webbrowser.open(f"file://{index_path}")
|
|
self.logger.info("Web interface opened in browser")
|
|
else:
|
|
self.logger.warning(
|
|
f"Web interface directory found but no index.html. Please check {web_dir} for setup instructions"
|
|
)
|
|
else:
|
|
self.logger.info(
|
|
"Web interface not yet implemented. This feature will open a web-based prompt management interface when the web_interface directory is created."
|
|
)
|
|
|
|
except Exception as e:
|
|
self.logger.error(f"Error opening web interface: {e}")
|
|
|
|
def search_prompts_api(self, search_text: str = "") -> List[Dict[str, Any]]:
|
|
"""
|
|
API method for JavaScript UI to search prompts.
|
|
|
|
Args:
|
|
search_text: Text to search for in prompts
|
|
|
|
Returns:
|
|
List of matching prompt dictionaries
|
|
"""
|
|
return self._search_prompts(search_text=search_text)
|
|
|
|
def get_recent_prompts_api(self, limit: int = 20) -> List[Dict[str, Any]]:
|
|
"""
|
|
API method for JavaScript UI to get recent prompts.
|
|
|
|
Args:
|
|
limit: Maximum number of recent prompts to retrieve
|
|
|
|
Returns:
|
|
List of recent prompt dictionaries ordered by creation time
|
|
"""
|
|
try:
|
|
return self.db.get_recent_prompts(limit=limit)
|
|
except Exception as e:
|
|
self.logger.error(f"Error getting recent prompts: {e}")
|
|
return []
|
|
|
|
def _start_gallery_system(self):
|
|
"""
|
|
Initialize and start the gallery monitoring system.
|
|
|
|
Starts the image monitor which watches for new generated images
|
|
and links them to their source prompts in the database.
|
|
"""
|
|
try:
|
|
self.logger.debug("Starting gallery system...")
|
|
|
|
# Start image monitoring
|
|
self.image_monitor.start_monitoring()
|
|
|
|
self.logger.debug("Gallery system started successfully")
|
|
|
|
except Exception as e:
|
|
self.logger.error(f"Failed to start gallery system: {e}")
|
|
self.logger.warning("Gallery features will be disabled")
|
|
|
|
def get_gallery_status(self) -> Dict[str, Any]:
|
|
"""
|
|
Get status of the gallery system.
|
|
|
|
Returns:
|
|
Dictionary containing status information for image monitor and prompt tracker
|
|
"""
|
|
return {
|
|
"image_monitor": self.image_monitor.get_status(),
|
|
"prompt_tracker": self.prompt_tracker.get_status(),
|
|
}
|
|
|
|
def cleanup_gallery_system(self):
|
|
"""
|
|
Clean up gallery system resources.
|
|
|
|
Stops image monitoring and releases associated resources.
|
|
Called automatically during object destruction.
|
|
"""
|
|
try:
|
|
if hasattr(self, "image_monitor"):
|
|
self.image_monitor.stop_monitoring()
|
|
self.logger.debug("Gallery system cleaned up")
|
|
except Exception as e:
|
|
self.logger.error(f"Error cleaning up gallery system: {e}")
|
|
|
|
def __del__(self):
|
|
"""
|
|
Cleanup when object is destroyed.
|
|
|
|
Ensures proper resource cleanup by stopping the gallery system.
|
|
"""
|
|
self.cleanup_gallery_system()
|
|
|
|
@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
|
|
|
|
# Combine all text inputs that affect the conditioning output
|
|
# Note: search_text doesn't affect output, so it's excluded
|
|
combined = f"{text}|{prepend_text}|{append_text}"
|
|
|
|
return hashlib.sha256(combined.encode()).hexdigest()
|