228 lines
9.4 KiB
Python
228 lines
9.4 KiB
Python
"""
|
|
ComfyUI integration utilities for PromptManager.
|
|
Provides hooks and patches to ensure PromptManager metadata appears in standard ComfyUI metadata.
|
|
"""
|
|
|
|
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."""
|
|
|
|
_instance = None
|
|
_lock = threading.Lock()
|
|
|
|
def __new__(cls):
|
|
if cls._instance is None:
|
|
with cls._lock:
|
|
if cls._instance is None:
|
|
cls._instance = super().__new__(cls)
|
|
return cls._instance
|
|
|
|
def __init__(self):
|
|
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.
|
|
|
|
Args:
|
|
node_id: The node ID (unique identifier for this prompt)
|
|
prompt_text: The actual prompt text that was encoded
|
|
metadata: Additional metadata from PromptManager
|
|
"""
|
|
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.
|
|
|
|
Args:
|
|
node_id: Optional specific node ID to get prompt for
|
|
|
|
Returns:
|
|
The prompt text or None
|
|
"""
|
|
# 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.
|
|
"""
|
|
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 fix them for standard parser compatibility
|
|
prompt_updated = False
|
|
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():
|
|
# Update the inputs to include our actual prompt text
|
|
if 'inputs' not in node_data:
|
|
node_data['inputs'] = {}
|
|
node_data['inputs']['text'] = current_prompt_text
|
|
|
|
# SIMPLE FIX: Change class_type to CLIPTextEncode for standard parser compatibility
|
|
# Keep original class_type in metadata for reference
|
|
if '_meta' not in node_data:
|
|
node_data['_meta'] = {}
|
|
node_data['_meta']['original_class_type'] = class_type
|
|
node_data['class_type'] = 'CLIPTextEncode'
|
|
|
|
prompt_updated = True
|
|
integration.logger.debug(f"Fixed PromptManager node {node_id} - changed class_type to CLIPTextEncode for compatibility")
|
|
|
|
# If no PromptManager nodes found, add a standalone one
|
|
if not prompt_updated:
|
|
virtual_node_id = "promptmanager_text"
|
|
prompt[virtual_node_id] = {
|
|
"class_type": "CLIPTextEncode", # Use CLIPTextEncode for compatibility
|
|
"inputs": {
|
|
"text": current_prompt_text
|
|
},
|
|
"_meta": {
|
|
"original_class_type": "PromptManager",
|
|
"virtual": True
|
|
}
|
|
}
|
|
integration.logger.debug("Added standalone CLIPTextEncode node with PromptManager 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.
|
|
|
|
Args:
|
|
max_age_seconds: Maximum age in seconds before cleanup
|
|
"""
|
|
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."""
|
|
global _integration_instance
|
|
if _integration_instance is None:
|
|
_integration_instance = ComfyUIMetadataIntegration()
|
|
return _integration_instance |