From 9a43b67008874611ec5f1577aa11ebe9a7aa9127 Mon Sep 17 00:00:00 2001 From: Vito Date: Sat, 31 Oct 2026 23:17:17 +0000 Subject: [PATCH 1/2] test: reset a disposed fixture package's exit actions (#580) (#581) load_package imports the package under a test-only name; its exit actions stayed registered after cleanup_package removed it from sys.modules, so the suite ended with a ModuleNotFoundError from _kill_encoders_at_exit. cleanup_package now resets the package's exit signal. --- tests/test_init_package.py | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/tests/test_init_package.py b/tests/test_init_package.py index 42df3b58..6483e809 100644 --- a/tests/test_init_package.py +++ b/tests/test_init_package.py @@ -136,6 +136,11 @@ def cleanup_package(package, require_stopped=True): if registration is not None: assert registration.dispose(3), "package history worker leaked" close_every_database() + # The package's exit actions outlive it otherwise: at interpreter exit + # they would import from a package no longer in sys.modules (#580). + exit_signal = getattr(package, "_exit_signal", None) + if exit_signal is not None: + exit_signal.reset() sys.modules.pop(package.__name__, None) @@ -402,6 +407,29 @@ class TestConsoleBanner(unittest.TestCase): self.assertNotIn("Loaded:", banner) +class TestFixtureExitActions(unittest.TestCase): + """A disposed fixture package runs nothing at interpreter exit (#580).""" + + def test_a_disposed_package_leaves_no_exit_action(self): + tmp = tempfile.TemporaryDirectory() + self.addCleanup(tmp.cleanup) + self.addCleanup(close_every_database) + output_dir = os.path.join(tmp.name, "output") + os.makedirs(output_dir) + name = "prompt_manager_pkg_exit_actions" + self.addCleanup(sys.modules.pop, name, None) + + with contextlib.redirect_stdout(io.StringIO()): + with load_package(name, os.path.join(tmp.name, "p.db"), output_dir) as ( + package, + _, + ): + signal = package._exit_signal + self.assertTrue(signal._on_exit, "the package installs exit actions") + + self.assertEqual(signal._on_exit, []) + + class PriorFileHandler(logging.FileHandler): """A caller-owned sink that stays closed once its owner closes it. From 3cb90b0c1e25d391ed5b9f6e7bc0a1b8898b1792 Mon Sep 17 00:00:00 2001 From: Vito Date: Sun, 1 Nov 2026 00:04:27 +0000 Subject: [PATCH 2/2] perf(images): read PNG text without decoding pixels (#523) (#577) * perf(images): read PNG text without decoding pixels (#523) Pillow's PngImageFile.text loads every pixel to reach text chunks written after IDAT. The scan's metadata batch and the image monitor now read PNG text through utils/png_text.py, which runs Pillow's own chunk handlers after the first IDAT and seeks over the pixel data. The legacy .text read is the floor: a differential test compares both readers on every fixture image, a curated corpus and a seeded fuzz. The only difference is a file whose pixel data is undecodable or cut short: it now keeps the text read before its pixels instead of reading nothing. Animated PNGs still use .text. * test(images): cover the text budget, XMP and layout guard of the PNG text walk (#523) --- TESTING.md | 13 ++ py/api/__init__.py | 6 +- py/api/admin.py | 6 +- tests/test_png_text.py | 455 ++++++++++++++++++++++++++++++++++++ utils/metadata_extractor.py | 21 +- utils/png_text.py | 116 +++++++++ 6 files changed, 604 insertions(+), 13 deletions(-) create mode 100644 tests/test_png_text.py create mode 100644 utils/png_text.py diff --git a/TESTING.md b/TESTING.md index df7c6695..d616c780 100644 --- a/TESTING.md +++ b/TESTING.md @@ -456,6 +456,19 @@ A hot SQLite read that returns many rows aggregates them into one step (`json_group_array`) instead of stepping once per row, so it does not wait for the GIL once per row while ComfyUI generates (#354, `get_image_link_matches`). +A metadata read never decodes pixels (#523). Pillow's `PngImageFile.text` loads +every pixel to reach text chunks after IDAT; the scan (`_extract_comfyui_metadata`) +and the image monitor (`ComfyUIMetadataExtractor.extract_metadata`) read PNG text +through `utils/png_text.py`, which runs Pillow's own chunk handlers after the +first IDAT and seeks over pixel data. The legacy `.text` read is the floor: +`tests/test_png_text.py` compares both readers on every image under +`tests/fixtures`, a curated corpus (text after IDAT, zTXt, iTXt, split IDAT, bad +checksums, truncations, oversized text, APNG) and a seeded fuzz, and asserts that +neither reader decodes pixels by patching `ImageFile.load`. The only reads that +differ are files whose pixel data is undecodable or cut short: legacy found +nothing, the walk keeps the text before the pixels. Animated PNGs still use +`.text`. + The `perf budget` job in `.github/workflows/test.yml` is the reference measurement: `ubuntu-latest`, one Python, no coverage, only the wall-clock tests. It runs on every push to `main` and on `workflow_dispatch`, never on pull diff --git a/py/api/__init__.py b/py/api/__init__.py index ce3272dd..9c71f7a7 100644 --- a/py/api/__init__.py +++ b/py/api/__init__.py @@ -85,6 +85,7 @@ try: from ...database.operations import PromptDatabase from ...utils.file_prompt import parse_comfyui_prompt from ...utils.logging_config import get_logger + from ...utils.png_text import png_text from ...utils.video_metadata import is_video_path, read_video_metadata except ImportError: import sys @@ -95,6 +96,7 @@ except ImportError: from database.operations import PromptDatabase from utils.file_prompt import parse_comfyui_prompt from utils.logging_config import get_logger + from utils.png_text import png_text from utils.video_metadata import is_video_path, read_video_metadata @@ -1048,6 +1050,8 @@ class PromptManagerAPI( def _extract_comfyui_metadata(self, image_path): """Embedded ComfyUI metadata of an image (PNG text chunks) or a video. + PNG text is read without decoding the pixels (``utils.png_text``, #523). + Videos go through ffprobe (see utils/video_metadata). Unreadable files are a warning, not an error: a scan over a large output tree meets truncated and foreign files, and it carries on past them. @@ -1056,7 +1060,7 @@ class PromptManagerAPI( return read_video_metadata(image_path) try: with Image.open(image_path) as img: - return dict(getattr(img, "text", None) or {}) + return png_text(img) except Exception as e: self.logger.warning( f"Unreadable image {os.path.basename(str(image_path))}: {e}" diff --git a/py/api/admin.py b/py/api/admin.py index 3488a056..dc70b9f0 100644 --- a/py/api/admin.py +++ b/py/api/admin.py @@ -1668,9 +1668,9 @@ class AdminRoutesMixin: def _extract_batch_metadata_sync(self, file_batch): """Extract metadata from a batch of files in one executor call. - Files are read on ``PromptManagerConfig.WORKER_THREADS`` threads: - Pillow decodes the whole image to reach trailing text chunks and - releases the GIL while doing so. + Files are read on ``PromptManagerConfig.WORKER_THREADS`` threads. + PNG text is read chunk by chunk without decoding pixels (#523); file + reads and ffprobe release the GIL. """ from ..config import PromptManagerConfig diff --git a/tests/test_png_text.py b/tests/test_png_text.py new file mode 100644 index 00000000..9633e506 --- /dev/null +++ b/tests/test_png_text.py @@ -0,0 +1,455 @@ +"""PNG text without a pixel decode (#523), against the legacy reader floor. + +Pillow's ``PngImageFile.text`` loads (decodes) every pixel to reach text +chunks written after IDAT. ``utils.png_text.png_text`` walks the remaining +chunks instead. The legacy reader (``Image.open(path).text``, every failure +read as ``{}``) is the floor: a file it read must read the same, and the +curated corpus below must read identically. The documented difference is +pixel data, which the walk never decodes: undecodable or cut-short pixel data +made the legacy reader fail, the walk still reads the text before it. +""" + +import io +import random +import struct +import tempfile +import unittest +import zlib +from pathlib import Path +from unittest.mock import patch + +from PIL import Image, ImageFile, PngImagePlugin + +from utils.metadata_extractor import ComfyUIMetadataExtractor +from utils.png_text import png_text + +FIXTURES = Path(__file__).resolve().parent / "fixtures" +IMAGE_SUFFIXES = {".png", ".jpg", ".jpeg", ".webp", ".gif", ".bmp"} +SIGNATURE = b"\x89PNG\r\n\x1a\n" +FUZZ_SEED = 523 +FUZZ_CASES = 400 + + +def legacy_text(path): + """The legacy scan reader: Pillow's ``.text`` (a full decode).""" + try: + with Image.open(path) as image: + return dict(getattr(image, "text", None) or {}) + except Exception: + return {} + + +def walked_text(path): + """The new reader under the same failure rule.""" + try: + with Image.open(path) as image: + return png_text(image) + except Exception: + return {} + + +def chunk(kind, data): + return ( + struct.pack(">I", len(data)) + + kind + + data + + struct.pack(">I", zlib.crc32(kind + data) & 0xFFFFFFFF) + ) + + +def chunks_of(raw): + """``[(kind, data)]`` of a well-formed PNG.""" + found, offset = [], 8 + while offset < len(raw): + length, kind = struct.unpack(">I4s", raw[offset : offset + 8]) + found.append((kind, raw[offset + 8 : offset + 8 + length])) + offset += 12 + length + return found + + +def assemble(parts): + return SIGNATURE + b"".join( + part if isinstance(part, bytes) else chunk(*part) for part in parts + ) + + +def base_png(mode="RGB", size=(64, 64), text=None, noise=False): + """A PNG from Pillow (text before IDAT), as ``[(kind, data)]``.""" + image = Image.new(mode, size) + if noise: + rng = random.Random(7) + image = Image.frombytes( + "RGB", size, bytes(rng.randrange(256) for _ in range(size[0] * size[1] * 3)) + ) + info = PngImagePlugin.PngInfo() + for key, value in (text or {}).items(): + info.add_text(key, value) + buffer = io.BytesIO() + image.save(buffer, format="PNG", pnginfo=info) + return chunks_of(buffer.getvalue()) + + +def split_idat(parts, pieces=3): + """Spread the IDAT payload over ``pieces`` consecutive IDAT chunks.""" + data = b"".join(d for k, d in parts if k == b"IDAT") + step = max(1, len(data) // pieces) + idats = [(b"IDAT", data[i : i + step]) for i in range(0, len(data), step)] + first = next(i for i, (k, _) in enumerate(parts) if k == b"IDAT") + rest = [p for p in parts if p[0] != b"IDAT"] + return rest[:first] + idats + rest[first:] + + +def before_iend(parts, *extra): + return parts[:-1] + list(extra) + parts[-1:] + + +def text_chunk(key, value): + return (b"tEXt", key.encode("latin-1") + b"\0" + value.encode("latin-1")) + + +def ztxt_chunk(key, value): + return (b"zTXt", key.encode() + b"\0\0" + zlib.compress(value.encode("latin-1"))) + + +def itxt_chunk(key, value, compressed=False): + body = value.encode("utf-8") + if compressed: + body = zlib.compress(body) + flag = b"\1\0" if compressed else b"\0\0" + return (b"iTXt", key.encode() + b"\0" + flag + b"en\0" + b"Title\0" + body) + + +def bad_crc(kind, data): + good = chunk(kind, data) + return good[:-4] + bytes(b ^ 0xFF for b in good[-4:]) + + +PROMPT = '{"3": {"class_type": "KSampler", "inputs": {"seed": 1}}}' +WORKFLOW = '{"nodes": [], "links": []}' + + +def curated(): + """Name to PNG bytes; each must read the same with both readers.""" + plain = base_png(text={"prompt": PROMPT}) + trailing = before_iend( + base_png(), text_chunk("prompt", PROMPT), text_chunk("workflow", WORKFLOW) + ) + cases = { + "no-text": assemble(base_png()), + "text-before-pixels": assemble(plain), + "text-after-pixels": assemble(trailing), + "ztxt-after-pixels": assemble( + before_iend(base_png(), ztxt_chunk("prompt", PROMPT)) + ), + "itxt-plain-and-compressed": assemble( + before_iend( + base_png(), + itxt_chunk("prompt", "café " + PROMPT), + itxt_chunk("workflow", WORKFLOW, compressed=True), + ) + ), + "xmp-after-pixels": assemble( + before_iend( + base_png(), + itxt_chunk("XML:com.adobe.xmp", ""), + text_chunk("prompt", PROMPT), + ) + ), + "itxt-bad-utf8": assemble( + before_iend( + base_png(), + (b"iTXt", b"prompt\0\0\0\0\0\xff\xfe"), + text_chunk("workflow", WORKFLOW), + ) + ), + "later-copy-wins": assemble( + before_iend(plain, text_chunk("prompt", '{"later": 1}')) + ), + "empty-key-and-no-separator": assemble( + before_iend(base_png(), (b"tEXt", b"\0value"), (b"tEXt", b"lonely")) + ), + "split-pixels-text-between": assemble( + split_idat(before_iend(base_png(noise=True), text_chunk("prompt", PROMPT))) + ), + "palette": assemble( + before_iend(base_png(mode="P"), text_chunk("prompt", PROMPT)) + ), + "sixteen-bit": assemble( + before_iend(base_png(mode="I;16"), text_chunk("prompt", PROMPT)) + ), + "bad-crc-before-pixels": assemble( + [base_png()[0], bad_crc(*text_chunk("prompt", PROMPT))] + base_png()[1:] + ), + "bad-crc-after-pixels": assemble( + base_png()[:-1] + [bad_crc(*text_chunk("prompt", PROMPT))] + base_png()[-1:] + ), + "unknown-chunk-after-pixels": assemble( + before_iend(base_png(), (b"prVt", b"x" * 40), text_chunk("prompt", PROMPT)) + ), + "garbage-chunk-id-after-pixels": assemble( + base_png()[:-1] + + [b"\0\0\0\4\x00\x01\x02\x03abcd\0\0\0\0"] + + [chunk(*text_chunk("prompt", PROMPT))] + + base_png()[-1:] + ), + "oversized-ztxt-after-pixels": assemble( + before_iend( + base_png(), + ( + b"zTXt", + b"prompt\0\0" + + zlib.compress(b"a" * (PngImagePlugin.MAX_TEXT_CHUNK + 10)), + ), + ) + ), + "pixels-length-past-end": assemble( + [p for p in base_png() if p[0] != b"IDAT" and p[0] != b"IEND"] + ) + + b"\xff\xff\x00\x00IDAT" + + b"\0" * 64, + } + full = assemble(trailing) + idat_at = full.index(b"IDAT") + cases["truncated-in-pixels"] = full[: idat_at + 20] + cases["no-end-after-pixels"] = full[: full.rindex(b"IEND") - 4] + cases["truncated-in-trailing-text"] = full[: full.rindex(b"workflow") + 4] + cases["unknown-chunk-past-end"] = ( + assemble(before_iend(base_png(), ztxt_chunk("prompt", PROMPT))[:-1]) + + chunk(b"prVt", b"x" * 40)[:-20] + ) + cases["not-a-png"] = b"nope, not an image at all" + cases["signature-only"] = SIGNATURE + return cases + + +def corrupt_pixels(): + """Complete IDAT chunks with valid CRCs over data zlib cannot inflate.""" + parts = before_iend(base_png(), text_chunk("prompt", PROMPT)) + return assemble( + [(k, b"\x78\x9c" + b"\xff" * 30) if k == b"IDAT" else (k, d) for k, d in parts] + ) + + +def apng(): + frames = [Image.new("RGB", (16, 16), colour) for colour in ("red", "blue")] + info = PngImagePlugin.PngInfo() + info.add_text("prompt", PROMPT) + buffer = io.BytesIO() + frames[0].save( + buffer, format="PNG", save_all=True, append_images=frames[1:], pnginfo=info + ) + return assemble( + before_iend(chunks_of(buffer.getvalue()), text_chunk("workflow", WORKFLOW)) + ) + + +def fixture_images(): + return sorted(p for p in FIXTURES.rglob("*") if p.suffix.lower() in IMAGE_SUFFIXES) + + +class _TempFiles(unittest.TestCase): + def setUp(self): + self._tmp = tempfile.TemporaryDirectory() + self.addCleanup(self._tmp.cleanup) + + def write(self, name, raw): + path = Path(self._tmp.name) / name + path.write_bytes(raw) + return path + + +class PngTextParityTests(_TempFiles): + def test_every_fixture_image_reads_like_the_legacy_reader(self): + images = fixture_images() + self.assertGreaterEqual(len(images), 25) + for path in images: + with self.subTest(path=path.relative_to(FIXTURES)): + self.assertEqual(walked_text(path), legacy_text(path)) + + def test_curated_corpus_reads_like_the_legacy_reader(self): + for name, raw in curated().items(): + with self.subTest(name): + path = self.write(name + ".png", raw) + self.assertEqual(walked_text(path), legacy_text(path)) + + def test_curated_corpus_exercises_trailing_text(self): + """The corpus is not vacuous: trailing text exists and is read.""" + cases = curated() + path = self.write("t.png", cases["text-after-pixels"]) + self.assertEqual(walked_text(path), {"prompt": PROMPT, "workflow": WORKFLOW}) + with Image.open(path) as image: + self.assertNotIn("prompt", image.info) + self.assertEqual( + walked_text(self.write("x.png", cases["truncated-in-pixels"])), {} + ) + + def test_itxt_values_keep_their_language_and_type(self): + path = self.write("i.png", curated()["itxt-plain-and-compressed"]) + new, old = walked_text(path), legacy_text(path) + self.assertIsInstance(new["prompt"], PngImagePlugin.iTXt) + self.assertEqual(new["prompt"].lang, old["prompt"].lang) + self.assertEqual(new["prompt"].tkey, old["prompt"].tkey) + + def test_animated_png_falls_back_to_the_legacy_read(self): + path = self.write("a.png", apng()) + self.assertEqual(walked_text(path), legacy_text(path)) + self.assertEqual(walked_text(path)["workflow"], WORKFLOW) + + def test_truncated_pixels_read_alike_when_truncated_images_load(self): + """With LOAD_TRUNCATED_IMAGES set (a node pack may set it) both readers + accept a short file alike.""" + path = self.write("t.png", curated()["truncated-in-pixels"]) + with patch.object(ImageFile, "LOAD_TRUNCATED_IMAGES", True): + self.assertEqual(walked_text(path), legacy_text(path)) + + def test_a_text_decode_error_ends_the_walk_like_legacy(self): + """Pillow's ``load_end`` stops at a chunk whose handler raises + ``UnicodeDecodeError`` and keeps the text read so far.""" + path = self.write( + "u.png", + assemble( + before_iend( + base_png(), + ztxt_chunk("prompt", PROMPT), + text_chunk("workflow", WORKFLOW), + ztxt_chunk("parameters", "steps 20"), + ) + ), + ) + + def undecodable(stream, position, length): + raise UnicodeDecodeError("latin-1", b"", 0, 1, "test") + + with patch.object(PngImagePlugin.PngStream, "chunk_tEXt", undecodable): + old = legacy_text(path) + self.assertEqual(walked_text(path), old) + self.assertEqual(old, {"prompt": PROMPT}) + + def test_the_text_memory_budget_counts_text_before_the_pixels(self): + """Pillow's MAX_TEXT_MEMORY spans the whole file: text before and + after the pixels, each under it, fail together in both readers.""" + path = self.write( + "m.png", + assemble( + before_iend( + base_png(text={"prompt": PROMPT}), text_chunk("workflow", WORKFLOW) + ) + ), + ) + budget = len(PROMPT) + len(WORKFLOW) - 1 + with patch.object(PngImagePlugin, "MAX_TEXT_MEMORY", budget): + self.assertEqual(legacy_text(path), {}) + self.assertEqual(walked_text(path), {}) + with patch.object(PngImagePlugin, "MAX_TEXT_MEMORY", budget + 1): + self.assertEqual(walked_text(path), legacy_text(path)) + self.assertEqual(len(walked_text(path)), 2) + + def test_complete_but_undecodable_pixels_read_more_than_legacy(self): + path = self.write("c.png", corrupt_pixels()) + self.assertEqual(legacy_text(path), {}) + self.assertEqual(walked_text(path), {"prompt": PROMPT}) + + def test_cut_short_pixels_read_the_text_before_them(self): + full = assemble(base_png(text={"prompt": PROMPT})) + path = self.write("s.png", full[: full.index(b"IDAT") + 12]) + self.assertEqual(legacy_text(path), {}) + self.assertEqual(walked_text(path), {"prompt": PROMPT}) + + def test_mutated_files_never_read_less_than_legacy(self): + """Seeded fuzz: truncations and byte flips of a file with text before, + between and after its pixel chunks. Wherever the legacy reader found + text, the walk finds exactly the same.""" + base = assemble( + split_idat( + before_iend( + base_png(text={"workflow": WORKFLOW}), + ztxt_chunk("prompt", PROMPT), + itxt_chunk("parameters", "steps 20"), + ) + ) + ) + rng = random.Random(FUZZ_SEED) + compared = 0 + for case in range(FUZZ_CASES): + raw = bytearray(base) + if case % 2: + del raw[rng.randrange(8, len(raw)) :] + else: + for _ in range(rng.randrange(1, 4)): + raw[rng.randrange(8, len(raw))] = rng.randrange(256) + path = self.write(f"fuzz{case}.png", bytes(raw)) + old = legacy_text(path) + if old: + compared += 1 + with self.subTest(case=case): + self.assertEqual(walked_text(path), old) + self.assertGreater(compared, FUZZ_CASES // 5) + + def test_non_png_images_have_no_text(self): + for name in ("tiny.jpg", "tiny.webp"): + with Image.open(FIXTURES / "local-inspect" / name) as image: + self.assertEqual(png_text(image), {}) + + +class NoPixelDecodeTests(_TempFiles): + """Metadata reads never decode pixels (the #523 memory finding).""" + + def png_fixtures(self): + return sorted((FIXTURES / "inspection_parity").glob("*.png")) + + def test_png_text_does_not_load_pixels_and_leaves_the_image_usable(self): + path = self.write("t.png", curated()["text-after-pixels"]) + with Image.open(path) as image: + with patch.object(ImageFile.ImageFile, "load", autospec=True) as load: + self.assertEqual(png_text(image)["workflow"], WORKFLOW) + load.assert_not_called() + image.load() + self.assertEqual(image.size, (64, 64)) + self.assertEqual(image.text["workflow"], WORKFLOW) + + def test_scan_reader_does_not_decode_pixels(self): + from py.api import PromptManagerAPI + + api = PromptManagerAPI.__new__(PromptManagerAPI) + fixtures = self.png_fixtures() + [ + self.write("t.png", curated()["text-after-pixels"]) + ] + with patch.object(ImageFile.ImageFile, "load", autospec=True) as load: + found = [api._extract_comfyui_metadata(str(p)) for p in fixtures] + load.assert_not_called() + self.assertEqual(found, [legacy_text(p) for p in fixtures]) + self.assertTrue(all(found)) + + def test_monitor_extractor_does_not_decode_pixels(self): + extractor = ComfyUIMetadataExtractor() + fixtures = self.png_fixtures() + with patch.object(ImageFile.ImageFile, "load", autospec=True) as load: + found = [extractor.extract_metadata(str(p)) for p in fixtures] + load.assert_not_called() + self.assertTrue(all(metadata and len(metadata) > 1 for metadata in found)) + + +class MonitorExtractorParityTests(_TempFiles): + """``extract_metadata`` answers exactly as with the legacy ``.text``.""" + + def legacy_extract(self, path): + def old_text(image): + return dict(getattr(image, "text", None) or {}) + + with patch("utils.metadata_extractor.png_text", old_text): + return ComfyUIMetadataExtractor().extract_metadata(str(path)) + + def test_every_fixture_and_curated_file_extracts_like_legacy(self): + paths = fixture_images() + [ + self.write(name + ".png", raw) for name, raw in curated().items() + ] + extractor = ComfyUIMetadataExtractor() + for path in paths: + with self.subTest(path=path.name): + self.assertEqual( + extractor.extract_metadata(str(path)), self.legacy_extract(path) + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/utils/metadata_extractor.py b/utils/metadata_extractor.py index c1c044a8..bd0868ea 100644 --- a/utils/metadata_extractor.py +++ b/utils/metadata_extractor.py @@ -37,6 +37,7 @@ from typing import Optional, Dict, Any from PIL import Image from .logging_config import get_logger +from .png_text import png_text class ComfyUIMetadataExtractor: @@ -88,12 +89,14 @@ class ComfyUIMetadataExtractor: # Add basic file information metadata["file_info"] = self.get_file_info(image_path, image) - # Extract ComfyUI-specific metadata from PNG text chunks - if hasattr(image, "text") and image.text: + # Extract ComfyUI-specific metadata from PNG text chunks, + # read without decoding the pixels (#523) + text = png_text(image) + if text: # Look for ComfyUI workflow data - if "workflow" in image.text: + if "workflow" in text: try: - workflow_data = json.loads(image.text["workflow"]) + workflow_data = json.loads(text["workflow"]) metadata["workflow"] = workflow_data # Extract text encoder nodes from workflow @@ -107,9 +110,9 @@ class ComfyUIMetadataExtractor: self.logger.warning(f"Failed to parse workflow JSON: {e}") # Look for prompt data - if "prompt" in image.text: + if "prompt" in text: try: - prompt_data = json.loads(image.text["prompt"]) + prompt_data = json.loads(text["prompt"]) metadata["prompt"] = prompt_data except json.JSONDecodeError as e: self.logger.warning(f"Failed to parse prompt JSON: {e}") @@ -129,13 +132,13 @@ class ComfyUIMetadataExtractor: ] for field in metadata_fields: - if field in image.text: + if field in text: try: # Try to parse as JSON first - metadata[field] = json.loads(image.text[field]) + metadata[field] = json.loads(text[field]) except (json.JSONDecodeError, TypeError): # Store as string if not valid JSON - metadata[field] = image.text[field] + metadata[field] = text[field] return ( metadata diff --git a/utils/png_text.py b/utils/png_text.py new file mode 100644 index 00000000..36f43e25 --- /dev/null +++ b/utils/png_text.py @@ -0,0 +1,116 @@ +"""PNG text chunks without decoding pixels (#523). + +Pillow's ``PngImageFile.text`` calls ``load()``, which decodes every pixel, to +reach text chunks written after the image data. On a scan or a busy image +monitor that is a full-size buffer per file for a few kilobytes of text. + +``png_text`` returns the same dictionary as ``image.text`` for an opened, +single-frame PNG, by walking the chunks after the first IDAT with Pillow's own +chunk handlers (same decoding, limits and duplicate rule) and seeking over the +pixel data. It follows Pillow's ``load_end`` loop: chunk checksums after the +pixels are not checked, an unreadable chunk header ends the walk, and another +chunk that runs past the end of the file fails like ``ImageFile._safe_read`` +(unless ``ImageFile.LOAD_TRUNCATED_IMAGES``). + +The one difference is pixel data, which is never decoded: pixel data that is +undecodable, or cut short, no longer hides the text read before it. Pillow +itself accepts a file cut inside its last pixel chunk once the decoder has +all rows, and only a decode can tell those apart, so the walk never reads +less than ``.text`` did (the PR #170 "legacy parity floor"). + +Anything else (other formats, animated PNGs, an image already loaded) is +answered by Pillow itself. +""" + +import os +import struct + +from PIL import ImageFile + +_PIXEL_CHUNKS = (b"IDAT", b"fdAT") +# Signature, IHDR chunk and IDAT chunk header: the earliest pixel data offset. +_FIRST_DATA_OFFSET = 8 + 25 + 8 + + +def png_text(image) -> dict: + """``dict(image.text)`` for an opened image, without a pixel decode. + + The image stays usable: its file position and Pillow state are restored. + Non-PNG images answer ``{}`` (only PNG has text chunks in Pillow). + """ + if not _walkable(image): + return dict(getattr(image, "text", None) or {}) + fp = image.fp + resume = fp.tell() + try: + return _walk_after_pixels(image.png, fp, image.tile[0][2] - 8) + finally: + fp.seek(resume) + + +def _walkable(image) -> bool: + """A plain single-frame PNG whose header pass stopped at its first IDAT.""" + tile = getattr(image, "tile", None) + return ( + image.format == "PNG" + and getattr(image, "png", None) is not None + and getattr(image, "fp", None) is not None + and not getattr(image, "is_animated", False) + and isinstance(tile, list) + and len(tile) == 1 + and tile[0][0] == "zip" + # The first IDAT's data offset; anything else is not a layout we know. + and isinstance(tile[0][2], int) + and tile[0][2] >= _FIRST_DATA_OFFSET + ) + + +def _walk_after_pixels(stream, fp, offset: int) -> dict: + """Pillow's ``load_end`` chunk loop from the first IDAT header, skipping pixels. + + ``stream`` is the image's own ``PngStream``: its text so far and its text + memory budget carry on exactly as a ``load()`` would carry them. The text + and info dictionaries are copied first, so ``image.text`` and + ``image.info`` are unchanged. Like ``load_end``, a handler for a chunk + misplaced after the pixels (a second IHDR or PLTE) may still update the + stream's other header fields; Pillow ignores them once the image is open. + """ + end = fp.seek(0, os.SEEK_END) + saved = stream.im_text, stream.im_info, stream.text_memory + stream.im_text, stream.im_info = dict(saved[0]), dict(saved[1]) + try: + fp.seek(offset) + while True: + try: + kind, position, length = stream.read() + except (struct.error, SyntaxError): + break + if kind == b"IEND": + break + try: + handled = kind not in _PIXEL_CHUNKS and _handle( + stream, kind, position, length + ) + except UnicodeDecodeError: + break + if not handled: + if position + length > end: + if kind in _PIXEL_CHUNKS or ImageFile.LOAD_TRUNCATED_IMAGES: + break + raise OSError("Truncated File Read") + fp.seek(position + length + 4) + continue + fp.read(4) # checksum: not verified after the pixels, as in Pillow + return stream.im_text + finally: + stream.im_text, stream.im_info, stream.text_memory = saved + + +def _handle(stream, kind: bytes, position: int, length: int) -> bool: + """Run Pillow's handler; ``False`` when the chunk's bytes must be skipped.""" + try: + stream.call(kind, position, length) + except (EOFError, AttributeError): + # A chunk Pillow reads past (EOFError) or has no handler for. + return False + return True