refactor: centralize shared normalization, timestamp, and search-result helpers
This commit is contained in:
+32
-35
@@ -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
@@ -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:
|
||||
|
||||
@@ -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
@@ -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
@@ -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
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
@@ -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
@@ -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
@@ -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()
|
||||
|
||||
|
||||
|
||||
@@ -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__":
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user