use a single function for saving images with metadata both in extra networks...

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

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))
-
-            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)
+        save_image_with_geninfo(image_to_save, info, temp_file_path, extension, params.pnginfo)
 
-        # 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)

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]