commit 7b972d6dfa26dcd8fb18d01f88d1de86b1cea431 Author: unknown Date: Sat Nov 1 19:15:03 2025 +0100 chore: Rename comfyui-Model-Linker to comfyui-model-linker (kebab-case) diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..c75cae6 --- /dev/null +++ b/.gitignore @@ -0,0 +1,31 @@ +# Python +__pycache__/ +*.py[cod] +*$py.class +*.so +.Python + +# Virtual environments +venv/ +env/ +ENV/ + +# IDE +.vscode/ +.idea/ +*.swp +*.swo +*~ + +# OS +.DS_Store +Thumbs.db + +# Logs +*.log + +# Temporary files +*.tmp +*.bak +*.cache + diff --git a/README.md b/README.md new file mode 100644 index 0000000..b47b6aa --- /dev/null +++ b/README.md @@ -0,0 +1,42 @@ +# ComfyUI Model Linker Extension + +A ComfyUI extension that helps users relink missing models in workflows using fuzzy matching. + +## Features + +- Scans all nodes in workflows to find missing models +- Uses fuzzy matching to suggest similar model files +- Updates workflow JSON in UI/memory (user saves themselves) +- Supports all node types (not just specific ones) +- Optional auto-resolve for 100% confidence matches + +## Installation + +1. Clone or download this repository +2. Place it in your ComfyUI `custom_nodes/` directory +3. Restart ComfyUI + +## Usage + +1. Open a workflow with missing models +2. Click the "🔗 Model Linker" button in ComfyUI's top menu bar +3. Review missing models and their suggested matches (with confidence scores) +4. Click "Resolve" for individual models or "Auto-Resolve 100% Matches" for perfect matches +5. Save your workflow when ready + +## Features + +- **Subgraph Support**: Automatically detects and handles missing models inside subgraphs +- **Smart Matching**: Shows 100% confidence matches when available, otherwise shows best matches (≥70% confidence) +- **Fuzzy Matching**: Uses intelligent similarity scoring to find model files even with different naming +- **Auto-Resolve**: One-click resolution for all perfect matches + +## Requirements + +- Python 3.8+ +- ComfyUI + +## License + +See LICENSE file for details. + diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..177db98 --- /dev/null +++ b/__init__.py @@ -0,0 +1,152 @@ +""" +@author: Model Linker Team +@title: ComfyUI Model Linker +@nickname: Model Linker +@version: 1.0.0 +@description: Extension for relinking missing models in ComfyUI workflows using fuzzy matching +""" + +import logging + +# Web directory for JavaScript interface +WEB_DIRECTORY = "./web" + +# Empty NODE_CLASS_MAPPINGS - we don't provide custom nodes, only web extension +# This prevents ComfyUI from showing "IMPORT FAILED" message +NODE_CLASS_MAPPINGS = {} + +__all__ = ["WEB_DIRECTORY"] + + +class ModelLinkerExtension: + """Main extension class for Model Linker.""" + + def __init__(self): + self.routes_setup = False + self.logger = logging.getLogger(__name__) + + def initialize(self): + """Initialize the extension and set up API routes.""" + try: + self.setup_routes() + self.logger.info("Model Linker: Extension initialized successfully") + except Exception as e: + self.logger.error(f"Model Linker: Extension initialization failed: {e}", exc_info=True) + + def setup_routes(self): + """Register API routes for the Model Linker extension.""" + if self.routes_setup: + return # Already set up + + try: + from aiohttp import web + + # Try to get routes from PromptServer + try: + from server import PromptServer + if not hasattr(PromptServer, 'instance') or PromptServer.instance is None: + self.logger.debug("Model Linker: PromptServer not available yet") + return False + + routes = PromptServer.instance.routes + except (ImportError, AttributeError) as e: + self.logger.debug(f"Model Linker: Could not access PromptServer: {e}") + return False + + # Import linker modules - use relative imports which should work for packages + try: + from .core.linker import analyze_and_find_matches, apply_resolution + from .core.scanner import get_model_files + except ImportError as e: + self.logger.error(f"Model Linker: Could not import core modules: {e}") + return False + + @routes.post("/model_linker/analyze") + async def analyze_workflow(request): + """Analyze workflow and return missing models with matches.""" + try: + data = await request.json() + workflow_json = data.get('workflow') + + if not workflow_json: + return web.json_response( + {'error': 'Workflow JSON is required'}, + status=400 + ) + + # Analyze and find matches + result = analyze_and_find_matches(workflow_json) + + return web.json_response(result) + except Exception as e: + self.logger.error(f"Model Linker analyze error: {e}", exc_info=True) + return web.json_response( + {'error': str(e)}, + status=500 + ) + + @routes.post("/model_linker/resolve") + async def resolve_models(request): + """Apply model resolution and return updated workflow.""" + try: + data = await request.json() + workflow_json = data.get('workflow') + resolutions = data.get('resolutions', []) + + if not workflow_json: + return web.json_response( + {'error': 'Workflow JSON is required'}, + status=400 + ) + + if not resolutions: + return web.json_response( + {'error': 'Resolutions array is required'}, + status=400 + ) + + # Apply resolutions + updated_workflow = apply_resolution(workflow_json, resolutions) + + return web.json_response({ + 'workflow': updated_workflow, + 'success': True + }) + except Exception as e: + self.logger.error(f"Model Linker resolve error: {e}", exc_info=True) + return web.json_response( + {'error': str(e), 'success': False}, + status=500 + ) + + @routes.get("/model_linker/models") + async def get_models(request): + """Get list of all available models (for debugging/UI display).""" + try: + models = get_model_files() + return web.json_response(models) + except Exception as e: + self.logger.error(f"Model Linker get_models error: {e}", exc_info=True) + return web.json_response( + {'error': str(e)}, + status=500 + ) + + self.routes_setup = True + self.logger.info("Model Linker: API routes registered successfully") + return True + + except ImportError as e: + self.logger.warning(f"Model Linker: Could not register routes (missing dependency): {e}") + return False + except Exception as e: + self.logger.error(f"Model Linker: Error setting up routes: {e}", exc_info=True) + return False + + +# Initialize the extension +try: + extension = ModelLinkerExtension() + extension.initialize() +except Exception as e: + logging.error(f"ComfyUI Model Linker extension initialization failed: {e}", exc_info=True) diff --git a/__pycache__/__init__.cpython-313.pyc b/__pycache__/__init__.cpython-313.pyc new file mode 100644 index 0000000..e5a9320 Binary files /dev/null and b/__pycache__/__init__.cpython-313.pyc differ diff --git a/core/__init__.py b/core/__init__.py new file mode 100644 index 0000000..995e30b --- /dev/null +++ b/core/__init__.py @@ -0,0 +1,4 @@ +""" +Core modules for Model Linker extension. +""" + diff --git a/core/__pycache__/__init__.cpython-313.pyc b/core/__pycache__/__init__.cpython-313.pyc new file mode 100644 index 0000000..1b630ff Binary files /dev/null and b/core/__pycache__/__init__.cpython-313.pyc differ diff --git a/core/__pycache__/linker.cpython-313.pyc b/core/__pycache__/linker.cpython-313.pyc new file mode 100644 index 0000000..85ef379 Binary files /dev/null and b/core/__pycache__/linker.cpython-313.pyc differ diff --git a/core/__pycache__/matcher.cpython-313.pyc b/core/__pycache__/matcher.cpython-313.pyc new file mode 100644 index 0000000..574f027 Binary files /dev/null and b/core/__pycache__/matcher.cpython-313.pyc differ diff --git a/core/__pycache__/scanner.cpython-313.pyc b/core/__pycache__/scanner.cpython-313.pyc new file mode 100644 index 0000000..a4a8c09 Binary files /dev/null and b/core/__pycache__/scanner.cpython-313.pyc differ diff --git a/core/__pycache__/workflow_analyzer.cpython-313.pyc b/core/__pycache__/workflow_analyzer.cpython-313.pyc new file mode 100644 index 0000000..00c98ef Binary files /dev/null and b/core/__pycache__/workflow_analyzer.cpython-313.pyc differ diff --git a/core/__pycache__/workflow_updater.cpython-313.pyc b/core/__pycache__/workflow_updater.cpython-313.pyc new file mode 100644 index 0000000..bfe9e9c Binary files /dev/null and b/core/__pycache__/workflow_updater.cpython-313.pyc differ diff --git a/core/linker.py b/core/linker.py new file mode 100644 index 0000000..e87f976 --- /dev/null +++ b/core/linker.py @@ -0,0 +1,197 @@ +""" +Core Linker Module + +Integrates all components to provide high-level API for model linking. +""" + +import os +import logging +from typing import Dict, Any, List, Optional + +from .scanner import get_model_files +from .workflow_analyzer import analyze_workflow_models, identify_missing_models +from .matcher import find_matches +from .workflow_updater import update_workflow_nodes + + +def analyze_and_find_matches( + workflow_json: Dict[str, Any], + similarity_threshold: float = 0.0, + max_matches_per_model: int = 10 +) -> Dict[str, Any]: + """ + Main entry point: analyze workflow and find matches for missing models. + + Args: + workflow_json: Complete workflow JSON dictionary + similarity_threshold: Minimum similarity score (0.0 to 1.0) for matches + max_matches_per_model: Maximum number of matches to return per missing model + + Returns: + Dictionary with analysis results: + { + 'missing_models': [ + { + 'node_id': node ID, + 'node_type': node type, + 'widget_index': widget index, + 'original_path': original path from workflow, + 'category': model category, + 'matches': [ + { + 'model': model dict from scanner, + 'filename': model filename, + 'similarity': similarity score (0.0-1.0), + 'confidence': confidence percentage (0-100) + }, + ... + ] + }, + ... + ], + 'total_missing': count of missing models, + 'total_models_analyzed': count of all models in workflow + } + """ + # Analyze workflow to find all model references + all_model_refs = analyze_workflow_models(workflow_json) + + # Get available models + available_models = get_model_files() + + # Identify missing models + missing_models = identify_missing_models(all_model_refs, available_models) + + # Find matches for each missing model + missing_with_matches = [] + for missing in missing_models: + original_path = missing.get('original_path', '') + + # Filter available models by category if known + # IMPORTANT: If category is 'unknown', we still try to find the right category + # by using node type hints + category = missing.get('category') + + # If category is unknown, try to use node type to infer category + if not category or category == 'unknown': + from .workflow_analyzer import NODE_TYPE_TO_CATEGORY_HINTS + node_type = missing.get('node_type', '') + category = NODE_TYPE_TO_CATEGORY_HINTS.get(node_type, 'unknown') + + candidates = available_models + if category and category != 'unknown': + # Prioritize models from the same category + candidates = [m for m in available_models if m.get('category') == category] + # Also include other categories as fallback + candidates.extend([m for m in available_models if m.get('category') != category]) + + # Find matches + matches = find_matches( + original_path, + candidates, + threshold=similarity_threshold, + max_results=max_matches_per_model + ) + + # Deduplicate matches by absolute path - same physical file should only appear once + # This handles cases where the same file exists in multiple base directories + # or has different relative_paths but is the same file + seen_absolute_paths = {} + deduplicated_matches = [] + for match in matches: + model_dict = match['model'] + absolute_path = model_dict.get('path', '') + + # Normalize absolute path for comparison + if absolute_path: + absolute_path = os.path.normpath(absolute_path) + + # If we haven't seen this absolute path, add it + if absolute_path not in seen_absolute_paths: + seen_absolute_paths[absolute_path] = match + deduplicated_matches.append(match) + else: + # If we've seen this absolute path before, replace with better match if confidence is higher + existing_match = seen_absolute_paths[absolute_path] + if match['confidence'] > existing_match['confidence']: + # Replace with better match + idx = deduplicated_matches.index(existing_match) + deduplicated_matches[idx] = match + seen_absolute_paths[absolute_path] = match + + missing_with_matches.append({ + **missing, + 'matches': deduplicated_matches + }) + + return { + 'missing_models': missing_with_matches, + 'total_missing': len(missing_with_matches), + 'total_models_analyzed': len(all_model_refs) + } + + +def apply_resolution( + workflow_json: Dict[str, Any], + resolutions: List[Dict[str, Any]] +) -> Dict[str, Any]: + """ + Apply model resolutions to workflow. + + Args: + workflow_json: Workflow JSON dictionary (will be modified) + resolutions: List of resolution dictionaries: + { + 'node_id': node ID, + 'widget_index': widget index, + 'resolved_path': absolute path to resolved model, + 'category': model category (optional), + 'resolved_model': model dict from scanner (optional) + } + + Returns: + Updated workflow JSON dictionary + """ + # Prepare mappings for workflow_updater + mappings = [] + for resolution in resolutions: + mapping = { + 'node_id': resolution.get('node_id'), + 'widget_index': resolution.get('widget_index'), + 'resolved_path': resolution.get('resolved_path'), + 'category': resolution.get('category'), + 'resolved_model': resolution.get('resolved_model'), + 'subgraph_id': resolution.get('subgraph_id'), # Include subgraph_id for subgraph nodes + 'is_top_level': resolution.get('is_top_level') # True for top-level nodes, False for nodes in subgraph definitions + } + + # If resolved_model provided, extract path if needed + if 'resolved_model' in resolution and resolution['resolved_model']: + resolved_model = resolution['resolved_model'] + if 'path' in resolved_model and not mapping.get('resolved_path'): + mapping['resolved_path'] = resolved_model['path'] + if 'base_directory' in resolved_model: + mapping['base_directory'] = resolved_model['base_directory'] + + mappings.append(mapping) + + # Update workflow + updated_workflow = update_workflow_nodes(workflow_json, mappings) + + return updated_workflow + + +def get_resolution_summary(workflow_json: Dict[str, Any]) -> Dict[str, Any]: + """ + Get summary of missing models and matches without applying resolutions. + + This is a convenience method that calls analyze_and_find_matches with defaults. + + Args: + workflow_json: Complete workflow JSON dictionary + + Returns: + Same format as analyze_and_find_matches + """ + return analyze_and_find_matches(workflow_json) + diff --git a/core/matcher.py b/core/matcher.py new file mode 100644 index 0000000..35deea7 --- /dev/null +++ b/core/matcher.py @@ -0,0 +1,167 @@ +""" +Fuzzy Matcher Module + +Implements fuzzy string matching to find similar model names. +""" + +import os +import re +from typing import List, Dict, Tuple +from difflib import SequenceMatcher + + +def normalize_filename(filename: str) -> str: + """ + Normalize a filename for comparison. + + Removes file extension, converts to lowercase, and normalizes + separators (underscores, hyphens, spaces). + + Args: + filename: Filename to normalize + + Returns: + Normalized string for comparison + """ + # Remove file extension + base = os.path.splitext(filename)[0] + + # Convert to lowercase + base = base.lower() + + # Normalize separators: replace underscores, hyphens, and spaces with a single space + base = re.sub(r'[_\-\s]+', ' ', base) + + # Strip whitespace + base = base.strip() + + return base + + +def calculate_similarity(str1: str, str2: str) -> float: + """ + Calculate similarity score between two strings (0.0 to 1.0). + + Uses SequenceMatcher to compute a ratio. + + Args: + str1: First string + str2: Second string + + Returns: + Similarity score from 0.0 (completely different) to 1.0 (identical) + """ + return SequenceMatcher(None, str1, str2).ratio() + + +def calculate_similarity_with_normalization(str1: str, str2: str) -> float: + """ + Calculate similarity score with filename normalization. + + Normalizes both strings before comparing. + + Args: + str1: First string (typically model filename) + str2: Second string (typically candidate model filename) + + Returns: + Similarity score from 0.0 to 1.0 + """ + norm1 = normalize_filename(str1) + norm2 = normalize_filename(str2) + return calculate_similarity(norm1, norm2) + + +def find_matches( + target_model: str, + candidate_models: List[Dict[str, str]], + threshold: float = 0.0, + max_results: int = 10 +) -> List[Dict[str, any]]: + """ + Find similar models using fuzzy matching. + + Args: + target_model: The target model filename/path to match + candidate_models: List of candidate model dictionaries with 'filename' or 'path' key + threshold: Minimum similarity score (0.0 to 1.0) to include in results + max_results: Maximum number of results to return + + Returns: + List of match dictionaries sorted by similarity (highest first): + { + 'model': original model dict from candidates, + 'filename': model filename, + 'similarity': similarity score (0.0 to 1.0), + 'confidence': confidence percentage (0 to 100) + } + """ + matches = [] + + # Extract just the filename from target_model (remove any subfolder paths) + # target_model might be just a filename or might include subfolder paths + target_filename = os.path.basename(target_model) + + # Normalize target filename once for exact match comparisons + target_norm = normalize_filename(target_filename) + + for candidate in candidate_models: + # Get filename from candidate (prefer 'filename' key, fallback to extracting from 'path' or 'relative_path') + candidate_filename = candidate.get('filename') + + # If no filename key, try to extract from path or relative_path + if not candidate_filename: + candidate_path = candidate.get('path', '') or candidate.get('relative_path', '') + if candidate_path: + candidate_filename = os.path.basename(candidate_path) + + if not candidate_filename: + continue + + # Calculate similarity comparing just filenames (not paths) + # This ensures we're comparing apples to apples + + # First check for exact match (after normalization) - should be 100% + # Only exact matches should get 100% confidence + candidate_norm = normalize_filename(candidate_filename) + + if target_norm == candidate_norm: + # Exact match after normalization = 100% confidence + similarity = 1.0 + else: + # Calculate similarity using SequenceMatcher + # This gives a ratio between 0.0 and 1.0 based on longest common subsequence + similarity = calculate_similarity_with_normalization(target_filename, candidate_filename) + + # Also try comparing without extensions for better matching + target_base = os.path.splitext(target_filename)[0] + candidate_base = os.path.splitext(candidate_filename)[0] + similarity_no_ext = calculate_similarity_with_normalization(target_base, candidate_base) + + # Use the higher of the two similarity scores + # But ensure we never get 1.0 unless it's an exact normalized match + similarity = max(similarity, similarity_no_ext) + + # Cap similarity at 0.999 for non-exact matches to prevent false 100% scores + # SequenceMatcher can sometimes give 1.0 for very similar but not identical strings + # due to normalization artifacts + if similarity >= 0.999 and target_norm != candidate_norm: + similarity = 0.999 + + # Only include if above threshold + if similarity >= threshold: + matches.append({ + 'model': candidate, + 'filename': candidate_filename, + 'similarity': similarity, + 'confidence': round(similarity * 100, 1) # Convert to percentage + }) + + # Sort by similarity (highest first) + matches.sort(key=lambda x: x['similarity'], reverse=True) + + # Limit to max_results + matches = matches[:max_results] + + return matches + diff --git a/core/scanner.py b/core/scanner.py new file mode 100644 index 0000000..8bd6ad6 --- /dev/null +++ b/core/scanner.py @@ -0,0 +1,148 @@ +""" +Directory Scanner Module + +Scans configured model directories and finds available model files. +""" + +import os +import logging +from typing import List, Dict, Tuple + +# Import folder_paths lazily - it may not be available until ComfyUI is initialized +try: + import folder_paths +except ImportError: + folder_paths = None + logging.warning("Model Linker: folder_paths not available yet - will retry later") + +# Model file extensions to look for +# This matches folder_paths.supported_pt_extensions +MODEL_EXTENSIONS = {'.ckpt', '.pt', '.pt2', '.bin', '.pth', '.safetensors', '.pkl', '.sft', '.onnx'} + + +def get_model_directories() -> Dict[str, Tuple[List[str], set]]: + """ + Get all configured model directories from folder_paths. + + Returns: + Dictionary mapping category name to (list of paths, set of extensions) + """ + global folder_paths + + if folder_paths is None: + # Try to import again + try: + import folder_paths as fp + folder_paths = fp + except ImportError: + logging.error("Model Linker: folder_paths still not available") + return {} + + return folder_paths.folder_names_and_paths.copy() + + +def scan_directory(directory: str, extensions: set, category: str) -> List[Dict[str, str]]: + """ + Recursively scan a single directory for model files. + + Args: + directory: Absolute path to directory to scan + extensions: Set of file extensions to look for + category: Model category name (e.g., 'checkpoints', 'loras') + + Returns: + List of dictionaries with model information: + { + 'filename': 'model.safetensors', + 'path': 'absolute/path/to/model.safetensors', + 'relative_path': 'subfolder/model.safetensors' or 'model.safetensors', + 'category': 'checkpoints', + 'base_directory': 'absolute/path/to/base' + } + """ + models = [] + + if not os.path.exists(directory) or not os.path.isdir(directory): + logging.debug(f"Directory does not exist or is not accessible: {directory}") + return models + + try: + # Get absolute path and normalize + base_directory = os.path.abspath(directory) + + # Walk through directory recursively + for root, dirs, files in os.walk(base_directory, followlinks=True): + # Skip hidden directories + dirs[:] = [d for d in dirs if not d.startswith('.')] + + for filename in files: + # Check if file has a model extension + file_ext = os.path.splitext(filename)[1].lower() + + # For categories with empty extension set, accept all files + # Otherwise, check if extension matches + if len(extensions) == 0 or file_ext in extensions or file_ext in MODEL_EXTENSIONS: + full_path = os.path.join(root, filename) + + # Calculate relative path from base directory + # IMPORTANT: Use OS-native path separators (backslashes on Windows) + # This matches ComfyUI's recursive_search format for get_filename_list + try: + relative_path = os.path.relpath(full_path, base_directory) + # DO NOT normalize - keep OS-native separators to match ComfyUI + # ComfyUI's get_filename_list uses os.path.relpath which returns + # backslashes on Windows, forward slashes on Unix + except ValueError: + # If paths are on different drives (Windows), use filename only + relative_path = filename + + models.append({ + 'filename': filename, + 'path': full_path, + 'relative_path': relative_path, + 'category': category, + 'base_directory': base_directory + }) + except (OSError, PermissionError) as e: + logging.warning(f"Error scanning directory {directory}: {e}") + + return models + + +def scan_all_directories() -> List[Dict[str, str]]: + """ + Scan all configured model directories and return list of available models. + + Returns: + List of dictionaries with model information (same format as scan_directory) + """ + all_models = [] + directories = get_model_directories() + + for category, (paths, extensions) in directories.items(): + # Skip categories that aren't typically model directories + if category in ['custom_nodes', 'configs']: + continue + + for directory_path in paths: + try: + models = scan_directory(directory_path, extensions, category) + all_models.extend(models) + logging.debug(f"Found {len(models)} models in {category}/{directory_path}") + except Exception as e: + logging.warning(f"Error scanning {category} directory {directory_path}: {e}") + + return all_models + + +def get_model_files() -> List[Dict[str, str]]: + """ + Get list of all available model files with metadata. + + This is the main entry point for getting model files. + + Returns: + List of model dictionaries (same format as scan_directory) + """ + return scan_all_directories() + diff --git a/core/workflow_analyzer.py b/core/workflow_analyzer.py new file mode 100644 index 0000000..84d9c88 --- /dev/null +++ b/core/workflow_analyzer.py @@ -0,0 +1,266 @@ +""" +Workflow Analyzer Module + +Extracts model references from workflow JSON and identifies missing models. +""" + +import os +import logging +from typing import List, Dict, Any, Optional + +# Import folder_paths lazily - it may not be available until ComfyUI is initialized +try: + import folder_paths +except ImportError: + folder_paths = None + logging.warning("Model Linker: folder_paths not available yet - will retry later") + + +# Common model file extensions +MODEL_EXTENSIONS = {'.ckpt', '.pt', '.pt2', '.bin', '.pth', '.safetensors', '.pkl', '.sft', '.onnx'} + +# Mapping of common node types to their expected model category +# This is used as hints but we don't rely solely on this +# UNETLoader uses 'diffusion_models' category (folder_paths maps 'unet' to 'diffusion_models') +NODE_TYPE_TO_CATEGORY_HINTS = { + 'CheckpointLoaderSimple': 'checkpoints', + 'CheckpointLoader': 'checkpoints', + 'unCLIPCheckpointLoader': 'checkpoints', + 'VAELoader': 'vae', + 'LoraLoader': 'loras', + 'LoraLoaderModelOnly': 'loras', + 'UNETLoader': 'diffusion_models', # UNETLoader uses diffusion_models category + 'ControlNetLoader': 'controlnet', + 'ControlNetLoaderAdvanced': 'controlnet', + 'CLIPVisionLoader': 'clip_vision', + 'UpscaleModelLoader': 'upscale_models', + 'HypernetworkLoader': 'hypernetworks', + 'EmbeddingLoader': 'embeddings', +} + + +def is_model_filename(value: Any) -> bool: + """ + Check if a value looks like a model filename. + + Args: + value: The value to check + + Returns: + True if it looks like a model filename + """ + if not isinstance(value, str): + return False + + # Check if it ends with a model extension + _, ext = os.path.splitext(value.lower()) + return ext in MODEL_EXTENSIONS + + +def try_resolve_model_path(value: str, categories: List[str] = None) -> Optional[tuple[str, str]]: + """ + Try to resolve a model path using folder_paths. + + Args: + value: The model filename/path to resolve + categories: Optional list of categories to try (if None, tries all) + + Returns: + Tuple of (category, full_path) if found, None otherwise + """ + if not isinstance(value, str) or not value.strip(): + return None + + # Remove any path separators that might indicate an absolute path prefix + # Workflows should store relative paths, but handle both cases + filename = value.strip() + + # Ensure folder_paths is available + global folder_paths + if folder_paths is None: + try: + import folder_paths as fp + folder_paths = fp + except ImportError: + logging.error("Model Linker: folder_paths not available") + return None + + # If categories not provided, try all categories + if categories is None: + categories = list(folder_paths.folder_names_and_paths.keys()) + + # Skip non-model categories + skip_categories = {'custom_nodes', 'configs'} + categories = [c for c in categories if c not in skip_categories] + + for category in categories: + try: + full_path = folder_paths.get_full_path(category, filename) + if full_path and os.path.exists(full_path): + return (category, full_path) + except Exception: + continue + + return None + + +def get_node_model_info(node: Dict[str, Any]) -> List[Dict[str, Any]]: + """ + Extract model references from a single node. + + This scans all widgets_values entries and tries to identify which ones + are model file references by attempting to resolve them. + + Args: + node: Node dictionary from workflow JSON + + Returns: + List of model reference dictionaries: + { + 'node_id': node id, + 'node_type': node type, + 'widget_index': index in widgets_values, + 'original_path': original path from workflow, + 'category': model category (if found), + 'exists': True if model exists + } + """ + model_refs = [] + node_id = node.get('id') + node_type = node.get('type', '') + widgets_values = node.get('widgets_values', []) + + if not widgets_values: + return model_refs + + # Get category hints for this node type + category_hint = NODE_TYPE_TO_CATEGORY_HINTS.get(node_type) + categories_to_try = [category_hint] if category_hint else None + + # For each widget value, check if it looks like a model file + for idx, value in enumerate(widgets_values): + if not is_model_filename(value): + continue + + # Try to resolve the model path + resolved = try_resolve_model_path(value, categories_to_try) + + if resolved: + category, full_path = resolved + exists = os.path.exists(full_path) + else: + # If we can't resolve it, check if it at least looks like a model filename + # This might be a missing model or a custom node's model + category = category_hint or 'unknown' + full_path = None + exists = False + + model_refs.append({ + 'node_id': node_id, + 'node_type': node_type, + 'widget_index': idx, + 'original_path': value, + 'category': category, + 'full_path': full_path, + 'exists': exists + }) + + return model_refs + + +def analyze_workflow_models(workflow_json: Dict[str, Any]) -> List[Dict[str, Any]]: + """ + Extract all model references from a workflow, including nested subgraphs. + + Args: + workflow_json: Complete workflow JSON dictionary + + Returns: + List of model reference dictionaries (same format as get_node_model_info) + Each dict includes 'subgraph_id' if the model is in a subgraph + """ + all_model_refs = [] + + # Get subgraph definitions first to check if node types are subgraph UUIDs + definitions = workflow_json.get('definitions', {}) + subgraphs = definitions.get('subgraphs', []) + subgraph_lookup = {sg.get('id'): sg.get('name', sg.get('id')) for sg in subgraphs} + + # Analyze top-level nodes + nodes = workflow_json.get('nodes', []) + for node in nodes: + try: + model_refs = get_node_model_info(node) + node_type = node.get('type', '') + + # Check if node type is a subgraph UUID + subgraph_name = None + subgraph_id = None + if node_type in subgraph_lookup: + subgraph_name = subgraph_lookup[node_type] + subgraph_id = node_type + + # Mark with subgraph info if it's a subgraph node + # For top-level subgraph instance nodes, subgraph_path is None + # This distinguishes them from nodes within subgraph definitions + for ref in model_refs: + ref['subgraph_id'] = subgraph_id + ref['subgraph_name'] = subgraph_name + ref['subgraph_path'] = None # Top-level, not in definitions.subgraphs + ref['is_top_level'] = True # Flag to indicate this is a top-level node + all_model_refs.extend(model_refs) + except Exception as e: + logging.warning(f"Error analyzing node {node.get('id', 'unknown')}: {e}") + continue + + # Recursively analyze subgraphs (definitions already loaded above) + if not subgraphs: # Re-get if not loaded above + subgraphs = definitions.get('subgraphs', []) + + for subgraph in subgraphs: + subgraph_id = subgraph.get('id') + subgraph_name = subgraph.get('name', subgraph_id) + subgraph_nodes = subgraph.get('nodes', []) + + logging.debug(f"Analyzing subgraph: {subgraph_name} (ID: {subgraph_id}) with {len(subgraph_nodes)} nodes") + + for node in subgraph_nodes: + try: + model_refs = get_node_model_info(node) + # Mark as belonging to this subgraph definition + for ref in model_refs: + ref['subgraph_id'] = subgraph_id + ref['subgraph_name'] = subgraph_name + ref['subgraph_path'] = ['definitions', 'subgraphs', subgraph_id, 'nodes'] + ref['is_top_level'] = False # This is inside a subgraph definition + all_model_refs.extend(model_refs) + except Exception as e: + logging.warning(f"Error analyzing subgraph node {node.get('id', 'unknown')}: {e}") + continue + + return all_model_refs + + +def identify_missing_models( + workflow_models: List[Dict[str, Any]], + available_models: List[Dict[str, str]] = None +) -> List[Dict[str, Any]]: + """ + Identify which models from the workflow are missing. + + Args: + workflow_models: List of model references from analyze_workflow_models + available_models: Optional list of available models (if None, checks via folder_paths) + + Returns: + List of missing model references (filtered to only missing ones) + """ + missing = [] + + for model_ref in workflow_models: + # If exists is False, it's missing + if not model_ref.get('exists', False): + missing.append(model_ref) + + return missing + diff --git a/core/workflow_updater.py b/core/workflow_updater.py new file mode 100644 index 0000000..735f053 --- /dev/null +++ b/core/workflow_updater.py @@ -0,0 +1,288 @@ +""" +Workflow Updater Module + +Updates workflow JSON by replacing model paths in nodes. +""" + +import os +import logging +from typing import Dict, Any, List, Optional + + +def convert_to_relative_path(absolute_path: str, category: str, base_directory: str = None) -> str: + """ + Convert an absolute path to a relative path for workflow storage. + + This should match the format that ComfyUI's get_filename_list() returns, + which uses relative paths from the category base directory with forward slashes. + + Args: + absolute_path: Full absolute path to the model file + category: Model category (e.g., 'checkpoints', 'loras') + base_directory: Optional base directory for the category + + Returns: + Relative path (filename or subfolder/filename) suitable for workflow storage + This MUST match the format ComfyUI uses for validation + """ + if not absolute_path or not os.path.isabs(absolute_path): + # Already relative or empty - return as-is (keep OS-native separators) + # Don't normalize path separators - must match ComfyUI's format exactly + return absolute_path + + # Use folder_paths.get_filename_list to find the exact format ComfyUI expects + # CRITICAL: ComfyUI uses OS-native path separators (backslashes on Windows, forward slashes on Unix) + # We must return the EXACT format from get_filename_list, not a normalized version + try: + import folder_paths + # Get all available filenames for this category + # This returns paths with OS-native separators (backslashes on Windows) + available_filenames = folder_paths.get_filename_list(category) + + # Try to find a matching entry in ComfyUI's list + # Compare by finding the file that resolves to our absolute path + for filename in available_filenames: + try: + full_path = folder_paths.get_full_path(category, filename) + if full_path and os.path.normpath(full_path) == os.path.normpath(absolute_path): + # Found exact match - return ComfyUI's format EXACTLY as-is + # This includes OS-native path separators + return filename + except Exception: + continue + except Exception: + # Fall back to manual calculation if folder_paths not available + pass + + # If base_directory is provided, calculate relative to it + # IMPORTANT: Use OS-native path separators (don't normalize to forward slashes) + # ComfyUI expects paths with backslashes on Windows, forward slashes on Unix + if base_directory: + try: + relative_path = os.path.relpath(absolute_path, base_directory) + # DO NOT normalize path separators - use OS-native format + # This matches what ComfyUI's recursive_search returns + return relative_path + except ValueError: + # Paths are on different drives (Windows) or can't be relativized + # Fall back to just filename + pass + + # Fallback: return just the filename + return os.path.basename(absolute_path) + + +def get_base_directory_for_model(model_dict: Dict[str, str], category: str) -> Optional[str]: + """ + Get the base directory for a model based on its metadata. + + Args: + model_dict: Model dictionary with 'base_directory' or 'path' key + category: Model category + + Returns: + Base directory path if found, None otherwise + """ + # Try to get base_directory from model dict + if 'base_directory' in model_dict: + return model_dict['base_directory'] + + # If we have the full path, try to find the category base directory + if 'path' in model_dict: + full_path = model_dict['path'] + # Import here to avoid circular dependency + import folder_paths + + # Try to get category directories + if category in folder_paths.folder_names_and_paths: + category_paths = folder_paths.get_folder_paths(category) + # Find which base directory this path belongs to + for base_dir in category_paths: + try: + if os.path.commonpath([full_path, base_dir]) == base_dir: + return base_dir + except (ValueError, FileNotFoundError): + continue + + return None + + +def update_model_path( + workflow: Dict[str, Any], + node_id: int, + widget_index: int, + resolved_path: str, + category: str = None, + base_directory: str = None, + resolved_model: Dict[str, Any] = None, + subgraph_id: str = None, + is_top_level: bool = None +) -> bool: + """ + Update a single model path in a workflow node, supporting both top-level and subgraph nodes. + + Args: + workflow: Workflow JSON dictionary + node_id: ID of the node to update + widget_index: Index in widgets_values array to update + resolved_path: Absolute path to the resolved model + category: Model category (optional, for calculating relative path) + base_directory: Base directory for the category (optional) + resolved_model: Model dict from scanner (optional) + subgraph_id: ID of the subgraph (UUID for subgraph type, or None) + is_top_level: True if this is a top-level node (even if it's a subgraph instance), + False if it's inside a subgraph definition, None to auto-detect + + Returns: + True if update was successful, False otherwise + """ + node = None + + # Determine if this is a top-level node or inside a subgraph definition + # - If is_top_level is True, it's a top-level node (even if it's a subgraph instance) + # - If is_top_level is False, it's inside a subgraph definition + # - If is_top_level is None and subgraph_id is set, check if node exists in top-level first + search_in_subgraph = False + + if is_top_level is False: + # Explicitly inside a subgraph definition + search_in_subgraph = True + elif is_top_level is True: + # Explicitly a top-level node + search_in_subgraph = False + elif subgraph_id: + # Auto-detect: Check if node exists in top-level nodes first + # (Top-level subgraph instances have subgraph_id set but are in workflow.nodes) + nodes = workflow.get('nodes', []) + for n in nodes: + if n.get('id') == node_id: + # Found in top-level - this is a subgraph instance node + search_in_subgraph = False + break + else: + # Not found in top-level - must be inside subgraph definition + search_in_subgraph = True + else: + # No subgraph_id - definitely top-level + search_in_subgraph = False + + # Search for the node + if search_in_subgraph: + # Find in subgraph definition + definitions = workflow.get('definitions', {}) + subgraphs = definitions.get('subgraphs', []) + + for subgraph in subgraphs: + if subgraph.get('id') == subgraph_id: + subgraph_nodes = subgraph.get('nodes', []) + for n in subgraph_nodes: + if n.get('id') == node_id: + node = n + break + break + else: + # Find in top-level nodes + nodes = workflow.get('nodes', []) + for n in nodes: + if n.get('id') == node_id: + node = n + break + + if not node: + location = f"subgraph {subgraph_id}" if subgraph_id else "top-level" + logging.warning(f"Node {node_id} not found in {location}") + return False + + widgets_values = node.get('widgets_values', []) + + if widget_index >= len(widgets_values): + logging.warning(f"Widget index {widget_index} out of range for node {node_id}") + return False + + # Get category from resolved_model if not provided + if not category and resolved_model: + category = resolved_model.get('category') + + # Convert absolute path to relative path for workflow storage + # IMPORTANT: Use the category from resolved_model, not the original missing model category + # This ensures we use the correct category for validation + if os.path.isabs(resolved_path): + # Use category from resolved_model for path conversion + effective_category = category + if resolved_model: + effective_category = resolved_model.get('category', category) + + relative_path = convert_to_relative_path(resolved_path, effective_category, base_directory) + else: + relative_path = resolved_path + + # Update the widget value + widgets_values[widget_index] = relative_path + + logging.debug(f"Updated node {node_id}, widget {widget_index} to: {relative_path}") + return True + + +def update_workflow_nodes( + workflow: Dict[str, Any], + mappings: List[Dict[str, Any]] +) -> Dict[str, Any]: + """ + Apply multiple model path changes to a workflow. + + Args: + workflow: Workflow JSON dictionary (will be modified in place) + mappings: List of mapping dictionaries: + { + 'node_id': node ID, + 'widget_index': widget index, + 'resolved_path': absolute path to resolved model, + 'category': model category (optional), + 'base_directory': base directory for category (optional), + 'resolved_model': model dict from scanner (optional, for base_directory) + } + + Returns: + Updated workflow dictionary (same reference, modified in place) + """ + updated_count = 0 + + for mapping in mappings: + node_id = mapping.get('node_id') + widget_index = mapping.get('widget_index') + resolved_path = mapping.get('resolved_path') + + if not all([node_id is not None, widget_index is not None, resolved_path]): + logging.warning(f"Invalid mapping: {mapping}") + continue + + # Try to get base_directory from resolved_model if provided + base_directory = mapping.get('base_directory') + if not base_directory and 'resolved_model' in mapping: + resolved_model = mapping['resolved_model'] + category = mapping.get('category', '') + base_directory = get_base_directory_for_model(resolved_model, category) + + category = mapping.get('category') + resolved_model = mapping.get('resolved_model') + subgraph_id = mapping.get('subgraph_id') + is_top_level = mapping.get('is_top_level') # True for top-level nodes, False for nodes in subgraph definitions + + success = update_model_path( + workflow, + node_id, + widget_index, + resolved_path, + category, + base_directory, + resolved_model, + subgraph_id, + is_top_level + ) + + if success: + updated_count += 1 + + logging.info(f"Updated {updated_count} model paths in workflow") + return workflow + diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..d89908b --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,7 @@ +[project] +name = "comfyui-model-linker" +version = "0.1.0" +description = "A ComfyUI extension that helps users relink missing models in workflows" +readme = "README.md" +requires-python = ">=3.8" + diff --git a/web/linker.js b/web/linker.js new file mode 100644 index 0000000..e4d0e01 --- /dev/null +++ b/web/linker.js @@ -0,0 +1,838 @@ +/** + * ComfyUI Model Linker Extension - Frontend + * + * Provides a menu button and dialog interface for relinking missing models in workflows. + */ + +// Import ComfyUI APIs +// These paths are relative to the ComfyUI web directory +import { app } from "../../../scripts/app.js"; +import { api } from "../../../scripts/api.js"; +import { $el, ComfyDialog } from "../../../scripts/ui.js"; + +// Check if ComfyButtonGroup is available (from newer ComfyUI versions) +let ComfyButtonGroup = null; +try { + // Try to import from scripts if available + if (typeof window !== 'undefined') { + try { + // Some ComfyUI versions expose this globally + if (window.ComfyButtonGroup) { + ComfyButtonGroup = window.ComfyButtonGroup; + } + } catch (e) { + // Ignore + } + } +} catch (e) { + // Fallback if ComfyButtonGroup not available +} + +class LinkerManagerDialog extends ComfyDialog { + constructor() { + super(); + this.currentWorkflow = null; + this.missingModels = []; + + // Create dialog element using $el + this.element = $el("div.comfy-modal", { + parent: document.body, + style: { + position: "fixed", + top: "50%", + left: "50%", + transform: "translate(-50%, -50%)", + width: "900px", + height: "700px", + maxWidth: "95vw", + maxHeight: "95vh", + backgroundColor: "var(--comfy-menu-bg, #202020)", + color: "var(--input-text, #ffffff)", + border: "2px solid var(--border-color, #555555)", + borderRadius: "8px", + padding: "0", + zIndex: "99999", + boxShadow: "0 4px 20px rgba(0,0,0,0.8)", + display: "none", + flexDirection: "column" + } + }, [ + this.createHeader(), + this.createContent(), + this.createFooter() + ]); + } + + createHeader() { + return $el("div", { + style: { + display: "flex", + justifyContent: "space-between", + alignItems: "center", + padding: "20px 20px 10px 20px", + borderBottom: "1px solid var(--border-color)", + backgroundColor: "var(--comfy-menu-bg, #202020)" + } + }, [ + $el("h2", { + textContent: "🔗 Model Linker", + style: { + margin: "0", + color: "var(--input-text)", + fontSize: "18px", + fontWeight: "600" + } + }), + $el("button", { + textContent: "×", + onclick: () => this.close(), + style: { + background: "none", + border: "none", + fontSize: "24px", + cursor: "pointer", + color: "var(--input-text)", + padding: "0", + width: "30px", + height: "30px", + borderRadius: "4px", + display: "flex", + alignItems: "center", + justifyContent: "center" + } + }) + ]); + } + + createContent() { + this.contentElement = $el("div", { + id: "model-linker-content", + style: { + padding: "16px", + overflowY: "auto", + flex: "1", + minHeight: "0" + } + }); + return this.contentElement; + } + + createFooter() { + return $el("div", { + style: { + padding: "16px", + borderTop: "1px solid var(--border-color)", + display: "flex", + justifyContent: "flex-end", + gap: "8px" + } + }, [ + $el("button", { + textContent: "Auto-Resolve 100% Matches", + onclick: () => this.autoResolve100Percent(), + className: "comfy-button", + style: { + padding: "8px 16px" + } + }) + ]); + } + + async show() { + this.element.style.display = "flex"; + await this.loadWorkflowData(); + } + + close() { + this.element.style.display = "none"; + } + + /** + * Load workflow data and display missing models + */ + async loadWorkflowData(workflow = null) { + if (!this.contentElement) return; + + // Show loading state + this.contentElement.innerHTML = '

