288 lines
12 KiB
Python
288 lines
12 KiB
Python
"""ComfyUI integration utilities for PromptManager.
|
|
|
|
This module provides deep integration with ComfyUI's metadata system to ensure that
|
|
PromptManager-generated prompts appear correctly in standard ComfyUI image metadata.
|
|
It patches core ComfyUI components to bridge the gap between PromptManager's custom
|
|
nodes and ComfyUI's standard metadata extraction.
|
|
|
|
Key features:
|
|
- Automatic patching of ComfyUI's SaveImage node
|
|
- 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. SaveImage node is patched to include registered prompts in metadata
|
|
3. Standard tools can then extract prompts from the generated images
|
|
|
|
Typical usage:
|
|
from utils.comfyui_integration import get_comfyui_integration
|
|
|
|
integration = get_comfyui_integration()
|
|
integration.register_prompt(node_id, prompt_text, metadata)
|
|
# Generated images will now include this prompt in their metadata
|
|
|
|
This integration is essential for:
|
|
- Third-party tool compatibility
|
|
- Standard metadata parsers
|
|
- Workflow sharing and reproduction
|
|
- Integration with existing ComfyUI ecosystems
|
|
"""
|
|
|
|
import threading
|
|
import time
|
|
import json
|
|
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
|
|
- Patch SaveImage to include PromptManager prompts in metadata
|
|
- 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 and attempts to patch the
|
|
SaveImage node. 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._saveimage_patched = False
|
|
self._initialized = True
|
|
|
|
# Try to patch SaveImage node on initialization
|
|
self._patch_saveimage_node()
|
|
|
|
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 _patch_saveimage_node(self):
|
|
"""
|
|
Patch ComfyUI's SaveImage node to include PromptManager prompts in metadata.
|
|
|
|
This method modifies ComfyUI's SaveImage.save_images method to automatically
|
|
include PromptManager prompts in the image metadata. The patching:
|
|
|
|
1. Wraps the original save_images method
|
|
2. Retrieves current PromptManager prompt text
|
|
3. Updates the text input in PromptManager nodes to reflect actual prompt
|
|
4. Calls the original method with the updated data
|
|
|
|
NOTE: We intentionally do NOT change class_type to CLIPTextEncode anymore.
|
|
That approach was corrupting saved workflows - when users saved and reloaded
|
|
workflows, ComfyUI would instantiate CLIPTextEncode instead of PromptManager,
|
|
causing errors with prepend_text/append_text inputs.
|
|
"""
|
|
try:
|
|
import nodes
|
|
|
|
if not hasattr(nodes, 'SaveImage'):
|
|
self.logger.warning("SaveImage node not found in ComfyUI nodes")
|
|
return
|
|
|
|
# Store original save_images method
|
|
original_save_images = nodes.SaveImage.save_images
|
|
integration = self # Capture self reference
|
|
|
|
def patched_save_images(self_node, images, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None):
|
|
"""Patched save_images method that includes PromptManager prompts."""
|
|
|
|
# Get current prompt text from PromptManager
|
|
current_prompt_text = integration.get_current_prompt_text()
|
|
|
|
if current_prompt_text:
|
|
integration.logger.debug(f"Including PromptManager prompt in SaveImage metadata: {current_prompt_text[:50]}...")
|
|
|
|
# If no prompt provided, create one with our text
|
|
if prompt is None:
|
|
prompt = {}
|
|
|
|
# Ensure prompt has the standard structure ComfyUI expects
|
|
if not isinstance(prompt, dict):
|
|
prompt = {}
|
|
|
|
# Find PromptManager nodes and ensure prompt text is captured
|
|
# NOTE: We do NOT change class_type anymore - that was corrupting saved workflows
|
|
# when users reload them. PromptManager stays as PromptManager.
|
|
for node_id, node_data in prompt.items():
|
|
if isinstance(node_data, dict):
|
|
class_type = node_data.get('class_type', '')
|
|
if 'promptmanager' in class_type.lower():
|
|
# Ensure the text input reflects the actual prompt used
|
|
if 'inputs' not in node_data:
|
|
node_data['inputs'] = {}
|
|
node_data['inputs']['text'] = current_prompt_text
|
|
integration.logger.debug(f"Updated PromptManager node {node_id} with prompt text")
|
|
|
|
# Call original method with potentially modified prompt
|
|
return original_save_images(self_node, images, filename_prefix, prompt, extra_pnginfo)
|
|
|
|
# Apply the patch
|
|
nodes.SaveImage.save_images = patched_save_images
|
|
self._saveimage_patched = True
|
|
self.logger.info("Successfully patched SaveImage node for PromptManager integration")
|
|
|
|
except Exception as e:
|
|
self.logger.error(f"Failed to patch SaveImage node: {e}")
|
|
self.logger.warning("PromptManager prompts may not appear in standard ComfyUI metadata")
|
|
|
|
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 |