Loading imgutils/sd/nai/extract.py +4 −1 Original line number Diff line number Diff line Loading @@ -164,7 +164,10 @@ class ImageLsbDataExtractor(object): raise ValueError(f'Image magic number mismatch, ' f'{self._magic_bytes!r} expected but {read_magic!r}.') read_len = reader.read_32bit_integer() // 8 next_int = reader.read_32bit_integer() if next_int is None: raise ValueError('No next int32 to read.') read_len = next_int // 8 json_data = reader.get_next_n_bytes(read_len) json_data = json.loads(gzip.decompress(json_data).decode("utf-8")) Loading imgutils/sd/nai/metadata.py +6 −1 Original line number Diff line number Diff line Loading @@ -15,6 +15,7 @@ This module is particularly useful for working with AI-generated images and thei import json import os import warnings import zlib from dataclasses import dataclass from typing import Optional, Union Loading Loading @@ -95,7 +96,11 @@ def _get_naimeta_raw(image: ImageTyping) -> dict: image = load_image(image, force_background=None, mode=None) try: return ImageLsbDataExtractor().extract_data(image) except (ValueError, json.JSONDecodeError): except (ValueError, json.JSONDecodeError, zlib.error, OSError, UnicodeDecodeError): # ValueError: binary data with wrong format # json.JSONDecodeError: zot a json-formatted data # zlib.error, OSError: not zlib compressed binary data # UnicodeDecodeError: cannot decode as utf-8 text return image.info or {} Loading test/sd/test_nai.py +3 −0 Original line number Diff line number Diff line Loading @@ -205,3 +205,6 @@ class TestSDNai: save_pnginfo=False, add_lsb_meta=False, ) assert get_naimeta_from_image('image.png') is None def test_image_error_with_wrong_format(self): assert get_naimeta_from_image(get_testfile('118519492_p0.png')) is None test/testfile/118519492_p0.png 0 → 100644 +1.33 MiB Loading image diff... Loading
imgutils/sd/nai/extract.py +4 −1 Original line number Diff line number Diff line Loading @@ -164,7 +164,10 @@ class ImageLsbDataExtractor(object): raise ValueError(f'Image magic number mismatch, ' f'{self._magic_bytes!r} expected but {read_magic!r}.') read_len = reader.read_32bit_integer() // 8 next_int = reader.read_32bit_integer() if next_int is None: raise ValueError('No next int32 to read.') read_len = next_int // 8 json_data = reader.get_next_n_bytes(read_len) json_data = json.loads(gzip.decompress(json_data).decode("utf-8")) Loading
imgutils/sd/nai/metadata.py +6 −1 Original line number Diff line number Diff line Loading @@ -15,6 +15,7 @@ This module is particularly useful for working with AI-generated images and thei import json import os import warnings import zlib from dataclasses import dataclass from typing import Optional, Union Loading Loading @@ -95,7 +96,11 @@ def _get_naimeta_raw(image: ImageTyping) -> dict: image = load_image(image, force_background=None, mode=None) try: return ImageLsbDataExtractor().extract_data(image) except (ValueError, json.JSONDecodeError): except (ValueError, json.JSONDecodeError, zlib.error, OSError, UnicodeDecodeError): # ValueError: binary data with wrong format # json.JSONDecodeError: zot a json-formatted data # zlib.error, OSError: not zlib compressed binary data # UnicodeDecodeError: cannot decode as utf-8 text return image.info or {} Loading
test/sd/test_nai.py +3 −0 Original line number Diff line number Diff line Loading @@ -205,3 +205,6 @@ class TestSDNai: save_pnginfo=False, add_lsb_meta=False, ) assert get_naimeta_from_image('image.png') is None def test_image_error_with_wrong_format(self): assert get_naimeta_from_image(get_testfile('118519492_p0.png')) is None