Preserve newlines when importing prompts from PNG metadata

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)

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:

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