mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-03-26 15:38:52 -03:00
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:
@@ -141,6 +141,28 @@ class ExifUtils:
|
|||||||
# Remove any existing recipe metadata
|
# Remove any existing recipe metadata
|
||||||
metadata = ExifUtils.remove_recipe_metadata(metadata)
|
metadata = ExifUtils.remove_recipe_metadata(metadata)
|
||||||
|
|
||||||
|
# Prepare checkpoint data
|
||||||
|
checkpoint_data = recipe_data.get("checkpoint") or {}
|
||||||
|
simplified_checkpoint = None
|
||||||
|
if isinstance(checkpoint_data, dict) and checkpoint_data:
|
||||||
|
simplified_checkpoint = {
|
||||||
|
"type": checkpoint_data.get("type", "checkpoint"),
|
||||||
|
"modelId": checkpoint_data.get("modelId", 0),
|
||||||
|
"modelVersionId": checkpoint_data.get("modelVersionId")
|
||||||
|
or checkpoint_data.get("id", 0),
|
||||||
|
"modelName": checkpoint_data.get(
|
||||||
|
"modelName", checkpoint_data.get("name", "")
|
||||||
|
),
|
||||||
|
"modelVersionName": checkpoint_data.get(
|
||||||
|
"modelVersionName", checkpoint_data.get("version", "")
|
||||||
|
),
|
||||||
|
"hash": checkpoint_data.get("hash", "").lower()
|
||||||
|
if checkpoint_data.get("hash")
|
||||||
|
else "",
|
||||||
|
"file_name": checkpoint_data.get("file_name", ""),
|
||||||
|
"baseModel": checkpoint_data.get("baseModel", ""),
|
||||||
|
}
|
||||||
|
|
||||||
# Prepare simplified loras data
|
# Prepare simplified loras data
|
||||||
simplified_loras = []
|
simplified_loras = []
|
||||||
for lora in recipe_data.get("loras", []):
|
for lora in recipe_data.get("loras", []):
|
||||||
@@ -160,7 +182,8 @@ class ExifUtils:
|
|||||||
'base_model': recipe_data.get('base_model', ''),
|
'base_model': recipe_data.get('base_model', ''),
|
||||||
'loras': simplified_loras,
|
'loras': simplified_loras,
|
||||||
'gen_params': recipe_data.get('gen_params', {}),
|
'gen_params': recipe_data.get('gen_params', {}),
|
||||||
'tags': recipe_data.get('tags', [])
|
'tags': recipe_data.get('tags', []),
|
||||||
|
**({'checkpoint': simplified_checkpoint} if simplified_checkpoint else {})
|
||||||
}
|
}
|
||||||
|
|
||||||
# Convert to JSON string
|
# Convert to JSON string
|
||||||
|
|||||||
61
tests/utils/test_exif_utils.py
Normal file
61
tests/utils/test_exif_utils.py
Normal 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"
|
||||||
Reference in New Issue
Block a user