refactor control to use new masking

CHANGELOG.md

@@ -17,6 +17,7 @@ And it also includes fixes for all reported issues so far
     - optional **live preview**  
     - optional **auto-segmentation** (e.g. segment-anything) using ml models  
       *note*: auto segmentation will automatically expand user-masked area to segments that include current user mask  
+    - can be combined with control processors in which case mask is applied before processor  
   - 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  

cli/install-sf.py

@@ -0,0 +1,57 @@
+import os
+import sys
+
+ver = '1.0.1.dev20240109'
+torch_supported = ['211', '212']
+cuda_supported = ['cu118', 'cu121']
+python_supported = ['39', '310', '311']
+repo_url = 'https://github.com/chengzeyi/stable-fast'
+path_url = '/releases/download/nightly'
+
+
+def install_pip(arg: str):
+    import subprocess
+    cmd = f'"{sys.executable}" -m pip install -U {arg}'
+    print(f'Running: {cmd}')
+    result = subprocess.run(cmd, shell=True, check=False, env=os.environ)
+    return result.returncode == 0
+
+
+def install_stable_fast():
+    import torch
+
+    python_ver = f'{sys.version_info.major}{sys.version_info.minor}'
+    if python_ver not in python_supported:
+        raise ValueError(f'StableFast unsupported python: {python_ver} required {python_supported}')
+    if sys.platform == 'linux':
+        bin_url = 'manylinux2014_x86_64.whl'
+    elif sys.platform == 'win32':
+        bin_url = 'win_amd64.whl'
+    else:
+        raise ValueError(f'StableFast unsupported platform: {sys.platform}')
+
+    torch_ver, cuda_ver = torch.__version__.split('+')
+    torch_ver = torch_ver.replace('.', '')
+
+    if torch_ver not in torch_supported:
+        print(f'StableFast unsupported torch: {torch_ver} required {torch_supported}')
+        print('Installing from source...')
+        url = 'git+https://github.com/chengzeyi/stable-fast.git@main#egg=stable-fast'
+    elif cuda_ver not in cuda_supported:
+        print(f'StableFast unsupported CUDA: {cuda_ver} required {cuda_supported}')
+        print('Installing from source...')
+        url = 'git+https://github.com/chengzeyi/stable-fast.git@main#egg=stable-fast'
+    else:
+        print('Installing wheel...')
+        file_url = f'stable_fast-{ver}+torch{torch_ver}{cuda_ver}-cp{python_ver}-cp{python_ver}-{bin_url}'
+        url = f'{repo_url}/{path_url}/{file_url}'
+
+    ok = install_pip(url)
+    if ok:
+        import sfast
+        print(f'StableFast installed: {sfast.__version__}')
+    else:
+        print('StableFast install failed')
+
+if __name__ == '__main__':
+    install_stable_fast()

modules/control/processors.py

@@ -235,5 +235,10 @@ class Processor():
             image_process = image_process.convert(mode)
         return image_process
 
-    def preview(self, image_input: Image):
-        return self.__call__(image_input)
+    def preview(self):
+        import modules.ui_control
+        input_image = modules.ui_control.input_source
+        if isinstance(input_image, list):
+            input_image = input_image[0]
+        if isinstance(input_image, Image.Image):
+            return self.__call__(input_image)

modules/control/run.py

@@ -15,7 +15,7 @@ from modules.control.units import lite # Kohya ControlLLLite
 from modules.control.units import t2iadapter # TencentARC T2I-Adapter
 from modules.control.units import reference # ControlNet-Reference
 from scripts import ipadapter # pylint: disable=no-name-in-module
-from modules import devices, shared, errors, processing, images, sd_models, scripts # pylint: disable=ungrouped-imports
+from modules import devices, shared, errors, processing, images, sd_models, scripts, masking # pylint: disable=ungrouped-imports
 
 
 debug = shared.log.trace if os.environ.get('SD_CONTROL_DEBUG', None) is not None else lambda *args, **kwargs: None
