diff --git a/py/utils/metadata_manager.py b/py/utils/metadata_manager.py index baac249e..43cfe976 100644 --- a/py/utils/metadata_manager.py +++ b/py/utils/metadata_manager.py @@ -55,6 +55,32 @@ class MetadataManager: 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]: """ @@ -96,6 +122,11 @@ class MetadataManager: 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 @@ -104,6 +135,14 @@ class MetadataManager: """ 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") @@ -111,12 +150,29 @@ class MetadataManager: 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 @@ -155,7 +211,12 @@ class MetadataManager: 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) diff --git a/tests/utils/test_metadata_manager.py b/tests/utils/test_metadata_manager.py new file mode 100644 index 00000000..0be417fe --- /dev/null +++ b/tests/utils/test_metadata_manager.py @@ -0,0 +1,214 @@ +from __future__ import annotations + +import json + +import pytest + +from py.utils.metadata_manager import MetadataManager +from py.utils.models import LoraMetadata + + +@pytest.mark.asyncio +async def test_save_metadata_fills_missing_file_facts(tmp_path) -> None: + """A payload missing file_name/size/modified is healed on write.""" + model_path = tmp_path / "MyModel.safetensors" + model_path.write_bytes(b"fake model data") + payload = {"file_path": str(model_path), "model_name": "My Model"} + + result = await MetadataManager.save_metadata(str(model_path), payload) + assert result is True + + metadata_path = tmp_path / "MyModel.metadata.json" + saved = json.loads(metadata_path.read_text(encoding="utf-8")) + stat_result = model_path.stat() + assert saved["file_name"] == "MyModel" + assert saved["size"] == stat_result.st_size + assert saved["modified"] == stat_result.st_mtime + + +@pytest.mark.asyncio +async def test_save_metadata_keeps_existing_file_facts(tmp_path) -> None: + """Existing file facts are never overwritten on write.""" + model_path = tmp_path / "MyModel.safetensors" + model_path.write_bytes(b"fake model data") + payload = { + "file_path": str(model_path), + "file_name": "CustomName", + "size": 123, + "modified": 456.0, + } + + result = await MetadataManager.save_metadata(str(model_path), payload) + assert result is True + + metadata_path = tmp_path / "MyModel.metadata.json" + saved = json.loads(metadata_path.read_text(encoding="utf-8")) + assert saved["file_name"] == "CustomName" + assert saved["size"] == 123 + assert saved["modified"] == 456.0 + + +@pytest.mark.asyncio +async def test_load_metadata_payload_restores_file_facts_when_sidecar_missing( + tmp_path, +) -> None: + """Missing sidecar: the payload is rebuilt with local file facts.""" + model_path = tmp_path / "MyModel.safetensors" + model_path.write_bytes(b"fake model data") + + payload = await MetadataManager.load_metadata_payload(str(model_path)) + + stat_result = model_path.stat() + assert payload["file_path"] == str(model_path) + assert payload["file_name"] == "MyModel" + assert payload["size"] == stat_result.st_size + assert payload["modified"] == stat_result.st_mtime + + +@pytest.mark.asyncio +async def test_hydrate_model_data_restores_required_fields_when_sidecar_missing( + tmp_path, +) -> None: + """Self-heal repro: sidecar deleted, cache entry hydrated, base fields kept.""" + model_path = tmp_path / "MyModel.safetensors" + model_path.write_bytes(b"fake model data") + model_data = { + "file_path": str(model_path), + "folder": "extra_loras", + "file_name": "MyModel", + "model_name": "My Model", + "size": 10, + "modified": 100.0, + "sha256": "abc123", + "base_model": "Illustrious", + "preview_url": "", + "civitai": {"id": 123}, + } + + await MetadataManager.hydrate_model_data(model_data) + + stat_result = model_path.stat() + assert model_data["file_name"] == "MyModel" + assert model_data["model_name"] == "My Model" + assert model_data["size"] == stat_result.st_size + # `modified` is the import timestamp per schema; the cache value wins. + assert model_data["modified"] == 100.0 + assert model_data["sha256"] == "abc123" + assert model_data["base_model"] == "Illustrious" + assert model_data["folder"] == "extra_loras" + assert model_data["civitai"] == {"id": 123} + + +@pytest.mark.asyncio +async def test_hydrate_model_data_disk_wins_when_sidecar_exists(tmp_path) -> None: + """Existing sidecar stays authoritative; no cache key is resurrected.""" + model_path = tmp_path / "MyModel.safetensors" + model_path.write_bytes(b"fake model data") + stat_result = model_path.stat() + sidecar = { + "file_path": str(model_path), + "file_name": "MyModel", + "model_name": "Disk Name", + "size": stat_result.st_size, + "modified": stat_result.st_mtime, + "sha256": "diskhash", + "base_model": "SDXL", + "preview_url": "", + } + (tmp_path / "MyModel.metadata.json").write_text( + json.dumps(sidecar), encoding="utf-8" + ) + model_data = { + "file_path": str(model_path), + "folder": "extra_loras", + "model_name": "Cache Name", + "sha256": "cachehash", + "civitai": {"id": 999}, + } + + await MetadataManager.hydrate_model_data(model_data) + + assert model_data["model_name"] == "Disk Name" + assert model_data["sha256"] == "diskhash" + # civitai is present only as the dataclass default {}; the cache's value + # ({"id": 999}) is not resurrected. + assert model_data["civitai"] == {} + assert model_data["folder"] == "extra_loras" + + +@pytest.mark.asyncio +async def test_hydrate_model_data_keeps_sha256_missing_for_caller_persist_fix( + tmp_path, +) -> None: + """Sidecar exists but lacks sha256: hydrate leaves it missing so the + caller's self-heal persist block still has work to do.""" + model_path = tmp_path / "MyModel.safetensors" + model_path.write_bytes(b"fake model data") + stat_result = model_path.stat() + sidecar = { + "file_path": str(model_path), + "file_name": "MyModel", + "model_name": "My Model", + "size": stat_result.st_size, + "modified": stat_result.st_mtime, + "base_model": "SDXL", + "preview_url": "", + # sha256 deliberately absent + } + (tmp_path / "MyModel.metadata.json").write_text( + json.dumps(sidecar), encoding="utf-8" + ) + model_data = { + "file_path": str(model_path), + "sha256": "cachehash", + "model_name": "Cache Name", + } + + await MetadataManager.hydrate_model_data(model_data) + + assert "sha256" not in model_data + + +@pytest.mark.asyncio +async def test_load_metadata_payload_tolerates_missing_model_file(tmp_path) -> None: + """A nonexistent model file must not crash payload loading.""" + missing_path = tmp_path / "Ghost.safetensors" + payload = await MetadataManager.load_metadata_payload(str(missing_path)) + assert payload["file_path"] == str(missing_path) + assert "file_name" not in payload + assert "size" not in payload + assert "modified" not in payload + + +@pytest.mark.asyncio +async def test_self_healed_sidecar_is_parseable(tmp_path) -> None: + """End-to-end repro: deleted sidecar + refresh recreates a parseable file.""" + model_path = tmp_path / "MyModel.safetensors" + model_path.write_bytes(b"fake model data") + model_data = { + "file_path": str(model_path), + "folder": "extra_loras", + "file_name": "MyModel", + "model_name": "My Model", + "size": 10, + "modified": 100.0, + "sha256": "abc123", + "base_model": "Illustrious", + "preview_url": "", + "civitai": {"id": 123}, + } + + # Simulate the self-heal flow: hydrate from (missing) sidecar, then persist. + await MetadataManager.hydrate_model_data(model_data) + data_to_save = model_data.copy() + data_to_save.pop("folder", None) + await MetadataManager.save_metadata(str(model_path), data_to_save) + + metadata, should_skip = await MetadataManager.load_metadata( + str(model_path), LoraMetadata + ) + assert should_skip is False + assert metadata is not None + assert metadata.file_name == "MyModel" + assert metadata.model_name == "My Model" + assert metadata.sha256 == "abc123"