From d906a14804b6844d4493c61bb6bc9cb2ea70a7df Mon Sep 17 00:00:00 2001 From: Dariusz L Date: Sat, 1 Aug 2026 22:32:20 +0200 Subject: [PATCH] refactor: split route registry into dedicated route modules --- core/routes/directories.py | 318 ++ core/routes/downloads.py | 349 +++ core/routes/metadata.py | 199 ++ core/routes/model_info.py | 1258 ++++++++ core/routes/registry.py | 4518 +---------------------------- core/routes/search.py | 1185 ++++++++ core/routes/workflow.py | 1392 +++++++++ tests/test_refactoring_targets.py | 113 + 8 files changed, 4837 insertions(+), 4495 deletions(-) create mode 100644 core/routes/directories.py create mode 100644 core/routes/downloads.py create mode 100644 core/routes/metadata.py create mode 100644 core/routes/model_info.py create mode 100644 core/routes/search.py create mode 100644 core/routes/workflow.py diff --git a/core/routes/directories.py b/core/routes/directories.py new file mode 100644 index 0000000..63aa843 --- /dev/null +++ b/core/routes/directories.py @@ -0,0 +1,318 @@ +"""Directories route registration.""" + +def register_directory_routes(context): + TEMPLATE_KEY_ALIASES = context.get('TEMPLATE_KEY_ALIASES') + asyncio = context.get('asyncio') + dedupe_local_base_directories = context.get('dedupe_local_base_directories') + get_category_folder_keys = context.get('get_category_folder_keys') + get_comfy_root_path = context.get('get_comfy_root_path') + get_default_root_for_category = context.get('get_default_root_for_category') + get_download_directory = context.get('get_download_directory') + get_enabled_download_categories = context.get('get_enabled_download_categories') + get_local_path_identity = context.get('get_local_path_identity') + get_model_files = context.get('get_model_files') + infer_download_path_templates = context.get('infer_download_path_templates') + is_civarchive_available = context.get('is_civarchive_available') + is_lora_manager_archive_available = context.get('is_lora_manager_archive_available') + json_api_endpoint = context.get('json_api_endpoint') + load_resolver_settings = context.get('load_resolver_settings') + normalize_download_category = context.get('normalize_download_category') + prefer_local_base_directory = context.get('prefer_local_base_directory') + register_version_routes = context.get('register_version_routes') + routes = context.get('routes') + self = context.get('self') + split_path_segments = context.get('split_path_segments') + web = context.get('web') + + @routes.get("/model_resolver/directories") + @json_api_endpoint("directories") + async def get_directories(request): + """Get available model directories.""" + import folder_paths + + categories = get_enabled_download_categories(list(folder_paths.folder_names_and_paths.keys())) + directories = {} + for cat in categories: + path = get_download_directory(cat) + if path: + directories[cat] = path + + return web.json_response(directories) + + @routes.get("/model_resolver/root-directories") + @json_api_endpoint("root directories") + async def get_root_directories(request): + """Get configured ComfyUI root directories for path settings.""" + import os + + import folder_paths + + known_categories = set(folder_paths.folder_names_and_paths.keys()) + skip_categories = {"custom_nodes", "configs"} + categories = get_enabled_download_categories(list(known_categories)) + category_source_keys = {} + for cat in [*known_categories]: + folder_key = normalize_download_category(cat) + if folder_key: + category_source_keys.setdefault(folder_key, []).append(cat) + + roots = {} + settings = load_resolver_settings() + comfy_root = get_comfy_root_path(folder_paths) + for cat in categories: + folder_key = normalize_download_category(cat) + if folder_key in known_categories: + raw_candidate_keys = [folder_key] + else: + raw_candidate_keys = [ + folder_key, + *TEMPLATE_KEY_ALIASES.get(folder_key, ()), + *category_source_keys.get(folder_key, []), + ] + if folder_key == "ultralytics": + raw_candidate_keys = [ + candidate_key + for candidate_key in raw_candidate_keys + if str(candidate_key or "").strip().lower() != "yolo" + ] + + candidate_keys = [] + for candidate_key in raw_candidate_keys: + if ( + candidate_key + and candidate_key in known_categories + and candidate_key not in candidate_keys + ): + candidate_keys.append(candidate_key) + paths = [] + for candidate_key in candidate_keys: + paths.extend(folder_paths.get_folder_paths(candidate_key) or []) + if folder_key == "ultralytics": + normalized_ultralytics_paths = [] + for path in paths: + normalized_path = os.path.normpath(str(path or "")) + basename = os.path.basename(normalized_path).lower() + if basename == "yolo": + continue + parent_dir = os.path.dirname(normalized_path) + if ( + basename in {"bbox", "segm"} + and os.path.basename(parent_dir).lower() + == "ultralytics" + ): + normalized_path = parent_dir + normalized_ultralytics_paths.append(normalized_path) + paths = normalized_ultralytics_paths + preferred_directory = get_default_root_for_category( + folder_key, settings + ) + normalized_paths = dedupe_local_base_directories( + paths, + preferred_directory=preferred_directory, + comfy_root=comfy_root, + ) + roots[folder_key] = normalized_paths + + for raw_key in known_categories: + if ( + not raw_key + or raw_key in skip_categories + or raw_key in roots + ): + continue + raw_paths = folder_paths.get_folder_paths(raw_key) or [] + roots[raw_key] = dedupe_local_base_directories( + raw_paths, + comfy_root=comfy_root, + ) + + return web.json_response(roots) + + @routes.get("/model_resolver/path-template-suggestions") + @json_api_endpoint("path template suggestions") + async def get_path_template_suggestions(request): + """Infer path template presets from existing local model folders.""" + from ..sources.popular import get_base_models_config + + force_rescan = request.query.get("force") == "1" + models = await asyncio.to_thread(get_model_files, force_rescan) + base_models_config = get_base_models_config() + suggestions = await asyncio.to_thread( + infer_download_path_templates, + models, + base_models_config, + ) + return web.json_response(suggestions) + + @routes.get("/model_resolver/capabilities") + @json_api_endpoint("capabilities") + async def get_capabilities(request): + """Get optional source capabilities available in this install.""" + from ..workflow_analyzer import NODE_TYPE_MODEL_WIDGET_CATEGORIES + return web.json_response( + { + "sources": { + "civarchive": is_civarchive_available(), + "lora_manager_archive": is_lora_manager_archive_available() + }, + "node_rules": NODE_TYPE_MODEL_WIDGET_CATEGORIES + } + ) + + register_version_routes(routes, web, json_api_endpoint) + + @routes.get("/model_resolver/subfolders/{category}") + @json_api_endpoint("subfolders") + async def get_subfolders(request): + """Get known subfolders for a category using ComfyUI folder_paths.""" + import os + + import folder_paths + + raw_category = (request.match_info.get("category") or "").strip() + category = normalize_download_category(raw_category) + + if not category or category == "unknown": + return web.json_response([]) + + known_categories = set(folder_paths.folder_names_and_paths.keys()) + folder_keys = get_category_folder_keys(category) + available_folder_keys = [ + folder_key for folder_key in folder_keys if folder_key in known_categories + ] + if not available_folder_keys: + self.logger.debug( + f"Model Resolver: skipping subfolder lookup for unknown category '{raw_category}' -> '{category}'" + ) + return web.json_response([]) + + subfolders = {} + settings = load_resolver_settings() + comfy_root = get_comfy_root_path(folder_paths) + preferred_directory = ( + get_default_root_for_category(category, settings) + or get_download_directory(category) + or "" + ) + + def add_subfolder(rel_path, base_dir=""): + rel_path = "/".join(split_path_segments(rel_path)) + if not rel_path or rel_path == ".": + return + base_dir = os.path.abspath(base_dir) if base_dir else "" + base_identity = ( + get_local_path_identity(base_dir) if base_dir else "" + ) + key = (rel_path.lower(), base_identity) + base_label = ( + os.path.basename(os.path.normpath(base_dir)) + if base_dir + else "" + ) + current = subfolders.get(key) + if current and not prefer_local_base_directory( + base_dir, + current.get("base_directory", ""), + preferred_directory, + comfy_root, + ): + return + subfolders[key] = { + "value": rel_path, + "label": rel_path, + "base_label": base_label, + "base_directory": base_dir, + } + + raw_base_dirs = [] + for folder_key in available_folder_keys: + for base_dir in folder_paths.get_folder_paths(folder_key) or []: + if not base_dir or not os.path.isdir(base_dir): + continue + raw_base_dirs.append(base_dir) + base_dirs = dedupe_local_base_directories( + raw_base_dirs, + preferred_directory=preferred_directory, + comfy_root=comfy_root, + ) + + def find_base_dir(full_path): + if not full_path: + return "" + full_path_identity = get_local_path_identity(full_path) + for base_dir in base_dirs: + base_identity = get_local_path_identity(base_dir) + try: + if ( + os.path.commonpath( + [full_path_identity, base_identity] + ) + == base_identity + ): + return base_dir + except Exception: + continue + return "" + + for folder_key in available_folder_keys: + filenames = folder_paths.get_filename_list(folder_key) or [] + for rel_path in filenames: + if not isinstance(rel_path, str): + continue + base_dir = "" + try: + base_dir = find_base_dir( + folder_paths.get_full_path(folder_key, rel_path) + ) + except Exception: + base_dir = "" + parts = split_path_segments(rel_path) + if len(parts) <= 1: + continue + current = "" + for part in parts[:-1]: + current = f"{current}/{part}" if current else part + add_subfolder(current, base_dir) + + for base_dir in base_dirs: + for root, dirs, _files in os.walk(base_dir): + rel_root = os.path.relpath(root, base_dir) + for dirname in dirs: + rel_path = ( + dirname + if rel_root in ("", ".") + else os.path.join(rel_root, dirname) + ) + add_subfolder(rel_path, base_dir) + + value_counts = {} + for item in subfolders.values(): + value_key = item.get("value", "").lower() + value_counts[value_key] = value_counts.get(value_key, 0) + 1 + + response_items = [] + for item in subfolders.values(): + value = item.get("value", "") + base_label = item.get("base_label", "") + label = ( + f"{value} ({base_label})" + if base_label and value_counts.get(value.lower(), 0) > 1 + else value + ) + response_items.append( + { + "value": value, + "label": label, + "base_directory": item.get("base_directory", ""), + } + ) + + return web.json_response( + sorted( + response_items, + key=lambda item: ( + item.get("value", "").lower(), + item.get("base_directory", "").lower(), + ), + ) + ) diff --git a/core/routes/downloads.py b/core/routes/downloads.py new file mode 100644 index 0000000..f38097d --- /dev/null +++ b/core/routes/downloads.py @@ -0,0 +1,349 @@ +"""Downloads route registration.""" + +def register_download_routes(context): + Aria2InstallError = context.get('Aria2InstallError') + UnsafeUrlError = context.get('UnsafeUrlError') + asyncio = context.get('asyncio') + cancel_download = context.get('cancel_download') + clear_completed_downloads = context.get('clear_completed_downloads') + first_non_empty = context.get('first_non_empty') + get_all_progress = context.get('get_all_progress') + get_aria2_status = context.get('get_aria2_status') + get_civarchive_model_details = context.get('get_civarchive_model_details') + get_civitai_model_details = context.get('get_civitai_model_details') + get_default_root_for_category = context.get('get_default_root_for_category') + get_download_directory = context.get('get_download_directory') + get_override_settings_from_request = context.get('get_override_settings_from_request') + get_progress = context.get('get_progress') + host_matches_domain = context.get('host_matches_domain') + install_aria2_engine = context.get('install_aria2_engine') + is_allowed_model_download_filename = context.get('is_allowed_model_download_filename') + json_api_endpoint = context.get('json_api_endpoint') + load_resolver_settings = context.get('load_resolver_settings') + normalize_download_category = context.get('normalize_download_category') + pause_download = context.get('pause_download') + resolve_download_subfolder = context.get('resolve_download_subfolder') + resume_download = context.get('resume_download') + routes = context.get('routes') + sanitize_download_filename = context.get('sanitize_download_filename') + save_resolver_settings = context.get('save_resolver_settings') + self = context.get('self') + split_path_segments = context.get('split_path_segments') + start_aria2_daemon = context.get('start_aria2_daemon') + start_background_download = context.get('start_background_download') + stop_aria2_daemon = context.get('stop_aria2_daemon') + to_bool = context.get('to_bool') + to_int = context.get('to_int') + validate_public_http_url = context.get('validate_public_http_url') + web = context.get('web') + + @routes.post("/model_resolver/download") + @json_api_endpoint("download", return_success_on_error=True) + async def download_model(request): + """Start downloading a model.""" + data = await request.json() + url = data.get("url", "") + filename = data.get("filename", "") + category = data.get("category", "checkpoints") + category = normalize_download_category(category) + subfolder = data.get("subfolder", "") + base_directory = data.get("base_directory", "") + path_metadata = data.get("path_metadata", {}) + if not isinstance(path_metadata, dict): + path_metadata = {} + download_metadata = data.get("download_metadata") or data.get( + "metadata", {} + ) + if not isinstance(download_metadata, dict): + download_metadata = {} + download_metadata = dict(download_metadata) + settings = load_resolver_settings() + if not base_directory: + base_directory = get_default_root_for_category(category, settings) + subfolder = resolve_download_subfolder( + category, + subfolder, + path_metadata, + settings, + ) + + if not url: + return web.json_response( + {"error": "URL is required"}, status=400 + ) + try: + url = await asyncio.to_thread(validate_public_http_url, url) + except UnsafeUrlError as exc: + return web.json_response({"error": str(exc)}, status=400) + + from urllib.parse import urlparse as _download_urlparse + + download_host = _download_urlparse(url).hostname + + if not filename: + # Extract filename from URL + from urllib.parse import unquote, urlparse + + parsed = urlparse(url) + filename = unquote(parsed.path.split("/")[-1]) + filename = sanitize_download_filename(filename) + + if not filename: + return web.json_response( + {"error": "Could not determine filename"}, status=400 + ) + if not is_allowed_model_download_filename(filename): + return web.json_response( + {"error": "Unsupported model file extension"}, status=400 + ) + + # Build headers if needed + headers = {} + if host_matches_domain(download_host, "huggingface.co"): + hf_token = data.get("hf_token", "") + if hf_token: + headers["Authorization"] = f"Bearer {hf_token}" + elif host_matches_domain( + download_host, + "civitai.com", + "civitai.red", + ): + civitai_key = data.get("civitai_key", "") + if civitai_key and "token=" not in url: + url += ( + f"{'&' if '?' in url else '?'}token={civitai_key}" + ) + civitai_session_token = str( + data.get("civitai_session_token", "") or "" + ).strip() + if civitai_session_token: + headers["Cookie"] = ( + f"__Secure-civitai-token={civitai_session_token}" + ) + + def _first_metadata_value(*values): + return first_non_empty(*values) + + def _metadata_int(value): + return to_int(value) + + inferred_source = "" + if host_matches_domain(download_host, "civitai.com", "civitai.red"): + inferred_source = "civitai" + elif host_matches_domain(download_host, "huggingface.co"): + inferred_source = "huggingface" + + download_metadata.setdefault("filename", filename) + download_metadata.setdefault("category", category) + download_metadata.setdefault("download_url", url) + download_metadata.setdefault("source_url", url) + download_metadata.setdefault("path_metadata", path_metadata) + download_metadata.setdefault( + "source", + _first_metadata_value( + download_metadata.get("details_source"), + path_metadata.get("source"), + inferred_source, + ), + ) + + model_id = _metadata_int( + _first_metadata_value( + download_metadata.get("model_id"), + download_metadata.get("modelId"), + path_metadata.get("model_id"), + ) + ) + version_id = _metadata_int( + _first_metadata_value( + download_metadata.get("version_id"), + download_metadata.get("versionId"), + path_metadata.get("version_id"), + ) + ) + source_name = str( + _first_metadata_value( + download_metadata.get("details_source"), + download_metadata.get("source"), + ) + ).lower() + try: + if ( + source_name == "civitai" + and model_id + and not download_metadata.get("civitai_details") + ): + details = await asyncio.to_thread( + get_civitai_model_details, + model_id, + version_id, + data.get("civitai_key", ""), + ) + if details: + download_metadata["civitai_details"] = details + elif ( + source_name == "civarchive" + and model_id + and not download_metadata.get("civitai_details") + ): + details = await asyncio.to_thread( + get_civarchive_model_details, + model_id, + version_id, + ) + if details: + download_metadata["civitai_details"] = details + except Exception as metadata_error: + self.logger.warning( + f"Model metadata lookup failed: {metadata_error}" + ) + + target_directory = "" + target_path = "" + try: + import os as _download_os + + target_directory = ( + get_download_directory(category, base_directory) or "" + ) + if target_directory and subfolder: + target_directory = _download_os.path.join( + target_directory, *split_path_segments(subfolder) + ) + if target_directory: + target_path = _download_os.path.join( + target_directory, filename + ) + except Exception: + target_directory = "" + target_path = "" + + # Start background download + download_id = start_background_download( + url=url, + filename=filename, + category=category, + headers=headers if headers else None, + subfolder=subfolder, + base_directory=base_directory, + metadata=download_metadata, + ) + + return web.json_response( + { + "success": True, + "download_id": download_id, + "filename": filename, + "category": category, + "path": target_path, + "directory": target_directory, + "download_backend": settings.get("download_backend", "python"), + } + ) + + @routes.get("/model_resolver/progress/{download_id}") + @json_api_endpoint("progress") + async def get_download_progress(request): + """Get progress for a specific download.""" + 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 + ) + + @routes.get("/model_resolver/progress") + @json_api_endpoint("progress") + async def get_all_downloads_progress(request): + """Get progress for all downloads.""" + progress = get_all_progress() + return web.json_response(progress) + + @routes.post("/model_resolver/cancel/{download_id}") + @json_api_endpoint("cancel", return_success_on_error=True) + async def cancel_download_route(request): + """Cancel a download in progress.""" + download_id = request.match_info["download_id"] + cancel_download(download_id) + return web.json_response({"success": True}) + + @routes.post("/model_resolver/pause/{download_id}") + @json_api_endpoint("pause", return_success_on_error=True) + async def pause_download_route(request): + """Pause an aria2 download.""" + download_id = request.match_info["download_id"] + result = pause_download(download_id) + status = 200 if result.get("success") else 400 + return web.json_response(result, status=status) + + @routes.post("/model_resolver/resume/{download_id}") + @json_api_endpoint("resume", return_success_on_error=True) + async def resume_download_route(request): + """Resume an aria2 download.""" + download_id = request.match_info["download_id"] + result = resume_download(download_id) + status = 200 if result.get("success") else 400 + return web.json_response(result, status=status) + + @routes.post("/model_resolver/clear_completed_downloads") + @json_api_endpoint("clear_completed_downloads", return_success_on_error=True) + async def clear_completed_downloads_route(request): + """Clear completed and failed downloads from progress memory.""" + clear_completed_downloads() + return web.json_response({"success": True}) + + @routes.get("/model_resolver/aria2/status") + @routes.post("/model_resolver/aria2/status") + @json_api_endpoint("aria2 status") + async def aria2_status_route(request): + """Report aria2 availability using saved settings, with optional override via POST.""" + settings = await get_override_settings_from_request(request) + return web.json_response(await asyncio.to_thread(get_aria2_status, settings)) + + + @routes.post("/model_resolver/aria2/start") + @json_api_endpoint("aria2 start", return_success_on_error=True) + async def aria2_start_route(request): + """Start the aria2 daemon without starting a download.""" + settings = await get_override_settings_from_request(request) + result = await asyncio.to_thread(start_aria2_daemon, settings) + status = 200 if result.get("success") else 400 + return web.json_response(result, status=status) + + @routes.get("/model_resolver/aria2/stop") + @routes.post("/model_resolver/aria2/stop") + @json_api_endpoint("aria2 stop", return_success_on_error=True) + async def aria2_stop_route(request): + """Stop the aria2 daemon started by Model Resolver.""" + result = await asyncio.to_thread(stop_aria2_daemon) + status = 200 if result.get("success") else 400 + return web.json_response(result, status=status) + + @routes.post("/model_resolver/aria2/install") + @json_api_endpoint("aria2 install") + async def aria2_install_route(request): + """Download official aria2c binary and save its path.""" + try: + payload = await request.json() + except Exception: + payload = {} + force = False + if isinstance(payload, dict): + force = to_bool(payload.get("force"), False) + try: + install_result = await asyncio.to_thread(install_aria2_engine, force) + except Aria2InstallError as exc: + self.logger.warning(f"Model Resolver aria2 install failed: {exc}") + return web.json_response({"success": False, "error": str(exc)}, status=500) + settings = await asyncio.to_thread( + save_resolver_settings, + { + "aria2c_path": install_result.get("aria2c_path", ""), + "download_backend": "aria2", + }, + ) + install_result["settings"] = settings + return web.json_response(install_result) diff --git a/core/routes/metadata.py b/core/routes/metadata.py new file mode 100644 index 0000000..1ef4651 --- /dev/null +++ b/core/routes/metadata.py @@ -0,0 +1,199 @@ +"""Metadata route registration.""" + +def register_metadata_routes(context): + asyncio = context.get('asyncio') + audit_metadata_sizes = context.get('audit_metadata_sizes') + build_missing_local_metadata = context.get('build_missing_local_metadata') + cancel_progress_response = context.get('cancel_progress_response') + get_metadata_build_capabilities = context.get('get_metadata_build_capabilities') + get_model_files = context.get('get_model_files') + get_progress_response = context.get('get_progress_response') + invalidate_local_hash_match_cache = context.get('invalidate_local_hash_match_cache') + json_api_endpoint = context.get('json_api_endpoint') + normalize_metadata_build_mode = context.get('normalize_metadata_build_mode') + routes = context.get('routes') + run_in_background_thread = context.get('run_in_background_thread') + self = context.get('self') + to_bool = context.get('to_bool') + to_int = context.get('to_int') + web = context.get('web') + + @routes.get("/model_resolver/models") + @json_api_endpoint("get_models") + async def get_models(request): + """Get list of all available models.""" + force_rescan = to_bool( + request.query.get("force") or request.query.get("force_rescan"), + False + ) + if force_rescan: + invalidate_local_hash_match_cache() + models = get_model_files(force_rescan=force_rescan) + return web.json_response(models) + + @routes.post("/model_resolver/metadata-size-audit") + @json_api_endpoint("metadata-size-audit") + async def metadata_size_audit_route(request): + """Check local model .metadata.json sidecars for stale size values.""" + payload = {} + try: + if request.can_read_body: + payload = await request.json() + except Exception: + payload = {} + if not isinstance(payload, dict): + payload = {} + + force_rescan = to_bool(payload.get("force_rescan"), True) + result = await asyncio.to_thread( + audit_metadata_sizes, + force_rescan=force_rescan, + worker_count=payload.get("worker_count"), + batch_size=payload.get("batch_size"), + ) + return web.json_response(result) + + @routes.get("/model_resolver/metadata-build/capabilities") + @json_api_endpoint("metadata-build-capabilities") + async def metadata_build_capabilities_route(request): + """Return local CPU and worker limits for the metadata builder.""" + return web.json_response(get_metadata_build_capabilities()) + + @routes.post("/model_resolver/metadata-build/start") + @json_api_endpoint("metadata-build-start") + async def metadata_build_start_route(request): + """Create missing local metadata sidecars and fill missing SHA256 values.""" + import uuid + + payload = {} + try: + if request.can_read_body: + payload = await request.json() + except Exception: + payload = {} + if not isinstance(payload, dict): + payload = {} + + force_rescan = to_bool(payload.get("force_rescan"), True) + worker_count = to_int(payload.get("worker_count"), 0) + metadata_mode = normalize_metadata_build_mode( + payload.get("metadata_mode") + ) + self.metadata_builder_progress.cleanup() + progress_id = f"metadata_build_{uuid.uuid4().hex}" + self.metadata_builder_progress.update( + progress_id, + status="queued", + stage="queued", + message="Preparing local metadata build...", + percent=0, + current=0, + total=0, + requested_worker_count=worker_count, + metadata_mode=metadata_mode, + ) + + def update_metadata_build_progress(progress_payload): + self._update_metadata_build_progress( + progress_id, progress_payload + ) + + def is_metadata_build_cancelled(): + return self.metadata_builder_progress.is_cancelled(progress_id) + + def run_metadata_build_task(): + result = build_missing_local_metadata( + force_rescan=force_rescan, + worker_count=worker_count, + metadata_mode=metadata_mode, + progress_callback=update_metadata_build_progress, + is_cancelled=is_metadata_build_cancelled, + ) + return result + + def on_success(result): + self.metadata_builder_progress.update( + progress_id, + status="done", + stage="done", + message="Local metadata build completed.", + percent=100, + active_models=[], + active_worker_count=0, + current_model="", + current_path="", + bytes_read=0, + total_bytes=0, + result=result, + **result, + ) + + def on_cancel(result=None): + self.metadata_builder_progress.update( + progress_id, + status="cancelled", + stage="cancelled", + message="Metadata build cancelled", + percent=100, + active_models=[], + active_worker_count=0, + current_model="", + current_path="", + bytes_read=0, + total_bytes=0, + result=result or {}, + **(result or {}), + ) + + def on_error(exc): + self.metadata_builder_progress.update( + progress_id, + status="error", + stage="error", + message=str(exc) or "Metadata build failed", + percent=100, + active_models=[], + active_worker_count=0, + current_model="", + current_path="", + bytes_read=0, + total_bytes=0, + error=str(exc) or "Metadata build failed", + ) + + run_in_background_thread( + self.metadata_builder_progress, + progress_id, + run_metadata_build_task, + on_success, + on_cancel, + on_error, + error_log_msg="Metadata build failed", + ) + return web.json_response( + { + "success": True, + "progress_id": progress_id, + "metadata_mode": metadata_mode, + } + ) + + @routes.get("/model_resolver/metadata-build/progress/{progress_id}") + @json_api_endpoint("metadata-build-progress") + async def metadata_build_progress_route(request): + """Return progress for local metadata sidecar creation.""" + return get_progress_response( + self.metadata_builder_progress, + request, + not_found_status=404, + ) + + @routes.post("/model_resolver/metadata-build/cancel/{progress_id}") + @json_api_endpoint("metadata-build-cancel") + async def metadata_build_cancel_route(request): + """Cancel local metadata sidecar creation.""" + return cancel_progress_response( + self.metadata_builder_progress, + request, + cancel_message="Stopping metadata build...", + ) diff --git a/core/routes/model_info.py b/core/routes/model_info.py new file mode 100644 index 0000000..a797b40 --- /dev/null +++ b/core/routes/model_info.py @@ -0,0 +1,1258 @@ +"""Model Info route registration.""" + +def register_model_info_routes(context): + UnsafeUrlError = context.get('UnsafeUrlError') + asyncio = context.get('asyncio') + build_civarchive_custom_result = context.get('build_civarchive_custom_result') + build_civitai_custom_result = context.get('build_civitai_custom_result') + build_huggingface_custom_result = context.get('build_huggingface_custom_result') + download_available = context.get('download_available') + extract_sha256_from_metadata = context.get('extract_sha256_from_metadata') + find_external_metadata_sidecar_path = context.get('find_external_metadata_sidecar_path') + find_local_file_path = context.get('find_local_file_path') + get_civarchive_model_details = context.get('get_civarchive_model_details') + get_civitai_download_url = context.get('get_civitai_download_url') + get_civitai_model_details = context.get('get_civitai_model_details') + get_existing_model_preview_path = context.get('get_existing_model_preview_path') + get_filename_from_path = context.get('get_filename_from_path') + get_huggingface_model_details = context.get('get_huggingface_model_details') + get_model_resolver_sidecar_path = context.get('get_model_resolver_sidecar_path') + host_matches_domain = context.get('host_matches_domain') + is_path_in_configured_model_roots = context.get('is_path_in_configured_model_roots') + json_api_endpoint = context.get('json_api_endpoint') + looks_like_model_file = context.get('looks_like_model_file') + normalize_category_to_model_type = context.get('normalize_category_to_model_type') + normalize_sha256 = context.get('normalize_sha256') + parse_civarchive_url = context.get('parse_civarchive_url') + parse_civitai_url = context.get('parse_civitai_url') + read_json_safe = context.get('read_json_safe') + request_public_url = context.get('request_public_url') + resolve_civarchive_by_hash = context.get('resolve_civarchive_by_hash') + resolve_civarchive_model_version = context.get('resolve_civarchive_model_version') + resolve_civitai_version_custom_result = context.get('resolve_civitai_version_custom_result') + routes = context.get('routes') + search_huggingface_for_file = context.get('search_huggingface_for_file') + search_local_matches_by_hash = context.get('search_local_matches_by_hash') + self = context.get('self') + time = context.get('time') + to_bool = context.get('to_bool') + validate_public_http_url = context.get('validate_public_http_url') + web = context.get('web') + write_model_resolver_metadata = context.get('write_model_resolver_metadata') + + @routes.post("/model_resolver/civitai-search") + @json_api_endpoint("civitai-search") + async def civitai_search(request): + """Fetch model metadata from trusted exact-match sources.""" + data = await request.json() + filename = data.get("filename", "") + category = data.get("category", "") + resolved_path = data.get("resolved_path", "") + local_only = to_bool(data.get("local_only"), False) + force_refresh = to_bool( + data.get("force_refresh") or data.get("force"), False + ) + provided_hash = normalize_sha256( + data.get("sha256") + or data.get("hash") + or data.get("file_hash") + or "" + ) + hf_token = data.get("hf_token", "") + brave_search_api_key = data.get("brave_search_api_key", "") + hf_use_brave_fallback = to_bool( + data.get("hf_use_brave_fallback", True), + True, + ) + + 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] + + local_metadata_fields = ( + "author", + "creator", + "license", + "usage_hint", + "usage_tips", + "preview_url", + "metadata_source", + "from_safetensors_header", + "local_metadata_available", + "header_metadata_keys", + "metadata_summary", + "base_model_raw", + "sha256_source", + "hashes", + "hash_status", + ) + + def add_local_metadata_fields(payload, result): + if not isinstance(payload, dict) or not isinstance(result, dict): + return payload + for key in local_metadata_fields: + value = result.get(key) + if value in (None, "", [], {}): + continue + payload[key] = value + return payload + + # Get the file path to hash + file_path = resolved_path if resolved_path else None + file_location = "" + + if not file_path: + file_path = find_local_file_path(filename, category) + + if file_path: + file_path = _os.path.realpath( + _os.path.abspath(_os.path.normpath(str(file_path))) + ) + if not is_path_in_configured_model_roots(file_path): + if resolved_path: + return web.json_response( + { + "error": ( + "resolved_path is outside configured model directories" + ) + }, + status=403, + ) + file_path = None + elif not _os.path.isfile(file_path): + if resolved_path: + return web.json_response( + {"error": "resolved_path is not an existing model file"}, + status=404, + ) + file_path = None + + if file_path and _os.path.exists(file_path): + file_location = _os.path.dirname(file_path).replace("\\", "/") + if file_location and not file_location.endswith("/"): + file_location += "/" + + external_model_description = "" + if file_path: + external_metadata_path = find_external_metadata_sidecar_path( + file_path + ) + external_metadata = ( + read_json_safe(external_metadata_path, {}) + if external_metadata_path + else {} + ) + if isinstance(external_metadata, dict): + external_model_description = str( + external_metadata.get("modelDescription") + or external_metadata.get("model_description") + or external_metadata.get("description") + or "" + ).strip() + + def infer_model_type_from_category(value): + return normalize_category_to_model_type(value) + + def split_result_descriptions(result, preferred_model_description=""): + result = result if isinstance(result, dict) else {} + explicit_model_description = ( + result.get("model_description") + or result.get("modelDescription") + or "" + ) + generic_description = result.get("description") or "" + model_description = ( + preferred_model_description + or explicit_model_description + or generic_description + ) + version_description = ( + result.get("version_description") + or result.get("versionDescription") + or "" + ) + if ( + not version_description + and explicit_model_description + and generic_description + and generic_description != explicit_model_description + ): + version_description = generic_description + if version_description == model_description: + version_description = "" + return model_description, version_description + + def build_info_response( + result=None, + *, + metadata_path="", + metadata_saved=False, + civitai_checked=False, + local_payload=False, + ): + result = result or {} + model_description, version_description = ( + split_result_descriptions( + result, + external_model_description, + ) + ) + size_value = result.get("size") + if not size_value and file_path and _os.path.exists(file_path): + try: + size_value = _os.path.getsize(file_path) + except Exception: + size_value = None + + model_type = ( + result.get("model_type") + or result.get("type") + or infer_model_type_from_category(category) + ) + + response = { + "filename": result.get("filename") or filename, + "category": category, + "file_path": result.get("file_path") or file_path or "", + "resolved_path": result.get("resolved_path") or file_path or "", + "metadata_path": metadata_path + or result.get("metadata_path") + or "", + "metadata_saved": bool(metadata_saved), + "location": result.get("location") or file_location, + "source": result.get("source") or "", + "details_source": result.get("details_source") + or result.get("source") + or "", + "url": result.get("url"), + "version_url": result.get("version_url"), + "download_url": result.get("download_url"), + "platform_url": result.get("platform_url"), + "repo_id": result.get("repo_id"), + "path": result.get("path"), + "model_id": result.get("model_id"), + "model_name": result.get("model_name") or clean_name, + "model_type": model_type, + "version_id": result.get("version_id"), + "version_name": result.get("version_name", ""), + "sha256": result.get("sha256") or provided_hash, + "size": size_value, + "base_model": result.get("base_model"), + "base_model_source": result.get("base_model_source"), + "base_model_inferred": bool(result.get("base_model_inferred")), + "tags": result.get("tags", []), + "trained_words": result.get("trained_words", []), + "images": result.get("images", []), + "clip_skip": result.get("clip_skip"), + "description": model_description, + "model_description": model_description, + "version_description": version_description, + "from_metadata": bool(result.get("from_metadata")), + "local_only": bool(local_payload), + "metadata_checked": bool(civitai_checked), + "civitai_checked": bool(civitai_checked), + } + return add_local_metadata_fields(response, result) + def extract_result_sha256(result): + return extract_sha256_from_metadata(result) + + def result_filename_matches(result): + if not isinstance(result, dict): + return False + expected = get_filename_from_path(filename).lower() + candidates = [ + result.get("filename"), + result.get("path"), + result.get("file_path"), + ] + for candidate in candidates: + basename = get_filename_from_path(candidate).lower() + if basename and basename == expected: + return True + return False + + def result_hash_matches(result, *, require_filename=False): + if not provided_hash or not isinstance(result, dict): + return False + result_hash = extract_result_sha256(result) + if result_hash != provided_hash: + return False + return not require_filename or result_filename_matches(result) + + def huggingface_page_url(result): + try: + from urllib.parse import quote as _quote + except Exception: + _quote = None + repo_id = str( + result.get("repo_id") or result.get("repo") or "" + ).strip() + hf_path = str( + result.get("path") or result.get("filename") or "" + ).strip() + if not repo_id or not hf_path: + return "" + if _quote: + hf_path = _quote(hf_path.replace("\\", "/"), safe="/") + return f"https://huggingface.co/{repo_id}/blob/main/{hf_path}" + + def prepare_remote_result(result, source_name): + result = dict(result or {}) + source_name = str(source_name or result.get("source") or "").lower() + if source_name == "huggingface": + download_url = result.get("download_url") or result.get("url") + page_url = result.get("page_url") or huggingface_page_url(result) + if page_url: + result["url"] = page_url + result["version_url"] = page_url + result["download_url"] = download_url + result["model_name"] = ( + result.get("model_name") + or result.get("name") + or result.get("repo_id") + or clean_name + ) + result["model_type"] = ( + result.get("model_type") + or infer_model_type_from_category(category) + ) + else: + result["model_name"] = ( + result.get("model_name") + or result.get("name") + or clean_name + ) + if result.get("url") and not result.get("version_url"): + result["version_url"] = result.get("url") + + result["source"] = source_name or result.get("source") or "" + result["details_source"] = ( + result.get("details_source") + or source_name + or result.get("source") + or "" + ) + result["file_path"] = file_path or result.get("file_path") or "" + result["resolved_path"] = ( + file_path or result.get("resolved_path") or "" + ) + result["location"] = result.get("location") or file_location + result["filename"] = result.get("filename") or filename + result["sha256"] = ( + extract_result_sha256(result) + or provided_hash + or result.get("sha256") + ) + if not result.get("size") and file_path and _os.path.exists(file_path): + try: + result["size"] = _os.path.getsize(file_path) + except Exception: + pass + return result + + def result_has_remote_identity(result): + if not isinstance(result, dict): + return False + return bool( + result.get("url") + or result.get("version_url") + or result.get("download_url") + or result.get("platform_url") + or result.get("repo_id") + or result.get("model_id") + or result.get("version_id") + or result.get("from_metadata") + ) + + def remote_link_is_marked_dead(item): + if not isinstance(item, dict): + return False + status = str(item.get("status") or "").lower() + return bool( + item.get("deletedAt") + or item.get("deleted_at") + or item.get("is_dead") + or item.get("isDead") + or item.get("likelyDead") + or item.get("likely_dead") + or item.get("dead") + or status in {"dead", "deleted", "unavailable", "missing"} + ) + + def collect_result_download_urls(result): + urls = [] + if remote_link_is_marked_dead(result): + return urls + expected_filename = result.get("filename") or filename + dead_urls = set() + mirrors = result.get("mirrors") or [] + if not isinstance(mirrors, list): + mirrors = [mirrors] + for mirror in mirrors: + if not isinstance(mirror, dict): + continue + if remote_link_is_marked_dead(mirror): + dead_url = str(mirror.get("url") or "").strip() + if dead_url.startswith(("http://", "https://")): + dead_urls.add(dead_url) + continue + url = str(mirror.get("url") or "").strip() + mirror_filename = ( + mirror.get("filename") + or mirror.get("name") + or expected_filename + ) + if ( + looks_like_model_file(url, mirror_filename) + and url not in urls + ): + urls.append(url) + raw_urls = result.get("download_urls") or [] + if not isinstance(raw_urls, list): + raw_urls = [raw_urls] + for raw_url in raw_urls: + url = str(raw_url or "").strip() + if ( + looks_like_model_file(url, expected_filename) + and url not in dead_urls + and url not in urls + ): + urls.append(url) + for key in ("download_url", "downloadUrl"): + url = str(result.get(key) or "").strip() + if ( + looks_like_model_file(url, expected_filename) + and url not in dead_urls + and url not in urls + ): + urls.append(url) + return urls + + def remote_download_url_is_alive(url): + headers = { + "User-Agent": "ComfyUI-Model-Resolver/1.0", + "Accept": "*/*", + } + try: + response, _final_url, _final_headers = request_public_url( + "HEAD", + url, + headers=headers, + timeout=8, + ) + try: + if response.status_code < 400: + return True + if response.status_code in {401, 403, 404, 410}: + return False + finally: + response.close() + except Exception: + pass + + try: + response, _final_url, _final_headers = request_public_url( + "GET", + url, + headers={**headers, "Range": "bytes=0-0"}, + stream=True, + timeout=8, + ) + try: + return response.status_code < 400 + finally: + response.close() + except Exception: + return False + + def civarchive_result_has_live_download(result): + urls = collect_result_download_urls(result) + if not urls: + return False + for url in urls[:3]: + if remote_download_url_is_alive(url): + result["download_url"] = url + result["download_urls"] = [ + url, + *[other for other in urls if other != url], + ] + return True + return False + + def save_remote_metadata(result, source_name): + metadata_path = result.get("metadata_path") or "" + metadata_saved = False + if ( + result.get("from_metadata") + or not file_path + or not _os.path.exists(file_path) + ): + return metadata_path, metadata_saved + try: + model_description, version_description = ( + split_result_descriptions(result) + ) + metadata_payload = { + "source": source_name, + "details_source": result.get("details_source") + or source_name, + "filename": filename, + "category": category, + "model_name": result.get("model_name", clean_name), + "name": result.get("model_name", clean_name), + "model_type": result.get("model_type", "") + or result.get("type", ""), + "type": result.get("model_type", "") + or result.get("type", ""), + "model_id": result.get("model_id"), + "version_id": result.get("version_id"), + "version_name": result.get("version_name", ""), + "sha256": result.get("sha256") or provided_hash, + "size": result.get("size"), + "base_model": result.get("base_model"), + "base_model_source": result.get("base_model_source"), + "base_model_inferred": bool(result.get("base_model_inferred")), + "tags": result.get("tags", []), + "trained_words": result.get("trained_words", []), + "images": result.get("images", []), + "clip_skip": result.get("clip_skip"), + "description": model_description, + "model_description": model_description, + "version_description": version_description, + "download_url": result.get("download_url"), + "source_url": result.get("version_url") or result.get("url"), + "version_url": result.get("version_url") or result.get("url"), + "model_url": result.get("url"), + "url": result.get("version_url") or result.get("url"), + "platform_url": result.get("platform_url"), + "repo_id": result.get("repo_id"), + "path": result.get("path"), + "path_metadata": { + "filename": filename, + "category": category, + "source": source_name, + "model_id": result.get("model_id"), + "version_id": result.get("version_id"), + "repo_id": result.get("repo_id"), + "path": result.get("path"), + }, + } + add_local_metadata_fields(metadata_payload, result) + metadata_path = write_model_resolver_metadata( + file_path, + metadata_payload, + category, + result.get("version_url") + or result.get("url") + or result.get("platform_url") + or result.get("download_url") + or "", + create_preview=True, + ) or "" + metadata_saved = bool(metadata_path) + if metadata_path: + result["metadata_path"] = metadata_path + except Exception as metadata_error: + self.logger.warning( + f"{source_name} metadata sidecar save failed: {metadata_error}" + ) + return metadata_path, metadata_saved + + if local_only: + result = None + metadata_path = "" + metadata_saved = False + if download_available and file_path and _os.path.exists(file_path): + try: + from ..sources.civitai import ( + get_model_info_for_file, + ) + + result = get_model_info_for_file( + file_path, + local_only=True, + ) + source_metadata_path = ( + result.get("metadata_path") + if isinstance(result, dict) + else "" + ) + canonical_metadata_path = get_model_resolver_sidecar_path( + file_path + ) + if ( + isinstance(result, dict) + and result.get("from_metadata") + and source_metadata_path + and canonical_metadata_path + and _os.path.abspath(source_metadata_path) + != _os.path.abspath(canonical_metadata_path) + and not _os.path.exists(canonical_metadata_path) + ): + preview_url = result.get("preview_url") or "" + if not preview_url: + preview_path = get_existing_model_preview_path(file_path) + if preview_path: + preview_url = preview_path.replace("\\", "/") + model_description, version_description = ( + split_result_descriptions(result) + ) + metadata_payload = { + "source": "metadata_import", + "details_source": result.get("source") or "metadata", + "filename": filename, + "category": category, + "model_name": result.get("model_name", clean_name), + "name": result.get("model_name", clean_name), + "model_type": result.get("model_type", ""), + "type": result.get("model_type", ""), + "model_id": result.get("model_id"), + "version_id": result.get("version_id"), + "version_name": result.get("version_name", ""), + "sha256": result.get("sha256"), + "size": result.get("size"), + "base_model": result.get("base_model"), + "base_model_source": result.get("base_model_source"), + "base_model_inferred": bool(result.get("base_model_inferred")), + "tags": result.get("tags", []), + "trained_words": result.get("trained_words", []), + "images": result.get("images", []), + "clip_skip": result.get("clip_skip"), + "description": model_description, + "model_description": model_description, + "version_description": version_description, + "download_url": result.get("download_url"), + "preview_url": preview_url, + "source_url": result.get("version_url") + or result.get("url"), + "version_url": result.get("version_url") + or result.get("url"), + "model_url": result.get("url"), + "url": result.get("version_url") or result.get("url"), + "path_metadata": { + "filename": filename, + "category": category, + "source": "metadata_import", + "model_id": result.get("model_id"), + "version_id": result.get("version_id"), + "imported_from": source_metadata_path, + }, + } + add_local_metadata_fields(metadata_payload, result) + metadata_path = write_model_resolver_metadata( + file_path, + metadata_payload, + category, + result.get("version_url") + or result.get("url") + or result.get("download_url") + or "", + create_preview=True, + ) or "" + metadata_saved = bool(metadata_path) + if metadata_path: + result["metadata_path"] = metadata_path + result["metadata_imported_from"] = source_metadata_path + except Exception as e: + self.logger.warning(f"Local model info error: {e}") + return web.json_response( + build_info_response( + result, + metadata_path=metadata_path, + metadata_saved=metadata_saved, + local_payload=True, + civitai_checked=False, + ) + ) + + # Search remote metadata. CivitAI and CivArchive are hash-only; + # HuggingFace name results must still confirm the exact SHA. + if download_available and file_path and _os.path.exists(file_path): + try: + from ..sources.civitai import ( + get_model_info_by_hash, + get_model_info_for_file, + ) + + if provided_hash: + result = get_model_info_by_hash( + provided_hash, + use_cache=not force_refresh, + ) + if result: + result["file_path"] = file_path + result["resolved_path"] = file_path + result["location"] = file_location + if not result.get("size"): + try: + result["size"] = _os.path.getsize(file_path) + except Exception: + pass + if not extract_result_sha256(result): + result["sha256"] = provided_hash + elif not force_refresh: + result = get_model_info_for_file(file_path) + else: + result = None + + if result and ( + result.get("url") + or result.get("version_url") + or result.get("from_metadata") + or result.get("trained_words") + or result.get("base_model_inferred") + or result.get("local_metadata_available") + or result.get("from_safetensors_header") + ) and ( + result.get("from_metadata") + or not provided_hash + or result_hash_matches(result) + ): + result_source = ( + "local" + if result.get("source") == "local" + and not result_has_remote_identity(result) + else "civitai" + ) + result = prepare_remote_result(result, result_source) + metadata_path, metadata_saved = save_remote_metadata( + result, + result_source, + ) + + return web.json_response( + build_info_response( + result, + metadata_path=metadata_path, + metadata_saved=metadata_saved, + civitai_checked=True, + ) + ) + except Exception as e: + self.logger.warning(f"CivitAI search error: {e}") + + if provided_hash: + try: + result = resolve_civarchive_by_hash( + provided_hash, + query=filename, + exact_only=False, + model_type=infer_model_type_from_category(category), + ) + if result: + if not extract_result_sha256(result): + result["sha256"] = provided_hash + result = prepare_remote_result(result, "civarchive") + if result_hash_matches(result) and civarchive_result_has_live_download(result): + metadata_path, metadata_saved = save_remote_metadata( + result, + "civarchive", + ) + return web.json_response( + build_info_response( + result, + metadata_path=metadata_path, + metadata_saved=metadata_saved, + civitai_checked=True, + ) + ) + if result_hash_matches(result): + self.logger.info( + f"CivArchive metadata candidate rejected: " + f"no live download link for {filename}" + ) + except Exception as e: + self.logger.warning(f"CivArchive hash lookup error: {e}") + + if provided_hash: + hf_attempts = [ + { + "label": "HuggingFace", + "use_api_search": True, + "use_comfy_org_fallback": True, + "use_brave_fallback": False, + } + ] + if hf_use_brave_fallback and brave_search_api_key: + hf_attempts.append( + { + "label": "HuggingFace Brave", + "use_api_search": False, + "use_comfy_org_fallback": False, + "use_brave_fallback": True, + } + ) + + for hf_attempt in hf_attempts: + try: + result = search_huggingface_for_file( + filename, + token=hf_token or None, + exact_only=True, + brave_api_key=brave_search_api_key or None, + use_api_search=hf_attempt["use_api_search"], + use_comfy_org_fallback=hf_attempt[ + "use_comfy_org_fallback" + ], + use_brave_fallback=hf_attempt[ + "use_brave_fallback" + ], + force_refresh=force_refresh, + ) + if result and result_hash_matches( + result, + require_filename=True, + ): + result = prepare_remote_result( + result, + "huggingface", + ) + metadata_path, metadata_saved = save_remote_metadata( + result, + "huggingface", + ) + return web.json_response( + build_info_response( + result, + metadata_path=metadata_path, + metadata_saved=metadata_saved, + civitai_checked=True, + ) + ) + if result: + self.logger.info( + f"{hf_attempt['label']} metadata candidate rejected: " + f"filename/hash mismatch for {filename}" + ) + except Exception as e: + self.logger.warning( + f"{hf_attempt['label']} metadata lookup error: {e}" + ) + + # No remote result found. Keep any useful local safetensors + # header metadata, including base-model fingerprints, model + # titles, descriptions, tags, trigger words, and author data. + local_fallback_result = None + if file_path and _os.path.exists(file_path): + try: + from ..sources.civitai import get_model_info_for_file + + local_candidate = get_model_info_for_file( + file_path, + local_only=True, + ) + if ( + isinstance(local_candidate, dict) + and ( + local_candidate.get("base_model") + or local_candidate.get("local_metadata_available") + or local_candidate.get("from_safetensors_header") + ) + ): + local_fallback_result = local_candidate + except Exception as local_fallback_error: + self.logger.debug( + f"Local base model inference fallback failed: {local_fallback_error}" + ) + + response = build_info_response( + local_fallback_result, + civitai_checked=True, + local_payload=bool(local_fallback_result), + ) + response["url"] = None + if force_refresh and file_path and _os.path.exists(file_path): + try: + metadata_payload = { + "filename": filename, + "category": category, + "model_name": response.get("model_name") or clean_name, + "name": response.get("model_name") or clean_name, + "model_type": response.get("model_type") + or infer_model_type_from_category(category), + "sha256": response.get("sha256") or provided_hash, + "size": response.get("size") or _os.path.getsize(file_path), + "base_model": response.get("base_model"), + "base_model_source": response.get("base_model_source"), + "base_model_inferred": bool(response.get("base_model_inferred")), + "description": response.get("model_description") + or response.get("description") + or "", + "model_description": response.get("model_description") + or response.get("description") + or "", + "version_description": response.get( + "version_description" + ) + or "", + "civitai_deleted": True, + "civitai_checked": True, + "remote_metadata_missing": True, + "source": "local", + "details_source": response.get("details_source") or "", + } + add_local_metadata_fields(metadata_payload, response) + metadata_path = write_model_resolver_metadata( + file_path, + metadata_payload, + category, + "", + ) or "" + response["metadata_path"] = metadata_path + response["metadata_saved"] = bool(metadata_path) + except Exception as metadata_error: + self.logger.warning( + f"Remote metadata no-match sidecar save failed: {metadata_error}" + ) + return web.json_response(response) + + + + def _custom_result_timestamp(): + return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) + + + + + + + + def _build_direct_custom_download_result(url, source, expected_filename=""): + from urllib.parse import urlparse + + if not looks_like_model_file(url, expected_filename): + return None + + parsed = urlparse(url) + filename = ( + get_filename_from_path(parsed.path) + or get_filename_from_path(expected_filename) + or "model" + ) + if expected_filename and "." in get_filename_from_path(expected_filename): + filename = get_filename_from_path(expected_filename) + source = str(source or "custom").strip().lower() + source_label = { + "civarchive": "CivArchive", + "civitai": "CivitAI", + "huggingface": "HuggingFace", + }.get(source, "Custom URL") + return { + "source": source, + "details_source": source, + "name": source_label, + "filename": filename, + "url": url, + "version_url": url, + "download_url": url, + "match_type": "custom_url", + "custom_url": True, + } + + def _collect_custom_url_local_hash_matches(result, category): + source_key = str(result.get("source") or "custom").strip().lower() + sha256 = normalize_sha256(extract_sha256_from_metadata(result)) + if not sha256: + return [] + + try: + matches = search_local_matches_by_hash( + sha256, + category=category or None, + max_matches=20, + ) + except Exception as hash_error: + self.logger.warning( + f"Custom URL local metadata hash lookup failed for {source_key}:{sha256}: {hash_error}" + ) + return [] + + return [ + { + **match, + "hash_lookup_source": source_key, + "hash_lookup_filename": result.get("filename") + or result.get("path") + or "", + "hash_lookup_sha256": sha256, + } + for match in matches + ] + + @routes.post("/model_resolver/custom-url") + @json_api_endpoint("custom-url") + async def custom_url(request): + """Resolve a user-provided provider URL into a normal search result.""" + data = await request.json() + raw_url = str(data.get("url") or data.get("custom_url") or "").strip() + if not raw_url: + return web.json_response( + {"error": "URL is required"}, status=400 + ) + + if raw_url.startswith("hf://"): + normalized_url = raw_url + else: + try: + normalized_url = await asyncio.to_thread( + validate_public_http_url, + raw_url, + ) + except UnsafeUrlError as exc: + return web.json_response( + {"error": str(exc)}, + status=400, + ) + + category = data.get("category") or "" + expected_filename = ( + data.get("filename") + or get_filename_from_path(data.get("original_path") or "") + or "" + ) + civitai_key = data.get("civitai_key", "") + hf_token = data.get("hf_token", "") + + result = None + source = "" + try: + civitai_parsed = parse_civitai_url(normalized_url) + except Exception: + civitai_parsed = None + try: + civarchive_parsed = parse_civarchive_url(normalized_url) + except Exception: + civarchive_parsed = None + + if civitai_parsed: + source = "civitai" + model_id = civitai_parsed.get("model_id") + version_id = civitai_parsed.get("version_id") + if model_id: + details = await asyncio.to_thread( + get_civitai_model_details, + model_id, + version_id, + civitai_key or None, + ) + result = build_civitai_custom_result( + details, + expected_filename=expected_filename, + api_key=civitai_key or None, + ) + elif version_id: + result = await asyncio.to_thread( + resolve_civitai_version_custom_result, + version_id, + expected_filename, + civitai_key or None, + ) + if not result and version_id: + result = { + "source": "civitai", + "details_source": "civitai", + "version_id": version_id, + "name": expected_filename or f"CivitAI version {version_id}", + "filename": expected_filename or f"civitai-{version_id}", + "url": normalized_url, + "version_url": normalized_url, + "download_url": get_civitai_download_url( + version_id, + civitai_key or None, + ), + "match_type": "custom_url", + "custom_url": True, + } + elif civarchive_parsed: + source = "civarchive" + if civarchive_parsed.get("sha256"): + result = await asyncio.to_thread( + resolve_civarchive_by_hash, + civarchive_parsed.get("sha256"), + expected_filename, + False, + normalize_category_to_model_type(category), + ) + else: + model_id = civarchive_parsed.get("model_id") + version_id = civarchive_parsed.get("version_id") + if model_id: + result = await asyncio.to_thread( + resolve_civarchive_model_version, + model_id, + version_id, + expected_filename or str(model_id), + False, + True, + ) + if not result: + details = await asyncio.to_thread( + get_civarchive_model_details, + model_id, + version_id, + True, + ) + result = build_civarchive_custom_result( + details, + expected_filename=expected_filename, + ) + if result: + result = dict(result) + result["source"] = "civarchive" + result["details_source"] = "civarchive" + result["match_type"] = "custom_url" + result["custom_url"] = True + else: + from urllib.parse import urlparse + + parsed_direct = urlparse(normalized_url) + direct_host = parsed_direct.hostname + if ( + host_matches_domain(direct_host, "civarchive.com") + and "/api/download/" in parsed_direct.path + ): + source = "civarchive" + result = _build_direct_custom_download_result( + normalized_url, + "civarchive", + expected_filename, + ) + else: + result = await asyncio.to_thread( + build_huggingface_custom_result, + normalized_url, + expected_filename, + hf_token or None, + ) + if result: + source = "huggingface" + + if not result: + return web.json_response( + { + "error": ( + "Unsupported or unresolved URL. Use a HuggingFace file URL, " + "CivitAI model/download URL, or CivArchive model/hash URL." + ) + }, + status=400, + ) + + result = dict(result) + result["provided_url"] = normalized_url + result["url_source"] = "custom" + result["searched_at"] = result.get("searched_at") or _custom_result_timestamp() + result.setdefault("category", category) + if expected_filename and not result.get("filename"): + result["filename"] = expected_filename + + if not (result.get("download_url") or result.get("url")): + return web.json_response( + {"error": "The URL resolved, but no download URL was found"}, + status=400, + ) + + source = source or result.get("source") or "custom" + local_hash_matches = _collect_custom_url_local_hash_matches( + result, + category, + ) + response = { + "success": True, + "source": source, + "result": result, + "custom": [result], + "searched_sources": ["custom"], + "local_hash_matches": local_hash_matches, + } + response[source] = result + return web.json_response(response) + + @routes.post("/model_resolver/model-details") + @json_api_endpoint("model-details") + async def model_details(request): + """Return normalized full model details for sources that expose model pages.""" + data = await request.json() + source = str(data.get("source", "")).strip().lower() + model_id = data.get("model_id") + version_id = data.get("version_id") + civitai_key = data.get("civitai_key", "") + hf_token = data.get("hf_token", "") + file_path = data.get("file_path", "") + branch = data.get("branch", "") + + if not download_available: + return web.json_response( + {"error": "Download providers are not available"}, status=503 + ) + + if source == "lora_manager_archive": + source = "civitai" + + if source not in {"civitai", "civarchive", "huggingface"}: + return web.json_response( + {"error": "Unsupported model details source"}, status=400 + ) + + if source == "huggingface": + model_id = str(model_id or "").strip() + branch = str(branch or version_id or "main").strip() or "main" + version_id = branch + else: + try: + model_id = ( + int(model_id) + if model_id is not None and str(model_id).strip() + else None + ) + except (TypeError, ValueError): + model_id = None + + try: + version_id = ( + int(version_id) + if version_id is not None and str(version_id).strip() + else None + ) + except (TypeError, ValueError): + version_id = None + + if not model_id: + return web.json_response( + {"error": "model_id is required"}, status=400 + ) + + if source == "civitai": + details = await asyncio.to_thread( + get_civitai_model_details, + model_id, + version_id, + civitai_key or None, + ) + elif source == "civarchive": + details = await asyncio.to_thread( + get_civarchive_model_details, + model_id, + version_id, + ) + else: + details = await asyncio.to_thread( + get_huggingface_model_details, + model_id, + file_path, + branch, + hf_token or None, + ) + + if not details: + return web.json_response( + {"error": "Model details not found"}, status=404 + ) + + return web.json_response(details) diff --git a/core/routes/registry.py b/core/routes/registry.py index 2970130..ec4900d 100644 --- a/core/routes/registry.py +++ b/core/routes/registry.py @@ -1,5 +1,8 @@ """Central route registration for the Model Resolver backend.""" +# Route modules receive these optional imports through the registration context. +# ruff: noqa: F401 + import asyncio import os import time @@ -74,6 +77,9 @@ def register_routes(self): split_path_segments, write_json_atomic, ) + from ..path_utils import ( + get_path_identity as get_local_path_identity, + ) from ..resolver import ( analyze_and_find_matches, apply_resolution, @@ -115,9 +121,18 @@ def register_routes(self): get_workflow_model_inventory, ) from .base_models import register_base_model_routes + from .directories import register_directory_routes + from .downloads import register_download_routes from .helpers import create_route_helpers + from .metadata import register_metadata_routes + from .model_info import register_model_info_routes + from .search import register_search_routes from .settings import register_settings_routes from .version import register_version_routes + from .workflow import ( + register_loaded_model_routes, + register_workflow_routes, + ) except ImportError as e: self.logger.error(f"Model Resolver: Could not import core modules: {e}") return False @@ -225,4503 +240,16 @@ def register_routes(self): ) register_base_model_routes(routes, web, json_api_endpoint) - - # ==================== ANALYZE ROUTES ==================== - - @routes.post("/model_resolver/analyze") - async def analyze_workflow(request): - """Analyze workflow and return missing models with matches.""" - try: - data = await request.json() - workflow_json = data.get("workflow") - analysis_id = str(data.get("analysis_id") or "").strip() - force_rescan = to_bool(data.get("force_rescan"), False) - if force_rescan: - invalidate_local_hash_match_cache() - - if workflow_json is None: - return web.json_response( - {"error": "Workflow JSON is required"}, status=400 - ) - if not isinstance(workflow_json, dict): - return web.json_response( - {"error": "Workflow JSON must be an object"}, status=400 - ) - - if analysis_id: - self._update_analysis_progress( - analysis_id, - { - "status": "starting", - "stage": "starting", - "message": "Starting analysis...", - "current": 0, - "total": 0, - }, - ) - - def update_analysis_progress(payload): - self._update_analysis_progress(analysis_id, payload) - - # Analyze and find matches - result = await asyncio.to_thread( - analyze_and_find_matches, - workflow_json, - 0.0, - 10, - update_analysis_progress if analysis_id else None, - force_rescan=force_rescan, - ) - - missing_models = result.get("missing_models", []) - filtered_missing = [] - for missing in missing_models: - name = missing.get("name") or missing.get("original_path", "") - if should_skip_existing_custom_node_reference(missing): - self.logger.info( - "Filtered existing custom-node model " - f"reference: {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 = get_filename_from_path( - missing.get("original_path", "") - ) - - # 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 using cached remote helper - file_size = fetch_remote_file_size_cached(workflow_url, timeout=5) - - missing["download_source"] = { - "source": source, - "url": workflow_url, - "model_url": missing.get( - "workflow_model_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: - popular_model_list_result = search_model_list( - filename, exact_only=True - ) - missing["download_source"] = { - "source": "popular", - "url": popular_info.get("url"), - "filename": filename, - "type": popular_info.get("type"), - "directory": popular_info.get("directory"), - "size": ( - popular_model_list_result.get("size") - if popular_model_list_result - else None - ) - or popular_info.get("size"), - "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_resolver/search endpoint - # when user clicks "Search Online" button, not automatically - - if analysis_id: - self.analysis_progress.update( - analysis_id, - status="completed", - stage="completed", - message="Analysis complete", - current=result.get("total_missing", 0), - total=result.get("total_missing", 0), - ) - - return web.json_response(result) - except Exception as e: - if "analysis_id" in locals() and analysis_id: - self.analysis_progress.update( - analysis_id, - status="error", - stage="error", - message=str(e), - current=0, - total=0, - ) - self.logger.error(f"Model Resolver analyze error: {e}", exc_info=True) - return web.json_response({"error": str(e)}, status=500) - - @routes.get("/model_resolver/analyze-progress/{analysis_id}") - @json_api_endpoint("analyze-progress") - async def get_analyze_progress(request): - """Get workflow analysis progress.""" - return get_progress_response( - self.analysis_progress, - request, - param_name="analysis_id", - not_found_payload={ - "status": "unknown", - "stage": "unknown", - "message": "No analysis progress available", - "current": 0, - "total": 0, - } - ) - - @routes.post("/model_resolver/resolve") - @json_api_endpoint("resolve", return_success_on_error=True) - async def resolve_models(request): - """Apply model resolution and return updated workflow.""" - 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} - ) - - @routes.post("/model_resolver/local-matches") - @json_api_endpoint("local-matches") - async def local_matches(request): - """Search local model files by filename/path.""" - data = await request.json() - filename = data.get("filename", "") - category = data.get("category", "") - force_rescan = to_bool(data.get("force_rescan"), False) - if force_rescan: - invalidate_local_hash_match_cache() - - if not filename: - return web.json_response( - {"error": "filename is required"}, status=400 - ) - - matches = search_local_matches( - filename, - category=category or None, - similarity_threshold=0.0, - max_matches_per_model=10, - force_rescan=force_rescan, - ) - return web.json_response({"matches": matches}) - - @routes.post("/model_resolver/local-model-hashes") - @json_api_endpoint("local-model-hashes") - async def local_model_hashes(request): - """Return SHA256 hashes already stored in local sidecar metadata.""" - data = await request.json() - model = data.get("model") if isinstance(data.get("model"), dict) else {} - path = ( - data.get("path") - or data.get("file_path") - or data.get("resolved_path") - or model.get("path") - or model.get("resolved_path") - or "" - ) - - if not path: - return web.json_response( - {"error": "path is required"}, status=400 - ) - - normalized_path = os.path.realpath( - os.path.abspath(os.path.normpath(str(path))) - ) - if not is_path_in_configured_model_roots(normalized_path): - return web.json_response( - {"error": "path is outside configured model directories"}, - status=403, - ) - - return web.json_response( - get_local_model_hash_metadata(normalized_path, model=model) - ) - - @routes.get("/model_resolver/model-preview") - async def get_model_preview(request): - """Serve an adjacent model preview from a configured model directory.""" - model_path = str(request.query.get("path") or "").strip() - if not model_path: - return web.Response(text="path is required", status=400) - - try: - normalized_path = os.path.realpath( - os.path.abspath(os.path.normpath(model_path)) - ) - except (OSError, TypeError, ValueError): - return web.Response(text="invalid model path", status=400) - if not os.path.isfile(normalized_path): - return web.Response(text="model file does not exist", status=404) - if not is_path_in_configured_model_roots(normalized_path): - return web.Response( - text="path is outside configured model directories", - status=403, - ) - - preview_path = get_existing_model_preview_path(normalized_path) - if not preview_path: - return web.Response(text="preview not found", status=404) - if not is_path_in_configured_model_roots(preview_path): - return web.Response( - text="preview is outside configured model directories", - status=403, - ) - - return web.FileResponse( - preview_path, - headers={"Cache-Control": "private, no-cache"}, - ) - - @routes.post("/model_resolver/workflow-model-hashes") - @json_api_endpoint("workflow-model-hashes") - async def workflow_model_hashes(request): - """Return hash metadata for existing local models used by a workflow.""" - data = await request.json() - workflow_json = data.get("workflow") - if not isinstance(workflow_json, dict): - return web.json_response( - {"error": "Workflow JSON must be an object"}, status=400 - ) - - settings = await asyncio.to_thread(load_resolver_settings) - if not resolver_bool_setting( - settings.get("workflow_hash_metadata_enabled"), True - ): - return web.json_response( - { - "success": True, - "enabled": False, - "models": [], - "by_node": {}, - "by_path": {}, - "count": 0, - } - ) - - inventory = await asyncio.to_thread( - get_workflow_model_inventory, - workflow_json, - ) - refs = inventory["model_refs"] - - by_node = {} - by_path = {} - models = [] - seen = set() - - for ref in refs: - if not isinstance(ref, dict) or not ref.get("exists"): - continue - full_path = str(ref.get("full_path") or "").strip() - if not full_path: - continue - model_info = { - "path": full_path, - "filename": get_filename_from_path(full_path), - "relative_path": ref.get("original_path") or "", - "category": ref.get("category") or "", - } - metadata = get_local_model_hash_metadata( - full_path, model=model_info - ) - sha256 = normalize_sha256(metadata.get("sha256")) - if not sha256: - continue - - entry = { - "node_id": ref.get("node_id"), - "node_type": ref.get("node_type") or "", - "widget_index": ref.get("widget_index"), - "widget_name": ref.get("widget_name") or "", - "path": ref.get("original_path") or "", - "filename": get_filename_from_path( - ref.get("original_path") or full_path - ), - "category": ref.get("category") or "", - "sha256": sha256, - "size": metadata.get("size") or 0, - } - entry_key = ( - str(entry.get("node_id")), - str(entry.get("widget_index")), - entry.get("path"), - sha256, - ) - if entry_key in seen: - continue - seen.add(entry_key) - models.append(entry) - - path_key = str(entry.get("path") or entry.get("filename") or "") - if path_key: - by_path[path_key] = entry - by_path[get_filename_from_path(path_key)] = entry - node_key = f"{entry.get('node_id')}:{entry.get('widget_index')}" - by_node[node_key] = entry - - return web.json_response( - { - "success": True, - "enabled": True, - "models": models, - "by_node": by_node, - "by_path": by_path, - "count": len(models), - } - ) - - @routes.post("/model_resolver/local-matches-by-hash") - @json_api_endpoint("local-matches-by-hash") - async def local_matches_by_hash(request): - """Search local model metadata sidecars for a remote SHA256.""" - data = await request.json() - sha256 = normalize_sha256( - data.get("sha256") - or data.get("hash") - or data.get("SHA256") - or "" - ) - if not sha256: - return web.json_response( - {"error": "sha256 is required"}, status=400 - ) - - category = data.get("category", "") - source = str( - data.get("source") - or data.get("hash_lookup_source") - or "download_source" - ).strip().lower().replace("-", "_") - filename = ( - data.get("filename") - or data.get("path") - or data.get("model_name") - or "" - ) - max_matches = to_int(data.get("max_matches"), 20) - force_rescan = to_bool(data.get("force_rescan"), False) - - matches = search_local_matches_by_hash( - sha256, - category=category or None, - max_matches=max_matches, - force_rescan=force_rescan, - ) - enriched_matches = [ - { - **match, - "hash_lookup_source": source or "download_source", - "hash_lookup_filename": filename, - "hash_lookup_sha256": sha256, - } - for match in matches - ] - return web.json_response( - { - "sha256": sha256, - "local_hash_matches": enriched_matches, - "matches": enriched_matches, - } - ) - - @routes.post("/model_resolver/open-containing-folder") - @json_api_endpoint("open-containing-folder") - async def open_containing_folder(request): - """Reveal a file or open a directory in the host file manager.""" - try: - data = await request.json() - except Exception as exc: - return web.json_response( - {"success": False, "error": f"invalid JSON body: {exc}"}, - status=400, - ) - if not isinstance(data, dict): - return web.json_response( - {"success": False, "error": "JSON body must be an object"}, - status=400, - ) - target_path = data.get("path", "") - - try: - normalized_path = normalize_file_manager_path(target_path) - except ValueError as exc: - return web.json_response( - {"success": False, "error": str(exc)}, status=400 - ) - - if not os.path.exists(normalized_path): - return web.json_response( - {"success": False, "error": "path does not exist"}, status=404 - ) - if not is_path_in_configured_model_roots(normalized_path): - return web.json_response( - { - "success": False, - "error": "path is outside configured model directories", - }, - status=403, - ) - - try: - result = await asyncio.to_thread( - open_in_file_manager, normalized_path - ) - except FileNotFoundError: - return web.json_response( - {"success": False, "error": "path no longer exists"}, - status=404, - ) - except UnsupportedFileManagerPlatformError as exc: - return web.json_response( - {"success": False, "error": str(exc)}, status=501 - ) - except FileManagerUnavailableError as exc: - return web.json_response( - { - "success": False, - "error": ( - f"{exc} The ComfyUI host may be running without " - "a graphical desktop session." - ), - }, - status=503, - ) - except FileManagerError as exc: - return web.json_response( - {"success": False, "error": str(exc)}, status=500 - ) - - return web.json_response({"success": True, **result}) - - # cleanup_hash_progress, update_hash_progress, is_hash_progress_cancelled, and mark_hash_progress_cancelled removed - - from ..path_utils import HashCalculationCancelled - - def resolve_hash_file_request(data): - import os as _os - - file_path = ( - data.get("file_path") - or data.get("resolved_path") - or data.get("path") - or "" - ) - if not file_path: - return "", "file_path is required" - - normalized_path = _os.path.realpath( - _os.path.abspath(_os.path.normpath(file_path)) - ) - if not _os.path.exists(normalized_path) or not _os.path.isfile(normalized_path): - return "", "file does not exist" - if not is_path_in_configured_model_roots(normalized_path): - return "", "file is outside configured model directories" - return normalized_path, "" - - def write_calculated_hash_metadata( - normalized_path, - sha256, - sha256_source="file", - ): - import os as _os - - resolved_metadata_path = get_safe_model_resolver_sidecar_path( - normalized_path - ) - - metadata_updated = False - try: - metadata = read_json_safe(resolved_metadata_path, {}) - if not isinstance(metadata, dict): - metadata = {} - - filename = get_filename_from_path(normalized_path) - stem, _ext = _os.path.splitext(filename) - hashes = metadata.get("hashes") - if not isinstance(hashes, dict): - hashes = {} - hashes["SHA256"] = sha256 - - metadata["schema"] = MODEL_RESOLVER_METADATA_SCHEMA - metadata["schema_version"] = ( - MODEL_RESOLVER_METADATA_SCHEMA_VERSION - ) - metadata["managed_by"] = MODEL_RESOLVER_METADATA_SCHEMA - metadata["sha256"] = sha256 - metadata["hashes"] = hashes - metadata["hash_status"] = "completed" - metadata["sha256_source"] = sha256_source or "file" - metadata["last_checked_at"] = time.time() - metadata.setdefault("file_name", stem) - metadata.setdefault("model_name", stem) - metadata.setdefault("file_path", normalized_path.replace("\\", "/")) - try: - metadata.setdefault("size", _os.path.getsize(normalized_path)) - except Exception: - pass - - write_json_atomic(resolved_metadata_path, metadata, indent=2) - metadata_updated = True - self.logger.info( - f"Stored SHA256 and updated metadata: {resolved_metadata_path}" - ) - except Exception as metadata_error: - self.logger.warning( - f"Could not update metadata with calculated SHA256 for {normalized_path}: {metadata_error}" - ) - - return resolved_metadata_path, metadata_updated - - def calculate_sha256_with_progress(normalized_path, progress_id=""): - import os as _os - - from ..path_utils import calculate_file_sha256 as _calculate_file_sha256_core - - total_bytes = max(0, _os.path.getsize(normalized_path)) - - self.hash_tracker.update( - progress_id, - status="running", - stage="header", - message="Checking safetensors header...", - percent=0, - bytes_read=0, - total_bytes=total_bytes, - ) - - if self.hash_tracker.is_cancelled(progress_id): - raise HashCalculationCancelled() - - hash_source = ["file"] - stage_transitioned = [False] - _bytes_read = [0] - _last_update = [0.0] - - def on_progress(bytes_read, total_bytes): - _bytes_read[0] = bytes_read - if not stage_transitioned[0]: - stage_transitioned[0] = True - self.logger.info( - "No SHA256 found in safetensors header; calculating full " - f"file SHA256: {normalized_path}" - ) - self.hash_tracker.update( - progress_id, - status="running", - stage="hashing", - message="Calculating SHA256...", - percent=0, - bytes_read=0, - total_bytes=total_bytes, - ) - - now = time.time() - if now - _last_update[0] >= 0.15 or bytes_read >= total_bytes: - percent = 98 if total_bytes <= 0 else min( - 98, - (bytes_read / total_bytes) * 98, - ) - self.hash_tracker.update( - progress_id, - status="running", - stage="hashing", - message="Calculating SHA256...", - percent=percent, - bytes_read=bytes_read, - total_bytes=total_bytes, - ) - _last_update[0] = now - - def is_cancelled(): - if self.hash_tracker.is_cancelled(progress_id): - percent = 0 if total_bytes <= 0 else min( - 98, - (_bytes_read[0] / total_bytes) * 98, - ) - self.hash_tracker.update( - progress_id, - status="cancelled", - stage="cancelled", - message="Hash calculation cancelled", - percent=percent, - bytes_read=_bytes_read[0], - total_bytes=total_bytes, - ) - return True - return False - - def on_hash_source(source_name): - hash_source[0] = source_name - if source_name == "safetensors_header": - self.logger.info( - f"Found SHA256 in safetensors header: {normalized_path}" - ) - self.hash_tracker.update( - progress_id, - status="running", - stage="header", - message="Found SHA256 in safetensors header", - percent=98, - bytes_read=0, - total_bytes=total_bytes, - sha256_source="safetensors_header", - ) - - sha256 = _calculate_file_sha256_core( - normalized_path, - chunk_size=1024 * 1024 * 4, - on_progress=on_progress, - is_cancelled=is_cancelled, - use_safetensors_header=True, - on_hash_source=on_hash_source, - ) - - if is_cancelled(): - raise HashCalculationCancelled() - - return sha256, hash_source[0] - - @routes.post("/model_resolver/calculate-file-hash") - @json_api_endpoint("calculate-file-hash") - async def calculate_file_hash_route(request): - """Calculate SHA256 for a local model and persist it to sidecar metadata.""" - data = await request.json() - - normalized_path, error = resolve_hash_file_request(data) - if error == "file_path is required": - return web.json_response( - {"error": "file_path is required"}, status=400 - ) - if error: - status = 403 if error == "file is outside configured model directories" else 404 - return web.json_response( - {"error": error}, status=status - ) - - sha256, sha256_source = calculate_sha256_with_progress(normalized_path) - if not sha256: - return web.json_response( - {"error": "could not calculate hash"}, status=500 - ) - resolved_metadata_path, metadata_updated = write_calculated_hash_metadata( - normalized_path, - sha256, - sha256_source, - ) - - return web.json_response( - { - "success": True, - "sha256": sha256, - "hash": sha256, - "sha256_source": sha256_source, - "file_path": normalized_path, - "metadata_path": resolved_metadata_path, - "metadata_updated": metadata_updated, - } - ) - - @routes.post("/model_resolver/calculate-file-hash/start") - @json_api_endpoint("calculate-file-hash-start") - async def calculate_file_hash_start_route(request): - """Start SHA256 calculation in a background thread.""" - import uuid - - data = await request.json() - normalized_path, error = resolve_hash_file_request(data) - if error == "file_path is required": - return web.json_response( - {"error": "file_path is required"}, status=400 - ) - if error: - status = 403 if error == "file is outside configured model directories" else 404 - return web.json_response( - {"error": error}, status=status - ) - - self.hash_tracker.cleanup() - progress_id = f"hash_{uuid.uuid4().hex}" - self.hash_tracker.update( - progress_id, - status="queued", - stage="queued", - message="Preparing hash calculation...", - percent=0, - file_path=normalized_path, - ) - - def run_hash_task(): - sha256, sha256_source = calculate_sha256_with_progress( - normalized_path, - progress_id=progress_id, - ) - if self.hash_tracker.is_cancelled(progress_id): - raise HashCalculationCancelled() - self.hash_tracker.update( - progress_id, - status="running", - stage="metadata", - message="Saving metadata...", - percent=99, - ) - resolved_metadata_path, metadata_updated = write_calculated_hash_metadata( - normalized_path, - sha256, - sha256_source, - ) - return sha256, sha256_source, resolved_metadata_path, metadata_updated - - def on_success(res): - sha256, sha256_source, resolved_metadata_path, metadata_updated = res - self.hash_tracker.update( - progress_id, - status="done", - stage="done", - message="Hash calculated", - percent=100, - sha256=sha256, - hash=sha256, - sha256_source=sha256_source, - file_path=normalized_path, - metadata_path=resolved_metadata_path, - metadata_updated=metadata_updated, - ) - - def on_cancel(res=None): - self.hash_tracker.update( - progress_id, - status="cancelled", - stage="cancelled", - message="Hash calculation cancelled", - ) - - run_in_background_thread( - self.hash_tracker, - progress_id, - run_hash_task, - on_success, - on_cancel, - error_log_msg=f"Hash calculation failed for {normalized_path}", - ) - return web.json_response( - { - "success": True, - "progress_id": progress_id, - } - ) - - @routes.get("/model_resolver/calculate-file-hash/progress/{progress_id}") - @json_api_endpoint("calculate-file-hash-progress") - async def calculate_file_hash_progress_route(request): - """Return progress for a background SHA256 calculation.""" - return get_progress_response( - self.hash_tracker, - request, - not_found_status=404 - ) - - @routes.post("/model_resolver/calculate-file-hash/cancel/{progress_id}") - @json_api_endpoint("calculate-file-hash-cancel") - async def calculate_file_hash_cancel_route(request): - """Cancel a background SHA256 calculation.""" - return cancel_progress_response( - self.hash_tracker, - request, - cancel_message="Stopping hash calculation..." - ) - - @routes.get("/model_resolver/models") - @json_api_endpoint("get_models") - async def get_models(request): - """Get list of all available models.""" - force_rescan = to_bool( - request.query.get("force") or request.query.get("force_rescan"), - False - ) - if force_rescan: - invalidate_local_hash_match_cache() - models = get_model_files(force_rescan=force_rescan) - return web.json_response(models) - - @routes.post("/model_resolver/metadata-size-audit") - @json_api_endpoint("metadata-size-audit") - async def metadata_size_audit_route(request): - """Check local model .metadata.json sidecars for stale size values.""" - payload = {} - try: - if request.can_read_body: - payload = await request.json() - except Exception: - payload = {} - if not isinstance(payload, dict): - payload = {} - - force_rescan = to_bool(payload.get("force_rescan"), True) - result = await asyncio.to_thread( - audit_metadata_sizes, - force_rescan=force_rescan, - worker_count=payload.get("worker_count"), - batch_size=payload.get("batch_size"), - ) - return web.json_response(result) - - @routes.get("/model_resolver/metadata-build/capabilities") - @json_api_endpoint("metadata-build-capabilities") - async def metadata_build_capabilities_route(request): - """Return local CPU and worker limits for the metadata builder.""" - return web.json_response(get_metadata_build_capabilities()) - - @routes.post("/model_resolver/metadata-build/start") - @json_api_endpoint("metadata-build-start") - async def metadata_build_start_route(request): - """Create missing local metadata sidecars and fill missing SHA256 values.""" - import uuid - - payload = {} - try: - if request.can_read_body: - payload = await request.json() - except Exception: - payload = {} - if not isinstance(payload, dict): - payload = {} - - force_rescan = to_bool(payload.get("force_rescan"), True) - worker_count = to_int(payload.get("worker_count"), 0) - metadata_mode = normalize_metadata_build_mode( - payload.get("metadata_mode") - ) - self.metadata_builder_progress.cleanup() - progress_id = f"metadata_build_{uuid.uuid4().hex}" - self.metadata_builder_progress.update( - progress_id, - status="queued", - stage="queued", - message="Preparing local metadata build...", - percent=0, - current=0, - total=0, - requested_worker_count=worker_count, - metadata_mode=metadata_mode, - ) - - def update_metadata_build_progress(progress_payload): - self._update_metadata_build_progress( - progress_id, progress_payload - ) - - def is_metadata_build_cancelled(): - return self.metadata_builder_progress.is_cancelled(progress_id) - - def run_metadata_build_task(): - result = build_missing_local_metadata( - force_rescan=force_rescan, - worker_count=worker_count, - metadata_mode=metadata_mode, - progress_callback=update_metadata_build_progress, - is_cancelled=is_metadata_build_cancelled, - ) - return result - - def on_success(result): - self.metadata_builder_progress.update( - progress_id, - status="done", - stage="done", - message="Local metadata build completed.", - percent=100, - active_models=[], - active_worker_count=0, - current_model="", - current_path="", - bytes_read=0, - total_bytes=0, - result=result, - **result, - ) - - def on_cancel(result=None): - self.metadata_builder_progress.update( - progress_id, - status="cancelled", - stage="cancelled", - message="Metadata build cancelled", - percent=100, - active_models=[], - active_worker_count=0, - current_model="", - current_path="", - bytes_read=0, - total_bytes=0, - result=result or {}, - **(result or {}), - ) - - def on_error(exc): - self.metadata_builder_progress.update( - progress_id, - status="error", - stage="error", - message=str(exc) or "Metadata build failed", - percent=100, - active_models=[], - active_worker_count=0, - current_model="", - current_path="", - bytes_read=0, - total_bytes=0, - error=str(exc) or "Metadata build failed", - ) - - run_in_background_thread( - self.metadata_builder_progress, - progress_id, - run_metadata_build_task, - on_success, - on_cancel, - on_error, - error_log_msg="Metadata build failed", - ) - return web.json_response( - { - "success": True, - "progress_id": progress_id, - "metadata_mode": metadata_mode, - } - ) - - @routes.get("/model_resolver/metadata-build/progress/{progress_id}") - @json_api_endpoint("metadata-build-progress") - async def metadata_build_progress_route(request): - """Return progress for local metadata sidecar creation.""" - return get_progress_response( - self.metadata_builder_progress, - request, - not_found_status=404, - ) - - @routes.post("/model_resolver/metadata-build/cancel/{progress_id}") - @json_api_endpoint("metadata-build-cancel") - async def metadata_build_cancel_route(request): - """Cancel local metadata sidecar creation.""" - return cancel_progress_response( - self.metadata_builder_progress, - request, - cancel_message="Stopping metadata build...", - ) - - @routes.post("/model_resolver/loaded") - @json_api_endpoint("get_loaded_models") - async def get_loaded_models(request): - """Get all currently loaded models in the workflow.""" - data = await request.json() - workflow_json = data.get("workflow") - loaded_id = str( - data.get("loaded_id") or data.get("progress_id") or "" - ).strip() - - if not workflow_json: - return web.json_response( - {"error": "Workflow JSON is required"}, status=400 - ) - if not isinstance(workflow_json, dict): - return web.json_response( - {"error": "Workflow JSON must be an object"}, status=400 - ) - - def update_loaded_progress(*args, **kwargs): - self._update_loaded_progress(loaded_id, *args, **kwargs) - - def get_workflow_node_count(): - node_count = 0 - nodes = workflow_json.get("nodes", []) - if isinstance(nodes, list): - node_count += len(nodes) - - definitions = workflow_json.get("definitions", {}) - subgraphs = ( - definitions.get("subgraphs", []) - if isinstance(definitions, dict) - else [] - ) - if isinstance(subgraphs, list): - for subgraph in subgraphs: - if not isinstance(subgraph, dict): - continue - subgraph_nodes = subgraph.get("nodes", []) - if isinstance(subgraph_nodes, list): - node_count += len(subgraph_nodes) - return node_count - - def interpolate_percent(start, end, current, total): - if not total: - return start - try: - ratio = max(0.0, min(1.0, float(current) / float(total))) - except (TypeError, ValueError, ZeroDivisionError): - ratio = 0.0 - return start + ((end - start) * ratio) - - def build_loaded_models_response(): - from ..workflow_analyzer import ( - URN_TYPE_MAP, - get_lora_model_strength, - ) - - update_loaded_progress( - "scanning", - "Scanning local model index...", - percent=5, - ) - - workflow_node_count = get_workflow_node_count() - - def update_workflow_analysis_progress(payload): - self._update_workflow_analysis_progress( - loaded_id, - workflow_node_count, - interpolate_percent, - payload, - start_percent=5, - end_percent=45, - ) - - inventory = get_workflow_model_inventory( - workflow_json, - progress_callback=update_workflow_analysis_progress, - ) - available_models = inventory["available_models"] - all_model_refs = inventory["model_refs"] - - # Create lookup for full paths by filename (with and without extension) - path_by_filename = {} - total_local_models = len(available_models) - update_loaded_progress( - "indexing", - f"Indexing {total_local_models} local models...", - percent=45, - current=0, - total=total_local_models, - ) - for index, model_info in enumerate(available_models, start=1): - rel_path = model_info.get("relative_path", "") - if rel_path: - filename = get_filename_from_path(rel_path) - path_by_filename[filename] = model_info.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] = model_info.get("path") - # Add the full relative path as key too - path_by_filename[rel_path] = model_info.get("path") - - if ( - total_local_models - and (index == total_local_models or index % 250 == 0) - ): - update_loaded_progress( - "indexing", - f"Indexing local model {index} of {total_local_models}", - percent=interpolate_percent( - 45, 55, index, total_local_models - ), - current=index, - total=total_local_models, - ) - - # Also use folder_paths.get_full_path() to get paths - try: - import folder_paths - except Exception: - folder_paths = None - - folder_categories = [ - "loras", - "checkpoints", - "vae", - "controlnet", - "upscale_models", - ] - if folder_paths: - for index, category_name in enumerate(folder_categories, start=1): - update_loaded_progress( - "indexing", - f"Reading {category_name} model list...", - percent=interpolate_percent( - 55, 65, index - 1, len(folder_categories) - ), - current=index - 1, - total=len(folder_categories), - ) - try: - filenames = folder_paths.get_filename_list(category_name) - for filename in filenames: - full_path = folder_paths.get_full_path( - category_name, filename - ) - if ( - full_path - and full_path not in path_by_filename.values() - ): - path_by_filename[filename] = full_path - filename_no_ext = ( - filename.rsplit(".", 1)[0] - if "." in filename - else filename - ) - if filename_no_ext not in path_by_filename: - path_by_filename[filename_no_ext] = full_path - except Exception: - pass - - # Also extract from node.properties.models - workflow_nodes = workflow_json.get("nodes", []) - nodes = list(workflow_nodes) if isinstance(workflow_nodes, list) else [] - node_scopes = { - id(node): { - "subgraph_id": "", - "subgraph_name": "", - "is_top_level": True, - } - for node in nodes - if isinstance(node, dict) - } - definitions = workflow_json.get("definitions", {}) - subgraphs = ( - definitions.get("subgraphs", []) - if isinstance(definitions, dict) - else [] - ) - for subgraph in subgraphs: - if isinstance(subgraph, dict): - subgraph_nodes = subgraph.get("nodes", []) - if isinstance(subgraph_nodes, list): - nodes.extend(subgraph_nodes) - for node in subgraph_nodes: - if not isinstance(node, dict): - continue - node_scopes[id(node)] = { - "subgraph_id": subgraph.get("id") or "", - "subgraph_name": ( - subgraph.get("name") - or subgraph.get("id") - or "" - ), - "is_top_level": False, - } - - def get_node_scope(node): - return node_scopes.get( - id(node), - { - "subgraph_id": "", - "subgraph_name": "", - "is_top_level": True, - }, - ) - - def node_matches_ref(node, ref): - if not isinstance(node, dict): - return False - if str(node.get("id")) != str(ref.get("node_id")): - return False - scope = get_node_scope(node) - if scope["is_top_level"] != ( - ref.get("is_top_level") is not False - ): - return False - if not scope["is_top_level"]: - return str(scope["subgraph_id"]) == str( - ref.get("subgraph_id") or "" - ) - return True - - # Collect all loaded models with their values - loaded_models = [] - total_refs = len(all_model_refs) - update_loaded_progress( - "building", - "Building loaded models list...", - percent=65, - current=0, - total=total_refs, - ) - - # Process each model reference from analyze_workflow_models - for index, ref in enumerate(all_model_refs, start=1): - if total_refs and (index == total_refs or index % 50 == 0): - update_loaded_progress( - "building", - f"Building loaded model {index} of {total_refs}", - percent=interpolate_percent(65, 94, index, total_refs), - current=index, - total=total_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 = get_filename_from_path(original_path) - strength = None - - if ( - ref.get("strength") is None - and category in {"lora", "loras"} - ): - for node in nodes: - if node_matches_ref(node, ref): - strength = get_lora_model_strength( - node, - widget_index, - ) - break - - if ref.get("strength") is not None: - strength = ref.get("strength") - - model_name, strength = adapt_custom_node_loaded_model( - ref, - model_name, - strength, - ) - - # 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, - "node_title": ref.get("node_title", ""), - "subgraph_id": ref.get("subgraph_id") or "", - "subgraph_name": ref.get("subgraph_name") or "", - "is_top_level": ref.get("is_top_level") is not False, - "locate_node_id": ref.get("locate_node_id"), - "locate_node_type": ref.get("locate_node_type", ""), - "locate_node_title": ref.get("locate_node_title", ""), - "locate_subgraph_id": ( - ref.get("locate_subgraph_id") or "" - ), - "locate_subgraph_name": ( - ref.get("locate_subgraph_name") or "" - ), - "locate_is_top_level": ( - ref.get("locate_is_top_level") is not False - ), - "locate_via_promoted_widget": ref.get( - "locate_via_promoted_widget", False - ), - "exists": exists, - "strength": strength, - "original_path": original_path, - "is_urn": ref.get("is_urn", False), - "custom_node_adapter": ref.get( - "custom_node_adapter" - ), - "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 - total_nodes = len(nodes) - for node_index, node in enumerate(nodes, start=1): - if total_nodes and ( - node_index == total_nodes or node_index % 50 == 0 - ): - update_loaded_progress( - "embedded", - "Checking embedded model metadata...", - percent=interpolate_percent( - 94, 98, node_index, total_nodes - ), - current=node_index, - total=total_nodes, - ) - - node_type = node.get("type", "") - properties = node.get("properties", {}) - if not isinstance(properties, dict): - properties = {} - models_list = properties.get("models", []) - if not isinstance(models_list, list): - models_list = [] - - for model_info in models_list: - if isinstance(model_info, dict): - name = model_info.get("name", "") - directory = model_info.get("directory", "") - node_scope = get_node_scope(node) - - if name: - # Check if this model is already in loaded_models - existing = next( - ( - model - for model in loaded_models - if model.get("original_path") == name - ), - None, - ) - if not existing: - loaded_models.append( - { - "name": get_filename_from_path(name), - "category": directory or "checkpoints", - "node_id": node.get("id"), - "widget_index": None, - "node_type": node_type, - "node_title": node.get("title", ""), - "subgraph_id": node_scope["subgraph_id"], - "subgraph_name": node_scope[ - "subgraph_name" - ], - "is_top_level": node_scope[ - "is_top_level" - ], - "exists": True, # Embedded models are loaded - "strength": None, - "original_path": name, - "is_urn": False, - } - ) - - update_loaded_progress( - "finalizing", - "Loaded models ready", - percent=99, - current=len(loaded_models), - total=len(loaded_models), - ) - return { - "loaded_models": loaded_models, - "total": len(loaded_models), - } - - update_loaded_progress( - "starting", - "Preparing loaded model scan...", - percent=0, - status="starting", - current=0, - total=0, - ) - - try: - result = await asyncio.to_thread(build_loaded_models_response) - except Exception as e: - update_loaded_progress( - "error", - str(e), - percent=100, - status="error", - current=0, - total=0, - ) - self.logger.error( - f"Model Resolver loaded models error: {e}", exc_info=True - ) - return web.json_response({"error": str(e)}, status=500) - - update_loaded_progress( - "completed", - "Loaded models ready", - percent=100, - status="completed", - current=result.get("total", 0), - total=result.get("total", 0), - ) - return web.json_response(result) - - @routes.get("/model_resolver/loaded-progress/{loaded_id}") - @json_api_endpoint("loaded-progress") - async def get_loaded_models_progress(request): - """Get loaded models inspection progress.""" - return get_progress_response( - self.loaded_progress, - request, - param_name="loaded_id", - not_found_payload={ - "status": "unknown", - "stage": "unknown", - "message": "No loaded models progress available", - "percent": 0, - "current": 0, - "total": 0, - } - ) - - # ==================== MODEL METADATA LOOKUP ROUTE ==================== - - @routes.post("/model_resolver/civitai-search") - @json_api_endpoint("civitai-search") - async def civitai_search(request): - """Fetch model metadata from trusted exact-match sources.""" - data = await request.json() - filename = data.get("filename", "") - category = data.get("category", "") - resolved_path = data.get("resolved_path", "") - local_only = to_bool(data.get("local_only"), False) - force_refresh = to_bool( - data.get("force_refresh") or data.get("force"), False - ) - provided_hash = normalize_sha256( - data.get("sha256") - or data.get("hash") - or data.get("file_hash") - or "" - ) - hf_token = data.get("hf_token", "") - brave_search_api_key = data.get("brave_search_api_key", "") - hf_use_brave_fallback = to_bool( - data.get("hf_use_brave_fallback", True), - True, - ) - - 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] - - local_metadata_fields = ( - "author", - "creator", - "license", - "usage_hint", - "usage_tips", - "preview_url", - "metadata_source", - "from_safetensors_header", - "local_metadata_available", - "header_metadata_keys", - "metadata_summary", - "base_model_raw", - "sha256_source", - "hashes", - "hash_status", - ) - - def add_local_metadata_fields(payload, result): - if not isinstance(payload, dict) or not isinstance(result, dict): - return payload - for key in local_metadata_fields: - value = result.get(key) - if value in (None, "", [], {}): - continue - payload[key] = value - return payload - - # Get the file path to hash - file_path = resolved_path if resolved_path else None - file_location = "" - - if not file_path: - file_path = find_local_file_path(filename, category) - - if file_path: - file_path = _os.path.realpath( - _os.path.abspath(_os.path.normpath(str(file_path))) - ) - if not is_path_in_configured_model_roots(file_path): - if resolved_path: - return web.json_response( - { - "error": ( - "resolved_path is outside configured model directories" - ) - }, - status=403, - ) - file_path = None - elif not _os.path.isfile(file_path): - if resolved_path: - return web.json_response( - {"error": "resolved_path is not an existing model file"}, - status=404, - ) - file_path = None - - if file_path and _os.path.exists(file_path): - file_location = _os.path.dirname(file_path).replace("\\", "/") - if file_location and not file_location.endswith("/"): - file_location += "/" - - external_model_description = "" - if file_path: - external_metadata_path = find_external_metadata_sidecar_path( - file_path - ) - external_metadata = ( - read_json_safe(external_metadata_path, {}) - if external_metadata_path - else {} - ) - if isinstance(external_metadata, dict): - external_model_description = str( - external_metadata.get("modelDescription") - or external_metadata.get("model_description") - or external_metadata.get("description") - or "" - ).strip() - - def infer_model_type_from_category(value): - return normalize_category_to_model_type(value) - - def split_result_descriptions(result, preferred_model_description=""): - result = result if isinstance(result, dict) else {} - explicit_model_description = ( - result.get("model_description") - or result.get("modelDescription") - or "" - ) - generic_description = result.get("description") or "" - model_description = ( - preferred_model_description - or explicit_model_description - or generic_description - ) - version_description = ( - result.get("version_description") - or result.get("versionDescription") - or "" - ) - if ( - not version_description - and explicit_model_description - and generic_description - and generic_description != explicit_model_description - ): - version_description = generic_description - if version_description == model_description: - version_description = "" - return model_description, version_description - - def build_info_response( - result=None, - *, - metadata_path="", - metadata_saved=False, - civitai_checked=False, - local_payload=False, - ): - result = result or {} - model_description, version_description = ( - split_result_descriptions( - result, - external_model_description, - ) - ) - size_value = result.get("size") - if not size_value and file_path and _os.path.exists(file_path): - try: - size_value = _os.path.getsize(file_path) - except Exception: - size_value = None - - model_type = ( - result.get("model_type") - or result.get("type") - or infer_model_type_from_category(category) - ) - - response = { - "filename": result.get("filename") or filename, - "category": category, - "file_path": result.get("file_path") or file_path or "", - "resolved_path": result.get("resolved_path") or file_path or "", - "metadata_path": metadata_path - or result.get("metadata_path") - or "", - "metadata_saved": bool(metadata_saved), - "location": result.get("location") or file_location, - "source": result.get("source") or "", - "details_source": result.get("details_source") - or result.get("source") - or "", - "url": result.get("url"), - "version_url": result.get("version_url"), - "download_url": result.get("download_url"), - "platform_url": result.get("platform_url"), - "repo_id": result.get("repo_id"), - "path": result.get("path"), - "model_id": result.get("model_id"), - "model_name": result.get("model_name") or clean_name, - "model_type": model_type, - "version_id": result.get("version_id"), - "version_name": result.get("version_name", ""), - "sha256": result.get("sha256") or provided_hash, - "size": size_value, - "base_model": result.get("base_model"), - "base_model_source": result.get("base_model_source"), - "base_model_inferred": bool(result.get("base_model_inferred")), - "tags": result.get("tags", []), - "trained_words": result.get("trained_words", []), - "images": result.get("images", []), - "clip_skip": result.get("clip_skip"), - "description": model_description, - "model_description": model_description, - "version_description": version_description, - "from_metadata": bool(result.get("from_metadata")), - "local_only": bool(local_payload), - "metadata_checked": bool(civitai_checked), - "civitai_checked": bool(civitai_checked), - } - return add_local_metadata_fields(response, result) - def extract_result_sha256(result): - return extract_sha256_from_metadata(result) - - def result_filename_matches(result): - if not isinstance(result, dict): - return False - expected = get_filename_from_path(filename).lower() - candidates = [ - result.get("filename"), - result.get("path"), - result.get("file_path"), - ] - for candidate in candidates: - basename = get_filename_from_path(candidate).lower() - if basename and basename == expected: - return True - return False - - def result_hash_matches(result, *, require_filename=False): - if not provided_hash or not isinstance(result, dict): - return False - result_hash = extract_result_sha256(result) - if result_hash != provided_hash: - return False - return not require_filename or result_filename_matches(result) - - def huggingface_page_url(result): - try: - from urllib.parse import quote as _quote - except Exception: - _quote = None - repo_id = str( - result.get("repo_id") or result.get("repo") or "" - ).strip() - hf_path = str( - result.get("path") or result.get("filename") or "" - ).strip() - if not repo_id or not hf_path: - return "" - if _quote: - hf_path = _quote(hf_path.replace("\\", "/"), safe="/") - return f"https://huggingface.co/{repo_id}/blob/main/{hf_path}" - - def prepare_remote_result(result, source_name): - result = dict(result or {}) - source_name = str(source_name or result.get("source") or "").lower() - if source_name == "huggingface": - download_url = result.get("download_url") or result.get("url") - page_url = result.get("page_url") or huggingface_page_url(result) - if page_url: - result["url"] = page_url - result["version_url"] = page_url - result["download_url"] = download_url - result["model_name"] = ( - result.get("model_name") - or result.get("name") - or result.get("repo_id") - or clean_name - ) - result["model_type"] = ( - result.get("model_type") - or infer_model_type_from_category(category) - ) - else: - result["model_name"] = ( - result.get("model_name") - or result.get("name") - or clean_name - ) - if result.get("url") and not result.get("version_url"): - result["version_url"] = result.get("url") - - result["source"] = source_name or result.get("source") or "" - result["details_source"] = ( - result.get("details_source") - or source_name - or result.get("source") - or "" - ) - result["file_path"] = file_path or result.get("file_path") or "" - result["resolved_path"] = ( - file_path or result.get("resolved_path") or "" - ) - result["location"] = result.get("location") or file_location - result["filename"] = result.get("filename") or filename - result["sha256"] = ( - extract_result_sha256(result) - or provided_hash - or result.get("sha256") - ) - if not result.get("size") and file_path and _os.path.exists(file_path): - try: - result["size"] = _os.path.getsize(file_path) - except Exception: - pass - return result - - def result_has_remote_identity(result): - if not isinstance(result, dict): - return False - return bool( - result.get("url") - or result.get("version_url") - or result.get("download_url") - or result.get("platform_url") - or result.get("repo_id") - or result.get("model_id") - or result.get("version_id") - or result.get("from_metadata") - ) - - def remote_link_is_marked_dead(item): - if not isinstance(item, dict): - return False - status = str(item.get("status") or "").lower() - return bool( - item.get("deletedAt") - or item.get("deleted_at") - or item.get("is_dead") - or item.get("isDead") - or item.get("likelyDead") - or item.get("likely_dead") - or item.get("dead") - or status in {"dead", "deleted", "unavailable", "missing"} - ) - - def collect_result_download_urls(result): - urls = [] - if remote_link_is_marked_dead(result): - return urls - expected_filename = result.get("filename") or filename - dead_urls = set() - mirrors = result.get("mirrors") or [] - if not isinstance(mirrors, list): - mirrors = [mirrors] - for mirror in mirrors: - if not isinstance(mirror, dict): - continue - if remote_link_is_marked_dead(mirror): - dead_url = str(mirror.get("url") or "").strip() - if dead_url.startswith(("http://", "https://")): - dead_urls.add(dead_url) - continue - url = str(mirror.get("url") or "").strip() - mirror_filename = ( - mirror.get("filename") - or mirror.get("name") - or expected_filename - ) - if ( - looks_like_model_file(url, mirror_filename) - and url not in urls - ): - urls.append(url) - raw_urls = result.get("download_urls") or [] - if not isinstance(raw_urls, list): - raw_urls = [raw_urls] - for raw_url in raw_urls: - url = str(raw_url or "").strip() - if ( - looks_like_model_file(url, expected_filename) - and url not in dead_urls - and url not in urls - ): - urls.append(url) - for key in ("download_url", "downloadUrl"): - url = str(result.get(key) or "").strip() - if ( - looks_like_model_file(url, expected_filename) - and url not in dead_urls - and url not in urls - ): - urls.append(url) - return urls - - def remote_download_url_is_alive(url): - headers = { - "User-Agent": "ComfyUI-Model-Resolver/1.0", - "Accept": "*/*", - } - try: - response, _final_url, _final_headers = request_public_url( - "HEAD", - url, - headers=headers, - timeout=8, - ) - try: - if response.status_code < 400: - return True - if response.status_code in {401, 403, 404, 410}: - return False - finally: - response.close() - except Exception: - pass - - try: - response, _final_url, _final_headers = request_public_url( - "GET", - url, - headers={**headers, "Range": "bytes=0-0"}, - stream=True, - timeout=8, - ) - try: - return response.status_code < 400 - finally: - response.close() - except Exception: - return False - - def civarchive_result_has_live_download(result): - urls = collect_result_download_urls(result) - if not urls: - return False - for url in urls[:3]: - if remote_download_url_is_alive(url): - result["download_url"] = url - result["download_urls"] = [ - url, - *[other for other in urls if other != url], - ] - return True - return False - - def save_remote_metadata(result, source_name): - metadata_path = result.get("metadata_path") or "" - metadata_saved = False - if ( - result.get("from_metadata") - or not file_path - or not _os.path.exists(file_path) - ): - return metadata_path, metadata_saved - try: - model_description, version_description = ( - split_result_descriptions(result) - ) - metadata_payload = { - "source": source_name, - "details_source": result.get("details_source") - or source_name, - "filename": filename, - "category": category, - "model_name": result.get("model_name", clean_name), - "name": result.get("model_name", clean_name), - "model_type": result.get("model_type", "") - or result.get("type", ""), - "type": result.get("model_type", "") - or result.get("type", ""), - "model_id": result.get("model_id"), - "version_id": result.get("version_id"), - "version_name": result.get("version_name", ""), - "sha256": result.get("sha256") or provided_hash, - "size": result.get("size"), - "base_model": result.get("base_model"), - "base_model_source": result.get("base_model_source"), - "base_model_inferred": bool(result.get("base_model_inferred")), - "tags": result.get("tags", []), - "trained_words": result.get("trained_words", []), - "images": result.get("images", []), - "clip_skip": result.get("clip_skip"), - "description": model_description, - "model_description": model_description, - "version_description": version_description, - "download_url": result.get("download_url"), - "source_url": result.get("version_url") or result.get("url"), - "version_url": result.get("version_url") or result.get("url"), - "model_url": result.get("url"), - "url": result.get("version_url") or result.get("url"), - "platform_url": result.get("platform_url"), - "repo_id": result.get("repo_id"), - "path": result.get("path"), - "path_metadata": { - "filename": filename, - "category": category, - "source": source_name, - "model_id": result.get("model_id"), - "version_id": result.get("version_id"), - "repo_id": result.get("repo_id"), - "path": result.get("path"), - }, - } - add_local_metadata_fields(metadata_payload, result) - metadata_path = write_model_resolver_metadata( - file_path, - metadata_payload, - category, - result.get("version_url") - or result.get("url") - or result.get("platform_url") - or result.get("download_url") - or "", - create_preview=True, - ) or "" - metadata_saved = bool(metadata_path) - if metadata_path: - result["metadata_path"] = metadata_path - except Exception as metadata_error: - self.logger.warning( - f"{source_name} metadata sidecar save failed: {metadata_error}" - ) - return metadata_path, metadata_saved - - if local_only: - result = None - metadata_path = "" - metadata_saved = False - if download_available and file_path and _os.path.exists(file_path): - try: - from ..sources.civitai import ( - get_model_info_for_file, - ) - - result = get_model_info_for_file( - file_path, - local_only=True, - ) - source_metadata_path = ( - result.get("metadata_path") - if isinstance(result, dict) - else "" - ) - canonical_metadata_path = get_model_resolver_sidecar_path( - file_path - ) - if ( - isinstance(result, dict) - and result.get("from_metadata") - and source_metadata_path - and canonical_metadata_path - and _os.path.abspath(source_metadata_path) - != _os.path.abspath(canonical_metadata_path) - and not _os.path.exists(canonical_metadata_path) - ): - preview_url = result.get("preview_url") or "" - if not preview_url: - preview_path = get_existing_model_preview_path(file_path) - if preview_path: - preview_url = preview_path.replace("\\", "/") - model_description, version_description = ( - split_result_descriptions(result) - ) - metadata_payload = { - "source": "metadata_import", - "details_source": result.get("source") or "metadata", - "filename": filename, - "category": category, - "model_name": result.get("model_name", clean_name), - "name": result.get("model_name", clean_name), - "model_type": result.get("model_type", ""), - "type": result.get("model_type", ""), - "model_id": result.get("model_id"), - "version_id": result.get("version_id"), - "version_name": result.get("version_name", ""), - "sha256": result.get("sha256"), - "size": result.get("size"), - "base_model": result.get("base_model"), - "base_model_source": result.get("base_model_source"), - "base_model_inferred": bool(result.get("base_model_inferred")), - "tags": result.get("tags", []), - "trained_words": result.get("trained_words", []), - "images": result.get("images", []), - "clip_skip": result.get("clip_skip"), - "description": model_description, - "model_description": model_description, - "version_description": version_description, - "download_url": result.get("download_url"), - "preview_url": preview_url, - "source_url": result.get("version_url") - or result.get("url"), - "version_url": result.get("version_url") - or result.get("url"), - "model_url": result.get("url"), - "url": result.get("version_url") or result.get("url"), - "path_metadata": { - "filename": filename, - "category": category, - "source": "metadata_import", - "model_id": result.get("model_id"), - "version_id": result.get("version_id"), - "imported_from": source_metadata_path, - }, - } - add_local_metadata_fields(metadata_payload, result) - metadata_path = write_model_resolver_metadata( - file_path, - metadata_payload, - category, - result.get("version_url") - or result.get("url") - or result.get("download_url") - or "", - create_preview=True, - ) or "" - metadata_saved = bool(metadata_path) - if metadata_path: - result["metadata_path"] = metadata_path - result["metadata_imported_from"] = source_metadata_path - except Exception as e: - self.logger.warning(f"Local model info error: {e}") - return web.json_response( - build_info_response( - result, - metadata_path=metadata_path, - metadata_saved=metadata_saved, - local_payload=True, - civitai_checked=False, - ) - ) - - # Search remote metadata. CivitAI and CivArchive are hash-only; - # HuggingFace name results must still confirm the exact SHA. - if download_available and file_path and _os.path.exists(file_path): - try: - from ..sources.civitai import ( - get_model_info_by_hash, - get_model_info_for_file, - ) - - if provided_hash: - result = get_model_info_by_hash( - provided_hash, - use_cache=not force_refresh, - ) - if result: - result["file_path"] = file_path - result["resolved_path"] = file_path - result["location"] = file_location - if not result.get("size"): - try: - result["size"] = _os.path.getsize(file_path) - except Exception: - pass - if not extract_result_sha256(result): - result["sha256"] = provided_hash - elif not force_refresh: - result = get_model_info_for_file(file_path) - else: - result = None - - if result and ( - result.get("url") - or result.get("version_url") - or result.get("from_metadata") - or result.get("trained_words") - or result.get("base_model_inferred") - or result.get("local_metadata_available") - or result.get("from_safetensors_header") - ) and ( - result.get("from_metadata") - or not provided_hash - or result_hash_matches(result) - ): - result_source = ( - "local" - if result.get("source") == "local" - and not result_has_remote_identity(result) - else "civitai" - ) - result = prepare_remote_result(result, result_source) - metadata_path, metadata_saved = save_remote_metadata( - result, - result_source, - ) - - return web.json_response( - build_info_response( - result, - metadata_path=metadata_path, - metadata_saved=metadata_saved, - civitai_checked=True, - ) - ) - except Exception as e: - self.logger.warning(f"CivitAI search error: {e}") - - if provided_hash: - try: - result = resolve_civarchive_by_hash( - provided_hash, - query=filename, - exact_only=False, - model_type=infer_model_type_from_category(category), - ) - if result: - if not extract_result_sha256(result): - result["sha256"] = provided_hash - result = prepare_remote_result(result, "civarchive") - if result_hash_matches(result) and civarchive_result_has_live_download(result): - metadata_path, metadata_saved = save_remote_metadata( - result, - "civarchive", - ) - return web.json_response( - build_info_response( - result, - metadata_path=metadata_path, - metadata_saved=metadata_saved, - civitai_checked=True, - ) - ) - if result_hash_matches(result): - self.logger.info( - f"CivArchive metadata candidate rejected: " - f"no live download link for {filename}" - ) - except Exception as e: - self.logger.warning(f"CivArchive hash lookup error: {e}") - - if provided_hash: - hf_attempts = [ - { - "label": "HuggingFace", - "use_api_search": True, - "use_comfy_org_fallback": True, - "use_brave_fallback": False, - } - ] - if hf_use_brave_fallback and brave_search_api_key: - hf_attempts.append( - { - "label": "HuggingFace Brave", - "use_api_search": False, - "use_comfy_org_fallback": False, - "use_brave_fallback": True, - } - ) - - for hf_attempt in hf_attempts: - try: - result = search_huggingface_for_file( - filename, - token=hf_token or None, - exact_only=True, - brave_api_key=brave_search_api_key or None, - use_api_search=hf_attempt["use_api_search"], - use_comfy_org_fallback=hf_attempt[ - "use_comfy_org_fallback" - ], - use_brave_fallback=hf_attempt[ - "use_brave_fallback" - ], - force_refresh=force_refresh, - ) - if result and result_hash_matches( - result, - require_filename=True, - ): - result = prepare_remote_result( - result, - "huggingface", - ) - metadata_path, metadata_saved = save_remote_metadata( - result, - "huggingface", - ) - return web.json_response( - build_info_response( - result, - metadata_path=metadata_path, - metadata_saved=metadata_saved, - civitai_checked=True, - ) - ) - if result: - self.logger.info( - f"{hf_attempt['label']} metadata candidate rejected: " - f"filename/hash mismatch for {filename}" - ) - except Exception as e: - self.logger.warning( - f"{hf_attempt['label']} metadata lookup error: {e}" - ) - - # No remote result found. Keep any useful local safetensors - # header metadata, including base-model fingerprints, model - # titles, descriptions, tags, trigger words, and author data. - local_fallback_result = None - if file_path and _os.path.exists(file_path): - try: - from ..sources.civitai import get_model_info_for_file - - local_candidate = get_model_info_for_file( - file_path, - local_only=True, - ) - if ( - isinstance(local_candidate, dict) - and ( - local_candidate.get("base_model") - or local_candidate.get("local_metadata_available") - or local_candidate.get("from_safetensors_header") - ) - ): - local_fallback_result = local_candidate - except Exception as local_fallback_error: - self.logger.debug( - f"Local base model inference fallback failed: {local_fallback_error}" - ) - - response = build_info_response( - local_fallback_result, - civitai_checked=True, - local_payload=bool(local_fallback_result), - ) - response["url"] = None - if force_refresh and file_path and _os.path.exists(file_path): - try: - metadata_payload = { - "filename": filename, - "category": category, - "model_name": response.get("model_name") or clean_name, - "name": response.get("model_name") or clean_name, - "model_type": response.get("model_type") - or infer_model_type_from_category(category), - "sha256": response.get("sha256") or provided_hash, - "size": response.get("size") or _os.path.getsize(file_path), - "base_model": response.get("base_model"), - "base_model_source": response.get("base_model_source"), - "base_model_inferred": bool(response.get("base_model_inferred")), - "description": response.get("model_description") - or response.get("description") - or "", - "model_description": response.get("model_description") - or response.get("description") - or "", - "version_description": response.get( - "version_description" - ) - or "", - "civitai_deleted": True, - "civitai_checked": True, - "remote_metadata_missing": True, - "source": "local", - "details_source": response.get("details_source") or "", - } - add_local_metadata_fields(metadata_payload, response) - metadata_path = write_model_resolver_metadata( - file_path, - metadata_payload, - category, - "", - ) or "" - response["metadata_path"] = metadata_path - response["metadata_saved"] = bool(metadata_path) - except Exception as metadata_error: - self.logger.warning( - f"Remote metadata no-match sidecar save failed: {metadata_error}" - ) - return web.json_response(response) - - - - def _custom_result_timestamp(): - return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) - - - - - - - - def _build_direct_custom_download_result(url, source, expected_filename=""): - from urllib.parse import urlparse - - if not looks_like_model_file(url, expected_filename): - return None - - parsed = urlparse(url) - filename = ( - get_filename_from_path(parsed.path) - or get_filename_from_path(expected_filename) - or "model" - ) - if expected_filename and "." in get_filename_from_path(expected_filename): - filename = get_filename_from_path(expected_filename) - source = str(source or "custom").strip().lower() - source_label = { - "civarchive": "CivArchive", - "civitai": "CivitAI", - "huggingface": "HuggingFace", - }.get(source, "Custom URL") - return { - "source": source, - "details_source": source, - "name": source_label, - "filename": filename, - "url": url, - "version_url": url, - "download_url": url, - "match_type": "custom_url", - "custom_url": True, - } - - def _collect_custom_url_local_hash_matches(result, category): - source_key = str(result.get("source") or "custom").strip().lower() - sha256 = normalize_sha256(extract_sha256_from_metadata(result)) - if not sha256: - return [] - - try: - matches = search_local_matches_by_hash( - sha256, - category=category or None, - max_matches=20, - ) - except Exception as hash_error: - self.logger.warning( - f"Custom URL local metadata hash lookup failed for {source_key}:{sha256}: {hash_error}" - ) - return [] - - return [ - { - **match, - "hash_lookup_source": source_key, - "hash_lookup_filename": result.get("filename") - or result.get("path") - or "", - "hash_lookup_sha256": sha256, - } - for match in matches - ] - - @routes.post("/model_resolver/custom-url") - @json_api_endpoint("custom-url") - async def custom_url(request): - """Resolve a user-provided provider URL into a normal search result.""" - data = await request.json() - raw_url = str(data.get("url") or data.get("custom_url") or "").strip() - if not raw_url: - return web.json_response( - {"error": "URL is required"}, status=400 - ) - - if raw_url.startswith("hf://"): - normalized_url = raw_url - else: - try: - normalized_url = await asyncio.to_thread( - validate_public_http_url, - raw_url, - ) - except UnsafeUrlError as exc: - return web.json_response( - {"error": str(exc)}, - status=400, - ) - - category = data.get("category") or "" - expected_filename = ( - data.get("filename") - or get_filename_from_path(data.get("original_path") or "") - or "" - ) - civitai_key = data.get("civitai_key", "") - hf_token = data.get("hf_token", "") - - result = None - source = "" - try: - civitai_parsed = parse_civitai_url(normalized_url) - except Exception: - civitai_parsed = None - try: - civarchive_parsed = parse_civarchive_url(normalized_url) - except Exception: - civarchive_parsed = None - - if civitai_parsed: - source = "civitai" - model_id = civitai_parsed.get("model_id") - version_id = civitai_parsed.get("version_id") - if model_id: - details = await asyncio.to_thread( - get_civitai_model_details, - model_id, - version_id, - civitai_key or None, - ) - result = build_civitai_custom_result( - details, - expected_filename=expected_filename, - api_key=civitai_key or None, - ) - elif version_id: - result = await asyncio.to_thread( - resolve_civitai_version_custom_result, - version_id, - expected_filename, - civitai_key or None, - ) - if not result and version_id: - result = { - "source": "civitai", - "details_source": "civitai", - "version_id": version_id, - "name": expected_filename or f"CivitAI version {version_id}", - "filename": expected_filename or f"civitai-{version_id}", - "url": normalized_url, - "version_url": normalized_url, - "download_url": get_civitai_download_url( - version_id, - civitai_key or None, - ), - "match_type": "custom_url", - "custom_url": True, - } - elif civarchive_parsed: - source = "civarchive" - if civarchive_parsed.get("sha256"): - result = await asyncio.to_thread( - resolve_civarchive_by_hash, - civarchive_parsed.get("sha256"), - expected_filename, - False, - normalize_category_to_model_type(category), - ) - else: - model_id = civarchive_parsed.get("model_id") - version_id = civarchive_parsed.get("version_id") - if model_id: - result = await asyncio.to_thread( - resolve_civarchive_model_version, - model_id, - version_id, - expected_filename or str(model_id), - False, - True, - ) - if not result: - details = await asyncio.to_thread( - get_civarchive_model_details, - model_id, - version_id, - True, - ) - result = build_civarchive_custom_result( - details, - expected_filename=expected_filename, - ) - if result: - result = dict(result) - result["source"] = "civarchive" - result["details_source"] = "civarchive" - result["match_type"] = "custom_url" - result["custom_url"] = True - else: - from urllib.parse import urlparse - - parsed_direct = urlparse(normalized_url) - direct_host = parsed_direct.hostname - if ( - host_matches_domain(direct_host, "civarchive.com") - and "/api/download/" in parsed_direct.path - ): - source = "civarchive" - result = _build_direct_custom_download_result( - normalized_url, - "civarchive", - expected_filename, - ) - else: - result = await asyncio.to_thread( - build_huggingface_custom_result, - normalized_url, - expected_filename, - hf_token or None, - ) - if result: - source = "huggingface" - - if not result: - return web.json_response( - { - "error": ( - "Unsupported or unresolved URL. Use a HuggingFace file URL, " - "CivitAI model/download URL, or CivArchive model/hash URL." - ) - }, - status=400, - ) - - result = dict(result) - result["provided_url"] = normalized_url - result["url_source"] = "custom" - result["searched_at"] = result.get("searched_at") or _custom_result_timestamp() - result.setdefault("category", category) - if expected_filename and not result.get("filename"): - result["filename"] = expected_filename - - if not (result.get("download_url") or result.get("url")): - return web.json_response( - {"error": "The URL resolved, but no download URL was found"}, - status=400, - ) - - source = source or result.get("source") or "custom" - local_hash_matches = _collect_custom_url_local_hash_matches( - result, - category, - ) - response = { - "success": True, - "source": source, - "result": result, - "custom": [result], - "searched_sources": ["custom"], - "local_hash_matches": local_hash_matches, - } - response[source] = result - return web.json_response(response) - - @routes.post("/model_resolver/model-details") - @json_api_endpoint("model-details") - async def model_details(request): - """Return normalized full model details for sources that expose model pages.""" - data = await request.json() - source = str(data.get("source", "")).strip().lower() - model_id = data.get("model_id") - version_id = data.get("version_id") - civitai_key = data.get("civitai_key", "") - hf_token = data.get("hf_token", "") - file_path = data.get("file_path", "") - branch = data.get("branch", "") - - if not download_available: - return web.json_response( - {"error": "Download providers are not available"}, status=503 - ) - - if source == "lora_manager_archive": - source = "civitai" - - if source not in {"civitai", "civarchive", "huggingface"}: - return web.json_response( - {"error": "Unsupported model details source"}, status=400 - ) - - if source == "huggingface": - model_id = str(model_id or "").strip() - branch = str(branch or version_id or "main").strip() or "main" - version_id = branch - else: - try: - model_id = ( - int(model_id) - if model_id is not None and str(model_id).strip() - else None - ) - except (TypeError, ValueError): - model_id = None - - try: - version_id = ( - int(version_id) - if version_id is not None and str(version_id).strip() - else None - ) - except (TypeError, ValueError): - version_id = None - - if not model_id: - return web.json_response( - {"error": "model_id is required"}, status=400 - ) - - if source == "civitai": - details = await asyncio.to_thread( - get_civitai_model_details, - model_id, - version_id, - civitai_key or None, - ) - elif source == "civarchive": - details = await asyncio.to_thread( - get_civarchive_model_details, - model_id, - version_id, - ) - else: - details = await asyncio.to_thread( - get_huggingface_model_details, - model_id, - file_path, - branch, - hf_token or None, - ) - - if not details: - return web.json_response( - {"error": "Model details not found"}, status=404 - ) - - return web.json_response(details) - - # ==================== DOWNLOAD ROUTES ==================== + route_context = {**globals(), **locals()} + register_workflow_routes(route_context) + register_metadata_routes(route_context) + register_loaded_model_routes(route_context) + register_model_info_routes(route_context) if download_available: - - from ..path_utils import ( - dedupe_local_base_directories, - get_comfy_root_path, - prefer_local_base_directory, - ) - from ..path_utils import ( - get_path_identity as get_local_path_identity, - ) - - # cleanup_search_progress, is_search_progress_cancelled, mark_search_progress_cancelled, and update_search_progress removed - - @routes.get("/model_resolver/search-progress/{progress_id}") - @json_api_endpoint("search-progress") - async def get_search_progress_route(request): - """Return live progress for an in-flight source search.""" - return get_progress_response( - self.search_tracker, - request, - not_found_payload={"exists": False}, - found_wrapper=lambda prog: {"exists": True, **prog} - ) - - @routes.post("/model_resolver/search-cancel/{progress_id}") - @json_api_endpoint("search-cancel", return_success_on_error=True) - async def cancel_search_progress_route(request): - """Mark an in-flight source search as cancelled.""" - return cancel_progress_response( - self.search_tracker, - request, - cancel_message="Cancelled" - ) - - @routes.post("/model_resolver/search") - async def search_sources(request): - """Search for model download sources.""" - try: - class SearchCancelled(BaseException): - # Progress helpers swallow ordinary callback errors; cancellation - # must bubble out of source loops and stop follow-up requests. - pass - - def raise_if_search_cancelled(source=""): - if self.search_tracker.is_cancelled(progress_id): - raise SearchCancelled("Search cancelled") - - def format_log_value(value): - if value is None or value == "": - return None - if isinstance(value, bool): - return "yes" if value else "no" - if isinstance(value, (list, tuple, set)): - return ",".join(str(item) for item in value) - - text = str(value) - if any(char.isspace() for char in text): - return '"' + text.replace("\\", "\\\\").replace('"', '\\"') + '"' - return text - - def format_log_size(value): - return format_size_bytes(value, include_space=False) - - def format_log_fields(**fields): - parts = [] - for key, value in fields.items(): - formatted = format_log_value(value) - if formatted is not None: - parts.append(f"{key}={formatted}") - return " ".join(parts) - - def normalize_result_extra(extra): - if not extra: - return {} - normalized = dict(extra) - model_id = normalized.pop("model_id", None) - version_id = normalized.pop("version_id", None) - if model_id or version_id: - normalized["ids"] = ( - f"{model_id}@{version_id}" - if model_id and version_id - else model_id or version_id - ) - if "files_count" in normalized: - normalized["files"] = normalized.pop("files_count") - return normalized - - def format_result_details(result, extra=None): - if isinstance(result, list): - return format_log_fields(count=len(result)) - if not isinstance(result, dict): - fields = {"result": "none" if result is None else result} - fields.update(normalize_result_extra(extra)) - return format_log_fields(**fields) - - model_id = result.get("model_id") - version_id = result.get("version_id") - ids = ( - f"{model_id}@{version_id}" - if model_id and version_id - else model_id or version_id - ) - fields = { - "name": result.get("name"), - "file": result.get("filename") or result.get("path"), - "match": result.get("match_type"), - "repo": result.get("repo_id") or result.get("repo"), - "ids": ids, - "size": format_log_size(result.get("size")), - "files": result.get("files_count"), - } - fields.update(normalize_result_extra(extra)) - return format_log_fields(**fields) - - def log_search_result(source_name, result, extra=None): - details = format_result_details(result, extra) - if result and ( - not isinstance(result, list) or len(result) > 0 - ): - self.logger.info(f"Search [{source_name}] found {details}") - else: - self.logger.info(f"Search [{source_name}] miss {details}") - - data = await request.json() - filename = data.get("filename", "") - category = data.get("category", "") - base_model_context = data.get("base_model_context", "") - progress_id = str(data.get("progress_id") or "").strip() - progress_source = str(data.get("progress_source") or "").strip() - civitai_candidate_limit = to_int(data.get("civitai_candidate_limit"), 5) - civitai_candidate_limit = max( - 1, min(civitai_candidate_limit, 20) - ) - civarchive_candidate_limit = to_int(data.get("civarchive_candidate_limit"), 10) - civarchive_candidate_limit = max( - 1, min(civarchive_candidate_limit, 30) - ) - # Handle both boolean and string forms - is_urn_raw = data.get("is_urn", False) - civitai_key = data.get("civitai_key", "") - civitai_session_token = data.get("civitai_session_token", "") - hf_token = data.get("hf_token", "") - brave_search_api_key = data.get("brave_search_api_key", "") - civitai_use_trpc_search = data.get( - "civitai_use_trpc_search", True - ) - civitai_use_api_search = data.get( - "civitai_use_api_search", True - ) - civitai_use_html_fallback = data.get( - "civitai_use_html_fallback", True - ) - hf_use_api_search = data.get("hf_use_api_search", True) - hf_use_comfy_org_fallback = data.get( - "hf_use_comfy_org_fallback", True - ) - hf_use_brave_fallback = data.get( - "hf_use_brave_fallback", True - ) - is_urn = to_bool(is_urn_raw, False) - hf_use_api_search = to_bool(hf_use_api_search, True) - civitai_use_trpc_search = to_bool(civitai_use_trpc_search, True) - civitai_use_api_search = to_bool(civitai_use_api_search, True) - civitai_use_html_fallback = to_bool(civitai_use_html_fallback, True) - hf_use_comfy_org_fallback = to_bool(hf_use_comfy_org_fallback, True) - hf_use_brave_fallback = to_bool(hf_use_brave_fallback, 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", - "civarchive", - "lora_manager_archive", - } - - search_local = "local" in normalized_sources - search_huggingface_source = "huggingface" in normalized_sources - search_civitai_source = "civitai" in normalized_sources - search_civarchive_source = "civarchive" in normalized_sources - search_lora_manager_archive_source = ( - "lora_manager_archive" in normalized_sources - ) - if not progress_source: - progress_source = ( - next(iter(normalized_sources)) - if len(normalized_sources) == 1 - else "all" - ) - force_search = to_bool(data.get("force_search"), False) - - self.search_tracker.update( - progress_id, - progress_source, - "starting", - "Preparing search", - 8, - ) - raise_if_search_cancelled(progress_source) - - if force_search: - self.search_tracker.update( - progress_id, - progress_source, - "cache", - "Refreshing search caches", - 12, - ) - if search_local: - reload_popular_databases() - reload_model_list() - if search_huggingface_source: - clear_huggingface_search_cache() - if search_civitai_source: - clear_civitai_search_cache() - if search_civarchive_source: - clear_civarchive_search_cache() - if search_lora_manager_archive_source: - clear_lora_manager_archive_search_cache() - self.logger.debug( - "Force search enabled: cleared cache " - + format_log_fields(sources=sorted(normalized_sources)) - ) - raise_if_search_cancelled(progress_source) - - self.logger.info( - f"Search [{','.join(sorted(normalized_sources))}] request " - + format_log_fields( - file=filename, - cat=category, - urn=is_urn, - ids=( - f"{data.get('model_id')}@{data.get('version_id')}" - if data.get("model_id") and data.get("version_id") - else data.get("model_id") - or data.get("version_id") - ), - base=base_model_context, - force=force_search, - ) - ) - - results = { - "popular": None, - "model_list": None, - "huggingface": None, - "civitai": None, - "civarchive": None, - "lora_manager_archive": None, - "local_hash_matches": [], - "found": False, - "searched_sources": sorted(normalized_sources), - "source_errors": {}, - } - - def current_search_timestamp(): - from datetime import datetime, timezone - - return ( - datetime.now(timezone.utc) - .replace(microsecond=0) - .isoformat() - ) - - def get_search_result_signature(source_key, result): - if isinstance(result, list): - return "|".join( - get_search_result_signature(source_key, item) - for item in result - ) - if not isinstance(result, dict): - return "" - - parts = [ - source_key, - result.get("download_url") - or result.get("url") - or result.get("model_url") - or "", - result.get("filename") or result.get("path") or "", - result.get("repo_id") or result.get("repo") or "", - result.get("model_id") or "", - result.get("version_id") or "", - result.get("name") or "", - ] - return "::".join(str(part).strip() for part in parts) - - def stamp_search_result(source_key, result): - if isinstance(result, list): - return [ - stamp_search_result(source_key, item) - for item in result - ] - if not isinstance(result, dict): - return result - - signature = get_search_result_signature(source_key, result) - if not signature: - return result - - timestamp = ( - result.get("searched_at") - or result.get("searchedAt") - or ( - None - if force_search - else self.search_result_timestamps.get(signature) - ) - ) - if not timestamp: - timestamp = current_search_timestamp() - if force_search: - self.search_result_timestamps[signature] = timestamp - else: - self.search_result_timestamps.setdefault( - signature, timestamp - ) - result["searched_at"] = timestamp - return result - - def stamp_search_results(payload): - for source_key in ( - "popular", - "model_list", - "huggingface", - "civitai", - "civarchive", - "lora_manager_archive", - ): - if payload.get(source_key): - payload[source_key] = stamp_search_result( - source_key, payload[source_key] - ) - return payload - - def mark_any_model_fallback(result): - if isinstance(result, list): - return [ - mark_any_model_fallback(item) - for item in result - ] - if not isinstance(result, dict): - return result - - marked = dict(result) - marked["any_model_match"] = True - marked["base_model_fallback"] = True - marked["requested_base_model"] = base_model_context - return marked - - def iter_result_items(result): - if isinstance(result, list): - for item in result: - if isinstance(item, dict): - yield item - elif isinstance(result, dict): - yield result - - def collect_local_hash_matches(payload): - matches = [] - seen_match_paths = set() - seen_hashes = set() - for source_key in ( - "huggingface", - "civitai", - "civarchive", - "lora_manager_archive", - "popular", - "model_list", - ): - for source_result in iter_result_items(payload.get(source_key)): - raise_if_search_cancelled(source_key) - sha256 = extract_sha256_from_metadata(source_result) - if not sha256: - continue - - if sha256 in seen_hashes: - continue - seen_hashes.add(sha256) - - self.search_tracker.update( - progress_id, - progress_source, - "local_hash", - "Checking local metadata hashes", - 94, - ) - try: - hash_matches = search_local_matches_by_hash( - sha256, - category=category or None, - max_matches=20, - force_rescan=force_search, - ) - except Exception as hash_error: - self.logger.warning( - f"Local metadata hash lookup failed for {source_key}:{sha256}: {hash_error}" - ) - continue - - for match in hash_matches: - model_path = ( - match.get("model", {}).get("path") - or match.get("path") - or "" - ) - path_key = model_path.lower() - if path_key and path_key in seen_match_paths: - continue - if path_key: - seen_match_paths.add(path_key) - enriched = { - **match, - "hash_lookup_source": source_key, - "hash_lookup_filename": source_result.get("filename") - or source_result.get("path") - or filename, - "hash_lookup_sha256": sha256, - } - matches.append(enriched) - - if matches: - self.logger.info( - "Search local hash matches " - + format_log_fields(count=len(matches)) - ) - return matches - - def make_source_progress_callback( - source_key, - percent_min=None, - percent_max=None, - ): - def source_progress_callback(payload): - raise_if_search_cancelled(source_key) - if not isinstance(payload, dict): - return - - progress_payload = dict(payload) - stage = progress_payload.pop("stage", "running") - message = progress_payload.pop( - "message", "Searching..." - ) - percent = progress_payload.pop("percent", None) - status = progress_payload.pop("status", "running") - progress_payload.pop("source", None) - - if ( - percent is not None - and percent_min is not None - and percent_max is not None - ): - try: - normalized_percent = max( - 0.0, min(100.0, float(percent)) - ) - percent = percent_min + ( - normalized_percent / 100.0 - ) * (percent_max - percent_min) - except (TypeError, ValueError): - percent = None - - self.search_tracker.update( - progress_id, - source_key, - stage, - message, - percent, - status=status, - **progress_payload, - ) - - return source_progress_callback - - def run_source_search( - source_key, - search_task_fn, - initial_stage="query", - initial_message=None, - initial_percent=30, - log_start_fields=None, - error_handlers=None, - ): - if initial_message is None: - initial_message = f"Querying {source_key.capitalize()}" - raise_if_search_cancelled(source_key) - self.search_tracker.update( - progress_id, - source_key, - initial_stage, - initial_message, - initial_percent, - ) - start_fields = log_start_fields or {"file": filename} - self.logger.info( - f"Search [{source_key}] start " - + format_log_fields(**start_fields) - ) - try: - raise_if_search_cancelled(source_key) - source_results, source_found = search_task_fn() - raise_if_search_cancelled(source_key) - done_messages = { - "local": "Local database checked", - "huggingface": "HuggingFace checked", - "civitai": "CivitAI checked", - "civarchive": "CivArchive checked", - "lora_manager_archive": "LoRA Manager archive checked", - } - done_msg = done_messages.get(source_key, f"{source_key} checked") - self.search_tracker.update( - progress_id, - source_key, - "done", - done_msg, - 92, - ) - return source_results, source_found - except Exception as e: - if error_handlers: - for exc_type, handler_fn in error_handlers.items(): - if isinstance(e, exc_type): - return handler_fn(e) - raise e - - def execute_search_with_fallback( - source_key, - search_fn, - any_model_label, - ): - raise_if_search_cancelled(source_key) - res = search_fn( - base_model_context or None, - make_source_progress_callback(source_key), - ) - raise_if_search_cancelled(source_key) - log_search_result(source_key, res) - - if not res and base_model_context: - raise_if_search_cancelled(source_key) - self.search_tracker.update( - progress_id, - source_key, - "any_model", - f"Retrying {any_model_label} any model", - 72, - ) - self.logger.info( - f"Search [{source_key}] retry any model " - + format_log_fields( - file=filename, - cat=category, - base=base_model_context, - ) - ) - res = search_fn( - None, - make_source_progress_callback(source_key, 72, 92), - ) - raise_if_search_cancelled(source_key) - log_search_result(f"{source_key}/any_model", res) - if res: - res = mark_any_model_fallback(res) - return res - - def search_local_sources(): - def task(): - source_results = {"popular": None, "model_list": None} - source_found = False - - popular_info = get_popular_model_url(filename) - log_search_result("popular", popular_info) - self.search_tracker.update( - progress_id, - "local", - "model_list", - "Checking local model database", - 58, - ) - model_list_result = search_model_list(filename) - log_search_result( - "model_list", - model_list_result, - { - "confidence": model_list_result.get("confidence") - if model_list_result - else None - }, - ) - if popular_info: - popular_result = { - "source": "popular", - "filename": filename, - **popular_info, - } - if ( - model_list_result - and model_list_result.get("filename", "").lower() - == filename.lower() - and model_list_result.get("size") - ): - popular_result["size"] = model_list_result.get( - "size" - ) - source_results["popular"] = popular_result - source_found = True - - if model_list_result: - confidence = model_list_result.get("confidence", 0) - if (is_urn and confidence >= 70) or not is_urn: - source_results["model_list"] = model_list_result - source_found = True - - return source_results, source_found - - return run_source_search( - "local", - task, - initial_stage="popular", - initial_message="Checking popular models", - initial_percent=28, - log_start_fields={"file": filename, "cat": category}, - ) - - def search_huggingface_source_task(): - def task(): - hf_result = search_huggingface_for_file( - filename, - token=hf_token or None, - brave_api_key=brave_search_api_key or None, - use_api_search=hf_use_api_search, - use_comfy_org_fallback=hf_use_comfy_org_fallback, - use_brave_fallback=hf_use_brave_fallback, - force_refresh=force_search, - progress_callback=make_source_progress_callback( - "huggingface" - ), - ) - log_search_result("huggingface", hf_result) - return {"huggingface": hf_result}, bool(hf_result) - - return run_source_search( - "huggingface", - task, - initial_stage="query", - initial_message="Querying HuggingFace", - initial_percent=32, - log_start_fields={"file": filename}, - ) - - def search_civitai_source_task(): - def task(): - source_results = {"civitai": None} - source_found = False - - if is_urn: - model_id_val = data.get("model_id") - version_id_val = data.get("version_id") - - if model_id_val and version_id_val: - self.search_tracker.update( - progress_id, - "civitai", - "urn", - "Resolving CivitAI URN", - 46, - ) - model_info = resolve_urn(model_id_val, version_id_val) - if model_info: - self.search_tracker.update( - progress_id, - "civitai", - "file", - "Selecting CivitAI file", - 76, - ) - 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_val - ) - source_results["civitai"] = build_search_result( - "civitai", - model_id=model_id_val, - version_id=version_id_val, - name=model_info.get("model_name"), - version_name=model_info.get("version_name"), - filename=model_info.get("expected_filename"), - type=category, - download_url=download_url, - url=f"https://civitai.com/models/{model_id_val}?modelVersionId={version_id_val}", - size=primary_file.get("size"), - base_model=model_info.get("base_model"), - tags=model_info.get("tags", []), - sha256=primary_file.get("sha256") - or (primary_file.get("hashes") or {}).get("SHA256") - or (primary_file.get("hashes") or {}).get("sha256"), - hashes=primary_file.get("hashes") or {}, - match_type="exact", - confidence=100.0, - ) - log_search_result( - "civitai/urn", - source_results["civitai"], - { - "files_count": len( - model_info.get("files", []) - ) - }, - ) - source_found = True - else: - log_search_result( - "civitai/urn", - None, - { - "model_id": model_id_val, - "version_id": version_id_val, - }, - ) - elif category: - self.search_tracker.update( - progress_id, - "civitai", - "fallback", - "Searching CivitAI fallback", - 58, - ) - self.logger.info( - "Search [civitai] URN ids missing; falling back" - ) - civitai_results = search_civitai( - filename, - model_type=category, - ) - log_search_result( - "civitai/fallback", - civitai_results[0] if civitai_results else None, - { - "results_count": len(civitai_results), - }, - ) - if civitai_results: - first_result = civitai_results[0] - source_results["civitai"] = build_search_result( - "civitai", - model_id=first_result.get("model_id"), - version_id=first_result.get("version_id"), - 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"), - base_model=first_result.get("base_model"), - tags=first_result.get("tags", []), - ) - source_found = True - else: - civitai_result = execute_search_with_fallback( - "civitai", - lambda base_ctx, cb: search_civitai_for_file( - filename, - api_key=civitai_key or None, - model_type=category, - base_model_context=base_ctx, - session_token=civitai_session_token or None, - candidate_limit=civitai_candidate_limit, - use_trpc_search=civitai_use_trpc_search, - use_api_search=civitai_use_api_search, - use_html_fallback=civitai_use_html_fallback, - progress_callback=cb, - ), - "CivitAI" - ) - if civitai_result: - source_results["civitai"] = civitai_result - source_found = True - - return source_results, source_found - - return run_source_search( - "civitai", - task, - initial_stage="query", - initial_message="Querying CivitAI", - initial_percent=30, - log_start_fields={ - "file": filename, - "cat": category, - "urn": is_urn, - }, - ) - - def search_civarchive_source_task(): - def task(): - source_results = {"civarchive": None} - source_found = False - - if is_urn: - model_id_val = data.get("model_id") - version_id_val = data.get("version_id") - if model_id_val and version_id_val: - self.search_tracker.update( - progress_id, - "civarchive", - "urn", - "Resolving CivArchive version", - 50, - ) - civarchive_result = resolve_civarchive_model_version( - model_id_val, - version_id_val, - query=filename, - ) - log_search_result( - "civarchive/urn", - civarchive_result, - { - "model_id": model_id_val, - "version_id": version_id_val, - }, - ) - if civarchive_result: - source_results["civarchive"] = civarchive_result - source_found = True - else: - log_search_result( - "civarchive/urn", - None, - { - "model_id": model_id_val, - "version_id": version_id_val, - }, - ) - else: - civarchive_result = execute_search_with_fallback( - "civarchive", - lambda base_ctx, cb: search_civarchive_for_file( - filename, - model_type=category, - base_model_context=base_ctx, - limit=civarchive_candidate_limit, - progress_callback=cb, - ), - "CivArchive" - ) - if civarchive_result: - source_results["civarchive"] = civarchive_result - source_found = True - - return source_results, source_found - - def handle_civarchive_error(e): - error_message = f"CivArchive search failed: {e}" - self.logger.warning(error_message) - self.search_tracker.update( - progress_id, - "civarchive", - "error", - error_message, - 100, - status="error", - ) - return { - "civarchive": None, - "source_errors": {"civarchive": error_message}, - }, False - - return run_source_search( - "civarchive", - task, - initial_stage="query", - initial_message="Querying CivArchive", - initial_percent=30, - log_start_fields={ - "file": filename, - "cat": category, - "urn": is_urn, - }, - error_handlers={CivArchiveSearchError: handle_civarchive_error}, - ) - - def search_lora_manager_archive_source_task(): - def task(): - lora_manager_archive_result = execute_search_with_fallback( - "lora_manager_archive", - lambda base_ctx, cb: search_lora_manager_archive_for_file( - filename, - model_type=category, - base_model_context=base_ctx, - progress_callback=cb, - ), - "LoRA archive" - ) - return ( - { - "lora_manager_archive": lora_manager_archive_result - }, - bool(lora_manager_archive_result), - ) - - return run_source_search( - "lora_manager_archive", - task, - initial_stage="query", - initial_message="Searching LoRA Manager archive", - initial_percent=36, - log_start_fields={"file": filename, "cat": category}, - ) - - - search_tasks = [] - if search_local: - search_tasks.append( - asyncio.to_thread(search_local_sources) - ) - if search_huggingface_source: - search_tasks.append( - asyncio.to_thread(search_huggingface_source_task) - ) - if search_civitai_source: - search_tasks.append( - asyncio.to_thread(search_civitai_source_task) - ) - if search_civarchive_source and ( - filename or (is_urn and model_id and version_id) - ): - search_tasks.append( - asyncio.to_thread(search_civarchive_source_task) - ) - if search_lora_manager_archive_source and filename: - search_tasks.append( - asyncio.to_thread( - search_lora_manager_archive_source_task - ) - ) - - if len(search_tasks) > 1: - self.logger.debug( - "Search sources async " - + format_log_fields(count=len(search_tasks)) - ) - self.search_tracker.update( - progress_id, - progress_source, - "running", - "Waiting for search sources", - 18, - ) - raise_if_search_cancelled(progress_source) - - for source_results, source_found in await asyncio.gather( - *search_tasks - ): - raise_if_search_cancelled(progress_source) - for source_key, source_result in source_results.items(): - if source_key == "source_errors": - results["source_errors"].update(source_result or {}) - continue - if source_result: - results[source_key] = source_result - if source_found: - results["found"] = True - - raise_if_search_cancelled(progress_source) - results["local_hash_matches"] = collect_local_hash_matches(results) - raise_if_search_cancelled(progress_source) - - self.logger.info( - f"Search [{','.join(results['searched_sources'])}] done " - + format_log_fields( - found=results["found"], - ) - ) - self.search_tracker.update( - progress_id, - progress_source, - "completed", - "Search complete", - 100, - status="completed", - ) - stamp_search_results(results) - return web.json_response(results) - - except SearchCancelled: - self.search_tracker.update( - progress_id if "progress_id" in locals() else "", - progress_source if "progress_source" in locals() else "", - "cancelled", - "Cancelled", - 100, - status="cancelled", - cancelled=True, - ) - self.logger.info( - "Search cancelled " - + format_log_fields( - source=progress_source if "progress_source" in locals() else "", - progress_id=progress_id if "progress_id" in locals() else "", - ) - ) - return web.json_response( - { - "cancelled": True, - "found": False, - "searched_sources": sorted(normalized_sources) - if "normalized_sources" in locals() - else [], - "source_errors": {}, - } - ) - - except Exception as e: - self.search_tracker.update( - progress_id if "progress_id" in locals() else "", - progress_source if "progress_source" in locals() else "", - "error", - str(e), - 100, - status="error", - ) - self.logger.exception(f"Model Resolver search error: {e}") - return web.json_response({"error": str(e)}, status=500) - - @routes.post("/model_resolver/clear-search-cache") - @json_api_endpoint("Clear search cache") - async def clear_search_cache_route(request): - """Clear backend search caches after token/settings changes.""" - clear_all_search_caches() - reload_popular_databases() - reload_model_list() - invalidate_model_files_cache() - invalidate_local_hash_match_cache() - self.search_result_timestamps.clear() - self.logger.info("Cleared backend search caches") - return web.json_response({"success": True, "cleared": "all"}) - - async def _check_credential_helper(request, payload_key, check_func, log_name): - try: - data = await request.json() - val = data.get(payload_key, "") - result = await asyncio.to_thread(check_func, val) - return web.json_response(result) - except Exception as e: - self.logger.exception(f"{log_name} check error: {e}") - return web.json_response({"error": str(e)}, status=500) - - @routes.post("/model_resolver/civitai/session-token/check") - async def civitai_session_token_check_route(request): - """Check whether a CivitAI browser session token is valid.""" - return await _check_credential_helper( - request, "civitai_session_token", check_civitai_session_token, "CivitAI session token" - ) - - @routes.post("/model_resolver/civitai/api-key/check") - async def civitai_api_key_check_route(request): - """Check whether a CivitAI API key is valid.""" - return await _check_credential_helper( - request, "civitai_key", check_civitai_api_key, "CivitAI API key" - ) - - @routes.post("/model_resolver/huggingface/token/check") - async def huggingface_token_check_route(request): - """Check whether a HuggingFace token is valid.""" - return await _check_credential_helper( - request, "hf_token", check_huggingface_token, "HuggingFace token" - ) - - @routes.post("/model_resolver/brave/api-key/check") - async def brave_api_key_check_route(request): - """Check whether a Brave Search API key is valid.""" - return await _check_credential_helper( - request, "brave_search_api_key", check_brave_search_api_key, "Brave Search API key" - ) - - @routes.get("/model_resolver/huggingface/author-index/status") - @json_api_endpoint("HuggingFace author index status") - async def huggingface_author_index_status_route(request): - """Return local HuggingFace author fallback index status.""" - return web.json_response( - get_known_author_fallback_indexes_status() - ) - - @routes.post("/model_resolver/huggingface/author-index/refresh") - @json_api_endpoint("HuggingFace author index refresh") - async def huggingface_author_index_refresh_route(request): - """Refresh HuggingFace author fallback index.""" - data = await request.json() - hf_token = data.get("hf_token", "") - result = await asyncio.to_thread( - refresh_known_author_fallback_indexes, hf_token or None - ) - clear_huggingface_search_cache() - return web.json_response(result) - - @routes.get("/model_resolver/model-list/status") - @json_api_endpoint("Model list status") - async def model_list_status_route(request): - """Return local model-list status and optionally compare with GitHub.""" - check_remote = ( - str(request.query.get("check_remote", "")).lower() - in {"1", "true", "yes"} - ) - return web.json_response( - get_model_list_update_status(check_remote=check_remote) - ) - - @routes.post("/model_resolver/model-list/update") - @json_api_endpoint("Model list update") - async def model_list_update_route(request): - """Download latest ComfyUI-Manager model-list.json.""" - result = await asyncio.to_thread(update_model_list_from_remote) - clear_all_search_caches() - self.search_result_timestamps.clear() - return web.json_response(result) - - @routes.post("/model_resolver/download") - @json_api_endpoint("download", return_success_on_error=True) - async def download_model(request): - """Start downloading a model.""" - data = await request.json() - url = data.get("url", "") - filename = data.get("filename", "") - category = data.get("category", "checkpoints") - category = normalize_download_category(category) - subfolder = data.get("subfolder", "") - base_directory = data.get("base_directory", "") - path_metadata = data.get("path_metadata", {}) - if not isinstance(path_metadata, dict): - path_metadata = {} - download_metadata = data.get("download_metadata") or data.get( - "metadata", {} - ) - if not isinstance(download_metadata, dict): - download_metadata = {} - download_metadata = dict(download_metadata) - settings = load_resolver_settings() - if not base_directory: - base_directory = get_default_root_for_category(category, settings) - subfolder = resolve_download_subfolder( - category, - subfolder, - path_metadata, - settings, - ) - - if not url: - return web.json_response( - {"error": "URL is required"}, status=400 - ) - try: - url = await asyncio.to_thread(validate_public_http_url, url) - except UnsafeUrlError as exc: - return web.json_response({"error": str(exc)}, status=400) - - from urllib.parse import urlparse as _download_urlparse - - download_host = _download_urlparse(url).hostname - - if not filename: - # Extract filename from URL - from urllib.parse import unquote, urlparse - - parsed = urlparse(url) - filename = unquote(parsed.path.split("/")[-1]) - filename = sanitize_download_filename(filename) - - if not filename: - return web.json_response( - {"error": "Could not determine filename"}, status=400 - ) - if not is_allowed_model_download_filename(filename): - return web.json_response( - {"error": "Unsupported model file extension"}, status=400 - ) - - # Build headers if needed - headers = {} - if host_matches_domain(download_host, "huggingface.co"): - hf_token = data.get("hf_token", "") - if hf_token: - headers["Authorization"] = f"Bearer {hf_token}" - elif host_matches_domain( - download_host, - "civitai.com", - "civitai.red", - ): - civitai_key = data.get("civitai_key", "") - if civitai_key and "token=" not in url: - url += ( - f"{'&' if '?' in url else '?'}token={civitai_key}" - ) - civitai_session_token = str( - data.get("civitai_session_token", "") or "" - ).strip() - if civitai_session_token: - headers["Cookie"] = ( - f"__Secure-civitai-token={civitai_session_token}" - ) - - def _first_metadata_value(*values): - return first_non_empty(*values) - - def _metadata_int(value): - return to_int(value) - - inferred_source = "" - if host_matches_domain(download_host, "civitai.com", "civitai.red"): - inferred_source = "civitai" - elif host_matches_domain(download_host, "huggingface.co"): - inferred_source = "huggingface" - - download_metadata.setdefault("filename", filename) - download_metadata.setdefault("category", category) - download_metadata.setdefault("download_url", url) - download_metadata.setdefault("source_url", url) - download_metadata.setdefault("path_metadata", path_metadata) - download_metadata.setdefault( - "source", - _first_metadata_value( - download_metadata.get("details_source"), - path_metadata.get("source"), - inferred_source, - ), - ) - - model_id = _metadata_int( - _first_metadata_value( - download_metadata.get("model_id"), - download_metadata.get("modelId"), - path_metadata.get("model_id"), - ) - ) - version_id = _metadata_int( - _first_metadata_value( - download_metadata.get("version_id"), - download_metadata.get("versionId"), - path_metadata.get("version_id"), - ) - ) - source_name = str( - _first_metadata_value( - download_metadata.get("details_source"), - download_metadata.get("source"), - ) - ).lower() - try: - if ( - source_name == "civitai" - and model_id - and not download_metadata.get("civitai_details") - ): - details = await asyncio.to_thread( - get_civitai_model_details, - model_id, - version_id, - data.get("civitai_key", ""), - ) - if details: - download_metadata["civitai_details"] = details - elif ( - source_name == "civarchive" - and model_id - and not download_metadata.get("civitai_details") - ): - details = await asyncio.to_thread( - get_civarchive_model_details, - model_id, - version_id, - ) - if details: - download_metadata["civitai_details"] = details - except Exception as metadata_error: - self.logger.warning( - f"Model metadata lookup failed: {metadata_error}" - ) - - target_directory = "" - target_path = "" - try: - import os as _download_os - - target_directory = ( - get_download_directory(category, base_directory) or "" - ) - if target_directory and subfolder: - target_directory = _download_os.path.join( - target_directory, *split_path_segments(subfolder) - ) - if target_directory: - target_path = _download_os.path.join( - target_directory, filename - ) - except Exception: - target_directory = "" - target_path = "" - - # Start background download - download_id = start_background_download( - url=url, - filename=filename, - category=category, - headers=headers if headers else None, - subfolder=subfolder, - base_directory=base_directory, - metadata=download_metadata, - ) - - return web.json_response( - { - "success": True, - "download_id": download_id, - "filename": filename, - "category": category, - "path": target_path, - "directory": target_directory, - "download_backend": settings.get("download_backend", "python"), - } - ) - - @routes.get("/model_resolver/progress/{download_id}") - @json_api_endpoint("progress") - async def get_download_progress(request): - """Get progress for a specific download.""" - 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 - ) - - @routes.get("/model_resolver/progress") - @json_api_endpoint("progress") - async def get_all_downloads_progress(request): - """Get progress for all downloads.""" - progress = get_all_progress() - return web.json_response(progress) - - @routes.post("/model_resolver/cancel/{download_id}") - @json_api_endpoint("cancel", return_success_on_error=True) - async def cancel_download_route(request): - """Cancel a download in progress.""" - download_id = request.match_info["download_id"] - cancel_download(download_id) - return web.json_response({"success": True}) - - @routes.post("/model_resolver/pause/{download_id}") - @json_api_endpoint("pause", return_success_on_error=True) - async def pause_download_route(request): - """Pause an aria2 download.""" - download_id = request.match_info["download_id"] - result = pause_download(download_id) - status = 200 if result.get("success") else 400 - return web.json_response(result, status=status) - - @routes.post("/model_resolver/resume/{download_id}") - @json_api_endpoint("resume", return_success_on_error=True) - async def resume_download_route(request): - """Resume an aria2 download.""" - download_id = request.match_info["download_id"] - result = resume_download(download_id) - status = 200 if result.get("success") else 400 - return web.json_response(result, status=status) - - @routes.post("/model_resolver/clear_completed_downloads") - @json_api_endpoint("clear_completed_downloads", return_success_on_error=True) - async def clear_completed_downloads_route(request): - """Clear completed and failed downloads from progress memory.""" - clear_completed_downloads() - return web.json_response({"success": True}) - - @routes.get("/model_resolver/aria2/status") - @routes.post("/model_resolver/aria2/status") - @json_api_endpoint("aria2 status") - async def aria2_status_route(request): - """Report aria2 availability using saved settings, with optional override via POST.""" - settings = await get_override_settings_from_request(request) - return web.json_response(await asyncio.to_thread(get_aria2_status, settings)) - - - @routes.post("/model_resolver/aria2/start") - @json_api_endpoint("aria2 start", return_success_on_error=True) - async def aria2_start_route(request): - """Start the aria2 daemon without starting a download.""" - settings = await get_override_settings_from_request(request) - result = await asyncio.to_thread(start_aria2_daemon, settings) - status = 200 if result.get("success") else 400 - return web.json_response(result, status=status) - - @routes.get("/model_resolver/aria2/stop") - @routes.post("/model_resolver/aria2/stop") - @json_api_endpoint("aria2 stop", return_success_on_error=True) - async def aria2_stop_route(request): - """Stop the aria2 daemon started by Model Resolver.""" - result = await asyncio.to_thread(stop_aria2_daemon) - status = 200 if result.get("success") else 400 - return web.json_response(result, status=status) - - @routes.post("/model_resolver/aria2/install") - @json_api_endpoint("aria2 install") - async def aria2_install_route(request): - """Download official aria2c binary and save its path.""" - try: - payload = await request.json() - except Exception: - payload = {} - force = False - if isinstance(payload, dict): - force = to_bool(payload.get("force"), False) - try: - install_result = await asyncio.to_thread(install_aria2_engine, force) - except Aria2InstallError as exc: - self.logger.warning(f"Model Resolver aria2 install failed: {exc}") - return web.json_response({"success": False, "error": str(exc)}, status=500) - settings = await asyncio.to_thread( - save_resolver_settings, - { - "aria2c_path": install_result.get("aria2c_path", ""), - "download_backend": "aria2", - }, - ) - install_result["settings"] = settings - return web.json_response(install_result) - - @routes.get("/model_resolver/directories") - @json_api_endpoint("directories") - async def get_directories(request): - """Get available model directories.""" - import folder_paths - - categories = get_enabled_download_categories(list(folder_paths.folder_names_and_paths.keys())) - directories = {} - for cat in categories: - path = get_download_directory(cat) - if path: - directories[cat] = path - - return web.json_response(directories) - - @routes.get("/model_resolver/root-directories") - @json_api_endpoint("root directories") - async def get_root_directories(request): - """Get configured ComfyUI root directories for path settings.""" - import os - - import folder_paths - - known_categories = set(folder_paths.folder_names_and_paths.keys()) - skip_categories = {"custom_nodes", "configs"} - categories = get_enabled_download_categories(list(known_categories)) - category_source_keys = {} - for cat in [*known_categories]: - folder_key = normalize_download_category(cat) - if folder_key: - category_source_keys.setdefault(folder_key, []).append(cat) - - roots = {} - settings = load_resolver_settings() - comfy_root = get_comfy_root_path(folder_paths) - for cat in categories: - folder_key = normalize_download_category(cat) - if folder_key in known_categories: - raw_candidate_keys = [folder_key] - else: - raw_candidate_keys = [ - folder_key, - *TEMPLATE_KEY_ALIASES.get(folder_key, ()), - *category_source_keys.get(folder_key, []), - ] - if folder_key == "ultralytics": - raw_candidate_keys = [ - candidate_key - for candidate_key in raw_candidate_keys - if str(candidate_key or "").strip().lower() != "yolo" - ] - - candidate_keys = [] - for candidate_key in raw_candidate_keys: - if ( - candidate_key - and candidate_key in known_categories - and candidate_key not in candidate_keys - ): - candidate_keys.append(candidate_key) - paths = [] - for candidate_key in candidate_keys: - paths.extend(folder_paths.get_folder_paths(candidate_key) or []) - if folder_key == "ultralytics": - normalized_ultralytics_paths = [] - for path in paths: - normalized_path = os.path.normpath(str(path or "")) - basename = os.path.basename(normalized_path).lower() - if basename == "yolo": - continue - parent_dir = os.path.dirname(normalized_path) - if ( - basename in {"bbox", "segm"} - and os.path.basename(parent_dir).lower() - == "ultralytics" - ): - normalized_path = parent_dir - normalized_ultralytics_paths.append(normalized_path) - paths = normalized_ultralytics_paths - preferred_directory = get_default_root_for_category( - folder_key, settings - ) - normalized_paths = dedupe_local_base_directories( - paths, - preferred_directory=preferred_directory, - comfy_root=comfy_root, - ) - roots[folder_key] = normalized_paths - - for raw_key in known_categories: - if ( - not raw_key - or raw_key in skip_categories - or raw_key in roots - ): - continue - raw_paths = folder_paths.get_folder_paths(raw_key) or [] - roots[raw_key] = dedupe_local_base_directories( - raw_paths, - comfy_root=comfy_root, - ) - - return web.json_response(roots) - - @routes.get("/model_resolver/path-template-suggestions") - @json_api_endpoint("path template suggestions") - async def get_path_template_suggestions(request): - """Infer path template presets from existing local model folders.""" - from ..sources.popular import get_base_models_config - - force_rescan = request.query.get("force") == "1" - models = await asyncio.to_thread(get_model_files, force_rescan) - base_models_config = get_base_models_config() - suggestions = await asyncio.to_thread( - infer_download_path_templates, - models, - base_models_config, - ) - return web.json_response(suggestions) - - @routes.get("/model_resolver/capabilities") - @json_api_endpoint("capabilities") - async def get_capabilities(request): - """Get optional source capabilities available in this install.""" - from ..workflow_analyzer import NODE_TYPE_MODEL_WIDGET_CATEGORIES - return web.json_response( - { - "sources": { - "civarchive": is_civarchive_available(), - "lora_manager_archive": is_lora_manager_archive_available() - }, - "node_rules": NODE_TYPE_MODEL_WIDGET_CATEGORIES - } - ) - - register_version_routes(routes, web, json_api_endpoint) - - @routes.get("/model_resolver/subfolders/{category}") - @json_api_endpoint("subfolders") - async def get_subfolders(request): - """Get known subfolders for a category using ComfyUI folder_paths.""" - import os - - import folder_paths - - raw_category = (request.match_info.get("category") or "").strip() - category = normalize_download_category(raw_category) - - if not category or category == "unknown": - return web.json_response([]) - - known_categories = set(folder_paths.folder_names_and_paths.keys()) - folder_keys = get_category_folder_keys(category) - available_folder_keys = [ - folder_key for folder_key in folder_keys if folder_key in known_categories - ] - if not available_folder_keys: - self.logger.debug( - f"Model Resolver: skipping subfolder lookup for unknown category '{raw_category}' -> '{category}'" - ) - return web.json_response([]) - - subfolders = {} - settings = load_resolver_settings() - comfy_root = get_comfy_root_path(folder_paths) - preferred_directory = ( - get_default_root_for_category(category, settings) - or get_download_directory(category) - or "" - ) - - def add_subfolder(rel_path, base_dir=""): - rel_path = "/".join(split_path_segments(rel_path)) - if not rel_path or rel_path == ".": - return - base_dir = os.path.abspath(base_dir) if base_dir else "" - base_identity = ( - get_local_path_identity(base_dir) if base_dir else "" - ) - key = (rel_path.lower(), base_identity) - base_label = ( - os.path.basename(os.path.normpath(base_dir)) - if base_dir - else "" - ) - current = subfolders.get(key) - if current and not prefer_local_base_directory( - base_dir, - current.get("base_directory", ""), - preferred_directory, - comfy_root, - ): - return - subfolders[key] = { - "value": rel_path, - "label": rel_path, - "base_label": base_label, - "base_directory": base_dir, - } - - raw_base_dirs = [] - for folder_key in available_folder_keys: - for base_dir in folder_paths.get_folder_paths(folder_key) or []: - if not base_dir or not os.path.isdir(base_dir): - continue - raw_base_dirs.append(base_dir) - base_dirs = dedupe_local_base_directories( - raw_base_dirs, - preferred_directory=preferred_directory, - comfy_root=comfy_root, - ) - - def find_base_dir(full_path): - if not full_path: - return "" - full_path_identity = get_local_path_identity(full_path) - for base_dir in base_dirs: - base_identity = get_local_path_identity(base_dir) - try: - if ( - os.path.commonpath( - [full_path_identity, base_identity] - ) - == base_identity - ): - return base_dir - except Exception: - continue - return "" - - for folder_key in available_folder_keys: - filenames = folder_paths.get_filename_list(folder_key) or [] - for rel_path in filenames: - if not isinstance(rel_path, str): - continue - base_dir = "" - try: - base_dir = find_base_dir( - folder_paths.get_full_path(folder_key, rel_path) - ) - except Exception: - base_dir = "" - parts = split_path_segments(rel_path) - if len(parts) <= 1: - continue - current = "" - for part in parts[:-1]: - current = f"{current}/{part}" if current else part - add_subfolder(current, base_dir) - - for base_dir in base_dirs: - for root, dirs, _files in os.walk(base_dir): - rel_root = os.path.relpath(root, base_dir) - for dirname in dirs: - rel_path = ( - dirname - if rel_root in ("", ".") - else os.path.join(rel_root, dirname) - ) - add_subfolder(rel_path, base_dir) - - value_counts = {} - for item in subfolders.values(): - value_key = item.get("value", "").lower() - value_counts[value_key] = value_counts.get(value_key, 0) + 1 - - response_items = [] - for item in subfolders.values(): - value = item.get("value", "") - base_label = item.get("base_label", "") - label = ( - f"{value} ({base_label})" - if base_label and value_counts.get(value.lower(), 0) > 1 - else value - ) - response_items.append( - { - "value": value, - "label": label, - "base_directory": item.get("base_directory", ""), - } - ) - - return web.json_response( - sorted( - response_items, - key=lambda item: ( - item.get("value", "").lower(), - item.get("base_directory", "").lower(), - ), - ) - ) + register_search_routes(route_context) + register_download_routes(route_context) + register_directory_routes(route_context) register_settings_routes(routes, web, json_api_endpoint) self.routes_setup = True diff --git a/core/routes/search.py b/core/routes/search.py new file mode 100644 index 0000000..30d9d7c --- /dev/null +++ b/core/routes/search.py @@ -0,0 +1,1185 @@ +"""Search route registration.""" + +def register_search_routes(context): + CivArchiveSearchError = context.get('CivArchiveSearchError') + asyncio = context.get('asyncio') + build_search_result = context.get('build_search_result') + cancel_progress_response = context.get('cancel_progress_response') + check_brave_search_api_key = context.get('check_brave_search_api_key') + check_civitai_api_key = context.get('check_civitai_api_key') + check_civitai_session_token = context.get('check_civitai_session_token') + check_huggingface_token = context.get('check_huggingface_token') + clear_all_search_caches = context.get('clear_all_search_caches') + clear_civarchive_search_cache = context.get('clear_civarchive_search_cache') + clear_civitai_search_cache = context.get('clear_civitai_search_cache') + clear_huggingface_search_cache = context.get('clear_huggingface_search_cache') + clear_lora_manager_archive_search_cache = context.get('clear_lora_manager_archive_search_cache') + extract_sha256_from_metadata = context.get('extract_sha256_from_metadata') + format_size_bytes = context.get('format_size_bytes') + get_civitai_download_url = context.get('get_civitai_download_url') + get_known_author_fallback_indexes_status = context.get('get_known_author_fallback_indexes_status') + get_model_list_update_status = context.get('get_model_list_update_status') + get_popular_model_url = context.get('get_popular_model_url') + get_progress_response = context.get('get_progress_response') + invalidate_local_hash_match_cache = context.get('invalidate_local_hash_match_cache') + invalidate_model_files_cache = context.get('invalidate_model_files_cache') + json_api_endpoint = context.get('json_api_endpoint') + refresh_known_author_fallback_indexes = context.get('refresh_known_author_fallback_indexes') + reload_model_list = context.get('reload_model_list') + reload_popular_databases = context.get('reload_popular_databases') + resolve_civarchive_model_version = context.get('resolve_civarchive_model_version') + resolve_urn = context.get('resolve_urn') + routes = context.get('routes') + search_civarchive_for_file = context.get('search_civarchive_for_file') + search_civitai = context.get('search_civitai') + search_civitai_for_file = context.get('search_civitai_for_file') + search_huggingface_for_file = context.get('search_huggingface_for_file') + search_local_matches_by_hash = context.get('search_local_matches_by_hash') + search_lora_manager_archive_for_file = context.get('search_lora_manager_archive_for_file') + search_model_list = context.get('search_model_list') + self = context.get('self') + to_bool = context.get('to_bool') + to_int = context.get('to_int') + update_model_list_from_remote = context.get('update_model_list_from_remote') + web = context.get('web') + + @routes.get("/model_resolver/search-progress/{progress_id}") + @json_api_endpoint("search-progress") + async def get_search_progress_route(request): + """Return live progress for an in-flight source search.""" + return get_progress_response( + self.search_tracker, + request, + not_found_payload={"exists": False}, + found_wrapper=lambda prog: {"exists": True, **prog} + ) + + @routes.post("/model_resolver/search-cancel/{progress_id}") + @json_api_endpoint("search-cancel", return_success_on_error=True) + async def cancel_search_progress_route(request): + """Mark an in-flight source search as cancelled.""" + return cancel_progress_response( + self.search_tracker, + request, + cancel_message="Cancelled" + ) + + @routes.post("/model_resolver/search") + async def search_sources(request): + """Search for model download sources.""" + try: + class SearchCancelled(BaseException): + # Progress helpers swallow ordinary callback errors; cancellation + # must bubble out of source loops and stop follow-up requests. + pass + + def raise_if_search_cancelled(source=""): + if self.search_tracker.is_cancelled(progress_id): + raise SearchCancelled("Search cancelled") + + def format_log_value(value): + if value is None or value == "": + return None + if isinstance(value, bool): + return "yes" if value else "no" + if isinstance(value, (list, tuple, set)): + return ",".join(str(item) for item in value) + + text = str(value) + if any(char.isspace() for char in text): + return '"' + text.replace("\\", "\\\\").replace('"', '\\"') + '"' + return text + + def format_log_size(value): + return format_size_bytes(value, include_space=False) + + def format_log_fields(**fields): + parts = [] + for key, value in fields.items(): + formatted = format_log_value(value) + if formatted is not None: + parts.append(f"{key}={formatted}") + return " ".join(parts) + + def normalize_result_extra(extra): + if not extra: + return {} + normalized = dict(extra) + model_id = normalized.pop("model_id", None) + version_id = normalized.pop("version_id", None) + if model_id or version_id: + normalized["ids"] = ( + f"{model_id}@{version_id}" + if model_id and version_id + else model_id or version_id + ) + if "files_count" in normalized: + normalized["files"] = normalized.pop("files_count") + return normalized + + def format_result_details(result, extra=None): + if isinstance(result, list): + return format_log_fields(count=len(result)) + if not isinstance(result, dict): + fields = {"result": "none" if result is None else result} + fields.update(normalize_result_extra(extra)) + return format_log_fields(**fields) + + model_id = result.get("model_id") + version_id = result.get("version_id") + ids = ( + f"{model_id}@{version_id}" + if model_id and version_id + else model_id or version_id + ) + fields = { + "name": result.get("name"), + "file": result.get("filename") or result.get("path"), + "match": result.get("match_type"), + "repo": result.get("repo_id") or result.get("repo"), + "ids": ids, + "size": format_log_size(result.get("size")), + "files": result.get("files_count"), + } + fields.update(normalize_result_extra(extra)) + return format_log_fields(**fields) + + def log_search_result(source_name, result, extra=None): + details = format_result_details(result, extra) + if result and ( + not isinstance(result, list) or len(result) > 0 + ): + self.logger.info(f"Search [{source_name}] found {details}") + else: + self.logger.info(f"Search [{source_name}] miss {details}") + + data = await request.json() + filename = data.get("filename", "") + category = data.get("category", "") + base_model_context = data.get("base_model_context", "") + progress_id = str(data.get("progress_id") or "").strip() + progress_source = str(data.get("progress_source") or "").strip() + civitai_candidate_limit = to_int(data.get("civitai_candidate_limit"), 5) + civitai_candidate_limit = max( + 1, min(civitai_candidate_limit, 20) + ) + civarchive_candidate_limit = to_int(data.get("civarchive_candidate_limit"), 10) + civarchive_candidate_limit = max( + 1, min(civarchive_candidate_limit, 30) + ) + # Handle both boolean and string forms + is_urn_raw = data.get("is_urn", False) + civitai_key = data.get("civitai_key", "") + civitai_session_token = data.get("civitai_session_token", "") + hf_token = data.get("hf_token", "") + brave_search_api_key = data.get("brave_search_api_key", "") + civitai_use_trpc_search = data.get( + "civitai_use_trpc_search", True + ) + civitai_use_api_search = data.get( + "civitai_use_api_search", True + ) + civitai_use_html_fallback = data.get( + "civitai_use_html_fallback", True + ) + hf_use_api_search = data.get("hf_use_api_search", True) + hf_use_comfy_org_fallback = data.get( + "hf_use_comfy_org_fallback", True + ) + hf_use_brave_fallback = data.get( + "hf_use_brave_fallback", True + ) + is_urn = to_bool(is_urn_raw, False) + hf_use_api_search = to_bool(hf_use_api_search, True) + civitai_use_trpc_search = to_bool(civitai_use_trpc_search, True) + civitai_use_api_search = to_bool(civitai_use_api_search, True) + civitai_use_html_fallback = to_bool(civitai_use_html_fallback, True) + hf_use_comfy_org_fallback = to_bool(hf_use_comfy_org_fallback, True) + hf_use_brave_fallback = to_bool(hf_use_brave_fallback, 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", + "civarchive", + "lora_manager_archive", + } + + search_local = "local" in normalized_sources + search_huggingface_source = "huggingface" in normalized_sources + search_civitai_source = "civitai" in normalized_sources + search_civarchive_source = "civarchive" in normalized_sources + search_lora_manager_archive_source = ( + "lora_manager_archive" in normalized_sources + ) + if not progress_source: + progress_source = ( + next(iter(normalized_sources)) + if len(normalized_sources) == 1 + else "all" + ) + force_search = to_bool(data.get("force_search"), False) + + self.search_tracker.update( + progress_id, + progress_source, + "starting", + "Preparing search", + 8, + ) + raise_if_search_cancelled(progress_source) + + if force_search: + self.search_tracker.update( + progress_id, + progress_source, + "cache", + "Refreshing search caches", + 12, + ) + if search_local: + reload_popular_databases() + reload_model_list() + if search_huggingface_source: + clear_huggingface_search_cache() + if search_civitai_source: + clear_civitai_search_cache() + if search_civarchive_source: + clear_civarchive_search_cache() + if search_lora_manager_archive_source: + clear_lora_manager_archive_search_cache() + self.logger.debug( + "Force search enabled: cleared cache " + + format_log_fields(sources=sorted(normalized_sources)) + ) + raise_if_search_cancelled(progress_source) + + self.logger.info( + f"Search [{','.join(sorted(normalized_sources))}] request " + + format_log_fields( + file=filename, + cat=category, + urn=is_urn, + ids=( + f"{data.get('model_id')}@{data.get('version_id')}" + if data.get("model_id") and data.get("version_id") + else data.get("model_id") + or data.get("version_id") + ), + base=base_model_context, + force=force_search, + ) + ) + + results = { + "popular": None, + "model_list": None, + "huggingface": None, + "civitai": None, + "civarchive": None, + "lora_manager_archive": None, + "local_hash_matches": [], + "found": False, + "searched_sources": sorted(normalized_sources), + "source_errors": {}, + } + + def current_search_timestamp(): + from datetime import datetime, timezone + + return ( + datetime.now(timezone.utc) + .replace(microsecond=0) + .isoformat() + ) + + def get_search_result_signature(source_key, result): + if isinstance(result, list): + return "|".join( + get_search_result_signature(source_key, item) + for item in result + ) + if not isinstance(result, dict): + return "" + + parts = [ + source_key, + result.get("download_url") + or result.get("url") + or result.get("model_url") + or "", + result.get("filename") or result.get("path") or "", + result.get("repo_id") or result.get("repo") or "", + result.get("model_id") or "", + result.get("version_id") or "", + result.get("name") or "", + ] + return "::".join(str(part).strip() for part in parts) + + def stamp_search_result(source_key, result): + if isinstance(result, list): + return [ + stamp_search_result(source_key, item) + for item in result + ] + if not isinstance(result, dict): + return result + + signature = get_search_result_signature(source_key, result) + if not signature: + return result + + timestamp = ( + result.get("searched_at") + or result.get("searchedAt") + or ( + None + if force_search + else self.search_result_timestamps.get(signature) + ) + ) + if not timestamp: + timestamp = current_search_timestamp() + if force_search: + self.search_result_timestamps[signature] = timestamp + else: + self.search_result_timestamps.setdefault( + signature, timestamp + ) + result["searched_at"] = timestamp + return result + + def stamp_search_results(payload): + for source_key in ( + "popular", + "model_list", + "huggingface", + "civitai", + "civarchive", + "lora_manager_archive", + ): + if payload.get(source_key): + payload[source_key] = stamp_search_result( + source_key, payload[source_key] + ) + return payload + + def mark_any_model_fallback(result): + if isinstance(result, list): + return [ + mark_any_model_fallback(item) + for item in result + ] + if not isinstance(result, dict): + return result + + marked = dict(result) + marked["any_model_match"] = True + marked["base_model_fallback"] = True + marked["requested_base_model"] = base_model_context + return marked + + def iter_result_items(result): + if isinstance(result, list): + for item in result: + if isinstance(item, dict): + yield item + elif isinstance(result, dict): + yield result + + def collect_local_hash_matches(payload): + matches = [] + seen_match_paths = set() + seen_hashes = set() + for source_key in ( + "huggingface", + "civitai", + "civarchive", + "lora_manager_archive", + "popular", + "model_list", + ): + for source_result in iter_result_items(payload.get(source_key)): + raise_if_search_cancelled(source_key) + sha256 = extract_sha256_from_metadata(source_result) + if not sha256: + continue + + if sha256 in seen_hashes: + continue + seen_hashes.add(sha256) + + self.search_tracker.update( + progress_id, + progress_source, + "local_hash", + "Checking local metadata hashes", + 94, + ) + try: + hash_matches = search_local_matches_by_hash( + sha256, + category=category or None, + max_matches=20, + force_rescan=force_search, + ) + except Exception as hash_error: + self.logger.warning( + f"Local metadata hash lookup failed for {source_key}:{sha256}: {hash_error}" + ) + continue + + for match in hash_matches: + model_path = ( + match.get("model", {}).get("path") + or match.get("path") + or "" + ) + path_key = model_path.lower() + if path_key and path_key in seen_match_paths: + continue + if path_key: + seen_match_paths.add(path_key) + enriched = { + **match, + "hash_lookup_source": source_key, + "hash_lookup_filename": source_result.get("filename") + or source_result.get("path") + or filename, + "hash_lookup_sha256": sha256, + } + matches.append(enriched) + + if matches: + self.logger.info( + "Search local hash matches " + + format_log_fields(count=len(matches)) + ) + return matches + + def make_source_progress_callback( + source_key, + percent_min=None, + percent_max=None, + ): + def source_progress_callback(payload): + raise_if_search_cancelled(source_key) + if not isinstance(payload, dict): + return + + progress_payload = dict(payload) + stage = progress_payload.pop("stage", "running") + message = progress_payload.pop( + "message", "Searching..." + ) + percent = progress_payload.pop("percent", None) + status = progress_payload.pop("status", "running") + progress_payload.pop("source", None) + + if ( + percent is not None + and percent_min is not None + and percent_max is not None + ): + try: + normalized_percent = max( + 0.0, min(100.0, float(percent)) + ) + percent = percent_min + ( + normalized_percent / 100.0 + ) * (percent_max - percent_min) + except (TypeError, ValueError): + percent = None + + self.search_tracker.update( + progress_id, + source_key, + stage, + message, + percent, + status=status, + **progress_payload, + ) + + return source_progress_callback + + def run_source_search( + source_key, + search_task_fn, + initial_stage="query", + initial_message=None, + initial_percent=30, + log_start_fields=None, + error_handlers=None, + ): + if initial_message is None: + initial_message = f"Querying {source_key.capitalize()}" + raise_if_search_cancelled(source_key) + self.search_tracker.update( + progress_id, + source_key, + initial_stage, + initial_message, + initial_percent, + ) + start_fields = log_start_fields or {"file": filename} + self.logger.info( + f"Search [{source_key}] start " + + format_log_fields(**start_fields) + ) + try: + raise_if_search_cancelled(source_key) + source_results, source_found = search_task_fn() + raise_if_search_cancelled(source_key) + done_messages = { + "local": "Local database checked", + "huggingface": "HuggingFace checked", + "civitai": "CivitAI checked", + "civarchive": "CivArchive checked", + "lora_manager_archive": "LoRA Manager archive checked", + } + done_msg = done_messages.get(source_key, f"{source_key} checked") + self.search_tracker.update( + progress_id, + source_key, + "done", + done_msg, + 92, + ) + return source_results, source_found + except Exception as e: + if error_handlers: + for exc_type, handler_fn in error_handlers.items(): + if isinstance(e, exc_type): + return handler_fn(e) + raise e + + def execute_search_with_fallback( + source_key, + search_fn, + any_model_label, + ): + raise_if_search_cancelled(source_key) + res = search_fn( + base_model_context or None, + make_source_progress_callback(source_key), + ) + raise_if_search_cancelled(source_key) + log_search_result(source_key, res) + + if not res and base_model_context: + raise_if_search_cancelled(source_key) + self.search_tracker.update( + progress_id, + source_key, + "any_model", + f"Retrying {any_model_label} any model", + 72, + ) + self.logger.info( + f"Search [{source_key}] retry any model " + + format_log_fields( + file=filename, + cat=category, + base=base_model_context, + ) + ) + res = search_fn( + None, + make_source_progress_callback(source_key, 72, 92), + ) + raise_if_search_cancelled(source_key) + log_search_result(f"{source_key}/any_model", res) + if res: + res = mark_any_model_fallback(res) + return res + + def search_local_sources(): + def task(): + source_results = {"popular": None, "model_list": None} + source_found = False + + popular_info = get_popular_model_url(filename) + log_search_result("popular", popular_info) + self.search_tracker.update( + progress_id, + "local", + "model_list", + "Checking local model database", + 58, + ) + model_list_result = search_model_list(filename) + log_search_result( + "model_list", + model_list_result, + { + "confidence": model_list_result.get("confidence") + if model_list_result + else None + }, + ) + if popular_info: + popular_result = { + "source": "popular", + "filename": filename, + **popular_info, + } + if ( + model_list_result + and model_list_result.get("filename", "").lower() + == filename.lower() + and model_list_result.get("size") + ): + popular_result["size"] = model_list_result.get( + "size" + ) + source_results["popular"] = popular_result + source_found = True + + if model_list_result: + confidence = model_list_result.get("confidence", 0) + if (is_urn and confidence >= 70) or not is_urn: + source_results["model_list"] = model_list_result + source_found = True + + return source_results, source_found + + return run_source_search( + "local", + task, + initial_stage="popular", + initial_message="Checking popular models", + initial_percent=28, + log_start_fields={"file": filename, "cat": category}, + ) + + def search_huggingface_source_task(): + def task(): + hf_result = search_huggingface_for_file( + filename, + token=hf_token or None, + brave_api_key=brave_search_api_key or None, + use_api_search=hf_use_api_search, + use_comfy_org_fallback=hf_use_comfy_org_fallback, + use_brave_fallback=hf_use_brave_fallback, + force_refresh=force_search, + progress_callback=make_source_progress_callback( + "huggingface" + ), + ) + log_search_result("huggingface", hf_result) + return {"huggingface": hf_result}, bool(hf_result) + + return run_source_search( + "huggingface", + task, + initial_stage="query", + initial_message="Querying HuggingFace", + initial_percent=32, + log_start_fields={"file": filename}, + ) + + def search_civitai_source_task(): + def task(): + source_results = {"civitai": None} + source_found = False + + if is_urn: + model_id_val = data.get("model_id") + version_id_val = data.get("version_id") + + if model_id_val and version_id_val: + self.search_tracker.update( + progress_id, + "civitai", + "urn", + "Resolving CivitAI URN", + 46, + ) + model_info = resolve_urn(model_id_val, version_id_val) + if model_info: + self.search_tracker.update( + progress_id, + "civitai", + "file", + "Selecting CivitAI file", + 76, + ) + 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_val + ) + source_results["civitai"] = build_search_result( + "civitai", + model_id=model_id_val, + version_id=version_id_val, + name=model_info.get("model_name"), + version_name=model_info.get("version_name"), + filename=model_info.get("expected_filename"), + type=category, + download_url=download_url, + url=f"https://civitai.com/models/{model_id_val}?modelVersionId={version_id_val}", + size=primary_file.get("size"), + base_model=model_info.get("base_model"), + tags=model_info.get("tags", []), + sha256=primary_file.get("sha256") + or (primary_file.get("hashes") or {}).get("SHA256") + or (primary_file.get("hashes") or {}).get("sha256"), + hashes=primary_file.get("hashes") or {}, + match_type="exact", + confidence=100.0, + ) + log_search_result( + "civitai/urn", + source_results["civitai"], + { + "files_count": len( + model_info.get("files", []) + ) + }, + ) + source_found = True + else: + log_search_result( + "civitai/urn", + None, + { + "model_id": model_id_val, + "version_id": version_id_val, + }, + ) + elif category: + self.search_tracker.update( + progress_id, + "civitai", + "fallback", + "Searching CivitAI fallback", + 58, + ) + self.logger.info( + "Search [civitai] URN ids missing; falling back" + ) + civitai_results = search_civitai( + filename, + model_type=category, + ) + log_search_result( + "civitai/fallback", + civitai_results[0] if civitai_results else None, + { + "results_count": len(civitai_results), + }, + ) + if civitai_results: + first_result = civitai_results[0] + source_results["civitai"] = build_search_result( + "civitai", + model_id=first_result.get("model_id"), + version_id=first_result.get("version_id"), + 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"), + base_model=first_result.get("base_model"), + tags=first_result.get("tags", []), + ) + source_found = True + else: + civitai_result = execute_search_with_fallback( + "civitai", + lambda base_ctx, cb: search_civitai_for_file( + filename, + api_key=civitai_key or None, + model_type=category, + base_model_context=base_ctx, + session_token=civitai_session_token or None, + candidate_limit=civitai_candidate_limit, + use_trpc_search=civitai_use_trpc_search, + use_api_search=civitai_use_api_search, + use_html_fallback=civitai_use_html_fallback, + progress_callback=cb, + ), + "CivitAI" + ) + if civitai_result: + source_results["civitai"] = civitai_result + source_found = True + + return source_results, source_found + + return run_source_search( + "civitai", + task, + initial_stage="query", + initial_message="Querying CivitAI", + initial_percent=30, + log_start_fields={ + "file": filename, + "cat": category, + "urn": is_urn, + }, + ) + + def search_civarchive_source_task(): + def task(): + source_results = {"civarchive": None} + source_found = False + + if is_urn: + model_id_val = data.get("model_id") + version_id_val = data.get("version_id") + if model_id_val and version_id_val: + self.search_tracker.update( + progress_id, + "civarchive", + "urn", + "Resolving CivArchive version", + 50, + ) + civarchive_result = resolve_civarchive_model_version( + model_id_val, + version_id_val, + query=filename, + ) + log_search_result( + "civarchive/urn", + civarchive_result, + { + "model_id": model_id_val, + "version_id": version_id_val, + }, + ) + if civarchive_result: + source_results["civarchive"] = civarchive_result + source_found = True + else: + log_search_result( + "civarchive/urn", + None, + { + "model_id": model_id_val, + "version_id": version_id_val, + }, + ) + else: + civarchive_result = execute_search_with_fallback( + "civarchive", + lambda base_ctx, cb: search_civarchive_for_file( + filename, + model_type=category, + base_model_context=base_ctx, + limit=civarchive_candidate_limit, + progress_callback=cb, + ), + "CivArchive" + ) + if civarchive_result: + source_results["civarchive"] = civarchive_result + source_found = True + + return source_results, source_found + + def handle_civarchive_error(e): + error_message = f"CivArchive search failed: {e}" + self.logger.warning(error_message) + self.search_tracker.update( + progress_id, + "civarchive", + "error", + error_message, + 100, + status="error", + ) + return { + "civarchive": None, + "source_errors": {"civarchive": error_message}, + }, False + + return run_source_search( + "civarchive", + task, + initial_stage="query", + initial_message="Querying CivArchive", + initial_percent=30, + log_start_fields={ + "file": filename, + "cat": category, + "urn": is_urn, + }, + error_handlers={CivArchiveSearchError: handle_civarchive_error}, + ) + + def search_lora_manager_archive_source_task(): + def task(): + lora_manager_archive_result = execute_search_with_fallback( + "lora_manager_archive", + lambda base_ctx, cb: search_lora_manager_archive_for_file( + filename, + model_type=category, + base_model_context=base_ctx, + progress_callback=cb, + ), + "LoRA archive" + ) + return ( + { + "lora_manager_archive": lora_manager_archive_result + }, + bool(lora_manager_archive_result), + ) + + return run_source_search( + "lora_manager_archive", + task, + initial_stage="query", + initial_message="Searching LoRA Manager archive", + initial_percent=36, + log_start_fields={"file": filename, "cat": category}, + ) + + + search_tasks = [] + if search_local: + search_tasks.append( + asyncio.to_thread(search_local_sources) + ) + if search_huggingface_source: + search_tasks.append( + asyncio.to_thread(search_huggingface_source_task) + ) + if search_civitai_source: + search_tasks.append( + asyncio.to_thread(search_civitai_source_task) + ) + if search_civarchive_source and ( + filename or (is_urn and model_id and version_id) + ): + search_tasks.append( + asyncio.to_thread(search_civarchive_source_task) + ) + if search_lora_manager_archive_source and filename: + search_tasks.append( + asyncio.to_thread( + search_lora_manager_archive_source_task + ) + ) + + if len(search_tasks) > 1: + self.logger.debug( + "Search sources async " + + format_log_fields(count=len(search_tasks)) + ) + self.search_tracker.update( + progress_id, + progress_source, + "running", + "Waiting for search sources", + 18, + ) + raise_if_search_cancelled(progress_source) + + for source_results, source_found in await asyncio.gather( + *search_tasks + ): + raise_if_search_cancelled(progress_source) + for source_key, source_result in source_results.items(): + if source_key == "source_errors": + results["source_errors"].update(source_result or {}) + continue + if source_result: + results[source_key] = source_result + if source_found: + results["found"] = True + + raise_if_search_cancelled(progress_source) + results["local_hash_matches"] = collect_local_hash_matches(results) + raise_if_search_cancelled(progress_source) + + self.logger.info( + f"Search [{','.join(results['searched_sources'])}] done " + + format_log_fields( + found=results["found"], + ) + ) + self.search_tracker.update( + progress_id, + progress_source, + "completed", + "Search complete", + 100, + status="completed", + ) + stamp_search_results(results) + return web.json_response(results) + + except SearchCancelled: + self.search_tracker.update( + progress_id if "progress_id" in locals() else "", + progress_source if "progress_source" in locals() else "", + "cancelled", + "Cancelled", + 100, + status="cancelled", + cancelled=True, + ) + self.logger.info( + "Search cancelled " + + format_log_fields( + source=progress_source if "progress_source" in locals() else "", + progress_id=progress_id if "progress_id" in locals() else "", + ) + ) + return web.json_response( + { + "cancelled": True, + "found": False, + "searched_sources": sorted(normalized_sources) + if "normalized_sources" in locals() + else [], + "source_errors": {}, + } + ) + + except Exception as e: + self.search_tracker.update( + progress_id if "progress_id" in locals() else "", + progress_source if "progress_source" in locals() else "", + "error", + str(e), + 100, + status="error", + ) + self.logger.exception(f"Model Resolver search error: {e}") + return web.json_response({"error": str(e)}, status=500) + + @routes.post("/model_resolver/clear-search-cache") + @json_api_endpoint("Clear search cache") + async def clear_search_cache_route(request): + """Clear backend search caches after token/settings changes.""" + clear_all_search_caches() + reload_popular_databases() + reload_model_list() + invalidate_model_files_cache() + invalidate_local_hash_match_cache() + self.search_result_timestamps.clear() + self.logger.info("Cleared backend search caches") + return web.json_response({"success": True, "cleared": "all"}) + + async def _check_credential_helper(request, payload_key, check_func, log_name): + try: + data = await request.json() + val = data.get(payload_key, "") + result = await asyncio.to_thread(check_func, val) + return web.json_response(result) + except Exception as e: + self.logger.exception(f"{log_name} check error: {e}") + return web.json_response({"error": str(e)}, status=500) + + @routes.post("/model_resolver/civitai/session-token/check") + async def civitai_session_token_check_route(request): + """Check whether a CivitAI browser session token is valid.""" + return await _check_credential_helper( + request, "civitai_session_token", check_civitai_session_token, "CivitAI session token" + ) + + @routes.post("/model_resolver/civitai/api-key/check") + async def civitai_api_key_check_route(request): + """Check whether a CivitAI API key is valid.""" + return await _check_credential_helper( + request, "civitai_key", check_civitai_api_key, "CivitAI API key" + ) + + @routes.post("/model_resolver/huggingface/token/check") + async def huggingface_token_check_route(request): + """Check whether a HuggingFace token is valid.""" + return await _check_credential_helper( + request, "hf_token", check_huggingface_token, "HuggingFace token" + ) + + @routes.post("/model_resolver/brave/api-key/check") + async def brave_api_key_check_route(request): + """Check whether a Brave Search API key is valid.""" + return await _check_credential_helper( + request, "brave_search_api_key", check_brave_search_api_key, "Brave Search API key" + ) + + @routes.get("/model_resolver/huggingface/author-index/status") + @json_api_endpoint("HuggingFace author index status") + async def huggingface_author_index_status_route(request): + """Return local HuggingFace author fallback index status.""" + return web.json_response( + get_known_author_fallback_indexes_status() + ) + + @routes.post("/model_resolver/huggingface/author-index/refresh") + @json_api_endpoint("HuggingFace author index refresh") + async def huggingface_author_index_refresh_route(request): + """Refresh HuggingFace author fallback index.""" + data = await request.json() + hf_token = data.get("hf_token", "") + result = await asyncio.to_thread( + refresh_known_author_fallback_indexes, hf_token or None + ) + clear_huggingface_search_cache() + return web.json_response(result) + + @routes.get("/model_resolver/model-list/status") + @json_api_endpoint("Model list status") + async def model_list_status_route(request): + """Return local model-list status and optionally compare with GitHub.""" + check_remote = ( + str(request.query.get("check_remote", "")).lower() + in {"1", "true", "yes"} + ) + return web.json_response( + get_model_list_update_status(check_remote=check_remote) + ) + + @routes.post("/model_resolver/model-list/update") + @json_api_endpoint("Model list update") + async def model_list_update_route(request): + """Download latest ComfyUI-Manager model-list.json.""" + result = await asyncio.to_thread(update_model_list_from_remote) + clear_all_search_caches() + self.search_result_timestamps.clear() + return web.json_response(result) diff --git a/core/routes/workflow.py b/core/routes/workflow.py new file mode 100644 index 0000000..a45eb25 --- /dev/null +++ b/core/routes/workflow.py @@ -0,0 +1,1392 @@ +"""Workflow and loaded-model route registration.""" + +def register_workflow_routes(context): + FileManagerError = context.get('FileManagerError') + FileManagerUnavailableError = context.get('FileManagerUnavailableError') + MODEL_RESOLVER_METADATA_SCHEMA = context.get('MODEL_RESOLVER_METADATA_SCHEMA') + MODEL_RESOLVER_METADATA_SCHEMA_VERSION = context.get('MODEL_RESOLVER_METADATA_SCHEMA_VERSION') + UnsupportedFileManagerPlatformError = context.get('UnsupportedFileManagerPlatformError') + analyze_and_find_matches = context.get('analyze_and_find_matches') + apply_resolution = context.get('apply_resolution') + asyncio = context.get('asyncio') + cancel_progress_response = context.get('cancel_progress_response') + download_available = context.get('download_available') + fetch_remote_file_size_cached = context.get('fetch_remote_file_size_cached') + get_existing_model_preview_path = context.get('get_existing_model_preview_path') + get_filename_from_path = context.get('get_filename_from_path') + get_local_model_hash_metadata = context.get('get_local_model_hash_metadata') + get_popular_model_url = context.get('get_popular_model_url') + get_progress_response = context.get('get_progress_response') + get_safe_model_resolver_sidecar_path = context.get('get_safe_model_resolver_sidecar_path') + get_workflow_model_inventory = context.get('get_workflow_model_inventory') + invalidate_local_hash_match_cache = context.get('invalidate_local_hash_match_cache') + is_path_in_configured_model_roots = context.get('is_path_in_configured_model_roots') + json_api_endpoint = context.get('json_api_endpoint') + load_resolver_settings = context.get('load_resolver_settings') + normalize_file_manager_path = context.get('normalize_file_manager_path') + normalize_sha256 = context.get('normalize_sha256') + open_in_file_manager = context.get('open_in_file_manager') + os = context.get('os') + read_json_safe = context.get('read_json_safe') + resolver_bool_setting = context.get('resolver_bool_setting') + routes = context.get('routes') + run_in_background_thread = context.get('run_in_background_thread') + search_local_matches = context.get('search_local_matches') + search_local_matches_by_hash = context.get('search_local_matches_by_hash') + search_model_list = context.get('search_model_list') + self = context.get('self') + should_skip_existing_custom_node_reference = context.get('should_skip_existing_custom_node_reference') + time = context.get('time') + to_bool = context.get('to_bool') + to_int = context.get('to_int') + web = context.get('web') + write_json_atomic = context.get('write_json_atomic') + + @routes.post("/model_resolver/analyze") + async def analyze_workflow(request): + """Analyze workflow and return missing models with matches.""" + try: + data = await request.json() + workflow_json = data.get("workflow") + analysis_id = str(data.get("analysis_id") or "").strip() + force_rescan = to_bool(data.get("force_rescan"), False) + if force_rescan: + invalidate_local_hash_match_cache() + + if workflow_json is None: + return web.json_response( + {"error": "Workflow JSON is required"}, status=400 + ) + if not isinstance(workflow_json, dict): + return web.json_response( + {"error": "Workflow JSON must be an object"}, status=400 + ) + + if analysis_id: + self._update_analysis_progress( + analysis_id, + { + "status": "starting", + "stage": "starting", + "message": "Starting analysis...", + "current": 0, + "total": 0, + }, + ) + + def update_analysis_progress(payload): + self._update_analysis_progress(analysis_id, payload) + + # Analyze and find matches + result = await asyncio.to_thread( + analyze_and_find_matches, + workflow_json, + 0.0, + 10, + update_analysis_progress if analysis_id else None, + force_rescan=force_rescan, + ) + + missing_models = result.get("missing_models", []) + filtered_missing = [] + for missing in missing_models: + name = missing.get("name") or missing.get("original_path", "") + if should_skip_existing_custom_node_reference(missing): + self.logger.info( + "Filtered existing custom-node model " + f"reference: {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 = get_filename_from_path( + missing.get("original_path", "") + ) + + # 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 using cached remote helper + file_size = fetch_remote_file_size_cached(workflow_url, timeout=5) + + missing["download_source"] = { + "source": source, + "url": workflow_url, + "model_url": missing.get( + "workflow_model_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: + popular_model_list_result = search_model_list( + filename, exact_only=True + ) + missing["download_source"] = { + "source": "popular", + "url": popular_info.get("url"), + "filename": filename, + "type": popular_info.get("type"), + "directory": popular_info.get("directory"), + "size": ( + popular_model_list_result.get("size") + if popular_model_list_result + else None + ) + or popular_info.get("size"), + "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_resolver/search endpoint + # when user clicks "Search Online" button, not automatically + + if analysis_id: + self.analysis_progress.update( + analysis_id, + status="completed", + stage="completed", + message="Analysis complete", + current=result.get("total_missing", 0), + total=result.get("total_missing", 0), + ) + + return web.json_response(result) + except Exception as e: + if "analysis_id" in locals() and analysis_id: + self.analysis_progress.update( + analysis_id, + status="error", + stage="error", + message=str(e), + current=0, + total=0, + ) + self.logger.error(f"Model Resolver analyze error: {e}", exc_info=True) + return web.json_response({"error": str(e)}, status=500) + + @routes.get("/model_resolver/analyze-progress/{analysis_id}") + @json_api_endpoint("analyze-progress") + async def get_analyze_progress(request): + """Get workflow analysis progress.""" + return get_progress_response( + self.analysis_progress, + request, + param_name="analysis_id", + not_found_payload={ + "status": "unknown", + "stage": "unknown", + "message": "No analysis progress available", + "current": 0, + "total": 0, + } + ) + + @routes.post("/model_resolver/resolve") + @json_api_endpoint("resolve", return_success_on_error=True) + async def resolve_models(request): + """Apply model resolution and return updated workflow.""" + 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} + ) + + @routes.post("/model_resolver/local-matches") + @json_api_endpoint("local-matches") + async def local_matches(request): + """Search local model files by filename/path.""" + data = await request.json() + filename = data.get("filename", "") + category = data.get("category", "") + force_rescan = to_bool(data.get("force_rescan"), False) + if force_rescan: + invalidate_local_hash_match_cache() + + if not filename: + return web.json_response( + {"error": "filename is required"}, status=400 + ) + + matches = search_local_matches( + filename, + category=category or None, + similarity_threshold=0.0, + max_matches_per_model=10, + force_rescan=force_rescan, + ) + return web.json_response({"matches": matches}) + + @routes.post("/model_resolver/local-model-hashes") + @json_api_endpoint("local-model-hashes") + async def local_model_hashes(request): + """Return SHA256 hashes already stored in local sidecar metadata.""" + data = await request.json() + model = data.get("model") if isinstance(data.get("model"), dict) else {} + path = ( + data.get("path") + or data.get("file_path") + or data.get("resolved_path") + or model.get("path") + or model.get("resolved_path") + or "" + ) + + if not path: + return web.json_response( + {"error": "path is required"}, status=400 + ) + + normalized_path = os.path.realpath( + os.path.abspath(os.path.normpath(str(path))) + ) + if not is_path_in_configured_model_roots(normalized_path): + return web.json_response( + {"error": "path is outside configured model directories"}, + status=403, + ) + + return web.json_response( + get_local_model_hash_metadata(normalized_path, model=model) + ) + + @routes.get("/model_resolver/model-preview") + async def get_model_preview(request): + """Serve an adjacent model preview from a configured model directory.""" + model_path = str(request.query.get("path") or "").strip() + if not model_path: + return web.Response(text="path is required", status=400) + + try: + normalized_path = os.path.realpath( + os.path.abspath(os.path.normpath(model_path)) + ) + except (OSError, TypeError, ValueError): + return web.Response(text="invalid model path", status=400) + if not os.path.isfile(normalized_path): + return web.Response(text="model file does not exist", status=404) + if not is_path_in_configured_model_roots(normalized_path): + return web.Response( + text="path is outside configured model directories", + status=403, + ) + + preview_path = get_existing_model_preview_path(normalized_path) + if not preview_path: + return web.Response(text="preview not found", status=404) + if not is_path_in_configured_model_roots(preview_path): + return web.Response( + text="preview is outside configured model directories", + status=403, + ) + + return web.FileResponse( + preview_path, + headers={"Cache-Control": "private, no-cache"}, + ) + + @routes.post("/model_resolver/workflow-model-hashes") + @json_api_endpoint("workflow-model-hashes") + async def workflow_model_hashes(request): + """Return hash metadata for existing local models used by a workflow.""" + data = await request.json() + workflow_json = data.get("workflow") + if not isinstance(workflow_json, dict): + return web.json_response( + {"error": "Workflow JSON must be an object"}, status=400 + ) + + settings = await asyncio.to_thread(load_resolver_settings) + if not resolver_bool_setting( + settings.get("workflow_hash_metadata_enabled"), True + ): + return web.json_response( + { + "success": True, + "enabled": False, + "models": [], + "by_node": {}, + "by_path": {}, + "count": 0, + } + ) + + inventory = await asyncio.to_thread( + get_workflow_model_inventory, + workflow_json, + ) + refs = inventory["model_refs"] + + by_node = {} + by_path = {} + models = [] + seen = set() + + for ref in refs: + if not isinstance(ref, dict) or not ref.get("exists"): + continue + full_path = str(ref.get("full_path") or "").strip() + if not full_path: + continue + model_info = { + "path": full_path, + "filename": get_filename_from_path(full_path), + "relative_path": ref.get("original_path") or "", + "category": ref.get("category") or "", + } + metadata = get_local_model_hash_metadata( + full_path, model=model_info + ) + sha256 = normalize_sha256(metadata.get("sha256")) + if not sha256: + continue + + entry = { + "node_id": ref.get("node_id"), + "node_type": ref.get("node_type") or "", + "widget_index": ref.get("widget_index"), + "widget_name": ref.get("widget_name") or "", + "path": ref.get("original_path") or "", + "filename": get_filename_from_path( + ref.get("original_path") or full_path + ), + "category": ref.get("category") or "", + "sha256": sha256, + "size": metadata.get("size") or 0, + } + entry_key = ( + str(entry.get("node_id")), + str(entry.get("widget_index")), + entry.get("path"), + sha256, + ) + if entry_key in seen: + continue + seen.add(entry_key) + models.append(entry) + + path_key = str(entry.get("path") or entry.get("filename") or "") + if path_key: + by_path[path_key] = entry + by_path[get_filename_from_path(path_key)] = entry + node_key = f"{entry.get('node_id')}:{entry.get('widget_index')}" + by_node[node_key] = entry + + return web.json_response( + { + "success": True, + "enabled": True, + "models": models, + "by_node": by_node, + "by_path": by_path, + "count": len(models), + } + ) + + @routes.post("/model_resolver/local-matches-by-hash") + @json_api_endpoint("local-matches-by-hash") + async def local_matches_by_hash(request): + """Search local model metadata sidecars for a remote SHA256.""" + data = await request.json() + sha256 = normalize_sha256( + data.get("sha256") + or data.get("hash") + or data.get("SHA256") + or "" + ) + if not sha256: + return web.json_response( + {"error": "sha256 is required"}, status=400 + ) + + category = data.get("category", "") + source = str( + data.get("source") + or data.get("hash_lookup_source") + or "download_source" + ).strip().lower().replace("-", "_") + filename = ( + data.get("filename") + or data.get("path") + or data.get("model_name") + or "" + ) + max_matches = to_int(data.get("max_matches"), 20) + force_rescan = to_bool(data.get("force_rescan"), False) + + matches = search_local_matches_by_hash( + sha256, + category=category or None, + max_matches=max_matches, + force_rescan=force_rescan, + ) + enriched_matches = [ + { + **match, + "hash_lookup_source": source or "download_source", + "hash_lookup_filename": filename, + "hash_lookup_sha256": sha256, + } + for match in matches + ] + return web.json_response( + { + "sha256": sha256, + "local_hash_matches": enriched_matches, + "matches": enriched_matches, + } + ) + + @routes.post("/model_resolver/open-containing-folder") + @json_api_endpoint("open-containing-folder") + async def open_containing_folder(request): + """Reveal a file or open a directory in the host file manager.""" + try: + data = await request.json() + except Exception as exc: + return web.json_response( + {"success": False, "error": f"invalid JSON body: {exc}"}, + status=400, + ) + if not isinstance(data, dict): + return web.json_response( + {"success": False, "error": "JSON body must be an object"}, + status=400, + ) + target_path = data.get("path", "") + + try: + normalized_path = normalize_file_manager_path(target_path) + except ValueError as exc: + return web.json_response( + {"success": False, "error": str(exc)}, status=400 + ) + + if not os.path.exists(normalized_path): + return web.json_response( + {"success": False, "error": "path does not exist"}, status=404 + ) + if not is_path_in_configured_model_roots(normalized_path): + return web.json_response( + { + "success": False, + "error": "path is outside configured model directories", + }, + status=403, + ) + + try: + result = await asyncio.to_thread( + open_in_file_manager, normalized_path + ) + except FileNotFoundError: + return web.json_response( + {"success": False, "error": "path no longer exists"}, + status=404, + ) + except UnsupportedFileManagerPlatformError as exc: + return web.json_response( + {"success": False, "error": str(exc)}, status=501 + ) + except FileManagerUnavailableError as exc: + return web.json_response( + { + "success": False, + "error": ( + f"{exc} The ComfyUI host may be running without " + "a graphical desktop session." + ), + }, + status=503, + ) + except FileManagerError as exc: + return web.json_response( + {"success": False, "error": str(exc)}, status=500 + ) + + return web.json_response({"success": True, **result}) + + # cleanup_hash_progress, update_hash_progress, is_hash_progress_cancelled, and mark_hash_progress_cancelled removed + + from ..path_utils import HashCalculationCancelled + + def resolve_hash_file_request(data): + import os as _os + + file_path = ( + data.get("file_path") + or data.get("resolved_path") + or data.get("path") + or "" + ) + if not file_path: + return "", "file_path is required" + + normalized_path = _os.path.realpath( + _os.path.abspath(_os.path.normpath(file_path)) + ) + if not _os.path.exists(normalized_path) or not _os.path.isfile(normalized_path): + return "", "file does not exist" + if not is_path_in_configured_model_roots(normalized_path): + return "", "file is outside configured model directories" + return normalized_path, "" + + def write_calculated_hash_metadata( + normalized_path, + sha256, + sha256_source="file", + ): + import os as _os + + resolved_metadata_path = get_safe_model_resolver_sidecar_path( + normalized_path + ) + + metadata_updated = False + try: + metadata = read_json_safe(resolved_metadata_path, {}) + if not isinstance(metadata, dict): + metadata = {} + + filename = get_filename_from_path(normalized_path) + stem, _ext = _os.path.splitext(filename) + hashes = metadata.get("hashes") + if not isinstance(hashes, dict): + hashes = {} + hashes["SHA256"] = sha256 + + metadata["schema"] = MODEL_RESOLVER_METADATA_SCHEMA + metadata["schema_version"] = ( + MODEL_RESOLVER_METADATA_SCHEMA_VERSION + ) + metadata["managed_by"] = MODEL_RESOLVER_METADATA_SCHEMA + metadata["sha256"] = sha256 + metadata["hashes"] = hashes + metadata["hash_status"] = "completed" + metadata["sha256_source"] = sha256_source or "file" + metadata["last_checked_at"] = time.time() + metadata.setdefault("file_name", stem) + metadata.setdefault("model_name", stem) + metadata.setdefault("file_path", normalized_path.replace("\\", "/")) + try: + metadata.setdefault("size", _os.path.getsize(normalized_path)) + except Exception: + pass + + write_json_atomic(resolved_metadata_path, metadata, indent=2) + metadata_updated = True + self.logger.info( + f"Stored SHA256 and updated metadata: {resolved_metadata_path}" + ) + except Exception as metadata_error: + self.logger.warning( + f"Could not update metadata with calculated SHA256 for {normalized_path}: {metadata_error}" + ) + + return resolved_metadata_path, metadata_updated + + def calculate_sha256_with_progress(normalized_path, progress_id=""): + import os as _os + + from ..path_utils import calculate_file_sha256 as _calculate_file_sha256_core + + total_bytes = max(0, _os.path.getsize(normalized_path)) + + self.hash_tracker.update( + progress_id, + status="running", + stage="header", + message="Checking safetensors header...", + percent=0, + bytes_read=0, + total_bytes=total_bytes, + ) + + if self.hash_tracker.is_cancelled(progress_id): + raise HashCalculationCancelled() + + hash_source = ["file"] + stage_transitioned = [False] + _bytes_read = [0] + _last_update = [0.0] + + def on_progress(bytes_read, total_bytes): + _bytes_read[0] = bytes_read + if not stage_transitioned[0]: + stage_transitioned[0] = True + self.logger.info( + "No SHA256 found in safetensors header; calculating full " + f"file SHA256: {normalized_path}" + ) + self.hash_tracker.update( + progress_id, + status="running", + stage="hashing", + message="Calculating SHA256...", + percent=0, + bytes_read=0, + total_bytes=total_bytes, + ) + + now = time.time() + if now - _last_update[0] >= 0.15 or bytes_read >= total_bytes: + percent = 98 if total_bytes <= 0 else min( + 98, + (bytes_read / total_bytes) * 98, + ) + self.hash_tracker.update( + progress_id, + status="running", + stage="hashing", + message="Calculating SHA256...", + percent=percent, + bytes_read=bytes_read, + total_bytes=total_bytes, + ) + _last_update[0] = now + + def is_cancelled(): + if self.hash_tracker.is_cancelled(progress_id): + percent = 0 if total_bytes <= 0 else min( + 98, + (_bytes_read[0] / total_bytes) * 98, + ) + self.hash_tracker.update( + progress_id, + status="cancelled", + stage="cancelled", + message="Hash calculation cancelled", + percent=percent, + bytes_read=_bytes_read[0], + total_bytes=total_bytes, + ) + return True + return False + + def on_hash_source(source_name): + hash_source[0] = source_name + if source_name == "safetensors_header": + self.logger.info( + f"Found SHA256 in safetensors header: {normalized_path}" + ) + self.hash_tracker.update( + progress_id, + status="running", + stage="header", + message="Found SHA256 in safetensors header", + percent=98, + bytes_read=0, + total_bytes=total_bytes, + sha256_source="safetensors_header", + ) + + sha256 = _calculate_file_sha256_core( + normalized_path, + chunk_size=1024 * 1024 * 4, + on_progress=on_progress, + is_cancelled=is_cancelled, + use_safetensors_header=True, + on_hash_source=on_hash_source, + ) + + if is_cancelled(): + raise HashCalculationCancelled() + + return sha256, hash_source[0] + + @routes.post("/model_resolver/calculate-file-hash") + @json_api_endpoint("calculate-file-hash") + async def calculate_file_hash_route(request): + """Calculate SHA256 for a local model and persist it to sidecar metadata.""" + data = await request.json() + + normalized_path, error = resolve_hash_file_request(data) + if error == "file_path is required": + return web.json_response( + {"error": "file_path is required"}, status=400 + ) + if error: + status = 403 if error == "file is outside configured model directories" else 404 + return web.json_response( + {"error": error}, status=status + ) + + sha256, sha256_source = calculate_sha256_with_progress(normalized_path) + if not sha256: + return web.json_response( + {"error": "could not calculate hash"}, status=500 + ) + resolved_metadata_path, metadata_updated = write_calculated_hash_metadata( + normalized_path, + sha256, + sha256_source, + ) + + return web.json_response( + { + "success": True, + "sha256": sha256, + "hash": sha256, + "sha256_source": sha256_source, + "file_path": normalized_path, + "metadata_path": resolved_metadata_path, + "metadata_updated": metadata_updated, + } + ) + + @routes.post("/model_resolver/calculate-file-hash/start") + @json_api_endpoint("calculate-file-hash-start") + async def calculate_file_hash_start_route(request): + """Start SHA256 calculation in a background thread.""" + import uuid + + data = await request.json() + normalized_path, error = resolve_hash_file_request(data) + if error == "file_path is required": + return web.json_response( + {"error": "file_path is required"}, status=400 + ) + if error: + status = 403 if error == "file is outside configured model directories" else 404 + return web.json_response( + {"error": error}, status=status + ) + + self.hash_tracker.cleanup() + progress_id = f"hash_{uuid.uuid4().hex}" + self.hash_tracker.update( + progress_id, + status="queued", + stage="queued", + message="Preparing hash calculation...", + percent=0, + file_path=normalized_path, + ) + + def run_hash_task(): + sha256, sha256_source = calculate_sha256_with_progress( + normalized_path, + progress_id=progress_id, + ) + if self.hash_tracker.is_cancelled(progress_id): + raise HashCalculationCancelled() + self.hash_tracker.update( + progress_id, + status="running", + stage="metadata", + message="Saving metadata...", + percent=99, + ) + resolved_metadata_path, metadata_updated = write_calculated_hash_metadata( + normalized_path, + sha256, + sha256_source, + ) + return sha256, sha256_source, resolved_metadata_path, metadata_updated + + def on_success(res): + sha256, sha256_source, resolved_metadata_path, metadata_updated = res + self.hash_tracker.update( + progress_id, + status="done", + stage="done", + message="Hash calculated", + percent=100, + sha256=sha256, + hash=sha256, + sha256_source=sha256_source, + file_path=normalized_path, + metadata_path=resolved_metadata_path, + metadata_updated=metadata_updated, + ) + + def on_cancel(res=None): + self.hash_tracker.update( + progress_id, + status="cancelled", + stage="cancelled", + message="Hash calculation cancelled", + ) + + run_in_background_thread( + self.hash_tracker, + progress_id, + run_hash_task, + on_success, + on_cancel, + error_log_msg=f"Hash calculation failed for {normalized_path}", + ) + return web.json_response( + { + "success": True, + "progress_id": progress_id, + } + ) + + @routes.get("/model_resolver/calculate-file-hash/progress/{progress_id}") + @json_api_endpoint("calculate-file-hash-progress") + async def calculate_file_hash_progress_route(request): + """Return progress for a background SHA256 calculation.""" + return get_progress_response( + self.hash_tracker, + request, + not_found_status=404 + ) + + @routes.post("/model_resolver/calculate-file-hash/cancel/{progress_id}") + @json_api_endpoint("calculate-file-hash-cancel") + async def calculate_file_hash_cancel_route(request): + """Cancel a background SHA256 calculation.""" + return cancel_progress_response( + self.hash_tracker, + request, + cancel_message="Stopping hash calculation..." + ) + +def register_loaded_model_routes(context): + adapt_custom_node_loaded_model = context.get('adapt_custom_node_loaded_model') + asyncio = context.get('asyncio') + get_filename_from_path = context.get('get_filename_from_path') + get_progress_response = context.get('get_progress_response') + get_workflow_model_inventory = context.get('get_workflow_model_inventory') + json_api_endpoint = context.get('json_api_endpoint') + routes = context.get('routes') + self = context.get('self') + web = context.get('web') + + @routes.post("/model_resolver/loaded") + @json_api_endpoint("get_loaded_models") + async def get_loaded_models(request): + """Get all currently loaded models in the workflow.""" + data = await request.json() + workflow_json = data.get("workflow") + loaded_id = str( + data.get("loaded_id") or data.get("progress_id") or "" + ).strip() + + if not workflow_json: + return web.json_response( + {"error": "Workflow JSON is required"}, status=400 + ) + if not isinstance(workflow_json, dict): + return web.json_response( + {"error": "Workflow JSON must be an object"}, status=400 + ) + + def update_loaded_progress(*args, **kwargs): + self._update_loaded_progress(loaded_id, *args, **kwargs) + + def get_workflow_node_count(): + node_count = 0 + nodes = workflow_json.get("nodes", []) + if isinstance(nodes, list): + node_count += len(nodes) + + definitions = workflow_json.get("definitions", {}) + subgraphs = ( + definitions.get("subgraphs", []) + if isinstance(definitions, dict) + else [] + ) + if isinstance(subgraphs, list): + for subgraph in subgraphs: + if not isinstance(subgraph, dict): + continue + subgraph_nodes = subgraph.get("nodes", []) + if isinstance(subgraph_nodes, list): + node_count += len(subgraph_nodes) + return node_count + + def interpolate_percent(start, end, current, total): + if not total: + return start + try: + ratio = max(0.0, min(1.0, float(current) / float(total))) + except (TypeError, ValueError, ZeroDivisionError): + ratio = 0.0 + return start + ((end - start) * ratio) + + def build_loaded_models_response(): + from ..workflow_analyzer import ( + URN_TYPE_MAP, + get_lora_model_strength, + ) + + update_loaded_progress( + "scanning", + "Scanning local model index...", + percent=5, + ) + + workflow_node_count = get_workflow_node_count() + + def update_workflow_analysis_progress(payload): + self._update_workflow_analysis_progress( + loaded_id, + workflow_node_count, + interpolate_percent, + payload, + start_percent=5, + end_percent=45, + ) + + inventory = get_workflow_model_inventory( + workflow_json, + progress_callback=update_workflow_analysis_progress, + ) + available_models = inventory["available_models"] + all_model_refs = inventory["model_refs"] + + # Create lookup for full paths by filename (with and without extension) + path_by_filename = {} + total_local_models = len(available_models) + update_loaded_progress( + "indexing", + f"Indexing {total_local_models} local models...", + percent=45, + current=0, + total=total_local_models, + ) + for index, model_info in enumerate(available_models, start=1): + rel_path = model_info.get("relative_path", "") + if rel_path: + filename = get_filename_from_path(rel_path) + path_by_filename[filename] = model_info.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] = model_info.get("path") + # Add the full relative path as key too + path_by_filename[rel_path] = model_info.get("path") + + if ( + total_local_models + and (index == total_local_models or index % 250 == 0) + ): + update_loaded_progress( + "indexing", + f"Indexing local model {index} of {total_local_models}", + percent=interpolate_percent( + 45, 55, index, total_local_models + ), + current=index, + total=total_local_models, + ) + + # Also use folder_paths.get_full_path() to get paths + try: + import folder_paths + except Exception: + folder_paths = None + + folder_categories = [ + "loras", + "checkpoints", + "vae", + "controlnet", + "upscale_models", + ] + if folder_paths: + for index, category_name in enumerate(folder_categories, start=1): + update_loaded_progress( + "indexing", + f"Reading {category_name} model list...", + percent=interpolate_percent( + 55, 65, index - 1, len(folder_categories) + ), + current=index - 1, + total=len(folder_categories), + ) + try: + filenames = folder_paths.get_filename_list(category_name) + for filename in filenames: + full_path = folder_paths.get_full_path( + category_name, filename + ) + if ( + full_path + and full_path not in path_by_filename.values() + ): + path_by_filename[filename] = full_path + filename_no_ext = ( + filename.rsplit(".", 1)[0] + if "." in filename + else filename + ) + if filename_no_ext not in path_by_filename: + path_by_filename[filename_no_ext] = full_path + except Exception: + pass + + # Also extract from node.properties.models + workflow_nodes = workflow_json.get("nodes", []) + nodes = list(workflow_nodes) if isinstance(workflow_nodes, list) else [] + node_scopes = { + id(node): { + "subgraph_id": "", + "subgraph_name": "", + "is_top_level": True, + } + for node in nodes + if isinstance(node, dict) + } + definitions = workflow_json.get("definitions", {}) + subgraphs = ( + definitions.get("subgraphs", []) + if isinstance(definitions, dict) + else [] + ) + for subgraph in subgraphs: + if isinstance(subgraph, dict): + subgraph_nodes = subgraph.get("nodes", []) + if isinstance(subgraph_nodes, list): + nodes.extend(subgraph_nodes) + for node in subgraph_nodes: + if not isinstance(node, dict): + continue + node_scopes[id(node)] = { + "subgraph_id": subgraph.get("id") or "", + "subgraph_name": ( + subgraph.get("name") + or subgraph.get("id") + or "" + ), + "is_top_level": False, + } + + def get_node_scope(node): + return node_scopes.get( + id(node), + { + "subgraph_id": "", + "subgraph_name": "", + "is_top_level": True, + }, + ) + + def node_matches_ref(node, ref): + if not isinstance(node, dict): + return False + if str(node.get("id")) != str(ref.get("node_id")): + return False + scope = get_node_scope(node) + if scope["is_top_level"] != ( + ref.get("is_top_level") is not False + ): + return False + if not scope["is_top_level"]: + return str(scope["subgraph_id"]) == str( + ref.get("subgraph_id") or "" + ) + return True + + # Collect all loaded models with their values + loaded_models = [] + total_refs = len(all_model_refs) + update_loaded_progress( + "building", + "Building loaded models list...", + percent=65, + current=0, + total=total_refs, + ) + + # Process each model reference from analyze_workflow_models + for index, ref in enumerate(all_model_refs, start=1): + if total_refs and (index == total_refs or index % 50 == 0): + update_loaded_progress( + "building", + f"Building loaded model {index} of {total_refs}", + percent=interpolate_percent(65, 94, index, total_refs), + current=index, + total=total_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 = get_filename_from_path(original_path) + strength = None + + if ( + ref.get("strength") is None + and category in {"lora", "loras"} + ): + for node in nodes: + if node_matches_ref(node, ref): + strength = get_lora_model_strength( + node, + widget_index, + ) + break + + if ref.get("strength") is not None: + strength = ref.get("strength") + + model_name, strength = adapt_custom_node_loaded_model( + ref, + model_name, + strength, + ) + + # 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, + "node_title": ref.get("node_title", ""), + "subgraph_id": ref.get("subgraph_id") or "", + "subgraph_name": ref.get("subgraph_name") or "", + "is_top_level": ref.get("is_top_level") is not False, + "locate_node_id": ref.get("locate_node_id"), + "locate_node_type": ref.get("locate_node_type", ""), + "locate_node_title": ref.get("locate_node_title", ""), + "locate_subgraph_id": ( + ref.get("locate_subgraph_id") or "" + ), + "locate_subgraph_name": ( + ref.get("locate_subgraph_name") or "" + ), + "locate_is_top_level": ( + ref.get("locate_is_top_level") is not False + ), + "locate_via_promoted_widget": ref.get( + "locate_via_promoted_widget", False + ), + "exists": exists, + "strength": strength, + "original_path": original_path, + "is_urn": ref.get("is_urn", False), + "custom_node_adapter": ref.get( + "custom_node_adapter" + ), + "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 + total_nodes = len(nodes) + for node_index, node in enumerate(nodes, start=1): + if total_nodes and ( + node_index == total_nodes or node_index % 50 == 0 + ): + update_loaded_progress( + "embedded", + "Checking embedded model metadata...", + percent=interpolate_percent( + 94, 98, node_index, total_nodes + ), + current=node_index, + total=total_nodes, + ) + + node_type = node.get("type", "") + properties = node.get("properties", {}) + if not isinstance(properties, dict): + properties = {} + models_list = properties.get("models", []) + if not isinstance(models_list, list): + models_list = [] + + for model_info in models_list: + if isinstance(model_info, dict): + name = model_info.get("name", "") + directory = model_info.get("directory", "") + node_scope = get_node_scope(node) + + if name: + # Check if this model is already in loaded_models + existing = next( + ( + model + for model in loaded_models + if model.get("original_path") == name + ), + None, + ) + if not existing: + loaded_models.append( + { + "name": get_filename_from_path(name), + "category": directory or "checkpoints", + "node_id": node.get("id"), + "widget_index": None, + "node_type": node_type, + "node_title": node.get("title", ""), + "subgraph_id": node_scope["subgraph_id"], + "subgraph_name": node_scope[ + "subgraph_name" + ], + "is_top_level": node_scope[ + "is_top_level" + ], + "exists": True, # Embedded models are loaded + "strength": None, + "original_path": name, + "is_urn": False, + } + ) + + update_loaded_progress( + "finalizing", + "Loaded models ready", + percent=99, + current=len(loaded_models), + total=len(loaded_models), + ) + return { + "loaded_models": loaded_models, + "total": len(loaded_models), + } + + update_loaded_progress( + "starting", + "Preparing loaded model scan...", + percent=0, + status="starting", + current=0, + total=0, + ) + + try: + result = await asyncio.to_thread(build_loaded_models_response) + except Exception as e: + update_loaded_progress( + "error", + str(e), + percent=100, + status="error", + current=0, + total=0, + ) + self.logger.error( + f"Model Resolver loaded models error: {e}", exc_info=True + ) + return web.json_response({"error": str(e)}, status=500) + + update_loaded_progress( + "completed", + "Loaded models ready", + percent=100, + status="completed", + current=result.get("total", 0), + total=result.get("total", 0), + ) + return web.json_response(result) + + @routes.get("/model_resolver/loaded-progress/{loaded_id}") + @json_api_endpoint("loaded-progress") + async def get_loaded_models_progress(request): + """Get loaded models inspection progress.""" + return get_progress_response( + self.loaded_progress, + request, + param_name="loaded_id", + not_found_payload={ + "status": "unknown", + "stage": "unknown", + "message": "No loaded models progress available", + "percent": 0, + "current": 0, + "total": 0, + } + ) + + # ==================== MODEL METADATA LOOKUP ROUTE ==================== diff --git a/tests/test_refactoring_targets.py b/tests/test_refactoring_targets.py index aa643d4..beb0e83 100644 --- a/tests/test_refactoring_targets.py +++ b/tests/test_refactoring_targets.py @@ -79,6 +79,77 @@ def _find_bound_extension(handler): class TestRefactoringTargets(unittest.IsolatedAsyncioTestCase): + ROUTE_GROUPS = { + "workflow": { + ("POST", "/model_resolver/analyze"), + ("GET", "/model_resolver/analyze-progress/{analysis_id}"), + ("POST", "/model_resolver/resolve"), + ("POST", "/model_resolver/local-matches"), + ("POST", "/model_resolver/local-model-hashes"), + ("GET", "/model_resolver/model-preview"), + ("POST", "/model_resolver/workflow-model-hashes"), + ("POST", "/model_resolver/local-matches-by-hash"), + ("POST", "/model_resolver/open-containing-folder"), + ("POST", "/model_resolver/calculate-file-hash"), + ("POST", "/model_resolver/calculate-file-hash/start"), + ("GET", "/model_resolver/calculate-file-hash/progress/{progress_id}"), + ("POST", "/model_resolver/calculate-file-hash/cancel/{progress_id}"), + }, + "metadata": { + ("GET", "/model_resolver/models"), + ("POST", "/model_resolver/metadata-size-audit"), + ("GET", "/model_resolver/metadata-build/capabilities"), + ("POST", "/model_resolver/metadata-build/start"), + ("GET", "/model_resolver/metadata-build/progress/{progress_id}"), + ("POST", "/model_resolver/metadata-build/cancel/{progress_id}"), + }, + "loaded": { + ("POST", "/model_resolver/loaded"), + ("GET", "/model_resolver/loaded-progress/{loaded_id}"), + }, + "model_info": { + ("POST", "/model_resolver/civitai-search"), + ("POST", "/model_resolver/custom-url"), + ("POST", "/model_resolver/model-details"), + }, + "search": { + ("GET", "/model_resolver/search-progress/{progress_id}"), + ("POST", "/model_resolver/search-cancel/{progress_id}"), + ("POST", "/model_resolver/search"), + ("POST", "/model_resolver/clear-search-cache"), + ("POST", "/model_resolver/civitai/session-token/check"), + ("POST", "/model_resolver/civitai/api-key/check"), + ("POST", "/model_resolver/huggingface/token/check"), + ("POST", "/model_resolver/brave/api-key/check"), + ("GET", "/model_resolver/huggingface/author-index/status"), + ("POST", "/model_resolver/huggingface/author-index/refresh"), + ("GET", "/model_resolver/model-list/status"), + ("POST", "/model_resolver/model-list/update"), + }, + "downloads": { + ("POST", "/model_resolver/download"), + ("GET", "/model_resolver/progress/{download_id}"), + ("GET", "/model_resolver/progress"), + ("POST", "/model_resolver/cancel/{download_id}"), + ("POST", "/model_resolver/pause/{download_id}"), + ("POST", "/model_resolver/resume/{download_id}"), + ("POST", "/model_resolver/clear_completed_downloads"), + ("GET", "/model_resolver/aria2/status"), + ("POST", "/model_resolver/aria2/status"), + ("POST", "/model_resolver/aria2/start"), + ("GET", "/model_resolver/aria2/stop"), + ("POST", "/model_resolver/aria2/stop"), + ("POST", "/model_resolver/aria2/install"), + }, + "directories": { + ("GET", "/model_resolver/directories"), + ("GET", "/model_resolver/root-directories"), + ("GET", "/model_resolver/path-template-suggestions"), + ("GET", "/model_resolver/capabilities"), + ("GET", "/model_resolver/subfolders/{category}"), + }, + } + def test_all_api_routes_are_registered(self): expected_routes = { ("GET", "/model_resolver/base-models"), @@ -145,6 +216,26 @@ class TestRefactoringTargets(unittest.IsolatedAsyncioTestCase): } self.assertEqual(expected_routes, set(routes_registered)) + def test_route_groups_cover_the_registered_feature_routes(self): + grouped_routes = set().union(*self.ROUTE_GROUPS.values()) + self.assertEqual( + grouped_routes, + set(routes_registered) + - { + ("GET", "/model_resolver/base-models"), + ("GET", "/model_resolver/base-models/status"), + ("POST", "/model_resolver/base-models/update"), + ("GET", "/model_resolver/version"), + ("GET", "/model_resolver/logs/backend/export"), + ("GET", "/model_resolver/settings"), + ("POST", "/model_resolver/settings"), + }, + ) + self.assertEqual( + sum(len(routes) for routes in self.ROUTE_GROUPS.values()), + len(grouped_routes), + ) + def test_project_version_helpers_preserve_version_comparison_behavior(self): self.assertEqual( node_mod._extract_project_version('version = "1.2.3"'), @@ -198,6 +289,28 @@ class TestRefactoringTargets(unittest.IsolatedAsyncioTestCase): payload = mock_json_response.call_args.args[0] self.assertIn("settings", payload) self.assertIn("schema", payload) + + async def test_metadata_capabilities_route_returns_capabilities(self): + get_handler = routes_registered[("GET", "/model_resolver/metadata-build/capabilities")] + with patch("aiohttp.web.json_response") as mock_json_response: + await get_handler(MagicMock()) + + payload = mock_json_response.call_args.args[0] + self.assertIn("metadata_modes", payload) + self.assertIn("default_metadata_mode", payload) + + async def test_custom_url_route_rejects_missing_url(self): + post_handler = routes_registered[("POST", "/model_resolver/custom-url")] + request = AsyncMock() + request.json.return_value = {} + + with patch("aiohttp.web.json_response") as mock_json_response: + await post_handler(request) + + mock_json_response.assert_called_once_with( + {"error": "URL is required"}, + status=400, + ) @classmethod def tearDownClass(cls):