from __future__ import annotations import platform import sys import traceback from collections.abc import Mapping from copy import copy, deepcopy from functools import partial from pathlib import Path from textwrap import dedent from typing import Any import gradio as gr import torch import modules # noqa: F401 from adetailer import ( AD_ENABLE, ALL_ARGS, ADetailerArgs, __version__, enable_check, get_models, mediapipe_predict, ultralytics_predict, ) from adetailer.common import mask_preprocess from controlnet_ext import ControlNetExt, controlnet_exists, get_cn_inpaint_models from modules import images, safe, script_callbacks, scripts, shared from modules.paths import data_path, models_path from modules.processing import ( StableDiffusionProcessingImg2Img, create_infotext, process_images, ) from modules.shared import cmd_opts, opts try: from rich import print from rich.traceback import install install(show_locals=True) except Exception: pass AFTER_DETAILER = "After Detailer" no_huggingface = getattr(cmd_opts, "ad_no_huggingface", False) adetailer_dir = Path(models_path, "adetailer") model_mapping = get_models(adetailer_dir, huggingface=not no_huggingface) txt2img_submit_button = img2img_submit_button = None print( f"[-] ADetailer initialized. version: {__version__}, num models: {len(model_mapping)}" ) class Widgets: def tolist(self): return [getattr(self, attr) for attr in ALL_ARGS.attrs] class ChangeTorchLoad: def __enter__(self): self.orig = torch.load torch.load = safe.unsafe_torch_load def __exit__(self, *args, **kwargs): torch.load = self.orig def gr_interactive(value: bool = True): return gr.update(interactive=value) def ordinal(n: int) -> str: d = {1: "st", 2: "nd", 3: "rd"} return str(n) + ("th" if 11 <= n % 100 <= 13 else d.get(n % 10, "th")) def suffix(n: int, c: str = " ") -> str: return "" if n == 0 else c + ordinal(n + 1) def on_widget_change(state: dict, value: Any, *, attr: str): state[attr] = value return state def on_generate_click(state: dict, *values: Any): for attr, value in zip(ALL_ARGS.attrs, values): state[attr] = value return state class AfterDetailerScript(scripts.Script): def __init__(self): super().__init__() self.controlnet_ext = None self.ultralytics_device = self.get_ultralytics_device() def title(self): return AFTER_DETAILER def show(self, is_img2img): return scripts.AlwaysVisible def ui(self, is_img2img): num_models = opts.data.get("ad_max_models", 2) widgets = [] states = [] self.infotext_fields = [] with gr.Accordion(AFTER_DETAILER, open=False, elem_id="AD_main_acc"): with gr.Row(): ad_enable = gr.Checkbox( label="Enable ADetailer", value=False, visible=True, ) self.infotext_fields.append((ad_enable, AD_ENABLE.name)) with gr.Group(), gr.Tabs(): for n in range(num_models): with gr.Tab(ordinal(n + 1)): w, state, infofields = self.one_ui_group(n, is_img2img) widgets.append(w) states.append(state) self.infotext_fields.extend(infofields) # return: [bool, dict, dict, ...] return [ad_enable] + states def one_ui_group(self, n: int, is_img2img: bool): model_list = list(model_mapping.keys()) w = Widgets() state = gr.State({}) with gr.Row(): model_choices = model_list if n == 0 else ["None"] + model_list w.ad_model = gr.Dropdown( label="ADetailer model" + suffix(n), choices=model_choices, value=model_choices[0], visible=True, type="value", ) with gr.Group(): with gr.Row(elem_id="AD_toprow_prompt" + suffix(n, "_")): w.ad_prompt = gr.Textbox( label="ad_prompt" + suffix(n), show_label=False, lines=3, placeholder="ADetailer prompt" + suffix(n), elem_id="AD_prompt" + suffix(n, "_"), ) with gr.Row(elem_id="AD_toprow_negative_prompt" + suffix(n, "_")): w.ad_negative_prompt = gr.Textbox( label="ad_negative_prompt" + suffix(n), show_label=False, lines=2, placeholder="ADetailer negative prompt" + suffix(n), elem_id="AD_negative_prompt" + suffix(n, "_"), ) with gr.Group(): with gr.Row(): w.ad_conf = gr.Slider( label="Detection model confidence threshold %" + suffix(n), minimum=0, maximum=100, step=1, value=30, visible=True, ) w.ad_dilate_erode = gr.Slider( label="Mask erosion (-) / dilation (+)" + suffix(n), minimum=-128, maximum=128, step=4, value=32, visible=True, ) with gr.Row(): w.ad_x_offset = gr.Slider( label="Mask x(→) offset" + suffix(n), minimum=-200, maximum=200, step=1, value=0, visible=True, ) w.ad_y_offset = gr.Slider( label="Mask y(↑) offset" + suffix(n), minimum=-200, maximum=200, step=1, value=0, visible=True, ) with gr.Row(): w.ad_mask_blur = gr.Slider( label="Inpaint mask blur" + suffix(n), minimum=0, maximum=64, step=1, value=4, visible=True, ) w.ad_denoising_strength = gr.Slider( label="Inpaint denoising strength" + suffix(n), minimum=0.0, maximum=1.0, step=0.01, value=0.4, visible=True, ) with gr.Group(): with gr.Row(): with gr.Column(variant="compact"): w.ad_inpaint_full_res = gr.Checkbox( label="Inpaint at full resolution " + suffix(n), value=True, visible=True, ) w.ad_inpaint_full_res_padding = gr.Slider( label="Inpaint at full resolution padding, pixels " + suffix(n), minimum=0, maximum=256, step=4, value=0, visible=True, ) w.ad_inpaint_full_res.change( gr_interactive, inputs=w.ad_inpaint_full_res, outputs=w.ad_inpaint_full_res_padding, queue=False, ) with gr.Column(variant="compact"): w.ad_use_inpaint_width_height = gr.Checkbox( label="Use separate width/height" + suffix(n), value=False, visible=True, ) w.ad_inpaint_width = gr.Slider( label="inpaint width" + suffix(n), minimum=64, maximum=2048, step=4, value=512, visible=True, ) w.ad_inpaint_height = gr.Slider( label="inpaint height" + suffix(n), minimum=64, maximum=2048, step=4, value=512, visible=True, ) w.ad_use_inpaint_width_height.change( lambda value: (gr_interactive(value), gr_interactive(value)), inputs=w.ad_use_inpaint_width_height, outputs=[w.ad_inpaint_width, w.ad_inpaint_height], queue=False, ) with gr.Row(): with gr.Column(variant="compact"): w.ad_use_steps = gr.Checkbox( label="Use separate steps" + suffix(n), value=False, visible=True, ) w.ad_steps = gr.Slider( label="ADetailer steps" + suffix(n), minimum=1, maximum=150, step=1, value=28, visible=True, ) w.ad_use_steps.change( gr_interactive, inputs=w.ad_use_steps, outputs=w.ad_steps, queue=False, ) with gr.Column(variant="compact"): w.ad_use_cfg_scale = gr.Checkbox( label="Use separate CFG scale" + suffix(n), value=False, visible=True, ) w.ad_cfg_scale = gr.Slider( label="ADetailer CFG scale" + suffix(n), minimum=0.0, maximum=30.0, step=0.5, value=7.0, visible=True, ) w.ad_use_cfg_scale.change( gr_interactive, inputs=w.ad_use_cfg_scale, outputs=w.ad_cfg_scale, queue=False, ) with gr.Group(), gr.Row(variant="panel"): cn_inpaint_models = ["None"] + get_cn_inpaint_models() w.ad_controlnet_model = gr.Dropdown( label="ControlNet model" + suffix(n), choices=cn_inpaint_models, value="None", visible=True, type="value", interactive=controlnet_exists, ) w.ad_controlnet_weight = gr.Slider( label="ControlNet weight" + suffix(n), minimum=0.0, maximum=1.0, step=0.05, value=1.0, visible=True, interactive=controlnet_exists, ) for attr in ALL_ARGS.attrs: widget = getattr(w, attr) on_change = partial(on_widget_change, attr=attr) widget.change( fn=on_change, inputs=[state, widget], outputs=[state], queue=False ) all_inputs = [state] + w.tolist() target_button = img2img_submit_button if is_img2img else txt2img_submit_button target_button.click( fn=on_generate_click, inputs=all_inputs, outputs=state, queue=False ) infotext_fields = [ (getattr(w, attr), name + suffix(n)) for attr, name in ALL_ARGS ] return w, state, infotext_fields def init_controlnet_ext(self) -> None: if self.controlnet_ext is not None: return self.controlnet_ext = ControlNetExt() if controlnet_exists: try: self.controlnet_ext.init_controlnet() except ImportError: error = traceback.format_exc() print( f"[-] ADetailer: ControlNetExt init failed:\n{error}", file=sys.stderr, ) def update_controlnet_args(self, p, args: ADetailerArgs) -> None: if self.controlnet_ext is None: self.init_controlnet_ext() if ( self.controlnet_ext is not None and self.controlnet_ext.cn_available and args.ad_controlnet_model != "None" ): self.controlnet_ext.update_scripts_args( p, args.ad_controlnet_model, args.ad_controlnet_weight ) def is_ad_enabled(self, *args_) -> bool: if len(args_) == 0 or (len(args_) == 1 and isinstance(args_[0], bool)): message = f""" [-] ADetailer: Not enough arguments passed to ADetailer. input: {args_!r} """ raise ValueError(dedent(message)) return enable_check(*args_) def get_args(self, *args_) -> list[ADetailerArgs]: """ `args_` is at least 1 in length by `is_ad_enabled` immediately above """ args = args_[1:] if isinstance(args_[0], bool) else args_ all_inputs = [] for n, arg_dict in enumerate(args, 1): try: inp = ADetailerArgs(**arg_dict) except ValueError as e: message = [ f"[-] ADetailer: ValidationError when validating {ordinal(n)} arguments: {e}\n" ] for attr in ALL_ARGS.attrs: arg = arg_dict.get(attr) dtype = type(arg) arg = "DEFAULT" if arg is None else repr(arg) message.append(f" {attr}: {arg} ({dtype})") raise ValueError("\n".join(message)) from e except TypeError as e: message = f"[-] ADetailer: {ordinal(n)} - Non-mapping arguments are sent: {arg_dict!r}\n{e}" raise TypeError(message) from e all_inputs.append(inp) return all_inputs def extra_params(self, arg_list: list[ADetailerArgs]) -> dict: params = {} for n, args in enumerate(arg_list): params.update(args.extra_params(suffix=suffix(n))) params["ADetailer version"] = __version__ return params @staticmethod def get_ultralytics_device() -> str: '`device = ""` means autodetect' device = "" if platform.system() == "Darwin": return device if any(getattr(cmd_opts, vram, False) for vram in ["lowvram", "medvram"]): device = "cpu" return device def get_prompt(self, p, args: ADetailerArgs) -> tuple[str, str]: i = p._idx if args.ad_prompt: prompt = args.ad_prompt elif not p.all_prompts: prompt = p.prompt elif i < len(p.all_prompts): prompt = p.all_prompts[i] else: j = i % len(p.all_prompts) prompt = p.all_prompts[j] if args.ad_negative_prompt: negative_prompt = args.ad_negative_prompt elif not p.all_negative_prompts: negative_prompt = p.negative_prompt elif i < len(p.all_negative_prompts): negative_prompt = p.all_negative_prompts[i] else: j = i % len(p.all_negative_prompts) negative_prompt = p.all_negative_prompts[j] return prompt, negative_prompt def get_seed(self, p) -> tuple[int, int]: i = p._idx if not p.all_seeds: seed = p.seed elif i < len(p.all_seeds): seed = p.all_seeds[i] else: j = i % len(p.all_seeds) seed = p.all_seeds[j] if not p.all_subseeds: subseed = p.subseed elif i < len(p.all_subseeds): subseed = p.all_subseeds[i] else: j = i % len(p.all_subseeds) subseed = p.all_subseeds[j] return seed, subseed def get_width_height(self, p, args: ADetailerArgs) -> tuple[int, int]: if args.ad_use_inpaint_width_height: width = args.ad_inpaint_width height = args.ad_inpaint_height else: width = p.width height = p.height return width, height def get_steps(self, p, args: ADetailerArgs) -> int: if args.ad_use_steps: return args.ad_steps return p.steps def get_cfg_scale(self, p, args: ADetailerArgs) -> float: if args.ad_use_cfg_scale: return args.ad_cfg_scale return p.cfg_scale def infotext(self, p) -> str: return create_infotext( p, p.all_prompts, p.all_seeds, p.all_subseeds, None, 0, 0 ) def write_params_txt(self, p) -> None: infotext = self.infotext(p) params_txt = Path(data_path, "params.txt") params_txt.write_text(infotext, encoding="utf-8") def script_filter(self, p, args: ADetailerArgs): script_runner = copy(p.scripts) ad_only_seleted_scripts = opts.data.get("ad_only_seleted_scripts", True) if not ad_only_seleted_scripts: return script_runner default = "dynamic_prompting,dynamic_thresholding,wildcards,wildcard_recursive" ad_script_names = opts.data.get("ad_script_names", default) script_names_set = { name for script_name in ad_script_names.split(",") for name in (script_name, script_name.strip()) } if args.ad_controlnet_model != "None": script_names_set.add("controlnet") filtered_alwayson = [] for script_object in script_runner.alwayson_scripts: filepath = script_object.filename filename = Path(filepath).stem if filename in script_names_set: filtered_alwayson.append(script_object) script_runner.alwayson_scripts = filtered_alwayson return script_runner def get_i2i_p(self, p, args: ADetailerArgs, image): prompt, negative_prompt = self.get_prompt(p, args) seed, subseed = self.get_seed(p) width, height = self.get_width_height(p, args) steps = self.get_steps(p, args) cfg_scale = self.get_cfg_scale(p, args) sampler_name = p.sampler_name if sampler_name in ["PLMS", "UniPC"]: sampler_name = "Euler" i2i = StableDiffusionProcessingImg2Img( init_images=[image], resize_mode=0, denoising_strength=args.ad_denoising_strength, mask=None, mask_blur=args.ad_mask_blur, inpainting_fill=1, inpaint_full_res=args.ad_inpaint_full_res, inpaint_full_res_padding=args.ad_inpaint_full_res_padding, inpainting_mask_invert=0, sd_model=p.sd_model, outpath_samples=p.outpath_samples, outpath_grids=p.outpath_grids, prompt=prompt, negative_prompt=negative_prompt, styles=p.styles, seed=seed, subseed=subseed, subseed_strength=p.subseed_strength, seed_resize_from_h=p.seed_resize_from_h, seed_resize_from_w=p.seed_resize_from_w, sampler_name=sampler_name, batch_size=1, n_iter=1, steps=steps, cfg_scale=cfg_scale, width=width, height=height, tiling=p.tiling, extra_generation_params=p.extra_generation_params, do_not_save_samples=True, do_not_save_grid=True, ) i2i.scripts = self.script_filter(p, args) i2i.script_args = deepcopy(p.script_args) i2i._disable_adetailer = True if args.ad_controlnet_model != "None": self.update_controlnet_args(i2i, args) return i2i def save_image(self, p, image, *, condition: str, suffix: str) -> None: i = p._idx seed, _ = self.get_seed(p) if opts.data.get(condition, False): images.save_image( image=image, path=p.outpath_samples, basename="", seed=seed, prompt=p.all_prompts[i] if i < len(p.all_prompts) else p.prompt, extension=opts.samples_format, info=self.infotext(p), p=p, suffix=suffix, ) def get_ad_model(self, name: str): if name not in model_mapping: msg = f"[-] ADetailer: Model {name!r} not found. Available models: {list(model_mapping.keys())}" raise ValueError(msg) return model_mapping[name] def process(self, p, *args_): if getattr(p, "_disable_adetailer", False): return if self.is_ad_enabled(*args_): arg_list = self.get_args(*args_) extra_params = self.extra_params(arg_list) p.extra_generation_params.update(extra_params) def _postprocess_image(self, p, pp, args: ADetailerArgs, *, n: int = 0) -> bool: """ Returns ------- bool `True` if image was processed, `False` otherwise. """ i = p._idx i2i = self.get_i2i_p(p, args, pp.image) seed, subseed = self.get_seed(p) is_mediapipe = args.ad_model.lower().startswith("mediapipe") kwargs = {} if is_mediapipe: predictor = mediapipe_predict ad_model = args.ad_model else: predictor = ultralytics_predict ad_model = self.get_ad_model(args.ad_model) kwargs["device"] = self.ultralytics_device with ChangeTorchLoad(): pred = predictor(ad_model, pp.image, args.ad_conf, **kwargs) masks = mask_preprocess( pred.masks, kernel=args.ad_dilate_erode, x_offset=args.ad_x_offset, y_offset=args.ad_y_offset, ) if not masks: print( f"[-] ADetailer: nothing detected on image {i + 1} with {ordinal(n + 1)} settings." ) return False self.save_image( p, pred.preview, condition="ad_save_previews", suffix="-ad-preview" + suffix(n, "-"), ) steps = len(masks) processed = None if is_mediapipe: print(f"mediapipe: {steps} detected.") p2 = copy(i2i) for j in range(steps): p2.image_mask = masks[j] processed = process_images(p2) p2 = copy(i2i) p2.init_images = [processed.images[0]] p2.seed = seed + j + 1 p2.subseed = subseed + j + 1 if processed is not None: pp.image = processed.images[0] return True return False def postprocess_image(self, p, pp, *args_): if getattr(p, "_disable_adetailer", False): return if not self.is_ad_enabled(*args_): return p._idx = getattr(p, "_idx", -1) + 1 init_image = copy(pp.image) arg_list = self.get_args(*args_) is_processed = False for n, args in enumerate(arg_list): if args.ad_model == "None": continue is_processed |= self._postprocess_image(p, pp, args, n=n) if is_processed: self.save_image( p, init_image, condition="ad_save_images_before", suffix="-ad-before" ) try: if p._idx == len(p.all_prompts) - 1: self.write_params_txt(p) except Exception: pass def on_after_component(component, **_kwargs): global txt2img_submit_button, img2img_submit_button if getattr(component, "elem_id", None) == "txt2img_generate": txt2img_submit_button = component return if getattr(component, "elem_id", None) == "img2img_generate": img2img_submit_button = component def on_ui_settings(): section = ("ADetailer", AFTER_DETAILER) shared.opts.add_option( "ad_max_models", shared.OptionInfo( default=2, label="Max models", component=gr.Slider, component_args={"minimum": 1, "maximum": 5, "step": 1}, section=section, ), ) shared.opts.add_option( "ad_save_previews", shared.OptionInfo(False, "Save mask previews", section=section), ) shared.opts.add_option( "ad_save_images_before", shared.OptionInfo(False, "Save images before ADetailer", section=section), ) shared.opts.add_option( "ad_only_seleted_scripts", shared.OptionInfo( True, "Apply only selected scripts to ADetailer", section=section ), ) textbox_args = { "placeholder": "comma-separated list of script names", "interactive": True, } shared.opts.add_option( "ad_script_names", shared.OptionInfo( default="dynamic_prompting,dynamic_thresholding,wildcards,wildcard_recursive", label="Script names to apply to ADetailer (separated by comma)", component=gr.Textbox, component_args=textbox_args, section=section, ), ) script_callbacks.on_ui_settings(on_ui_settings) script_callbacks.on_after_component(on_after_component)