From 482a6ce8cbdde82fecfabd4a47a7720b54676a54 Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Mon, 12 Sep 2022 18:59:53 +0300
Subject: [PATCH]  [Feature Request] Save defaults for extras & keep image
 parameters after using extras #251

---
 modules/extras.py |  4 +++-
 modules/images.py | 11 +++++++++--
 2 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/modules/extras.py b/modules/extras.py
index 40935f988..596cd1728 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -13,6 +13,8 @@ cached_images = {}
 def run_extras(image, gfpgan_visibility, codeformer_visibility, codeformer_weight, upscaling_resize, extras_upscaler_1, extras_upscaler_2, extras_upscaler_2_visibility):
     devices.torch_gc()
 
+    existing_pnginfo = image.info or {}
+
     image = image.convert("RGB")
     info = ""
 
@@ -65,7 +67,7 @@ def run_extras(image, gfpgan_visibility, codeformer_visibility, codeformer_weigh
     while len(cached_images) > 2:
         del cached_images[next(iter(cached_images.keys()))]
 
-    images.save_image(image, outpath, "", None, info=info, extension=opts.samples_format, short_filename=True, no_prompt=True, pnginfo_section_name="extras")
+    images.save_image(image, outpath, "", None, info=info, extension=opts.samples_format, short_filename=True, no_prompt=True, pnginfo_section_name="extras", existing_info=existing_pnginfo)
 
     return image, plaintext_to_html(info), ''
 
diff --git a/modules/images.py b/modules/images.py
index 334f8fec4..05fc42060 100644
--- a/modules/images.py
+++ b/modules/images.py
@@ -247,7 +247,7 @@ def sanitize_filename_part(text, replace_spaces=True):
     return text.translate({ord(x): '' for x in invalid_filename_chars})[:128]
 
 
-def save_image(image, path, basename, seed=None, prompt=None, extension='png', info=None, short_filename=False, no_prompt=False, pnginfo_section_name='parameters', p=None):
+def save_image(image, path, basename, seed=None, prompt=None, extension='png', info=None, short_filename=False, no_prompt=False, pnginfo_section_name='parameters', p=None, existing_info=None):
     # would be better to add this as an argument in future, but will do for now
     is_a_grid = basename != ""
 
@@ -258,7 +258,9 @@ def save_image(image, path, basename, seed=None, prompt=None, extension='png', i
     else:
         file_decoration = opts.samples_filename_format or "[seed]-[prompt_spaces]"
 
-    file_decoration = "-" + file_decoration.lower()
+    if file_decoration != "":
+        file_decoration = "-" + file_decoration.lower()
+
     if seed is not None:
         file_decoration = file_decoration.replace("[seed]", str(seed))
     if prompt is not None:
@@ -273,6 +275,11 @@ def save_image(image, path, basename, seed=None, prompt=None, extension='png', i
 
     if extension == 'png' and opts.enable_pnginfo and info is not None:
         pnginfo = PngImagePlugin.PngInfo()
+
+        if existing_info is not None:
+            for k, v in existing_info.items():
+                pnginfo.add_text(k, v)
+
         pnginfo.add_text(pnginfo_section_name, info)
     else:
         pnginfo = None