mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-08-09 07:20:15 -03:00
feat(metadata): add CivitAI AutoV3 hash support across all storage layers
- Three-state autov3 field (not-checked / checked-unavailable / 12-hex value) in .metadata.json sidecars, in-memory ModelHashIndex, and SQLite (models.autov3 column + autov3_index table) with column-presence migration - Background self-terminating backfill for legacy rows: per-model-type concurrency guard, executor-offloaded I/O, Civitai-first resolution (SHA256-matched version file) falling back to the embedded safetensors header hash - Civitai-first propagation on metadata refresh, scan, and download paths; reject the empty-string SHA256 placeholder and strip OneTrainer 0x prefix - List API hash filters and hash index lookups accept 12-char AutoV3 - Cap safetensors header reads at 64 MiB to prevent crafted-file allocation - Prevent stale AutoV3 mappings on file replacement while preserving them on same-file re-registration (lazy-hash completion)
This commit is contained in:
@@ -1,10 +1,18 @@
|
||||
import json
|
||||
import struct
|
||||
import pytest
|
||||
|
||||
from py.utils.metadata_manager import MetadataManager
|
||||
from py.utils.models import BaseModelMetadata
|
||||
|
||||
|
||||
def _write_safetensors(path, metadata, payload=b"payload-bytes"):
|
||||
"""Write a minimal real safetensors file: 8-byte little-endian header
|
||||
length, a JSON header containing ``__metadata__``, then arbitrary payload."""
|
||||
header = json.dumps({"__metadata__": metadata}).encode("utf-8")
|
||||
path.write_bytes(struct.pack("<Q", len(header)) + header + payload)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_base_model_metadata_sets_empty_civitai_dict():
|
||||
metadata = BaseModelMetadata(
|
||||
@@ -35,3 +43,35 @@ async def test_create_default_metadata_uses_empty_civitai(tmp_path):
|
||||
payload = json.loads(metadata_path.read_text(encoding="utf-8"))
|
||||
|
||||
assert payload.get("civitai") == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_default_metadata_writes_autov3_from_safetensors(tmp_path):
|
||||
model_path = tmp_path / "example.safetensors"
|
||||
_write_safetensors(model_path, {"sshs_model_hash": "abcdef1234567890abcdef"})
|
||||
|
||||
metadata = await MetadataManager.create_default_metadata(str(model_path))
|
||||
|
||||
assert metadata is not None
|
||||
assert metadata.autov3 == "abcdef123456"
|
||||
|
||||
metadata_path = model_path.with_suffix(".metadata.json")
|
||||
payload = json.loads(metadata_path.read_text(encoding="utf-8"))
|
||||
|
||||
assert payload.get("autov3") == "abcdef123456"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_default_metadata_writes_autov3_null_for_non_safetensors(tmp_path):
|
||||
model_path = tmp_path / "example.safetensors"
|
||||
model_path.write_bytes(b"stub")
|
||||
|
||||
metadata = await MetadataManager.create_default_metadata(str(model_path))
|
||||
|
||||
assert metadata is not None
|
||||
assert metadata.autov3 == ""
|
||||
|
||||
metadata_path = model_path.with_suffix(".metadata.json")
|
||||
payload = json.loads(metadata_path.read_text(encoding="utf-8"))
|
||||
|
||||
assert payload.get("autov3") is None
|
||||
|
||||
Reference in New Issue
Block a user