diff --git a/frontend/frontend.py b/frontend/frontend.py index 020901be1..de34742b5 100644 --- a/frontend/frontend.py +++ b/frontend/frontend.py @@ -200,9 +200,13 @@ def draw_gradio_ui(opt, img2img=lambda x: x, txt2img=lambda x: x, imgproc=lambda value=img2img_mask_modes[img2img_defaults['mask_mode']], visible=True) - img2img_mask_blur_strength = gr.Slider(minimum=1, maximum=10, step=1, + img2img_mask_restore = gr.Checkbox(label="Only modify regenerated parts of image", + value=img2img_defaults['mask_restore'], + visible=True) + + img2img_mask_blur_strength = gr.Slider(minimum=1, maximum=100, step=1, label="How much blurry should the mask be? (to avoid hard edges)", - value=3, visible=False) + value=3, visible=True) img2img_resize = gr.Radio(label="Resize mode", choices=["Just resize", "Crop and resize", @@ -294,7 +298,7 @@ def draw_gradio_ui(opt, img2img=lambda x: x, txt2img=lambda x: x, imgproc=lambda img2img_height ], [img2img_image_editor, img2img_image_mask, img2img_btn_editor, img2img_btn_mask, - img2img_painterro_btn, img2img_mask, img2img_mask_blur_strength] + img2img_painterro_btn, img2img_mask, img2img_mask_blur_strength, img2img_mask_restore] ) # img2img_image_editor_mode.change( @@ -335,8 +339,8 @@ def draw_gradio_ui(opt, img2img=lambda x: x, txt2img=lambda x: x, imgproc=lambda ) img2img_func = img2img - img2img_inputs = [img2img_prompt, img2img_image_editor_mode, img2img_mask, - img2img_mask_blur_strength, img2img_steps, img2img_sampling, img2img_toggles, + img2img_inputs = [img2img_prompt, img2img_image_editor_mode, img2img_mask, img2img_mask_blur_strength, + img2img_mask_restore, img2img_steps, img2img_sampling, img2img_toggles, img2img_realesrgan_model_name, img2img_batch_count, img2img_cfg, img2img_denoising, img2img_seed, img2img_height, img2img_width, img2img_resize, img2img_image_editor, img2img_image_mask, img2img_embeddings] diff --git a/frontend/ui_functions.py b/frontend/ui_functions.py index 6557841ee..ee6af8d07 100644 --- a/frontend/ui_functions.py +++ b/frontend/ui_functions.py @@ -9,10 +9,10 @@ def change_image_editor_mode(choice, cropped_image, masked_image, resize_mode, width, height): if choice == "Mask": update_image_result = update_image_mask(cropped_image, resize_mode, width, height) - return [gr.update(visible=False), update_image_result, gr.update(visible=False), gr.update(visible=True), gr.update(visible=False), gr.update(visible=True), gr.update(visible=True)] + return [gr.update(visible=False), update_image_result, gr.update(visible=False), gr.update(visible=True), gr.update(visible=False), gr.update(visible=True), gr.update(visible=True), gr.update(visible=True)] update_image_result = update_image_mask(masked_image["image"] if masked_image is not None else None, resize_mode, width, height) - return [update_image_result, gr.update(visible=False), gr.update(visible=True), gr.update(visible=False), gr.update(visible=True), gr.update(visible=False), gr.update(visible=False)] + return [update_image_result, gr.update(visible=False), gr.update(visible=True), gr.update(visible=False), gr.update(visible=True), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)] def update_image_mask(cropped_image, resize_mode, width, height): resized_cropped_image = resize_image(resize_mode, cropped_image, width, height) if cropped_image else None diff --git a/scripts/webui.py b/scripts/webui.py index 21b507abc..c0e5b7a5a 100644 --- a/scripts/webui.py +++ b/scripts/webui.py @@ -790,7 +790,7 @@ def process_images( outpath, func_init, func_sample, prompt, seed, sampler_name, skip_grid, skip_save, batch_size, n_iter, steps, cfg_scale, width, height, prompt_matrix, use_GFPGAN, use_RealESRGAN, realesrgan_model_name, fp, ddim_eta=0.0, do_not_save_grid=False, normalize_prompt_weights=True, init_img=None, init_mask=None, - keep_mask=False, mask_blur_strength=3, denoising_strength=0.75, resize_mode=None, uses_loopback=False, + keep_mask=False, mask_blur_strength=3, mask_restore=False, denoising_strength=0.75, resize_mode=None, uses_loopback=False, uses_random_seed_loopback=False, sort_samples=True, write_info_files=True, write_sample_info_to_log_file=False, jpg_sample=False, variant_amount=0.0, variant_seed=None,imgProcessorTask=False, job_info: JobInfo = None): """this is the main loop that both txt2img and img2img use; it calls func_init once inside all the scopes and func_sample once per batch""" @@ -1045,6 +1045,26 @@ def process_images( if imgProcessorTask == True: output_images.append(image) + if mask_restore and init_mask: + #init_mask = init_mask if keep_mask else ImageOps.invert(init_mask) + init_mask = init_mask.filter(ImageFilter.GaussianBlur(mask_blur_strength)) + init_mask = init_mask.convert('L') + init_img = init_img.convert('RGB') + image = image.convert('RGB') + + if use_RealESRGAN and RealESRGAN is not None: + if RealESRGAN.model.name != realesrgan_model_name: + try_loading_RealESRGAN(realesrgan_model_name) + output, img_mode = RealESRGAN.enhance(np.array(init_img, dtype=np.uint8)) + init_img = Image.fromarray(output) + init_img = init_img.convert('RGB') + + output, img_mode = RealESRGAN.enhance(np.array(init_mask, dtype=np.uint8)) + init_mask = Image.fromarray(output) + init_mask = init_mask.convert('L') + + image = Image.composite(init_img, image, init_mask) + if not skip_save: save_sample(image, sample_path_i, filename, jpg_sample, write_info_files, write_sample_info_to_log_file, prompt_matrix, init_img, uses_loopback, uses_random_seed_loopback, skip_save, skip_grid, sort_samples, sampler_name, ddim_eta, n_iter, batch_size, i, denoising_strength, resize_mode, False) @@ -1259,7 +1279,7 @@ def blurArr(a,r=8): -def img2img(prompt: str, image_editor_mode: str, mask_mode: str, mask_blur_strength: int, ddim_steps: int, sampler_name: str, +def img2img(prompt: str, image_editor_mode: str, mask_mode: str, mask_blur_strength: int, mask_restore: bool, ddim_steps: int, sampler_name: str, toggles: List[int], realesrgan_model_name: str, n_iter: int, cfg_scale: float, denoising_strength: float, seed: int, height: int, width: int, resize_mode: int, init_info: any = None, init_info_mask: any = None, fp = None, job_info: JobInfo = None): # print([prompt, image_editor_mode, init_info, init_info_mask, mask_mode, @@ -1504,6 +1524,7 @@ def sample(init_data, x, conditioning, unconditional_conditioning, sampler_name) init_mask=init_mask, keep_mask=keep_mask, mask_blur_strength=mask_blur_strength, + mask_restore=mask_restore, denoising_strength=denoising_strength, resize_mode=resize_mode, uses_loopback=loopback, @@ -1574,6 +1595,7 @@ def sample(init_data, x, conditioning, unconditional_conditioning, sampler_name) keep_mask=keep_mask, mask_blur_strength=mask_blur_strength, denoising_strength=denoising_strength, + mask_restore=mask_restore, resize_mode=resize_mode, uses_loopback=loopback, sort_samples=sort_samples, @@ -1722,6 +1744,7 @@ def processGoBig(image): init_img = result init_mask = None keep_mask = False + mask_restore = False assert 0. <= denoising_strength <= 1., 'can only work with strength in [0.0, 1.0]' def init(): @@ -1868,6 +1891,7 @@ def make_mask_image(r): keep_mask=False, mask_blur_strength=None, denoising_strength=denoising_strength, + mask_restore=mask_restore, resize_mode=resize_mode, uses_loopback=False, sort_samples=True, @@ -2186,6 +2210,7 @@ def run_RealESRGAN(image, model_name: str): 'cfg_scale': 5.0, 'denoising_strength': 0.75, 'mask_mode': 0, + 'mask_restore': False, 'resize_mode': 0, 'seed': '', 'height': 512, @@ -2199,24 +2224,6 @@ def run_RealESRGAN(image, model_name: str): img2img_toggle_defaults = [img2img_toggles[i] for i in img2img_defaults['toggles']] img2img_image_mode = 'sketch' -def change_image_editor_mode(choice, cropped_image, resize_mode, width, height): - if choice == "Mask": - return [gr.update(visible=False), gr.update(visible=True), gr.update(visible=False), gr.update(visible=True), gr.update(visible=False), gr.update(visible=False), gr.update(visible=True), gr.update(visible=True)] - return [gr.update(visible=True), gr.update(visible=False), gr.update(visible=True), gr.update(visible=False), gr.update(visible=True), gr.update(visible=True), gr.update(visible=False), gr.update(visible=False)] - -def update_image_mask(cropped_image, resize_mode, width, height): - resized_cropped_image = resize_image(resize_mode, cropped_image, width, height) if cropped_image else None - return gr.update(value=resized_cropped_image) - - - -def copy_img_to_upscale_esrgan(img): - update = gr.update(selected='realesrgan_tab') - image_data = re.sub('^data:image/.+;base64,', '', img) - processed_image = Image.open(BytesIO(base64.b64decode(image_data))) - return {'realesrgan_source': processed_image, 'tabs': update} - - help_text = """ ## Mask/Crop * The masking/cropping is very temperamental. diff --git a/scripts/webui_streamlit.py b/scripts/webui_streamlit.py index 26017149e..811536b85 100644 --- a/scripts/webui_streamlit.py +++ b/scripts/webui_streamlit.py @@ -913,7 +913,7 @@ def process_images( outpath, func_init, func_sample, prompt, seed, sampler_name, save_grid, batch_size, n_iter, steps, cfg_scale, width, height, prompt_matrix, use_GFPGAN, use_RealESRGAN, realesrgan_model_name, fp=None, ddim_eta=0.0, normalize_prompt_weights=True, init_img=None, init_mask=None, - keep_mask=False, mask_blur_strength=3, denoising_strength=0.75, resize_mode=None, uses_loopback=False, + keep_mask=False, mask_blur_strength=3, mask_restore=False, denoising_strength=0.75, resize_mode=None, uses_loopback=False, uses_random_seed_loopback=False, sort_samples=True, write_info_files=True, jpg_sample=False, variant_amount=0.0, variant_seed=None, save_individual_images: bool = True): """this is the main loop that both txt2img and img2img use; it calls func_init once inside all the scopes and func_sample once per batch""" @@ -1156,7 +1156,29 @@ def process_images( if simple_templating: grid_captions.append( captions[i] + "\ngfpgan_esrgan" ) - + + if mask_restore and init_mask: + #init_mask = init_mask if keep_mask else ImageOps.invert(init_mask) + init_mask = init_mask.filter(ImageFilter.GaussianBlur(mask_blur_strength)) + init_mask = init_mask.convert('L') + init_img = init_img.convert('RGB') + image = image.convert('RGB') + + if use_RealESRGAN and st.session_state["RealESRGAN"] is not None: + if st.session_state["RealESRGAN"].model.name != realesrgan_model_name: + #try_loading_RealESRGAN(realesrgan_model_name) + load_models(use_GFPGAN=use_GFPGAN, use_RealESRGAN=use_RealESRGAN, RealESRGAN_model=realesrgan_model_name) + + output, img_mode = st.session_state["RealESRGAN"].enhance(np.array(init_img, dtype=np.uint8)) + init_img = Image.fromarray(output) + init_img = init_img.convert('RGB') + + output, img_mode = st.session_state["RealESRGAN"].enhance(np.array(init_mask, dtype=np.uint8)) + init_mask = Image.fromarray(output) + init_mask = init_mask.convert('L') + + image = Image.composite(init_img, image, init_mask) + if save_individual_images: save_sample(image, sample_path_i, filename, jpg_sample, prompts, seeds, width, height, steps, cfg_scale, normalize_prompt_weights, use_GFPGAN, write_info_files, prompt_matrix, init_img, uses_loopback, uses_random_seed_loopback, @@ -1257,7 +1279,7 @@ def resize_image(resize_mode, im, width, height): return res def img2img(prompt: str = '', init_info: any = None, init_info_mask: any = None, mask_mode: int = 0, mask_blur_strength: int = 3, - ddim_steps: int = 50, sampler_name: str = 'DDIM', + mask_restore: bool = False, ddim_steps: int = 50, sampler_name: str = 'DDIM', n_iter: int = 1, cfg_scale: float = 7.5, denoising_strength: float = 0.8, seed: int = -1, height: int = 512, width: int = 512, resize_mode: int = 0, fp = None, variant_amount: float = None, variant_seed: int = None, ddim_eta:float = 0.0, @@ -1426,6 +1448,7 @@ def sample(init_data, x, conditioning, unconditional_conditioning, sampler_name) init_mask=init_mask, keep_mask=keep_mask, mask_blur_strength=mask_blur_strength, + mask_restore=mask_restore, denoising_strength=denoising_strength, resize_mode=resize_mode, uses_loopback=loopback, @@ -1486,8 +1509,9 @@ def sample(init_data, x, conditioning, unconditional_conditioning, sampler_name) init_img=init_img, init_mask=init_mask, keep_mask=keep_mask, - mask_blur_strength=2, + mask_blur_strength=mask_blur_strength, denoising_strength=denoising_strength, + mask_restore=mask_restore, resize_mode=resize_mode, uses_loopback=loopback, sort_samples=group_by_prompt,