From 15b4bad80afa7b8ce0e90bdc9501907a8d9a7a09 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Mon, 15 Jan 2024 12:11:55 -0500 Subject: [PATCH] refactor control to use new masking --- CHANGELOG.md | 1 + cli/install-sf.py | 57 +++++++++++++++++++++++ extensions-builtin/sd-webui-controlnet | 2 +- modules/control/processors.py | 9 +++- modules/control/run.py | 64 +++++++++++--------------- modules/control/unit.py | 4 +- modules/control/units/controlnet.py | 2 +- modules/control/units/lite.py | 2 +- modules/control/units/xs.py | 2 +- modules/errors.py | 4 +- modules/masking.py | 16 +++++-- modules/processing.py | 3 +- modules/ui_control.py | 12 ++--- 13 files changed, 116 insertions(+), 62 deletions(-) create mode 100644 cli/install-sf.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 9f59f4e17..3477f8222 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/cli/install-sf.py b/cli/install-sf.py new file mode 100644 index 000000000..3cedb65d2 --- /dev/null +++ b/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() diff --git a/extensions-builtin/sd-webui-controlnet b/extensions-builtin/sd-webui-controlnet index 84ef20567..8b5f7c1d0 160000 --- a/extensions-builtin/sd-webui-controlnet +++ b/extensions-builtin/sd-webui-controlnet @@ -1 +1 @@ -Subproject commit 84ef2056737dab2c63e708464c547926168b1088 +Subproject commit 8b5f7c1d0e6b857b9033a54226ad83714553770d diff --git a/modules/control/processors.py b/modules/control/processors.py index 76a3b6e19..6f30f579c 100644 --- a/modules/control/processors.py +++ b/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) diff --git a/modules/control/run.py b/modules/control/run.py index 7c4ccc16f..01ac6d38a 100644 --- a/modules/control/run.py +++ b/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 diff --git a/modules/control/unit.py b/modules/control/unit.py index b4505e91d..de7decdd4 100644 --- a/modules/control/unit.py +++ b/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: diff --git a/modules/control/units/controlnet.py b/modules/control/units/controlnet.py index 444f151c4..c339a227c 100644 --- a/modules/control/units/controlnet.py +++ b/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 diff --git a/modules/control/units/lite.py b/modules/control/units/lite.py index cc385610c..eb97165dd 100644 --- a/modules/control/units/lite.py +++ b/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 diff --git a/modules/control/units/xs.py b/modules/control/units/xs.py index 673f4b90d..e1c7a561e 100644 --- a/modules/control/units/xs.py +++ b/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 diff --git a/modules/errors.py b/modules/errors.py index 8cb2eb599..0e5a0b094 100644 --- a/modules/errors.py +++ b/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): diff --git a/modules/masking.py b/modules/masking.py index 1c33e9724..9d0015668 100644 --- a/modules/masking.py +++ b/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]) diff --git a/modules/processing.py b/modules/processing.py index a567a81d7..201b759fe 100644 --- a/modules/processing.py +++ b/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') diff --git a/modules/ui_control.py b/modules/ui_control.py index 4ad193ef6..fdc1c108c 100644 --- a/modules/ui_control.py +++ b/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 = [