Skip to content

Commit

Permalink
Merge pull request #100 from deepghs/fix/nai
Browse files Browse the repository at this point in the history
dev(narugo): fix nai decode error
  • Loading branch information
narugo1992 authored Sep 8, 2024
2 parents 51fe908 + 6218299 commit 28db387
Show file tree
Hide file tree
Showing 5 changed files with 17 additions and 2 deletions.
5 changes: 4 additions & 1 deletion imgutils/sd/nai/extract.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"))
Expand Down
7 changes: 6 additions & 1 deletion imgutils/sd/nai/metadata.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
import json
import os
import warnings
import zlib
from dataclasses import dataclass
from typing import Optional, Union

Expand Down Expand Up @@ -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 {}


Expand Down
7 changes: 7 additions & 0 deletions test/sd/test_nai.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Binary file added test/testfile/118438300_p1.png
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Binary file added test/testfile/118519492_p0.png
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.

0 comments on commit 28db387

Please sign in to comment.