mrq/stable-diffusion-webui
1
1
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Support reusing filename stems during batch processing in Extras tab

Browse Source
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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
modules/extras.py
View File

@ -1,3 +1,5 @@
import os
import numpy as np import numpy as np
from PIL import Image from PIL import Image
@ -17,27 +19,31 @@ def run_extras(image, image_folder, gfpgan_visibility, codeformer_visibility, co
devices.torch_gc() devices.torch_gc()
imageArr = [] imageArr = []
# Also keep track of original file names
imageNameArr = []
if image_folder != None: if image_folder is not None:
if image != None: if image is not None:
print("Batch detected and single image detected, please only use one of the two. Aborting.") print("Batch detected and single image detected, please only use one of the two. Aborting.")
return None return None
#convert file to pillow image #convert file to pillow image
for img in image_folder: for img in image_folder:
image = Image.fromarray(np.array(Image.open(img))) image = Image.fromarray(np.array(Image.open(img)))
imageArr.append(image) imageArr.append(image)
imageNameArr.append(os.path.splitext(img.orig_name)[0])
elif image != None: elif image is not None:
if image_folder != None: if image_folder is not None:
print("Batch detected and single image detected, please only use one of the two. Aborting.") print("Batch detected and single image detected, please only use one of the two. Aborting.")
return None return None
else: else:
imageArr.append(image) imageArr.append(image)
imageNameArr.append(None)
outpath = opts.outdir_samples or opts.outdir_extras_samples outpath = opts.outdir_samples or opts.outdir_extras_samples
outputs = [] outputs = []
for image in imageArr: for image, image_name in zip(imageArr, imageNameArr):
existing_pnginfo = image.info or {} existing_pnginfo = image.info or {}
image = image.convert("RGB") image = image.convert("RGB")
@ -90,7 +96,9 @@ def run_extras(image, image_folder, gfpgan_visibility, codeformer_visibility, co
while len(cached_images) > 2: while len(cached_images) > 2:
del cached_images[next(iter(cached_images.keys()))] 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) outputs.append(image)

6
modules/images.py
View File

@ -303,7 +303,7 @@ def get_next_sequence_number(path, basename):
return result + 1 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: if short_filename or prompt is None or seed is None:
file_decoration = "" file_decoration = ""
elif opts.save_to_dirs: 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) os.makedirs(path, exist_ok=True)
if forced_filename is None:
basecount = get_next_sequence_number(path, basename) basecount = get_next_sequence_number(path, basename)
fullfn = "a.png" fullfn = "a.png"
fullfn_without_extension = "a" 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}") fullfn_without_extension = os.path.join(path, f"{fn}{file_decoration}")
if not os.path.exists(fullfn): if not os.path.exists(fullfn):
break break
else:
fullfn = os.path.join(path, f"{forced_filename}.{extension}")
fullfn_without_extension = os.path.join(path, forced_filename)
def exif_bytes(): def exif_bytes():
return piexif.dump({ return piexif.dump({

1
modules/shared.py
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()]}), "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": 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"), "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): def __init__(self):