added extras batch work from directory

This commit is contained in:
winterspringsummer 2022-10-16 13:50:55 +09:00 committed by AUTOMATIC1111
parent 0c5fa9a681
commit 2ce27728f6
2 changed files with 30 additions and 5 deletions

View file

@ -20,12 +20,13 @@ import gradio as gr
cached_images = {}
def run_extras(extras_mode, resize_mode, image, image_folder, gfpgan_visibility, codeformer_visibility, codeformer_weight, upscaling_resize, upscaling_resize_w, upscaling_resize_h, upscaling_crop, extras_upscaler_1, extras_upscaler_2, extras_upscaler_2_visibility):
def run_extras(extras_mode, resize_mode, image, image_folder, input_dir, output_dir, show_extras_results, gfpgan_visibility, codeformer_visibility, codeformer_weight, upscaling_resize, upscaling_resize_w, upscaling_resize_h, upscaling_crop, extras_upscaler_1, extras_upscaler_2, extras_upscaler_2_visibility):
devices.torch_gc()
imageArr = []
# Also keep track of original file names
imageNameArr = []
outputs = []
if extras_mode == 1:
#convert file to pillow image
@ -33,13 +34,24 @@ def run_extras(extras_mode, resize_mode, image, image_folder, gfpgan_visibility,
image = Image.open(img)
imageArr.append(image)
imageNameArr.append(os.path.splitext(img.orig_name)[0])
elif extras_mode == 2:
if input_dir == '':
return outputs, "Please select an input directory.", ''
image_list = [file for file in [os.path.join(input_dir, x) for x in os.listdir(input_dir)] if os.path.isfile(file)]
for img in image_list:
image = Image.open(img)
imageArr.append(image)
imageNameArr.append(img)
else:
imageArr.append(image)
imageNameArr.append(None)
outpath = opts.outdir_samples or opts.outdir_extras_samples
if extras_mode == 2 and output_dir != '':
outpath = output_dir
else:
outpath = opts.outdir_samples or opts.outdir_extras_samples
outputs = []
for image, image_name in zip(imageArr, imageNameArr):
if image is None:
return outputs, "Please select an input image.", ''
@ -112,7 +124,8 @@ def run_extras(extras_mode, resize_mode, image, image_folder, gfpgan_visibility,
image.info = existing_pnginfo
image.info["extras"] = info
outputs.append(image)
if extras_mode != 2 or show_extras_results :
outputs.append(image)
devices.torch_gc()

View file

@ -1016,6 +1016,15 @@ def create_ui(wrap_gradio_gpu_call):
with gr.TabItem('Batch Process'):
image_batch = gr.File(label="Batch Process", file_count="multiple", interactive=True, type="file")
with gr.TabItem('Batch from Directory'):
extras_batch_input_dir = gr.Textbox(label="Input directory", **shared.hide_dirs,
placeholder="A directory on the same machine where the server is running."
)
extras_batch_output_dir = gr.Textbox(label="Output directory", **shared.hide_dirs,
placeholder="Leave blank to save images to the default path."
)
show_extras_results = gr.Checkbox(label='Show result images', value=True)
with gr.Tabs(elem_id="extras_resize_mode"):
with gr.TabItem('Scale by'):
upscaling_resize = gr.Slider(minimum=1.0, maximum=4.0, step=0.05, label="Resize", value=2)
@ -1060,6 +1069,9 @@ def create_ui(wrap_gradio_gpu_call):
dummy_component,
extras_image,
image_batch,
extras_batch_input_dir,
extras_batch_output_dir,
show_extras_results,
gfpgan_visibility,
codeformer_visibility,
codeformer_weight,