diff --git a/core/metadata_audit.py b/core/metadata_audit.py index 60718d2..773b659 100644 --- a/core/metadata_audit.py +++ b/core/metadata_audit.py @@ -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: diff --git a/core/metadata_builder.py b/core/metadata_builder.py index ac7d0fa..bd0acec 100644 --- a/core/metadata_builder.py +++ b/core/metadata_builder.py @@ -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) diff --git a/core/metadata_model_utils.py b/core/metadata_model_utils.py new file mode 100644 index 0000000..0758726 --- /dev/null +++ b/core/metadata_model_utils.py @@ -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 diff --git a/tests/test_metadata_model_helpers.py b/tests/test_metadata_model_helpers.py new file mode 100644 index 0000000..53d5f54 --- /dev/null +++ b/tests/test_metadata_model_helpers.py @@ -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()