Support reusing filename stems during batch processing in Extras tab

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