add lama to control masking options

CHANGELOG.md

@@ -47,6 +47,9 @@ And it also includes fixes for all reported issues so far
       if you dont provide mask or mask is empty, you can instead use auto-mask to automatically generate mask  
       this is especially useful if you want to use advanced masking on batch or video inputs and dont want to manually mask each image  
       *note*: such auto-created mask is also subject to all other selected settings such as auto-segmentation, blur, erode and dilate  
+    - optional **object removal** using LaMA model  
+      remove selected objects from images with a single click  
+      works best when combined with auto-segmentation to remove smaller objects  
     - masking can be combined with control processors in which case mask is applied before processor  
   - support for many additional controlnet models  
     now built-in models include 30+ SD15 models and 15+ SDXL models  

html/locale_en.json

@@ -20,6 +20,7 @@
   {"id":"","label":"☲","localized":"","hint":"Change view type"},
   {"id":"","label":"📐","localized":"","hint":"Measure"},
   {"id":"","label":"🔍","localized":"","hint":"Search"},
+  {"id":"","label":"🖼️","localized":"","hint":"LaMa remove selected object from image"},
   {"id":"","label":"🖼️","localized":"","hint":"Show preview"},
   {"id":"","label":"✎","localized":"","hint":"Interrogate image using BLIP model"},
   {"id":"","label":"✐","localized":"","hint":"Interrogate image using DeepBooru model"}

modules/control/run.py

