feat: add checkpoint metadata to EXIF recipe data

Add support for storing checkpoint information in image EXIF metadata. The checkpoint data is simplified and includes fields like model ID, version, name, hash, and base model. This allows for better tracking of AI model checkpoints used in image generation workflows.
This commit is contained in:
Will Miao
2025-11-29 08:46:38 +08:00
parent 073fb3a94a
commit df93670598
2 changed files with 86 additions and 2 deletions

View File

@@ -0,0 +1,61 @@
import json
from py.utils.exif_utils import ExifUtils
def test_append_recipe_metadata_includes_checkpoint(monkeypatch, tmp_path):
captured = {}
monkeypatch.setattr(
ExifUtils, "extract_image_metadata", staticmethod(lambda _path: None)
)
def fake_update_image_metadata(image_path, metadata):
captured["path"] = image_path
captured["metadata"] = metadata
return image_path
monkeypatch.setattr(
ExifUtils, "update_image_metadata", staticmethod(fake_update_image_metadata)
)
checkpoint = {
"type": "checkpoint",
"modelId": 827184,
"modelVersionId": 2167369,
"modelName": "WAI-illustrious-SDXL",
"modelVersionName": "v15.0",
"hash": "ABC123",
"file_name": "WAI-illustrious-SDXL",
"baseModel": "Illustrious",
}
recipe_data = {
"title": "Semi-realism",
"base_model": "Illustrious",
"loras": [],
"tags": [],
"checkpoint": checkpoint,
}
image_path = tmp_path / "image.webp"
image_path.write_bytes(b"")
ExifUtils.append_recipe_metadata(str(image_path), recipe_data)
assert captured["path"] == str(image_path)
assert captured["metadata"].startswith("Recipe metadata: ")
payload = json.loads(captured["metadata"].split("Recipe metadata: ", 1)[1])
assert payload["checkpoint"] == {
"type": "checkpoint",
"modelId": 827184,
"modelVersionId": 2167369,
"modelName": "WAI-illustrious-SDXL",
"modelVersionName": "v15.0",
"hash": "abc123",
"file_name": "WAI-illustrious-SDXL",
"baseModel": "Illustrious",
}
assert payload["base_model"] == "Illustrious"