dev(narugo): add support for webp metadata read

imgutils/metadata/__init__.py

@@ -1 +1,2 @@
+from .geninfo import read_geninfo_gif, read_geninfo_parameters, read_geninfo_exif
 from .lsb import read_lsb_raw_bytes, read_lsb_metadata, write_lsb_raw_bytes, write_lsb_metadata, LSBReadError

imgutils/metadata/geninfo.py

@@ -0,0 +1,43 @@
+from typing import Optional
+
+import piexif
+from piexif.helper import UserComment
+
+from ..data import ImageTyping, load_image
+
+
+def read_geninfo_parameters(image: ImageTyping) -> Optional[str]:
+    image = load_image(image, mode=None, force_background=None)
+    infos = image.info or {}
+    return infos.get('parameters')
+
+
+def read_geninfo_exif(image: ImageTyping) -> Optional[str]:
+    image = load_image(image, mode=None, force_background=None)
+    infos = image.info or {}
+    if "exif" in infos:
+        exif_data = infos["exif"]
+        try:
+            exif = piexif.load(exif_data)
+        except OSError:
+            # memory / exif was not valid so piexif tried to read from a file
+            exif = None
+
+        exif_comment = (exif or {}).get("Exif", {}).get(piexif.ExifIFD.UserComment, b"")
+        try:
+            exif_comment = UserComment.load(exif_comment)
+        except ValueError:
+            exif_comment = exif_comment.decode("utf8", errors="ignore")
+
+        return exif_comment
+    else:
+        return None
+
+
+def read_geninfo_gif(image: ImageTyping) -> Optional[str]:
+    image = load_image(image, mode=None, force_background=None)
+    infos = image.info or {}
+    if "comment" in infos:  # for gif
+        return infos["comment"].decode("utf8", errors="ignore")
+    else:
+        return None

imgutils/sd/metadata.py

@@ -15,6 +15,7 @@
 from PIL.PngImagePlugin import PngInfo
 
 from ..data import ImageTyping, load_image
+from ..metadata import read_geninfo_parameters, read_geninfo_exif
 
 _PARAM_PATTERN = re.compile(r'\s*(?P<key>[\w ]+):\s*(?P<value>"(?:\\.|[^\\"])+"|[^,]*)(?:,|$)')
 _SIZE_PATTERN = re.compile(r"^(?P<size1>-?\d+)\s*x\s*(?P<size2>-?\d+)$")
@@ -256,7 +257,9 @@ def get_sdmeta_from_image(image: ImageTyping) -> Optional[SDMetaData]:
         <class 'imgutils.sd.metadata.SDMetaData'>
     """
     image = load_image(image, mode=None, force_background=None)
-    pnginfo_text = image.info.get('parameters')
+    pnginfo_text = (read_geninfo_parameters(image) or
+                    read_geninfo_exif(image) or
+                    read_geninfo_parameters(image))
     if pnginfo_text:
         return parse_sdmeta_from_text(pnginfo_text)
     else:

imgutils/sd/nai.py

@@ -21,8 +21,9 @@
 from PIL import Image
 from PIL.PngImagePlugin import PngInfo
 
-from imgutils.data import load_image, ImageTyping
-from imgutils.metadata import read_lsb_metadata, write_lsb_metadata, LSBReadError
+from ..data import load_image, ImageTyping
+from ..metadata import read_lsb_metadata, write_lsb_metadata, LSBReadError, read_geninfo_parameters, \
+    read_geninfo_exif, read_geninfo_gif
 
 
 @dataclass
@@ -78,6 +79,17 @@ def pnginfo(self) -> PngInfo:
         return info
 
 
+class _InvalidNAIMetaError(Exception):
+    pass
+
+
+def _naimeta_validate(data):
+    if isinstance(data, dict) and data.get('Software') and data.get('Source') and data.get('Comment'):
+        return data
+    else:
+        raise _InvalidNAIMetaError
+
+
 def _get_naimeta_raw(image: ImageTyping) -> dict:
     """
     Extract raw NAI metadata from an image.
@@ -93,9 +105,29 @@ def _get_naimeta_raw(image: ImageTyping) -> dict:
     """
     image = load_image(image, force_background=None, mode=None)
     try:
-        return read_lsb_metadata(image)
-    except LSBReadError:
-        return image.info or {}
+        return _naimeta_validate(read_lsb_metadata(image))
+    except (LSBReadError, _InvalidNAIMetaError):
+        pass
+
+    try:
+        return _naimeta_validate(image.info or {})
+    except (LSBReadError, _InvalidNAIMetaError):
+        pass
+
+    try:
+        return _naimeta_validate(json.loads(read_geninfo_parameters(image)))
+    except (TypeError, json.JSONDecodeError, _InvalidNAIMetaError):
+        pass
+
+    try:
+        return _naimeta_validate(json.loads(read_geninfo_exif(image)))
+    except (TypeError, json.JSONDecodeError, _InvalidNAIMetaError):
+        pass
+
+    try:
+        return _naimeta_validate(json.loads(read_geninfo_gif(image)))
+    except (TypeError, json.JSONDecodeError, _InvalidNAIMetaError):
+        raise _InvalidNAIMetaError
 
 
 def get_naimeta_from_image(image: ImageTyping) -> Optional[NAIMetadata]:
@@ -111,8 +143,11 @@ def get_naimeta_from_image(image: ImageTyping) -> Optional[NAIMetadata]:
     :return: A NAIMetadata object if successful, None otherwise.
     :rtype: Optional[NAIMetadata]
     """
