diff --git a/CHANGELOG.md b/CHANGELOG.md index 4ec9a5279..d8bc7679a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,20 +1,21 @@ # Change Log for SD.Next -## Update for 2023-01-13 +## Update for 2023-01-14 Another release with a lot more functionality in new Control module and FaceID & IPAdapter modules Plus welcome additions to UI performance and accessibility and flexibility of deployment And it also includes fixes for all reported issues so far -However, - - **Control**: - add **inpaint** support applies to both *img2img* and *controlnet* workflows - *note*: set blur to level you desire - add **outpaint** support applies to both *img2img* and *controlnet* workflows *note*: increase denoising strength since outpainted area is blank by default + - new **mask** module + - granular blur (gaussian), errode (reduce or remove noise) and dilate (pad or expand) with **live preview** + - *optional* **auto-segmentation** (e.g. segment-anything) using ml models + auto segmentation will automatically expand masked area to segments that include current user mask - allow **resize** both *before* and *after* generate operation this allows for workflows such as: *image -> upscale or downscale -> generate -> upscale or downscale -> output* providing more flexibility and than standard hires workflow diff --git a/modules/control/run.py b/modules/control/run.py index 5c5d6da30..7c4ccc16f 100644 --- a/modules/control/run.py +++ b/modules/control/run.py @@ -80,7 +80,7 @@ def control_run(units: List[unit.Unit], inputs, inits, mask, unit_type: str, is_ hdr_clamp, hdr_boundary, hdr_threshold, hdr_center, hdr_channel_shift, hdr_full_shift, hdr_maximize, hdr_max_center, hdr_max_boundry, resize_mode_before, resize_name_before, width_before, height_before, scale_by_before, selected_scale_tab_before, resize_mode_after, resize_name_after, width_after, height_after, scale_by_after, selected_scale_tab_after, - denoising_strength, batch_count, batch_size, mask_blur, mask_overlap, + denoising_strength, batch_count, batch_size, video_skip_frames, video_type, video_duration, video_loop, video_pad, video_interpolate, ip_adapter, ip_scale, ip_image, *input_script_args @@ -135,7 +135,6 @@ def control_run(units: List[unit.Unit], inputs, inits, mask, unit_type: str, is_ denoising_strength = denoising_strength, n_iter = batch_count, batch_size = batch_size, - mask_blur=mask_blur, outpath_samples=shared.opts.outdir_samples or shared.opts.outdir_control_samples, outpath_grids=shared.opts.outdir_grids or shared.opts.outdir_control_grids, ) diff --git a/modules/masking.py b/modules/masking.py index 6c19276c8..1c33e9724 100644 --- a/modules/masking.py +++ b/modules/masking.py @@ -1,4 +1,12 @@ +from types import SimpleNamespace +import os +import time +import gradio as gr +import numpy as np +import cv2 from PIL import Image, ImageFilter, ImageOps +from transformers import SamModel, SamImageProcessor, MaskGenerationPipeline +from modules import shared, errors, devices, ui_components, ui_symbols def get_crop_region(mask, pad=0): @@ -83,3 +91,266 @@ def fill(image, mask): for _ in range(repeats): image_mod.alpha_composite(blurred) return image_mod.convert("RGB") + + +""" +[docs](https://huggingface.co/docs/transformers/v4.36.1/en/model_doc/sam#overview) +TODO: +- PerSAM +- https://huggingface.co/docs/transformers/tasks/semantic_segmentation +- transformers.pipeline.MaskGenerationPipeline: https://huggingface.co/models?pipeline_tag=mask-generation +- transformers.pipeline.ImageSegmentationPipeline: https://huggingface.co/models?pipeline_tag=image-segmentation +""" + +MODELS = { + 'None': None, + 'Facebook SAM ViT Base': 'facebook/sam-vit-base', + 'Facebook SAM ViT Large': 'facebook/sam-vit-large', + 'Facebook SAM ViT Huge': 'facebook/sam-vit-huge', + 'SlimSAM Uniform': 'Zigeng/SlimSAM-uniform-50', + 'SlimSAM Uniform Tiny': 'Zigeng/SlimSAM-uniform-77', + # 'Tiny Random': 'fxmarty/sam-vit-tiny-random', +} +COLORMAP = ['autumn', 'bone', 'jet', 'winter', 'rainbow', 'ocean', 'summer', 'spring', 'cool', 'hsv', 'pink', 'hot', 'parula', 'magma', 'inferno', 'plasma', 'viridis', 'cividis', 'twilight', 'shifted', 'turbo', 'deepgreen'] +cache_dir = 'models/control/segment' +loaded_model = None +model: SamModel = None +processor: SamImageProcessor = None +generator: MaskGenerationPipeline = None +debug = shared.log.trace if os.environ.get('SD_MASK_DEBUG', None) is not None else lambda *args, **kwargs: None +debug('Trace: MASK') +busy = False +btn_segment = None +controls = [] +opts = SimpleNamespace(**{ + 'mask_blur': 0.01, + 'mask_erode': 0.01, + 'mask_dilate': 0.01, + 'seg_iou_thresh': 0.5, + 'seg_score_thresh': 0.5, + 'seg_nms_thresh': 0.5, + 'seg_overlap_ratio': 0.3, + 'seg_points_per_batch': 64, + 'seg_topK': 50, + 'seg_colormap': 'pink', + 'preview_type': 'composite', + 'seg_live': True, + 'weight_original': 0.5, + 'weight_mask': 0.5, + 'kernel_iterations': 1, +}) + + +def init_model(selected_model: str): + global busy, loaded_model, model, processor, generator # pylint: disable=global-statement + if selected_model == "None": + if model is not None: + shared.log.debug('Segment unloading model') + model = None + loaded_model = None + processor = None + generator = None + devices.torch_gc() + return selected_model + model_path = MODELS[selected_model] + if model_path is not None and (loaded_model != selected_model or model is None or processor is None): + busy = True + t0 = time.time() + shared.log.debug(f'Segment loading: model={selected_model} path={model_path}') + model = SamModel.from_pretrained(model_path, cache_dir=cache_dir).to(device=devices.device) + processor = SamImageProcessor.from_pretrained(model_path, cache_dir=cache_dir) + generator = MaskGenerationPipeline( + model=model, + image_processor=processor, + device=devices.device, + # output_bboxes_mask=False, + # output_rle_masks=False, + ) + devices.torch_gc() + shared.log.debug(f'Segment loaded: model={selected_model} path={model_path} time={time.time()-t0:.2f}s') + busy = False + return selected_model + + +def run_segment(input_image: gr.Image, input_mask: np.ndarray): + outputs = None + with devices.inference_context(): + try: + outputs = generator( + input_image, + points_per_batch=opts.seg_points_per_batch, + pred_iou_thresh=opts.seg_iou_thresh, + stability_score_thresh=opts.seg_score_thresh, + crops_nms_thresh=opts.seg_nms_thresh, + crop_overlap_ratio=opts.seg_overlap_ratio, + crops_n_layers=0, + crop_n_points_downscale_factor=1, + ) + except Exception as e: + shared.log.error(f'Segment error: {e}') + errors.display(e, 'Segment') + return outputs + devices.torch_gc() + i = 1 + combined_mask = np.zeros(input_mask.shape, dtype='uint8') + input_mask_size = np.count_nonzero(input_mask) + debug(f'Segment: {vars(opts)}') + for mask in outputs['masks']: + mask = mask.astype('uint8') + mask_size = np.count_nonzero(mask) + if mask_size == 0: + continue + overlap = 0 + if input_mask_size > 0: + overlap = cv2.bitwise_and(mask, input_mask) + overlap = np.count_nonzero(overlap) + if overlap == 0: + continue + mask = (opts.seg_topK + 1 - i) * mask * (255 // opts.seg_topK) # set grayscale intensity so we can recolor + combined_mask = combined_mask + mask + debug(f'Segment mask: i={i} size={input_image.width}x{input_image.height} masked={mask_size}px overlap={overlap} score={outputs["scores"][i-1]:.2f}') + i += 1 + if i > opts.seg_topK: + break + return combined_mask + + +def run_mask(input_image: gr.Image, input_mask: gr.Image = None, return_type: str = None, mask_blur: int = None, mask_padding: int = None, segment_enable=True): + if input_image is None: + return input_mask + if isinstance(input_image, list): + input_image = input_image[0] + if isinstance(input_image, dict): + input_mask = input_image.get('mask', None) + input_image = input_image.get('image', None) + if input_mask is None: + input_mask = input_image.convert('L') + input_mask = input_mask.point(lambda x: 255 if x > 127 else 0) + else: + input_mask = input_mask.convert('L') + shared.log.debug(f'Segment mask: input={input_image} mask={input_mask} type={return_type}') + + input_mask = np.array(input_mask) // 255 + t0 = time.time() + + if mask_blur is not None: + opts.mask_blur = mask_blur / min(input_image.width, input_image.height) + if mask_padding is not None: + opts.mask_dilate = mask_padding / min(input_image.width, input_image.height) + + if generator is None or not segment_enable: + mask = input_mask * 255 + else: + mask = run_segment(input_image, input_mask) + + if mask is None: + shared.log.error('Segment error: no mask') + return input_mask + + debug(f'Segment mask: mask={mask.shape}') + if opts.mask_erode > 0: + try: + kernel = np.ones((int(opts.mask_erode * input_image.height / 4) + 1, int(opts.mask_erode * input_image.width / 4) + 1), np.uint8) + cv2_mask = cv2.erode(mask, kernel, iterations=opts.kernel_iterations) # remove noise + mask = cv2_mask + debug(f'Segment erode={opts.mask_erode} kernel={kernel} mask={mask.shape}') + except Exception as e: + shared.log.error(f'Segment erode: {e}') + if opts.mask_dilate > 0: + try: + kernel = np.ones((int(opts.mask_dilate * input_image.height / 4) + 1, int(opts.mask_dilate * input_image.width / 4) + 1), np.uint8) + cv2_mask = cv2.dilate(mask, kernel, iterations=opts.kernel_iterations) # expand area + mask = cv2_mask + debug(f'Segment dilate={opts.mask_dilate} kernel={kernel} mask={mask.shape}') + except Exception as e: + shared.log.error(f'Segment dilate: {e}') + if opts.mask_blur > 0: + try: + sigmax, sigmay = 1 + int(opts.mask_blur * input_image.width / 4), 1 + int(opts.mask_blur * input_image.height / 4) + cv2_mask = cv2.GaussianBlur(mask, (0, 0), sigmaX=sigmax, sigmaY=sigmay) # blur mask + mask = cv2_mask + debug(f'Segment blur={opts.mask_blur} x={sigmax} y={sigmay} mask={mask.shape}') + except Exception as e: + shared.log.error(f'Segment blur: {e}') + + mask_size = np.count_nonzero(mask) + total_size = np.prod(mask.shape) + area_size = np.count_nonzero(mask) + colored_mask = cv2.applyColorMap(mask, COLORMAP.index(opts.seg_colormap)) # recolor mask + combined_image = cv2.addWeighted(np.array(input_image), opts.weight_original, colored_mask, opts.weight_mask, 0) + binary_mask = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)[1] # otsu uses mean instead of threshold + t1 = time.time() + + return_type = return_type or opts.preview_type + shared.log.debug(f'Segment mask opts: size={input_image.width}x{input_image.height} masked={mask_size}px area={area_size/total_size:.2f} time={t1-t0:.2f}') + if return_type == 'none': + return input_mask + elif return_type == 'binary': + return Image.fromarray(binary_mask) + elif return_type == 'grayscale': + return Image.fromarray(mask) + elif return_type == 'color': + return Image.fromarray(colored_mask) + elif return_type == 'composite': + return Image.fromarray(combined_image) + return input_mask + + +def run_mask_live(input_image: gr.Image): + global busy # pylint: disable=global-statement + if opts.seg_live: + if not busy: + busy = True + res = run_mask(input_image) + busy = False + return res + else: + return None + + +def create_segment_ui(): + def update_opts(*args): + opts.seg_live = args[0] + opts.mask_blur = args[1] + opts.mask_erode = args[2] + opts.mask_dilate = args[3] + opts.seg_score_thresh = args[4] + opts.seg_iou_thresh = args[5] + opts.seg_nms_thresh = args[6] + opts.preview_type = args[7] + opts.seg_colormap = args[8] + + def display_controls(selected_model): + return 4 * [gr.update(visible=True)] + (len(controls) - 4) * [gr.update(visible=selected_model != 'None')] + + global btn_segment # pylint: disable=global-statement + with gr.Accordion(open=False, label="Mask", elem_id="control_mask", elem_classes=["small-accordion"]): + controls.clear() + with gr.Row(): + controls.append(gr.Checkbox(label="Live update", value=True)) + with gr.Row(): + controls.append(gr.Slider(minimum=0.0, maximum=1.0, step=0.01, label='Blur', value=0.01, elem_id="control_mask_blur")) + controls.append(gr.Slider(minimum=0.0, maximum=1.0, step=0.01, label='Erode', value=0.01, elem_id="control_mask_erode")) + controls.append(gr.Slider(minimum=0.0, maximum=1.0, step=0.01, label='Dilate', value=0.01, elem_id="control_mask_dilate")) + with gr.Row(): + selected_model = gr.Dropdown(label="Auto-segment", choices=MODELS.keys(), value='None') + btn_segment = ui_components.ToolButton(value=ui_symbols.refresh, visible=False) + with gr.Row(): + controls.append(gr.Slider(minimum=0.0, maximum=1.0, step=0.01, label='Score', value=0.5, visible=False)) + controls.append(gr.Slider(minimum=0.0, maximum=1.0, step=0.01, label='IOU', value=0.5, visible=False)) + controls.append(gr.Slider(minimum=0.0, maximum=1.0, step=0.01, label='NMS', value=0.5, visible=False)) + with gr.Row(): + controls.append(gr.Dropdown(label="Preview", choices=['none', 'binary', 'grayscale', 'color', 'composite'], value='composite')) + controls.append(gr.Dropdown(label="Colormap", choices=COLORMAP, value='pink')) + + selected_model.change(fn=init_model, inputs=[selected_model], outputs=[selected_model]) + selected_model.change(fn=display_controls, inputs=[selected_model], outputs=controls) + for control in controls: + control.change(fn=update_opts, inputs=controls, outputs=[]) + + +def bind_controls(input_image: gr.Image, preview_image: gr.Image): + btn_segment.click(run_mask, inputs=[input_image], outputs=[preview_image]) + input_image.edit(fn=run_mask_live, inputs=[input_image], outputs=[preview_image]) + for control in controls: + control.change(fn=run_mask_live, inputs=[input_image], outputs=[preview_image]) diff --git a/modules/processing.py b/modules/processing.py index ca0147ce3..a567a81d7 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -1290,15 +1290,18 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing): if self.image_mask is not None: if type(self.image_mask) == list: self.image_mask = self.image_mask[0] - self.image_mask = create_binary_mask(self.image_mask) - if self.inpainting_mask_invert: - self.image_mask = ImageOps.invert(self.image_mask) - if self.mask_blur > 0: - np_mask = np.array(self.image_mask) - kernel_size = 2 * int(2.5 * self.mask_blur + 0.5) + 1 - np_mask = cv2.GaussianBlur(np_mask, (kernel_size, 1), self.mask_blur) - np_mask = cv2.GaussianBlur(np_mask, (1, kernel_size), self.mask_blur) - self.image_mask = Image.fromarray(np_mask) + if shared.backend == shared.Backend.ORIGINAL: + self.image_mask = create_binary_mask(self.image_mask) + if self.inpainting_mask_invert: + self.image_mask = ImageOps.invert(self.image_mask) + if self.mask_blur > 0: + np_mask = np.array(self.image_mask) + kernel_size = 2 * int(2.5 * self.mask_blur + 0.5) + 1 + np_mask = cv2.GaussianBlur(np_mask, (kernel_size, 1), self.mask_blur) + np_mask = cv2.GaussianBlur(np_mask, (1, kernel_size), self.mask_blur) + self.image_mask = Image.fromarray(np_mask) + else: + self.image_mask = modules.masking.run_mask(input_image=self.init_images, input_mask=self.image_mask, return_type='grayscale', mask_blur=self.mask_blur, mask_padding=self.inpaint_full_res_padding, segment_enable=False) if self.inpaint_full_res: self.mask_for_overlay = self.image_mask mask = self.image_mask.convert('L') diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py index fc98ed76f..0fe7e3a87 100644 --- a/modules/processing_diffusers.py +++ b/modules/processing_diffusers.py @@ -17,6 +17,7 @@ import modules.prompt_parser_diffusers as prompt_parser_diffusers from modules.sd_hijack_hypertile import hypertile_set from modules.processing_correction import correction_callback from modules.processing_vae import vae_encode, vae_decode +from modules.masking import run_mask debug = shared.log.trace if os.environ.get('SD_DIFFUSERS_DEBUG', None) is not None else lambda *args, **kwargs: None @@ -167,9 +168,10 @@ def process_diffusers(p: StableDiffusionProcessing): elif (sd_models.get_diffusers_task(model) == sd_models.DiffusersTaskType.INPAINTING or is_img2img_model) and len(getattr(p, 'init_images' ,[])) > 0: p.ops.append('inpaint') if p.task_args.get('mask_image', None) is not None: # provided as override by a module - p.mask = shared.sd_model.mask_processor.blur(p.task_args['mask_image'], blur_factor=p.mask_blur) if p.mask_blur > 0 else p.task_args['mask_image'] + # p.mask = shared.sd_model.mask_processor.blur(p.task_args['mask_image'], blur_factor=p.mask_blur) if p.mask_blur > 0 else p.task_args['mask_image'] + p.mask = run_mask(input_image=p.init_images, input_mask=p.task_args['mask_image'], return_type='grayscale') elif getattr(p, 'image_mask', None) is not None: # standard - p.mask = p.image_mask + p.mask = run_mask(input_image=p.init_images, input_mask=p.image_mask, return_type='grayscale') elif getattr(p, 'mask', None) is not None: # backward compatibility pass else: # fallback diff --git a/modules/segment.py b/modules/segment.py deleted file mode 100644 index 261ecf56e..000000000 --- a/modules/segment.py +++ /dev/null @@ -1,157 +0,0 @@ -""" -[docs](https://huggingface.co/docs/transformers/v4.36.1/en/model_doc/sam#overview) -TODO: -- PerSAM -- transformers.pipeline.MaskGenerationPipeline: https://huggingface.co/models?pipeline_tag=mask-generation -- transformers.pipeline.ImageSegmentationPipeline: https://huggingface.co/models?pipeline_tag=image-segmentation -""" - -from transformers import SamModel, SamImageProcessor, MaskGenerationPipeline -from PIL import Image -import gradio as gr -import numpy as np -import cv2 -from modules import shared, devices - - -MODELS = { - 'None': None, - 'Facebook SAM ViT Base': 'facebook/sam-vit-base', - 'Facebook SAM ViT Large': 'facebook/sam-vit-large', - 'Facebook SAM ViT Huge': 'facebook/sam-vit-huge', - 'SlimSAM Uniform': 'Zigeng/SlimSAM-uniform-50', -} -COLORMAP = ['autumn', 'bone', 'jet', 'winter', 'rainbow', 'ocean', 'summer', 'spring', 'cool', 'hsv', 'pink', 'hot', 'parula', 'magma', 'inferno', 'plasma', 'viridis', 'cividis', 'twilight', 'shifted', 'turbo', 'deepgreen'] -cache_dir = 'models/control/segment' -loaded_model = None -model: SamModel = None -processor: SamImageProcessor = None - - -def init(selected_model: str, input_image: gr.Image): - global loaded_model, model, processor # pylint: disable=global-statement - if input_image is None or input_image.get('image', None) is None: - return False - if selected_model == "None": - return False - if selected_model == "None": - model = None - loaded_model = None - processor = None - model_path = MODELS[selected_model] - if model_path is not None and (loaded_model != selected_model or model is None or processor is None): - shared.log.debug(f'Segment loading: model={selected_model} path={model_path}') - model = SamModel.from_pretrained(model_path, cache_dir=cache_dir).to(device=devices.device) - processor = SamImageProcessor.from_pretrained(model_path, cache_dir=cache_dir) - shared.log.debug(f'Segment loaded: model={selected_model} path={model_path}') - if model is None or processor is None: - return False - return True - - -# run as auto-mask with all possible masks -def run_segment(selected_model: str, input_image: gr.Image, points_per_batch=64, pred_iou_thresh=0.75, stability_score_thresh=0.85, crops_nms_thresh=0.5, crop_overlap_ratio=0.3, topK=25, colormap='jet', erode=0, dilate=0): - if not init(selected_model, input_image): - return gr.update(), None - input_mask = input_image.get('mask', None) or Image.new('L', input_image.get('image', None).size, 255) - input_mask = input_mask.convert('L') - input_image = input_image.get('image', None) - generator: MaskGenerationPipeline = MaskGenerationPipeline(model=model, image_processor=processor, device=devices.device) - with devices.inference_context(): - outputs = generator( - input_image, - points_per_batch=points_per_batch, - pred_iou_thresh=pred_iou_thresh, - stability_score_thresh=stability_score_thresh, - crops_nms_thresh=crops_nms_thresh, - crop_overlap_ratio=crop_overlap_ratio, - ) - combined_mask = np.zeros(input_mask.size, dtype='uint8') - input_mask = np.array(input_mask) // 255 - input_mask_size = np.count_nonzero(input_mask) - print('HERE', input_mask.shape, input_mask_size) - i = 1 - for mask in outputs['masks']: - mask = mask.astype('uint8') - mask_size = np.count_nonzero(mask) - if mask_size == 0: - continue - overlap = 0 - if input_mask_size > 0: - overlap = cv2.bitwise_and(mask, input_mask) - overlap = np.count_nonzero(overlap) - if overlap == 0: - continue - # TODO erode,dilate - if erode > 0: - mask = cv2.erode(mask, np.ones((erode, erode), np.uint8), iterations=2) # remove noise - if dilate > 0: - mask = cv2.dilate(mask, np.ones((dilate, dilate), np.uint8), iterations=2) # expand area - mask = (topK + 1 - i) * mask * (255 // topK) # set grayscale intensity so we can recolor - combined_mask = combined_mask + mask - i += 1 - if i > topK: - break - mask_size = np.count_nonzero(combined_mask) - total_size = np.prod(combined_mask.shape) - area_size = np.count_nonzero(combined_mask) - shared.log.debug(f'Segment mask: size={input_image.width}x{input_image.height} input={input_mask_size}px masked={mask_size}px area={area_size/total_size:.2f}') - colored_mask = cv2.applyColorMap(combined_mask, COLORMAP.index(colormap)) # recolor mask - combined_image = cv2.addWeighted(np.array(input_image), 0.6, colored_mask, 0.4, 0) - _thres, binary_mask = cv2.threshold(combined_mask, 1, 255, cv2.THRESH_BINARY_INV) # create mask - binary_mask = np.invert(binary_mask) - - binary_mask = Image.fromarray(binary_mask) - combined_mask = Image.fromarray(combined_mask) - colored_mask = Image.fromarray(colored_mask) - overlay_image = Image.fromarray(combined_image) - - # TODO return type - binary_mask.save('/tmp/mask-binary.png') - combined_mask.save('/tmp/mask-combined.png') - combined_mask.save('/tmp/mask-colored.png') - overlay_image.save('/tmp/mask-overlay.png') - return input_image, overlay_image - - -# run with sam model directly needing set of points -def run_segment_points(selected_model: str, input_image: gr.Image): - if not init(selected_model, input_image): - return input_image - # input_mask = input_image.get('mask', None) or Image.new('L', input_image.get('image', None).size, 0) - input_image = input_image.get('image', None) - with devices.inference_context(): - inputs = processor( - input_image, - input_points=[[[256, 256]]], # TODO calculate points based on mask - return_tensors="pt" - ).to(device=devices.device) - outputs = model( - pixel_values=inputs['pixel_values'], - multimask_output=True, - ) - masks = processor.post_process_masks( - outputs.pred_masks.cpu(), - inputs["original_sizes"].cpu(), - inputs["reshaped_input_sizes"].cpu() - ) - scores = outputs.iou_scores - mask = masks[0].squeeze(0) - scores = scores[0].squeeze(0) - masks = mask.unbind(0) - output_masks = [] - for i, mask in enumerate(masks): - mask = mask.detach().cpu().numpy() - mask = mask.astype('uint8') * 255 - mask = cv2.dilate(mask, np.ones((3, 3), np.uint8), iterations=2) - total_size = np.prod(mask.shape) - area_size = np.count_nonzero(mask) - shared.log.debug(f'Segment mask: area={area_size/total_size:.2f} score={scores[i].item():.2f}') - mask = Image.fromarray(mask) - output_masks.append(mask) - - -def create_segment_ui(input_image: gr.Image, preview_image: gr.Image): - selected = gr.Dropdown(label="Segment", choices=MODELS.keys(), value='None') - selected.change(fn=run_segment, inputs=[selected, input_image], outputs=[input_image, preview_image]) - return selected diff --git a/modules/ui_control.py b/modules/ui_control.py index 50d3c2817..4ad193ef6 100644 --- a/modules/ui_control.py +++ b/modules/ui_control.py @@ -2,7 +2,6 @@ import os import time import gradio as gr import matplotlib.pyplot as plt -import numpy as np from PIL import Image from modules.control import unit from modules.control import processors # patrickvonplaten controlnet_aux @@ -12,7 +11,7 @@ from modules.control.units import lite # vislearn ControlNet-XS from modules.control.units import t2iadapter # TencentARC T2I-Adapter from modules.control.units import reference # reference pipeline from scripts import ipadapter # pylint: disable=no-name-in-module -from modules import errors, shared, progress, sd_samplers, ui_components, ui_symbols, ui_common, ui_sections, generation_parameters_copypaste, call_queue, scripts, segment # pylint: disable=ungrouped-imports +from modules import errors, shared, progress, sd_samplers, ui_components, ui_symbols, ui_common, ui_sections, generation_parameters_copypaste, call_queue, scripts, masking # pylint: disable=ungrouped-imports gr_height = 512 @@ -34,6 +33,7 @@ def initialize(): lite.cache_dir = os.path.join(shared.opts.control_dir, 'lite') t2iadapter.cache_dir = os.path.join(shared.opts.control_dir, 'adapter') processors.cache_dir = os.path.join(shared.opts.control_dir, 'processor') + masking.cache_dir = os.path.join(shared.opts.control_dir, 'segment') unit.default_device = devices.device unit.default_dtype = devices.dtype os.makedirs(shared.opts.control_dir, exist_ok=True) @@ -42,6 +42,7 @@ def initialize(): os.makedirs(lite.cache_dir, exist_ok=True) os.makedirs(t2iadapter.cache_dir, exist_ok=True) os.makedirs(processors.cache_dir, exist_ok=True) + os.makedirs(masking.cache_dir, exist_ok=True) scripts.scripts_current = scripts.scripts_control scripts.scripts_current.initialize_scripts(is_img2img=True) @@ -116,35 +117,7 @@ def get_video(filepath: str): return msg -def select_mask(image: Image.Image, negative: bool = False): - if image is None: - return image - image_mask = image.convert("L") - if negative: - image_mask = image_mask.point(lambda x: 255 if x < 4 else 0) - else: - image_mask = image_mask.point(lambda x: 255 if x > 127 else 0) - return image_mask - - -def expand_mask(image: Image.Image, expand: int = 64): - import cv2 - if image is None: - return image - pil_mask = image.convert("L") - np_mask = np.array(pil_mask) - erode, dilate, threshold = 3, 8, 4 - if threshold > 0: - _thres, np_mask = cv2.threshold(np_mask, threshold, 255, cv2.THRESH_BINARY_INV) # create mask - if erode > 0: - np_mask = cv2.erode(np_mask, np.ones((erode, erode), np.uint8), iterations=expand//dilate) # remove noise - if dilate > 0: - np_mask = cv2.dilate(np_mask, np.ones((dilate, dilate), np.uint8), iterations=expand//dilate) # expand area - image_mask = Image.fromarray(np_mask.astype(np.uint8)) - return image_mask - - -def select_input(input_mode, input_image, selected_init, init_type, input_resize, input_inpaint, input_video, input_batch, input_folder, _mask_blur, mask_overlap): +def select_input(input_mode, input_image, selected_init, init_type, input_resize, input_inpaint, input_video, input_batch, input_folder): global busy, input_source, input_init, input_mask # pylint: disable=global-statement busy = True if input_mode == 'Select': @@ -172,14 +145,14 @@ def select_input(input_mode, input_image, selected_init, init_type, input_resize # control inputs if isinstance(selected_input, Image.Image): # image via upload -> image if input_mode == 'Outpaint': - input_mask = expand_mask(image=selected_input, expand=mask_overlap) + input_mask = masking.run_mask(input_image=selected_input, input_mask=None, return_type='grayscale') input_source = [selected_input] input_type = 'PIL.Image' shared.log.debug(f'Control input: type={input_type} input={input_source}') status = f'Control input | Image | Size {selected_input.width}x{selected_input.height} | Mode {selected_input.mode}' res = [gr.Tabs.update(selected='out-gallery'), status] elif isinstance(selected_input, dict): # inpaint -> dict image+mask - input_mask = select_mask(image=selected_input['mask'], negative=False) + input_mask = masking.run_mask(input_image=selected_input['image'], input_mask=selected_input['mask'], return_type='grayscale') selected_input = selected_input['image'] input_source = [selected_input] input_type = 'PIL.Image' @@ -217,7 +190,7 @@ def select_input(input_mode, input_image, selected_init, init_type, input_resize elif init_type == 2: # Separate init image if isinstance(selected_init, Image.Image): # image via upload -> image if input_mode == 'Outpaint': - input_mask = expand_mask(image=selected_init, expand=mask_overlap) + input_mask = masking.run_mask(input_image=selected_init, input_mask=None, return_type='grayscale') input_source = [selected_init] input_init = [selected_init] input_type = 'PIL.Image' @@ -225,7 +198,7 @@ def select_input(input_mode, input_image, selected_init, init_type, input_resize status = f'Control input | Image | Size {selected_init.width}x{selected_init.height} | Mode {selected_init.mode}' res = [gr.Tabs.update(selected='out-gallery'), status] elif isinstance(selected_init, dict): # inpaint -> dict image+mask - input_mask = select_mask(image=selected_init['mask']) + input_mask = masking.run_mask(input_image=selected_init['image'], input_mask=selected_init['mask'], return_type='grayscale') input_init = selected_init['image'] input_source = [selected_init] input_type = 'PIL.Image' @@ -313,9 +286,6 @@ def create_ui(_blocks: gr.Blocks=None): input_type = gr.Radio(label="Input type", choices=['Control only', 'Init image same as control', 'Separate init image'], value='Control only', type='index', elem_id='control_input_type') with gr.Row(): denoising_strength = gr.Slider(minimum=0.01, maximum=1.0, step=0.01, label='Denoising strength', value=0.50, elem_id="control_denoising_strength") - with gr.Row(): - mask_blur = gr.Slider(minimum=0, maximum=100, step=1, label='Blur', value=8, elem_id="control_mask_blur") - mask_overlap = gr.Slider(minimum=0, maximum=100, step=1, label='Overlap', value=64, elem_id="control_mask_overlap") with gr.Accordion(open=False, label="Size", elem_id="control_size", elem_classes=["small-accordion"]): with gr.Tabs(): @@ -329,7 +299,11 @@ def create_ui(_blocks: gr.Blocks=None): steps, sampler_index = ui_sections.create_sampler_and_steps_selection(sd_samplers.samplers, "control") batch_count, batch_size = ui_sections.create_batch_inputs('control') + seed, _reuse_seed, subseed, _reuse_subseed, subseed_strength, seed_resize_from_h, seed_resize_from_w = ui_sections.create_seed_inputs('control', reuse_visible=False) + + masking.create_segment_ui() + cfg_scale, clip_skip, image_cfg_scale, diffusers_guidance_rescale, sag_scale, full_quality, restore_faces, tiling, hdr_clamp, hdr_boundary, hdr_threshold, hdr_center, hdr_channel_shift, hdr_full_shift, hdr_maximize, hdr_max_center, hdr_max_boundry = ui_sections.create_advanced_inputs('control') with gr.Accordion(open=False, label="Video", elem_id="control_video", elem_classes=["small-accordion"]): @@ -366,7 +340,7 @@ def create_ui(_blocks: gr.Blocks=None): input_mode = gr.Label(value='select', visible=False) input_image = gr.Image(label="Input", show_label=False, type="pil", source="upload", interactive=True, tool="editor", height=gr_height, visible=True, image_mode='RGB', elem_id='control_input_select') input_resize = gr.Image(label="Input", show_label=False, type="pil", source="upload", interactive=True, tool="select", height=gr_height, visible=False, image_mode='RGB', elem_id='control_input_resize') - input_inpaint = gr.Image(label="Input", show_label=False, type="pil", source="upload", interactive=True, tool="sketch", height=gr_height, visible=False, image_mode='RGB', elem_id='control_input_inpaint', brush_radius=64, mask_opacity=0.6) + input_inpaint = gr.Image(label="Input", show_label=False, type="pil", source="upload", interactive=True, tool="sketch", height=gr_height, visible=False, image_mode='RGB', elem_id='control_input_inpaint', brush_radius=32, mask_opacity=0.6) interrogate_clip, interrogate_booru = ui_sections.create_interrogate_buttons('control') with gr.Row(): input_buttons = [gr.Button('Select', visible=True, interactive=False), gr.Button('Inpaint', visible=True, interactive=True), gr.Button('Outpaint', visible=True, interactive=True)] @@ -402,10 +376,6 @@ def create_ui(_blocks: gr.Blocks=None): with gr.Tab('Preview', id='preview-image') as tab_image: preview_process = gr.Image(label="Input", show_label=False, type="pil", source="upload", interactive=False, height=gr_height, visible=True) - # TODO segment as accordian - # with gr.Row(): - # segment_ui = segment.create_segment_ui(input_inpaint, preview_process) - with gr.Tabs(elem_id='control-tabs') as _tabs_control_type: with gr.Tab('ControlNet') as _tab_controlnet: @@ -684,7 +654,7 @@ def create_ui(_blocks: gr.Blocks=None): interrogate_clip.click(fn=ui_common.interrogate_clip, inputs=[input_image], outputs=[prompt]) interrogate_booru.click(fn=ui_common.interrogate_booru, inputs=[input_image], outputs=[prompt]) - select_fields = [input_mode, input_image, init_image, input_type, input_resize, input_inpaint, input_video, input_batch, input_folder, mask_blur, mask_overlap] + select_fields = [input_mode, input_image, init_image, input_type, input_resize, input_inpaint, input_video, input_batch, input_folder] select_output = [output_tabs, result_txt] select_dict = dict( fn=select_input, @@ -711,7 +681,7 @@ def create_ui(_blocks: gr.Blocks=None): cfg_scale, clip_skip, image_cfg_scale, diffusers_guidance_rescale, sag_scale, full_quality, restore_faces, tiling, hdr_clamp, hdr_boundary, hdr_threshold, hdr_center, hdr_channel_shift, hdr_full_shift, hdr_maximize, hdr_max_center, hdr_max_boundry, resize_mode_before, resize_name_before, width_before, height_before, scale_by_before, selected_scale_tab_before, resize_mode_after, resize_name_after, width_after, height_after, scale_by_after, selected_scale_tab_after, - denoising_strength, batch_count, batch_size, mask_blur, mask_overlap, + denoising_strength, batch_count, batch_size, video_skip_frames, video_type, video_duration, video_loop, video_pad, video_interpolate, ip_adapter, ip_scale, ip_image, ] @@ -736,6 +706,8 @@ def create_ui(_blocks: gr.Blocks=None): generation_parameters_copypaste.add_paste_fields("control", input_image, paste_fields, override_settings) bindings = generation_parameters_copypaste.ParamBinding(paste_button=btn_paste, tabname="control", source_text_component=prompt, source_image_component=output_gallery) generation_parameters_copypaste.register_paste_params_button(bindings) + masking.bind_controls(input_inpaint, preview_process) + if os.environ.get('SD_CONTROL_DEBUG', None) is not None: # debug only from modules.control.test import test_processors, test_controlnets, test_adapters, test_xs, test_lite