""" @author: Model Linker Team @title: ComfyUI Model Linker @nickname: Model Linker @version: 1.1.0 @description: Extension for relinking missing models and downloading from HuggingFace/CivitAI """ import logging from .core.log_system.log_funcs import ( log_debug, log_info, log_warn, log_error, log_exception, ) # 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 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 # Import download modules try: from .core.downloader import ( start_background_download, get_progress, get_all_progress, cancel_download, get_download_directory, ) from .core.sources.popular import ( get_popular_model_url, search_popular_models, ) from .core.sources.model_list import ( search_model_list, search_model_list_multiple, ) from .core.sources.huggingface import search_huggingface_for_file from .core.sources.civitai import ( search_civitai_for_file, search_civitai, get_civitai_download_url, resolve_urn, ) download_available = True except ImportError as e: self.logger.warning( f"Model Linker: Download features not available: {e}" ) download_available = False # ==================== ANALYZE ROUTES ==================== @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) # Filter out LoraManager lorAs that already exist locally (exists=True) # These should not appear in missing models at all missing_models = result.get("missing_models", []) filtered_missing = [] for missing in missing_models: is_lora = missing.get("is_lora_v2") exists = missing.get("exists") name = missing.get("name") or missing.get("original_path", "") # Log for debugging import logging logging.getLogger(__name__).debug( f"Filtering: {name} is_lora_v2={is_lora} exists={exists}" ) # Skip LoraManager lorAs that already exist locally if is_lora and exists: logging.getLogger(__name__).info( f"Filtered out LoraManager lora: {name}" ) continue filtered_missing.append(missing) result["missing_models"] = filtered_missing result["total_missing"] = len(filtered_missing) # If download available, check for download sources only from LOCAL sources # (workflow_url, popular, model-list.json) - skip automatic online search # Online search is now only triggered on-demand via search button if download_available: for missing in result.get("missing_models", []): # Check if there's a 100% local match matches = missing.get("matches", []) has_perfect_match = any( m.get("confidence", 0) == 100 for m in matches ) if not has_perfect_match: filename = ( missing.get("original_path", "") .split("/")[-1] .split("\\")[-1] ) # 0. Check workflow URL first (highest priority - directly from workflow) workflow_url = missing.get("workflow_url", "") if workflow_url: # Determine source from URL if "huggingface.co" in workflow_url: source = "huggingface" elif "civitai.com" in workflow_url: source = "civitai" else: source = "workflow" # Try to get file size with HEAD request (non-blocking, timeout quickly) file_size = None try: import requests head_response = requests.head( workflow_url, allow_redirects=True, timeout=5, ) if head_response.status_code == 200: file_size = int( head_response.headers.get( "content-length", 0 ) ) except Exception: pass # Size unknown is fine missing["download_source"] = { "source": source, "url": workflow_url, "filename": filename, "directory": missing.get( "workflow_directory", "" ) or missing.get("category", "checkpoints"), "match_type": "exact", "url_source": "workflow", "size": file_size, } continue # 1. Check popular models (always exact match) popular_info = get_popular_model_url(filename) if popular_info: missing["download_source"] = { "source": "popular", "url": popular_info.get("url"), "filename": filename, "type": popular_info.get("type"), "directory": popular_info.get("directory"), "match_type": "exact", } continue # 2. Check model list (ComfyUI Manager database) # Use exact_only=True to avoid confusing fuzzy matches for downloads model_list_result = search_model_list( filename, exact_only=True ) if model_list_result: missing["download_source"] = { "source": "model_list", "url": model_list_result.get("url"), "filename": model_list_result.get("filename"), "name": model_list_result.get("name"), "type": model_list_result.get("type"), "directory": model_list_result.get("directory"), "size": model_list_result.get("size"), "match_type": model_list_result.get( "match_type" ), "confidence": model_list_result.get( "confidence" ), } continue # NOTE: Search for online sources (HuggingFace, CivitAI) is # now done on-demand via /model_linker/search endpoint # when user clicks "Search Online" button, not automatically 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.""" 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) @routes.post("/model_linker/loaded") async def get_loaded_models(request): """Get all currently loaded models in the workflow.""" 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 ) # Import workflow analyzer to extract models from .core.workflow_analyzer import ( analyze_workflow_models, try_resolve_model_path, is_model_filename, URN_REGEX, URN_TYPE_MAP, ) # Get available models for existence checking available_models = get_model_files() available_paths = {m.get("path") for m in available_models} # Create lookup for full paths by filename (with and without extension) path_by_filename = {} for m in available_models: rel_path = m.get("relative_path", "") if rel_path: filename = rel_path.split("/")[-1].split("\\")[-1] path_by_filename[filename] = m.get("path") # Also add without extension for matching (simple approach) if "." in filename: filename_no_ext = filename.rsplit(".", 1)[0] if filename_no_ext not in path_by_filename: path_by_filename[filename_no_ext] = m.get("path") # Add the full relative path as key too path_by_filename[rel_path] = m.get("path") # Also use folder_paths.get_full_path() to get paths import folder_paths for cat in [ "loras", "checkpoints", "vae", "controlnet", "upscale_models", ]: try: filenames = folder_paths.get_filename_list(cat) for fn in filenames: full_path = folder_paths.get_full_path(cat, fn) if ( full_path and full_path not in path_by_filename.values() ): path_by_filename[fn] = full_path fn_no_ext = ( fn.rsplit(".", 1)[0] if "." in fn else fn ) if fn_no_ext not in path_by_filename: path_by_filename[fn_no_ext] = full_path except Exception: pass # Analyze workflow to get all model references all_model_refs = analyze_workflow_models(workflow_json) # Also extract from node.properties.models nodes = list(workflow_json.get("nodes", [])) definitions = workflow_json.get("definitions", {}) subgraphs = definitions.get("subgraphs", []) for subgraph in subgraphs: nodes.extend(subgraph.get("nodes", [])) # Collect all loaded models with their values loaded_models = [] # Process each model reference from analyze_workflow_models for ref in all_model_refs: original_path = ref.get("original_path", "") node_id = ref.get("node_id") widget_index = ref.get("widget_index") node_type = ref.get("node_type", "") category = ref.get("category", "unknown") # Determine model name and strength model_name = original_path.split("/")[-1].split("\\")[-1] strength = None # For standard LoraLoader nodes, strength is in next widget_value if node_type in ["LoraLoader", "LoraLoaderModelOnly"]: # Find the node in workflow to get strength value for node in nodes: if str(node.get("id")) == str(node_id): widgets_values = node.get("widgets_values", []) if len(widgets_values) > widget_index + 1: try: strength = float( widgets_values[widget_index + 1] ) except (ValueError, TypeError): strength = 1.0 break # For text-based lora loaders (LoraLoaderV2, LoraManager), get strength from ref if ref.get("is_lora_v2"): strength = ref.get("strength") model_name = ref.get("name", model_name) # Check if model exists locally exists = ref.get("exists", False) # If URN, resolve to display name if ref.get("is_urn"): urn = ref.get("urn", {}) # Use model name from URN as display name model_name = ( f"urn:{urn.get('type', 'model')}:{urn.get('model_id')}" ) category = urn.get("type", category) if category in URN_TYPE_MAP: category = URN_TYPE_MAP[category] loaded_models.append( { "name": model_name, "category": category, "node_id": node_id, "widget_index": widget_index, "node_type": node_type, "exists": exists, "strength": strength, "original_path": original_path, "is_urn": ref.get("is_urn", False), "is_lora_v2": ref.get("is_lora_v2", False), "active": ref.get("active"), "connected": ref.get("connected", True), "resolved_path": ( path_by_filename.get(model_name) or path_by_filename.get(original_path) ), } ) # Also check node.properties.models for embedded models for node in nodes: node_type = node.get("type", "") properties = node.get("properties", {}) models_list = properties.get("models", []) for model_info in models_list: if isinstance(model_info, dict): name = model_info.get("name", "") url = model_info.get("url", "") directory = model_info.get("directory", "") if name: # Check if this model is already in loaded_models existing = next( ( m for m in loaded_models if m.get("original_path") == name ), None, ) if not existing: loaded_models.append( { "name": name.split("/")[-1].split("\\")[ -1 ], "category": directory or "checkpoints", "node_id": node.get("id"), "widget_index": None, "node_type": node_type, "exists": True, # Embedded models are loaded "strength": None, "original_path": name, "is_urn": False, } ) return web.json_response( {"loaded_models": loaded_models, "total": len(loaded_models)} ) except Exception as e: self.logger.error( f"Model Linker get_loaded_models error: {e}", exc_info=True ) return web.json_response({"error": str(e)}, status=500) # ==================== CIVITAI SEARCH ROUTE ==================== @routes.post("/model_linker/civitai-search") async def civitai_search(request): """Search CivitAI for a model using file hash.""" try: data = await request.json() filename = data.get("filename", "") category = data.get("category", "") resolved_path = data.get("resolved_path", "") if not filename: return web.json_response( {"error": "Filename is required"}, status=400 ) # Clean filename for display import os as _os clean_name = _os.path.splitext(filename)[0] # Get the file path to hash file_path = resolved_path if resolved_path else None if not file_path and category: # Try to find the file in the model directories using folder_paths try: import folder_paths # Map category to folder_paths type category_map = { "loras": "loras", "checkpoints": "checkpoints", "vae": "vae", "controlnet": "controlnet", "upscale_models": "upscale_models", } folder_type = category_map.get( category.lower(), category.lower() ) file_path = folder_paths.get_full_path( folder_type, filename ) except Exception: pass # If not found, try scanner if not file_path: try: from .core.scanner import get_model_files available_models = get_model_files() for m in available_models: if ( m.get("relative_path", "").endswith(filename) or m.get("filename", "") == filename ): file_path = m.get("path") break except Exception: pass # Search CivitAI for the model using hash if download_available and file_path and _os.path.exists(file_path): try: from .core.sources.civitai import ( get_model_info_for_file, ) result = get_model_info_for_file(file_path) if result and result.get("url"): return web.json_response( { "filename": filename, "url": result.get("url"), "version_url": result.get("version_url"), "model_id": result.get("model_id"), "model_name": result.get( "model_name", clean_name ), "model_type": result.get("model_type", ""), "version_id": result.get("version_id"), "version_name": result.get("version_name", ""), "sha256": result.get("sha256"), "base_model": result.get("base_model"), "tags": result.get("tags", []), "trained_words": result.get( "trained_words", [] ), "images": result.get("images", []), "clip_skip": result.get("clip_skip"), "description": result.get("description", ""), "model_description": result.get( "model_description", "" ), } ) except Exception as e: self.logger.warning(f"CivitAI search error: {e}") # No result found - try fallback to filename search if download_available: try: from .core.sources.civitai import ( search_civitai_for_file, ) result = search_civitai_for_file(filename) if result and result.get("url"): return web.json_response( { "url": result["url"], "model_name": result.get("name", clean_name), "version_id": result.get("version_id"), } ) except Exception as e: self.logger.warning(f"CivitAI fallback search error: {e}") # No result found return web.json_response({"url": None}) except Exception as e: self.logger.error( f"Model Linker civitai-search error: {e}", exc_info=True ) return web.json_response({"error": str(e)}, status=500) # ==================== DOWNLOAD ROUTES ==================== if download_available: @routes.post("/model_linker/search") async def search_sources(request): """Search for model download sources.""" try: data = await request.json() filename = data.get("filename", "") category = data.get("category", "") # Handle both boolean and string forms is_urn_raw = data.get("is_urn", False) is_urn = ( is_urn_raw if isinstance(is_urn_raw, bool) else (str(is_urn_raw).lower() == "true") ) # For URN-only requests, model_id and version_id are required instead of filename model_id = data.get("model_id") version_id = data.get("version_id") if not filename and not (is_urn and model_id and version_id): return web.json_response( { "error": "Filename is required for non-URN, or model_id+version_id for URN" }, status=400, ) raw_sources = data.get("sources", ["all"]) if isinstance(raw_sources, str): raw_sources = [raw_sources] elif not isinstance(raw_sources, list): raw_sources = ["all"] normalized_sources = { str(source).strip().lower() for source in raw_sources if str(source).strip() } if not normalized_sources: normalized_sources = {"all"} if "all" in normalized_sources: normalized_sources = {"local", "huggingface", "civitai"} search_local = "local" in normalized_sources search_huggingface_source = "huggingface" in normalized_sources search_civitai_source = "civitai" in normalized_sources # Debug logging self.logger.info( f"Search request: filename={filename}, category={category}, is_urn={is_urn}, model_id={data.get('model_id')}, version_id={data.get('version_id')}, sources={sorted(normalized_sources)}" ) results = { "popular": None, "model_list": None, "huggingface": None, "civitai": None, "found": False, "searched_sources": sorted(normalized_sources), } # 1. Search local databases (curated + model-list) if search_local: popular_info = get_popular_model_url(filename) if popular_info: results["popular"] = { "source": "popular", "filename": filename, **popular_info, } results["found"] = True model_list_result = search_model_list(filename) if model_list_result: confidence = model_list_result.get("confidence", 0) if is_urn and confidence >= 70: results["model_list"] = model_list_result results["found"] = True elif not is_urn: results["model_list"] = model_list_result results["found"] = True # 2. Search HuggingFace for exact file match if search_huggingface_source: hf_result = search_huggingface_for_file(filename) if hf_result: results["huggingface"] = hf_result results["found"] = True # 3. Search CivitAI - use direct download for URNs if search_civitai_source: # For URNs, use direct model_id/version_id to get download URL if is_urn: # Get model_id and version_id from request data model_id = data.get("model_id") version_id = data.get("version_id") if model_id and version_id: # Use resolve_urn to get model info (cached) model_info = resolve_urn(model_id, version_id) if model_info: primary_file = None for file_info in model_info.get("files", []): if ( file_info.get("name") == model_info.get("expected_filename") ): primary_file = file_info break if primary_file is None: primary_file = ( model_info.get("files") or [{}] )[0] download_url = get_civitai_download_url( version_id ) results["civitai"] = { "source": "civitai", "name": model_info.get("model_name"), "filename": model_info.get( "expected_filename" ), "type": category, "download_url": download_url, "url": f"https://civitai.com/models/{model_id}?modelVersionId={version_id}", "model_id": model_id, "version_id": version_id, "size": primary_file.get("size"), } results["found"] = True elif category: # Fallback to search if no IDs self.logger.info( f"URN: No model_id/version_id, falling back to CivitAI search" ) civitai_results = search_civitai( filename, model_type=category ) if civitai_results: first_result = civitai_results[0] results["civitai"] = { "source": "civitai", "name": first_result.get("name"), "filename": first_result.get("filename"), "type": first_result.get("type"), "download_url": first_result.get( "download_url" ), "url": first_result.get("url"), "size": first_result.get("size"), } results["found"] = True else: civitai_result = search_civitai_for_file(filename) if civitai_result: results["civitai"] = civitai_result results["found"] = True return web.json_response(results) except Exception as e: self.logger.error( f"Model Linker search error: {e}", exc_info=True ) return web.json_response({"error": str(e)}, status=500) @routes.post("/model_linker/download") async def download_model(request): """Start downloading a model.""" try: data = await request.json() url = data.get("url", "") filename = data.get("filename", "") category = data.get("category", "checkpoints") subfolder = data.get("subfolder", "") if not url: return web.json_response( {"error": "URL is required"}, status=400 ) if not filename: # Extract filename from URL from urllib.parse import urlparse, unquote parsed = urlparse(url) filename = unquote(parsed.path.split("/")[-1]) if not filename: return web.json_response( {"error": "Could not determine filename"}, status=400 ) # Build headers if needed headers = {} if "huggingface.co" in url: hf_token = data.get("hf_token", "") if hf_token: headers["Authorization"] = f"Bearer {hf_token}" elif "civitai.com" in url: civitai_key = data.get("civitai_key", "") if civitai_key and "token=" not in url: url += ( f"{'&' if '?' in url else '?'}token={civitai_key}" ) # Start background download download_id = start_background_download( url=url, filename=filename, category=category, headers=headers if headers else None, subfolder=subfolder, ) return web.json_response( { "success": True, "download_id": download_id, "filename": filename, "category": category, } ) except Exception as e: self.logger.error( f"Model Linker download error: {e}", exc_info=True ) return web.json_response( {"error": str(e), "success": False}, status=500 ) @routes.get("/model_linker/progress/{download_id}") async def get_download_progress(request): """Get progress for a specific download.""" try: download_id = request.match_info["download_id"] progress = get_progress(download_id) if progress: return web.json_response(progress) else: return web.json_response( {"error": "Download not found"}, status=404 ) except Exception as e: self.logger.error( f"Model Linker progress error: {e}", exc_info=True ) return web.json_response({"error": str(e)}, status=500) @routes.get("/model_linker/progress") async def get_all_downloads_progress(request): """Get progress for all downloads.""" try: progress = get_all_progress() return web.json_response(progress) except Exception as e: self.logger.error( f"Model Linker progress error: {e}", exc_info=True ) return web.json_response({"error": str(e)}, status=500) @routes.post("/model_linker/cancel/{download_id}") async def cancel_download_route(request): """Cancel a download in progress.""" try: download_id = request.match_info["download_id"] cancel_download(download_id) return web.json_response({"success": True}) except Exception as e: self.logger.error( f"Model Linker cancel error: {e}", exc_info=True ) return web.json_response( {"error": str(e), "success": False}, status=500 ) @routes.get("/model_linker/directories") async def get_directories(request): """Get available model directories.""" try: categories = [ "checkpoints", "loras", "vae", "controlnet", "clip", "clip_vision", "embeddings", "upscale_models", "diffusion_models", "text_encoders", "ipadapter", "sams", ] directories = {} for cat in categories: path = get_download_directory(cat) if path: directories[cat] = path return web.json_response(directories) except Exception as e: self.logger.error( f"Model Linker directories 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: log_error( f"ComfyUI Model Linker extension initialization failed: {e}", exc_info=True )