diff --git a/utils/comfyui_integration.py b/utils/comfyui_integration.py new file mode 100644 index 0000000..5d778a4 --- /dev/null +++ b/utils/comfyui_integration.py @@ -0,0 +1,228 @@ +""" +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 \ No newline at end of file