aboutsummaryrefslogtreecommitdiff
path: root/modules
diff options
context:
space:
mode:
authorcatboxanon <122327233+catboxanon@users.noreply.github.com>2023-03-22 17:51:40 +0000
committercatboxanon <122327233+catboxanon@users.noreply.github.com>2023-03-22 17:51:40 +0000
commitcaf84e8233236cd36fcd42e39e4199262dc3f4e7 (patch)
tree0f31fc3392baec8b76f092121e3a28ce63932c15 /modules
parenta9fed7c364061ae6efb37f797b6b522cb3cf7aa2 (diff)
Expose inpainting mask and composite
For inpainting, this exposes the mask and masked composite and gives the user the ability to display these in the web UI, save to disk, or both.
Diffstat (limited to 'modules')
-rw-r--r--modules/processing.py16
-rw-r--r--modules/shared.py4
2 files changed, 20 insertions, 0 deletions
diff --git a/modules/processing.py b/modules/processing.py
index 59717b4c..2e5a363f 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -689,6 +689,22 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
image.info["parameters"] = text
output_images.append(image)
+ if hasattr(p, 'mask_for_overlay') and p.mask_for_overlay:
+ image_mask = p.mask_for_overlay.convert('RGB')
+ image_mask_composite = Image.composite(image.convert('RGBA').convert('RGBa'), Image.new('RGBa', image.size), p.mask_for_overlay.convert('L')).convert('RGBA')
+
+ if opts.save_mask:
+ images.save_image(image_mask, p.outpath_samples, "", seeds[i], prompts[i], opts.samples_format, info=infotext(n, i), p=p, suffix="-mask")
+
+ if opts.save_mask_composite:
+ images.save_image(image_mask_composite, p.outpath_samples, "", seeds[i], prompts[i], opts.samples_format, info=infotext(n, i), p=p, suffix="-mask-composite")
+
+ if opts.return_mask:
+ output_images.append(image_mask)
+
+ if opts.return_mask_composite:
+ output_images.append(image_mask_composite)
+
del x_samples_ddim
devices.torch_gc()
diff --git a/modules/shared.py b/modules/shared.py
index f28a12cc..d26e9111 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -332,6 +332,8 @@ options_templates.update(options_section(('saving-images', "Saving images/grids"
"save_images_before_face_restoration": OptionInfo(False, "Save a copy of image before doing face restoration."),
"save_images_before_highres_fix": OptionInfo(False, "Save a copy of image before applying highres fix."),
"save_images_before_color_correction": OptionInfo(False, "Save a copy of image before applying color correction to img2img results"),
+ "save_mask": OptionInfo(False, "For inpainting, save a copy of the greyscale mask"),
+ "save_mask_composite": OptionInfo(False, "For inpainting, save a masked composite"),
"jpeg_quality": OptionInfo(80, "Quality for saved jpeg images", gr.Slider, {"minimum": 1, "maximum": 100, "step": 1}),
"webp_lossless": OptionInfo(False, "Use lossless compression for webp images"),
"export_for_4chan": OptionInfo(True, "If the saved image file size is above the limit, or its either width or height are above the limit, save a downscaled copy as JPG"),
@@ -454,6 +456,8 @@ options_templates.update(options_section(('extra_networks', "Extra Networks"), {
options_templates.update(options_section(('ui', "User interface"), {
"return_grid": OptionInfo(True, "Show grid in results for web"),
+ "return_mask": OptionInfo(False, "For inpainting, include the greyscale mask in results for web"),
+ "return_mask_composite": OptionInfo(False, "For inpainting, include masked composite in results for web"),
"do_not_show_images": OptionInfo(False, "Do not show any images in results for web"),
"add_model_hash_to_info": OptionInfo(True, "Add model hash to generation information"),
"add_model_name_to_info": OptionInfo(True, "Add model name to generation information"),