Vito Sansevero
ea7696bb9a
Re-wrap module docstrings, comments and f-string log messages; remove the unused typing.List and PngInfo imports; mark the PIL availability probe in diagnostics with noqa: F401; turn two placeholder-less f-strings into plain strings; strip trailing whitespace inside the diagnostics SQL.
221 lines
7.7 KiB
Python
221 lines
7.7 KiB
Python
"""ComfyUI integration utilities for PromptManager.
|
|
|
|
This module tracks the prompts PromptManager nodes encode during execution so
|
|
other components (e.g. image/prompt linking) can look them up.
|
|
|
|
ComfyUI already writes every node's own inputs into saved image metadata, so
|
|
SaveImage is deliberately NOT patched. An earlier patch wrote the last-run prompt
|
|
into every PromptManager node, saving negative prompts as copies of the
|
|
positive (#75).
|
|
|
|
Key features:
|
|
- Thread-safe prompt registration and retrieval
|
|
- Cross-thread prompt context sharing
|
|
- Standard metadata format compatibility
|
|
- Automatic cleanup of old prompt registrations
|
|
|
|
The integration works by:
|
|
1. PromptManager nodes register their prompts during execution
|
|
2. Callers retrieve the current prompt text by node id or recency
|
|
|
|
Typical usage:
|
|
from utils.comfyui_integration import get_comfyui_integration
|
|
|
|
integration = get_comfyui_integration()
|
|
integration.register_prompt(node_id, prompt_text, metadata)
|
|
text = integration.get_current_prompt_text(node_id)
|
|
"""
|
|
|
|
import threading
|
|
import time
|
|
from typing import Dict, Any, Optional
|
|
|
|
try:
|
|
from .logging_config import get_logger
|
|
except ImportError:
|
|
import sys
|
|
import os
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
from utils.logging_config import get_logger
|
|
|
|
|
|
class ComfyUIMetadataIntegration:
|
|
"""Integrates PromptManager with ComfyUI's standard metadata system.
|
|
|
|
This singleton class manages the integration between PromptManager custom nodes
|
|
and ComfyUI's standard metadata system. It ensures that prompts generated by
|
|
PromptManager appear in the standard ComfyUI metadata format for compatibility
|
|
with third-party tools and parsers.
|
|
|
|
The integration uses thread-local storage combined with global tracking to
|
|
handle prompt context across different execution threads, which is necessary
|
|
because ComfyUI executions may span multiple threads.
|
|
|
|
Key responsibilities:
|
|
- Register prompts from PromptManager nodes during execution
|
|
- Manage prompt lifecycle and cleanup
|
|
"""
|
|
|
|
_instance = None
|
|
_lock = threading.Lock()
|
|
|
|
def __new__(cls):
|
|
"""Ensure singleton pattern with thread safety.
|
|
|
|
Returns:
|
|
The single instance of ComfyUIMetadataIntegration
|
|
"""
|
|
if cls._instance is None:
|
|
with cls._lock:
|
|
if cls._instance is None:
|
|
cls._instance = super().__new__(cls)
|
|
return cls._instance
|
|
|
|
def __init__(self):
|
|
"""Initialize the ComfyUI integration system.
|
|
|
|
Sets up prompt tracking data structures. Uses _initialized flag to
|
|
prevent duplicate initialization.
|
|
"""
|
|
if hasattr(self, "_initialized"):
|
|
return
|
|
|
|
self.logger = get_logger("prompt_manager.comfyui_integration")
|
|
self._current_prompts = {}
|
|
self._thread_local = threading.local()
|
|
self._initialized = True
|
|
|
|
def register_prompt(self, node_id: str, prompt_text: str, metadata: Dict[str, Any]):
|
|
"""
|
|
Register a prompt from PromptManager for inclusion in ComfyUI metadata.
|
|
|
|
This method is called by PromptManager nodes during execution to register
|
|
their prompt text for later inclusion in image metadata. The prompt is stored
|
|
in both thread-local and global storage to handle cross-thread access scenarios.
|
|
|
|
Args:
|
|
node_id: Unique identifier for this prompt (typically the node ID)
|
|
prompt_text: The actual prompt text that was encoded by PromptManager
|
|
metadata: Additional metadata from PromptManager (category, tags, etc.)
|
|
"""
|
|
thread_id = threading.current_thread().ident
|
|
|
|
# Store in thread-local storage
|
|
if not hasattr(self._thread_local, "prompts"):
|
|
self._thread_local.prompts = {}
|
|
|
|
self._thread_local.prompts[node_id] = {
|
|
"text": prompt_text,
|
|
"metadata": metadata,
|
|
"timestamp": time.time(),
|
|
"thread_id": thread_id,
|
|
}
|
|
|
|
# Also store globally for cross-thread access
|
|
with self._lock:
|
|
self._current_prompts[f"{thread_id}_{node_id}"] = {
|
|
"text": prompt_text,
|
|
"metadata": metadata,
|
|
"timestamp": time.time(),
|
|
"thread_id": thread_id,
|
|
}
|
|
|
|
self.logger.debug(
|
|
f"Registered prompt for node {node_id}: {prompt_text[:50]}..."
|
|
)
|
|
|
|
def get_current_prompt_text(self, node_id: str = None) -> Optional[str]:
|
|
"""
|
|
Get the current prompt text for metadata inclusion.
|
|
|
|
Retrieves the most appropriate prompt text for inclusion in image metadata.
|
|
Uses a fallback strategy: thread-local storage first, then global storage,
|
|
with recency checks to avoid stale prompts.
|
|
|
|
Args:
|
|
node_id: Optional specific node ID to retrieve prompt for
|
|
|
|
Returns:
|
|
The prompt text string if available, None if no suitable prompt found
|
|
"""
|
|
# First try thread-local storage
|
|
if hasattr(self._thread_local, "prompts"):
|
|
if node_id and node_id in self._thread_local.prompts:
|
|
return self._thread_local.prompts[node_id]["text"]
|
|
elif self._thread_local.prompts:
|
|
# Return the most recent prompt from this thread
|
|
latest = max(
|
|
self._thread_local.prompts.values(), key=lambda x: x["timestamp"]
|
|
)
|
|
return latest["text"]
|
|
|
|
# Fallback to global storage
|
|
thread_id = threading.current_thread().ident
|
|
with self._lock:
|
|
# Look for prompts from current thread
|
|
thread_prompts = {
|
|
k: v
|
|
for k, v in self._current_prompts.items()
|
|
if v["thread_id"] == thread_id
|
|
}
|
|
|
|
if thread_prompts:
|
|
latest = max(thread_prompts.values(), key=lambda x: x["timestamp"])
|
|
return latest["text"]
|
|
|
|
# Last resort: return the most recent prompt from any thread
|
|
if self._current_prompts:
|
|
latest = max(
|
|
self._current_prompts.values(), key=lambda x: x["timestamp"]
|
|
)
|
|
# Only return if it's recent (within last 5 minutes)
|
|
if time.time() - latest["timestamp"] < 300:
|
|
return latest["text"]
|
|
|
|
return None
|
|
|
|
def cleanup_old_prompts(self, max_age_seconds: int = 600):
|
|
"""
|
|
Clean up old prompt registrations.
|
|
|
|
Removes prompt registrations that are older than the specified age to
|
|
prevent memory leaks and ensure that only recent, relevant prompts are
|
|
used for metadata inclusion.
|
|
|
|
Args:
|
|
max_age_seconds: Maximum age in seconds before a prompt registration
|
|
is considered stale and removed (default: 600
|
|
seconds/10 minutes)
|
|
"""
|
|
current_time = time.time()
|
|
|
|
with self._lock:
|
|
old_keys = [
|
|
key
|
|
for key, prompt_data in self._current_prompts.items()
|
|
if current_time - prompt_data["timestamp"] > max_age_seconds
|
|
]
|
|
|
|
for key in old_keys:
|
|
del self._current_prompts[key]
|
|
|
|
if old_keys:
|
|
self.logger.debug(f"Cleaned up {len(old_keys)} old prompt registrations")
|
|
|
|
|
|
# Global instance
|
|
_integration_instance = None
|
|
|
|
|
|
def get_comfyui_integration() -> ComfyUIMetadataIntegration:
|
|
"""Get the global ComfyUI integration instance.
|
|
|
|
Returns:
|
|
The singleton ComfyUIMetadataIntegration instance, creating it if necessary
|
|
"""
|
|
global _integration_instance
|
|
if _integration_instance is None:
|
|
_integration_instance = ComfyUIMetadataIntegration()
|
|
return _integration_instance
|