From 14d7ae49972b3c4db7d69daca6ba9b206b1707eb Mon Sep 17 00:00:00 2001 From: d3x-at <70750382+d3x-at@users.noreply.github.com> Date: Wed, 2 Oct 2024 21:20:43 +0200 Subject: [PATCH] prepare pluggable png metadata extraction --- src/sd_parsers/_parser_manager.py | 21 +++++++++++---------- src/sd_parsers/parser.py | 4 ++-- src/sd_parsers/parsers/_automatic1111.py | 11 ++++++++--- src/sd_parsers/parsers/_comfyui.py | 10 +++++++--- src/sd_parsers/parsers/_dummy_parser.py | 13 ++++++++----- src/sd_parsers/parsers/_fooocus.py | 11 ++++++++--- src/sd_parsers/parsers/_invokeai/parser.py | 8 ++++++-- src/sd_parsers/parsers/_novelai.py | 10 +++++++--- 8 files changed, 57 insertions(+), 31 deletions(-) diff --git a/src/sd_parsers/_parser_manager.py b/src/sd_parsers/_parser_manager.py index 1cbd446..7ed2bdf 100644 --- a/src/sd_parsers/_parser_manager.py +++ b/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,7 +46,6 @@ class ParserManager: def __init__( self, *, - two_pass: bool = True, normalize_parameters: bool = True, managed_parsers: Optional[List[Type[Parser]]] = None, ): @@ -46,15 +53,10 @@ def __init__( 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: diff --git a/src/sd_parsers/parser.py b/src/sd_parsers/parser.py index a07c743..bd21d9c 100644 --- a/src/sd_parsers/parser.py +++ b/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. diff --git a/src/sd_parsers/parsers/_automatic1111.py b/src/sd_parsers/parsers/_automatic1111.py index 0d8540c..71e3314 100644 --- a/src/sd_parsers/parsers/_automatic1111.py +++ b/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: diff --git a/src/sd_parsers/parsers/_comfyui.py b/src/sd_parsers/parsers/_comfyui.py index c8da5c2..a97468f 100644 --- a/src/sd_parsers/parsers/_comfyui.py +++ b/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)} diff --git a/src/sd_parsers/parsers/_dummy_parser.py b/src/sd_parsers/parsers/_dummy_parser.py index 5cc0425..0bfa286 100644 --- a/src/sd_parsers/parsers/_dummy_parser.py +++ b/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( diff --git a/src/sd_parsers/parsers/_fooocus.py b/src/sd_parsers/parsers/_fooocus.py index 58b893f..0fc327a 100644 --- a/src/sd_parsers/parsers/_fooocus.py +++ b/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: diff --git a/src/sd_parsers/parsers/_invokeai/parser.py b/src/sd_parsers/parsers/_invokeai/parser.py index f36075b..98b1376 100644 --- a/src/sd_parsers/parsers/_invokeai/parser.py +++ b/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) diff --git a/src/sd_parsers/parsers/_novelai.py b/src/sd_parsers/parsers/_novelai.py index ba0cea7..ebdc566 100644 --- a/src/sd_parsers/parsers/_novelai.py +++ b/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"]