diff --git a/imgutils/sd/nai/extract.py b/imgutils/sd/nai/extract.py index 29efbf457f3..17cb4f0eee4 100644 --- a/imgutils/sd/nai/extract.py +++ b/imgutils/sd/nai/extract.py @@ -164,7 +164,10 @@ def extract_data(self, image: Image.Image) -> dict: 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")) diff --git a/imgutils/sd/nai/metadata.py b/imgutils/sd/nai/metadata.py index 3c05969d4b0..91b14320240 100644 --- a/imgutils/sd/nai/metadata.py +++ b/imgutils/sd/nai/metadata.py @@ -15,6 +15,7 @@ import json import os import warnings +import zlib from dataclasses import dataclass from typing import Optional, Union @@ -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 {} diff --git a/test/sd/test_nai.py b/test/sd/test_nai.py index 181610f1fd7..f2d51b2b226 100644 --- a/test/sd/test_nai.py +++ b/test/sd/test_nai.py @@ -205,3 +205,10 @@ def test_save_image_with_naimeta_both_no_with_title(self, nai3_clear_file, nai3_ save_pnginfo=False, add_lsb_meta=False, ) assert get_naimeta_from_image('image.png') is None + + @pytest.mark.parametrize(['file'], [ + ('118519492_p0.png',), + ('118438300_p1.png',), + ]) + def test_image_error_with_wrong_format(self, file): + assert get_naimeta_from_image(get_testfile(file)) is None diff --git a/test/testfile/118438300_p1.png b/test/testfile/118438300_p1.png new file mode 100644 index 00000000000..994a1b74cbd Binary files /dev/null and b/test/testfile/118438300_p1.png differ diff --git a/test/testfile/118519492_p0.png b/test/testfile/118519492_p0.png new file mode 100644 index 00000000000..f0195e925b3 Binary files /dev/null and b/test/testfile/118519492_p0.png differ