382 lines
16 KiB
Python
382 lines
16 KiB
Python
import hashlib
|
|
import json
|
|
import os
|
|
import tempfile
|
|
import unittest
|
|
|
|
from core.metadata_builder import (
|
|
METADATA_BUILD_MODE_CALCULATE_FRESH,
|
|
METADATA_BUILD_MODE_IMPORT_EXISTING,
|
|
build_missing_local_metadata,
|
|
normalize_metadata_build_mode,
|
|
)
|
|
from core.path_utils import get_model_resolver_sidecar_path
|
|
|
|
|
|
class MetadataBuilderTests(unittest.TestCase):
|
|
def _write_model(self, directory, filename, content):
|
|
path = os.path.join(directory, filename)
|
|
with open(path, "wb") as handle:
|
|
handle.write(content)
|
|
return path
|
|
|
|
def _write_fake_safetensors(self, directory, filename, metadata):
|
|
header = json.dumps({"__metadata__": metadata}).encode("utf-8")
|
|
content = len(header).to_bytes(8, "little") + header
|
|
return self._write_model(directory, filename, content)
|
|
|
|
def _write_metadata(self, model_path, payload):
|
|
metadata_path = get_model_resolver_sidecar_path(model_path)
|
|
with open(metadata_path, "w", encoding="utf-8") as handle:
|
|
json.dump(payload, handle)
|
|
return metadata_path
|
|
|
|
def _read_metadata(self, model_path):
|
|
metadata_path = get_model_resolver_sidecar_path(model_path)
|
|
with open(metadata_path, "r", encoding="utf-8") as handle:
|
|
return metadata_path, json.load(handle)
|
|
|
|
def _model_info(self, model_path, category="loras"):
|
|
return {
|
|
"filename": os.path.basename(model_path),
|
|
"relative_path": os.path.basename(model_path),
|
|
"path": model_path,
|
|
"category": category,
|
|
"base_directory": os.path.dirname(model_path),
|
|
}
|
|
|
|
def test_creates_missing_metadata_from_local_safetensors_header(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
header_sha = "a" * 64
|
|
model_path = self._write_fake_safetensors(
|
|
tmpdir,
|
|
"header_model.safetensors",
|
|
{
|
|
"modelspec.title": "Header Model",
|
|
"modelspec.description": "Local description",
|
|
"modelspec.tags": "tag-a, tag-b",
|
|
"modelspec.trigger_phrase": "trigger-a",
|
|
"sha256": header_sha,
|
|
},
|
|
)
|
|
|
|
result = build_missing_local_metadata(
|
|
models=[self._model_info(model_path)],
|
|
metadata_mode=METADATA_BUILD_MODE_IMPORT_EXISTING,
|
|
)
|
|
metadata_path, metadata = self._read_metadata(model_path)
|
|
|
|
self.assertTrue(result["success"])
|
|
self.assertEqual(1, result["created_metadata"])
|
|
self.assertEqual(0, result["calculated_hashes"])
|
|
self.assertEqual(1, result["header_hashes"])
|
|
self.assertEqual(metadata_path, result["updated"][0]["metadata_path"])
|
|
self.assertEqual("Header Model", metadata["model_name"])
|
|
self.assertEqual(header_sha, metadata["sha256"])
|
|
self.assertEqual("safetensors_header", metadata["sha256_source"])
|
|
self.assertIn("tag-a", metadata["tags"])
|
|
self.assertIn("trigger-a", metadata["trained_words"])
|
|
self.assertIn("safetensors_header_metadata", metadata)
|
|
self.assertIn("modelspec.title", metadata["safetensors_header_metadata"]["metadata"])
|
|
|
|
def test_updates_existing_metadata_without_sha256_and_preserves_user_fields(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
content = b"local model bytes"
|
|
expected_sha = hashlib.sha256(content).hexdigest()
|
|
model_path = self._write_model(tmpdir, "needs_hash.ckpt", content)
|
|
self._write_metadata(
|
|
model_path,
|
|
{
|
|
"model_name": "Existing Name",
|
|
"notes": "keep me",
|
|
"favorite": True,
|
|
},
|
|
)
|
|
|
|
events = []
|
|
result = build_missing_local_metadata(
|
|
models=[self._model_info(model_path, "checkpoints")],
|
|
progress_callback=events.append,
|
|
)
|
|
_metadata_path, metadata = self._read_metadata(model_path)
|
|
|
|
self.assertTrue(result["success"])
|
|
self.assertEqual(1, result["updated_metadata"])
|
|
self.assertEqual(1, result["calculated_hashes"])
|
|
self.assertEqual(expected_sha, metadata["sha256"])
|
|
self.assertEqual("file", metadata["sha256_source"])
|
|
self.assertEqual("Existing Name", metadata["model_name"])
|
|
self.assertEqual("keep me", metadata["notes"])
|
|
self.assertTrue(metadata["favorite"])
|
|
self.assertTrue(any(event.get("stage") == "hashing" for event in events))
|
|
self.assertTrue(any(event.get("current_model") == "needs_hash.ckpt" for event in events))
|
|
|
|
def test_skips_existing_metadata_that_already_has_sha256(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
model_path = self._write_model(tmpdir, "complete.safetensors", b"abc")
|
|
existing = {
|
|
"model_name": "Complete",
|
|
"sha256": "b" * 64,
|
|
"notes": "do not rewrite",
|
|
}
|
|
self._write_metadata(model_path, existing)
|
|
|
|
result = build_missing_local_metadata(models=[self._model_info(model_path)])
|
|
_metadata_path, metadata = self._read_metadata(model_path)
|
|
|
|
self.assertTrue(result["success"])
|
|
self.assertEqual(0, result["created_metadata"])
|
|
self.assertEqual(0, result["updated_metadata"])
|
|
self.assertEqual(1, result["skipped_complete"])
|
|
self.assertEqual(existing, metadata)
|
|
|
|
def test_same_stem_models_receive_extension_qualified_sidecars(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
bin_content = b"bin adapter weights"
|
|
safetensors_content = b"safetensors adapter weights"
|
|
bin_path = self._write_model(tmpdir, "ip-adapter_sd15.bin", bin_content)
|
|
safetensors_path = self._write_model(
|
|
tmpdir,
|
|
"ip-adapter_sd15.safetensors",
|
|
safetensors_content,
|
|
)
|
|
|
|
result = build_missing_local_metadata(
|
|
models=[
|
|
self._model_info(bin_path),
|
|
self._model_info(safetensors_path),
|
|
],
|
|
worker_count=1,
|
|
metadata_mode=METADATA_BUILD_MODE_CALCULATE_FRESH,
|
|
)
|
|
|
|
bin_metadata_path = get_model_resolver_sidecar_path(bin_path)
|
|
safetensors_metadata_path = get_model_resolver_sidecar_path(
|
|
safetensors_path
|
|
)
|
|
legacy_metadata_path = os.path.splitext(bin_path)[0] + ".metadata.json"
|
|
with open(bin_metadata_path, "r", encoding="utf-8") as handle:
|
|
bin_metadata = json.load(handle)
|
|
with open(safetensors_metadata_path, "r", encoding="utf-8") as handle:
|
|
safetensors_metadata = json.load(handle)
|
|
legacy_metadata_exists = os.path.exists(legacy_metadata_path)
|
|
|
|
self.assertTrue(result["success"])
|
|
self.assertEqual(2, result["created_metadata"])
|
|
self.assertFalse(legacy_metadata_exists)
|
|
self.assertEqual("ip-adapter_sd15.bin", bin_metadata["filename"])
|
|
self.assertEqual(
|
|
hashlib.sha256(bin_content).hexdigest(),
|
|
bin_metadata["sha256"],
|
|
)
|
|
self.assertEqual(
|
|
"ip-adapter_sd15.safetensors",
|
|
safetensors_metadata["filename"],
|
|
)
|
|
self.assertEqual(
|
|
hashlib.sha256(safetensors_content).hexdigest(),
|
|
safetensors_metadata["sha256"],
|
|
)
|
|
self.assertEqual("comfyui-model-resolver", bin_metadata["managed_by"])
|
|
self.assertNotIn("favorite", bin_metadata)
|
|
self.assertNotIn("notes", bin_metadata)
|
|
self.assertEqual(
|
|
METADATA_BUILD_MODE_CALCULATE_FRESH,
|
|
bin_metadata["metadata_build_mode"],
|
|
)
|
|
|
|
def test_external_metadata_remains_untouched_and_is_not_copied(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
bin_content = b"bin"
|
|
bin_path = self._write_model(tmpdir, "adapter.bin", bin_content)
|
|
safetensors_path = self._write_model(
|
|
tmpdir,
|
|
"adapter.safetensors",
|
|
b"safe",
|
|
)
|
|
legacy_sha = "c" * 64
|
|
legacy_path = os.path.splitext(bin_path)[0] + ".metadata.json"
|
|
legacy_payload = {
|
|
"filename": "adapter.bin",
|
|
"file_path": bin_path,
|
|
"sha256": legacy_sha,
|
|
"notes": "preserve external user details",
|
|
"favorite": True,
|
|
}
|
|
with open(legacy_path, "w", encoding="utf-8") as handle:
|
|
json.dump(legacy_payload, handle)
|
|
|
|
result = build_missing_local_metadata(
|
|
models=[
|
|
self._model_info(bin_path),
|
|
self._model_info(safetensors_path),
|
|
],
|
|
worker_count=1,
|
|
metadata_mode=METADATA_BUILD_MODE_CALCULATE_FRESH,
|
|
)
|
|
|
|
with open(
|
|
get_model_resolver_sidecar_path(bin_path),
|
|
"r",
|
|
encoding="utf-8",
|
|
) as handle:
|
|
bin_metadata = json.load(handle)
|
|
with open(
|
|
get_model_resolver_sidecar_path(safetensors_path),
|
|
"r",
|
|
encoding="utf-8",
|
|
) as handle:
|
|
safetensors_metadata = json.load(handle)
|
|
with open(legacy_path, "r", encoding="utf-8") as handle:
|
|
legacy_after = json.load(handle)
|
|
|
|
self.assertTrue(result["success"])
|
|
self.assertEqual(legacy_payload, legacy_after)
|
|
self.assertEqual(
|
|
hashlib.sha256(bin_content).hexdigest(),
|
|
bin_metadata["sha256"],
|
|
)
|
|
self.assertNotEqual(legacy_sha, bin_metadata["sha256"])
|
|
self.assertNotIn("notes", bin_metadata)
|
|
self.assertNotIn("favorite", bin_metadata)
|
|
self.assertEqual("adapter.bin", bin_metadata["filename"])
|
|
self.assertEqual("adapter.safetensors", safetensors_metadata["filename"])
|
|
self.assertEqual(
|
|
hashlib.sha256(b"safe").hexdigest(),
|
|
safetensors_metadata["sha256"],
|
|
)
|
|
|
|
def test_import_mode_uses_external_metadata_without_modifying_it(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
model_path = self._write_model(
|
|
tmpdir,
|
|
"imported.safetensors",
|
|
b"local model bytes",
|
|
)
|
|
external_path = os.path.join(tmpdir, "imported.metadata.json")
|
|
external_sha = "d" * 64
|
|
external_payload = {
|
|
"filename": "imported.safetensors",
|
|
"file_path": model_path,
|
|
"model_name": "Imported model name",
|
|
"base_model": "SDXL 1.0",
|
|
"tags": ["external-tag"],
|
|
"sha256": external_sha,
|
|
"hash_status": "completed",
|
|
"notes": "user-owned note",
|
|
"favorite": True,
|
|
}
|
|
with open(external_path, "w", encoding="utf-8") as handle:
|
|
json.dump(external_payload, handle)
|
|
|
|
result = build_missing_local_metadata(
|
|
models=[self._model_info(model_path)],
|
|
worker_count=1,
|
|
metadata_mode=METADATA_BUILD_MODE_IMPORT_EXISTING,
|
|
)
|
|
_metadata_path, metadata = self._read_metadata(model_path)
|
|
with open(external_path, "r", encoding="utf-8") as handle:
|
|
external_after = json.load(handle)
|
|
|
|
self.assertTrue(result["success"])
|
|
self.assertEqual(1, result["created_metadata"])
|
|
self.assertEqual(0, result["calculated_hashes"])
|
|
self.assertEqual(external_payload, external_after)
|
|
self.assertEqual(external_sha, metadata["sha256"])
|
|
self.assertEqual("external_metadata", metadata["sha256_source"])
|
|
self.assertEqual("Imported model name", metadata["model_name"])
|
|
self.assertEqual("SDXL 1.0", metadata["base_model"])
|
|
self.assertEqual(["external-tag"], metadata["tags"])
|
|
self.assertNotIn("notes", metadata)
|
|
self.assertNotIn("favorite", metadata)
|
|
self.assertEqual(
|
|
external_path.replace(os.sep, "/"),
|
|
metadata["imported_from"],
|
|
)
|
|
self.assertEqual(
|
|
METADATA_BUILD_MODE_IMPORT_EXISTING,
|
|
metadata["metadata_build_mode"],
|
|
)
|
|
|
|
def test_import_mode_calculates_hash_when_external_hash_is_pending(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
content = b"calculate missing external hash"
|
|
model_path = self._write_model(
|
|
tmpdir,
|
|
"pending.safetensors",
|
|
content,
|
|
)
|
|
external_path = os.path.join(tmpdir, "pending.metadata.json")
|
|
with open(external_path, "w", encoding="utf-8") as handle:
|
|
json.dump(
|
|
{
|
|
"filename": "pending.safetensors",
|
|
"sha256": "e" * 64,
|
|
"hash_status": "pending",
|
|
},
|
|
handle,
|
|
)
|
|
|
|
result = build_missing_local_metadata(
|
|
models=[self._model_info(model_path)],
|
|
worker_count=1,
|
|
metadata_mode=METADATA_BUILD_MODE_IMPORT_EXISTING,
|
|
)
|
|
_metadata_path, metadata = self._read_metadata(model_path)
|
|
|
|
self.assertEqual(1, result["calculated_hashes"])
|
|
self.assertEqual(hashlib.sha256(content).hexdigest(), metadata["sha256"])
|
|
self.assertEqual("file", metadata["sha256_source"])
|
|
|
|
def test_invalid_metadata_mode_falls_back_to_fresh_calculation(self):
|
|
self.assertEqual(
|
|
METADATA_BUILD_MODE_CALCULATE_FRESH,
|
|
normalize_metadata_build_mode("unsupported"),
|
|
)
|
|
|
|
def test_exact_sidecar_remains_canonical_after_sibling_is_removed(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
bin_path = self._write_model(tmpdir, "adapter.bin", b"bin")
|
|
safetensors_path = self._write_model(
|
|
tmpdir,
|
|
"adapter.safetensors",
|
|
b"safe",
|
|
)
|
|
exact_metadata_path = get_model_resolver_sidecar_path(bin_path)
|
|
with open(exact_metadata_path, "w", encoding="utf-8") as handle:
|
|
json.dump({"filename": "adapter.bin", "size": 3}, handle)
|
|
|
|
os.remove(safetensors_path)
|
|
selected_path = get_model_resolver_sidecar_path(bin_path)
|
|
|
|
self.assertEqual(exact_metadata_path, selected_path)
|
|
|
|
def test_builds_missing_metadata_with_multiple_workers(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
model_paths = [
|
|
self._write_model(tmpdir, f"parallel_{index}.ckpt", f"model-{index}".encode("utf-8"))
|
|
for index in range(4)
|
|
]
|
|
events = []
|
|
|
|
result = build_missing_local_metadata(
|
|
models=[self._model_info(path, "checkpoints") for path in model_paths],
|
|
worker_count=2,
|
|
progress_callback=events.append,
|
|
)
|
|
|
|
metadata_payloads = [self._read_metadata(path)[1] for path in model_paths]
|
|
|
|
self.assertTrue(result["success"])
|
|
self.assertEqual(2, result["worker_count"])
|
|
self.assertGreaterEqual(result["cpu_count"], 1)
|
|
self.assertEqual(4, result["created_metadata"])
|
|
self.assertEqual(4, result["calculated_hashes"])
|
|
self.assertEqual(4, len(metadata_payloads))
|
|
self.assertTrue(all(payload.get("sha256") for payload in metadata_payloads))
|
|
self.assertTrue(any(event.get("worker_count") == 2 for event in events))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|