From caf84e8233236cd36fcd42e39e4199262dc3f4e7 Mon Sep 17 00:00:00 2001
From: catboxanon <122327233+catboxanon@users.noreply.github.com>
Date: Wed, 22 Mar 2023 17:51:40 +0000
Subject: 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.
---
 modules/processing.py | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

(limited to 'modules/processing.py')

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()
-- 
cgit v1.2.1