@@ -363,45 +363,21 @@ def control_run(units: List[unit.Unit], inputs, inits, mask, unit_type: str, is_
 
                     # process
                     if input_image is None:
-                        p.image = None
-                        processed_image = None
+                        p.image = []
                         debug(f'Control: process=None image={p.image} mask={mask}')
-                    elif mask is not None and not has_models:
-                        processed_image = mask
-                        debug(f'Control: process=None image={p.image} mask={mask}')
-                    elif len(active_process) == 0 and unit_type == 'reference':
-                        p.ref_image = p.override or input_image
-                        p.task_args['ref_image'] = p.ref_image
-                        debug(f'Control: process=None image={p.ref_image}')
-                        if p.ref_image is None:
-                            msg = 'Control: attempting reference mode but image is none'
-                            shared.log.error(msg)
-                            restore_pipeline()
-                            return msg
-                        processed_image = p.ref_image
-                    elif len(active_process) == 1:
-                        image_mode = 'L' if unit_type == 'adapter' and len(active_model) > 0 and ('Canny' in active_model[0].model_id or 'Sketch' in active_model[0].model_id) else 'RGB'
-                        p.image = active_process[0](input_image, image_mode)
-                        p.task_args['image'] = p.image
-                        p.extra_generation_params["Control process"] = active_process[0].processor_id
-                        debug(f'Control: process={active_process[0].processor_id} image={p.image}')
-                        if p.image is None:
-                            msg = 'Control: attempting process but output is none'
-                            shared.log.error(msg)
-                            restore_pipeline()
-                            return msg
-                        processed_image = p.image
-                    else:
-                        if len(active_process) > 0:
-                            p.image = []
-                            for i, process in enumerate(active_process): # list[image]
-                                image_mode = 'L' if unit_type == 'adapter' and len(active_model) > i and ('Canny' in active_model[i].model_id or 'Sketch' in active_model[i].model_id) else 'RGB'
-                                p.image.append(process(input_image, image_mode))
-                        else:
-                            p.image = [input_image]
+                    elif len(active_process) == 0:
+                        p.image = [masking.run_mask(input_image=input_image, input_mask=mask, return_type='masked') if mask is not None else input_image]
+                    elif len(active_process) > 0:
+                        p.image = []
+                        masked_image = masking.run_mask(input_image=input_image, input_mask=mask, return_type='masked') if mask is not None else input_image
+                        for i, process in enumerate(active_process): # list[image]
+                            image_mode = 'L' if unit_type == 'adapter' and len(active_model) > i and ('Canny' in active_model[i].model_id or 'Sketch' in active_model[i].model_id) else 'RGB' # t2iadapter canny and sketch work in grayscale only
+                            debug(f'Control: process={[process.processor_id for p in active_process]} i={i} image={p.image}')
+                            p.image.append(process(masked_image, image_mode))
+
+                    if len(p.image) > 0:
                         p.task_args['image'] = p.image
                         p.extra_generation_params["Control process"] = [p.processor_id for p in active_process]
-                        debug(f'Control: process={[p.processor_id for p in active_process]} image={p.image}')
                         if any(img is None for img in p.image):
                             msg = 'Control: attempting process but output is none'
                             shared.log.error(msg)
@@ -410,8 +386,20 @@ def control_run(units: List[unit.Unit], inputs, inits, mask, unit_type: str, is_
                         processed_image = [np.array(i) for i in p.image]
                         processed_image = util.blend(processed_image) # blend all processed images into one
                         processed_image = Image.fromarray(processed_image)
+                    else:
+                        processed_image = input_image
 
-                    if unit_type == 'controlnet' and input_type == 1: # Init image same as control
+                    if unit_type == 'reference':
+                        p.ref_image = p.override or input_image
+                        p.task_args.pop('image', None)
+                        p.task_args['ref_image'] = p.ref_image
+                        debug(f'Control: process=None image={p.ref_image}')
+                        if p.ref_image is None:
+                            msg = 'Control: attempting reference mode but image is none'
+                            shared.log.error(msg)
+                            restore_pipeline()
+                            return msg
+                    elif unit_type == 'controlnet' and input_type == 1: # Init image same as control
                         p.task_args['image'] = input_image
                         p.task_args['control_image'] = p.image
                         p.task_args['strength'] = p.denoising_strength

modules/control/unit.py

@@ -29,7 +29,6 @@ class Unit(): # mashup of gradio controls and mapping to actual implementation c
                  preview_btn = None,
                  model_id = None,
                  model_strength = None,
-                 image_input = None,
                  preview_process = None,
                  image_upload = None,
                  image_preview = None,
@@ -50,7 +49,6 @@ class Unit(): # mashup of gradio controls and mapping to actual implementation c
         self.adapter: t2iadapter.Adapter = None
         self.controlnet: Union[controlnet.ControlNet, xs.ControlNetXS] = None
         # map to input image
-        self.input: Image = image_input
         self.override: Image = None
         # global settings but passed per-unit
         self.factor = 1.0
@@ -161,7 +159,7 @@ class Unit(): # mashup of gradio controls and mapping to actual implementation c
         if reset_btn is not None:
             reset_btn.click(fn=reset, inputs=[], outputs=[enabled_cb, model_id, process_id, model_strength])
         if preview_btn is not None:
-            preview_btn.click(fn=self.process.preview, inputs=[self.input], outputs=[preview_process]) # return list of images for gallery
+            preview_btn.click(fn=self.process.preview, inputs=[], outputs=[preview_process]) # return list of images for gallery
         if image_upload is not None:
             image_upload.upload(fn=upload_image, inputs=[image_upload], outputs=[image_preview]) # return list of images for gallery
         if control_start is not None and control_end is not None:

modules/control/units/controlnet.py

@@ -48,7 +48,7 @@ def find_models():
     files = [f for f in files if f.endswith('.safetensors')]
     downloaded_models = {}
     for f in files:
-        basename = os.path.splitext(f)[0]
+        basename = os.path.splitext(os.path.relpath(f, path))[0]
         downloaded_models[basename] = os.path.join(path, f)
     all_models.update(downloaded_models)
     return downloaded_models

modules/control/units/lite.py

@@ -35,7 +35,7 @@ def find_models():
     files = [f for f in files if f.endswith('.safetensors')]
     downloaded_models = {}
     for f in files:
-        basename = os.path.splitext(f)[0]
+        basename = os.path.splitext(os.path.relpath(f, path))[0]
         downloaded_models[basename] = os.path.join(path, f)
     all_models.update(downloaded_models)
     return downloaded_models

modules/control/units/xs.py

@@ -31,7 +31,7 @@ def find_models():
     files = [f for f in files if f.endswith('.safetensors')]
     downloaded_models = {}
     for f in files:
-        basename = os.path.splitext(f)[0]
+        basename = os.path.splitext(os.path.relpath(f, path))[0]
         downloaded_models[basename] = os.path.join(path, f)
     all_models.update(downloaded_models)
     return downloaded_models

modules/errors.py

@@ -9,7 +9,7 @@ from installer import log as installer_log, setup_logging
 
 setup_logging()
 log = installer_log
-console = Console(log_time=True, log_time_format='%H:%M:%S-%f', theme=Theme({
+console = Console(log_time=True, tab_size=4, log_time_format='%H:%M:%S-%f', soft_wrap=True, safe_box=True, theme=Theme({
     "traceback.border": "black",
     "traceback.border.syntax_error": "black",
     "inspect.value.border": "black",
@@ -37,7 +37,7 @@ def print_error_explanation(message):
 
 def display(e: Exception, task, suppress=[]): # noqa: B006
     log.error(f"{task or 'error'}: {type(e).__name__}")
-    console.print_exception(show_locals=False, max_frames=10, extra_lines=1, suppress=suppress, theme="ansi_dark", word_wrap=False, width=min([console.width, 200]))
+    console.print_exception(show_locals=False, max_frames=10, extra_lines=1, suppress=suppress, theme="ansi_dark", word_wrap=False, width=console.width)
 
 
 def display_once(e: Exception, task):

modules/masking.py

@@ -242,6 +242,7 @@ def run_mask(input_image: gr.Image, input_mask: gr.Image = None, return_type: st
         mask = input_mask * 255
     else:
         mask = run_segment(input_image, input_mask)
+    mask = cv2.resize(mask, (input_image.width, input_image.height), interpolation=cv2.INTER_LINEAR)
 
     if mask is None:
         shared.log.error('Segment error: no mask')
@@ -276,9 +277,6 @@ def run_mask(input_image: gr.Image, input_mask: gr.Image = None, return_type: st
     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
@@ -286,12 +284,22 @@ def run_mask(input_image: gr.Image, input_mask: gr.Image = None, return_type: st
     if return_type == 'none':
         return input_mask
     elif return_type == 'binary':
+        binary_mask = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)[1] # otsu uses mean instead of threshold
         return Image.fromarray(binary_mask)
+    elif return_type == 'masked':
+        orig = np.array(input_image)
+        mask = cv2.cvtColor(mask, cv2.COLOR_GRAY2RGB)
+        masked_image = cv2.bitwise_and(orig, mask)
+        return Image.fromarray(masked_image)
     elif return_type == 'grayscale':
         return Image.fromarray(mask)
     elif return_type == 'color':
+        colored_mask = cv2.applyColorMap(mask, COLORMAP.index(opts.seg_colormap)) # recolor mask
         return Image.fromarray(colored_mask)
     elif return_type == 'composite':
+        colored_mask = cv2.applyColorMap(mask, COLORMAP.index(opts.seg_colormap)) # recolor mask
+        orig = np.array(input_image)
+        combined_image = cv2.addWeighted(orig, opts.weight_original, colored_mask, opts.weight_mask, 0)
         return Image.fromarray(combined_image)
     return input_mask
 
@@ -340,7 +348,7 @@ def create_segment_ui():
             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="Preview", choices=['none', 'masked', '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])

modules/processing.py

@@ -1301,7 +1301,8 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
                     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 hasattr(self, 'init_images'):
+                    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')

modules/ui_control.py

@@ -404,7 +404,6 @@ def create_ui(_blocks: gr.Blocks=None):
                         units.append(unit.Unit(
                             unit_type = 'controlnet',
                             result_txt = result_txt,
-                            image_input = input_image,
                             enabled_cb = enabled_cb,
                             reset_btn = reset_btn,
                             process_id = process_id,
@@ -456,7 +455,6 @@ def create_ui(_blocks: gr.Blocks=None):
                         units.append(unit.Unit(
                             unit_type = 'adapter',
                             result_txt = result_txt,
-                            image_input = input_image,
                             enabled_cb = enabled_cb,
                             reset_btn = reset_btn,
                             process_id = process_id,
@@ -499,7 +497,6 @@ def create_ui(_blocks: gr.Blocks=None):
                         units.append(unit.Unit(
                             unit_type = 'xs',
                             result_txt = result_txt,
-                            image_input = input_image,
                             enabled_cb = enabled_cb,
                             reset_btn = reset_btn,
                             process_id = process_id,
@@ -541,7 +538,6 @@ def create_ui(_blocks: gr.Blocks=None):
                         units.append(unit.Unit(
                             unit_type = 'lite',
                             result_txt = result_txt,
-                            image_input = input_image,
                             enabled_cb = enabled_cb,
                             reset_btn = reset_btn,
                             process_id = process_id,
@@ -580,7 +576,6 @@ def create_ui(_blocks: gr.Blocks=None):
                         units.append(unit.Unit(
                             unit_type = 'reference',
                             result_txt = result_txt,
-                            image_input = input_image,
                             enabled_cb = enabled_cb,
                             reset_btn = reset_btn,
                             process_id = process_id,
@@ -666,11 +661,12 @@ def create_ui(_blocks: gr.Blocks=None):
                 )
                 prompt.submit(**select_dict)
                 btn_generate.click(**select_dict)
-                for ctrl in [input_image, 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]:
+                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)
-                    elif hasattr(ctrl, 'select'):
-                        ctrl.select(**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)
 
                 tabs_state = gr.Text(value='none', visible=False)
                 input_fields = [