StableDiffusion/WebUI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

use a single function for saving images with metadata both in extra networks and main mode for #10395

Browse Source
This commit is contained in:
AUTOMATIC 2023-05-17 21:03:41 +03:00
parent 9c91a86720
commit a6b618d072
2 changed files with 43 additions and 46 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

70
modules/images.py
View File

@ -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)

19
modules/ui_extra_networks.py
View File

@ -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]