""" 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