From 759aef99f1c82f5122ac501069733e729d86a6da Mon Sep 17 00:00:00 2001 From: Dariusz L Date: Sat, 18 Jul 2026 12:26:17 +0200 Subject: [PATCH] Unify model result building and progress payload updates --- __init__.py | 68 ++++++++++------ core/aria2_installer.py | 8 +- core/network_utils.py | 15 ++++ core/sources/civarchive.py | 8 +- core/sources/civitai.py | 6 +- core/sources/common.py | 88 +++++++++++++++++++++ core/sources/huggingface.py | 3 +- tests/test_refactoring_unification.py | 50 ++++++++++++ web/resolver/utils/render_format_methods.js | 4 + web/resolver/views/model_info_methods.js | 5 +- 10 files changed, 221 insertions(+), 34 deletions(-) create mode 100644 core/sources/common.py diff --git a/__init__.py b/__init__.py index ccc340b..4fdea4e 100644 --- a/__init__.py +++ b/__init__.py @@ -166,7 +166,28 @@ class JobProgressTracker: "updated_at": now, } + def update_from_payload( + self, + progress_id: Optional[str], + progress_payload: Dict[str, Any], + default_stage: str = "running", + ) -> None: + if not progress_id or not isinstance(progress_payload, dict): + return + data = dict(progress_payload) + status = data.pop("status", None) + stage = data.pop("stage", None) or default_stage + if not status: + if stage in ("done", "completed"): + status = "completed" + elif stage == "cancelled": + status = "cancelled" + else: + status = "running" + self.update(progress_id, status=status, stage=stage, **data) + def is_cancelled(self, progress_id) -> bool: + if not progress_id: return False with self.lock: @@ -214,29 +235,15 @@ class ModelResolverExtension: analysis_id: Optional[str], payload: Dict[str, Any] ) -> None: - if not analysis_id: - return - status = "running" if payload.get("stage") != "completed" else "completed" - self.analysis_progress.update(analysis_id, status=status, **payload) + self.analysis_progress.update_from_payload(analysis_id, payload) def _update_metadata_build_progress( self, progress_id: Optional[str], progress_payload: Dict[str, Any] ) -> None: - if not progress_id or not isinstance(progress_payload, dict): - return - data = dict(progress_payload) - status = data.pop("status", None) - stage = data.get("stage") or "" - if not status: - if stage == "done": - status = "done" - elif stage == "cancelled": - status = "cancelled" - else: - status = "running" - self.metadata_builder_progress.update(progress_id, status=status, **data) + self.metadata_builder_progress.update_from_payload(progress_id, progress_payload) + def _update_loaded_progress( self, @@ -334,9 +341,18 @@ class ModelResolverExtension: ) from .core.path_templates import infer_download_path_templates from .core.path_utils import ( + HashCalculationCancelled, + calculate_file_sha256, + dedupe_local_base_directories, extract_safetensors_header_sha256, + get_comfy_root_path, get_filename_from_path, + get_path_abs, + get_path_identity, + get_path_key, is_path_in_configured_model_roots, + is_path_within, + prefer_local_base_directory, read_json_safe, write_json_atomic, ) @@ -352,6 +368,7 @@ class ModelResolverExtension: from .core.settings import ( TEMPLATE_KEY_ALIASES, get_default_root_for_category, + get_settings_schema, resolve_download_subfolder, ) from .core.settings import ( @@ -378,7 +395,11 @@ class ModelResolverExtension: to_bool, to_int, ) - from .core.workflow_analyzer import analyze_workflow_models + from .core.workflow_analyzer import ( + NODE_TYPE_MODEL_WIDGET_CATEGORIES, + URN_TYPE_MAP, + analyze_workflow_models, + ) except ImportError as e: self.logger.error(f"Model Resolver: Could not import core modules: {e}") return False @@ -422,6 +443,8 @@ class ModelResolverExtension: check_civitai_session_token, get_civitai_download_url, get_civitai_model_details, + get_model_info_by_hash, + get_model_info_for_file, parse_civitai_url, resolve_civitai_version_custom_result, resolve_urn, @@ -459,8 +482,11 @@ class ModelResolverExtension: update_model_list_from_remote, ) from .core.sources.popular import ( + get_base_models_config, + get_base_models_status, get_popular_model_url, search_popular_models, + update_base_models_from_remote, ) from .core.sources.popular import ( reload_databases as reload_popular_databases, @@ -583,12 +609,12 @@ class ModelResolverExtension: @routes.get("/model_resolver/base-models") @json_api_endpoint("base-models") async def get_base_models(request): - """Return the base models config from metadata/base-models.json. + """ + Return popular base models configuration dictionary. Returns {base_models: [{name, aliases}]} so the frontend dropdown can populate correctly via baseModels.base_models. """ - from .core.sources.popular import get_base_models_config data = get_base_models_config() return web.json_response(data) @@ -597,7 +623,6 @@ class ModelResolverExtension: async def get_base_models_status_route(request): """Get local and optional remote base models status.""" check_remote = request.query.get("check_remote") == "1" - from .core.sources.popular import get_base_models_status status = await asyncio.to_thread(get_base_models_status, check_remote) return web.json_response(status) @@ -605,7 +630,6 @@ class ModelResolverExtension: @json_api_endpoint("base-models update") async def update_base_models_route(request): """Update base models list from CivitAI.""" - from .core.sources.popular import update_base_models_from_remote status = await asyncio.to_thread(update_base_models_from_remote) return web.json_response(status) diff --git a/core/aria2_installer.py b/core/aria2_installer.py index 5d1b7b9..899ec8c 100644 --- a/core/aria2_installer.py +++ b/core/aria2_installer.py @@ -13,6 +13,7 @@ from typing import Any, Dict, Optional import requests +from .network_utils import fetch_json_from_public_url from .path_utils import is_path_within ARIA2_RELEASE_API_URL = "https://api.github.com/repos/aria2/aria2/releases/latest" @@ -178,20 +179,19 @@ def _select_release_asset(release: Dict[str, Any], tokens: Dict[str, str]) -> Di def _fetch_latest_release() -> Dict[str, Any]: try: - response = requests.get(ARIA2_RELEASE_API_URL, headers=REQUEST_HEADERS, timeout=(10, 30)) - response.raise_for_status() - except requests.RequestException as exc: + data = fetch_json_from_public_url(ARIA2_RELEASE_API_URL, headers=REQUEST_HEADERS, timeout=30) + except Exception as exc: raise Aria2InstallError( "Could not contact GitHub to check the latest aria2 release. " "Check your internet connection or firewall/proxy settings. " f"Details: {exc}" ) from exc - data = response.json() if not isinstance(data, dict): raise Aria2InstallError("GitHub returned an invalid aria2 release response.") return data + def _assert_within_install_root(path: Path) -> None: install_root = ARIA2_INSTALL_ROOT.resolve() target = path.resolve() diff --git a/core/network_utils.py b/core/network_utils.py index 06b7df9..7bca8ae 100644 --- a/core/network_utils.py +++ b/core/network_utils.py @@ -285,3 +285,18 @@ def request_source_json( return None +def fetch_json_from_public_url( + url: str, + headers: Optional[Dict[str, str]] = None, + timeout: int = 30, +) -> Dict[str, Any]: + """ + Convenience wrapper to fetch and parse JSON from a public HTTP(S) URL. + """ + res = request_source_json(url, headers=headers, timeout=timeout, raise_on_error=True) + if not isinstance(res, dict): + raise ValueError(f"URL did not return a valid JSON object: {url}") + return res + + + diff --git a/core/sources/civarchive.py b/core/sources/civarchive.py index 21a84de..532227c 100644 --- a/core/sources/civarchive.py +++ b/core/sources/civarchive.py @@ -33,7 +33,6 @@ from ..progress import get_progress_reporter from ..type_utils import ( CIVARCHIVE_API_TYPE_MAP, DEFAULT_BROWSER_USER_AGENT, - build_search_result, clear_remote_size_cache, extract_file_size, extract_trained_words, @@ -46,6 +45,7 @@ from ..type_utils import ( select_primary_model_file, to_int, ) +from .common import build_unified_search_result log = create_module_logger(__name__) @@ -1155,7 +1155,7 @@ def _build_result_from_normalized_version( if size is None: size = _resolve_file_size_bytes(file_info, download_urls) - return build_search_result( + return build_unified_search_result( source="civarchive", model_id=model_id, version_id=version_id, @@ -1612,7 +1612,7 @@ def _build_result_from_payload( else "" ) - return build_search_result( + return build_unified_search_result( source="civarchive", model_id=model_id, version_id=version_id, @@ -1769,7 +1769,7 @@ def _build_result_from_search_candidate( open_url = urljoin(CIVARCHIVE_BASE_URL, str(candidate.get("url") or "")) sha256 = (parsed or {}).get("sha256") - return build_search_result( + return build_unified_search_result( source="civarchive", model_id=candidate.get("model_id") or candidate.get("modelId"), version_id=candidate.get("version_id") or candidate.get("modelVersionId"), diff --git a/core/sources/civitai.py b/core/sources/civitai.py index 7eb3518..bf38751 100644 --- a/core/sources/civitai.py +++ b/core/sources/civitai.py @@ -43,7 +43,6 @@ from ..type_utils import ( DEFAULT_BROWSER_USER_AGENT, as_dict, as_list, - build_search_result, check_credential_http, extract_trained_words, first_non_empty, @@ -150,6 +149,9 @@ def check_civitai_api_key(api_key: Optional[str]) -> Dict[str, Any]: ) +from .common import build_unified_search_result + + def _build_civitai_result_from_version( model_id: int, model_name: str, @@ -165,7 +167,7 @@ def _build_civitai_result_from_version( sha256 = file_info.get("sha256") or hashes.get("SHA256") or hashes.get("sha256") size = file_info.get("sizeKB", 0) * 1024 if file_info.get("sizeKB") else file_info.get("size") download_url = file_info.get("downloadUrl") or get_civitai_download_url(version_id) - return build_search_result( + return build_unified_search_result( source="civitai", model_id=model_id, version_id=version_id, diff --git a/core/sources/common.py b/core/sources/common.py new file mode 100644 index 0000000..c2fa335 --- /dev/null +++ b/core/sources/common.py @@ -0,0 +1,88 @@ +""" +Common utilities and result builder functions for external model metadata sources. +""" + +from typing import Any, Dict, List, Optional + +from ..type_utils import build_search_result, normalize_sha256 + + +def normalize_hashes_dict(hashes: Optional[Dict[str, Any]]) -> Dict[str, str]: + """ + Normalize a hashes dictionary (e.g. from CivitAI API or CivArchive) so that + hash algorithm keys are standardized and SHA256 values are uppercase. + """ + if not isinstance(hashes, dict): + return {} + + normalized: Dict[str, str] = {} + for k, v in hashes.items(): + if not k or not v: + continue + key_str = str(k).strip() + val_str = str(v).strip() + if key_str.lower() in ("sha256", "sha-256"): + normalized["sha256"] = normalize_sha256(val_str) + elif key_str.lower() in ("autov2", "auto_v2"): + normalized["autoV2"] = val_str + elif key_str.lower() in ("autov1", "auto_v1"): + normalized["autoV1"] = val_str + elif key_str.lower() == "blake3": + normalized["blake3"] = val_str + else: + normalized[key_str] = val_str + + return normalized + + +def build_unified_search_result( + source: str, + *, + model_id: Any, + version_id: Any, + name: str = "", + version_name: str = "", + type: str = "", + filename: str = "", + url: str = "", + download_url: Optional[str] = None, + size: Optional[int] = None, + base_model: Optional[str] = None, + tags: Optional[List[str]] = None, + match_type: str = "similar", + confidence: float = 0.0, + sha256: Optional[str] = None, + hashes: Optional[Dict[str, Any]] = None, + trained_words: Optional[List[str]] = None, + images: Optional[List[Dict[str, Any]]] = None, + **extra: Any, +) -> Dict[str, Any]: + """ + Unified result builder for all model metadata sources. + Normalizes hash dictionaries and formats a standard model search result. + """ + norm_hashes = normalize_hashes_dict(hashes) + if not sha256 and "sha256" in norm_hashes: + sha256 = norm_hashes["sha256"] + + return build_search_result( + source, + model_id=model_id, + version_id=version_id, + name=name, + version_name=version_name, + type=type, + filename=filename, + url=url, + download_url=download_url, + size=size, + base_model=base_model, + tags=tags, + match_type=match_type, + confidence=confidence, + sha256=sha256, + hashes=norm_hashes, + trained_words=trained_words, + images=images, + **extra, + ) diff --git a/core/sources/huggingface.py b/core/sources/huggingface.py index 026cfa4..960e30a 100644 --- a/core/sources/huggingface.py +++ b/core/sources/huggingface.py @@ -28,6 +28,7 @@ from ..type_utils import ( looks_like_model_file, prepare_remote_size_probe_url, ) +from .common import build_unified_search_result HF_API_URL = "https://huggingface.co/api" HF_AUTHOR_FALLBACKS = ["Comfy-Org"] @@ -371,7 +372,7 @@ def _build_huggingface_result( size = extract_file_size(file_info) if not size: size = _fetch_remote_file_size_bytes(download_url, headers=headers) - return build_search_result( + return build_unified_search_result( "huggingface", model_id=repo_id, version_id=None, diff --git a/tests/test_refactoring_unification.py b/tests/test_refactoring_unification.py index c3f2471..7b1eef0 100644 --- a/tests/test_refactoring_unification.py +++ b/tests/test_refactoring_unification.py @@ -266,6 +266,55 @@ class TestRefactoringUnification(unittest.TestCase): self.assertEqual(as_list(["a", "", None, "b"]), ["a", "b"]) self.assertEqual(list(_listify_tags(["a", "", None, "b"])), ["a", "b"]) + def test_build_unified_search_result(self): + from core.sources.common import build_unified_search_result, normalize_hashes_dict + + valid_sha = "a" * 64 + norm_hashes = normalize_hashes_dict({"sha256": valid_sha, "autoV2": "1234567890"}) + self.assertEqual(norm_hashes["sha256"], valid_sha.lower()) + self.assertEqual(norm_hashes["autoV2"], "1234567890") + + res = build_unified_search_result( + "civitai", + model_id=123, + version_id=456, + name="Test Model", + hashes={"sha256": valid_sha}, + download_url="https://civitai.com/api/download/123", + ) + self.assertEqual(res["source"], "civitai") + self.assertEqual(res["model_id"], 123) + self.assertEqual(res["version_id"], 456) + self.assertEqual(res["sha256"], valid_sha.lower()) + self.assertEqual(res["hashes"]["sha256"], valid_sha.lower()) + + def test_job_progress_tracker_update_from_payload(self): + import importlib + node_mod = importlib.import_module("comfyui-model-resolver") + JobProgressTracker = node_mod.JobProgressTracker + + tracker = JobProgressTracker("Test Job") + tracker.update_from_payload("job99", {"stage": "processing", "current": 5, "total": 10, "message": "Working..."}) + data = tracker.get("job99") + self.assertEqual(data["status"], "running") + self.assertEqual(data["stage"], "processing") + self.assertEqual(data["current"], 5) + self.assertEqual(data["message"], "Working...") + + tracker.update_from_payload("job99", {"stage": "completed", "message": "Done!"}) + data_done = tracker.get("job99") + self.assertEqual(data_done["status"], "completed") + + @patch("core.network_utils.request_source_json") + def test_fetch_json_from_public_url(self, mock_request_source_json): + from core.network_utils import fetch_json_from_public_url + mock_request_source_json.return_value = {"key": "value"} + + data = fetch_json_from_public_url("https://example.com/api.json") + self.assertEqual(data, {"key": "value"}) + mock_request_source_json.assert_called_once() + + if __name__ == "__main__": unittest.main() @@ -274,3 +323,4 @@ if __name__ == "__main__": + diff --git a/web/resolver/utils/render_format_methods.js b/web/resolver/utils/render_format_methods.js index f27b234..bff94f7 100644 --- a/web/resolver/utils/render_format_methods.js +++ b/web/resolver/utils/render_format_methods.js @@ -5,6 +5,10 @@ import { getSvgIcon } from "../../utils/icon_utils.js"; import { escapeHtml, escapeJsString, getFilenameFromPath, pollBackgroundTask } from "./html_utils.js"; export const renderFormatMethods = { + escapeHtml, + escapeJsString, + getFilenameFromPath, + encodeContextMenuModel(context = null) { if (!context) return ''; try { diff --git a/web/resolver/views/model_info_methods.js b/web/resolver/views/model_info_methods.js index d966b51..d08b3f7 100644 --- a/web/resolver/views/model_info_methods.js +++ b/web/resolver/views/model_info_methods.js @@ -2,10 +2,13 @@ import { app } from "../../../../../scripts/app.js"; import { api } from "../../../../../scripts/api.js"; import { $el } from "../../../../../scripts/ui.js"; import { getSvgIcon } from "../../utils/icon_utils.js"; -import { sanitizeDescriptionHtml, pollBackgroundTask, safeStorage, copyTextWithFeedback } from "../utils/html_utils.js"; +import { escapeHtml, escapeJsString, getFilenameFromPath, sanitizeDescriptionHtml, pollBackgroundTask, safeStorage, copyTextWithFeedback } from "../utils/html_utils.js"; import { getModelCardUrl } from "../utils/url_utils.js"; import { getCivitaiModelUrl } from "../globals.js"; export const modelInfoMethods = { + escapeHtml, + escapeJsString, + getFilenameFromPath, /** * Handle click outside context menu to hide it */