Refactor shared metadata model helpers

This commit is contained in:
Dariusz L
2026-08-03 12:46:48 +02:00
parent 3353423015
commit 128a6ccacf
4 changed files with 141 additions and 82 deletions
+10 -42
View File
@@ -10,13 +10,22 @@ from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Any, Dict, Iterable, List, Optional, Tuple
from .log_system import create_module_logger
from .metadata_model_utils import (
dedupe_models as _dedupe_models,
)
from .metadata_model_utils import (
is_model_file_path as _is_model_file_path,
)
from .metadata_model_utils import (
model_identity_key as _model_identity_key, # noqa: F401
)
from .path_utils import (
find_metadata_sidecar_path,
get_filename_from_path,
get_path_identity,
read_merged_model_metadata,
)
from .type_utils import MODEL_EXTENSIONS, format_size_bytes
from .type_utils import format_size_bytes
log = create_module_logger(__name__)
@@ -158,18 +167,6 @@ def extract_metadata_size(metadata: Dict[str, Any], model_filename: str = "") ->
return None
def _is_model_file_path(path: str) -> bool:
if not path or not os.path.isfile(path):
return False
filename = get_filename_from_path(path).lower()
if filename.endswith((".metadata.json", ".civitai.info")):
return False
return os.path.splitext(filename)[1].lower() in MODEL_EXTENSIONS
def _format_signed_size(bytes_value: int) -> str:
if bytes_value == 0:
return "0 B"
@@ -189,35 +186,6 @@ def _model_key(model_path: str, metadata_path: str) -> Tuple[str, str]:
return model_identity, metadata_identity
def _model_identity_key(model: Dict[str, Any]) -> str:
model_path = str(model.get("path") or "").strip()
if not model_path:
return ""
try:
return get_path_identity(model_path)
except (OSError, ValueError):
try:
return os.path.normcase(os.path.abspath(model_path))
except (OSError, ValueError):
return os.path.normcase(model_path)
def _dedupe_models(models: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
result: List[Dict[str, Any]] = []
seen = set()
for model in models or []:
if not isinstance(model, dict):
continue
identity = _model_identity_key(model)
if not identity:
continue
if identity in seen:
continue
seen.add(identity)
result.append(model)
return result
def _get_worker_count(total_models: int, requested_workers: Optional[int] = None) -> Tuple[int, int]:
cpu_count = os.cpu_count() or 4
if requested_workers is not None:
+9 -40
View File
@@ -13,6 +13,15 @@ from concurrent.futures import CancelledError, ThreadPoolExecutor, as_completed
from typing import Any, Callable, Dict, List, Optional, Tuple
from .log_system import create_module_logger
from .metadata_model_utils import (
dedupe_models as _dedupe_models,
)
from .metadata_model_utils import (
is_model_file_path as _is_model_file_path,
)
from .metadata_model_utils import (
model_identity_key as _model_identity_key,
)
from .path_utils import (
MODEL_RESOLVER_METADATA_SCHEMA,
MODEL_RESOLVER_METADATA_SCHEMA_VERSION,
@@ -22,13 +31,11 @@ from .path_utils import (
find_external_metadata_sidecar_path,
get_filename_from_path,
get_model_resolver_sidecar_path,
get_path_identity,
read_json_safe,
read_safetensors_header,
write_json_atomic,
)
from .type_utils import (
MODEL_EXTENSIONS,
extract_sha256_from_metadata,
format_size_bytes,
normalize_category_to_model_type,
@@ -60,44 +67,6 @@ def _select_metadata_path(model_path: str) -> Tuple[str, str, bool]:
return target_path, target_path if target_exists else "", target_exists
def _is_model_file_path(path: str) -> bool:
if not path or not os.path.isfile(path):
return False
filename = get_filename_from_path(path).lower()
if filename.endswith((".metadata.json", ".civitai.info")):
return False
return os.path.splitext(filename)[1].lower() in MODEL_EXTENSIONS
def _model_identity_key(model: Dict[str, Any]) -> str:
model_path = str(model.get("path") or "").strip()
if not model_path:
return ""
try:
return get_path_identity(model_path)
except (OSError, ValueError):
try:
return os.path.normcase(os.path.abspath(model_path))
except (OSError, ValueError):
return os.path.normcase(model_path)
def _dedupe_models(models: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
result: List[Dict[str, Any]] = []
seen = set()
for model in models or []:
if not isinstance(model, dict):
continue
identity = _model_identity_key(model)
if not identity or identity in seen:
continue
seen.add(identity)
result.append(model)
return result
def _coerce_positive_int(value: Any) -> int:
try:
number = int(value)
+48
View File
@@ -0,0 +1,48 @@
"""Shared helpers for collections of local model metadata entries."""
import os
from typing import Any, Dict, List
from .path_utils import get_filename_from_path, get_path_identity
from .type_utils import MODEL_EXTENSIONS
def is_model_file_path(path: str) -> bool:
"""Return whether ``path`` points to an existing supported model file."""
if not path or not os.path.isfile(path):
return False
filename = get_filename_from_path(path).lower()
if filename.endswith((".metadata.json", ".civitai.info")):
return False
return os.path.splitext(filename)[1].lower() in MODEL_EXTENSIONS
def model_identity_key(model: Dict[str, Any]) -> str:
"""Return the stable filesystem identity for a model entry."""
model_path = str(model.get("path") or "").strip()
if not model_path:
return ""
try:
return get_path_identity(model_path)
except (OSError, ValueError):
try:
return os.path.normcase(os.path.abspath(model_path))
except (OSError, ValueError):
return os.path.normcase(model_path)
def dedupe_models(models: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
"""Keep the first valid model entry for every filesystem identity."""
result: List[Dict[str, Any]] = []
seen = set()
for model in models or []:
if not isinstance(model, dict):
continue
identity = model_identity_key(model)
if not identity or identity in seen:
continue
seen.add(identity)
result.append(model)
return result
+74
View File
@@ -0,0 +1,74 @@
import os
import tempfile
import unittest
from core.metadata_audit import (
_dedupe_models as audit_dedupe_models,
_is_model_file_path as audit_is_model_file_path,
_model_identity_key as audit_model_identity_key,
)
from core.metadata_builder import (
_dedupe_models as builder_dedupe_models,
_is_model_file_path as builder_is_model_file_path,
_model_identity_key as builder_model_identity_key,
)
class MetadataModelHelperTests(unittest.TestCase):
def test_model_file_path_validation_matches_between_metadata_flows(self):
with tempfile.TemporaryDirectory() as tmpdir:
model_path = os.path.join(tmpdir, "model.safetensors")
metadata_path = os.path.join(tmpdir, "model.metadata.json")
civitai_path = os.path.join(tmpdir, "model.civitai.info")
text_path = os.path.join(tmpdir, "notes.txt")
for path in (model_path, metadata_path, civitai_path, text_path):
with open(path, "wb") as handle:
handle.write(b"data")
expected = {
model_path: True,
metadata_path: False,
civitai_path: False,
text_path: False,
os.path.join(tmpdir, "missing.safetensors"): False,
"": False,
}
for path, is_model in expected.items():
self.assertEqual(is_model, builder_is_model_file_path(path))
self.assertEqual(is_model, audit_is_model_file_path(path))
def test_model_identity_and_dedupe_match_between_metadata_flows(self):
with tempfile.TemporaryDirectory() as tmpdir:
model_path = os.path.join(tmpdir, "models", "model.safetensors")
os.makedirs(os.path.dirname(model_path))
with open(model_path, "wb") as handle:
handle.write(b"data")
equivalent_path = os.path.join(tmpdir, "models", ".", "model.safetensors")
first = {"path": model_path, "name": "first"}
duplicate = {"path": equivalent_path, "name": "duplicate"}
empty = {"path": ""}
models = [first, duplicate, empty, "not a model"]
self.assertEqual(
builder_model_identity_key(first),
audit_model_identity_key(first),
)
self.assertEqual(
builder_model_identity_key(first),
builder_model_identity_key(duplicate),
)
self.assertEqual(
audit_model_identity_key(first),
audit_model_identity_key(duplicate),
)
self.assertEqual("", builder_model_identity_key(empty))
self.assertEqual("", audit_model_identity_key(empty))
self.assertEqual([first], builder_dedupe_models(models))
self.assertEqual([first], audit_dedupe_models(models))
if __name__ == "__main__":
unittest.main()