-    data = _get_naimeta_raw(image)
-    if data.get('Software') and data.get('Source') and data.get('Comment'):
+    try:
+        data = _get_naimeta_raw(image)
+    except _InvalidNAIMetaError:
+        return None
+    else:
         return NAIMetadata(
             software=data['Software'],
             source=data['Source'],
@@ -121,8 +156,6 @@ def get_naimeta_from_image(image: ImageTyping) -> Optional[NAIMetadata]:
             generation_time=float(data['Generation time']) if data.get('Generation time') else None,
             description=data.get('Description'),
         )
-    else:
-        return None
 
 
 def _get_pnginfo(metadata: Union[NAIMetadata, PngInfo]) -> PngInfo:

test/sd/test_nai.py

@@ -6,6 +6,11 @@
 from ..testings import get_testfile
 
 
+@pytest.fixture()
+def nai3_webp_file():
+    return get_testfile('nai3_webp.webp')
+
+
 @pytest.fixture()
 def nai3_file():
     return get_testfile('nai3.png')
@@ -40,6 +45,37 @@ def nai3_clear_rgba_image():
     return image
 
 
+@pytest.fixture()
+def nai3_webp_meta():
+    return NAIMetadata(
+        software='NovelAI',
+        source='Stable Diffusion XL C1E1DE52',
+        parameters={
+            'prompt': '2girls,side-by-side,nekomata okayu,shiina mahiru,symmetrical pose,general,masterpiece,, '
+                      'best quality, amazing quality, very aesthetic, absurdres',
+            'steps': 28, 'height': 832, 'width': 1216, 'scale': 5.0, 'uncond_scale': 0.0, 'cfg_rescale': 0.0,
+            'seed': 210306140, 'n_samples': 1, 'hide_debug_overlay': False, 'noise_schedule': 'native',
+            'legacy_v3_extend': False, 'reference_information_extracted_multiple': [],
+            'reference_strength_multiple': [], 'sampler': 'k_euler_ancestral', 'controlnet_strength': 1.0,
+            'controlnet_model': None, 'dynamic_thresholding': False, 'dynamic_thresholding_percentile': 0.999,
+            'dynamic_thresholding_mimic_scale': 10.0, 'sm': False, 'sm_dyn': False, 'skip_cfg_above_sigma': None,
+            'skip_cfg_below_sigma': 0.0, 'lora_unet_weights': None, 'lora_clip_weights': None,
+            'deliberate_euler_ancestral_bug': True, 'prefer_brownian': False,
+            'cfg_sched_eligibility': 'enable_for_post_summer_samplers', 'explike_fine_detail': False,
+            'minimize_sigma_inf': False, 'uncond_per_vibe': True, 'wonky_vibe_correlation': True, 'version': 1,
+            'uc': 'lowres, {bad}, error, fewer, extra, missing, worst quality, jpeg artifacts, bad quality, '
+                  'watermark, unfinished, displeasing, chromatic aberration, signature, extra digits, '
+                  'artistic error, username, scan, [abstract], ',
+            'request_type': 'PromptGenerateRequest',
+            'signed_hash': 'nM6vZLFGJWW7SH2xc4lpRY9sJGbPQKXaUzhUVX/u2NvCAyLg9abn90XBCiNmwqh1hK5hk+o7wYHkPJvhkfAnBg=='
+        },
+        title=None,
+        generation_time=6.494704299024306,
+        description='2girls,side-by-side,nekomata okayu,shiina mahiru,symmetrical pose,general,masterpiece,, '
+                    'best quality, amazing quality, very aesthetic, absurdres'
+    )
+
+
 @pytest.fixture()
 def nai3_meta_without_title():
     return NAIMetadata(
@@ -212,3 +248,6 @@ def test_save_image_with_naimeta_both_no_with_title(self, nai3_clear_file, nai3_
     ])
     def test_image_error_with_wrong_format(self, file):
         assert get_naimeta_from_image(get_testfile(file)) is None
+
+    def test_get_naimeta_from_image_webp(self, nai3_webp_file, nai3_webp_meta):
+        assert get_naimeta_from_image(nai3_webp_file) == pytest.approx(nai3_webp_meta)