From 14bb53aa686f78d7a3be3ef40ee05f5b310fef33 Mon Sep 17 00:00:00 2001 From: LuqP2 Date: Sat, 8 Aug 2026 21:07:41 -0300 Subject: [PATCH] fix: expose tool attribution in JPEG and WebP --- ...ivitai-auto-tool-attribution-2026-08-08.md | 6 +-- metadata_utils_impl.py | 40 ++++++++++++------- tests/test_metadata_utils.py | 25 ++++++++++++ 3 files changed, 54 insertions(+), 17 deletions(-) diff --git a/docs/research/civitai-auto-tool-attribution-2026-08-08.md b/docs/research/civitai-auto-tool-attribution-2026-08-08.md index 937d4e6..5cb47a5 100644 --- a/docs/research/civitai-auto-tool-attribution-2026-08-08.md +++ b/docs/research/civitai-auto-tool-attribution-2026-08-08.md @@ -28,16 +28,16 @@ The workflow and prompt payloads are preserved and the active save node is rewri ## Metadata contract -New PNG outputs add two generic `tEXt` chunks: +New outputs expose two generic metadata fields: PNG uses `tEXt` chunks, while JPEG and WebP use XMP properties: ```text engine = ComfyUI tools = ["Image MetaHub"] ``` -`tools` is a JSON array, not a delimiter-based string. This keeps engine attribution semantically distinct from additional creation, processing, editing, or saving tools. It also survives `ExifReader` independently of which generation parser wins. +`tools` is a JSON array, not a delimiter-based string. This keeps engine attribution semantically distinct from additional creation, processing, editing, or saving tools. The fields survive `ExifReader` independently of which generation parser wins. -The same values are included in `imagemetahub_data` for MetaHub consumers, while the separate chunks are the upstream interchange contract. Existing `parameters`, `workflow`, `prompt`, and `imagemetahub_data` fields are unchanged. +The same values are included in `imagemetahub_data` for MetaHub consumers, while the separate PNG/XMP fields are the upstream interchange contract. Existing `parameters`, `workflow`, `prompt`, and `imagemetahub_data` fields are unchanged. ## Upstream behavior diff --git a/metadata_utils_impl.py b/metadata_utils_impl.py index 9f9a381..ba5cd50 100644 --- a/metadata_utils_impl.py +++ b/metadata_utils_impl.py @@ -614,10 +614,10 @@ def _serialize_metadata_json_ascii(value: Any) -> Optional[str]: def build_comfyui_xmp_packet(imh_metadata: dict) -> Optional[bytes]: workflow_text, prompt_text = _get_workflow_prompt_texts(imh_metadata) - if not workflow_text and not prompt_text: - return None - - fields = [] + fields = [ + f"{xml_escape(METAHUB_GENERATION_ENGINE)}", + f"{xml_escape(json.dumps(METAHUB_ADDITIONAL_TOOLS))}", + ] if workflow_text: fields.append(f"{xml_escape(workflow_text)}") if prompt_text: @@ -899,6 +899,25 @@ def _build_exif_bytes(a1111_metadata: str, imh_metadata: dict) -> Optional[bytes return None +def _inject_jpeg_xmp(image_path: str, xmp_bytes: bytes) -> bool: + try: + xmp_payload = b"http://ns.adobe.com/xap/1.0/\x00" + xmp_bytes + segment_length = len(xmp_payload) + 2 + if segment_length > 0xFFFF: + return False + + path = Path(image_path) + jpeg_data = path.read_bytes() + if not jpeg_data.startswith(b"\xff\xd8"): + return False + + xmp_segment = b"\xff\xe1" + segment_length.to_bytes(2, "big") + xmp_payload + path.write_bytes(jpeg_data[:2] + xmp_segment + jpeg_data[2:]) + return True + except Exception: + return False + + def save_jpeg_with_metadata( image: Image.Image, image_path: str, @@ -919,16 +938,9 @@ def save_jpeg_with_metadata( save_kwargs["exif"] = exif_bytes elif a1111_metadata: save_kwargs["comment"] = make_civitai_safe_text(a1111_metadata).encode("utf-8", errors="replace") - if xmp_bytes: - save_kwargs["xmp"] = xmp_bytes - try: - image.save(image_path, **save_kwargs) - except TypeError: - if "xmp" in save_kwargs: - save_kwargs.pop("xmp", None) - image.save(image_path, **save_kwargs) - else: - raise + image.save(image_path, **save_kwargs) + if xmp_bytes and not _inject_jpeg_xmp(image_path, xmp_bytes): + print(f"[MetaHub] Warning: JPEG XMP metadata could not be embedded in {image_path}") except Exception as e: print(f"[MetaHub] Warning: JPEG metadata save failed for {image_path}: {e}") image.save(image_path, "JPEG", quality=quality) diff --git a/tests/test_metadata_utils.py b/tests/test_metadata_utils.py index 0621d87..991c291 100644 --- a/tests/test_metadata_utils.py +++ b/tests/test_metadata_utils.py @@ -6,6 +6,7 @@ import copy from metadata_utils import ( build_a1111_metadata, + build_comfyui_xmp_packet, build_imh_metadata, build_metadata_sources, build_metadata_status, @@ -18,7 +19,9 @@ from metadata_utils import ( get_workflow_json, make_civitai_safe_text, resolve_filename_pattern_path, + save_jpeg_with_metadata, save_png_with_metadata, + save_webp_with_metadata, build_ui_preview, ) from PIL import Image @@ -63,6 +66,28 @@ def _save_simple_png(path: Path) -> None: image.save(path, "PNG") +def test_build_comfyui_xmp_packet_includes_tool_attribution(): + packet = build_comfyui_xmp_packet({}) + + assert packet is not None + assert b"ComfyUI" in packet + assert b'["Image MetaHub"]' in packet + + +def test_jpeg_and_webp_saves_embed_tool_attribution_xmp(tmp_path): + image = Image.new("RGB", (2, 2), color=(0, 0, 0)) + outputs = [ + (tmp_path / "attribution.jpg", save_jpeg_with_metadata), + (tmp_path / "attribution.webp", save_webp_with_metadata), + ] + + for path, save in outputs: + save(image, str(path), 95, "params", {}) + data = path.read_bytes() + assert b"ComfyUI" in data + assert b'["Image MetaHub"]' in data + + def test_output_path_relative_resolves_under_comfy_output(tmp_path, monkeypatch): class FakeFolderPaths: @staticmethod