Merge pull request #7295 from askaliuk/askaliuk-inpaint-batch-support

Basic inpainting batch support

2 files changed, 25 insertions, 4 deletions

diff --git a/modules/img2img.py b/modules/img2img.py
index 2168c8e2..fe9447c7 100644
--- a/modules/img2img.py
+++ b/modules/img2img.py
@@ -16,11 +16,16 @@ import modules.images as images
 import modules.scripts

 

 

-def process_batch(p, input_dir, output_dir, args):

+def process_batch(p, input_dir, output_dir, inpaint_mask_dir, args):

     processing.fix_seed(p)

 

     images = shared.listfiles(input_dir)

 

+    inpaint_masks = shared.listfiles(inpaint_mask_dir)

+    is_inpaint_batch = inpaint_mask_dir and len(inpaint_masks) > 0

+    if is_inpaint_batch:

+        print(f"\nInpaint batch is enabled. {len(inpaint_masks)} masks found.")

+

     print(f"Will process {len(images)} images, creating {p.n_iter * p.batch_size} new images for each.")

 

     save_normally = output_dir == ''

@@ -43,6 +48,15 @@ def process_batch(p, input_dir, output_dir, args):
         img = ImageOps.exif_transpose(img)

         p.init_images = [img] * p.batch_size

 

+        if is_inpaint_batch:

+            # try to find corresponding mask for an image using simple filename matching

+            mask_image_path = os.path.join(inpaint_mask_dir, os.path.basename(image))

+            # if not found use first one ("same mask for all images" use-case)

+            if not mask_image_path in inpaint_masks:

+                mask_image_path = inpaint_masks[0]

+            mask_image = Image.open(mask_image_path)

+            p.image_mask = mask_image

+

         proc = modules.scripts.scripts_img2img.run(p, *args)

         if proc is None:

             proc = process_images(p)

@@ -59,7 +73,7 @@ def process_batch(p, input_dir, output_dir, args):
                 processed_image.save(os.path.join(output_dir, filename))

 

 

-def img2img(id_task: str, mode: int, prompt: str, negative_prompt: str, prompt_styles, init_img, sketch, init_img_with_mask, inpaint_color_sketch, inpaint_color_sketch_orig, init_img_inpaint, init_mask_inpaint, steps: int, sampler_index: int, mask_blur: int, mask_alpha: float, inpainting_fill: int, restore_faces: bool, tiling: bool, n_iter: int, batch_size: int, cfg_scale: float, denoising_strength: float, seed: int, subseed: int, subseed_strength: float, seed_resize_from_h: int, seed_resize_from_w: int, seed_enable_extras: bool, height: int, width: int, resize_mode: int, inpaint_full_res: bool, inpaint_full_res_padding: int, inpainting_mask_invert: int, img2img_batch_input_dir: str, img2img_batch_output_dir: str, *args):

+def img2img(id_task: str, mode: int, prompt: str, negative_prompt: str, prompt_styles, init_img, sketch, init_img_with_mask, inpaint_color_sketch, inpaint_color_sketch_orig, init_img_inpaint, init_mask_inpaint, steps: int, sampler_index: int, mask_blur: int, mask_alpha: float, inpainting_fill: int, restore_faces: bool, tiling: bool, n_iter: int, batch_size: int, cfg_scale: float, denoising_strength: float, seed: int, subseed: int, subseed_strength: float, seed_resize_from_h: int, seed_resize_from_w: int, seed_enable_extras: bool, height: int, width: int, resize_mode: int, inpaint_full_res: bool, inpaint_full_res_padding: int, inpainting_mask_invert: int, img2img_batch_input_dir: str, img2img_batch_output_dir: str, img2img_batch_inpaint_mask_dir: str, *args):

     is_batch = mode == 5

 

     if mode == 0:  # img2img

@@ -139,7 +153,7 @@ def img2img(id_task: str, mode: int, prompt: str, negative_prompt: str, prompt_s
     if is_batch:

         assert not shared.cmd_opts.hide_ui_dir_config, "Launched with --hide-ui-dir-config, batch img2img disabled"

 

-        process_batch(p, img2img_batch_input_dir, img2img_batch_output_dir, args)

+        process_batch(p, img2img_batch_input_dir, img2img_batch_output_dir, img2img_batch_inpaint_mask_dir, args)

 

         processed = Processed(p, [], p.seed, "")

     else:

diff --git a/modules/ui.py b/modules/ui.py
index 85ae62c7..fddb9177 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -691,9 +691,15 @@ def create_ui():
 

                     with gr.TabItem('Batch', id='batch', elem_id="img2img_batch_tab") as tab_batch:

                         hidden = '<br>Disabled when launched with --hide-ui-dir-config.' if shared.cmd_opts.hide_ui_dir_config else ''

-                        gr.HTML(f"<p style='padding-bottom: 1em;' class=\"text-gray-500\">Process images in a directory on the same machine where the server is running.<br>Use an empty output directory to save pictures normally instead of writing to the output directory.{hidden}</p>")

+                        gr.HTML(

+                            f"<p style='padding-bottom: 1em;' class=\"text-gray-500\">Process images in a directory on the same machine where the server is running." +

+                            f"<br>Use an empty output directory to save pictures normally instead of writing to the output directory." +

+                            f"<br>Add inpaint batch mask directory to enable inpaint batch processing."

+                            f"{hidden}</p>"

+                        )

                         img2img_batch_input_dir = gr.Textbox(label="Input directory", **shared.hide_dirs, elem_id="img2img_batch_input_dir")

                         img2img_batch_output_dir = gr.Textbox(label="Output directory", **shared.hide_dirs, elem_id="img2img_batch_output_dir")

+                        img2img_batch_inpaint_mask_dir = gr.Textbox(label="Inpaint batch mask directory (required for inpaint batch processing only)", **shared.hide_dirs, elem_id="img2img_batch_inpaint_mask_dir")

 

                 def copy_image(img):

                     if isinstance(img, dict) and 'image' in img:

@@ -838,6 +844,7 @@ def create_ui():
                     inpainting_mask_invert,

                     img2img_batch_input_dir,

                     img2img_batch_output_dir,

+                    img2img_batch_inpaint_mask_dir

                 ] + custom_inputs,

                 outputs=[

                     img2img_gallery,