""" NodeValidator: Validates and corrects node names in workflow diagrams. Enhanced with GPU layer management, error handling, and ComfyUI integration. """ import json import logging import time import traceback from typing import Any import numpy as np import torch from sentence_transformers import SentenceTransformer from sklearn.metrics.pairwise import cosine_similarity from tinydb import Query, TinyDB try: import nodes as comfy_nodes COMFY_NODES_AVAILABLE = True except ImportError: COMFY_NODES_AVAILABLE = False logging.warning("ComfyUI nodes module not available, validation will use catalog only") from ..utils.model_loaders.device_utils import load_model from ..utils.process_utils import fetch_node_input_names, fetch_node_output_names, format, json_format class NODE_MISS_EXCEPTION(Exception): """Exception raised when a node type is not found in the database.""" pass class NodeBase: """ Base class for node information management and semantic search. Maintains a database of available ComfyUI nodes and provides semantic search capabilities for finding similar nodes. """ def __init__(self, cfg: Any): """ Initialize NodeBase with node database and optionally embedding model. Args: cfg: Configuration object with node_base settings """ logging.info("Initializing NodeBase...") node_list_path = cfg.node_base.node_list_path logging.debug(f"Loading node list from: {node_list_path}") try: with open(node_list_path) as f: self.node_list = json.load(f) logging.info(f"Loaded {len(self.node_list)} node types") except Exception as e: logging.error(f"Failed to load node list: {e}") raise node_info_db_path = cfg.node_base.node_info_db_path logging.debug(f"Loading node info database from: {node_info_db_path}") try: self.node_info_db = TinyDB(node_info_db_path) node_info_count = len(self.node_info_db.all()) logging.info(f"Loaded {node_info_count} node info entries") except (TypeError, ValueError) as e: # Handle case where file is in old format (plain JSON array instead of TinyDB format) logging.warning("Node info database appears to be in old format, attempting to convert...") try: with open(node_info_db_path, encoding="utf-8") as f: old_data = json.load(f) if isinstance(old_data, list): db = TinyDB(node_info_db_path) db.truncate() db.insert_multiple(old_data) db.close() logging.info(f"Converted {len(old_data)} entries from old format to TinyDB format") self.node_info_db = TinyDB(node_info_db_path) node_info_count = len(self.node_info_db.all()) logging.info(f"Loaded {node_info_count} node info entries") else: raise ValueError(f"Unexpected data format in {node_info_db_path}") except Exception as conv_e: logging.error(f"Failed to convert node info database: {conv_e}") logging.error("Please run 'UpdateNodeCatalog' node to regenerate the catalog files") raise ValueError( f"Node info database format is invalid. Please run 'UpdateNodeCatalog' node to regenerate it.\n" f"Original error: {e}\n" f"Conversion error: {conv_e}" ) from conv_e except Exception as e: logging.error(f"Failed to load node info database: {e}") raise self.cache_node_meta_info = {} # Load embedding model (optional - only needed for semantic search) embedding_model_path = getattr(cfg.node_base, "embedding_model_path", None) if embedding_model_path: logging.debug(f"Loading embedding model: {embedding_model_path}") try: device = "cuda" if torch.cuda.is_available() else "cpu" logging.debug(f"Using {device} for embedding model") self.embedding_model = SentenceTransformer(embedding_model_path, device=device) logging.debug("Computing embeddings for semantic search...") self.list_embeddings = self.embedding_model.encode(self.node_list, batch_size=32, show_progress_bar=False) logging.info(f"Embeddings computed for {len(self.list_embeddings)} nodes") except Exception as e: logging.error(f"Failed to load embedding model: {e}") raise else: self.embedding_model = None self.list_embeddings = None logging.info("Embedding model not loaded (semantic search disabled)") def fetch_node_meta_info(self, node_type: str) -> dict | None: """ Fetch metadata information for a node type. Args: node_type: Node type name Returns: Node metadata dictionary or None if not found Raises: NODE_MISS_EXCEPTION: If node type is not in the database """ if node_type not in self.node_list: logging.warning(f"Node type '{node_type}' not in node list") return None node_meta_info = self.cache_node_meta_info.get(node_type, None) if node_meta_info is not None: return node_meta_info query = Query() node_meta_info_list = self.node_info_db.search(query.node_type == node_type) if len(node_meta_info_list) == 0: logging.warning(f"Node type '{node_type}' not found in database") raise NODE_MISS_EXCEPTION(node_type) node_meta_info = node_meta_info_list[0] self.cache_node_meta_info[node_type] = node_meta_info return node_meta_info def find_most_similar(self, node_type: str, k: int = 5) -> list[dict]: """ Find the k most similar nodes to the given node type using semantic search. Args: node_type: Node type name to search for k: Number of similar nodes to return Returns: List of dictionaries containing node information Raises: ValueError: If embedding model is not loaded """ if self.embedding_model is None or self.list_embeddings is None: raise ValueError("Embedding model not loaded. Semantic search requires embedding_model_path to be provided.") query_embedding = self.embedding_model.encode([node_type]) similarities = cosine_similarity(query_embedding, self.list_embeddings)[0] top_k_indices = np.argsort(similarities)[-k:][::-1] similar_node_list = [self.node_list[i] for i in top_k_indices] similar_node_infos = [] for similar_node_type in similar_node_list: try: similar_node_infos.append(self.get_node_diagram_info(similar_node_type)) except Exception as e: logging.warning(f"Failed to get info for node '{similar_node_type}': {e}") continue return similar_node_infos def not_in_sql(self, node_type: str) -> bool: """ Check if a node type exists in the node list. Args: node_type: Node type name to check Returns: True if node exists, False otherwise """ return node_type in self.node_list def is_valid_comfy_node(self, node_type: str) -> bool: """ Check if a node type is a valid ComfyUI node class name. This checks against ComfyUI's actual NODE_CLASS_MAPPINGS, not just the catalog. The catalog may contain display names or aliases that aren't actual node class names. Args: node_type: Node type name to check Returns: True if node is a valid ComfyUI node class, False otherwise """ if not COMFY_NODES_AVAILABLE: # Fallback to catalog check if ComfyUI nodes aren't available return node_type in self.node_list # Check against ComfyUI's actual node registry return node_type in comfy_nodes.NODE_CLASS_MAPPINGS def get_node_diagram_info(self, node_type: str) -> dict: """ Get diagram information for a node type with enhanced metadata. Args: node_type: Node type name Returns: Dictionary with node_name, input_names, output_names, and additional metadata """ meta_info = self.fetch_node_meta_info(node_type) input_names = fetch_node_input_names(meta_info) output_names = fetch_node_output_names(meta_info) info = {"node_name": node_type, "input_names": input_names, "output_names": output_names} if meta_info: if "category" in meta_info: info["category"] = meta_info["category"] if "display_name" in meta_info: info["display_name"] = meta_info["display_name"] if "description" in meta_info: info["description"] = meta_info["description"] if "is_deprecated" in meta_info: info["is_deprecated"] = meta_info["is_deprecated"] if "is_experimental" in meta_info: info["is_experimental"] = meta_info["is_experimental"] if "is_output_node" in meta_info: info["is_output_node"] = meta_info["is_output_node"] if "is_api_node" in meta_info: info["is_api_node"] = meta_info["is_api_node"] if "input_metadata" in meta_info: info["input_metadata"] = meta_info["input_metadata"] if "output_display_names" in meta_info: info["output_display_names"] = meta_info["output_display_names"] if "output_tooltips" in meta_info: info["output_tooltips"] = meta_info["output_tooltips"] return info class NodeValidator: """ Validator responsible for validating and correcting node names in workflow diagrams. Uses semantic search and optionally a language model to fix invalid node names by finding the most similar valid node names. """ def __init__( self, node_base: NodeBase, cfg: Any, n_gpu_layers: int | str = "auto", dtype: str | torch.dtype = "auto", device_preference: str = "auto", attn_implementation: str = "auto", context_size: int = 4096, temperature: float = 0.95, top_p: float = 0.7, use_mmap: bool = True, use_mlock: bool = False, n_batch: int = 512, n_threads: int | None = None, ): """ Initialize NodeValidator. Args: node_base: NodeBase instance for node information cfg: Configuration object with node_validator settings n_gpu_layers: For GGUF: "auto", -1 for all, 0 for CPU, or specific count dtype: For HuggingFace: "auto", "fp16", "bf16", "fp32", "fp8", or torch.dtype device_preference: "auto", "cuda", or "cpu" """ logging.info("WorkflowGenerator NodeValidator: Initializing...") self.k = getattr(cfg.node_validator, "k", 5) self.node_base = node_base self.temperature = temperature self.top_p = top_p self.local = getattr(cfg.node_validator, "local", False) if self.local: self.max_new_tokens = getattr(cfg.node_validator, "max_new_tokens", 4096) model_format = getattr(cfg.node_validator, "model_format", "huggingface") model_path = cfg.node_validator.model_path tokenizer_path = getattr(cfg.node_validator, "tokenizer_path", None) logging.info(f"WorkflowGenerator NodeValidator: Loading local model: {model_format}, path: {model_path}") load_kwargs = {} if model_format.lower() == "huggingface": load_kwargs = { "dtype": dtype, "device_preference": device_preference, "attn_implementation": attn_implementation, } else: # GGUF models use n_gpu_layers, context_size, and performance parameters load_kwargs = { "n_gpu_layers": n_gpu_layers, "device_preference": device_preference, "context_size": context_size, "use_mmap": use_mmap, "use_mlock": use_mlock, "n_batch": n_batch, "n_threads": n_threads, } load_start = time.time() try: self.model_wrapper = load_model( model_format=model_format, model_path=model_path, tokenizer_path=tokenizer_path, **load_kwargs ) load_time = time.time() - load_start logging.info(f"WorkflowGenerator NodeValidator: Model loaded in {load_time:.2f}s") except Exception as e: logging.error(f"WorkflowGenerator NodeValidator: Failed to load model: {e}") raise self.model = self.model_wrapper.model self.tokenizer = self.model_wrapper.tokenizer try: if torch.cuda.is_available() and model_format.lower() == "huggingface": vram_used = torch.cuda.memory_allocated() / 1024**3 logging.debug(f"GPU memory allocated: {vram_used:.2f} GB") except Exception: pass else: logging.info("WorkflowGenerator NodeValidator: Using non-local mode (semantic search only)") def refine_diagram(self, diagram: list, prompt: str, seed: int | None = None) -> tuple[list, dict]: """ Refine a workflow diagram by validating and correcting node names. Args: diagram: List of edges representing the workflow diagram prompt: Original user prompt/description seed: Random seed for reproducible refinement (only used when local=True) Returns: Tuple of (refined_diagram, debug_info) where: - refined_diagram: Refined diagram with corrected node names - debug_info: Dictionary containing: - llm_prompts: List of prompts sent to LLM (one per invalid node) - candidate_nodes: Dictionary mapping invalid node names to their candidate nodes """ if seed is not None: max_seed = 2**32 - 1 if seed < 0 or seed > max_seed: seed = seed % (max_seed + 1) # Wrap around if out of range logging.info("Refining workflow diagram...") invalid_nodes = [] swap_map = {} debug_info = {"llm_prompts": [], "candidate_nodes": {}} logging.debug("Scanning diagram for invalid node names...") for edge in diagram: src_type, src_name, dst_type, dst_name = edge src_node = format(src_type) dst_node = format(dst_type) is_in_catalog = self.node_base.not_in_sql(src_node) is_valid_comfy = self.node_base.is_valid_comfy_node(src_node) if (not is_in_catalog or not is_valid_comfy) and src_node not in swap_map: invalid_nodes.append(src_node) if not is_in_catalog: logging.debug(f"Source node '{src_node}' not found in node database") elif not is_valid_comfy: logging.debug(f"Source node '{src_node}' is not a valid ComfyUI node class (may be display name/alias)") try: candidate_nodes = self.node_base.find_most_similar(src_node, self.k) debug_info["candidate_nodes"][src_node] = candidate_nodes except Exception as e: logging.warning(f"Failed to get candidate nodes for '{src_node}': {e}") debug_info["candidate_nodes"][src_node] = [] if self.local: swap_name, llm_prompt = self.fix_node_name_local(prompt, diagram, src_node, seed=seed, return_prompt=True) if llm_prompt: debug_info["llm_prompts"].append({"invalid_node": src_node, "prompt": llm_prompt}) else: swap_name = self.fix_node_name(prompt, diagram, src_node) if swap_name is not None: swap_map[src_node] = swap_name logging.info(f"Fixed node name: '{src_node}' -> '{swap_name}'") is_in_catalog = self.node_base.not_in_sql(dst_node) is_valid_comfy = self.node_base.is_valid_comfy_node(dst_node) if (not is_in_catalog or not is_valid_comfy) and dst_node not in swap_map: invalid_nodes.append(dst_node) if not is_in_catalog: logging.debug(f"Destination node '{dst_node}' not found in node database") elif not is_valid_comfy: logging.debug(f"Destination node '{dst_node}' is not a valid ComfyUI node class (may be display name/alias)") try: candidate_nodes = self.node_base.find_most_similar(dst_node, self.k) debug_info["candidate_nodes"][dst_node] = candidate_nodes except Exception as e: logging.warning(f"Failed to get candidate nodes for '{dst_node}': {e}") debug_info["candidate_nodes"][dst_node] = [] if self.local: swap_name, llm_prompt = self.fix_node_name_local(prompt, diagram, dst_node, seed=seed, return_prompt=True) if llm_prompt: debug_info["llm_prompts"].append({"invalid_node": dst_node, "prompt": llm_prompt}) else: swap_name = self.fix_node_name(prompt, diagram, dst_node) if swap_name is not None: swap_map[dst_node] = swap_name logging.info(f"Fixed node name: '{dst_node}' -> '{swap_name}'") logging.debug("Validating connections (input/output names)...") connection_issues = [] for i, edge in enumerate(diagram): src_type, src_name, dst_type, dst_name = edge src_node = format(src_type) dst_node = format(dst_type) if src_node in swap_map: src_node = swap_map[src_node] edge[0] = src_node if dst_node in swap_map: dst_node = swap_map[dst_node] edge[2] = dst_node try: src_meta = self.node_base.fetch_node_meta_info(src_node) if src_meta: src_output_names = fetch_node_output_names(src_meta) if src_name not in src_output_names: connection_issues.append( {"edge_index": i, "node": src_node, "type": "output", "name": src_name, "available": src_output_names} ) logging.debug(f"Source node '{src_node}' has no output '{src_name}'. Available: {src_output_names[:5]}...") except Exception as e: logging.warning(f"Failed to validate source node '{src_node}' output: {e}") try: dst_meta = self.node_base.fetch_node_meta_info(dst_node) if dst_meta: dst_input_names = fetch_node_input_names(dst_meta) if dst_name not in dst_input_names: connection_issues.append( {"edge_index": i, "node": dst_node, "type": "input", "name": dst_name, "available": dst_input_names} ) logging.debug(f"Destination node '{dst_node}' has no input '{dst_name}'. Available: {dst_input_names[:5]}...") if dst_node not in swap_map: try: candidate_nodes = self.node_base.find_most_similar(dst_node, self.k) debug_info["candidate_nodes"][dst_node] = candidate_nodes for candidate in candidate_nodes: candidate_meta = self.node_base.fetch_node_meta_info(candidate["node_name"]) if candidate_meta: candidate_inputs = fetch_node_input_names(candidate_meta) if dst_name in candidate_inputs: logging.debug(f"Candidate node '{candidate['node_name']}' has input '{dst_name}'") if self.local: connection_issue_info = {"type": "input", "name": dst_name, "available": dst_input_names} swap_name, llm_prompt = self.fix_node_name_local( prompt, diagram, dst_node, seed=seed, return_prompt=True, preferred_candidate=candidate["node_name"], connection_issue=connection_issue_info, ) if llm_prompt: debug_info["llm_prompts"].append( { "invalid_node": dst_node, "invalid_connection": f"input '{dst_name}'", "prompt": llm_prompt, } ) if swap_name and swap_name != dst_node: swap_map[dst_node] = swap_name logging.info( f"Fixed node name: '{dst_node}' -> '{swap_name}' (has required input '{dst_name}')" ) break else: swap_map[dst_node] = candidate["node_name"] logging.info( f"Fixed node name: '{dst_node}' -> '{candidate['node_name']}' (has required input '{dst_name}')" ) break except Exception as e: logging.warning(f"Failed to find candidate nodes for '{dst_node}': {e}") except Exception as e: logging.warning(f"Failed to validate destination node '{dst_node}' input: {e}") if len(swap_map) == 0 and len(connection_issues) == 0: logging.info("Diagram refinement complete: no corrections needed") else: if len(swap_map) > 0: logging.debug(f"Applying {len(swap_map)} node name corrections...") for i, edge in enumerate(diagram): src_type, src_name, dst_type, dst_name = edge src_node = format(src_type) dst_node = format(dst_type) if src_node in swap_map: swap_name = swap_map[src_node] edge[0] = edge[0].replace(src_node, swap_name) if dst_node in swap_map: swap_name = swap_map[dst_node] edge[2] = edge[2].replace(dst_node, swap_name) diagram[i] = edge if len(connection_issues) > 0: logging.warning(f"Found {len(connection_issues)} connection issues") logging.info(f"Diagram refinement complete: {len(swap_map)} node corrections, {len(connection_issues)} connection issues") return diagram, debug_info def fix_node_name(self, desc: str, diagram: list, error_name: str) -> str | None: """ Fix node name using semantic search only (non-local mode). Args: desc: User description diagram: Workflow diagram error_name: Invalid node name to fix Returns: Corrected node name or None if no suitable match found """ logging.info(f"Fixing node name '{error_name}' using semantic search") try: candidate_nodes = self.node_base.find_most_similar(error_name, self.k) if candidate_nodes: return candidate_nodes[0]["node_name"] except Exception as e: logging.warning(f"Failed to find similar node for '{error_name}': {e}") return None def fix_node_name_local( self, desc: str, diagram: list, error_name: str, seed: int | None = None, return_prompt: bool = False, preferred_candidate: str | None = None, connection_issue: dict | None = None, ) -> str | None | tuple[str | None, str | None]: """ Fix node name using local language model. Args: desc: User description diagram: Workflow diagram error_name: Invalid node name to fix seed: Random seed for reproducible generation (optional) return_prompt: If True, return tuple of (result, full_prompt_text), else just result preferred_candidate: Specific candidate to prioritize (e.g., one that has required input/output) connection_issue: Dictionary describing the connection problem, e.g.: {"type": "input", "name": "text_list", "available": ["text_a", "text_b", ...]} Returns: If return_prompt=False: Corrected node name or None if no suitable match found If return_prompt=True: Tuple of (corrected_node_name, full_prompt_text) where full_prompt_text is the complete prompt sent to LLM """ if not self.local: result = self.fix_node_name(desc, diagram, error_name) if return_prompt: return result, None return result logging.info(f"Fixing node name '{error_name}' using local model") try: candidate_nodes = self.node_base.find_most_similar(error_name, self.k) candidate_node_names = [e["node_name"] for e in candidate_nodes] if preferred_candidate and preferred_candidate in candidate_node_names: preferred_nodes = [] other_nodes = [] for c in candidate_nodes: if c["node_name"] == preferred_candidate: c_copy = c.copy() c_copy["_recommended"] = True if connection_issue: c_copy["_reason"] = f"Has required {connection_issue['type']}: '{connection_issue['name']}'" preferred_nodes.append(c_copy) else: other_nodes.append(c) candidate_nodes = preferred_nodes + other_nodes candidate_node_names = [e["node_name"] for e in candidate_nodes] if connection_issue: issue_type = connection_issue.get("type", "input") missing_name = connection_issue.get("name", "unknown") available_names = connection_issue.get("available", []) available_str = ", ".join([f"'{n}'" for n in available_names[:10]]) if len(available_names) > 10: available_str += ", ..." prompt = ( f"I would like you to act as an expert in ComfyUI platform. " f"I will provide a workflow description and a logical diagram in JSON format representing the ComfyUI workflow. " f"The logical diagram is a list of edges [edge_1, edge_2, ...], where each edge is " f"[source_node_name, output_name, destination_node_name, input_name], representing a connection from source to destination.\n\n" f"**Workflow Description:** {desc if desc else 'Not provided'}\n\n" f"**Logical Diagram:** {str(diagram)}\n\n" f"**PROBLEM DETECTED:**\n" f"The workflow uses node '{error_name}', but there is a CONNECTION ERROR:\n" f"- The workflow tries to connect to {issue_type} '{missing_name}'\n" f"- However, node '{error_name}' does NOT have this {issue_type}\n" f"- Available {issue_type}s on '{error_name}': {available_str}\n\n" f"**YOUR TASK:**\n" f"Select a replacement node from the candidates below that:\n" f"1. Has the required {issue_type} '{missing_name}' (CRITICAL)\n" f"2. Is semantically similar to '{error_name}'\n" f"3. Fits the workflow's purpose\n\n" f"**Candidate Nodes:** {str(candidate_nodes)}\n\n" f"**IMPORTANT NOTES:**\n" f"- Look for candidates with '_recommended': True - these have the required {issue_type}\n" f"- Check each candidate's '{issue_type}_names' list carefully\n" f"- The replacement MUST have '{missing_name}' in its {issue_type}s\n\n" f"**Response Format:**\n" f"Return ONLY the chosen node name in pure JSON format:\n" f"'''json{{\"candidate_node_name\":\"...\"}}'''" ) else: prompt = ( f"I would like you to act as an expert in ComfyUI platform. " f"I will provide a workflow description and a logical diagram in JSON format representing the ComfyUI workflow. " f"The logical diagram is a list of edges [edge_1, edge_2, ...], where each edge is " f"[source_node_name, output_name, destination_node_name, input_name], representing a connection from source to destination.\n\n" f"**Workflow Description:** {desc if desc else 'Not provided'}\n\n" f"**Logical Diagram:** {str(diagram)}\n\n" f"**PROBLEM DETECTED:**\n" f"The node name '{error_name}' is not valid in ComfyUI.\n\n" f"**YOUR TASK:**\n" f"Select the most suitable replacement node from the candidates below that:\n" f"1. Is semantically similar to '{error_name}'\n" f"2. Fits the workflow's purpose based on the description and diagram\n" f"3. Has compatible inputs/outputs for the connections in the diagram\n\n" f"**Candidate Nodes:** {str(candidate_nodes)}\n\n" f"**Response Format:**\n" f"Return ONLY the chosen node name in pure JSON format:\n" f"'''json{{\"candidate_node_name\":\"...\"}}'''" ) messages = [ {"role": "system", "content": "You are Qwen, created by Alibaba Cloud. You are a helpful assistant."}, {"role": "user", "content": prompt}, ] # Use HuggingFace tokenizer even for GGUF models (embedded tokenizer doesn't expose apply_chat_template) text = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) full_prompt_text = text if return_prompt else None model_inputs = self.tokenizer([text], return_tensors="pt") if hasattr(self.model_wrapper, "device") and isinstance(self.model_wrapper.device, torch.device): model_inputs = model_inputs.to(self.model_wrapper.device) generate_kwargs = {"max_new_tokens": self.max_new_tokens, "temperature": self.temperature, "top_p": self.top_p} if seed is not None: max_seed = 2**32 - 1 if seed < 0 or seed > max_seed: seed = seed % (max_seed + 1) # Wrap around if out of range generate_kwargs["seed"] = seed generated_ids = self.model_wrapper.generate(model_inputs.input_ids, **generate_kwargs) generated_ids = [output_ids[len(input_ids) :] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)] result = self.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0] try: result = json_format(result) if isinstance(result, dict): result = result.get("candidate_node_name", None) except Exception as e: logging.warning(f"Failed to parse JSON response: {e}") traceback.print_exc() if return_prompt: return None, full_prompt_text return None if result not in candidate_node_names: logging.warning(f"Model returned '{result}' which is not in candidate list") if return_prompt: return None, full_prompt_text return None if return_prompt: return result, full_prompt_text return result except Exception as e: logging.error(f"Failed to fix node name using local model: {e}") traceback.print_exc() if return_prompt: return None, None return None def cleanup(self): """Clean up model resources.""" if hasattr(self, "model_wrapper"): del self.model_wrapper if hasattr(self, "model"): del self.model if hasattr(self, "tokenizer"): del self.tokenizer if torch.cuda.is_available(): torch.cuda.empty_cache() logging.info("WorkflowGenerator NodeValidator: Resources cleaned up")