diff --git a/modules/images.py b/modules/images.py index b2de3662..4e8cd993 100644 --- a/modules/images.py +++ b/modules/images.py @@ -482,6 +482,43 @@ def get_next_sequence_number(path, basename): return result + 1 +def save_image_with_geninfo(image, geninfo, filename, extension=None, existing_pnginfo=None): + if extension is None: + extension = os.path.splitext(filename)[1] + + image_format = Image.registered_extensions()[extension] + + existing_pnginfo = existing_pnginfo or {} + if opts.enable_pnginfo: + existing_pnginfo['parameters'] = geninfo + + if extension.lower() == '.png': + pnginfo_data = PngImagePlugin.PngInfo() + for k, v in (existing_pnginfo or {}).items(): + pnginfo_data.add_text(k, str(v)) + + image.save(filename, format=image_format, quality=opts.jpeg_quality, pnginfo=pnginfo_data) + + elif extension.lower() in (".jpg", ".jpeg", ".webp"): + if image.mode == 'RGBA': + image = image.convert("RGB") + elif image.mode == 'I;16': + image = image.point(lambda p: p * 0.0038910505836576).convert("RGB" if extension.lower() == ".webp" else "L") + + image.save(filename, format=image_format, quality=opts.jpeg_quality, lossless=opts.webp_lossless) + + if opts.enable_pnginfo and geninfo is not None: + exif_bytes = piexif.dump({ + "Exif": { + piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(geninfo or "", encoding="unicode") + }, + }) + + piexif.insert(exif_bytes, filename) + else: + image.save(filename, format=image_format, quality=opts.jpeg_quality) + + def save_image(image, path, basename, seed=None, prompt=None, extension='png', info=None, short_filename=False, no_prompt=False, grid=False, pnginfo_section_name='parameters', p=None, existing_info=None, forced_filename=None, suffix="", save_to_dirs=None): """Save an image. @@ -566,38 +603,13 @@ def save_image(image, path, basename, seed=None, prompt=None, extension='png', i info = params.pnginfo.get(pnginfo_section_name, None) def _atomically_save_image(image_to_save, filename_without_extension, extension): - # save image with .tmp extension to avoid race condition when another process detects new image in the directory + """ + save image with .tmp extension to avoid race condition when another process detects new image in the directory + """ temp_file_path = f"{filename_without_extension}.tmp" - image_format = Image.registered_extensions()[extension] - if extension.lower() == '.png': - pnginfo_data = PngImagePlugin.PngInfo() - if opts.enable_pnginfo: - for k, v in params.pnginfo.items(): - pnginfo_data.add_text(k, str(v)) + save_image_with_geninfo(image_to_save, info, temp_file_path, extension, params.pnginfo) - image_to_save.save(temp_file_path, format=image_format, quality=opts.jpeg_quality, pnginfo=pnginfo_data) - - elif extension.lower() in (".jpg", ".jpeg", ".webp"): - if image_to_save.mode == 'RGBA': - image_to_save = image_to_save.convert("RGB") - elif image_to_save.mode == 'I;16': - image_to_save = image_to_save.point(lambda p: p * 0.0038910505836576).convert("RGB" if extension.lower() == ".webp" else "L") - - image_to_save.save(temp_file_path, format=image_format, quality=opts.jpeg_quality, lossless=opts.webp_lossless) - - if opts.enable_pnginfo and info is not None: - exif_bytes = piexif.dump({ - "Exif": { - piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(info or "", encoding="unicode") - }, - }) - - piexif.insert(exif_bytes, temp_file_path) - else: - image_to_save.save(temp_file_path, format=image_format, quality=opts.jpeg_quality) - - # atomically rename the file with correct extension os.replace(temp_file_path, filename_without_extension + extension) fullfn_without_extension, extension = os.path.splitext(params.filename) diff --git a/modules/ui_extra_networks.py b/modules/ui_extra_networks.py index 471df23b..c6e45fb1 100644 --- a/modules/ui_extra_networks.py +++ b/modules/ui_extra_networks.py @@ -4,7 +4,7 @@ from pathlib import Path from PIL import PngImagePlugin from modules import shared -from modules.images import read_info_from_image +from modules.images import read_info_from_image, save_image_with_geninfo import gradio as gr import json import html @@ -343,22 +343,7 @@ def setup_ui(ui, gallery): assert is_allowed, f'writing to {filename} is not allowed' - if geninfo: - ext = os.path.splitext(filename)[1].lower() - if ext == '.png': - pnginfo_data = PngImagePlugin.PngInfo() - pnginfo_data.add_text('parameters', geninfo) - image.save(filename, pnginfo=pnginfo_data) - elif ext in ('.jpg', '.jpeg', '.webp'): - exif_bytes = piexif.dump({ - 'Exif': {piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(geninfo or '', - encoding='unicode')} - }) - image.save(filename, exif=exif_bytes, quality=shared.opts.jpeg_quality) - else: - image.save(filename) - else: - image.save(filename) + save_image_with_geninfo(image, geninfo, filename) return [page.create_html(ui.tabname) for page in ui.stored_extra_pages]