From 2ce27728f6433911274efa67856315d22df56629 Mon Sep 17 00:00:00 2001 From: winterspringsummer Date: Sun, 16 Oct 2022 13:50:55 +0900 Subject: added extras batch work from directory --- modules/extras.py | 23 ++++++++++++++++++----- 1 file changed, 18 insertions(+), 5 deletions(-) (limited to 'modules/extras.py') diff --git a/modules/extras.py b/modules/extras.py index f2f5a7b0..5b52b27d 100644 --- a/modules/extras.py +++ b/modules/extras.py @@ -20,26 +20,38 @@ 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 for img in image_folder: 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() -- cgit v1.2.1