From aff97c5e2908d6f046f58b63f86b5b00b218b9c2 Mon Sep 17 00:00:00 2001 From: Dariusz L Date: Sat, 18 Jul 2026 10:02:34 +0200 Subject: [PATCH] refactor: centralize shared normalization, timestamp, and search-result helpers --- __init__.py | 67 +++++++++++++------------- core/matcher.py | 3 +- core/path_templates.py | 3 +- core/path_utils.py | 15 ++++-- core/settings.py | 21 +++----- core/sources/huggingface.py | 58 ++++++++++++---------- core/sources/model_list.py | 4 +- core/sources/popular.py | 69 +++++++++++++++++---------- core/type_utils.py | 19 ++++++-- tests/test_refactoring_unification.py | 47 ++++++++++++++++++ 10 files changed, 196 insertions(+), 110 deletions(-) diff --git a/__init__.py b/__init__.py index 8872a4d..dc3c5e0 100644 --- a/__init__.py +++ b/__init__.py @@ -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( diff --git a/core/matcher.py b/core/matcher.py index b55f57b..85897b0 100644 --- a/core/matcher.py +++ b/core/matcher.py @@ -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: diff --git a/core/path_templates.py b/core/path_templates.py index f1bc1ad..db0c4d8 100644 --- a/core/path_templates.py +++ b/core/path_templates.py @@ -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]: diff --git a/core/path_utils.py b/core/path_utils.py index 0cb81a8..83a8392 100644 --- a/core/path_utils.py +++ b/core/path_utils.py @@ -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, diff --git a/core/settings.py b/core/settings.py index a020f43..6e65cd1 100644 --- a/core/settings.py +++ b/core/settings.py @@ -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: diff --git a/core/sources/huggingface.py b/core/sources/huggingface.py index 458cb5e..026cfa4 100644 --- a/core/sources/huggingface.py +++ b/core/sources/huggingface.py @@ -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, + ) diff --git a/core/sources/model_list.py b/core/sources/model_list.py index c3965ca..c33ff58 100644 --- a/core/sources/model_list.py +++ b/core/sources/model_list.py @@ -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]: diff --git a/core/sources/popular.py b/core/sources/popular.py index a266019..7cb0214 100644 --- a/core/sources/popular.py +++ b/core/sources/popular.py @@ -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, diff --git a/core/type_utils.py b/core/type_utils.py index e8ed032..3abc989 100644 --- a/core/type_utils.py +++ b/core/type_utils.py @@ -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() + + diff --git a/tests/test_refactoring_unification.py b/tests/test_refactoring_unification.py index fdc7322..c3f2471 100644 --- a/tests/test_refactoring_unification.py +++ b/tests/test_refactoring_unification.py @@ -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__": +