diff --git a/imgutils/sd/nai.py b/imgutils/sd/nai.py index 5a120331ac..03d17bd1f6 100644 --- a/imgutils/sd/nai.py +++ b/imgutils/sd/nai.py @@ -189,15 +189,14 @@ def _save_gif_with_naimeta(image: Image.Image, dst_file: Union[str, os.PathLike] 'image/png': _save_png_with_naimeta, 'image/jpeg': _save_exif_with_naimeta, 'image/webp': _save_exif_with_naimeta, - 'image/tiff': _save_exif_with_naimeta, 'image/gif': _save_gif_with_naimeta, } -_LSB_ALLOWED_TYPES = {'image/png', 'image/tiff', 'image/gif', 'image/bmp'} +_LSB_ALLOWED_TYPES = {'image/png', 'image/tiff'} def save_image_with_naimeta( image: ImageTyping, dst_file: Union[str, os.PathLike], metadata: NAIMetadata, - add_lsb_meta: Union[str, bool] = 'auto', save_metainfo: bool = True, **kwargs) -> Image.Image: + add_lsb_meta: Union[str, bool] = 'auto', save_metainfo: Union[str, bool] = 'auto', **kwargs) -> Image.Image: mimetype, _ = mimetypes.guess_type(dst_file) if add_lsb_meta == 'auto': if mimetype in _LSB_ALLOWED_TYPES: @@ -206,9 +205,18 @@ def save_image_with_naimeta( add_lsb_meta = False else: if add_lsb_meta and mimetype not in _LSB_ALLOWED_TYPES: - raise ValueError('LSB metadata cannot be saved to lossy image format, ' + raise ValueError('LSB metadata cannot be saved to lossy image format or RGBA-incompatible format, ' 'add_lsb_meta will be disabled. ' f'Only {", ".join(sorted(_LSB_ALLOWED_TYPES))} images supported.') + if save_metainfo == 'auto': + if mimetype in _FN_IMG_SAVE: + save_metainfo = True + else: + save_metainfo = False + else: + if save_metainfo and mimetype not in _FN_IMG_SAVE: + raise SystemError(f'Not supported to save as a {mimetype!r} type, ' + f'supported mimetypes are {sorted(_FN_IMG_SAVE.keys())!r}.') if not add_lsb_meta and not save_metainfo: warnings.warn(f'Both LSB meta and pnginfo is disabled, no metadata will be saved to {dst_file!r}.') @@ -216,12 +224,7 @@ def save_image_with_naimeta( if add_lsb_meta: image = add_naimeta_to_image(image, metadata=metadata) if save_metainfo: - mimetype, _ = mimetypes.guess_type(dst_file) - if mimetype not in _FN_IMG_SAVE: - raise SystemError(f'Not supported to save as a {mimetype!r} type, ' - f'supported mimetypes are {sorted(_FN_IMG_SAVE.keys())!r}.') - else: - _FN_IMG_SAVE[mimetype](image, dst_file, metadata, **kwargs) + _FN_IMG_SAVE[mimetype](image, dst_file, metadata, **kwargs) else: image.save(dst_file, **kwargs) return image diff --git a/requirements.txt b/requirements.txt index 55b9abc1ad..e73a4cb83a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -15,4 +15,5 @@ pyclipper deprecation>=2.0.0 hfutils>=0.2.3 filelock -bchlib \ No newline at end of file +bchlib +piexif \ No newline at end of file diff --git a/test/sd/test_nai.py b/test/sd/test_nai.py index aa8cfedc48..e0929bf945 100644 --- a/test/sd/test_nai.py +++ b/test/sd/test_nai.py @@ -178,18 +178,6 @@ def test_add_naimeta_to_image_rgba(self, nai3_clear_rgba_image, nai3_meta_withou image = add_naimeta_to_image(nai3_clear_rgba_image, metadata=nai3_meta_without_title) assert get_naimeta_from_image(image) == pytest.approx(nai3_meta_without_title) - def test_save_image_with_naimeta_metainfo_only(self, nai3_clear_file, nai3_meta_without_title): - with isolated_directory(): - save_image_with_naimeta(nai3_clear_file, 'image.png', - metadata=nai3_meta_without_title, add_lsb_meta=False) - assert get_naimeta_from_image('image.png') == pytest.approx(nai3_meta_without_title) - - def test_save_image_with_naimeta_lsbmeta_only(self, nai3_clear_file, nai3_meta_without_title): - with isolated_directory(): - save_image_with_naimeta(nai3_clear_file, 'image.png', - metadata=nai3_meta_without_title, save_metainfo=False) - assert get_naimeta_from_image('image.png') == pytest.approx(nai3_meta_without_title) - def test_save_image_with_naimeta_both_no(self, nai3_clear_file, nai3_meta_without_title): with isolated_directory(): with pytest.warns(Warning): @@ -270,48 +258,54 @@ def test_save_image_with_naimeta_rgba(self, nai3_clear_rgba_file, nai3_meta_with assert get_naimeta_from_image(f'image{ext}') == \ (pytest.approx(nai3_meta_without_title) if okay else None) - @pytest.mark.parametrize(['ext'], [ - ('.webp',), - ('.jpg',), - ('.jpeg',), + @pytest.mark.parametrize(['ext', 'okay'], [ + ('.png', True), + ('.webp', False), + ('.jpg', False), + ('.jpeg', False), + ('.tiff', True), + ('.gif', False), ]) - def test_save_image_with_naimeta_exifs_lsb_true_lossy(self, nai3_clear_file, nai3_meta_without_title, ext): - with isolated_directory(), pytest.raises(ValueError): - save_image_with_naimeta(nai3_clear_file, f'image{ext}', - add_lsb_meta=True, metadata=nai3_meta_without_title) - - @pytest.mark.parametrize(['ext'], [ - ('.tiff',), - ('.gif',), - ]) - def test_save_image_with_naimeta_exifs_lsb_true_non_lossy(self, nai3_clear_file, nai3_meta_without_title, ext): + def test_save_image_with_naimeta_lsb_true(self, nai3_clear_file, nai3_meta_without_title, + ext, okay): with isolated_directory(): - save_image_with_naimeta(nai3_clear_file, f'image{ext}', - add_lsb_meta=True, metadata=nai3_meta_without_title) - assert get_naimeta_from_image(f'image{ext}') == pytest.approx(nai3_meta_without_title) - - @pytest.mark.parametrize(['ext'], [ - ('.webp',), - ('.jpg',), - ('.jpeg',), - ('.tiff',), - ('.gif',), + if okay: + save_image_with_naimeta(nai3_clear_file, f'image{ext}', + add_lsb_meta=True, metadata=nai3_meta_without_title) + assert get_naimeta_from_image(f'image{ext}') == pytest.approx(nai3_meta_without_title) + else: + with pytest.raises(ValueError): + save_image_with_naimeta(nai3_clear_file, f'image{ext}', + add_lsb_meta=True, metadata=nai3_meta_without_title) + + @pytest.mark.parametrize(['ext', 'warns', 'okay'], [ + ('.png', False, True), + ('.webp', False, True), + ('.jpg', False, True), + ('.jpeg', False, True), + ('.tiff', True, False), + ('.gif', False, True), ]) - def test_save_image_with_naimeta_metainfo_only_exifs(self, nai3_clear_file, nai3_meta_without_title, ext): - with isolated_directory(), pytest.warns(None): + def test_save_image_with_naimeta_metainfo_only(self, nai3_clear_file, nai3_meta_without_title, + ext, warns, okay): + with isolated_directory(), pytest.warns(Warning if warns else None): save_image_with_naimeta(nai3_clear_file, f'image{ext}', metadata=nai3_meta_without_title, add_lsb_meta=False) - assert get_naimeta_from_image(f'image{ext}') == pytest.approx(nai3_meta_without_title) - - @pytest.mark.parametrize(['ext'], [ - ('.webp',), - ('.jpg',), - ('.jpeg',), - ('.tiff',), - ('.gif',), + assert get_naimeta_from_image(f'image{ext}') == \ + (pytest.approx(nai3_meta_without_title) if okay else None) + + @pytest.mark.parametrize(['ext', 'warns', 'okay'], [ + ('.png', False, True), + ('.webp', True, False), + ('.jpg', True, False), + ('.jpeg', True, False), + ('.tiff', False, True), + ('.gif', True, False), ]) - def test_save_image_with_naimeta_lsbmeta_only_exifs(self, nai3_clear_file, nai3_meta_without_title, ext): - with isolated_directory(), pytest.warns(Warning): + def test_save_image_with_naimeta_lsbmeta_only(self, nai3_clear_file, nai3_meta_without_title, + ext, warns, okay): + with isolated_directory(), pytest.warns(Warning if warns else None): save_image_with_naimeta(nai3_clear_file, f'image{ext}', metadata=nai3_meta_without_title, save_metainfo=False) - assert get_naimeta_from_image(f'image{ext}') is None + assert get_naimeta_from_image(f'image{ext}') == \ + (pytest.approx(nai3_meta_without_title) if okay else None)