refactor: centralize shared normalization, timestamp, and search-result helpers

This commit is contained in:
Dariusz L
2026-07-18 10:02:34 +02:00
parent 719b3b3f43
commit aff97c5e29
10 changed files with 196 additions and 110 deletions
+32 -35
View File
@@ -364,6 +364,7 @@ class ModelResolverExtension:
save_settings as save_resolver_settings,
)
from .core.type_utils import (
build_search_result,
extract_sha256_from_metadata,
fetch_remote_file_size_cached,
first_non_empty,
@@ -3860,30 +3861,26 @@ class ModelResolverExtension:
download_url = get_civitai_download_url(
version_id_val
)
source_results["civitai"] = {
"source": "civitai",
"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}",
"model_id": model_id_val,
"version_id": 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")
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,
}
hashes=primary_file.get("hashes") or {},
match_type="exact",
confidence=100.0,
)
log_search_result(
"civitai/urn",
source_results["civitai"],
@@ -3927,19 +3924,19 @@ class ModelResolverExtension:
)
if civitai_results:
first_result = civitai_results[0]
source_results["civitai"] = {
"source": "civitai",
"name": first_result.get("name"),
"filename": first_result.get("filename"),
"type": first_result.get("type"),
"download_url": first_result.get(
"download_url"
),
"url": first_result.get("url"),
"size": first_result.get("size"),
"base_model": first_result.get("base_model"),
"tags": first_result.get("tags", []),
}
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(
+2 -1
View File
@@ -239,7 +239,8 @@ def normalize_base_model(value: str) -> str:
Returns:
Normalized base model name
"""
return re.sub(r"[^a-z0-9]+", "", str(value or "").lower())
from .type_utils import normalize_alphanumeric_lower
return normalize_alphanumeric_lower(value)
def base_model_matches(candidate: str, preferred: Optional[str]) -> bool:
+2 -1
View File
@@ -46,7 +46,8 @@ TAG_FOLDER_TOKENS = {
def _normalize_token(value: Any) -> str:
return re.sub(r"[^a-z0-9]+", "", str(value or "").lower())
from .type_utils import normalize_alphanumeric_lower
return normalize_alphanumeric_lower(value)
def _base_model_token_variants(value: Any) -> List[str]:
+12 -3
View File
@@ -383,7 +383,8 @@ def extract_safetensors_header_sha256(
def _normalize_base_model_token(value: Any) -> str:
return re.sub(r"[^a-z0-9]+", "", str(value or "").lower())
from .type_utils import normalize_alphanumeric_lower
return normalize_alphanumeric_lower(value)
def _load_base_model_aliases() -> List[Tuple[str, str]]:
@@ -394,8 +395,9 @@ def _load_base_model_aliases() -> List[Tuple[str, str]]:
aliases: List[Tuple[str, str]] = []
try:
data = read_json_safe(os.path.join(METADATA_DIR, "base-models.json"), {})
for entry in data.get("base_models", []) if isinstance(data, dict) else []:
from .sources.popular import load_raw_base_models
base_models = load_raw_base_models()
for entry in base_models:
if not isinstance(entry, dict):
continue
name = str(entry.get("name") or "").strip()
@@ -418,6 +420,12 @@ def _load_base_model_aliases() -> List[Tuple[str, str]]:
return aliases
def clear_base_model_alias_cache() -> None:
"""Clear the base model alias cache."""
global _BASE_MODEL_ALIAS_CACHE
_BASE_MODEL_ALIAS_CACHE = None
def _canonical_base_model_from_text(value: Any) -> str:
normalized = _normalize_base_model_token(value)
if not normalized:
@@ -429,6 +437,7 @@ def _canonical_base_model_from_text(value: Any) -> str:
return ""
def infer_safetensors_base_model(
file_path: str,
max_tensor_keys: int = 500,
+7 -14
View File
@@ -5,7 +5,7 @@ from __future__ import annotations
import json
import re
from pathlib import Path
from typing import Any, Dict, Iterable, Mapping, Optional
from typing import Any, Dict, List, Mapping, Optional
from .path_utils import read_json_safe, write_json_atomic
@@ -199,22 +199,15 @@ def _template_for_category(settings: Mapping[str, Any], category: str) -> str:
return templates.get(category_key, "")
def _listify_tags(value: Any) -> Iterable[str]:
if isinstance(value, (list, tuple, set)):
for item in value:
text = str(item or "").strip()
if text:
yield text
return
if isinstance(value, str):
for item in re.split(r"[,;]+", value):
text = item.strip()
if text:
yield text
def _listify_tags(value: Any) -> List[str]:
from .type_utils import as_list
return as_list(value)
def _normalize_tag(value: str) -> str:
return re.sub(r"[^a-z0-9]+", "", str(value or "").lower())
from .type_utils import normalize_alphanumeric_lower
return normalize_alphanumeric_lower(value)
def _resolve_base_model_mapping(mappings: Mapping[str, str], base_model: Any) -> str:
+32 -26
View File
@@ -20,6 +20,7 @@ from ..matcher import build_filename_search_queries
from ..network_utils import host_matches_domain, request_source_json, request_source_response
from ..path_utils import METADATA_DIR, get_filename_from_path, read_json_safe, write_json_atomic
from ..type_utils import (
build_search_result,
check_credential_http,
clear_remote_size_cache,
extract_file_size,
@@ -370,16 +371,19 @@ def _build_huggingface_result(
size = extract_file_size(file_info)
if not size:
size = _fetch_remote_file_size_bytes(download_url, headers=headers)
return {
"source": "huggingface",
"repo_id": repo_id,
"filename": get_filename_from_path(file_path),
"path": file_path,
"url": download_url,
"size": size,
"match_type": match_type,
"sha256": sha256,
}
return build_search_result(
"huggingface",
model_id=repo_id,
version_id=None,
filename=get_filename_from_path(file_path),
url=download_url,
download_url=download_url,
size=size,
match_type=match_type,
sha256=sha256,
repo_id=repo_id,
path=file_path,
)
def _get_repo_tree(
@@ -1070,20 +1074,22 @@ def build_huggingface_custom_result(
page_url = (
f"https://huggingface.co/{repo_id}/blob/{branch}/{quote_url_path(file_path)}"
)
return {
"source": "huggingface",
"details_source": "huggingface",
"repo_id": repo_id,
"repo": repo_id,
"path": file_path,
"filename": filename,
"name": repo_id,
"url": download_url,
"download_url": download_url,
"page_url": page_url,
"version_url": page_url,
"size": size,
"match_type": "custom_url",
"custom_url": True,
}
return build_search_result(
"huggingface",
model_id=repo_id,
version_id=None,
name=repo_id,
filename=filename,
url=download_url,
download_url=download_url,
size=size,
match_type="custom_url",
details_source="huggingface",
repo_id=repo_id,
repo=repo_id,
path=file_path,
page_url=page_url,
version_url=page_url,
custom_url=True,
)
+2 -2
View File
@@ -6,7 +6,6 @@ Search the ComfyUI Manager model-list.json database with fuzzy matching.
import hashlib
import os
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional
from ..log_system import create_module_logger
@@ -36,7 +35,8 @@ _model_list_cache: Optional[List[Dict]] = None
def _utc_now_iso() -> str:
return datetime.now(timezone.utc).replace(microsecond=0).isoformat()
from ..type_utils import utc_now_iso
return utc_now_iso()
def _fetch_json_url(url: str) -> Dict[str, Any]:
+44 -25
View File
@@ -6,7 +6,6 @@ Curated list of common models with known download URLs.
import os
import re
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional
from ..log_system import create_module_logger
@@ -58,6 +57,22 @@ from ..matcher import normalize_base_model as _normalize_base_model
base_models_mgr = CatalogManager(BASE_MODELS_FILE, BASE_MODELS_META_FILE, "base_models")
_raw_base_models_cache: Optional[List[Dict[str, Any]]] = None
def load_raw_base_models() -> List[Dict[str, Any]]:
"""Loads and caches the raw base models database."""
global _raw_base_models_cache
if _raw_base_models_cache is None:
try:
data = base_models_mgr.read_data()
_raw_base_models_cache = data.get("base_models", [])
except Exception as e:
log.error(f"Error loading raw base models: {e}")
_raw_base_models_cache = []
return _raw_base_models_cache
def load_base_model_aliases() -> Dict[str, List[str]]:
"""Load and normalize base model aliases from base-models.json."""
global _base_models_aliases_cache
@@ -67,31 +82,30 @@ def load_base_model_aliases() -> Dict[str, List[str]]:
aliases_dict = {}
try:
data = base_models_mgr.read_data()
base_models = data.get("base_models", [])
base_models = load_raw_base_models()
for model in base_models:
name = model.get("name", "")
normalized_name = _normalize_base_model(name)
if not normalized_name:
continue
name = model.get("name", "")
normalized_name = _normalize_base_model(name)
if not normalized_name:
continue
aliases = model.get("aliases", [])
# Normalize each alias token and deduplicate
normalized_tokens = []
seen_tokens = set()
for alias in aliases:
if not alias:
continue
normalized_alias = _normalize_base_model(alias)
if normalized_alias and normalized_alias not in seen_tokens:
seen_tokens.add(normalized_alias)
normalized_tokens.append(normalized_alias)
aliases = model.get("aliases", [])
# Normalize each alias token and deduplicate
normalized_tokens = []
seen_tokens = set()
for alias in aliases:
if not alias:
continue
normalized_alias = _normalize_base_model(alias)
if normalized_alias and normalized_alias not in seen_tokens:
seen_tokens.add(normalized_alias)
normalized_tokens.append(normalized_alias)
# Ensure normalized name itself is in preferred tokens if not already
if normalized_name not in seen_tokens:
normalized_tokens.append(normalized_name)
# Ensure normalized name itself is in preferred tokens if not already
if normalized_name not in seen_tokens:
normalized_tokens.append(normalized_name)
aliases_dict[normalized_name] = normalized_tokens
aliases_dict[normalized_name] = normalized_tokens
except Exception as e:
log.error(f"Error loading base model aliases: {e}")
@@ -186,10 +200,13 @@ def get_all_popular_models() -> Dict[str, Any]:
def reload_databases():
"""Force reload of all databases."""
global _popular_models_cache, _model_aliases_cache, _base_models_aliases_cache
global _popular_models_cache, _model_aliases_cache, _base_models_aliases_cache, _raw_base_models_cache
_popular_models_cache = None
_model_aliases_cache = None
_base_models_aliases_cache = None
_raw_base_models_cache = None
from ..path_utils import clear_base_model_alias_cache
clear_base_model_alias_cache()
_load_popular_models()
_load_model_aliases()
load_base_model_aliases()
@@ -283,8 +300,9 @@ def get_base_models_status(check_remote: bool = False) -> Dict[str, Any]:
new_models_found = True
break
from ..type_utils import utc_now_iso
status.update({
"remote_checked_at": datetime.now(timezone.utc).replace(microsecond=0).isoformat(),
"remote_checked_at": utc_now_iso(),
"update_available": new_models_found
})
except Exception as e:
@@ -349,7 +367,8 @@ def update_base_models_from_remote() -> Dict[str, Any]:
new_added_count += 1
new_added_names.append(name)
now_str = datetime.now(timezone.utc).replace(microsecond=0).isoformat()
from ..type_utils import utc_now_iso
now_str = utc_now_iso()
meta = {
"updated_at": now_str,
"last_checked_at": now_str,
+16 -3
View File
@@ -4,6 +4,7 @@ Type Utilities Module
Unified helper functions for safe data type casting and normalization.
"""
import re
from typing import Any, Dict, List
@@ -27,7 +28,7 @@ def as_list(value: Any) -> List[Any]:
Supports:
- lists (returns list filtered of None/empty strings)
- tuples and sets (casts to list and filters)
- comma-separated strings (splits, strips, and filters)
- comma-separated or semicolon-separated strings (splits, strips, and filters)
Args:
value: Value to cast
@@ -40,10 +41,11 @@ def as_list(value: Any) -> List[Any]:
if isinstance(value, (tuple, set)):
return [item for item in value if item not in (None, "")]
if isinstance(value, str):
return [item.strip() for item in value.split(",") if item.strip()]
return [item.strip() for item in re.split(r"[,;]+", value) if item.strip()]
return []
MODEL_EXTENSIONS = {
".ckpt",
".pt",
@@ -262,7 +264,6 @@ CIVARCHIVE_API_TYPE_MAP = {
"motion_modules": "MotionModule",
}
import re
from typing import Any, Callable, Dict, List, Optional, Tuple
@@ -1170,3 +1171,15 @@ def parse_provider_model_url(url: str, allowed_domains: list[str]) -> Optional[D
return parse_civitai_model_path(parsed.path, parsed.query)
def normalize_alphanumeric_lower(value: Any) -> str:
"""Normalize a string or token to standard lowercase alphanumeric format."""
return re.sub(r"[^a-z0-9]+", "", str(value or "").lower())
def utc_now_iso() -> str:
"""Return the current UTC timestamp formatted as ISO 8601 string without microseconds."""
from datetime import datetime, timezone
return datetime.now(timezone.utc).replace(microsecond=0).isoformat()
+47
View File
@@ -19,6 +19,10 @@ from core.type_utils import (
class TestRefactoringUnification(unittest.TestCase):
def tearDown(self):
from core.sources.popular import reload_databases
reload_databases()
def test_thresholds(self):
self.assertEqual(MODEL_TITLE_MATCH_THRESHOLD, 82.0)
@@ -220,6 +224,48 @@ class TestRefactoringUnification(unittest.TestCase):
self.assertEqual(get_filename_from_path(""), "")
self.assertEqual(get_filename_from_path(None), "")
def test_alphanumeric_normalizers(self):
from core.matcher import normalize_base_model
from core.path_utils import _normalize_base_model_token
from core.path_templates import _normalize_token
from core.settings import _normalize_tag
inputs = ["SD 1.5", "sd-xl_1.0!!", "Flux.1-dev", "", None]
for val in inputs:
res_matcher = normalize_base_model(val)
res_path = _normalize_base_model_token(val)
res_temp = _normalize_token(val)
res_set = _normalize_tag(val)
# verify they all behave identically
self.assertEqual(res_matcher, res_path)
self.assertEqual(res_matcher, res_temp)
self.assertEqual(res_matcher, res_set)
# verify correct regex behavior
import re
expected = re.sub(r"[^a-z0-9]+", "", str(val or "").lower())
self.assertEqual(res_matcher, expected)
def test_utc_now_iso(self):
from core.sources.model_list import _utc_now_iso
res = _utc_now_iso()
self.assertTrue(isinstance(res, str))
self.assertEqual(len(res), 25) # YYYY-MM-DDTHH:MM:SS+00:00
self.assertIn("T", res)
self.assertTrue(res.endswith("+00:00"))
# Check no microsecond
self.assertNotIn(".", res)
def test_listification(self):
from core.type_utils import as_list
from core.settings import _listify_tags
self.assertEqual(as_list("a,b, c"), ["a", "b", "c"])
self.assertEqual(list(_listify_tags("a,b; c")), ["a", "b", "c"])
self.assertEqual(as_list(["a", "", None, "b"]), ["a", "b"])
self.assertEqual(list(_listify_tags(["a", "", None, "b"])), ["a", "b"])
if __name__ == "__main__":
unittest.main()
@@ -227,3 +273,4 @@ if __name__ == "__main__":