diff --git a/core/metadata_audit.py b/core/metadata_audit.py index 0cef847..2be78ba 100644 --- a/core/metadata_audit.py +++ b/core/metadata_audit.py @@ -17,14 +17,12 @@ from .path_utils import ( get_path_identity, read_merged_model_metadata, ) -from .type_utils import format_size_bytes +from .type_utils import FILE_SIZE_BYTE_KEYS, FILE_SIZE_KIB_KEYS, format_size_bytes from .worker_utils import resolve_worker_count log = create_module_logger(__name__) -BYTE_SIZE_KEYS = ("size", "sizeBytes", "size_bytes", "fileSize", "file_size", "bytes") -KIB_SIZE_KEYS = ("sizeKB", "size_kb") MAX_AUDIT_WORKERS = 64 MIN_AUDIT_WORKERS = 1 MIN_AUDIT_BATCH_SIZE = 16 @@ -68,10 +66,10 @@ def _coerce_size_bytes(value: Any, multiplier: int = 1) -> Optional[int]: def _iter_size_fields(data: Dict[str, Any], prefix: str = "") -> Iterable[Tuple[str, Any, int]]: - for key in BYTE_SIZE_KEYS: + for key in FILE_SIZE_BYTE_KEYS: if key in data: yield (f"{prefix}{key}", data.get(key), 1) - for key in KIB_SIZE_KEYS: + for key in FILE_SIZE_KIB_KEYS: if key in data: yield (f"{prefix}{key}", data.get(key), 1024) diff --git a/core/type_utils.py b/core/type_utils.py index 23b6a3c..326f6de 100644 --- a/core/type_utils.py +++ b/core/type_utils.py @@ -118,6 +118,11 @@ QUANTIZATION_LEVEL_SUFFIXES = ( ) +FILE_SIZE_BYTE_KEYS = ("size", "sizeBytes", "size_bytes", "fileSize", "file_size", "bytes") +FILE_SIZE_KIB_KEYS = ("sizeKB", "size_kb") +FILE_SIZE_LFS_KEYS = FILE_SIZE_BYTE_KEYS[:3] + + QUANTIZATION_SCHEME_SUFFIXES = ( "gptq", "awq", @@ -1119,7 +1124,7 @@ def extract_file_size(file_info: Dict[str, Any]) -> Optional[int]: return None # 1. Check direct sizeKB or size_kb keys (e.g. from CivArchive or CivitAI) - for key in ("sizeKB", "size_kb"): + for key in FILE_SIZE_KIB_KEYS: val = file_info.get(key) if val is not None and val != "": try: @@ -1128,7 +1133,7 @@ def extract_file_size(file_info: Dict[str, Any]) -> Optional[int]: pass # 2. Check direct size bytes keys - for key in ("sizeBytes", "size_bytes", "fileSize", "file_size", "bytes", "size"): + for key in (*FILE_SIZE_BYTE_KEYS[1:], FILE_SIZE_BYTE_KEYS[0]): val = file_info.get(key) if val is not None and val != "": size = parse_size_header(val) @@ -1138,7 +1143,7 @@ def extract_file_size(file_info: Dict[str, Any]) -> Optional[int]: # 3. Check nested LFS info (common in HuggingFace API metadata) lfs_info = file_info.get("lfs") if isinstance(lfs_info, dict): - for key in ("size", "sizeBytes", "size_bytes"): + for key in FILE_SIZE_LFS_KEYS: size = parse_size_header(lfs_info.get(key)) if size: return size diff --git a/tests/test_file_size_aliases.py b/tests/test_file_size_aliases.py new file mode 100644 index 0000000..95cda3a --- /dev/null +++ b/tests/test_file_size_aliases.py @@ -0,0 +1,36 @@ +import unittest + +from core.metadata_audit import extract_metadata_size +from core.type_utils import extract_file_size + + +class FileSizeAliasTests(unittest.TestCase): + + def test_byte_aliases_are_supported_by_both_extractors(self): + byte_aliases = ("size", "sizeBytes", "size_bytes", "fileSize", "file_size", "bytes") + + for key in byte_aliases: + with self.subTest(key=key): + self.assertEqual(2048, extract_file_size({key: 2048})) + self.assertEqual((2048, key), extract_metadata_size({key: 2048}, "model.safetensors")) + + def test_kibibyte_aliases_are_supported_by_both_extractors(self): + for key in ("sizeKB", "size_kb"): + with self.subTest(key=key): + self.assertEqual(2048, extract_file_size({key: 2})) + self.assertEqual((2048, key), extract_metadata_size({key: 2}, "model.safetensors")) + + def test_lfs_byte_aliases_remain_supported(self): + for key in ("size", "sizeBytes", "size_bytes"): + with self.subTest(key=key): + self.assertEqual(2048, extract_file_size({"lfs": {key: 2048}})) + + def test_existing_extractor_precedence_is_preserved_when_aliases_collide(self): + metadata = {"size": 1024, "sizeBytes": 2048} + + self.assertEqual(2048, extract_file_size(metadata)) + self.assertEqual((2048, "sizeBytes"), extract_metadata_size(metadata, "model.safetensors")) + + +if __name__ == "__main__": + unittest.main()