From 5fe2c75da0de5ebc21c00f641aafab5bbdae037e Mon Sep 17 00:00:00 2001
From: olivier <olivier@pcedev.com>
Date: Tue, 20 Sep 2022 00:13:12 +0200
Subject: [PATCH] Support reusing filename stems during batch processing in
 Extras tab

---
 modules/extras.py | 20 ++++++++++++++------
 modules/images.py | 24 ++++++++++++++----------
 modules/shared.py |  1 +
 3 files changed, 29 insertions(+), 16 deletions(-)

diff --git a/modules/extras.py b/modules/extras.py
index a38310f0..96bd2118 100644
--- a/modules/extras.py
+++ b/modules/extras.py
@@ -1,3 +1,5 @@
+import os
+
 import numpy as np
 from PIL import Image
 
@@ -17,27 +19,31 @@ def run_extras(image, image_folder, gfpgan_visibility, codeformer_visibility, co
     devices.torch_gc()
 
     imageArr = []
+    # Also keep track of original file names
+    imageNameArr = []
 
-    if image_folder != None:
-        if image != None:
+    if image_folder is not None:
+        if image is not None:
             print("Batch detected and single image detected, please only use one of the two. Aborting.")
             return None
         #convert file to pillow image
         for img in image_folder:
             image = Image.fromarray(np.array(Image.open(img)))
             imageArr.append(image)
+            imageNameArr.append(os.path.splitext(img.orig_name)[0])
 
-    elif image != None:
-        if image_folder != None:
+    elif image is not None:
+        if image_folder is not None:
             print("Batch detected and single image detected, please only use one of the two. Aborting.")
             return None
         else:
             imageArr.append(image)
+            imageNameArr.append(None)
 
     outpath = opts.outdir_samples or opts.outdir_extras_samples
 
     outputs = []
-    for image in imageArr:
+    for image, image_name in zip(imageArr, imageNameArr):
         existing_pnginfo = image.info or {}
 
         image = image.convert("RGB")
@@ -90,7 +96,9 @@ def run_extras(image, image_folder, gfpgan_visibility, codeformer_visibility, co
         while len(cached_images) > 2:
             del cached_images[next(iter(cached_images.keys()))]
 
-        images.save_image(image, path=outpath, basename="", seed=None, prompt=None, extension=opts.samples_format, info=info, short_filename=True, no_prompt=True, grid=False, pnginfo_section_name="extras", existing_info=existing_pnginfo)
+        images.save_image(image, path=outpath, basename="", seed=None, prompt=None, extension=opts.samples_format, info=info, short_filename=True,
+                          no_prompt=True, grid=False, pnginfo_section_name="extras", existing_info=existing_pnginfo,
+                          forced_filename=image_name if opts.use_original_name_batch else None)
 
         outputs.append(image)
 
diff --git a/modules/images.py b/modules/images.py
index e287d0df..530a8440 100644
--- a/modules/images.py
+++ b/modules/images.py
@@ -303,7 +303,7 @@ def get_next_sequence_number(path, basename):
 
     return result + 1
 
-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):
+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):
     if short_filename or prompt is None or seed is None:
         file_decoration = ""
     elif opts.save_to_dirs:
@@ -335,15 +335,19 @@ def save_image(image, path, basename, seed=None, prompt=None, extension='png', i
 
     os.makedirs(path, exist_ok=True)
 
-    basecount = get_next_sequence_number(path, basename)
-    fullfn = "a.png"
-    fullfn_without_extension = "a"
-    for i in range(500):
-        fn = f"{basecount+i:05}" if basename == '' else f"{basename}-{basecount+i:04}"
-        fullfn = os.path.join(path, f"{fn}{file_decoration}.{extension}")
-        fullfn_without_extension = os.path.join(path, f"{fn}{file_decoration}")
-        if not os.path.exists(fullfn):
-            break
+    if forced_filename is None:
+        basecount = get_next_sequence_number(path, basename)
+        fullfn = "a.png"
+        fullfn_without_extension = "a"
+        for i in range(500):
+            fn = f"{basecount+i:05}" if basename == '' else f"{basename}-{basecount+i:04}"
+            fullfn = os.path.join(path, f"{fn}{file_decoration}.{extension}")
+            fullfn_without_extension = os.path.join(path, f"{fn}{file_decoration}")
+            if not os.path.exists(fullfn):
+                break
+    else:
+        fullfn = os.path.join(path, f"{forced_filename}.{extension}")
+        fullfn_without_extension = os.path.join(path, forced_filename)
 
     def exif_bytes():
         return piexif.dump({
diff --git a/modules/shared.py b/modules/shared.py
index 38f54a80..6f513ab7 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -161,6 +161,7 @@ class Options:
         "sd_model_checkpoint": OptionInfo(None, "Stable Diffusion checkpoint", gr.Radio, lambda: {"choices": [x.title for x in modules.sd_models.checkpoints_list.values()]}),
         "js_modal_lightbox": OptionInfo(True, "Enable full page image viewer"),
         "js_modal_lightbox_initialy_zoomed": OptionInfo(True, "Show images zoomed in by default in full page image viewer"),
+        "use_original_name_batch": OptionInfo(True, "Use original name for output filename during batch process"),
     }
 
     def __init__(self):