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