prepare pluggable png metadata extraction

src/sd_parsers/_parser_manager.py

@@ -4,7 +4,7 @@
 
 import logging
 from contextlib import contextmanager
-from typing import TYPE_CHECKING, Callable, List, Optional, Type, Union
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Type, Union
 
 from PIL import Image
 
@@ -20,6 +20,14 @@
 
 logger = logging.getLogger(__name__)
 
+GET_PNG_METADATA: List[Callable[[Image.Image], Dict[str, Any]]] = [
+    # use image.info
+    lambda i: i.info,
+    # use image.text property (iTxt, tEXt and zTXt chunks may appear at the end of the file)
+    lambda i: i.text,  # type: ignore
+]
+"""A list of metadata retrieval functions to provide multiple metadata entrypoints for each parser module."""
+
 
 @contextmanager
 def _get_image(image: Union[str, bytes, Path, SupportsRead[bytes], Image.Image]):
@@ -38,23 +46,17 @@ class ParserManager:
     def __init__(
         self,
         *,
-        two_pass: bool = True,
         normalize_parameters: bool = True,
         managed_parsers: Optional[List[Type[Parser]]] = None,
     ):
         """
         Initializes a ParserManager object.
 
         Optional Parameters:
-            two_pass: for PNG images, use `Image.info` before using `Image.text` as metadata source.
             normalize_parameters: Try to unify the parameter keys of the parser outputs.
             managed_parsers: A list of parsers to be managed.
 
-        The performance effects of two-pass parsing depends on the given image files.
-        If the image files are correctly formed and can be read with one of the supported parser modules,
-        setting `two_pass` to `True` will considerably shorten the time needed to read the image parameters.
         """
-        self.two_pass = two_pass
         self.managed_parsers: List[Parser] = [
             parser(normalize_parameters) for parser in managed_parsers or MANAGED_PARSERS
         ]
@@ -112,14 +114,13 @@ def _read_parameters(
     ):
         with _get_image(image) as image:
             # two_pass only makes sense with PNG images
-            two_pass = image.format == "PNG" if self.two_pass else False
 
-            for use_text in [False, True] if two_pass else [True]:
+            for get_metadata in GET_PNG_METADATA:
                 for parser in (
                     sorted(self.managed_parsers, key=key) if key else self.managed_parsers
                 ):
                     try:
-                        parameters, parsing_context = parser.read_parameters(image, use_text)
+                        parameters, parsing_context = parser.read_parameters(image, get_metadata)
                         yield parser, parameters, parsing_context
 
                     except ParserError as error:

src/sd_parsers/parser.py

@@ -4,7 +4,7 @@
 
 from abc import ABC, abstractmethod
 from contextlib import suppress
-from typing import Any, Dict, Iterable, Iterator, List, Optional, Tuple, Union
+from typing import Any, Callable, Dict, Iterable, Iterator, List, Optional, Tuple, Union
 
 from PIL import ExifTags
 from PIL.Image import Image
@@ -40,7 +40,7 @@ def __init__(self, normalize_parameters: bool = True):
     def read_parameters(
         self,
         image: Image,
-        use_text: bool = True,
+        get_png_metadata: Callable[[Image], Dict[str, Any]] | None = None,
     ) -> Tuple[dict[str, Any], Any]:
         """
         Read generation parameters from image.

src/sd_parsers/parsers/_automatic1111.py

@@ -3,7 +3,7 @@
 import json
 import re
 from contextlib import suppress
-from typing import Any, Dict
+from typing import Any, Callable, Dict
 
 from PIL.Image import Image
 
@@ -21,10 +21,15 @@ class AUTOMATIC1111Parser(Parser):
 
     _COMPLEXITY_INDEX = 100
 
-    def read_parameters(self, image: Image, use_text: bool = True):
+    def read_parameters(
+        self,
+        image: Image,
+        get_png_metadata: Callable[[Image], Dict[str, Any]] | None = None,
+    ):
         try:
             if image.format == "PNG":
-                parameters = image.text["parameters"] if use_text else image.info["parameters"]  # type: ignore
+                metadata = get_png_metadata(image) if get_png_metadata else image.info
+                parameters = metadata["parameters"]
             elif image.format in ("JPEG", "WEBP"):
                 parameters = get_exif_value(image, "UserComment")
             else:

src/sd_parsers/parsers/_comfyui.py

@@ -4,7 +4,7 @@
 import logging
 from collections import defaultdict
 from contextlib import suppress
-from typing import Any, Dict, Generator, List, Optional, Set, Tuple
+from typing import Any, Callable, Dict, Generator, List, Optional, Set, Tuple
 
 from PIL.Image import Image
 
@@ -25,12 +25,16 @@
 class ComfyUIParser(Parser):
     """Parser for images generated by ComfyUI"""
 
-    def read_parameters(self, image: Image, use_text: bool = True):
+    def read_parameters(
+        self,
+        image: Image,
+        get_png_metadata: Callable[[Image], Dict[str, Any]] | None = None,
+    ):
         if image.format != "PNG":
             raise MetadataError("unsupported image format", image.format)
 
         try:
-            metadata = image.text if use_text else image.info  # type: ignore
+            metadata = get_png_metadata(image) if get_png_metadata else image.info
             prompt = metadata["prompt"]
             workflow = metadata["workflow"]
             parameters = {"prompt": json.loads(prompt), "workflow": json.loads(workflow)}

src/sd_parsers/parsers/_dummy_parser.py

@@ -1,7 +1,7 @@
 """Example stub for additional parsers"""
 
 import json
-from typing import Any, Dict
+from typing import Any, Callable, Dict
 
 from PIL.Image import Image
 
@@ -15,7 +15,11 @@ class DummyParser(Parser):
     Example stub for additional parsers
     """
 
