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

Unify provider search result builders

Browse Source
This commit is contained in:
Dariusz L
2026-08-05 06:13:18 +02:00
parent 5f899316ce
commit bb01732b89
6 changed files with 62 additions and 97 deletions
Download Patch File Download Diff File Expand all files Collapse all files
core/sources/civarchive.py
+7 -4
View File
@@ -34,6 +34,7 @@ from ..path_utils import get_filename_from_path
from ..progress import get_progress_reporter
from ..type_utils import (
DEFAULT_BROWSER_USER_AGENT,
build_search_result,
clear_remote_size_cache,
extract_file_size,
extract_trained_words,
@@ -50,7 +51,6 @@ from ..type_utils import (
)
from .common import (
build_custom_url_result,
build_unified_search_result,
collect_download_urls,
is_remote_link_marked_dead,
)
@@ -1208,7 +1208,7 @@ def _build_result_from_normalized_version(
if size is None:
size = _resolve_file_size_bytes(file_info, download_urls)
return build_unified_search_result(
return build_search_result(
source="civarchive",
model_id=model_id,
version_id=version_id,
@@ -1231,6 +1231,7 @@ def _build_result_from_normalized_version(
platform=model_details.get("platform"),
is_deleted=False,
match_type=match_type,
normalize_hashes=True,
)
@@ -1658,7 +1659,7 @@ def _build_result_from_payload(
else ""
)
return build_unified_search_result(
return build_search_result(
source="civarchive",
model_id=model_id,
version_id=version_id,
@@ -1689,6 +1690,7 @@ def _build_result_from_payload(
sha256=sha256,
hash=sha256,
hashes=hashes,
normalize_hashes=True,
)
@@ -1815,7 +1817,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_unified_search_result(
return build_search_result(
source="civarchive",
model_id=candidate.get("model_id") or candidate.get("modelId"),
version_id=candidate.get("version_id") or candidate.get("modelVersionId"),
@@ -1839,6 +1841,7 @@ def _build_result_from_search_candidate(
sha256=sha256,
hash=sha256,
hashes={"SHA256": sha256} if sha256 else {},
normalize_hashes=True,
)
core/sources/civitai.py
+6 -3
View File
@@ -44,6 +44,7 @@ from ..type_utils import (
DEFAULT_BROWSER_USER_AGENT,
as_dict,
as_list,
build_search_result,
check_credential_http,
extract_trained_words,
first_non_empty,
@@ -143,7 +144,7 @@ def check_civitai_api_key(api_key: Optional[str]) -> Dict[str, Any]:
)
from .common import build_custom_url_result, build_unified_search_result
from .common import build_custom_url_result
def _build_civitai_result_from_version(
@@ -161,7 +162,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_unified_search_result(
return build_search_result(
source="civitai",
model_id=model_id,
version_id=version_id,
@@ -176,6 +177,7 @@ def _build_civitai_result_from_version(
match_type=match_type,
sha256=sha256,
hashes=hashes,
normalize_hashes=True,
)
@@ -934,7 +936,7 @@ def search_civitai_for_file(
api_key=api_key,
)
if hash_result:
result = build_unified_search_result(
result = build_search_result(
"civitai",
model_id=hash_result.get("model_id"),
version_id=hash_result.get("version_id"),
@@ -953,6 +955,7 @@ def search_civitai_for_file(
confidence=100.0,
sha256=requested_sha256,
hashes=hash_result.get("hashes") or {"SHA256": requested_sha256},
normalize_hashes=True,
description=hash_result.get("description") or "",
model_description=hash_result.get("model_description") or "",
)
core/sources/common.py
-83
View File
@@ -4,8 +4,6 @@ Common utilities and result builder functions for external model metadata source
from typing import Any, Callable, Dict, List, Optional, Tuple
from ..type_utils import build_search_result, normalize_sha256
def is_remote_link_marked_dead(item: Any) -> bool:
"""Return whether remote metadata marks a link as unavailable."""
@@ -95,87 +93,6 @@ def collect_download_urls(
return urls
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,
)
def build_custom_url_result(
source: str,
*,
core/sources/huggingface.py
+2 -2
View File
@@ -31,7 +31,6 @@ from ..type_utils import (
normalize_sha256,
prepare_remote_size_probe_url,
)
from .common import build_unified_search_result
HF_API_URL = "https://huggingface.co/api"
HF_AUTHOR_FALLBACKS = ["Comfy-Org", "Kijai"]
@@ -717,7 +716,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_unified_search_result(
return build_search_result(
"huggingface",
model_id=repo_id,
version_id=None,
@@ -727,6 +726,7 @@ def _build_huggingface_result(
size=size,
match_type=match_type,
sha256=sha256,
normalize_hashes=True,
repo_id=repo_id,
path=file_path,
branch="main",
core/type_utils.py
+32 -1
View File
@@ -874,6 +874,31 @@ def normalize_sha256(value: Any) -> str:
return text.lower() if SHA256_PATTERN.match(text) else ""
def normalize_hashes_dict(hashes: Optional[Dict[str, Any]]) -> Dict[str, str]:
"""Normalize common model hash keys and values into a stable mapping."""
if not isinstance(hashes, dict):
return {}
normalized: Dict[str, str] = {}
for key, value in hashes.items():
if not key or not value:
continue
key_str = str(key).strip()
value_str = str(value).strip()
if key_str.lower() in ("sha256", "sha-256"):
normalized["sha256"] = normalize_sha256(value_str)
elif key_str.lower() in ("autov2", "auto_v2"):
normalized["autoV2"] = value_str
elif key_str.lower() in ("autov1", "auto_v1"):
normalized["autoV1"] = value_str
elif key_str.lower() == "blake3":
normalized["blake3"] = value_str
else:
normalized[key_str] = value_str
return normalized
def unique_ordered_strings(values: List[Any]) -> List[str]:
"""Return unique, non-empty strings while preserving their original order."""
seen = set()
@@ -1245,12 +1270,18 @@ def build_search_result(
match_type: str = "similar",
confidence: float = 0.0,
sha256: Optional[str] = None,
hashes: Optional[Dict[str, str]] = None,
hashes: Optional[Dict[str, Any]] = None,
normalize_hashes: bool = False,
trained_words: Optional[List[str]] = None,
images: Optional[List[Dict[str, Any]]] = None,
**extra: Any,
) -> Dict[str, Any]:
"""Helper to unify search result dictionary creation across different sources."""
if normalize_hashes:
hashes = normalize_hashes_dict(hashes)
if not sha256 and "sha256" in hashes:
sha256 = hashes["sha256"]
result = {
"source": source,
"model_id": model_id,
tests/test_refactoring_unification.py
+15 -4
View File
@@ -262,22 +262,33 @@ 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
def test_build_search_result_normalizes_hashes(self):
from core.type_utils import 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(
res = build_search_result(
"civitai",
model_id=123,
version_id=456,
name="Test Model",
hashes={"sha256": valid_sha},
download_url="https://civitai.com/api/download/123",
normalize_hashes=True,
)
expected = build_search_result(
"civitai",
model_id=123,
version_id=456,
name="Test Model",
download_url="https://civitai.com/api/download/123",
sha256=valid_sha.lower(),
hashes={"sha256": valid_sha.lower()},
)
self.assertEqual(expected, res)
self.assertEqual(res["source"], "civitai")
self.assertEqual(res["model_id"], 123)
self.assertEqual(res["version_id"], 456)