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

Merge pull request #5119 from 0xb8/master

Browse Source 
Atomically rename saved image to avoid race condition with other processes
This commit is contained in:
AUTOMATIC1111 2022-12-10 13:26:07 +03:00 committed by GitHub
commit 89237852f4
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 33 additions and 26 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

35
modules/images.py
View File

@ -501,14 +501,11 @@ def save_image(image, path, basename, seed=None, prompt=None, extension='png', i
image = params.image image = params.image
fullfn = params.filename fullfn = params.filename
info = params.pnginfo.get(pnginfo_section_name, None) info = params.pnginfo.get(pnginfo_section_name, None)
fullfn_without_extension, extension = os.path.splitext(params.filename)
def exif_bytes(): def _atomically_save_image(image_to_save, filename_without_extension, extension):
return piexif.dump({ # save image with .tmp extension to avoid race condition when another process detects new image in the directory
"Exif": { temp_file_path = filename_without_extension + ".tmp"
piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(info or "", encoding="unicode") image_format = Image.registered_extensions()[extension]
},
})
if extension.lower() == '.png': if extension.lower() == '.png':
pnginfo_data = PngImagePlugin.PngInfo() pnginfo_data = PngImagePlugin.PngInfo()
@ -516,15 +513,27 @@ def save_image(image, path, basename, seed=None, prompt=None, extension='png', i
for k, v in params.pnginfo.items(): for k, v in params.pnginfo.items():
pnginfo_data.add_text(k, str(v)) pnginfo_data.add_text(k, str(v))
image.save(fullfn, quality=opts.jpeg_quality, pnginfo=pnginfo_data) image_to_save.save(temp_file_path, format=image_format, quality=opts.jpeg_quality, pnginfo=pnginfo_data)
elif extension.lower() in (".jpg", ".jpeg", ".webp"): elif extension.lower() in (".jpg", ".jpeg", ".webp"):
image.save(fullfn, quality=opts.jpeg_quality) image_to_save.save(temp_file_path, format=image_format, quality=opts.jpeg_quality)
if opts.enable_pnginfo and info is not None: if opts.enable_pnginfo and info is not None:
piexif.insert(exif_bytes(), fullfn) exif_bytes = piexif.dump({
"Exif": {
piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(info or "", encoding="unicode")
},
})
piexif.insert(exif_bytes, temp_file_path)
else: else:
image.save(fullfn, quality=opts.jpeg_quality) 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)
_atomically_save_image(image, fullfn_without_extension, extension)
image.already_saved_as = fullfn image.already_saved_as = fullfn
@ -538,9 +547,7 @@ def save_image(image, path, basename, seed=None, prompt=None, extension='png', i
elif oversize: elif oversize:
image = image.resize((image.width * target_side_length // image.height, target_side_length), LANCZOS) image = image.resize((image.width * target_side_length // image.height, target_side_length), LANCZOS)
image.save(fullfn_without_extension + ".jpg", quality=opts.jpeg_quality) _atomically_save_image(image, fullfn_without_extension, ".jpg")
if opts.enable_pnginfo and info is not None:
piexif.insert(exif_bytes(), fullfn_without_extension + ".jpg")
if opts.save_txt and info is not None: if opts.save_txt and info is not None:
txt_fullfn = f"{fullfn_without_extension}.txt" txt_fullfn = f"{fullfn_without_extension}.txt"