diff --git a/__init__.py b/__init__.py index 43c56ad..0973d58 100644 --- a/__init__.py +++ b/__init__.py @@ -20,39 +20,45 @@ __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) - + 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: + + 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 @@ -60,179 +66,221 @@ class ModelLinkerExtension: 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 + 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.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 + from .core.sources.civitai import ( + search_civitai_for_file, + search_civitai, + get_civitai_download_url, + ) + download_available = True except ImportError as e: - self.logger.warning(f"Model Linker: Download features not available: {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') - + workflow_json = data.get("workflow") + if not workflow_json: return web.json_response( - {'error': 'Workflow JSON is required'}, - status=400 + {"error": "Workflow JSON is required"}, status=400 ) - + # Analyze and find matches result = analyze_and_find_matches(workflow_json) - + # If download available, auto-search for download sources when no 100% local match if download_available: - for missing in result.get('missing_models', []): + 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) - + 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] - + 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', '') + 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' + if "huggingface.co" in workflow_url: + source = "huggingface" + elif "civitai.com" in workflow_url: + source = "civitai" else: - source = 'workflow' - + 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) + + 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)) + 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 + + 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' + 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) + 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') + 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 - + # 3. Search HuggingFace (exact_only=True for downloads) - hf_result = search_huggingface_for_file(filename, exact_only=True) + hf_result = search_huggingface_for_file( + filename, exact_only=True + ) if hf_result: - missing['download_source'] = { - 'source': 'huggingface', - 'url': hf_result.get('url'), - 'filename': hf_result.get('filename'), - 'name': hf_result.get('repo_id', ''), - 'size': hf_result.get('size', ''), - 'match_type': hf_result.get('match_type', 'exact') + missing["download_source"] = { + "source": "huggingface", + "url": hf_result.get("url"), + "filename": hf_result.get("filename"), + "name": hf_result.get("repo_id", ""), + "size": hf_result.get("size", ""), + "match_type": hf_result.get( + "match_type", "exact" + ), } continue - + # 4. Search CivitAI (exact_only=True for downloads) - civitai_result = search_civitai_for_file(filename, exact_only=True) + civitai_result = search_civitai_for_file( + filename, exact_only=True + ) if civitai_result: - missing['download_source'] = { - 'source': 'civitai', - 'url': civitai_result.get('download_url'), # CivitAI uses download_url - 'filename': civitai_result.get('filename'), - 'name': civitai_result.get('name', ''), - 'size': civitai_result.get('size', ''), - 'match_type': civitai_result.get('match_type', 'exact') + missing["download_source"] = { + "source": "civitai", + "url": civitai_result.get( + "download_url" + ), # CivitAI uses download_url + "filename": civitai_result.get("filename"), + "name": civitai_result.get("name", ""), + "size": civitai_result.get("size", ""), + "match_type": civitai_result.get( + "match_type", "exact" + ), } - + 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 - ) - + 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', []) - + workflow_json = data.get("workflow") + resolutions = data.get("resolutions", []) + if not workflow_json: return web.json_response( - {'error': 'Workflow JSON is required'}, - status=400 + {"error": "Workflow JSON is required"}, status=400 ) - + if not resolutions: return web.json_response( - {'error': 'Resolutions array is required'}, - status=400 + {"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 - }) + + 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 + {"error": str(e), "success": False}, status=500 ) - + @routes.get("/model_linker/models") async def get_models(request): """Get list of all available models.""" @@ -240,161 +288,224 @@ class ModelLinkerExtension: 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.logger.error( + f"Model Linker get_models 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', '') - + 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") + ) + if not filename: return web.json_response( - {'error': 'Filename is required'}, - status=400 + {"error": "Filename is required"}, status=400 ) - + + # 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')}" + ) + results = { - 'popular': None, - 'model_list': None, - 'huggingface': None, - 'civitai': None, - 'found': False + "popular": None, + "model_list": None, + "huggingface": None, + "civitai": None, + "found": False, } - + # 1. Check popular models first (curated database) popular_info = get_popular_model_url(filename) if popular_info: - results['popular'] = { - 'source': 'popular', - 'filename': filename, - **popular_info + results["popular"] = { + "source": "popular", + "filename": filename, + **popular_info, } - results['found'] = True - + results["found"] = True + # 2. Search model-list.json (ComfyUI Manager database with fuzzy matching) - if not results['found']: + # Only accept if confidence >= 70% (or if not URN, allow any) + if not results["found"]: model_list_result = search_model_list(filename) if model_list_result: - results['model_list'] = model_list_result - results['found'] = True - + confidence = model_list_result.get("confidence", 0) + # For URNs, only accept high-confidence matches (>70%) + # Otherwise continue to CivitAI search for exact model + 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 + # 3. Search HuggingFace for exact file match - if not results['found']: + if not results["found"]: hf_result = search_huggingface_for_file(filename) if hf_result: - results['huggingface'] = hf_result - results['found'] = True - - # 4. Search CivitAI for exact file match - if not results['found']: - civitai_result = search_civitai_for_file(filename) - if civitai_result: - results['civitai'] = civitai_result - results['found'] = True - + results["huggingface"] = hf_result + results["found"] = True + + # 4. Search CivitAI - use direct download for URNs + if not results["found"]: + # 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 direct download URL - no search needed + self.logger.info( + f"URN: Using direct download for model_id={model_id}, version_id={version_id}" + ) + download_url = get_civitai_download_url(version_id) + results["civitai"] = { + "source": "civitai", + "name": filename, + "filename": filename, + "type": category, + "download_url": download_url, + "url": f"https://civitai.com/models/{model_id}", + } + 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 + 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', '') - + 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 + {"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]) - + filename = unquote(parsed.path.split("/")[-1]) + if not filename: return web.json_response( - {'error': 'Could not determine filename'}, - status=400 + {"error": "Could not determine filename"}, status=400 ) - + # Build headers if needed headers = {} - if 'huggingface.co' in url: - hf_token = data.get('hf_token', '') + 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}" - + 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 + 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 + { + "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'] + 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 + {"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 + 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.""" @@ -402,59 +513,71 @@ class ModelLinkerExtension: 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 + 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'] + download_id = request.match_info["download_id"] cancel_download(download_id) - return web.json_response({'success': True}) + 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 + 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' + "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.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}") + 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) + self.logger.error( + f"Model Linker: Error setting up routes: {e}", exc_info=True + ) return False @@ -463,4 +586,6 @@ try: extension = ModelLinkerExtension() extension.initialize() except Exception as e: - logging.error(f"ComfyUI Model Linker extension initialization failed: {e}", exc_info=True) + logging.error( + f"ComfyUI Model Linker extension initialization failed: {e}", exc_info=True + ) diff --git a/core/downloader.py b/core/downloader.py index c4560db..6521000 100644 --- a/core/downloader.py +++ b/core/downloader.py @@ -36,7 +36,7 @@ def format_bytes(bytes_value: int) -> str: if bytes_value == 0: return "0 B" k = 1024 - sizes = ['B', 'KB', 'MB', 'GB', 'TB'] + sizes = ["B", "KB", "MB", "GB", "TB"] i = 0 while bytes_value >= k and i < len(sizes) - 1: bytes_value /= k @@ -47,54 +47,73 @@ def format_bytes(bytes_value: int) -> str: def get_download_directory(category: str) -> Optional[str]: """ Get the appropriate download directory for a model category. - + Args: category: Model category (e.g., 'checkpoints', 'loras', 'vae') - + Returns: Absolute path to the download directory, or None if not found """ + global folder_paths + if folder_paths is None: - return None - + # Try to import again - ComfyUI might have initialized since last check + try: + import folder_paths as fp + + folder_paths = fp + except ImportError: + return None + # Map common category names to folder_paths keys category_map = { - 'checkpoint': 'checkpoints', - 'checkpoints': 'checkpoints', - 'lora': 'loras', - 'loras': 'loras', - 'vae': 'vae', - 'controlnet': 'controlnet', - 'clip': 'clip', - 'clip_vision': 'clip_vision', - 'upscaler': 'upscale_models', - 'upscale_models': 'upscale_models', - 'embeddings': 'embeddings', - 'embedding': 'embeddings', - 'diffusion_models': 'diffusion_models', - 'unet': 'diffusion_models', - 'text_encoders': 'text_encoders', - 'text_encoder': 'text_encoders', - 'ipadapter': 'ipadapter', - 'ip-adapter': 'ipadapter', + "checkpoint": "checkpoints", + "checkpoints": "checkpoints", + "lora": "loras", + "loras": "loras", + "vae": "vae", + "controlnet": "controlnet", + "clip": "clip", + "clip_vision": "clip_vision", + "upscaler": "upscale_models", + "upscale_models": "upscale_models", + "embeddings": "embeddings", + "embedding": "embeddings", + "diffusion_models": "diffusion_models", + "unet": "diffusion_models", + "text_encoders": "text_encoders", + "text_encoder": "text_encoders", + "ipadapter": "ipadapter", + "ip-adapter": "ipadapter", + # Legacy/unsupported categories map to checkpoints as fallback + "default": "upscale_models", # Model list uses "default" for upscalers } - + folder_key = category_map.get(category.lower(), category.lower()) - + try: paths = folder_paths.get_folder_paths(folder_key) if paths: # Return the first path (usually the main models directory) return paths[0] + + # If category not found, try to get any models directory as fallback + all_names = folder_paths.get_folder_names() + if all_names: + # Fall back to first available directory + fallback_paths = folder_paths.get_folder_paths(all_names[0]) + if fallback_paths: + return fallback_paths[0] except Exception as e: logger.debug(f"Could not get folder path for {folder_key}: {e}") - + return None def generate_download_id() -> str: """Generate a unique download ID.""" import uuid + return str(uuid.uuid4())[:8] @@ -104,11 +123,11 @@ def download_file( download_id: str, headers: Optional[Dict[str, str]] = None, chunk_size: int = None, - progress_callback: Optional[Callable[[int, int], None]] = None + progress_callback: Optional[Callable[[int, int], None]] = None, ) -> Dict[str, Any]: """ Download a file from URL with progress tracking and speed calculation. - + Args: url: URL to download from dest_path: Destination file path @@ -116,213 +135,244 @@ def download_file( headers: Optional HTTP headers (for auth tokens) chunk_size: Download chunk size in bytes (defaults to 1MB) progress_callback: Optional callback(downloaded_bytes, total_bytes) - + Returns: Result dictionary with status and info """ global download_progress, cancelled_downloads - + # Use default 1MB chunk size if not specified if chunk_size is None: chunk_size = CHUNK_SIZE - + result = { - 'success': False, - 'download_id': download_id, - 'path': dest_path, - 'error': None, - 'size': 0 + "success": False, + "download_id": download_id, + "path": dest_path, + "error": None, + "size": 0, } - + # Initialize progress tracking with speed calculation start_time = time.time() speed_history: deque = deque(maxlen=SPEED_HISTORY_SIZE) last_speed_update = start_time last_downloaded = 0 last_cli_log = start_time # Track when we last logged to CLI - + with download_lock: download_progress[download_id] = { - 'status': 'starting', - 'progress': 0, - 'total_size': 0, - 'downloaded': 0, - 'filename': os.path.basename(dest_path), - 'url': url, - 'error': None, - 'speed': 0, # bytes per second - 'start_time': start_time + "status": "starting", + "progress": 0, + "total_size": 0, + "downloaded": 0, + "filename": os.path.basename(dest_path), + "url": url, + "error": None, + "speed": 0, # bytes per second + "start_time": start_time, } - + try: # Ensure destination directory exists os.makedirs(os.path.dirname(dest_path), exist_ok=True) - + # Verbose logging - what model and from where filename = os.path.basename(dest_path) - source = "HuggingFace" if "huggingface.co" in url else "CivitAI" if "civitai.com" in url else "URL" + source = ( + "HuggingFace" + if "huggingface.co" in url + else "CivitAI" + if "civitai.com" in url + else "URL" + ) print(f"\n[Model Linker] Starting download: {filename}") print(f"[Model Linker] Source: {source}") print(f"[Model Linker] URL: {url}") - + # Start download - response = requests.get( - url, - headers=headers, - stream=True, - timeout=30 - ) + response = requests.get(url, headers=headers, stream=True, timeout=30) response.raise_for_status() - + # Get total size - total_size = int(response.headers.get('content-length', 0)) + total_size = int(response.headers.get("content-length", 0)) total_size_str = format_bytes(total_size) if total_size > 0 else "unknown" print(f"[Model Linker] Size: {total_size_str}") - + with download_lock: - download_progress[download_id]['total_size'] = total_size - download_progress[download_id]['status'] = 'downloading' - + download_progress[download_id]["total_size"] = total_size + download_progress[download_id]["status"] = "downloading" + downloaded = 0 - + # Download with progress and speed calculation cancelled = False - with open(dest_path, 'wb') as f: + with open(dest_path, "wb") as f: for chunk in response.iter_content(chunk_size=chunk_size): # Check for cancellation if download_id in cancelled_downloads: cancelled = True break - + if chunk: f.write(chunk) downloaded += len(chunk) - + # Calculate speed with smoothing current_time = time.time() time_delta = current_time - last_speed_update - + # Update speed every 0.5 seconds to avoid too frequent calculations if time_delta >= 0.5: bytes_delta = downloaded - last_downloaded - instant_speed = bytes_delta / time_delta if time_delta > 0 else 0 + instant_speed = ( + bytes_delta / time_delta if time_delta > 0 else 0 + ) speed_history.append(instant_speed) - + # Calculate smoothed speed (average of recent samples) - smoothed_speed = sum(speed_history) / len(speed_history) if speed_history else 0 - + smoothed_speed = ( + sum(speed_history) / len(speed_history) + if speed_history + else 0 + ) + last_speed_update = current_time last_downloaded = downloaded - + # Update progress with speed with download_lock: - download_progress[download_id]['downloaded'] = downloaded - download_progress[download_id]['speed'] = int(smoothed_speed) + download_progress[download_id]["downloaded"] = downloaded + download_progress[download_id]["speed"] = int( + smoothed_speed + ) if total_size > 0: - download_progress[download_id]['progress'] = int((downloaded / total_size) * 100) - + download_progress[download_id]["progress"] = int( + (downloaded / total_size) * 100 + ) + # CLI progress logging (every CLI_LOG_INTERVAL seconds) if current_time - last_cli_log >= CLI_LOG_INTERVAL: last_cli_log = current_time - progress_pct = int((downloaded / total_size) * 100) if total_size > 0 else 0 + progress_pct = ( + int((downloaded / total_size) * 100) + if total_size > 0 + else 0 + ) downloaded_str = format_bytes(downloaded) - total_str = format_bytes(total_size) if total_size > 0 else "?" + total_str = ( + format_bytes(total_size) if total_size > 0 else "?" + ) speed_str = format_bytes(int(smoothed_speed)) + "/s" - print(f"[Model Linker] Progress: {downloaded_str} / {total_str} ({progress_pct}%) - {speed_str}") + print( + f"[Model Linker] Progress: {downloaded_str} / {total_str} ({progress_pct}%) - {speed_str}" + ) else: # Just update downloaded bytes without recalculating speed with download_lock: - download_progress[download_id]['downloaded'] = downloaded + download_progress[download_id]["downloaded"] = downloaded if total_size > 0: - download_progress[download_id]['progress'] = int((downloaded / total_size) * 100) - + download_progress[download_id]["progress"] = int( + (downloaded / total_size) * 100 + ) + if progress_callback: progress_callback(downloaded, total_size) - + # Handle cancellation after file is closed (so we can delete it on Windows) # Also check if cancellation was requested while we were finishing up if cancelled or download_id in cancelled_downloads: with download_lock: - download_progress[download_id]['status'] = 'cancelled' + download_progress[download_id]["status"] = "cancelled" # Clean up partial/incomplete file try: if os.path.exists(dest_path): os.remove(dest_path) - print(f"[Model Linker] Cancelled: {filename} - incomplete file deleted") + print( + f"[Model Linker] Cancelled: {filename} - incomplete file deleted" + ) else: print(f"[Model Linker] Cancelled: {filename} - no file to delete") except Exception as e: - print(f"[Model Linker] Warning: Could not delete incomplete file {dest_path}: {e}") + print( + f"[Model Linker] Warning: Could not delete incomplete file {dest_path}: {e}" + ) # Try harder on Windows - sometimes the file handle takes a moment to release try: time.sleep(0.5) # time is already imported at module level if os.path.exists(dest_path): os.remove(dest_path) - print(f"[Model Linker] Cancelled: {filename} - incomplete file deleted (delayed)") + print( + f"[Model Linker] Cancelled: {filename} - incomplete file deleted (delayed)" + ) except Exception: pass - result['error'] = 'Download cancelled' + result["error"] = "Download cancelled" cancelled_downloads.discard(download_id) return result - + # Success with download_lock: - download_progress[download_id]['status'] = 'completed' - download_progress[download_id]['progress'] = 100 - download_progress[download_id]['speed'] = 0 # Reset speed on completion - - result['success'] = True - result['size'] = downloaded - + download_progress[download_id]["status"] = "completed" + download_progress[download_id]["progress"] = 100 + download_progress[download_id]["speed"] = 0 # Reset speed on completion + + result["success"] = True + result["size"] = downloaded + # CLI completion log elapsed = time.time() - start_time avg_speed = downloaded / elapsed if elapsed > 0 else 0 print(f"[Model Linker] ✓ Download complete: {filename}") - print(f"[Model Linker] Size: {format_bytes(downloaded)}, Time: {elapsed:.1f}s, Avg speed: {format_bytes(int(avg_speed))}/s") - + print( + f"[Model Linker] Size: {format_bytes(downloaded)}, Time: {elapsed:.1f}s, Avg speed: {format_bytes(int(avg_speed))}/s" + ) + except requests.exceptions.RequestException as e: error_msg = str(e) # Check for specific HTTP errors - if hasattr(e, 'response') and e.response is not None: + if hasattr(e, "response") and e.response is not None: status_code = e.response.status_code if status_code in [401, 403]: - if 'huggingface.co' in url: + if "huggingface.co" in url: error_msg = f"Unauthorized (HTTP {status_code}): HuggingFace token may be required." - elif 'civitai.com' in url: + elif "civitai.com" in url: error_msg = f"Unauthorized (HTTP {status_code}): CivitAI API key may be required." else: - error_msg = f"Unauthorized (HTTP {status_code}): Authentication required." + error_msg = ( + f"Unauthorized (HTTP {status_code}): Authentication required." + ) elif status_code == 404: error_msg = "Model not found (HTTP 404): The file may have been moved or deleted." - + with download_lock: - download_progress[download_id]['status'] = 'error' - download_progress[download_id]['error'] = error_msg - result['error'] = error_msg - + download_progress[download_id]["status"] = "error" + download_progress[download_id]["error"] = error_msg + result["error"] = error_msg + # CLI error log print(f"[Model Linker] ✗ Download failed: {os.path.basename(dest_path)}") print(f"[Model Linker] Error: {error_msg}") - + # Clean up partial file try: if os.path.exists(dest_path): os.remove(dest_path) except: pass - + except Exception as e: error_msg = str(e) with download_lock: - download_progress[download_id]['status'] = 'error' - download_progress[download_id]['error'] = error_msg - result['error'] = error_msg - + download_progress[download_id]["status"] = "error" + download_progress[download_id]["error"] = error_msg + result["error"] = error_msg + # CLI error log print(f"[Model Linker] ✗ Download failed: {os.path.basename(dest_path)}") print(f"[Model Linker] Error: {error_msg}") logger.error(f"Download error: {e}", exc_info=True) - + return result @@ -332,11 +382,11 @@ def download_model( category: str, download_id: Optional[str] = None, headers: Optional[Dict[str, str]] = None, - subfolder: str = "" + subfolder: str = "", ) -> Dict[str, Any]: """ Download a model to the appropriate directory. - + Args: url: URL to download from filename: Filename to save as @@ -344,37 +394,37 @@ def download_model( download_id: Optional download ID (generated if not provided) headers: Optional HTTP headers subfolder: Optional subfolder within category directory - + Returns: Result dictionary """ if download_id is None: download_id = generate_download_id() - + # Get destination directory dest_dir = get_download_directory(category) if not dest_dir: return { - 'success': False, - 'download_id': download_id, - 'error': f'Could not find directory for category: {category}' + "success": False, + "download_id": download_id, + "error": f"Could not find directory for category: {category}", } - + # Add subfolder if specified if subfolder: dest_dir = os.path.join(dest_dir, subfolder) - + dest_path = os.path.join(dest_dir, filename) - + # Check if file already exists if os.path.exists(dest_path): return { - 'success': False, - 'download_id': download_id, - 'error': f'File already exists: {dest_path}', - 'path': dest_path + "success": False, + "download_id": download_id, + "error": f"File already exists: {dest_path}", + "path": dest_path, } - + return download_file(url, dest_path, download_id, headers) @@ -400,8 +450,9 @@ def clear_completed_downloads(): """Clear completed/failed downloads from progress tracking.""" with download_lock: to_remove = [ - did for did, info in download_progress.items() - if info.get('status') in ('completed', 'error', 'cancelled') + did + for did, info in download_progress.items() + if info.get("status") in ("completed", "error", "cancelled") ] for did in to_remove: del download_progress[did] @@ -413,20 +464,60 @@ def start_background_download( filename: str, category: str, headers: Optional[Dict[str, str]] = None, - subfolder: str = "" + subfolder: str = "", ) -> str: """ Start a download in a background thread. - + Returns: download_id for tracking progress """ download_id = generate_download_id() - + + # Pre-initialize progress dict so it's always available for polling + # even if download fails before download_file is called + with download_lock: + download_progress[download_id] = { + "status": "starting", + "progress": 0, + "total_size": 0, + "downloaded": 0, + "filename": filename, + "url": url, + "error": None, + "speed": 0, + "start_time": time.time(), + } + def run_download(): - download_model(url, filename, category, download_id, headers, subfolder) - + try: + result = download_model( + url, filename, category, download_id, headers, subfolder + ) + if not result.get("success"): + # Mark as error if download failed + with download_lock: + if download_id in download_progress: + download_progress[download_id]["status"] = "error" + download_progress[download_id]["error"] = result.get( + "error", "Download failed" + ) + except Exception as e: + # Ensure any exception is captured and logged + with download_lock: + download_progress[download_id] = { + "status": "error", + "progress": 0, + "total_size": 0, + "downloaded": 0, + "filename": filename, + "url": url, + "error": str(e), + "speed": 0, + "start_time": time.time(), + } + thread = threading.Thread(target=run_download, daemon=True) thread.start() - + return download_id diff --git a/core/linker.py b/core/linker.py index 423e2cb..1ec8504 100644 --- a/core/linker.py +++ b/core/linker.py @@ -23,74 +23,76 @@ logger = logging.getLogger(__name__) URL_PATTERN = re.compile(r'(https?://(?:huggingface\.co|civitai\.com)[^\s"\'<>\)\\]+)') # Model file extensions to look for -MODEL_EXTENSIONS = ('.safetensors', '.ckpt', '.pt', '.pth', '.bin', '.onnx') +MODEL_EXTENSIONS = (".safetensors", ".ckpt", ".pt", ".pth", ".bin", ".onnx") def extract_workflow_urls(workflow_json: Dict[str, Any]) -> Dict[str, Dict[str, Any]]: """ Extract model URLs from workflow JSON. - + Sources: 1. node.properties.models array - contains {name, url, directory} 2. Regex extraction from workflow JSON string - finds HuggingFace/CivitAI URLs - + Args: workflow_json: Complete workflow JSON dictionary - + Returns: Dict mapping model filename -> {url, directory, source} """ url_map = {} - + # Convert to string for regex search workflow_str = json.dumps(workflow_json) - + # Collect all nodes including from subgraphs - all_nodes = list(workflow_json.get('nodes', [])) - definitions = workflow_json.get('definitions', {}) - subgraphs = definitions.get('subgraphs', []) + all_nodes = list(workflow_json.get("nodes", [])) + definitions = workflow_json.get("definitions", {}) + subgraphs = definitions.get("subgraphs", []) for subgraph in subgraphs: - subgraph_nodes = subgraph.get('nodes', []) + subgraph_nodes = subgraph.get("nodes", []) all_nodes.extend(subgraph_nodes) - + # 1. Extract from node.properties.models (authoritative source) for node in all_nodes: - node_type = node.get('type', '') - properties = node.get('properties', {}) - models_list = properties.get('models', []) - + 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', '') - + name = model_info.get("name", "") + url = model_info.get("url", "") + directory = model_info.get("directory", "") + if name and name not in url_map: url_map[name] = { - 'url': url, - 'directory': directory, - 'node_type': node_type, - 'source': 'node_properties' + "url": url, + "directory": directory, + "node_type": node_type, + "source": "node_properties", } - + # 2. Extract URLs via regex from workflow JSON urls_found = URL_PATTERN.findall(workflow_str) - + # Clean URLs (remove trailing characters that may have been captured) cleaned_urls = [] for url in urls_found: - url = url.split(')')[0].replace('\\n', '').replace('\n', '').strip() + url = url.split(")")[0].replace("\\n", "").replace("\n", "").strip() if url: cleaned_urls.append(url) - + # 3. Extract model filenames via regex - model_pattern = re.compile(r'([\w\-\.%]+\.(?:safetensors|ckpt|pt|pth|bin|onnx))', re.IGNORECASE) + model_pattern = re.compile( + r"([\w\-\.%]+\.(?:safetensors|ckpt|pt|pth|bin|onnx))", re.IGNORECASE + ) model_files_raw = model_pattern.findall(workflow_str) - + # Clean and decode filenames model_files = set() model_name_map = {} # decoded -> original - + for model in model_files_raw: cleaned = model.strip() if cleaned and cleaned[0].isalnum(): @@ -100,79 +102,81 @@ def extract_workflow_urls(workflow_json: Dict[str, Any]) -> Dict[str, Dict[str, decoded = cleaned model_files.add(decoded) model_name_map[decoded] = cleaned - + # 4. Match URLs to model filenames for model in model_files: # Skip if already found in node.properties.models - if model in url_map and url_map[model].get('url'): + if model in url_map and url_map[model].get("url"): continue - + original_name = model_name_map.get(model, model) - + for url in cleaned_urls: # Check decoded name in URL if model in url: if model not in url_map: - url_map[model] = {'url': url, 'directory': '', 'source': 'regex'} - elif not url_map[model].get('url'): - url_map[model]['url'] = url - url_map[model]['source'] = 'regex' + url_map[model] = {"url": url, "directory": "", "source": "regex"} + elif not url_map[model].get("url"): + url_map[model]["url"] = url + url_map[model]["source"] = "regex" break # Check original (possibly URL-encoded) name in URL if original_name in url: if model not in url_map: - url_map[model] = {'url': url, 'directory': '', 'source': 'regex'} - elif not url_map[model].get('url'): - url_map[model]['url'] = url - url_map[model]['source'] = 'regex' + url_map[model] = {"url": url, "directory": "", "source": "regex"} + elif not url_map[model].get("url"): + url_map[model]["url"] = url + url_map[model]["source"] = "regex" break # Check without extension model_base = os.path.splitext(model)[0] if model_base in url or unquote(model_base) in url: if model not in url_map: - url_map[model] = {'url': url, 'directory': '', 'source': 'regex'} - elif not url_map[model].get('url'): - url_map[model]['url'] = url - url_map[model]['source'] = 'regex' + url_map[model] = {"url": url, "directory": "", "source": "regex"} + elif not url_map[model].get("url"): + url_map[model]["url"] = url + url_map[model]["source"] = "regex" break - + return url_map def parse_huggingface_url(url: str) -> Tuple[Optional[str], Optional[str]]: """ Extract HuggingFace repo and path from URL. - + Args: url: HuggingFace URL - + Returns: Tuple of (repo_id, file_path) or (None, None) if not valid """ - if not url or 'huggingface.co' not in url: + if not url or "huggingface.co" not in url: return None, None - + # Pattern: https://huggingface.co/user/repo/resolve/main/path/to/file.safetensors - match = re.match(r'https?://huggingface\.co/([^/]+/[^/]+)/(?:resolve|blob)/[^/]+/(.+)', url) + match = re.match( + r"https?://huggingface\.co/([^/]+/[^/]+)/(?:resolve|blob)/[^/]+/(.+)", url + ) if match: return match.group(1), match.group(2) - + return None, None def analyze_and_find_matches( workflow_json: Dict[str, Any], similarity_threshold: float = 0.0, - max_matches_per_model: int = 10 + 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: { @@ -204,109 +208,114 @@ def analyze_and_find_matches( # Extract URLs from workflow (node.properties.models + regex) workflow_urls = extract_workflow_urls(workflow_json) logger.debug(f"Extracted {len(workflow_urls)} URLs from 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) - + # Enrich missing models with workflow URLs for missing in missing_models: - original_path = missing.get('original_path', '') + original_path = missing.get("original_path", "") filename = os.path.basename(original_path) - + if filename in workflow_urls: url_info = workflow_urls[filename] - missing['workflow_url'] = url_info.get('url', '') - missing['workflow_directory'] = url_info.get('directory', '') - missing['url_source'] = url_info.get('source', '') - + missing["workflow_url"] = url_info.get("url", "") + missing["workflow_directory"] = url_info.get("directory", "") + missing["url_source"] = url_info.get("source", "") + # Handle URNs: fetch expected filename from CivitAI for missing in missing_models: - if missing.get('is_urn'): - urn = missing['urn'] - model_info = resolve_urn(urn['model_id'], urn['version_id']) + if missing.get("is_urn"): + urn = missing["urn"] + model_info = resolve_urn(urn["model_id"], urn["version_id"]) if model_info: - missing['civitai_info'] = { - 'model_name': model_info.get('model_name'), - 'version_name': model_info.get('version_name') + missing["civitai_info"] = { + "model_name": model_info.get("model_name"), + "version_name": model_info.get("version_name"), } if model_info: - missing['civitai_info'] = model_info - missing['expected_filename'] = model_info['expected_filename'] - logger.debug(f"URN {missing['original_path']} → {missing['expected_filename']}") - + missing["civitai_info"] = model_info + missing["expected_filename"] = model_info["expected_filename"] + logger.debug( + f"URN {missing['original_path']} → {missing['expected_filename']}" + ) + + # Add URN type to missing model for search functionality + if "urn" in missing: + missing["urn_type"] = missing["urn"].get("type", "") + # Find matches for each missing model missing_with_matches = [] for missing in missing_models: - target_for_matching = missing.get('original_path', '') - + target_for_matching = missing.get("original_path", "") + # For URNs, prefer expected_filename for matching - if missing.get('is_urn') and missing.get('expected_filename'): - target_for_matching = missing['expected_filename'] - + if missing.get("is_urn") and missing.get("expected_filename"): + target_for_matching = missing["expected_filename"] + # Filter available models by category if known - category = missing.get('category') - if not category or category == 'unknown': + category = missing.get("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') - + + node_type = missing.get("node_type", "") + category = NODE_TYPE_TO_CATEGORY_HINTS.get(node_type, "unknown") + candidates = available_models - if category and category != 'unknown': - candidates = [m for m in available_models if m.get('category') == category] - candidates.extend([m for m in available_models if m.get('category') != category]) - + if category and category != "unknown": + candidates = [m for m in available_models if m.get("category") == category] + candidates.extend( + [m for m in available_models if m.get("category") != category] + ) + # Find matches matches = find_matches( target_for_matching, candidates, threshold=similarity_threshold, - max_results=max_matches_per_model + max_results=max_matches_per_model, ) - + # Deduplicate matches by absolute path seen_absolute_paths = {} deduplicated_matches = [] for match in matches: - model_dict = match['model'] - absolute_path = model_dict.get('path', '') + model_dict = match["model"] + absolute_path = model_dict.get("path", "") if absolute_path: absolute_path = os.path.normpath(absolute_path) - + if absolute_path not in seen_absolute_paths: seen_absolute_paths[absolute_path] = match deduplicated_matches.append(match) else: existing_match = seen_absolute_paths[absolute_path] - if match['confidence'] > existing_match['confidence']: + if match["confidence"] > existing_match["confidence"]: idx = deduplicated_matches.index(existing_match) deduplicated_matches[idx] = match seen_absolute_paths[absolute_path] = match - - missing_with_matches.append({ - **missing, - 'matches': deduplicated_matches - }) - + + 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) + "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]] + 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: @@ -317,7 +326,7 @@ def apply_resolution( 'category': model category (optional), 'resolved_model': model dict from scanner (optional) } - + Returns: Updated workflow JSON dictionary """ @@ -325,42 +334,45 @@ def apply_resolution( 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 + "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'] - + 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/web/linker.js b/web/linker.js index 2da04df..7716287 100644 --- a/web/linker.js +++ b/web/linker.js @@ -1861,8 +1861,36 @@ class LinkerManagerDialog extends ComfyDialog { * Search online for a model */ async searchOnline(missing) { - const filename = missing.original_path?.split('/').pop()?.split('\\').pop() || ''; - const category = missing.category || ''; + let filename = missing.original_path?.split('/').pop()?.split('\\').pop() || ''; + let category = missing.category || ''; + + // For URNs, use the CivitAI model name for searching instead of the URN itself + // and pass the URN type as category (CivitAI expects specific type names) + if (missing.is_urn) { + if (missing.civitai_info?.model_name) { + filename = missing.civitai_info.model_name; + } + // Pass URN type directly - CivitAPI expects types like 'Upscaler', 'Checkpoint' + const urnType = missing.urn_type || ''; + if (urnType) { + // Map URN types to CivitAI type names + const typeMap = { + 'checkpoint': 'Checkpoint', + 'lora': 'LORA', + 'vae': 'VAE', + 'upscaler': 'Upscaler', + 'upscale_model': 'Upscaler', + 'embedding': 'TextualInversion', + 'controlnet': 'Controlnet' + }; + const civitaiType = typeMap[urnType.toLowerCase()]; + if (civitaiType) { + category = civitaiType; + } + } + } + + const isUrn = missing.is_urn || false; const resultsId = `search-results-${missing.node_id}-${missing.widget_index}`; const resultsDiv = this.contentElement?.querySelector(`#${resultsId}`); const searchBtn = this.contentElement?.querySelector(`#search-${missing.node_id}-${missing.widget_index}`); @@ -1877,10 +1905,19 @@ class LinkerManagerDialog extends ComfyDialog { resultsDiv.innerHTML = 'Searching HuggingFace and CivitAI...'; } + // For URNs, include model_id and version_id for direct download + const searchData = { filename, category, is_urn: isUrn }; + if (isUrn && missing.urn) { + searchData.model_id = missing.urn.model_id; + searchData.version_id = missing.urn.version_id; + } + + console.log('Model Linker: Search request:', JSON.stringify(searchData)); + const response = await api.fetchApi('/model_linker/search', { method: 'POST', headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ filename, category }) + body: JSON.stringify(searchData) }); if (!response.ok) { @@ -1888,6 +1925,7 @@ class LinkerManagerDialog extends ComfyDialog { } const data = await response.json(); + console.log('Model Linker: Search response:', JSON.stringify(data)); this.displaySearchResults(missing, data, resultsDiv); } catch (error) { @@ -1911,8 +1949,9 @@ class LinkerManagerDialog extends ComfyDialog { const popular = data.popular; const modelListResult = data.model_list; - const hfResult = data.huggingface?.[0]; - const civitaiResult = data.civitai?.[0]; + // Backend sends these as objects, not arrays - handle both + const hfResult = data.huggingface ? (Array.isArray(data.huggingface) ? data.huggingface[0] : data.huggingface) : null; + const civitaiResult = data.civitai ? (Array.isArray(data.civitai) ? data.civitai[0] : data.civitai) : null; const hasResults = popular || modelListResult || hfResult || civitaiResult; if (!hasResults) { @@ -1975,14 +2014,20 @@ class LinkerManagerDialog extends ComfyDialog { // CivitAI result if (civitaiResult && civitaiResult.download_url) { + const modelUrl = civitaiResult.url || `https://civitai.com/models/${civitaiResult.model_id || ''}`; + // Use expected_filename from URN resolution, or fallback to civitaiResult filename + const downloadFilename = missing.civitai_info?.expected_filename || civitaiResult.filename || civitaiResult.name; html += `
`; html += `Found on CivitAI`; html += `
`; html += `${civitaiResult.filename || civitaiResult.name} `; html += `${civitaiResult.type || civitaiResult.name}`; html += `
`; - html += `
`; - html += ``; html += `
`;