-    def read_parameters(self, image: Image, use_text: bool = True):
+    def read_parameters(
+        self,
+        image: Image,
+        get_png_metadata: Callable[[Image], Dict[str, Any]] | None = None,
+    ):
         """
         Read the relevant generation parameters from the given image.
 
@@ -42,9 +46,8 @@ def read_parameters(self, image: Image, use_text: bool = True):
                 parsing_context["parameters_key"] = "user_comment"
 
             elif image.format == "PNG":
-                # Use `image.text` as parameters source if use_text is True.
-                # Use `image.info` otherwise.
-                metadata = image.text if use_text else image.info  # type: ignore
+                # use metadata retrieval function if given, otherwise the Image.info field
+                metadata = get_png_metadata(image) if get_png_metadata else image.info
 
                 # deserialize parameters in json format
                 parameters["some_image_parameter"] = json.loads(

src/sd_parsers/parsers/_fooocus.py

@@ -2,7 +2,7 @@
 
 import copy
 import json
-from typing import Any, Dict
+from typing import Any, Callable, Dict
 
 from PIL.Image import Image
 
@@ -20,10 +20,15 @@ class FooocusParser(Parser):
 
     _COMPLEXITY_INDEX = 90
 
-    def read_parameters(self, image: Image, use_text: bool = True):
+    def read_parameters(
+        self,
+        image: Image,
+        get_png_metadata: Callable[[Image], Dict[str, Any]] | None = None,
+    ):
         try:
             if image.format == "PNG":
-                parameters = image.text["parameters"] if use_text else image.info["parameters"]  # type: ignore
+                metadata = get_png_metadata(image) if get_png_metadata else image.info
+                parameters = metadata["parameters"]
             elif image.format in ("JPEG", "WEBP"):
                 parameters = get_exif_value(image, "UserComment")
             else:

src/sd_parsers/parsers/_invokeai/parser.py

@@ -43,11 +43,15 @@ class VariantParser(NamedTuple):
 class InvokeAIParser(Parser):
     """parser for images generated by invokeai"""
 
-    def read_parameters(self, image: Image, use_text: bool = True):
+    def read_parameters(
+        self,
+        image: Image,
+        get_png_metadata: Callable[[Image], Dict[str, Any]] | None = None,
+    ):
         if image.format != "PNG":
             raise MetadataError("unsupported image format", image.format)
 
-        metadata = image.text if use_text else image.info  # type: ignore
+        metadata = get_png_metadata(image) if get_png_metadata else image.info
         for variant in VARIANT_PARSERS:
             try:
                 parameters = variant.read_parameters(metadata)

src/sd_parsers/parsers/_novelai.py

@@ -4,7 +4,7 @@
 import json
 import re
 from contextlib import suppress
-from typing import Any, Dict
+from typing import Any, Callable, Dict
 
 from PIL.Image import Image
 
@@ -19,11 +19,15 @@
 class NovelAIParser(Parser):
     """parser for images generated by NovelAI"""
 
-    def read_parameters(self, image: Image, use_text: bool = True):
+    def read_parameters(
+        self,
+        image: Image,
+        get_png_metadata: Callable[[Image], Dict[str, Any]] | None = None,
+    ):
         if image.format != "PNG":
             raise MetadataError("unsupported image format", image.format)
 
-        metadata = image.text if use_text else image.info  # type: ignore
+        metadata = get_png_metadata(image) if get_png_metadata else image.info
         try:
             description = metadata["Description"]
             software = metadata["Software"]