Support reusing filename stems during batch processing in Extras tab

This commit is contained in:
olivier 2022-09-20 00:13:12 +02:00 committed by AUTOMATIC1111
parent 53651696dc
commit 5fe2c75da0
3 changed files with 29 additions and 16 deletions

View File

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

View File

@ -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,6 +335,7 @@ def save_image(image, path, basename, seed=None, prompt=None, extension='png', i
os.makedirs(path, exist_ok=True)
if forced_filename is None:
basecount = get_next_sequence_number(path, basename)
fullfn = "a.png"
fullfn_without_extension = "a"
@ -344,6 +345,9 @@ def save_image(image, path, basename, seed=None, prompt=None, extension='png', i
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({

View File

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