Vito Sansevero
7d65ee5516
utils.hashing.generate_prompt_hash is now the only prompt hash: PromptManagerBase._generate_hash delegates to it and the inline hashlib.sha256 copies in the prompt tracker and image monitor are gone. Remove PromptManagerBase._open_web_interface (it targeted a web_interface folder that never existed) and the unused restart_gallery.py script. IS_CHANGED on PromptManager and PromptManagerText now includes category and tags, so editing either re-runs the node and updates the saved metadata.
315 lines
12 KiB
Python
315 lines
12 KiB
Python
"""
|
|
PromptManagerBase: Shared logic for PromptManager node variants.
|
|
|
|
Provides database initialization, prompt saving, hashing, tag parsing,
|
|
search, gallery system management, and cleanup — extracted from the
|
|
duplicate code in prompt_manager.py and prompt_manager_text.py.
|
|
"""
|
|
|
|
import os
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
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 .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 get_prompt_tracker
|
|
from .utils.hashing import generate_prompt_hash
|
|
from .utils.prompt_graph import resolve_text, run_prompt_nodes
|
|
from .utils.usage_tracking import PENDING_FIRST_USE
|
|
except ImportError:
|
|
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 get_prompt_tracker
|
|
from utils.hashing import generate_prompt_hash
|
|
from utils.prompt_graph import resolve_text, run_prompt_nodes
|
|
from utils.usage_tracking import PENDING_FIRST_USE
|
|
|
|
|
|
class PromptManagerBase:
|
|
"""Mixin providing shared prompt management logic for ComfyUI nodes.
|
|
|
|
Handles database connection, prompt saving with deduplication,
|
|
tag parsing, search, gallery system lifecycle, and cleanup.
|
|
Subclasses only need to define ComfyUI-specific class attributes
|
|
(INPUT_TYPES, RETURN_TYPES, FUNCTION) and their execution method.
|
|
"""
|
|
|
|
def __init__(self, logger_name: str = "prompt_manager.node"):
|
|
self.logger = get_logger(logger_name)
|
|
self.logger.debug(f"Initializing {self.__class__.__name__} node")
|
|
|
|
self.db = PromptDatabase()
|
|
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()
|
|
|
|
self._start_gallery_system()
|
|
self.logger.debug(f"{self.__class__.__name__} node initialization completed")
|
|
|
|
def _save_prompt_to_database(
|
|
self,
|
|
text: str,
|
|
category: Optional[str] = None,
|
|
tags: Optional[list] = None,
|
|
count_run: bool = False,
|
|
) -> Optional[int]:
|
|
"""Save the prompt to the SQLite database.
|
|
|
|
Runs are normally counted by the queue hook (utils/usage_tracking.py), so a
|
|
save only counts one when asked to (batch items whose text the hook could not
|
|
see) or when the hook queued this prompt before it existed.
|
|
|
|
Args:
|
|
text: The prompt text
|
|
category: Optional category
|
|
tags: List of tags
|
|
count_run: Count this execution as a run of the prompt
|
|
|
|
Returns:
|
|
The prompt ID if saved successfully, None otherwise
|
|
"""
|
|
try:
|
|
prompt_hash = self._generate_hash(text)
|
|
self.logger.debug(f"Generated hash for prompt: {prompt_hash[:16]}...")
|
|
|
|
existing = self.db.get_prompt_by_hash(prompt_hash)
|
|
if existing:
|
|
self.logger.info(
|
|
f"Found existing prompt with ID {existing['id']}, updating metadata"
|
|
)
|
|
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")
|
|
if count_run:
|
|
self._record_use(existing["id"])
|
|
return existing["id"]
|
|
|
|
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}")
|
|
pending = getattr(self, "pending_first_use", PENDING_FIRST_USE)
|
|
runs = pending.consume(prompt_hash) + (1 if count_run else 0)
|
|
if runs:
|
|
self._record_use(prompt_id, times=runs)
|
|
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}")
|
|
return None
|
|
|
|
def _prompt_role(
|
|
self, prompt_graph: Any, unique_id: Any, text: str
|
|
) -> Dict[str, bool]:
|
|
"""How this node's execution relates to the run, from the queued graph.
|
|
|
|
Returns flags:
|
|
is_run_prompt: feeds a sampler's positive input (images belong to it)
|
|
seen_at_queue: the queue hook resolved exactly this text, so it already
|
|
counted the run and the image's metadata can identify the prompt
|
|
known: the graph and node id were available
|
|
"""
|
|
if not isinstance(prompt_graph, dict) or unique_id is None:
|
|
return {"is_run_prompt": True, "seen_at_queue": False, "known": False}
|
|
node_id = str(unique_id)
|
|
resolved = resolve_text(prompt_graph, node_id)
|
|
return {
|
|
"is_run_prompt": node_id in run_prompt_nodes(prompt_graph),
|
|
"seen_at_queue": resolved is not None
|
|
and generate_prompt_hash(resolved) == generate_prompt_hash(text),
|
|
"known": True,
|
|
}
|
|
|
|
def _track_prompt_execution(
|
|
self,
|
|
text: str,
|
|
encoding_text: str,
|
|
category: Optional[str],
|
|
tags: Optional[list],
|
|
additional_data: Dict[str, Any],
|
|
prompt_graph: Any = None,
|
|
unique_id: Any = None,
|
|
) -> Optional[int]:
|
|
"""Save the prompt and, for the run's positive prompt, register it for linking.
|
|
|
|
Negative-prompt nodes are saved but never become the prompt their images are
|
|
linked to. Text the queue hook could resolve is counted there and linked from
|
|
the image's own metadata; anything else (PromptSearchList batches, unknown
|
|
text nodes, or a resolution that doesn't match) is counted here and pushed
|
|
onto the FIFO queue. Without graph context the pre-3.2.4 behaviour is kept.
|
|
"""
|
|
role = self._prompt_role(prompt_graph, unique_id, text)
|
|
batch_item = (
|
|
role["known"] and role["is_run_prompt"] and not role["seen_at_queue"]
|
|
)
|
|
prompt_id = self._save_prompt_to_database(
|
|
text=text, category=category, tags=tags, count_run=batch_item
|
|
)
|
|
if prompt_id and role["is_run_prompt"]:
|
|
self.prompt_tracker.set_current_prompt(
|
|
prompt_text=encoding_text,
|
|
additional_data={**additional_data, "prompt_id": prompt_id},
|
|
push_to_queue=batch_item or not role["known"],
|
|
)
|
|
return prompt_id
|
|
|
|
def _record_use(self, prompt_id: int, times: int = 1) -> None:
|
|
"""Count runs; a tracking failure must never block generation."""
|
|
try:
|
|
self.db.record_prompt_use(prompt_id, times=times)
|
|
except Exception as e:
|
|
self.logger.warning(f"Could not record use of prompt {prompt_id}: {e}")
|
|
|
|
def _inject_lora_trigger_words(self, text: str) -> str:
|
|
"""Append LoRA trigger words if integration is enabled.
|
|
|
|
Returns the text unchanged if the integration is disabled or
|
|
LoraManager is not installed.
|
|
"""
|
|
try:
|
|
from .py.config import IntegrationConfig
|
|
except ImportError:
|
|
try:
|
|
from py.config import IntegrationConfig
|
|
except ImportError:
|
|
return text
|
|
|
|
if (
|
|
not IntegrationConfig.LORA_MANAGER_ENABLED
|
|
or not IntegrationConfig.LORA_TRIGGER_WORDS_ENABLED
|
|
):
|
|
return text
|
|
|
|
try:
|
|
from .py.lora_utils import get_trigger_cache, inject_trigger_words
|
|
except ImportError:
|
|
try:
|
|
from py.lora_utils import get_trigger_cache, inject_trigger_words
|
|
except ImportError:
|
|
return text
|
|
|
|
cache = get_trigger_cache()
|
|
|
|
# Lazy-load cache on first use
|
|
if not cache.is_loaded and IntegrationConfig.LORA_MANAGER_PATH:
|
|
cache.load(IntegrationConfig.LORA_MANAGER_PATH)
|
|
|
|
modified, injected = inject_trigger_words(text, cache)
|
|
if injected:
|
|
self.logger.info(f"Injected trigger words: {', '.join(injected)}")
|
|
return modified
|
|
|
|
def _generate_hash(self, text: str) -> str:
|
|
"""The prompt's deduplication hash (see utils.hashing.generate_prompt_hash)."""
|
|
return generate_prompt_hash(text)
|
|
|
|
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 []
|
|
|
|
try:
|
|
from .py.config import PromptManagerConfig
|
|
|
|
max_results = PromptManagerConfig.MAX_SEARCH_RESULTS
|
|
except Exception:
|
|
max_results = 100
|
|
|
|
results = self.db.search_prompts(
|
|
text=search_text.strip(),
|
|
category=None,
|
|
tags=None,
|
|
rating_min=None,
|
|
limit=max_results,
|
|
)
|
|
|
|
return results
|
|
|
|
except Exception as e:
|
|
self.logger.error(f"Error searching prompts: {e}")
|
|
return []
|
|
|
|
def search_prompts_api(self, search_text: str = "") -> List[Dict[str, Any]]:
|
|
"""API method for JavaScript UI to search prompts."""
|
|
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."""
|
|
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."""
|
|
try:
|
|
self.logger.debug("Starting gallery system...")
|
|
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."""
|
|
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.
|
|
|
|
Note: The image monitor is a singleton shared across all node instances,
|
|
so we intentionally do NOT stop it here. It should run for the entire
|
|
ComfyUI session. Stopping it would kill monitoring for all nodes.
|
|
"""
|
|
self.logger.debug("Gallery system cleaned up (monitor continues running)")
|