Compare commits
2
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
d92b2eabd9 | ||
|
|
e5868d96f2 |
+13
@@ -445,6 +445,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
|
||||
|
||||
+5
-1
@@ -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}"
|
||||
|
||||
+3
-3
@@ -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
|
||||
|
||||
|
||||
@@ -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", "<x:xmpmeta/>"),
|
||||
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()
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user