From 0c0e80d90935fbec1796af22016bbb4873208274 Mon Sep 17 00:00:00 2001
From: unnamedplugins <79282950+unnamedplugins@users.noreply.github.com>
Date: Tue, 13 Sep 2022 10:57:07 -0400
Subject: [PATCH] Add mask restore feature to streamlit, prevent color
 correction from modifying initial image when mask_restore is turned on

---
 scripts/webui_streamlit.py | 48 ++++++++++++++++++++++++--------------
 1 file changed, 31 insertions(+), 17 deletions(-)

diff --git a/scripts/webui_streamlit.py b/scripts/webui_streamlit.py
index 341fa2a..f3ddd30 100644
--- a/scripts/webui_streamlit.py
+++ b/scripts/webui_streamlit.py
@@ -1644,17 +1644,28 @@ def img2img(prompt: str = '', init_info: any = None, init_info_mask: any = None,
 			if initial_seed is None:
 				initial_seed = seed
 
+			input_image = init_img
 			init_img = output_images[0]
 
 			if do_color_correction and correction_target is not None:
-				init_img = Image.fromarray(cv2.cvtColor(exposure.match_histograms(
-                                        cv2.cvtColor(
-                                                np.asarray(init_img),
-                                            cv2.COLOR_RGB2LAB
-                                        ),
-                                        correction_target,
-                                    channel_axis=2
-                                    ), cv2.COLOR_LAB2RGB).astype("uint8"))
+				init_img = Image.fromarray(
+					cv2.cvtColor(
+						exposure.match_histograms(
+							cv2.cvtColor(
+								np.asarray(init_img),
+								cv2.COLOR_RGB2LAB
+							),
+							correction_target,
+							channel_axis=2
+						),
+						cv2.COLOR_LAB2RGB).astype("uint8"))
+				if mask_restore is True:
+					color_mask = init_mask.filter(ImageFilter.GaussianBlur(mask_blur_strength))
+					color_mask = color_mask.convert('L')
+					source_image = input_image.convert('RGB')
+					target_image = init_img.convert('RGB')
+
+					init_img = Image.composite(source_image, target_image, color_mask)
 
 			if not random_seed_loopback:
 				seed = seed + 1
@@ -2377,6 +2388,9 @@ def layout():
 						                                       help="Ensure the sum of all weights add up to 1.0")
 						loopback = st.checkbox("Loopback.", value=defaults.img2img.loopback, help="Use images from previous batch when creating next batch.")
 						random_seed_loopback = st.checkbox("Random loopback seed.", value=defaults.img2img.random_seed_loopback, help="Random loopback seed")
+						img2img_mask_restore = st.checkbox("Only modify regenerated parts of image",
+														   value=defaults.img2img.mask_restore,
+														   help="Enable to restore the unmasked parts of the image with the input, may not blend as well but preserves detail")
 						save_individual_images = st.checkbox("Save individual images.", value=defaults.img2img.save_individual_images,
 						                                     help="Save each image generated before any filter or enhancement is applied.")
 						save_grid = st.checkbox("Save grid",value=defaults.img2img.save_grid, help="Save a grid with all the images generated into a single image.")
@@ -2517,15 +2531,15 @@ def layout():
 							new_mask = mask.resize((width, height))
 	
 						try:
-							output_images, seed, info, stats = img2img(prompt=prompt, init_info=new_img, init_info_mask=new_mask, mask_mode=mask_mode, ddim_steps=st.session_state["sampling_steps"],
-								                                   sampler_name=st.session_state["sampler_name"], n_iter=batch_count,
-								                                   cfg_scale=cfg_scale, denoising_strength=st.session_state["denoising_strength"], variant_seed=variant_seed,
-								                                   seed=seed, noise_mode=noise_mode, find_noise_steps=find_noise_steps, width=width, height=height, fp=defaults.general.fp, variant_amount=variant_amount, 
-								                                   ddim_eta=0.0, write_info_files=write_info_files, RealESRGAN_model=RealESRGAN_model,
-								                                   separate_prompts=separate_prompts, normalize_prompt_weights=normalize_prompt_weights,
-								                                   save_individual_images=save_individual_images, save_grid=save_grid, 
-								                                   group_by_prompt=group_by_prompt, save_as_jpg=save_as_jpg, use_GFPGAN=use_GFPGAN,
-								                                   use_RealESRGAN=use_RealESRGAN if not loopback else False, loopback=loopback
+							output_images, seed, info, stats = img2img(prompt=prompt, init_info=new_img, init_info_mask=new_mask, mask_mode=mask_mode, mask_restore=img2img_mask_restore,
+																	   ddim_steps=st.session_state["sampling_steps"], sampler_name=st.session_state["sampler_name"], n_iter=batch_count,
+																	   cfg_scale=cfg_scale, denoising_strength=st.session_state["denoising_strength"], variant_seed=variant_seed,
+																	   seed=seed, noise_mode=noise_mode, find_noise_steps=find_noise_steps, width=width, height=height, fp=defaults.general.fp,
+																	   variant_amount=variant_amount, ddim_eta=0.0, write_info_files=write_info_files, RealESRGAN_model=RealESRGAN_model,
+																	   separate_prompts=separate_prompts, normalize_prompt_weights=normalize_prompt_weights,
+																	   save_individual_images=save_individual_images, save_grid=save_grid,
+																	   group_by_prompt=group_by_prompt, save_as_jpg=save_as_jpg, use_GFPGAN=use_GFPGAN,
+																	   use_RealESRGAN=use_RealESRGAN if not loopback else False, loopback=loopback
 								                                   )
 		
 							#show a message when the generation is complete.