nodes/StableLlama-ComfyUI-basic_data_handling
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix_too_large_EXIF

Browse Source
This commit is contained in:
StableLlama
2026-10-16 17:27:58 +02:00
parent 3e1d77627c
commit 9d53d63725
2 changed files with 189 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/basic_data_handling/path_nodes.py
+94 -21
View File
@@ -1272,6 +1272,50 @@ def build_xmp_packet(metadata_text: str) -> bytes:
return packet.encode("utf-8")
# A JPEG file carries EXIF in a single APP1 marker, whose length field is 16 bit:
# Pillow refuses to write anything bigger than this with "EXIF data is too long".
# The full workflow JSON of a non-trivial graph exceeds that easily, so JPEG EXIF
# must be trimmed (see ``_fit_jpeg_exif``) instead of failing the whole save.
_JPEG_MAX_EXIF_BYTES = 65533
def _fit_jpeg_exif(metadata_text: str, workflow_metadata: Optional[dict[str, str]] = None) -> Optional[bytes]:
"""
Build a JPEG EXIF block that fits into Pillow's single-APP1-marker limit.
A JPEG's EXIF lives in one 16-bit-sized segment, so anything beyond
``_JPEG_MAX_EXIF_BYTES`` makes Pillow raise "EXIF data is too long" and fail
the entire save. As the full-graph workflow JSON of a real graph is much
larger than that, the metadata channels are dropped one by one - least
essential first - until the block fits, and what was dropped is reported.
The image itself is always written.
Returns the EXIF bytes, or ``None`` when even the plain prompt does not fit
(nothing is embedded then).
"""
workflow_metadata = workflow_metadata or {}
# Richest first: everything, then without the (usually huge) full-graph
# "workflow" entry, then without any workflow info, then without the prompt.
without_workflow_graph = {key: value for key, value in workflow_metadata.items() if key != "workflow"}
candidates: list[tuple[str, dict[str, str], str]] = [(metadata_text, dict(workflow_metadata), "")]
if len(without_workflow_graph) != len(workflow_metadata):
candidates.append((metadata_text, without_workflow_graph, 'had to drop the full "workflow" graph'))
if workflow_metadata:
candidates.append((metadata_text, {}, "had to drop the workflow prompt info"))
if metadata_text:
candidates.append(("", {}, "had to drop the prompt metadata"))
for text, wf_meta, note in candidates:
exif = build_image_exif(text, include_description=True, workflow_metadata=wf_meta)
if exif is not None and len(exif) <= _JPEG_MAX_EXIF_BYTES:
if note:
print(f"Basic data handling: The metadata is too large for a JPEG EXIF block; {note}.")
return exif
print("Basic data handling: The metadata is too large for a JPEG EXIF block; saving the image without metadata.")
return None
def metadata_save_kwargs(metadata_text: str, fmt: str, workflow_metadata: Optional[dict[str, str]] = None) -> dict:
"""
Return the extra keyword arguments that embed ``metadata_text`` (the explicit
@@ -1288,7 +1332,7 @@ def metadata_save_kwargs(metadata_text: str, fmt: str, workflow_metadata: Option
pnginfo = build_png_info(metadata_text, workflow_metadata)
return {"pnginfo": pnginfo} if pnginfo is not None else {}
if fmt in ("jpg", "jpeg"):
exif = build_image_exif(metadata_text, include_description=True, workflow_metadata=workflow_metadata)
exif = _fit_jpeg_exif(metadata_text, workflow_metadata)
return {"exif": exif} if exif is not None else {}
if fmt in ("webp", "jxl"):
exif = build_image_exif(metadata_text, include_description=False, workflow_metadata=workflow_metadata)
@@ -1626,34 +1670,57 @@ def _has_jxl_support() -> bool:
return False
def _save_pil_image(
pil_img, path: str, fmt: str, quality: int, metadata_text: str, workflow_metadata: Optional[dict[str, str]] = None
) -> None:
"""
Write ``pil_img`` to ``path`` once, letting any error propagate. ``fmt`` is a
canonical save token (see ``_PIL_FORMAT_BY_TOKEN``). For the lossy/metadata-
capable set (png/jpg/webp/jxl) the prompt text and the workflow metadata are
embedded and ``quality`` is honoured; all other discovered formats are
written plainly.
"""
if fmt == "png":
pil_img.save(path, format="PNG", **metadata_save_kwargs(metadata_text, fmt, workflow_metadata))
elif fmt == "jpg":
# JPEG cannot carry alpha; drop it so an RGBA input still saves.
pil_img.convert("RGB").save(path, format="JPEG", quality=quality, **metadata_save_kwargs(metadata_text, fmt, workflow_metadata))
elif fmt == "webp":
pil_img.save(path, format="WEBP", quality=quality, **metadata_save_kwargs(metadata_text, fmt, workflow_metadata))
elif fmt == "jxl":
pil_img.save(path, format="JXL", quality=quality, **metadata_save_kwargs(metadata_text, fmt, workflow_metadata))
else:
if metadata_text or workflow_metadata:
print("Basic data handling: Prompt/workflow metadata is not supported for this format; skipping it.")
pil_img.save(path, format=_PIL_FORMAT_BY_TOKEN.get(fmt, fmt.upper()))
def _write_pil_image(
pil_img, path: str, fmt: str, quality: int, metadata_text: str, workflow_metadata: Optional[dict[str, str]] = None
) -> bool:
"""
Write ``pil_img`` to ``path``. ``fmt`` is a canonical save token (see
``_PIL_FORMAT_BY_TOKEN``). For the lossy/metadata-capable set (png/jpg/webp/
jxl) the prompt text and the workflow metadata are embedded and ``quality``
is honoured; all other discovered formats are written plainly. Returns True
on success.
Write ``pil_img`` to ``path``, returning True on success.
Embedded metadata must never cost the user their image: when the first
attempt fails and metadata was involved, the image is written again without
any metadata (a format/library limit on the metadata side, e.g. a container
that cannot hold the requested block) and that is reported. Formats without
metadata support are written plainly from the start.
"""
try:
if fmt == "png":
pil_img.save(path, format="PNG", **metadata_save_kwargs(metadata_text, fmt, workflow_metadata))
elif fmt == "jpg":
# JPEG cannot carry alpha; drop it so an RGBA input still saves.
pil_img.convert("RGB").save(path, format="JPEG", quality=quality, **metadata_save_kwargs(metadata_text, fmt, workflow_metadata))
elif fmt == "webp":
pil_img.save(path, format="WEBP", quality=quality, **metadata_save_kwargs(metadata_text, fmt, workflow_metadata))
elif fmt == "jxl":
pil_img.save(path, format="JXL", quality=quality, **metadata_save_kwargs(metadata_text, fmt, workflow_metadata))
else:
if metadata_text or workflow_metadata:
print("Basic data handling: Prompt/workflow metadata is not supported for this format; skipping it.")
pil_img.save(path, format=_PIL_FORMAT_BY_TOKEN.get(fmt, fmt.upper()))
_save_pil_image(pil_img, path, fmt, quality, metadata_text, workflow_metadata)
return True
except Exception as e:
print(f"Basic data handling: Error saving image: {e}")
return False
if not metadata_text and not workflow_metadata:
print(f"Basic data handling: Error saving image: {e}")
return False
try:
print(f"Basic data handling: Saving with metadata failed ({e}); retrying without metadata.")
_save_pil_image(pil_img, path, fmt, quality, "", None)
return True
except Exception as retry_error:
print(f"Basic data handling: Error saving image: {retry_error}")
return False
def _ui_image_entry(path: str) -> Optional[dict[str, str]]:
@@ -1715,6 +1782,9 @@ class PathSaveImageRGB(ComfyNodeABC):
When a batch is saved from a data list the workflow necessarily contains
every prompt of the graph; the prompt of the individual image is the one in
``parameters``. Formats that cannot carry text metadata ignore both.
A JPEG can only hold one 16-bit-sized EXIF block, so for very large
workflows the workflow info is reduced (largest channel first) instead of
failing the save; the image itself is always written.
"""
@classmethod
@@ -1861,6 +1931,9 @@ class PathSaveImageRGBA(ComfyNodeABC):
When a batch is saved from a data list the workflow necessarily contains
every prompt of the graph; the prompt of the individual image is the one in
``parameters``. Formats that cannot carry text metadata ignore both.
A JPEG can only hold one 16-bit-sized EXIF block, so for very large
workflows the workflow info is reduced (largest channel first) instead of
failing the save; the image itself is always written.
"""
@classmethod
tests/test_path_nodes.py
+95
View File
@@ -977,6 +977,101 @@ def test_path_save_image_workflow_metadata_in_exif(tmp_path, monkeypatch):
assert exif.get(0x010F) == "workflow:" + json.dumps(extra_pnginfo["workflow"])
def _oversized_workflow_metadata():
"""A workflow whose full graph JSON comfortably exceeds one JPEG EXIF block
(a JPEG carries EXIF in a single 16-bit-sized APP1 marker)."""
big_graph = {"nodes": [{"id": i, "type": "SomeNode", "widgets_values": ["x" * 200]} for i in range(300)]}
workflow_prompt = {"3": {"class_type": "KSampler", "inputs": {"text": "a" * 500}}}
return workflow_prompt, {"workflow": big_graph}
def test_path_save_image_oversized_jpeg_metadata_keeps_prompt(tmp_path, capsys):
"""Too much metadata for a JPEG EXIF block must not fail the save: the
largest (full-graph) channel is dropped, the rest is kept, and it is
reported."""
save_node = PathSaveImageRGB()
img = _batch_rgb(1)
workflow_prompt, extra_pnginfo = _oversized_workflow_metadata()
out = str(tmp_path / "oversized")
result = save_node.save_image(
img, out, format="jpg", prompt="explicit", workflow_prompt=workflow_prompt, workflow_extra_pnginfo=extra_pnginfo
)
assert save_success(result) is True
assert os.path.exists(out + ".jpg")
with Image.open(out + ".jpg") as loaded:
loaded.load()
assert loaded.format == "JPEG"
exif = loaded.getexif()
# The explicit prompt (and the smaller API prompt) survive ...
assert decode_user_comment(exif.get_ifd(0x8769).get(0x9286)) == "explicit"
assert exif.get(0x0110) == "prompt:" + json.dumps(workflow_prompt)
# ... while the oversized full workflow graph was dropped.
assert 0x010F not in exif
assert 'had to drop the full "workflow" graph' in capsys.readouterr().out
def test_path_save_image_oversized_jpeg_drops_all_metadata_but_saves(tmp_path, capsys):
"""When even the remaining channels do not fit, the image is still written -
just without metadata."""
save_node = PathSaveImageRGB()
img = _batch_rgb(1)
_, extra_pnginfo = _oversized_workflow_metadata()
out = str(tmp_path / "huge")
# The explicit prompt itself is far beyond a single EXIF block.
result = save_node.save_image(img, out, format="jpg", prompt="x" * 70000, workflow_extra_pnginfo=extra_pnginfo)
assert save_success(result) is True
assert os.path.exists(out + ".jpg")
with Image.open(out + ".jpg") as loaded:
loaded.load()
assert loaded.format == "JPEG"
assert not loaded.getexif()
assert "saving the image without metadata" in capsys.readouterr().out
def test_path_save_image_png_keeps_oversized_workflow(tmp_path):
"""Formats with room enough (PNG text chunks) keep the metadata untouched."""
save_node = PathSaveImageRGB()
img = _batch_rgb(1)
workflow_prompt, extra_pnginfo = _oversized_workflow_metadata()
out = str(tmp_path / "big_png")
result = save_node.save_image(img, out, format="png", workflow_prompt=workflow_prompt, workflow_extra_pnginfo=extra_pnginfo)
assert save_success(result) is True
with Image.open(out + ".png") as loaded:
loaded.load()
assert json.loads(loaded.text["workflow"]) == extra_pnginfo["workflow"]
def test_write_pil_image_retries_without_metadata(tmp_path, monkeypatch):
"""If embedding metadata fails for an unexpected reason, the image is saved
again without it instead of losing the image."""
import src.basic_data_handling.path_nodes as pn
attempts = []
def fake_save(pil_img, path, fmt, quality, metadata_text, workflow_metadata=None):
attempts.append(bool(metadata_text or workflow_metadata))
if metadata_text or workflow_metadata:
raise ValueError("EXIF data is too long")
pil_img.save(path, format="PNG")
monkeypatch.setattr(pn, "_save_pil_image", fake_save)
img = Image.new("RGB", (8, 8), (255, 0, 0))
out = str(tmp_path / "retry.png")
assert pn._write_pil_image(img, out, "png", 95, "prompt", {"workflow": "{}"}) is True
assert attempts == [True, False]
with Image.open(out) as loaded:
loaded.load()
def test_path_save_image_metadata_can_be_disabled(tmp_path, monkeypatch):
"""ComfyUI's ``--disable-metadata`` suppresses both metadata channels, like
it does in the built-in save node."""