from datetime import datetime import os import json import logging import time from typing import Any, Dict, Optional, Type, Union, cast from .models import BaseModelMetadata, CheckpointMetadata, EmbeddingMetadata, LoraMetadata from .file_utils import normalize_path, find_preview_file, calculate_sha256, calculate_autov3 from .lora_metadata import extract_lora_metadata, extract_checkpoint_metadata logger = logging.getLogger(__name__) class MetadataManager: """ Centralized manager for all metadata operations. This class is responsible for: 1. Loading metadata safely with fallback mechanisms 2. Saving metadata with atomic operations 3. Creating default metadata for models 4. Handling unknown fields gracefully """ @staticmethod async def load_metadata(file_path: str, model_class: Type[BaseModelMetadata] = LoraMetadata) -> tuple[Optional[BaseModelMetadata], bool]: """ Load metadata safely. Returns: tuple: (metadata, should_skip) - metadata: BaseModelMetadata instance or None - should_skip: True if corrupted metadata file exists and model should be skipped """ metadata_path = f"{os.path.splitext(file_path)[0]}.metadata.json" # Check if metadata file exists if not os.path.exists(metadata_path): return None, False try: with open(metadata_path, 'r', encoding='utf-8') as f: data = json.load(f) # Create model instance metadata = model_class.from_dict(data) # Normalize paths await MetadataManager._normalize_metadata_paths(metadata, file_path) return metadata, False except (json.JSONDecodeError, Exception) as e: error_type = "Invalid JSON" if isinstance(e, json.JSONDecodeError) else "Parse error" logger.error(f"{error_type} in metadata file: {metadata_path}. Error: {str(e)}. Skipping model to preserve existing data.") return None, True # should_skip = True @staticmethod def _fill_local_file_facts(payload: Dict[str, Any], file_path: str) -> None: """Fill missing local file facts (``file_name``/``size``/``modified``) from disk. These three fields are part of the required metadata schema but describe the local file, not remote metadata. Payloads rebuilt by the self-heal refresh flow (sidecar deleted, then recreated from remote data) lack them, which makes the recreated sidecar unparseable by ``BaseModelMetadata.from_dict`` and causes the scanner to skip the model. Fill them from the actual file whenever absent. """ if not file_path: return if payload.get("file_name") and "size" in payload and "modified" in payload: return try: stat_result = os.stat(file_path) except OSError: return if not payload.get("file_name"): payload["file_name"] = os.path.splitext(os.path.basename(file_path))[0] if "size" not in payload: payload["size"] = stat_result.st_size if "modified" not in payload: payload["modified"] = stat_result.st_mtime @staticmethod async def load_metadata_payload(file_path: str) -> Dict[str, Any]: """ Load metadata and return it as a dictionary, including any unknown fields. Falls back to reading the raw JSON file if parsing into a model class fails. """ payload: Dict[str, Any] = {} metadata_obj, should_skip = await MetadataManager.load_metadata(file_path) if metadata_obj: payload = metadata_obj.to_dict() unknown_fields = getattr(metadata_obj, "_unknown_fields", None) if isinstance(unknown_fields, dict): payload.update(unknown_fields) else: if not should_skip: metadata_path = ( file_path if file_path.endswith(".metadata.json") else f"{os.path.splitext(file_path)[0]}.metadata.json" ) if os.path.exists(metadata_path): try: with open(metadata_path, "r", encoding="utf-8") as handle: raw = json.load(handle) if isinstance(raw, dict): payload = raw except json.JSONDecodeError: logger.warning( "Failed to parse metadata file %s while loading payload", metadata_path, ) except Exception as exc: # pragma: no cover - defensive logging logger.warning("Failed to read metadata file %s: %s", metadata_path, exc) if not isinstance(payload, dict): payload = {} if file_path: payload.setdefault("file_path", normalize_path(file_path)) # Required schema fields that are local filesystem facts. When the # sidecar is missing (e.g. deleted and being recreated by the # self-heal refresh flow), restore them so the recreated sidecar # and cache entries stay parseable. MetadataManager._fill_local_file_facts(payload, file_path) return payload @staticmethod async def hydrate_model_data(model_data: Dict[str, Any]) -> Dict[str, Any]: """ Replace the provided model data with the authoritative payload from disk. Preserves the cached folder entry if present. When the sidecar is missing entirely (self-heal after manual deletion), the disk payload is nearly empty and the cache snapshot is the only source for the schema fields required by ``BaseModelMetadata.from_dict`` (file_name/model_name/size/modified/sha256/base_model/preview_url), so every missing key is restored from it to keep any recreated sidecar parseable and avoid data loss on failed refreshes. When the sidecar exists, disk data stays authoritative and no cache key is resurrected. """ file_path = model_data.get("file_path") if not file_path: return model_data folder = model_data.get("folder") metadata_path = f"{os.path.splitext(file_path)[0]}.metadata.json" sidecar_exists = os.path.exists(metadata_path) cached = model_data.copy() payload = await MetadataManager.load_metadata_payload(file_path) if folder is not None: payload["folder"] = folder model_data.clear() model_data.update(payload) if not sidecar_exists: for key, value in cached.items(): if key not in model_data and key != "folder": model_data[key] = value # The schema defines `modified` as the import timestamp; keep the # cache's value over the stat-derived fallback from # load_metadata_payload. if "modified" in cached: model_data["modified"] = cached["modified"] # file_name/size are local file facts; prefer fresh stat values over # the possibly stale cache snapshot. MetadataManager._fill_local_file_facts(model_data, file_path) return model_data @staticmethod async def save_metadata(path: str, metadata: Union[BaseModelMetadata, Dict[str, Any]]) -> bool: """ Save metadata with atomic write operations. Args: path: Path to the model file or directly to the metadata file metadata: Metadata to save (either BaseModelMetadata object or dict) Returns: bool: Success or failure """ # Determine if the input is a metadata path or a model file path if path.endswith('.metadata.json'): metadata_path = path else: # Use existing logic for model file paths file_path = path metadata_path = f"{os.path.splitext(file_path)[0]}.metadata.json" temp_path = f"{metadata_path}.tmp" try: # Convert to dict if needed if isinstance(metadata, BaseModelMetadata): metadata_dict = metadata.to_dict() # Preserve unknown fields if present if hasattr(metadata, '_unknown_fields'): metadata_dict.update(metadata._unknown_fields) else: metadata_dict = metadata.copy() # Normalize paths if 'file_path' in metadata_dict: metadata_dict['file_path'] = normalize_path(metadata_dict['file_path']) if 'preview_url' in metadata_dict: metadata_dict['preview_url'] = normalize_path(metadata_dict['preview_url']) # Local file facts are required schema fields; fill them when a # payload rebuilt without them (e.g. self-heal) is being persisted. if metadata_dict.get("file_path"): MetadataManager._fill_local_file_facts(metadata_dict, metadata_dict["file_path"]) # Write to temporary file first with open(temp_path, 'w', encoding='utf-8') as f: json.dump(metadata_dict, f, indent=2, ensure_ascii=False) # Atomic rename operation os.replace(temp_path, metadata_path) return True except Exception as e: logger.error(f"Error saving metadata to {metadata_path}: {str(e)}") # Clean up temporary file if it exists if os.path.exists(temp_path): try: os.remove(temp_path) except: pass return False @staticmethod async def create_default_metadata(file_path: str, model_class: Type[BaseModelMetadata] = LoraMetadata) -> Optional[BaseModelMetadata]: """ Create basic metadata structure for a model file. This replaces the old get_file_info function with a more appropriately named method. Args: file_path: Path to the model file model_class: Class to instantiate Returns: BaseModelMetadata instance or None if file doesn't exist """ # First check if file actually exists and resolve symlinks try: real_path = os.path.realpath(file_path) if not os.path.exists(real_path): return None except Exception as e: logger.error(f"Error checking file existence for {file_path}: {e}") return None try: base_name = os.path.splitext(os.path.basename(file_path))[0] dir_path = os.path.dirname(file_path) # Find preview image preview_url = find_preview_file(base_name, dir_path) # Calculate file hash start_hash_time = time.perf_counter() logger.debug(f"Calculating SHA256 hash for {real_path}...") sha256 = await calculate_sha256(real_path) hash_duration = time.perf_counter() - start_hash_time logger.info(f"SHA256 hash calculated for {real_path} in {hash_duration:.3f}s") # AutoV3 reads only the safetensors header, so it is cheap even for # large files. At creation time we always know the checked state: # store "" when no recognized hash is embedded (checked-unavailable). autov3 = calculate_autov3(real_path) # Create instance based on model type if model_class.__name__ == "CheckpointMetadata": metadata = cast(Type[CheckpointMetadata], model_class)( file_name=base_name, model_name=base_name, file_path=normalize_path(file_path), size=os.path.getsize(real_path), modified=datetime.now().timestamp(), sha256=sha256, base_model="Unknown", preview_url=normalize_path(preview_url), tags=[], modelDescription="", sub_type="checkpoint", from_civitai=True ) elif model_class.__name__ == "EmbeddingMetadata": metadata = cast(Type[EmbeddingMetadata], model_class)( file_name=base_name, model_name=base_name, file_path=normalize_path(file_path), size=os.path.getsize(real_path), modified=datetime.now().timestamp(), sha256=sha256, base_model="Unknown", preview_url=normalize_path(preview_url), tags=[], modelDescription="", sub_type="embedding", from_civitai=True ) else: # Default to LoraMetadata metadata = cast(Type[LoraMetadata], model_class)( file_name=base_name, model_name=base_name, file_path=normalize_path(file_path), size=os.path.getsize(real_path), modified=datetime.now().timestamp(), sha256=sha256, base_model="Unknown", preview_url=normalize_path(preview_url), tags=[], modelDescription="", from_civitai=True, usage_tips="{}" ) # Record the AutoV3 state explicitly ("" = checked, no value). metadata.autov3 = autov3 or "" # Try to extract model-specific metadata # await MetadataManager._enrich_metadata(metadata, real_path) # Save the created metadata logger.info(f"Creating new .metadata.json for {file_path} (Reason: No existing metadata found)") await MetadataManager.save_metadata(file_path, metadata) return metadata except Exception as e: logger.error(f"Error creating default metadata for {file_path}: {e}") return None @staticmethod async def _enrich_metadata(metadata: BaseModelMetadata, file_path: str) -> None: """ Enrich metadata with model-specific information Args: metadata: Metadata to enrich file_path: Path to the model file """ try: if metadata.__class__.__name__ == "LoraMetadata": model_info = await extract_lora_metadata(file_path) metadata.base_model = model_info['base_model'] # elif metadata.__class__.__name__ == "CheckpointMetadata": # model_info = await extract_checkpoint_metadata(file_path) # metadata.base_model = model_info['base_model'] # if 'model_type' in model_info: # metadata.model_type = model_info['model_type'] except Exception as e: logger.error(f"Error enriching metadata: {str(e)}") @staticmethod async def _normalize_metadata_paths(metadata: BaseModelMetadata, file_path: str) -> None: """ Normalize paths in metadata object Args: metadata: Metadata object to update file_path: Current file path for the model """ need_update = False # Check if file_name matches the actual file name base_name = os.path.splitext(os.path.basename(file_path))[0] if metadata.file_name != base_name: metadata.file_name = base_name need_update = True # Check if file path is different from what's in metadata if normalize_path(file_path) != metadata.file_path: metadata.file_path = normalize_path(file_path) need_update = True # Check if preview exists at the current location preview_url = metadata.preview_url if preview_url: # Get directory parts of both paths file_dir = os.path.dirname(file_path) preview_dir = os.path.dirname(preview_url) # Update preview if it doesn't exist OR if model and preview are in different directories if not os.path.exists(preview_url) or file_dir != preview_dir: base_name = os.path.splitext(os.path.basename(file_path))[0] dir_path = os.path.dirname(file_path) new_preview_url = find_preview_file(base_name, dir_path) if new_preview_url: metadata.preview_url = normalize_path(new_preview_url) need_update = True # If path attributes were changed, save the metadata back to disk if need_update: await MetadataManager.save_metadata(file_path, metadata)