From 13fe1d32f19984cbd6f85c7b6cb1fba527bc8d7d Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Mon, 19 Oct 2026 23:32:19 -0700 Subject: [PATCH] fix(images): card reads PNG records with the inspector's reader The card read PNG text through Pillow's header info, where a later duplicate prompt chunk silently wins and text after IDAT is missed. The inspector's bounded reader refuses conflicting duplicates and reads text after IDAT, so a file with two different prompt chunks showed a card seed the inspector refused. PNG cards now read their records with read_inspection_file, the reader the inspector uses. When it refuses a file the card keeps its dimensions from Pillow and shows no seed. Other formats still use Pillow's header info, as the inspector's header read does. --- TESTING.md | 9 ++- py/api/images.py | 35 +++++++++-- tests/test_api_workspace_images.py | 94 +++++++++++++++++++++++++++--- 3 files changed, 124 insertions(+), 14 deletions(-) diff --git a/TESTING.md b/TESTING.md index 411e14e..4ac2c90 100644 --- a/TESTING.md +++ b/TESTING.md @@ -449,8 +449,13 @@ and skip the parse. - `seed`: a decimal string, or null. Always a string, because `JSON.parse` rounds numbers past 2**53 and ComfyUI's randomized seeds are often that large. It is `extract_generation_params` from `utils/generation_params.py` (the inspector's - legacy parser, #251), run on the `prompt`/`workflow` header records read by the - same open that reads the dimensions. The records go through the inspection + legacy parser, #251), run on the `prompt`/`workflow` records read with the + dimensions. A PNG is read by the inspector's own reader + (`utils.prompt_inspector.read_inspection_file`): text after IDAT counts, the + shared byte budgets apply, and a file with conflicting duplicate `prompt` or + `workflow` chunks gets no card seed, as the inspector refuses it (the card + keeps its dimensions). Other formats use Pillow's header `info`, as the + inspector's header read does. The records go through the inspection projection first (`utils/inspection_response.project_graph`), which carries integers past 2**53 as their exact digits, so the card and the inspector read the same digits. The parse runs after the file handle closes. A record over diff --git a/py/api/images.py b/py/api/images.py index da9d7d6..0613f4e 100644 --- a/py/api/images.py +++ b/py/api/images.py @@ -37,6 +37,7 @@ try: from ...utils.card_seed import card_seed, header_params, parsed_positive from ...utils.parallel import map_parallel from ...utils.inspection_response import encode_inspection_response + from ...utils.prompt_inspector import MetadataReadError, read_inspection_file except ImportError: from database.image_favorites import MAX_KEY_LENGTH from database.library import display_title @@ -44,11 +45,40 @@ except ImportError: from utils.card_seed import card_seed, header_params, parsed_positive from utils.parallel import map_parallel from utils.inspection_response import encode_inspection_response + from utils.prompt_inspector import MetadataReadError, read_inspection_file IMAGE_EXTENSIONS = _IMAGE_EXTENSIONS # The header records the card parses (see utils.card_seed.header_params). _CARD_RECORDS = ("prompt", "workflow") +_PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n" + + +def _card_header(path: Path) -> tuple: + """``(width, height, records)`` of one image card, every handle closed. + + A PNG's records come from the inspector's own reader + (``read_inspection_file``: text after IDAT, the shared byte budgets, and + the refusal of conflicting duplicate prompt or workflow chunks), so the + card reads exactly what the inspector reads. When that reader refuses the + file, the card keeps its dimensions and gets no records. Other formats + keep Pillow's header ``info``, as the inspector's header read does. + Raises OSError or ValueError for a file that cannot be read at all. + """ + with path.open("rb") as stream: + if stream.read(len(_PNG_SIGNATURE)) != _PNG_SIGNATURE: + stream.seek(0) + with Image.open(stream) as source: + width, height = source.size + return width, height, {k: source.info.get(k) for k in _CARD_RECORDS} + try: + inspection = read_inspection_file(path) + except MetadataReadError: + with path.open("rb") as stream, Image.open(stream) as source: + width, height = source.size + return width, height, None + records = {k: inspection.metadata.get(k) for k in _CARD_RECORDS} + return inspection.width, inspection.height, records def _unavailable_link() -> dict: @@ -2399,10 +2429,7 @@ class ImageRoutesMixin: records = None if not image["is_video"]: try: - with path.open("rb") as stream: - with Image.open(stream) as source: - image["width"], image["height"] = source.size - records = {k: source.info.get(k) for k in _CARD_RECORDS} + image["width"], image["height"], records = _card_header(path) except (OSError, ValueError, Image.DecompressionBombError): # A corrupt or disappearing preview must not hide the listing. pass diff --git a/tests/test_api_workspace_images.py b/tests/test_api_workspace_images.py index a4774e5..6964a4e 100644 --- a/tests/test_api_workspace_images.py +++ b/tests/test_api_workspace_images.py @@ -74,6 +74,43 @@ def records_png(path, records, mtime=100): return path +def png_chunk(kind, data): + """One PNG chunk: length, type, data and CRC.""" + import struct + import zlib + + return ( + struct.pack(">I", len(data)) + + kind + + data + + struct.pack(">I", zlib.crc32(kind + data)) + ) + + +def png_with_text_chunks(path, before=(), after=(), mtime=100): + """A PNG with ``(keyword, text)`` tEXt chunks before and after its IDAT.""" + import io + + buffer = io.BytesIO() + with Image.new("RGB", (4, 4)) as image: + image.save(buffer, format="PNG") + data = buffer.getvalue() + idat = data.index(b"IDAT") - 4 + iend = data.index(b"IEND") - 4 + + def chunks(records): + return b"".join( + png_chunk(b"tEXt", key.encode("latin-1") + b"\0" + text.encode("latin-1")) + for key, text in records + ) + + path.write_bytes( + data[:idat] + chunks(before) + data[idat:iend] + chunks(after) + data[iend:] + ) + os.utime(path, (mtime, mtime)) + return path + + def seeded_records(positive, seed): """API prompt records for a one-sampler graph with *positive* and *seed*.""" graph = { @@ -605,15 +642,54 @@ class TestWorkspaceImages(ImageAPITestCase): ) async def test_card_reads_each_header_once(self): + from py.api import images + for index in range(3): records_png( self.output_dir / f"{index}.png", seeded_records("a fox", index) ) self.media("clip.mp4") - with patch("PIL.Image.open", wraps=Image.open) as opened: + with ( + patch.object( + images, "read_inspection_file", wraps=images.read_inspection_file + ) as read, + patch("PIL.Image.open", wraps=Image.open) as opened, + ): _, data = await self.get() self.assertEqual(len(data["images"]), 4) - self.assertEqual(opened.call_count, 3) + # One inspector read per PNG (dimensions and records); no decoder open. + self.assertEqual(read.call_count, 3) + self.assertEqual(opened.call_count, 0) + + async def test_card_reads_png_records_as_the_inspector_does(self): + one, two = seeded_records("a fox", 1), seeded_records("a fox", 2) + # Conflicting duplicate prompt chunks: the inspector refuses the file, + # so the card shows no seed (Pillow alone would keep the last one). + png_with_text_chunks( + self.output_dir / "conflict.png", + before=[("prompt", one["prompt"]), ("prompt", two["prompt"])], + mtime=300, + ) + # Text after IDAT: Pillow's header open misses it; the inspector reads it. + png_with_text_chunks( + self.output_dir / "late.png", after=[("prompt", one["prompt"])], mtime=200 + ) + png_with_text_chunks( + self.output_dir / "repeat.png", + before=[("prompt", one["prompt"]), ("prompt", one["prompt"])], + ) + response, data = await self.get() + self.assertEqual(response.status, 200, data) + items = {item["filename"]: item for item in data["images"]} + self.assertIsNone(items["conflict.png"]["seed"]) + self.assertEqual( + (items["conflict.png"]["width"], items["conflict.png"]["height"]), (4, 4) + ) + self.assertEqual(items["late.png"]["seed"], "1") + self.assertEqual(items["repeat.png"]["seed"], "1") + shown = await self._inspector_seeds(data["images"]) + self.assertNotEqual(shown["conflict.png"]["status"], "known") + self.assertEqual(shown["late.png"], {"status": "known", "value": "1"}) async def test_deeply_nested_record_still_lists_and_inspects(self): nested = "[" * 900 + "1" + "]" * 900 @@ -699,18 +775,20 @@ class TestWorkspaceImages(ImageAPITestCase): async def test_card_facts_parse_after_the_file_is_closed(self): records_png(self.output_dir / "a.png", seeded_records("a fox", 5)) streams, closed = [], [] - real_open, real_parse = Image.open, header_params_of_images() + real_open, real_parse = Path.open, header_params_of_images() - def opening(stream, *args, **kwargs): - streams.append(stream) - return real_open(stream, *args, **kwargs) + def opening(path, *args, **kwargs): + stream = real_open(path, *args, **kwargs) + if path.name == "a.png": + streams.append(stream) + return stream def parsing(info): - closed.append(all(stream.closed for stream in streams)) + closed.append(bool(streams) and all(s.closed for s in streams)) return real_parse(info) with ( - patch("PIL.Image.open", side_effect=opening), + patch.object(Path, "open", autospec=True, side_effect=opening), patch("py.api.images.header_params", side_effect=parsing), ): _, data = await self.get()