diff --git a/ai_diffusion/image.py b/ai_diffusion/image.py index 12f896656..bbd73c0e4 100644 --- a/ai_diffusion/image.py +++ b/ai_diffusion/image.py @@ -492,6 +492,52 @@ def save_png_w_itxt(img_path: str | Path, png_data: bytes, keyword: str, text: s f.write(struct.pack(">I", zlib.crc32(b"iTXt" + itxt_data) & 0xFFFFFFFF)) ihdr_inserted = True + @staticmethod + def read_png_text(img_path: str | Path) -> dict[str, str]: + """Read tEXt/zTXt/iTXt chunks from a PNG file, keyed by keyword. + + QImageReader.text() collapses newlines in PNG text chunks, which destroys the line + structure of prompts written by other tools. Parsing the chunks directly keeps the + text byte-for-byte as it was written. + """ + result: dict[str, str] = {} + data = Path(img_path).read_bytes() + if data[:8] != b"\x89PNG\r\n\x1a\n": + return result + + offset = 8 + while offset + 8 <= len(data): + length = struct.unpack(">I", data[offset : offset + 4])[0] + chunk_type = data[offset + 4 : offset + 8] + chunk_data = data[offset + 8 : offset + 8 + length] + offset += 12 + length # length + type + data + crc + + if chunk_type == b"IEND": + break + if chunk_type not in (b"tEXt", b"zTXt", b"iTXt"): + continue + + try: + keyword, rest = chunk_data.split(b"\x00", 1) + if chunk_type == b"tEXt": + text = rest.decode("utf-8", errors="replace") + elif chunk_type == b"zTXt": + # rest = compression method (1 byte) + compressed text + text = zlib.decompress(rest[1:]).decode("utf-8", errors="replace") + else: # iTXt + # rest = flag + method + language\0 + translated keyword\0 + text + compressed = rest[0] == 1 + _, _, tail = rest[2:].split(b"\x00", 2) + if compressed: + tail = zlib.decompress(tail) + text = tail.decode("utf-8", errors="replace") + result[keyword.decode("latin1")] = text + except Exception as e: + log.warning(f"Skipping malformed PNG text chunk in {img_path}: {e}") + continue # keep reading the rest of the file + + return result + @classmethod def mask_subtract(cls, lhs: Image, rhs: Image): return cls._mask_op(rhs, lhs, QPainter.CompositionMode.CompositionMode_SourceOut) diff --git a/ai_diffusion/persistence.py b/ai_diffusion/persistence.py index 0535255a2..03a2df701 100644 --- a/ai_diffusion/persistence.py +++ b/ai_diffusion/persistence.py @@ -13,7 +13,7 @@ from . import eventloop from .backend.api import FillMode, InpaintMode -from .image import ImageCollection +from .image import Image, ImageCollection from .localization import translate as _ from .model.control import ControlLayer, ControlLayerList from .model.custom_workflow import CustomWorkspace @@ -345,20 +345,33 @@ def _find_annotation(document, name: str): return None +def _read_image_text(filename: str) -> dict[str, str]: + """Text chunks of an image, keyed by keyword ('parameters', 'prompt', ...). + + PNGs are parsed directly: QImageReader collapses newlines in text chunks, which would + flatten a multi-line prompt into a single paragraph. Other formats fall back to Qt. + """ + if filename.lower().endswith(".png"): + if text := Image.read_png_text(filename): + return text + reader = QImageReader(filename) + return {key: reader.text(key) for key in reader.textKeys()} + + def import_prompt_from_file(model: DocumentModel): exts = (".png", ".jpg", ".jpeg", ".webp") filename = model.document.filename if model.regions.positive == "" and model.regions.negative == "" and filename.endswith(exts): try: - reader = QImageReader(filename) + image_text = _read_image_text(filename) # A1111 - if text := reader.text("parameters"): + if text := image_text.get("parameters"): if "Negative prompt:" in text: positive, negative = text.split("Negative prompt:", 1) model.regions.positive = positive.strip() model.regions.negative = negative.split("Steps:", 1)[0].strip() # ComfyUI - elif text := reader.text("prompt"): + elif text := image_text.get("prompt"): prompt: dict[str, dict] = json.loads(text) for node in prompt.values(): if node["class_type"] in _comfy_sampler_types: diff --git a/tests/test_image.py b/tests/test_image.py index 38ef98b5f..7ddbb7724 100644 --- a/tests/test_image.py +++ b/tests/test_image.py @@ -406,3 +406,14 @@ def test_save_png_with_metadata(tmp_path): data = file_path.read_bytes() assert data.startswith(b"\x89PNG\r\n\x1a\n") assert b"my test metadata in the png" in data + + +def test_read_png_text_preserves_newlines(tmp_path): + # Newlines in PNG text chunks must survive: QImageReader.text() collapses them, so the + # plugin parses the chunks directly instead. + img = Image.create(Extent(2, 2), Qt.GlobalColor.red) + file_path = tmp_path / "test_newlines.png" + text = "line one\nline two\nline three" + img.save_png_with_metadata(file_path, text) + + assert Image.read_png_text(file_path)["parameters"] == text