Analyzing workflow...

'; + + try { + // Use provided workflow, or get current workflow from ComfyUI + if (!workflow) { + workflow = this.getCurrentWorkflow(); + } + + if (!workflow) { + this.contentElement.innerHTML = '

No workflow loaded. Please load a workflow first.

'; + return; + } + + // Call analyze endpoint + const response = await api.fetchApi('/model_linker/analyze', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ workflow }) + }); + + if (!response.ok) { + throw new Error(`API error: ${response.status}`); + } + + const data = await response.json(); + this.displayMissingModels(this.contentElement, data); + + } catch (error) { + console.error('Model Linker: Error loading workflow data:', error); + if (this.contentElement) { + this.contentElement.innerHTML = `

Error: ${error.message}

`; + } + } + } + + /** + * Get current workflow from ComfyUI + */ + getCurrentWorkflow() { + // Try to get workflow from app + if (app?.graph) { + try { + // Use ComfyUI's workflow serialization + const workflow = app.graph.serialize(); + return workflow; + } catch (e) { + console.warn('Model Linker: Could not serialize workflow from graph:', e); + } + } + return null; + } + + /** + * Display missing models in the dialog + */ + displayMissingModels(container, data) { + const missingModels = data.missing_models || []; + const totalMissing = data.total_missing || 0; + + if (totalMissing === 0) { + container.innerHTML = '

✓ No missing models found. All models are available!

'; + return; + } + + let html = `

