From 3baff5906b897aa4bc314b48a889cc76c3f7fa72 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Tue, 30 Jan 2024 11:26:34 -0500 Subject: [PATCH] control enable image drag, prompt parse, add metadata --- CHANGELOG.md | 2 - javascript/imageParams.js | 7 +- modules/control/run.py | 8 + modules/images.py | 12 +- modules/masking.py | 18 ++- modules/processing.py | 5 +- modules/ui_control.py | 291 +++++++--------------------------- modules/ui_control_helpers.py | 233 +++++++++++++++++++++++++++ 8 files changed, 324 insertions(+), 252 deletions(-) create mode 100644 modules/ui_control_helpers.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 89541a311..a71cd777b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,10 +9,8 @@ OPTIONAL: - pending `diffusers==0.26.0` - wuerstchen v3 [pr](https://github.com/huggingface/diffusers/pull/6487) - style aligned [pr](https://github.com/huggingface/diffusers/pull/6489) -- mixture tiling [pr](https://github.com/huggingface/diffusers/tree/main/examples/community#stable-diffusion-mixture-tiling) - instaflow [pr](https://github.com/huggingface/diffusers/pull/6057)[repo](https://github.com/gnobitab/RectifiedFlow) - control api -- control params paste - masking api - preprocess api diff --git a/javascript/imageParams.js b/javascript/imageParams.js index 0b739cd00..e31aaa667 100644 --- a/javascript/imageParams.js +++ b/javascript/imageParams.js @@ -8,7 +8,12 @@ async function initDragDrop() { const target = e.composedPath()[0]; if (!target.placeholder) return; if (target.placeholder.indexOf('Prompt') === -1) return; - const promptTarget = get_tab_index('tabs') === 1 ? 'img2img_prompt_image' : 'txt2img_prompt_image'; + const tab = get_tab_index('tabs'); + let promptTarget = ''; + if (tab === 0) promptTarget = 'txt2img_prompt_image'; + else if (tab === 1) promptTarget = 'img2img_prompt_image'; + else if (tab === 2) promptTarget = 'control_prompt_image'; + else return; const imgParent = gradioApp().getElementById(promptTarget); const fileInput = imgParent.querySelector('input[type="file"]'); if (!imgParent || !fileInput) return; diff --git a/modules/control/run.py b/modules/control/run.py index 7132453fe..57186eb9b 100644 --- a/modules/control/run.py +++ b/modules/control/run.py @@ -383,6 +383,14 @@ def control_run(units: List[unit.Unit], inputs, inits, mask, unit_type: str, is_ processed_images = [] masked_image = masking.run_mask(input_image=input_image, input_mask=mask, return_type='Masked') if mask is not None else input_image + if mask is not None: + p.extra_generation_params["Mask only"] = masking.opts.mask_only if masking.opts.mask_only else None + p.extra_generation_params["Mask auto"] = masking.opts.auto_mask if masking.opts.auto_mask != 'None' else None + p.extra_generation_params["Mask invert"] = masking.opts.invert if masking.opts.invert else None + p.extra_generation_params["Mask blur"] = masking.opts.mask_blur if masking.opts.mask_blur > 0 else None + p.extra_generation_params["Mask erode"] = masking.opts.mask_erode if masking.opts.mask_erode > 0 else None + p.extra_generation_params["Mask dilate"] = masking.opts.mask_dilate if masking.opts.mask_dilate > 0 else None + p.extra_generation_params["Mask model"] = masking.opts.model if masking.opts.model is not None else None 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: i={i+1} process="{process.processor_id}" input={masked_image} override={process.override}') diff --git a/modules/images.py b/modules/images.py index f9ecee143..7488affdd 100644 --- a/modules/images.py +++ b/modules/images.py @@ -805,18 +805,18 @@ def image_data(data): err2 = None try: image = Image.open(io.BytesIO(data)) - errors.log.debug(f'Decoded object: image={image}') - textinfo, _ = read_info_from_image(image) - return textinfo, None + info, _ = read_info_from_image(image) + errors.log.debug(f'Decoded object: image={image} metadata={info}') + return info, None except Exception as e: err1 = e try: if len(data) > 1024 * 10: errors.log.warning(f'Error decoding object: data too long: {len(data)}') return gr.update(), None - text = data.decode('utf8') - errors.log.debug(f'Decoded object: size={len(text)}') - return text, None + info = data.decode('utf8') + errors.log.debug(f'Decoded object: data={len(data)} metadata={info}') + return info, None except Exception as e: err2 = e errors.log.error(f'Error decoding object: {err1 or err2}') diff --git a/modules/masking.py b/modules/masking.py index c98fb21b6..a15f881ed 100644 --- a/modules/masking.py +++ b/modules/masking.py @@ -121,7 +121,6 @@ MODELS = { } 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 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') @@ -131,6 +130,7 @@ btn_lama = None lama_model = None controls = [] opts = SimpleNamespace(**{ + 'model': None, 'auto_mask': 'None', 'mask_only': False, 'mask_blur': 0.01, @@ -154,21 +154,21 @@ opts = SimpleNamespace(**{ def init_model(selected_model: str): - global busy, loaded_model, generator # pylint: disable=global-statement + global busy, generator # pylint: disable=global-statement model_path = MODELS[selected_model] if model_path is None: # none if generator is not None: shared.log.debug('Segment unloading model') - loaded_model = None + opts.model = None generator = None devices.torch_gc() return selected_model if 'Rembg' in selected_model: # rembg - loaded_model = model_path + opts.model = model_path generator = None devices.torch_gc() return selected_model - if loaded_model != selected_model or generator is None: # sam pipeline + if opts.model != selected_model or generator is None: # sam pipeline busy = True t0 = time.time() shared.log.debug(f'Segment loading: model={selected_model} path={model_path}') @@ -183,7 +183,7 @@ def init_model(selected_model: str): ) devices.torch_gc() shared.log.debug(f'Segment loaded: model={selected_model} path={model_path} time={time.time()-t0:.2f}s') - loaded_model = selected_model + opts.model = selected_model busy = False return selected_model @@ -250,7 +250,7 @@ def run_rembg(input_image: Image, input_mask: np.ndarray): 'alpha_matting_foreground_threshold': 240, 'alpha_matting_background_threshold': 10, 'alpha_matting_erode_size': int(opts.mask_erode * 40), - 'session': rembg.new_session(loaded_model), + 'session': rembg.new_session(opts.model), } mask = rembg.remove(**args) mask = np.array(mask) @@ -339,7 +339,7 @@ def run_mask(input_image: gr.Image, input_mask: gr.Image = None, return_type: st else: opts.mask_padding = int(opts.mask_dilate * input_image.height / 4) + 1 - if loaded_model is None or not segment_enable: + if opts.model is None or not segment_enable: mask = input_mask elif generator is None: mask = run_rembg(input_image, input_mask) @@ -347,6 +347,7 @@ def run_mask(input_image: gr.Image, input_mask: gr.Image = None, return_type: st mask = run_segment(input_image, input_mask) mask = cv2.resize(mask, (input_image.width, input_image.height), interpolation=cv2.INTER_LINEAR) + debug(f'Mask opts: {opts}') debug(f'Segment mask: mask={mask.shape}') if opts.mask_erode > 0: try: @@ -483,6 +484,7 @@ def create_segment_ui(): selected_model.change(fn=init_model, inputs=[selected_model], outputs=[selected_model]) for control in controls: control.change(fn=update_opts, inputs=controls, outputs=[]) + return controls def bind_controls(image_controls: List[gr.Image], preview_image: gr.Image, output_image: gr.Image): diff --git a/modules/processing.py b/modules/processing.py index 48560fdb5..edb53b9f4 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -646,15 +646,14 @@ def create_infotext(p: StableDiffusionProcessing, all_prompts=None, all_seeds=No if 'img2img' in p.ops or 'inpaint' in p.ops: args["Init image size"] = f"{getattr(p, 'init_img_width', 0)}x{getattr(p, 'init_img_height', 0)}" args["Init image hash"] = getattr(p, 'init_img_hash', None) - args["Mask weight"] = getattr(p, "inpainting_mask_weight", shared.opts.inpainting_mask_weight) if p.is_using_inpainting_conditioning else None args['Resize scale'] = getattr(p, 'scale_by', None) - args["Mask blur"] = p.mask_blur if getattr(p, 'mask', None) is not None and getattr(p, 'mask_blur', 0) > 0 else None + args["Mask weight"] = getattr(p, "inpainting_mask_weight", shared.opts.inpainting_mask_weight) if p.is_using_inpainting_conditioning else None args["Denoising strength"] = getattr(p, 'denoising_strength', None) if args["Size"] is None: args["Size"] = args["Init image size"] # lookup by index if getattr(p, 'resize_mode', None) is not None: - args['Resize mode'] = shared.resize_modes[p.resize_mode] + args['Resize mode'] = shared.resize_modes[p.resize_mode] if shared.resize_modes[p.resize_mode] != 'None' else None if 'face' in p.ops: args["Face restoration"] = shared.opts.face_restoration_model if 'color' in p.ops: diff --git a/modules/ui_control.py b/modules/ui_control.py index ac49964e2..08802951b 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 -from PIL import Image from modules.control import unit from modules.control import processors # patrickvonplaten controlnet_aux from modules.control.units import controlnet # lllyasviel ControlNet @@ -10,59 +9,15 @@ from modules.control.units import xs # vislearn ControlNet-XS 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 modules import errors, shared, progress, sd_samplers, ui_components, ui_symbols, ui_common, ui_sections, generation_parameters_copypaste, call_queue, scripts, masking, ipadapter # 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, ipadapter, images # pylint: disable=ungrouped-imports +from modules import ui_control_helpers as helpers gr_height = None max_units = shared.opts.control_max_units units: list[unit.Unit] = [] # main state variable -input_source = None -input_init = None -input_mask = None debug = shared.log.trace if os.environ.get('SD_CONTROL_DEBUG', None) is not None else lambda *args, **kwargs: None debug('Trace: CONTROL') -busy = False # used to synchronize select_input and generate_click - - -def initialize(): - from modules import devices - shared.log.debug(f'UI initialize: control models={shared.opts.control_dir}') - controlnet.cache_dir = os.path.join(shared.opts.control_dir, 'controlnet') - xs.cache_dir = os.path.join(shared.opts.control_dir, 'xs') - 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) - os.makedirs(controlnet.cache_dir, exist_ok=True) - os.makedirs(xs.cache_dir, exist_ok=True) - 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) - - -def interrogate_clip(): - prompt = None - try: - prompt = shared.interrogator.interrogate(input_source[0]) - except Exception: - pass - return gr.update() if prompt is None else prompt - - -def interrogate_booru(): - prompt = None - try: - from modules import deepbooru - prompt = deepbooru.model.tag(input_source[0]) - except Exception: - pass - return gr.update() if prompt is None else prompt def return_controls(res): @@ -85,7 +40,7 @@ def return_controls(res): def generate_click(job_id: str, active_tab: str, *args): - while busy: + while helpers.busy: time.sleep(0.01) from modules.control.run import control_run debug(f'Control: tab="{active_tab}" job={job_id} args={args}') @@ -96,7 +51,7 @@ def generate_click(job_id: str, active_tab: str, *args): shared.mem_mon.reset() progress.start_task(job_id) try: - for results in control_run(units, input_source, input_init, input_mask, active_tab, True, *args): + for results in control_run(units, helpers.input_source, helpers.input_init, helpers.input_mask, active_tab, True, *args): progress.record_results(job_id, results) yield return_controls(results) except Exception as e: @@ -107,178 +62,8 @@ def generate_click(job_id: str, active_tab: str, *args): shared.state.end() -def display_units(num_units): - return (num_units * [gr.update(visible=True)]) + ((max_units - num_units) * [gr.update(visible=False)]) - - -def get_video(filepath: str): - try: - import cv2 - from modules.control.util import decode_fourcc - video = cv2.VideoCapture(filepath) - if not video.isOpened(): - msg = f'Control: video open failed: path="{filepath}"' - shared.log.error(msg) - return msg - frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT)) - fps = video.get(cv2.CAP_PROP_FPS) - duration = float(frames) / fps - w, h = int(video.get(cv2.CAP_PROP_FRAME_WIDTH)), int(video.get(cv2.CAP_PROP_FRAME_HEIGHT)) - codec = decode_fourcc(video.get(cv2.CAP_PROP_FOURCC)) - video.release() - shared.log.debug(f'Control: input video: path={filepath} frames={frames} fps={fps} size={w}x{h} codec={codec}') - msg = f'Control input | Video | Size {w}x{h} | Frames {frames} | FPS {fps:.2f} | Duration {duration:.2f} | Codec {codec}' - return msg - except Exception as e: - msg = f'Control: video open failed: path={filepath} {e}' - shared.log.error(msg) - return msg - - -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': - selected_input = input_image - elif input_mode == 'Outpaint': - selected_input = input_resize - elif input_mode == 'Inpaint': - selected_input = input_inpaint - elif input_mode == 'Video': - selected_input = input_video - elif input_mode == 'Batch': - selected_input = input_batch - elif input_mode == 'Folder': - selected_input = input_folder - else: - selected_input = None - if selected_input is None: - input_source = None - busy = False - debug('Control input: none') - 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) - input_mask = None - status = 'Control input | Unknown' - res = [gr.Tabs.update(selected='out-gallery'), status] - # control inputs - if isinstance(selected_input, Image.Image): # image via upload -> image - if input_mode == 'Outpaint': - input_mask = masking.run_mask(input_image=selected_input, input_mask=None, return_type='Grayscale') - input_source = [selected_input] - input_type = 'PIL.Image' - 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 = 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' - 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, gr.components.image.Image): # not likely - input_source = [selected_input.value] - input_type = 'gr.Image' - res = [gr.Tabs.update(selected='out-gallery'), status] - elif isinstance(selected_input, str): # video via upload > tmp filepath to video - input_source = selected_input - input_type = 'gr.Video' - status = get_video(input_source) - res = [gr.Tabs.update(selected='out-video'), status] - elif isinstance(selected_input, list): # batch or folder via upload -> list of tmp filepaths - if hasattr(selected_input[0], 'name'): - input_type = 'tempfiles' - input_source = [f.name for f in selected_input] # tempfile - else: - input_type = 'files' - input_source = selected_input - status = f'Control input | Images | Files {len(input_source)}' - res = [gr.Tabs.update(selected='out-gallery'), status] - else: # unknown - input_source = None - shared.log.debug(f'Control input: type={input_type} input={input_source}') - # init inputs: optional - if init_type == 0: # Control only - input_init = None - elif init_type == 1: # Init image same as control assigned during runtime - input_init = None - elif init_type == 2: # Separate init image - if isinstance(selected_init, Image.Image): # image via upload -> image - if input_mode == 'Outpaint': - 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' - 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 = 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' - status = f'Control input | Image | Size {selected_init.width}x{selected_init.height} | Mode {selected_input.mode}' - res = [gr.Tabs.update(selected='out-gallery'), status] - elif isinstance(selected_init, gr.components.image.Image): # not likely - input_init = [selected_init.value] - input_type = 'gr.Image' - res = [gr.Tabs.update(selected='out-gallery'), status] - elif isinstance(selected_init, str): # video via upload > tmp filepath to video - input_init = selected_init - input_type = 'gr.Video' - status = get_video(input_init) - res = [gr.Tabs.update(selected='out-video'), status] - elif isinstance(selected_init, list): # batch or folder via upload -> list of tmp filepaths - if hasattr(selected_init[0], 'name'): - input_type = 'tempfiles' - input_init = [f.name for f in selected_init] # tempfile - else: - input_type = 'files' - input_init = selected_init - status = f'Control input | Images | Files {len(input_init)}' - res = [gr.Tabs.update(selected='out-gallery'), status] - else: # unknown - input_init = None - debug(f'Control select input: source={input_source} init={input_init} mode={input_mode}') - busy = False - return res - - -def video_type_change(video_type): - return [ - gr.update(visible=video_type != 'None'), - gr.update(visible=video_type == 'GIF' or video_type == 'PNG'), - gr.update(visible=video_type == 'MP4'), - gr.update(visible=video_type == 'MP4'), - ] - - -def copy_input(mode_from, mode_to, input_image, input_resize, input_inpaint): - debug(f'Control transfter input: from={mode_from} to={mode_to} image={input_image} resize={input_resize} inpaint={input_inpaint}') - def getimg(ctrl): - if ctrl is None: - return None - return ctrl.get('image', None) if isinstance(ctrl, dict) else ctrl - - if mode_from == mode_to: - return [gr.update(), gr.update(), gr.update()] - elif mode_to == 'Select': - return [getimg(input_resize) if mode_from == 'Outpaint' else getimg(input_inpaint), None, None] - elif mode_to == 'Inpaint': - return [None, None, getimg(input_image) if mode_from == 'Select' else getimg(input_resize)] - elif mode_to == 'Outpaint': - return [None, getimg(input_image) if mode_from == 'Select' else getimg(input_inpaint), None] - else: - shared.log.error(f'Control transfer unknown input: from={mode_from} to={mode_to}') - return [gr.update(), gr.update(), gr.update()] - - -def transfer_input(dst): - return [gr.update(visible=dst=='Select'), gr.update(visible=dst=='Outpaint'), gr.update(visible=dst=='Inpaint'), gr.update(interactive=dst!='Select'), gr.update(interactive=dst!='Inpaint'), gr.update(interactive=dst!='Outpaint')] - - def create_ui(_blocks: gr.Blocks=None): - initialize() + helpers.initialize() if shared.backend == shared.Backend.ORIGINAL: with gr.Blocks(analytics_enabled = False) as control_ui: @@ -287,6 +72,9 @@ def create_ui(_blocks: gr.Blocks=None): with gr.Blocks(analytics_enabled = False) as control_ui: prompt, styles, negative, btn_generate, btn_paste, btn_extra, prompt_counter, btn_prompt_counter, negative_counter, btn_negative_counter = ui_sections.create_toprow(is_img2img=False, id_part='control') + txt_prompt_img = gr.File(label="", elem_id="control_prompt_image", file_count="single", type="binary", visible=False) + txt_prompt_img.change(fn=images.image_data, inputs=[txt_prompt_img], outputs=[prompt, txt_prompt_img]) + with gr.Group(elem_id="control_interface", equal_height=False): with gr.Row(elem_id='control_settings'): @@ -313,7 +101,7 @@ def create_ui(_blocks: gr.Blocks=None): 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() + mask_controls = 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') @@ -327,7 +115,7 @@ def create_ui(_blocks: gr.Blocks=None): video_loop = gr.Checkbox(label='Loop', value=True, visible=False) video_pad = gr.Slider(label='Pad frames', minimum=0, maximum=24, step=1, value=1, visible=False) video_interpolate = gr.Slider(label='Interpolate frames', minimum=0, maximum=24, step=1, value=0, visible=False) - video_type.change(fn=video_type_change, inputs=[video_type], outputs=[video_duration, video_loop, video_pad, video_interpolate]) + video_type.change(fn=helpers.video_type_change, inputs=[video_type], outputs=[video_duration, video_loop, video_pad, video_interpolate]) with gr.Accordion(open=False, label="Extensions", elem_id="control_extensions", elem_classes=["small-accordion"]): input_script_args = scripts.scripts_current.setup_ui(parent='control', accordion=False) @@ -433,7 +221,7 @@ def create_ui(_blocks: gr.Blocks=None): ) if i == 0: units[-1].enabled = True # enable first unit in group - num_controlnet_units.change(fn=display_units, inputs=[num_controlnet_units], outputs=controlnet_ui_units) + num_controlnet_units.change(fn=helpers.display_units, inputs=[num_controlnet_units], outputs=controlnet_ui_units) with gr.Tab('IP Adapter') as _tab_ipadapter: with gr.Row(): @@ -484,7 +272,7 @@ def create_ui(_blocks: gr.Blocks=None): ) if i == 0: units[-1].enabled = True # enable first unit in group - num_adapter_units.change(fn=display_units, inputs=[num_adapter_units], outputs=adapter_ui_units) + num_adapter_units.change(fn=helpers.display_units, inputs=[num_adapter_units], outputs=adapter_ui_units) with gr.Tab('XS') as _tab_controlnetxs: gr.HTML('ControlNet XS') @@ -530,7 +318,7 @@ def create_ui(_blocks: gr.Blocks=None): ) if i == 0: units[-1].enabled = True # enable first unit in group - num_controlnet_units.change(fn=display_units, inputs=[num_controlnet_units], outputs=controlnetxs_ui_units) + num_controlnet_units.change(fn=helpers.display_units, inputs=[num_controlnet_units], outputs=controlnetxs_ui_units) with gr.Tab('Lite') as _tab_lite: gr.HTML('Control LLLite') @@ -571,7 +359,7 @@ def create_ui(_blocks: gr.Blocks=None): ) if i == 0: units[-1].enabled = True # enable first unit in group - num_lite_units.change(fn=display_units, inputs=[num_lite_units], outputs=lite_ui_units) + num_lite_units.change(fn=helpers.display_units, inputs=[num_lite_units], outputs=lite_ui_units) with gr.Tab('Reference') as _tab_reference: gr.HTML('ControlNet reference-only control') @@ -662,20 +450,20 @@ def create_ui(_blocks: gr.Blocks=None): setting.change(fn=processors.update_settings, inputs=settings, outputs=[]) for btn in input_buttons: - btn.click(fn=copy_input, inputs=[input_mode, btn, input_image, input_resize, input_inpaint], outputs=[input_image, input_resize, input_inpaint], _js='controlInputMode') - btn.click(fn=transfer_input, inputs=[btn], outputs=[input_image, input_resize, input_inpaint] + input_buttons) + btn.click(fn=helpers.copy_input, inputs=[input_mode, btn, input_image, input_resize, input_inpaint], outputs=[input_image, input_resize, input_inpaint], _js='controlInputMode') + btn.click(fn=helpers.transfer_input, inputs=[btn], outputs=[input_image, input_resize, input_inpaint] + input_buttons) show_preview.change(fn=lambda x: gr.update(visible=x), inputs=[show_preview], outputs=[column_preview]) input_type.change(fn=lambda x: gr.update(visible=x == 2), inputs=[input_type], outputs=[column_init]) btn_prompt_counter.click(fn=call_queue.wrap_queued_call(ui_common.update_token_counter), inputs=[prompt, steps], outputs=[prompt_counter]) btn_negative_counter.click(fn=call_queue.wrap_queued_call(ui_common.update_token_counter), inputs=[negative, steps], outputs=[negative_counter]) - btn_interrogate_clip.click(fn=interrogate_clip, inputs=[], outputs=[prompt]) - btn_interrogate_booru.click(fn=interrogate_booru, inputs=[], outputs=[prompt]) + btn_interrogate_clip.click(fn=helpers.interrogate_clip, inputs=[], outputs=[prompt]) + btn_interrogate_booru.click(fn=helpers.interrogate_booru, inputs=[], outputs=[prompt]) 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, + fn=helpers.select_input, _js="controlInputMode", inputs=select_fields, outputs=select_output, @@ -723,7 +511,46 @@ def create_ui(_blocks: gr.Blocks=None): prompt.submit(**control_dict) btn_generate.click(**control_dict) - paste_fields = [] # TODO paste fields + paste_fields = [ + # prompt + (prompt, "Prompt"), + (negative, "Negative prompt"), + # input + (denoising_strength, "Denoising strength"), + # resize # TODO resize params + (width_before, "Size-1"), + (height_before, "Size-2"), + (resize_mode_before, "Resize mode"), + (scale_by_before, "Resize scale"), + # sampler + (sampler_index, "Sampler"), + (steps, "Steps"), + # batch + (batch_count, "Batch-1"), + (batch_size, "Batch-2"), + # seed + (seed, "Seed"), + # mask + (mask_controls[1], "Mask only"), + (mask_controls[2], "Mask invert"), + (mask_controls[3], "Mask blur"), + (mask_controls[4], "Mask erode"), + (mask_controls[5], "Mask dilate"), + (mask_controls[6], "Mask auto"), + # advanced + (cfg_scale, "CFG scale"), + (clip_skip, "Clip skip"), + (image_cfg_scale, "Image CFG scale"), + (diffusers_guidance_rescale, "CFG rescale"), + (full_quality, "Full quality"), + (restore_faces, "Face restoration"), + (tiling, "Tiling"), + # second pass # TODO second pass params + # hidden + (seed_resize_from_w, "Seed resize from-1"), + (seed_resize_from_h, "Seed resize from-2"), + *scripts.scripts_control.infotext_fields + ] 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) diff --git a/modules/ui_control_helpers.py b/modules/ui_control_helpers.py new file mode 100644 index 000000000..ded840b39 --- /dev/null +++ b/modules/ui_control_helpers.py @@ -0,0 +1,233 @@ +import os +import gradio as gr +from PIL import Image +from modules import shared, scripts, masking # pylint: disable=ungrouped-imports + + +gr_height = None +max_units = shared.opts.control_max_units +debug = shared.log.trace if os.environ.get('SD_CONTROL_DEBUG', None) is not None else lambda *args, **kwargs: None +debug('Trace: CONTROL') + +# state variables +busy = False # used to synchronize select_input and generate_click +input_source = None +input_init = None +input_mask = None + + +def initialize(): + from modules import devices + from modules.control import unit + from modules.control import processors # patrickvonplaten controlnet_aux + from modules.control.units import controlnet # lllyasviel ControlNet + from modules.control.units import xs # vislearn ControlNet-XS + from modules.control.units import lite # vislearn ControlNet-XS + from modules.control.units import t2iadapter # TencentARC T2I-Adapter + shared.log.debug(f'UI initialize: control models={shared.opts.control_dir}') + controlnet.cache_dir = os.path.join(shared.opts.control_dir, 'controlnet') + xs.cache_dir = os.path.join(shared.opts.control_dir, 'xs') + 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) + os.makedirs(controlnet.cache_dir, exist_ok=True) + os.makedirs(xs.cache_dir, exist_ok=True) + 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) + + +def interrogate_clip(): + prompt = None + try: + prompt = shared.interrogator.interrogate(input_source[0]) + except Exception: + pass + return gr.update() if prompt is None else prompt + + +def interrogate_booru(): + prompt = None + try: + from modules import deepbooru + prompt = deepbooru.model.tag(input_source[0]) + except Exception: + pass + return gr.update() if prompt is None else prompt + + +def display_units(num_units): + return (num_units * [gr.update(visible=True)]) + ((max_units - num_units) * [gr.update(visible=False)]) + + +def get_video(filepath: str): + try: + import cv2 + from modules.control.util import decode_fourcc + video = cv2.VideoCapture(filepath) + if not video.isOpened(): + msg = f'Control: video open failed: path="{filepath}"' + shared.log.error(msg) + return msg + frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT)) + fps = video.get(cv2.CAP_PROP_FPS) + duration = float(frames) / fps + w, h = int(video.get(cv2.CAP_PROP_FRAME_WIDTH)), int(video.get(cv2.CAP_PROP_FRAME_HEIGHT)) + codec = decode_fourcc(video.get(cv2.CAP_PROP_FOURCC)) + video.release() + shared.log.debug(f'Control: input video: path={filepath} frames={frames} fps={fps} size={w}x{h} codec={codec}') + msg = f'Control input | Video | Size {w}x{h} | Frames {frames} | FPS {fps:.2f} | Duration {duration:.2f} | Codec {codec}' + return msg + except Exception as e: + msg = f'Control: video open failed: path={filepath} {e}' + shared.log.error(msg) + return msg + + +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': + selected_input = input_image + elif input_mode == 'Outpaint': + selected_input = input_resize + elif input_mode == 'Inpaint': + selected_input = input_inpaint + elif input_mode == 'Video': + selected_input = input_video + elif input_mode == 'Batch': + selected_input = input_batch + elif input_mode == 'Folder': + selected_input = input_folder + else: + selected_input = None + if selected_input is None: + input_source = None + busy = False + debug('Control input: none') + 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) + input_mask = None + status = 'Control input | Unknown' + res = [gr.Tabs.update(selected='out-gallery'), status] + # control inputs + if isinstance(selected_input, Image.Image): # image via upload -> image + if input_mode == 'Outpaint': + input_mask = masking.run_mask(input_image=selected_input, input_mask=None, return_type='Grayscale') + input_source = [selected_input] + input_type = 'PIL.Image' + 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 = 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' + 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, gr.components.image.Image): # not likely + input_source = [selected_input.value] + input_type = 'gr.Image' + res = [gr.Tabs.update(selected='out-gallery'), status] + elif isinstance(selected_input, str): # video via upload > tmp filepath to video + input_source = selected_input + input_type = 'gr.Video' + status = get_video(input_source) + res = [gr.Tabs.update(selected='out-video'), status] + elif isinstance(selected_input, list): # batch or folder via upload -> list of tmp filepaths + if hasattr(selected_input[0], 'name'): + input_type = 'tempfiles' + input_source = [f.name for f in selected_input] # tempfile + else: + input_type = 'files' + input_source = selected_input + status = f'Control input | Images | Files {len(input_source)}' + res = [gr.Tabs.update(selected='out-gallery'), status] + else: # unknown + input_source = None + shared.log.debug(f'Control input: type={input_type} input={input_source}') + # init inputs: optional + if init_type == 0: # Control only + input_init = None + elif init_type == 1: # Init image same as control assigned during runtime + input_init = None + elif init_type == 2: # Separate init image + if isinstance(selected_init, Image.Image): # image via upload -> image + if input_mode == 'Outpaint': + 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' + 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 = 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' + status = f'Control input | Image | Size {selected_init.width}x{selected_init.height} | Mode {selected_input.mode}' + res = [gr.Tabs.update(selected='out-gallery'), status] + elif isinstance(selected_init, gr.components.image.Image): # not likely + input_init = [selected_init.value] + input_type = 'gr.Image' + res = [gr.Tabs.update(selected='out-gallery'), status] + elif isinstance(selected_init, str): # video via upload > tmp filepath to video + input_init = selected_init + input_type = 'gr.Video' + status = get_video(input_init) + res = [gr.Tabs.update(selected='out-video'), status] + elif isinstance(selected_init, list): # batch or folder via upload -> list of tmp filepaths + if hasattr(selected_init[0], 'name'): + input_type = 'tempfiles' + input_init = [f.name for f in selected_init] # tempfile + else: + input_type = 'files' + input_init = selected_init + status = f'Control input | Images | Files {len(input_init)}' + res = [gr.Tabs.update(selected='out-gallery'), status] + else: # unknown + input_init = None + debug(f'Control select input: source={input_source} init={input_init} mode={input_mode}') + busy = False + return res + + +def video_type_change(video_type): + return [ + gr.update(visible=video_type != 'None'), + gr.update(visible=video_type == 'GIF' or video_type == 'PNG'), + gr.update(visible=video_type == 'MP4'), + gr.update(visible=video_type == 'MP4'), + ] + + +def copy_input(mode_from, mode_to, input_image, input_resize, input_inpaint): + debug(f'Control transfter input: from={mode_from} to={mode_to} image={input_image} resize={input_resize} inpaint={input_inpaint}') + def getimg(ctrl): + if ctrl is None: + return None + return ctrl.get('image', None) if isinstance(ctrl, dict) else ctrl + + if mode_from == mode_to: + return [gr.update(), gr.update(), gr.update()] + elif mode_to == 'Select': + return [getimg(input_resize) if mode_from == 'Outpaint' else getimg(input_inpaint), None, None] + elif mode_to == 'Inpaint': + return [None, None, getimg(input_image) if mode_from == 'Select' else getimg(input_resize)] + elif mode_to == 'Outpaint': + return [None, getimg(input_image) if mode_from == 'Select' else getimg(input_inpaint), None] + else: + shared.log.error(f'Control transfer unknown input: from={mode_from} to={mode_to}') + return [gr.update(), gr.update(), gr.update()] + + +def transfer_input(dst): + return [gr.update(visible=dst=='Select'), gr.update(visible=dst=='Outpaint'), gr.update(visible=dst=='Inpaint'), gr.update(interactive=dst!='Select'), gr.update(interactive=dst!='Inpaint'), gr.update(interactive=dst!='Outpaint')]