Refactor shared metadata model helpers
This commit is contained in:
+10
-42
@@ -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:
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user