Add batch processing to Extras tab

modules/extras.py

@@ -13,16 +13,35 @@ import piexif.helper
 cached_images = {}
 
 
-def run_extras(image, gfpgan_visibility, codeformer_visibility, codeformer_weight, upscaling_resize, extras_upscaler_1, extras_upscaler_2, extras_upscaler_2_visibility):
+def run_extras(image, image_folder, gfpgan_visibility, codeformer_visibility, codeformer_weight, upscaling_resize, extras_upscaler_1, extras_upscaler_2, extras_upscaler_2_visibility):
     devices.torch_gc()
 
+    imageArr = []
+
+    if image_folder != None:
+        if image != 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)
+
+    elif image != None:
+        if image_folder != None:
+            print("Batch detected and single image detected, please only use one of the two. Aborting.")
+            return None
+        else:
+            imageArr.append(image)
+
+    outpath = opts.outdir_samples or opts.outdir_extras_samples
+
+    for image in imageArr:
         existing_pnginfo = image.info or {}
 
         image = image.convert("RGB")
         info = ""
 
-    outpath = opts.outdir_samples or opts.outdir_extras_samples
-
         if gfpgan_visibility > 0:
             restored_img = modules.gfpgan_model.gfpgan_fix_faces(np.array(image, dtype=np.uint8))
             res = Image.fromarray(restored_img)
@@ -72,7 +91,7 @@ def run_extras(image, gfpgan_visibility, codeformer_visibility, codeformer_weigh
 
         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)
 
-    return image, plaintext_to_html(info), ''
+    return imageArr, plaintext_to_html(info), ''
 
 
 def run_pnginfo(image):

modules/ui.py

@@ -644,9 +644,13 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
     with gr.Blocks(analytics_enabled=False) as extras_interface:
         with gr.Row().style(equal_height=False):
             with gr.Column(variant='panel'):
-                with gr.Group():
+                with gr.Tabs():
+                    with gr.TabItem('Single Image'):
                         image = gr.Image(label="Source", source="upload", interactive=True, type="pil")
 
+                    with gr.TabItem('Batch Process'):
+                        image_batch = gr.File(label="Batch Process", file_count="multiple", source="upload", interactive=True, type="file")
+
                 upscaling_resize = gr.Slider(minimum=1.0, maximum=4.0, step=0.05, label="Resize", value=2)
 
                 with gr.Group():
@@ -666,7 +670,7 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
                 submit = gr.Button('Generate', elem_id="extras_generate", variant='primary')
 
             with gr.Column(variant='panel'):
-                result_image = gr.Image(label="Result")
+                result_images = gr.Gallery(label="Result")
                 html_info_x = gr.HTML()
                 html_info = gr.HTML()
 
@@ -674,6 +678,7 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
             fn=run_extras,
             inputs=[
                 image,
+                image_batch,
                 gfpgan_visibility,
                 codeformer_visibility,
                 codeformer_weight,
@@ -683,7 +688,7 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo):
                 extras_upscaler_2_visibility,
             ],
             outputs=[
-                result_image,
+                result_images,
                 html_info_x,
                 html_info,
             ]