fix: harden nodes, validators, and utility modules
Phase 0-3 fixes from codebase audit: - Extract shared node logic into prompt_manager_base.py (DRY) - Add input validation guards to validators.py - Fix logging_config.py buffer management - Harden image_monitor.py against race conditions - Fix prompt_tracker.py cleanup edge cases - Fix diagnostics.py import path - Wire config.py server instance correctly
This commit is contained in:
Vito Sansevero
+2
-2
@@ -79,7 +79,7 @@ except Exception as e:
|
||||
|
||||
logger = get_logger("prompt_manager.init")
|
||||
logger.error(f"Failed to register API routes: {e}")
|
||||
except:
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Start image monitoring globally at module import time
|
||||
@@ -107,7 +107,7 @@ except Exception as e:
|
||||
from .utils.logging_config import get_logger
|
||||
_init_logger = get_logger("prompt_manager.init")
|
||||
_init_logger.error(f"Failed to start image monitoring: {e}")
|
||||
except:
|
||||
except Exception:
|
||||
print(f"[ComfyUI-PromptManager] Warning: Failed to start image monitoring: {e}")
|
||||
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]
|
||||
|
||||
+9
-270
@@ -3,22 +3,8 @@ 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
|
||||
from typing import Any, Tuple
|
||||
|
||||
try:
|
||||
from comfy.comfy_types import IO, ComfyNodeABC, InputTypeDict
|
||||
@@ -35,23 +21,16 @@ except ImportError:
|
||||
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
|
||||
from .prompt_manager_base import PromptManagerBase
|
||||
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
|
||||
from prompt_manager_base import PromptManagerBase
|
||||
|
||||
|
||||
class PromptManager(ComfyNodeABC):
|
||||
class PromptManager(PromptManagerBase, ComfyNodeABC):
|
||||
"""
|
||||
A ComfyUI custom node that functions like CLIPTextEncode but adds:
|
||||
- Persistent storage of all prompts in SQLite database
|
||||
@@ -61,18 +40,7 @@ class PromptManager(ComfyNodeABC):
|
||||
"""
|
||||
|
||||
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")
|
||||
super().__init__(logger_name="prompt_manager.node")
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls) -> InputTypeDict:
|
||||
@@ -183,9 +151,6 @@ class PromptManager(ComfyNodeABC):
|
||||
# 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 = (
|
||||
@@ -210,7 +175,7 @@ class PromptManager(ComfyNodeABC):
|
||||
|
||||
try:
|
||||
prompt_id = self._save_prompt_to_database(
|
||||
text=storage_text.strip(), # Always strip whitespace
|
||||
text=storage_text.strip(),
|
||||
category=category.strip() if category else None,
|
||||
tags=extended_tags if extended_tags else None,
|
||||
)
|
||||
@@ -218,7 +183,7 @@ class PromptManager(ComfyNodeABC):
|
||||
# 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
|
||||
prompt_text=encoding_text.strip(),
|
||||
additional_data={
|
||||
"category": category.strip() if category else None,
|
||||
"tags": extended_tags,
|
||||
@@ -227,7 +192,7 @@ class PromptManager(ComfyNodeABC):
|
||||
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
|
||||
"final_text": encoding_text.strip(),
|
||||
},
|
||||
)
|
||||
self.logger.debug(
|
||||
@@ -237,7 +202,6 @@ class PromptManager(ComfyNodeABC):
|
||||
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(
|
||||
@@ -247,7 +211,7 @@ class PromptManager(ComfyNodeABC):
|
||||
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
|
||||
node_id = f"promptmanager_{int(time.time() * 1000)}"
|
||||
self.comfyui_integration.register_prompt(
|
||||
node_id,
|
||||
encoding_text.strip(),
|
||||
@@ -263,228 +227,6 @@ class PromptManager(ComfyNodeABC):
|
||||
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):
|
||||
@@ -496,8 +238,5 @@ class PromptManager(ComfyNodeABC):
|
||||
"""
|
||||
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()
|
||||
|
||||
@@ -0,0 +1,231 @@
|
||||
"""
|
||||
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 hashlib
|
||||
import os
|
||||
import webbrowser
|
||||
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
|
||||
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
|
||||
|
||||
|
||||
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
|
||||
) -> 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:
|
||||
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")
|
||||
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}")
|
||||
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 _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
|
||||
"""
|
||||
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 []
|
||||
|
||||
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 _open_web_interface(self):
|
||||
"""Open the web interface in the default browser."""
|
||||
try:
|
||||
current_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
web_dir = os.path.join(current_dir, "web_interface")
|
||||
|
||||
if os.path.exists(web_dir):
|
||||
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. "
|
||||
f"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."""
|
||||
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."""
|
||||
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."""
|
||||
self.cleanup_gallery_system()
|
||||
+8
-267
@@ -3,22 +3,8 @@ PromptManagerText: A text-only version of PromptManager that outputs STRING
|
||||
without CLIP encoding, while maintaining all database and search features.
|
||||
"""
|
||||
|
||||
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
|
||||
from typing import Tuple
|
||||
|
||||
try:
|
||||
from comfy.comfy_types import IO, ComfyNodeABC, InputTypeDict
|
||||
@@ -35,23 +21,16 @@ except ImportError:
|
||||
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
|
||||
from .prompt_manager_base import PromptManagerBase
|
||||
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
|
||||
from prompt_manager_base import PromptManagerBase
|
||||
|
||||
|
||||
class PromptManagerText(ComfyNodeABC):
|
||||
class PromptManagerText(PromptManagerBase, ComfyNodeABC):
|
||||
"""
|
||||
A ComfyUI custom node that provides all PromptManager features but outputs
|
||||
only a STRING without CLIP encoding. Includes:
|
||||
@@ -63,18 +42,7 @@ class PromptManagerText(ComfyNodeABC):
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.logger = get_logger("prompt_manager_text.node")
|
||||
self.logger.debug("Initializing PromptManagerText 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("PromptManagerText node initialization completed")
|
||||
super().__init__(logger_name="prompt_manager_text.node")
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls) -> InputTypeDict:
|
||||
@@ -174,9 +142,6 @@ class PromptManagerText(ComfyNodeABC):
|
||||
# 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
|
||||
|
||||
# Save prompt to database and set execution context for gallery tracking
|
||||
prompt_id = None
|
||||
if storage_text and storage_text.strip():
|
||||
@@ -191,7 +156,7 @@ class PromptManagerText(ComfyNodeABC):
|
||||
|
||||
try:
|
||||
prompt_id = self._save_prompt_to_database(
|
||||
text=storage_text.strip(), # Always strip whitespace
|
||||
text=storage_text.strip(),
|
||||
category=category.strip() if category else None,
|
||||
tags=extended_tags if extended_tags else None,
|
||||
)
|
||||
@@ -199,7 +164,7 @@ class PromptManagerText(ComfyNodeABC):
|
||||
# Set current prompt for image tracking
|
||||
if prompt_id:
|
||||
execution_id = self.prompt_tracker.set_current_prompt(
|
||||
prompt_text=final_text.strip(), # Use final combined text for tracking
|
||||
prompt_text=final_text.strip(),
|
||||
additional_data={
|
||||
"category": category.strip() if category else None,
|
||||
"tags": extended_tags,
|
||||
@@ -219,7 +184,7 @@ class PromptManagerText(ComfyNodeABC):
|
||||
self.logger.warning(f"Failed to save prompt to database: {e}")
|
||||
|
||||
# Register with ComfyUI integration for standard metadata compatibility
|
||||
node_id = f"promptmanagertext_{int(time.time() * 1000)}" # Unique node ID
|
||||
node_id = f"promptmanagertext_{int(time.time() * 1000)}"
|
||||
self.comfyui_integration.register_prompt(
|
||||
node_id,
|
||||
final_text.strip(),
|
||||
@@ -235,227 +200,6 @@ class PromptManagerText(ComfyNodeABC):
|
||||
self.logger.debug(f"Text processing completed: {final_text[:100]}...")
|
||||
return (final_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}")
|
||||
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, text="", category="", tags="", search_text="",
|
||||
prepend_text="", append_text="", **kwargs):
|
||||
@@ -467,8 +211,5 @@ class PromptManagerText(ComfyNodeABC):
|
||||
"""
|
||||
import hashlib
|
||||
|
||||
# Combine all text inputs that affect the 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()
|
||||
|
||||
+1
-1
@@ -79,7 +79,7 @@ class GalleryConfig:
|
||||
PROCESSING_DELAY = 2.0 # Seconds to wait before processing new files
|
||||
|
||||
# Prompt tracking settings
|
||||
PROMPT_TIMEOUT = 120 # Seconds to keep prompt context active
|
||||
PROMPT_TIMEOUT = 600 # Seconds to keep prompt context active (10 min for long generations)
|
||||
CLEANUP_INTERVAL = 300 # Seconds between cleanup of expired prompts
|
||||
|
||||
# Database settings
|
||||
|
||||
@@ -246,7 +246,7 @@ class GalleryDiagnostics:
|
||||
f.write("test")
|
||||
os.remove(test_file)
|
||||
can_write = True
|
||||
except:
|
||||
except (IOError, OSError):
|
||||
can_write = False
|
||||
|
||||
self.logger.info(f" [EDIT] Can write to directory: {can_write}")
|
||||
|
||||
+19
-2
@@ -61,8 +61,16 @@ class ImageGenerationHandler(FileSystemEventHandler):
|
||||
self.db_manager = db_manager
|
||||
self.prompt_tracker = prompt_tracker
|
||||
self.metadata_extractor = ComfyUIMetadataExtractor()
|
||||
self.processing_delay = 2.0 # Wait 2 seconds before processing
|
||||
self.logger = get_logger('prompt_manager.image_monitor')
|
||||
|
||||
# Read from GalleryConfig if available, otherwise use defaults
|
||||
try:
|
||||
from ..py.config import GalleryConfig
|
||||
self.processing_delay = GalleryConfig.PROCESSING_DELAY
|
||||
self.supported_extensions = tuple(GalleryConfig.SUPPORTED_EXTENSIONS)
|
||||
except Exception:
|
||||
self.processing_delay = 2.0
|
||||
self.supported_extensions = ('.png', '.jpg', '.jpeg', '.webp', '.gif')
|
||||
|
||||
def on_created(self, event):
|
||||
"""Handle filesystem creation events.
|
||||
@@ -96,7 +104,7 @@ class ImageGenerationHandler(FileSystemEventHandler):
|
||||
# Skip files in thumbnails directory - those are derivatives, not generated images
|
||||
if '/thumbnails/' in filepath or '\\thumbnails\\' in filepath:
|
||||
return False
|
||||
return filepath.lower().endswith(('.png', '.jpg', '.jpeg', '.webp', '.gif'))
|
||||
return filepath.lower().endswith(self.supported_extensions)
|
||||
|
||||
def process_new_image(self, image_path: str):
|
||||
"""Process a newly created image file for gallery integration.
|
||||
@@ -288,6 +296,15 @@ class ImageMonitor:
|
||||
self.logger.warning("Image monitoring already running")
|
||||
return
|
||||
|
||||
# Check if monitoring is enabled in config
|
||||
try:
|
||||
from ..py.config import GalleryConfig
|
||||
if not GalleryConfig.MONITORING_ENABLED:
|
||||
self.logger.info("Image monitoring disabled in config")
|
||||
return
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Check config first, then auto-detect if not configured
|
||||
if not output_directories:
|
||||
try:
|
||||
|
||||
@@ -9,6 +9,7 @@ This module provides centralized logging configuration with support for:
|
||||
- Log viewer API integration
|
||||
"""
|
||||
|
||||
import collections
|
||||
import logging
|
||||
import logging.handlers
|
||||
import os
|
||||
@@ -70,7 +71,7 @@ class PromptManagerLogger:
|
||||
}
|
||||
|
||||
# Memory buffer for recent logs (for web viewer)
|
||||
self._log_buffer = []
|
||||
self._log_buffer = collections.deque(maxlen=self.config['buffer_size'])
|
||||
self._buffer_lock = threading.Lock()
|
||||
|
||||
# Initialize loggers
|
||||
@@ -176,10 +177,6 @@ class PromptManagerLogger:
|
||||
}
|
||||
|
||||
self._log_buffer.append(log_entry)
|
||||
|
||||
# Maintain buffer size limit
|
||||
if len(self._log_buffer) > self.config['buffer_size']:
|
||||
self._log_buffer = self._log_buffer[-self.config['buffer_size']:]
|
||||
|
||||
def get_recent_logs(self, limit: int = 100, level: Optional[str] = None) -> List[Dict[str, Any]]:
|
||||
"""Get recent log entries from memory buffer.
|
||||
@@ -192,14 +189,14 @@ class PromptManagerLogger:
|
||||
List of log entry dictionaries, most recent first
|
||||
"""
|
||||
with self._buffer_lock:
|
||||
logs = self._log_buffer[:]
|
||||
|
||||
logs = list(self._log_buffer)
|
||||
|
||||
# Filter by level if specified
|
||||
if level:
|
||||
level_num = getattr(logging, level.upper(), None)
|
||||
if level_num:
|
||||
logs = [log for log in logs if getattr(logging, log['level']) >= level_num]
|
||||
|
||||
|
||||
# Return most recent first
|
||||
return list(reversed(logs[-limit:]))
|
||||
|
||||
|
||||
@@ -75,8 +75,15 @@ class PromptTracker:
|
||||
self._local = threading.local()
|
||||
self.active_prompts = {} # Global tracking for multiple threads
|
||||
self.lock = threading.Lock()
|
||||
self.cleanup_interval = 300 # 5 minutes
|
||||
self.prompt_timeout = 600 # 10 minutes (increased for longer generations)
|
||||
|
||||
# Read from GalleryConfig if available, otherwise use defaults
|
||||
try:
|
||||
from ..py.config import GalleryConfig
|
||||
self.cleanup_interval = GalleryConfig.CLEANUP_INTERVAL
|
||||
self.prompt_timeout = GalleryConfig.PROMPT_TIMEOUT
|
||||
except Exception:
|
||||
self.cleanup_interval = 300 # 5 minutes
|
||||
self.prompt_timeout = 600 # 10 minutes
|
||||
|
||||
# Start cleanup thread
|
||||
self.cleanup_thread = threading.Thread(target=self._cleanup_expired_prompts, daemon=True)
|
||||
|
||||
+8
-8
@@ -135,10 +135,10 @@ def validate_tags(tags: Union[str, List[str], None]) -> bool:
|
||||
|
||||
if len(tag.strip()) > 50:
|
||||
raise ValueError("Individual tags cannot exceed 50 characters")
|
||||
|
||||
# Check for invalid characters (optional - you can adjust this)
|
||||
if not re.match(r'^[a-zA-Z0-9\s\-_]+$', tag.strip()):
|
||||
raise ValueError(f"Tag '{tag}' contains invalid characters")
|
||||
|
||||
# Reject control characters and null bytes
|
||||
if re.search(r'[\x00-\x1f]', tag.strip()):
|
||||
raise ValueError(f"Tag '{tag}' contains invalid control characters")
|
||||
|
||||
if len(tags) > 20: # Reasonable limit
|
||||
raise ValueError("Maximum 20 tags allowed")
|
||||
@@ -177,10 +177,10 @@ def validate_category(category: Optional[str]) -> bool:
|
||||
|
||||
if len(category) > 100:
|
||||
raise ValueError("Category cannot exceed 100 characters")
|
||||
|
||||
# Check for invalid characters (adjust as needed)
|
||||
if not re.match(r'^[a-zA-Z0-9\s\-_]+$', category):
|
||||
raise ValueError("Category contains invalid characters")
|
||||
|
||||
# Reject control characters and null bytes
|
||||
if re.search(r'[\x00-\x1f]', category):
|
||||
raise ValueError("Category contains invalid control characters")
|
||||
|
||||
return True
|
||||
|
||||
|
||||
Reference in New Issue
Block a user