Found ${totalMissing} missing model(s):

`; + html += '
'; + + // Sort missing models: those with 100% confidence matches first, then others + const sortedMissingModels = missingModels.sort((a, b) => { + const aMatches = a.matches || []; + const bMatches = b.matches || []; + + // Filter to 70%+ confidence + const aFiltered = aMatches.filter(m => m.confidence >= 70); + const bFiltered = bMatches.filter(m => m.confidence >= 70); + + // Check if they have 100% matches + const aHas100 = aFiltered.some(m => m.confidence === 100); + const bHas100 = bFiltered.some(m => m.confidence === 100); + + // If one has 100% and the other doesn't, prioritize the one with 100% + if (aHas100 && !bHas100) return -1; + if (!aHas100 && bHas100) return 1; + + // If both have 100% or neither has 100%, sort by best confidence + const aBestConf = aFiltered.length > 0 ? Math.max(...aFiltered.map(m => m.confidence)) : 0; + const bBestConf = bFiltered.length > 0 ? Math.max(...bFiltered.map(m => m.confidence)) : 0; + + return bBestConf - aBestConf; // Higher confidence first + }); + + for (const missing of sortedMissingModels) { + html += this.renderMissingModel(missing); + } + + html += '
'; + container.innerHTML = html; + + // Attach event listeners for resolve buttons (use sorted order) + sortedMissingModels.forEach((missing, missingIndex) => { + const allMatches = missing.matches || []; + + // Filter out matches below 70% confidence threshold + const filteredMatches = allMatches.filter(m => m.confidence >= 70); + + // Filter to only 100% matches if available, otherwise use filtered matches (>=70%) + const perfectMatches = filteredMatches.filter(m => m.confidence === 100); + const matchesToUse = perfectMatches.length > 0 ? perfectMatches : filteredMatches; + + matchesToUse.forEach((match, matchIndex) => { + const buttonId = `resolve-${missing.node_id}-${missing.widget_index}-${matchIndex}`; + const resolveButton = container.querySelector(`#${buttonId}`); + if (resolveButton) { + resolveButton.addEventListener('click', () => { + this.resolveModel(missing, match.model); + }); + } + }); + }); + } + + /** + * Render a single missing model entry + */ + renderMissingModel(missing) { + const allMatches = missing.matches || []; + + // Filter out matches below 70% confidence threshold + const filteredMatches = allMatches.filter(m => m.confidence >= 70); + const hasMatches = filteredMatches.length > 0; + + let html = `
`; + + // Display subgraph name as primary identifier if available, otherwise show node type + // A node type that's a UUID indicates it's a subgraph instance + const isSubgraphNode = missing.node_type && missing.node_type.match(/^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/i); + + if (missing.subgraph_name) { + // Show subgraph name as primary identifier + html += `
Subgraph: ${missing.subgraph_name} (ID: ${missing.node_id})
`; + } else if (isSubgraphNode) { + // Node type is a UUID (subgraph) but we don't have the name (shouldn't happen, but handle gracefully) + html += `
Node: Subgraph (ID: ${missing.node_id})
`; + } else { + // Regular node + html += `
Node: ${missing.node_type} (ID: ${missing.node_id})
`; + } + html += `
Missing Model: ${missing.original_path}
`; + html += `
Category: ${missing.category || 'unknown'}
`; + + if (hasMatches) { + // Filter out matches below 70% confidence threshold + const filteredMatches = allMatches.filter(m => m.confidence >= 70); + + // Separate 100% matches from others (from filtered list) + const perfectMatches = filteredMatches.filter(m => m.confidence === 100); + const otherMatches = filteredMatches.filter(m => m.confidence < 100 && m.confidence >= 70); + + // If we have 100% matches, only show those. Otherwise, show other matches sorted by confidence + const matchesToShow = perfectMatches.length > 0 + ? perfectMatches + : otherMatches.sort((a, b) => b.confidence - a.confidence).slice(0, 5); + + html += `
Suggested Matches:
`; + html += ''; + + // Add note if only showing 100% matches + if (perfectMatches.length > 0 && otherMatches.length > 0) { + html += `
Showing only 100% confidence matches. ${otherMatches.length} other match${otherMatches.length > 1 ? 'es' : ''} available.
`; + } + } else if (allMatches.length > 0 && filteredMatches.length === 0) { + // Had matches but all were below 70% threshold + html += `
No matches found above 70% confidence threshold.
`; + } else { + html += `
No matches found.
`; + } + + html += '
'; + return html; + } + + /** + * Show a notification banner (similar to ComfyUI's "Reconnecting" banner) + */ + showNotification(message, type = 'success') { + // Create notification banner + const notification = $el("div", { + style: { + position: "fixed", + top: "0", + left: "50%", + transform: "translateX(-50%)", + backgroundColor: type === 'success' ? '#28a745' : type === 'error' ? '#dc3545' : '#007acc', + color: "#ffffff", + padding: "12px 24px", + borderRadius: "0 0 8px 8px", + fontSize: "14px", + fontWeight: "500", + zIndex: "100000", + boxShadow: "0 4px 12px rgba(0,0,0,0.4)", + display: "flex", + alignItems: "center", + gap: "12px", + minWidth: "300px", + maxWidth: "600px", + textAlign: "center", + animation: "slideDown 0.3s ease" + } + }, [ + type === 'success' ? $el("span", { + textContent: "✓", + style: { + fontSize: "18px", + fontWeight: "bold" + } + }) : type === 'error' ? $el("span", { + textContent: "×", + style: { + fontSize: "18px", + fontWeight: "bold" + } + }) : null, + $el("span", { + textContent: message + }), + $el("button", { + textContent: "×", + onclick: () => { + if (notification.parentNode) { + notification.style.opacity = "0"; + notification.style.transform = "translateX(-50%) translateY(-100%)"; + setTimeout(() => { + if (notification.parentNode) { + notification.parentNode.removeChild(notification); + } + }, 300); + } + }, + style: { + background: "none", + border: "none", + color: "#ffffff", + fontSize: "20px", + cursor: "pointer", + padding: "0", + marginLeft: "auto", + opacity: "0.8", + width: "24px", + height: "24px", + display: "flex", + alignItems: "center", + justifyContent: "center", + borderRadius: "4px" + } + }) + ]); + + // Add CSS animation if not already added + if (!document.getElementById('model-linker-notification-style')) { + const style = $el("style", { + id: 'model-linker-notification-style', + textContent: ` + @keyframes slideDown { + from { + opacity: 0; + transform: translateX(-50%) translateY(-100%); + } + to { + opacity: 1; + transform: translateX(-50%) translateY(0); + } + } + ` + }); + document.head.appendChild(style); + } + + document.body.appendChild(notification); + + // Auto-dismiss after 4 seconds for success, 6 seconds for errors + const dismissTime = type === 'success' ? 4000 : 6000; + setTimeout(() => { + if (notification.parentNode) { + notification.style.opacity = "0"; + notification.style.transform = "translateX(-50%) translateY(-100%)"; + setTimeout(() => { + if (notification.parentNode) { + notification.parentNode.removeChild(notification); + } + }, 300); + } + }, dismissTime); + } + + /** + * Resolve a single model + */ + async resolveModel(missing, resolvedModel) { + if (!resolvedModel) { + this.showNotification('No resolved model selected', 'error'); + return; + } + + try { + const workflow = this.getCurrentWorkflow(); + if (!workflow) { + this.showNotification('No workflow loaded', 'error'); + return; + } + + const resolution = { + node_id: missing.node_id, + widget_index: missing.widget_index, + resolved_path: resolvedModel.path, + category: missing.category, + resolved_model: resolvedModel, + subgraph_id: missing.subgraph_id, // Include subgraph_id for subgraph nodes + is_top_level: missing.is_top_level // True for top-level nodes, False for nodes in subgraph definitions + }; + + const response = await api.fetchApi('/model_linker/resolve', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + workflow, + resolutions: [resolution] + }) + }); + + if (!response.ok) { + throw new Error(`API error: ${response.status}`); + } + + const data = await response.json(); + + if (data.success) { + // Update workflow in ComfyUI + await this.updateWorkflowInComfyUI(data.workflow); + + // Show success notification + const modelName = resolvedModel.relative_path || resolvedModel.filename || 'model'; + this.showNotification(`✓ Model linked successfully: ${modelName}`, 'success'); + + // Reload dialog using the updated workflow from API response + // This ensures we're analyzing the correct updated workflow + await this.loadWorkflowData(data.workflow); + } else { + this.showNotification('Failed to resolve model: ' + (data.error || 'Unknown error'), 'error'); + } + + } catch (error) { + console.error('Model Linker: Error resolving model:', error); + this.showNotification('Error resolving model: ' + error.message, 'error'); + } + } + + /** + * Auto-resolve all 100% confidence matches + */ + async autoResolve100Percent() { + if (!this.contentElement) return; + + try { + const workflow = this.getCurrentWorkflow(); + if (!workflow) { + this.showNotification('No workflow loaded', 'error'); + return; + } + + // Analyze workflow first + const analyzeResponse = await api.fetchApi('/model_linker/analyze', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ workflow }) + }); + + if (!analyzeResponse.ok) { + throw new Error(`API error: ${analyzeResponse.status}`); + } + + const analyzeData = await analyzeResponse.json(); + const missingModels = analyzeData.missing_models || []; + + // Collect all 100% matches + const resolutions = []; + for (const missing of missingModels) { + const matches = missing.matches || []; + const perfectMatch = matches.find((m) => m.confidence === 100); + + if (perfectMatch && perfectMatch.model) { + resolutions.push({ + node_id: missing.node_id, + widget_index: missing.widget_index, + resolved_path: perfectMatch.model.path, + category: missing.category, + resolved_model: perfectMatch.model, + subgraph_id: missing.subgraph_id, // Include subgraph_id for subgraph nodes + is_top_level: missing.is_top_level // True for top-level nodes, False for nodes in subgraph definitions + }); + } + } + + if (resolutions.length === 0) { + this.showNotification('No 100% confidence matches found to auto-resolve.', 'error'); + return; + } + + // Apply resolutions + const resolveResponse = await api.fetchApi('/model_linker/resolve', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + workflow, + resolutions + }) + }); + + if (!resolveResponse.ok) { + throw new Error(`API error: ${resolveResponse.status}`); + } + + const resolveData = await resolveResponse.json(); + + if (resolveData.success) { + // Update workflow in ComfyUI + await this.updateWorkflowInComfyUI(resolveData.workflow); + + // Show success notification + this.showNotification( + `✓ Successfully linked ${resolutions.length} model${resolutions.length > 1 ? 's' : ''}!`, + 'success' + ); + + // Reload dialog using the updated workflow from API response + // This ensures we're analyzing the correct updated workflow + await this.loadWorkflowData(resolveData.workflow); + } else { + this.showNotification('Failed to resolve models: ' + (resolveData.error || 'Unknown error'), 'error'); + } + + } catch (error) { + console.error('Model Linker: Error auto-resolving:', error); + this.showNotification('Error auto-resolving: ' + error.message, 'error'); + } + } + + /** + * Update workflow in ComfyUI's UI/memory + */ + async updateWorkflowInComfyUI(workflow) { + if (app && app.loadGraphData) { + try { + await app.loadGraphData(workflow, false, false); + } catch (error) { + console.error('Model Linker: Error updating workflow in ComfyUI:', error); + throw error; + } + } else { + console.warn('Model Linker: Could not update workflow - app.loadGraphData not available'); + } + } +} + +// Main extension class +class ModelLinker { + constructor() { + this.linkerButton = null; + this.buttonId = "model-linker-button"; + this.dialog = null; + } + + setup = () => { + // Remove any existing button + this.removeExistingButton(); + + // Find a visible menu element + const allMenus = document.querySelectorAll("[class*='menu']"); + + // Try to find a visible menu + let visibleMenu = null; + for (const menu of allMenus) { + const style = window.getComputedStyle(menu); + if (style.display !== 'none' && style.visibility !== 'hidden') { + visibleMenu = menu; + break; + } + } + + // Try alternative: app.menu.settingsGroup + if (!visibleMenu && app.menu?.settingsGroup?.element) { + visibleMenu = app.menu.settingsGroup.element.parentElement; + } + + // Try alternative selectors for the top bar + if (!visibleMenu) { + const alternatives = [ + 'header', + '.header', + '.top-bar', + '.toolbar', + '.nav', + '.navigation', + '[role="toolbar"]', + '[role="menubar"]' + ]; + + for (const selector of alternatives) { + const element = document.querySelector(selector); + if (element) { + const style = window.getComputedStyle(element); + if (style.display !== 'none' && style.visibility !== 'hidden') { + visibleMenu = element; + break; + } + } + } + } + + if (visibleMenu) { + this.createLinkerButton(visibleMenu); + } else { + this.createFloatingButton(); + } + + // Create dialog instance (will be created on demand) + if (!this.dialog) { + this.dialog = new LinkerManagerDialog(); + } + } + + removeExistingButton() { + // Remove any existing button by ID + const existingButton = document.getElementById(this.buttonId); + if (existingButton) { + existingButton.remove(); + } + + // Also remove the stored reference if it exists + if (this.linkerButton && this.linkerButton.parentNode) { + this.linkerButton.remove(); + this.linkerButton = null; + } + } + + createLinkerButton(menu) { + this.linkerButton = $el("button", { + id: this.buttonId, + textContent: "🔗 Model Linker", + title: "Open Model Linker to resolve missing models in workflow", + onclick: () => { + this.openLinkerManager(); + }, + style: { + backgroundColor: "var(--comfy-input-bg, #353535)", + color: "var(--input-text, #ffffff)", + border: "2px solid var(--border-color, #555555)", + padding: "8px 16px", + margin: "4px", + borderRadius: "6px", + cursor: "pointer", + fontSize: "14px", + fontWeight: "600", + display: "inline-block", + minWidth: "80px", + textAlign: "center", + zIndex: "1000", + position: "relative", + transition: "all 0.2s ease" + } + }); + + // Add hover effects + this.linkerButton.addEventListener("mouseenter", () => { + this.linkerButton.style.backgroundColor = "var(--comfy-input-bg-hover, #4a4a4a)"; + this.linkerButton.style.borderColor = "var(--primary-color, #007acc)"; + this.linkerButton.style.transform = "translateY(-1px)"; + this.linkerButton.style.boxShadow = "0 2px 4px rgba(0,0,0,0.2)"; + }); + + this.linkerButton.addEventListener("mouseleave", () => { + this.linkerButton.style.backgroundColor = "var(--comfy-input-bg, #353535)"; + this.linkerButton.style.borderColor = "var(--border-color, #555555)"; + this.linkerButton.style.transform = "translateY(0)"; + this.linkerButton.style.boxShadow = "none"; + }); + + // Try to insert before settings group if using app.menu + if (app.menu?.settingsGroup?.element && menu === app.menu.settingsGroup.element.parentElement) { + app.menu.settingsGroup.element.before(this.linkerButton); + } else { + menu.appendChild(this.linkerButton); + } + } + + createFloatingButton() { + // Create a floating button as fallback + this.linkerButton = $el("button", { + id: this.buttonId, + textContent: "🔗 Model Linker", + title: "Open Model Linker to resolve missing models in workflow", + onclick: () => { + this.openLinkerManager(); + }, + style: { + position: "fixed", + top: "10px", + right: "10px", + zIndex: "10000", + backgroundColor: "var(--comfy-input-bg, #353535)", + color: "var(--input-text, #ffffff)", + border: "2px solid var(--primary-color, #007acc)", + padding: "8px 16px", + borderRadius: "6px", + cursor: "pointer", + fontSize: "14px", + fontWeight: "600", + boxShadow: "0 4px 12px rgba(0,0,0,0.3)", + transition: "all 0.2s ease" + } + }); + + // Add hover effects + this.linkerButton.addEventListener("mouseenter", () => { + this.linkerButton.style.backgroundColor = "var(--primary-color, #007acc)"; + this.linkerButton.style.transform = "scale(1.05)"; + }); + + this.linkerButton.addEventListener("mouseleave", () => { + this.linkerButton.style.backgroundColor = "var(--comfy-input-bg, #353535)"; + this.linkerButton.style.transform = "scale(1)"; + }); + + document.body.appendChild(this.linkerButton); + } + + openLinkerManager() { + try { + if (!this.dialog) { + this.dialog = new LinkerManagerDialog(); + } + this.dialog.show(); + } catch (error) { + console.error("🔗 Model Linker: Error creating/showing dialog:", error); + alert("Error opening Model Linker: " + error.message); + } + } +} + +const modelLinker = new ModelLinker(); + +// Register the extension +app.registerExtension({ + name: "Model Linker", + setup: modelLinker.setup +}); +