nodes/Azornes-Comfyui-Model-Resolver
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Unify model result building and progress payload updates

Browse Source
This commit is contained in:
Dariusz L
2026-07-18 12:26:17 +02:00
parent 92dbfc7932
commit 759aef99f1
10 changed files with 221 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+46 -22
View File
@@ -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)
core/aria2_installer.py
+4 -4
View File
@@ -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()
core/network_utils.py
+15
View File
@@ -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
core/sources/civarchive.py
+4 -4
View File
@@ -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"),
core/sources/civitai.py
+4 -2
View File
@@ -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,
core/sources/common.py
+88
View File
@@ -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,
)
core/sources/huggingface.py
+2 -1
View File
@@ -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,
tests/test_refactoring_unification.py
+50
View File
@@ -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__":
web/resolver/utils/render_format_methods.js
+4
View File
@@ -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 {
web/resolver/views/model_info_methods.js
+4 -1
View File
@@ -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
*/