fallback to original file retrieving; skip img if mask not found

usage of `shared.walk_files` breaks controlnet extension
images are processed in different order 
which leads to unmatched img file used for img2img and img file used for controlnet 
(if no folder is specified for control net
or the same as img2img input dir used for it)
This commit is contained in:
Artem Kotov 2023-06-01 15:44:55 +04:00 committed by GitHub
parent 49f4b4be67
commit ba110bf093
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 7 additions and 12 deletions

View File

@ -17,12 +17,11 @@ import modules.scripts
def process_batch(p, input_dir, output_dir, inpaint_mask_dir, args, to_scale=False, scale_by=1.0): def process_batch(p, input_dir, output_dir, inpaint_mask_dir, args, to_scale=False, scale_by=1.0):
processing.fix_seed(p) processing.fix_seed(p)
allowed_extensions = [ext for ext, f in Image.registered_extensions().items() if f in Image.OPEN] images = shared.listfiles(input_dir)
images = list(shared.walk_files(input_dir, allowed_extensions=allowed_extensions))
is_inpaint_batch = False is_inpaint_batch = False
if inpaint_mask_dir: if inpaint_mask_dir:
inpaint_masks = list(shared.walk_files(inpaint_mask_dir, allowed_extensions=allowed_extensions)) inpaint_masks = shared.listfiles(inpaint_mask_dir)
is_inpaint_batch = len(inpaint_masks) > 0 is_inpaint_batch = len(inpaint_masks) > 0
if is_inpaint_batch: if is_inpaint_batch:
print(f"\nInpaint batch is enabled. {len(inpaint_masks)} masks found.") print(f"\nInpaint batch is enabled. {len(inpaint_masks)} masks found.")
@ -59,22 +58,18 @@ def process_batch(p, input_dir, output_dir, inpaint_mask_dir, args, to_scale=Fal
p.init_images = [img] * p.batch_size p.init_images = [img] * p.batch_size
image_path = Path(image) image_path = Path(image)
if image_path.parent == Path(input_dir):
image_subpath = ""
else:
image_subpath = str(image_path.parent).replace(f"{input_dir}/", "")
if is_inpaint_batch: if is_inpaint_batch:
# try to find corresponding mask for an image using simple filename matching # try to find corresponding mask for an image using simple filename matching
if len(inpaint_masks) == 1: if len(inpaint_masks) == 1:
mask_image_path = inpaint_masks[0] mask_image_path = inpaint_masks[0]
else: else:
# try to find corresponding mask for an image using simple filename matching # try to find corresponding mask for an image using simple filename matching
mask_image_dir = Path(inpaint_mask_dir).joinpath(image_subpath) mask_image_dir = Path(inpaint_mask_dir)
masks_found = list(mask_image_dir.glob(f"{image_path.stem}.*")) masks_found = list(mask_image_dir.glob(f"{image_path.stem}.*"))
if len(masks_found) == 0: if len(masks_found) == 0:
raise ValueError(f"Mask is not found for {image_path} in {mask_image_dir}") print(f"Warning: mask is not found for {image_path} in {mask_image_dir}. Skipping it.")
continue
# it should contain only 1 matching mask # it should contain only 1 matching mask
# otherwise user has many masks with the same name but different extensions # otherwise user has many masks with the same name but different extensions
@ -95,10 +90,10 @@ def process_batch(p, input_dir, output_dir, inpaint_mask_dir, args, to_scale=Fal
filename = f"{left}-{n}{right}" filename = f"{left}-{n}{right}"
if not save_normally: if not save_normally:
os.makedirs(os.path.join(output_dir, image_subpath), exist_ok=True) os.makedirs(output_dir, exist_ok=True)
if processed_image.mode == 'RGBA': if processed_image.mode == 'RGBA':
processed_image = processed_image.convert("RGB") processed_image = processed_image.convert("RGB")
processed_image.save(os.path.join(output_dir, image_subpath, filename)) 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, image_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, selected_scale_tab: int, height: int, width: int, scale_by: float, 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, override_settings_texts, *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, image_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, selected_scale_tab: int, height: int, width: int, scale_by: float, 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, override_settings_texts, *args):