From 2ce27728f6433911274efa67856315d22df56629 Mon Sep 17 00:00:00 2001
From: winterspringsummer <linoo@naver.com>
Date: Sun, 16 Oct 2022 13:50:55 +0900
Subject: [PATCH] added extras batch work from directory

---
 modules/extras.py | 23 ++++++++++++++++++-----
 modules/ui.py     | 12 ++++++++++++
 2 files changed, 30 insertions(+), 5 deletions(-)

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()
 
diff --git a/modules/ui.py b/modules/ui.py
index b867d40f..08fa72c6 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -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,