@@ -260,7 +260,11 @@ def control_run(units: List[unit.Unit], inputs, inits, mask, unit_type: str, is_
     else: # run in txt2img/img2img mode
         if len(active_strength) > 0:
             p.strength = active_strength[0]
-        pipe = diffusers.AutoPipelineForText2Image.from_pipe(shared.sd_model) # use set_diffuser_pipe
+        try:
+            pipe = diffusers.AutoPipelineForText2Image.from_pipe(shared.sd_model) # use set_diffuser_pipe
+        except Exception as e:
+            shared.log.warning(f'Control pipeline create: {e}')
+            pipe = shared.sd_model
         instance = None
 
     debug(f'Control pipeline: class={pipe.__class__} args={vars(p)}')

modules/lama.py

@@ -0,0 +1,96 @@
+import os
+from urllib.parse import urlparse
+import cv2
+import torch
+import numpy as np
+from torch.hub import download_url_to_file, get_dir
+from PIL import Image
+from modules import devices
+from modules.shared import log
+
+
+LAMA_MODEL_URL = "https://github.com/enesmsahin/simple-lama-inpainting/releases/download/v0.1.0/big-lama.pt"
+
+
+def prepare_img_and_mask(image, mask, device, pad_out_to_modulo=8, scale_factor=None):
+    def ceil_modulo(x, mod):
+        if x % mod == 0:
+            return x
+        return (x // mod + 1) * mod
+
+    def get_image(img):
+        if isinstance(img, Image.Image):
+            img = np.array(img)
+        if img.ndim == 3:
+            img = np.transpose(img, (2, 0, 1))  # chw
+        elif img.ndim == 2:
+            img = img[np.newaxis, ...]
+        img = img.astype(np.float32) / 255
+        return img
+
+    def pad_img_to_modulo(img, mod):
+        _channels, height, width = img.shape
+        out_height = ceil_modulo(height, mod)
+        out_width = ceil_modulo(width, mod)
+        return np.pad(
+            img,
+            ((0, 0), (0, out_height - height), (0, out_width - width)),
+            mode="symmetric",
+        )
+
+    def scale_image(img, factor, interpolation=cv2.INTER_AREA):
+        if img.shape[0] == 1:
+            img = img[0]
+        else:
+            img = np.transpose(img, (1, 2, 0))
+        img = cv2.resize(img, dsize=None, fx=factor, fy=factor, interpolation=interpolation)
+        if img.ndim == 2:
+            img = img[None, ...]
+        else:
+            img = np.transpose(img, (2, 0, 1))
+        return img
+
+    out_image = get_image(image)
+    out_mask = get_image(mask)
+    if scale_factor is not None:
+        out_image = scale_image(out_image, scale_factor)
+        out_mask = scale_image(out_mask, scale_factor, interpolation=cv2.INTER_NEAREST)
+    if pad_out_to_modulo is not None and pad_out_to_modulo > 1:
+        out_image = pad_img_to_modulo(out_image, pad_out_to_modulo)
+        out_mask = pad_img_to_modulo(out_mask, pad_out_to_modulo)
+    out_image = torch.from_numpy(out_image).unsqueeze(0).to(device)
+    out_mask = torch.from_numpy(out_mask).unsqueeze(0).to(device)
+    out_mask = (out_mask > 0) * 1
+    return out_image, out_mask
+
+
+def download_model():
+    parts = urlparse(LAMA_MODEL_URL)
+    hub_dir = get_dir()
+    model_dir = os.path.join(hub_dir, "checkpoints")
+    os.makedirs(os.path.join(model_dir, "hub", "checkpoints"), exist_ok=True)
+    filename = os.path.basename(parts.path)
+    cached_file = os.path.join(model_dir, filename)
+    if not os.path.exists(cached_file):
+        log.info(f'LaMa download: url={LAMA_MODEL_URL} file={cached_file}')
+        hash_prefix = None
+        download_url_to_file(LAMA_MODEL_URL, cached_file, hash_prefix, progress=True)
+    return cached_file
+
+
+class SimpleLama:
+    def __init__(self):
+        self.device = devices.device
+        model_path = download_model()
+        self.model = torch.jit.load(model_path, map_location=self.device)
+        self.model.eval()
+        self.model.to(self.device)
+
+    def __call__(self, image: Image.Image | np.ndarray, mask: Image.Image | np.ndarray):
+        image, mask = prepare_img_and_mask(image, mask, self.device)
+        with devices.inference_context():
+            inpainted = self.model(image, mask)
+            cur_res = inpainted[0].permute(1, 2, 0).detach().cpu().numpy()
+            cur_res = np.clip(cur_res * 255, 0, 255).astype(np.uint8)
+            cur_res = Image.fromarray(cur_res)
+            return cur_res

modules/masking.py

@@ -125,7 +125,9 @@ 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
+btn_mask = None
+btn_lama = None
+lama_model = None
 controls = []
 opts = SimpleNamespace(**{
     'auto_mask': 'None',
@@ -316,6 +318,8 @@ def run_mask(input_image: gr.Image, input_mask: gr.Image = None, return_type: st
     if isinstance(input_image, dict):
         input_mask = input_image.get('mask', None)
         input_image = input_image.get('image', None)
+    if input_image is None:
+        return input_mask
 
     t0 = time.time()
     input_mask = get_mask(input_image, input_mask) # perform optional auto-masking
@@ -397,6 +401,21 @@ def run_mask(input_image: gr.Image, input_mask: gr.Image = None, return_type: st
     return input_mask
 
 
+def run_lama(input_image: gr.Image, input_mask: gr.Image = None):
+    global lama_model # pylint: disable=global-statement
+    if isinstance(input_image, dict):
+        input_mask = input_image.get('mask', None)
+        input_image = input_image.get('image', None)
+    if input_image is None:
+        return None
+    input_mask = run_mask(input_image, input_mask, return_type='Grayscale')
+    if lama_model is None:
+        from modules.lama import SimpleLama
+        lama_model = SimpleLama()
+    result = lama_model(input_image, input_mask)
+    return result
+
+
 def run_mask_live(input_image: gr.Image):
     global busy # pylint: disable=global-statement
     if opts.seg_live:
@@ -423,12 +442,13 @@ def create_segment_ui():
         opts.preview_type = args[9]
         opts.seg_colormap = args[10]
 
-    global btn_segment # pylint: disable=global-statement
+    global btn_mask, btn_lama # 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))
-            btn_segment = ui_components.ToolButton(value=ui_symbols.refresh, visible=True)
+            btn_mask = ui_components.ToolButton(value=ui_symbols.refresh, visible=True)
+            btn_lama = ui_components.ToolButton(value=ui_symbols.image, visible=True)
         with gr.Row():
             controls.append(gr.Checkbox(label="Inpaint masked only", value=False))
         with gr.Row():
@@ -451,9 +471,10 @@ def create_segment_ui():
             control.change(fn=update_opts, inputs=controls, outputs=[])
 
 
-def bind_controls(image_controls: List[gr.Image], preview_image: gr.Image):
+def bind_controls(image_controls: List[gr.Image], preview_image: gr.Image, output_image: gr.Image):
     for image_control in image_controls:
-        btn_segment.click(run_mask, inputs=[image_control], outputs=[preview_image])
+        btn_mask.click(run_mask, inputs=[image_control], outputs=[preview_image])
+        btn_lama.click(run_lama, inputs=[image_control], outputs=[output_image])
         image_control.edit(fn=run_mask_live, inputs=[image_control], outputs=[preview_image])
         for control in controls:
             control.change(fn=run_mask_live, inputs=[image_control], outputs=[preview_image])

modules/ui_control.py

@@ -135,7 +135,9 @@ def select_input(input_mode, input_image, selected_init, init_type, input_resize
     else:
         selected_input = None
     if selected_input is None:
+        input_source = None
         busy = False
+        debug('Control clear input')
         return [gr.Tabs.update(), '']
     debug(f'Control select input: source={selected_input} init={selected_init} type={init_type} mode={input_mode}')
     input_type = type(selected_input)
@@ -655,6 +657,8 @@ def create_ui(_blocks: gr.Blocks=None):
                 for ctrl in [input_image, input_resize, input_video, input_batch, input_folder, init_image, init_video, init_batch, init_folder, tab_image, tab_video, tab_batch, tab_folder, tab_image_init, tab_video_init, tab_batch_init, tab_folder_init]:
                     if hasattr(ctrl, 'change'):
                         ctrl.change(**select_dict)
+                    if hasattr(ctrl, 'clear'):
+                        ctrl.clear(**select_dict)
                 for ctrl in [input_inpaint]: # gradio image mode inpaint triggeres endless loop on change event
                     if hasattr(ctrl, 'upload'):
                         ctrl.upload(**select_dict)
@@ -693,7 +697,7 @@ 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_image, input_inpaint, input_resize], preview_process)
+                masking.bind_controls([input_image, input_inpaint, input_resize], preview_process, output_image)
 
 
                 if os.environ.get('SD_CONTROL_DEBUG', None) is not None: # debug only

modules/ui_symbols.py

@@ -22,6 +22,7 @@ reset = '🔄'
 upload = '⬆️'
 search = '🔍'
 preview = '🖼️'
+image = '🖌️'
 mark_diag = '※'
 mark_flag = '⁜'
 int_clip = '✎'