379 lines
12 KiB
Python
379 lines
12 KiB
Python
"""
|
|
Metadata maintenance helpers.
|
|
|
|
Provides audits for local model sidecar metadata files.
|
|
"""
|
|
|
|
import math
|
|
import os
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
from typing import Any, Dict, List, Optional, Tuple
|
|
|
|
from .contracts import ResolvedModel
|
|
from .log_system import create_module_logger
|
|
from .metadata_model_utils import ModelRecord, dedupe_models, is_model_file_path
|
|
from .metadata_utils import merge_counted_payload
|
|
from .path_utils import (
|
|
find_metadata_sidecar_path,
|
|
get_filename_from_path,
|
|
get_path_identity,
|
|
read_merged_model_metadata,
|
|
)
|
|
from .type_utils import extract_file_size_detail, format_size_bytes
|
|
from .worker_utils import resolve_worker_count
|
|
|
|
log = create_module_logger(__name__)
|
|
|
|
|
|
MAX_AUDIT_WORKERS = 64
|
|
MIN_AUDIT_WORKERS = 1
|
|
MIN_AUDIT_BATCH_SIZE = 16
|
|
MAX_AUDIT_BATCH_SIZE = 256
|
|
_AUDIT_COUNT_KEYS = (
|
|
"scanned_models",
|
|
"metadata_files",
|
|
"checked_metadata",
|
|
"missing_metadata",
|
|
"missing_size",
|
|
"invalid_metadata",
|
|
"skipped_directories",
|
|
"skipped_non_model_files",
|
|
)
|
|
_AUDIT_LIST_KEYS = ("errors", "mismatches")
|
|
|
|
|
|
def _empty_audit_counts() -> Dict[str, Any]:
|
|
return {
|
|
"scanned_models": 0,
|
|
"metadata_files": 0,
|
|
"checked_metadata": 0,
|
|
"missing_metadata": 0,
|
|
"missing_size": 0,
|
|
"invalid_metadata": 0,
|
|
"skipped_directories": 0,
|
|
"skipped_non_model_files": 0,
|
|
"errors": [],
|
|
"mismatches": [],
|
|
}
|
|
|
|
|
|
def _extract_size_from_object(data: Any, prefix: str = "") -> Optional[Tuple[int, str]]:
|
|
if not isinstance(data, dict):
|
|
return None
|
|
|
|
detail = extract_file_size_detail(data)
|
|
if detail is not None:
|
|
size, field_name = detail
|
|
if field_name:
|
|
return size, f"{prefix}{field_name}"
|
|
return size, f"{prefix}size"
|
|
|
|
return None
|
|
|
|
|
|
def _file_entry_matches_model(entry: Dict[str, Any], model_filename: str) -> bool:
|
|
entry_name = get_filename_from_path(
|
|
entry.get("name")
|
|
or entry.get("filename")
|
|
or entry.get("file_name")
|
|
or entry.get("fileName")
|
|
or ""
|
|
).lower()
|
|
return bool(entry_name and entry_name == model_filename.lower())
|
|
|
|
|
|
def _select_file_entry(entries: Any, model_filename: str) -> Optional[Dict[str, Any]]:
|
|
if not isinstance(entries, list):
|
|
return None
|
|
|
|
candidates = [entry for entry in entries if isinstance(entry, dict)]
|
|
if not candidates:
|
|
return None
|
|
|
|
for entry in candidates:
|
|
if _file_entry_matches_model(entry, model_filename):
|
|
return entry
|
|
|
|
for entry in candidates:
|
|
if entry.get("primary") and str(entry.get("type") or "").lower() == "model":
|
|
return entry
|
|
|
|
for entry in candidates:
|
|
if entry.get("primary"):
|
|
return entry
|
|
|
|
if len(candidates) == 1:
|
|
return candidates[0]
|
|
|
|
return None
|
|
|
|
|
|
def extract_metadata_size(metadata: Dict[str, Any], model_filename: str = "") -> Optional[Tuple[int, str]]:
|
|
"""
|
|
Return the stored model size and the field path it came from.
|
|
|
|
Top-level size is preferred because Model Resolver writes LoRA
|
|
Manager-compatible sidecars with the authoritative value there.
|
|
"""
|
|
top_level = _extract_size_from_object(metadata)
|
|
if top_level:
|
|
return top_level
|
|
|
|
for key in ("file", "file_info", "selected_file", "selectedFile"):
|
|
nested = metadata.get(key)
|
|
found = _extract_size_from_object(nested, f"{key}.")
|
|
if found:
|
|
return found
|
|
|
|
nested_sources = [
|
|
("files", metadata.get("files")),
|
|
("selected_version.files", (metadata.get("selected_version") or {}).get("files") if isinstance(metadata.get("selected_version"), dict) else None),
|
|
("selectedVersion.files", (metadata.get("selectedVersion") or {}).get("files") if isinstance(metadata.get("selectedVersion"), dict) else None),
|
|
("civitai.files", (metadata.get("civitai") or {}).get("files") if isinstance(metadata.get("civitai"), dict) else None),
|
|
]
|
|
for prefix, entries in nested_sources:
|
|
entry = _select_file_entry(entries, model_filename)
|
|
found = _extract_size_from_object(entry, f"{prefix}[].")
|
|
if found:
|
|
return found
|
|
|
|
return None
|
|
|
|
|
|
def _format_signed_size(bytes_value: int) -> str:
|
|
if bytes_value == 0:
|
|
return "0 B"
|
|
sign = "+" if bytes_value > 0 else "-"
|
|
return f"{sign}{format_size_bytes(abs(bytes_value))}"
|
|
|
|
|
|
def _model_key(model_path: str, metadata_path: str) -> Tuple[str, str]:
|
|
return get_path_identity(model_path), get_path_identity(metadata_path)
|
|
|
|
|
|
def _get_batch_size(
|
|
total_models: int,
|
|
worker_count: int,
|
|
requested_batch_size: Optional[int] = None,
|
|
) -> int:
|
|
if requested_batch_size is not None:
|
|
try:
|
|
batch_size = int(requested_batch_size)
|
|
except (TypeError, ValueError):
|
|
batch_size = 0
|
|
if batch_size > 0:
|
|
return max(1, batch_size)
|
|
|
|
if total_models <= 0:
|
|
return MIN_AUDIT_BATCH_SIZE
|
|
|
|
target_batches = max(1, worker_count * 4)
|
|
batch_size = math.ceil(total_models / target_batches)
|
|
return max(MIN_AUDIT_BATCH_SIZE, min(MAX_AUDIT_BATCH_SIZE, batch_size))
|
|
|
|
|
|
def _make_batches(models: List[ResolvedModel], batch_size: int) -> List[List[ResolvedModel]]:
|
|
if not models:
|
|
return []
|
|
safe_batch_size = max(1, int(batch_size or 1))
|
|
return [
|
|
models[index:index + safe_batch_size]
|
|
for index in range(0, len(models), safe_batch_size)
|
|
]
|
|
|
|
|
|
def _audit_one_model(model: ResolvedModel) -> Dict[str, Any]:
|
|
result = _empty_audit_counts()
|
|
model_path = model.path.strip()
|
|
if not model_path:
|
|
return result
|
|
|
|
if os.path.isdir(model_path):
|
|
result["skipped_directories"] += 1
|
|
return result
|
|
|
|
if not is_model_file_path(model_path):
|
|
result["skipped_non_model_files"] += 1
|
|
return result
|
|
|
|
result["scanned_models"] += 1
|
|
|
|
try:
|
|
actual_size = os.path.getsize(model_path)
|
|
except OSError as exc:
|
|
result["errors"].append(
|
|
{
|
|
"model_path": model_path,
|
|
"metadata_path": "",
|
|
"message": str(exc),
|
|
}
|
|
)
|
|
return result
|
|
|
|
selected_metadata_path = find_metadata_sidecar_path(model_path)
|
|
existing_metadata_paths = (
|
|
[selected_metadata_path]
|
|
if selected_metadata_path and os.path.isfile(selected_metadata_path)
|
|
else []
|
|
)
|
|
if not existing_metadata_paths:
|
|
result["missing_metadata"] += 1
|
|
return result
|
|
|
|
model_filename = get_filename_from_path(model_path)
|
|
seen_entries = set()
|
|
for metadata_path in existing_metadata_paths:
|
|
entry_key = _model_key(model_path, metadata_path)
|
|
if entry_key in seen_entries:
|
|
continue
|
|
seen_entries.add(entry_key)
|
|
result["metadata_files"] += 1
|
|
|
|
metadata = read_merged_model_metadata(model_path, None)
|
|
if not isinstance(metadata, dict):
|
|
result["invalid_metadata"] += 1
|
|
result["errors"].append(
|
|
{
|
|
"model_path": model_path,
|
|
"metadata_path": metadata_path,
|
|
"message": "Metadata JSON is missing, invalid, or not an object.",
|
|
}
|
|
)
|
|
continue
|
|
|
|
size_info = extract_metadata_size(metadata, model_filename)
|
|
if not size_info:
|
|
result["missing_size"] += 1
|
|
continue
|
|
|
|
result["checked_metadata"] += 1
|
|
metadata_size, size_field = size_info
|
|
if metadata_size == actual_size:
|
|
continue
|
|
|
|
difference = actual_size - metadata_size
|
|
result["mismatches"].append(
|
|
{
|
|
"filename": model_filename,
|
|
"relative_path": model.relative_path or model_filename,
|
|
"category": model.category,
|
|
"base_directory": model.base_directory,
|
|
"model_path": model_path,
|
|
"metadata_path": metadata_path,
|
|
"metadata_size": metadata_size,
|
|
"actual_size": actual_size,
|
|
"difference": difference,
|
|
"metadata_size_label": format_size_bytes(metadata_size),
|
|
"actual_size_label": format_size_bytes(actual_size),
|
|
"difference_label": _format_signed_size(difference),
|
|
"size_field": size_field,
|
|
}
|
|
)
|
|
|
|
return result
|
|
|
|
|
|
def _audit_model_batch(batch: List[ResolvedModel]) -> Dict[str, Any]:
|
|
result = _empty_audit_counts()
|
|
for model in batch:
|
|
merge_counted_payload(
|
|
result,
|
|
_audit_one_model(model),
|
|
numeric_keys=_AUDIT_COUNT_KEYS,
|
|
list_keys=_AUDIT_LIST_KEYS,
|
|
)
|
|
return result
|
|
|
|
|
|
def audit_metadata_sizes(
|
|
models: Optional[List[ModelRecord]] = None,
|
|
*,
|
|
force_rescan: bool = True,
|
|
worker_count: Optional[int] = None,
|
|
batch_size: Optional[int] = None,
|
|
) -> Dict[str, Any]:
|
|
"""
|
|
Compare sidecar metadata size values against actual local model file sizes.
|
|
"""
|
|
if models is None:
|
|
from .scanner import get_model_files
|
|
|
|
models = get_model_files(force_rescan=force_rescan)
|
|
|
|
audit_models = dedupe_models(models)
|
|
total_models = len(audit_models)
|
|
resolved_worker_count, cpu_count = resolve_worker_count(
|
|
total_models,
|
|
worker_count,
|
|
default_worker_multiplier=4,
|
|
default_worker_limit=MAX_AUDIT_WORKERS,
|
|
cpu_count_fallback=4,
|
|
minimum_workers=MIN_AUDIT_WORKERS,
|
|
maximum_workers=MAX_AUDIT_WORKERS,
|
|
empty_total_workers=1,
|
|
)
|
|
resolved_batch_size = _get_batch_size(total_models, resolved_worker_count, batch_size)
|
|
batches = _make_batches(audit_models, resolved_batch_size)
|
|
active_worker_count = min(resolved_worker_count, len(batches) or 1)
|
|
aggregate = _empty_audit_counts()
|
|
|
|
if batches:
|
|
if active_worker_count <= 1:
|
|
for batch in batches:
|
|
merge_counted_payload(
|
|
aggregate,
|
|
_audit_model_batch(batch),
|
|
numeric_keys=_AUDIT_COUNT_KEYS,
|
|
list_keys=_AUDIT_LIST_KEYS,
|
|
)
|
|
else:
|
|
with ThreadPoolExecutor(max_workers=active_worker_count) as executor:
|
|
futures = [executor.submit(_audit_model_batch, batch) for batch in batches]
|
|
for future in as_completed(futures):
|
|
merge_counted_payload(
|
|
aggregate,
|
|
future.result(),
|
|
numeric_keys=_AUDIT_COUNT_KEYS,
|
|
list_keys=_AUDIT_LIST_KEYS,
|
|
)
|
|
|
|
mismatches = aggregate["mismatches"]
|
|
errors = aggregate["errors"]
|
|
|
|
mismatches.sort(
|
|
key=lambda item: (
|
|
str(item.get("category") or "").lower(),
|
|
str(item.get("relative_path") or item.get("filename") or "").lower(),
|
|
str(item.get("metadata_path") or "").lower(),
|
|
)
|
|
)
|
|
|
|
result = {
|
|
"success": True,
|
|
"cpu_count": cpu_count,
|
|
"worker_count": active_worker_count,
|
|
"batch_size": resolved_batch_size,
|
|
"batch_count": len(batches),
|
|
"scanned_models": aggregate["scanned_models"],
|
|
"metadata_files": aggregate["metadata_files"],
|
|
"checked_metadata": aggregate["checked_metadata"],
|
|
"missing_metadata": aggregate["missing_metadata"],
|
|
"missing_size": aggregate["missing_size"],
|
|
"invalid_metadata": aggregate["invalid_metadata"],
|
|
"skipped_directories": aggregate["skipped_directories"],
|
|
"skipped_non_model_files": aggregate["skipped_non_model_files"],
|
|
"error_count": len(errors),
|
|
"errors": errors[:50],
|
|
"mismatch_count": len(mismatches),
|
|
"mismatches": mismatches,
|
|
}
|
|
|
|
log.info(
|
|
"Metadata size audit finished: "
|
|
f"models={aggregate['scanned_models']}, metadata={aggregate['metadata_files']}, "
|
|
f"size_values={aggregate['checked_metadata']}, mismatches={len(mismatches)}, "
|
|
f"missing_size={aggregate['missing_size']}, errors={len(errors)}, "
|
|
f"workers={active_worker_count}, batches={len(batches)}, batch_size={resolved_batch_size}"
|
|
)
|
|
return result
|