diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index af459ce..421c213 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -8,7 +8,7 @@ repos: - id: mixed-line-ending - repo: https://github.com/astral-sh/ruff-pre-commit - rev: "v0.0.290" + rev: "v0.0.292" hooks: - id: ruff args: [--fix, --exit-non-zero-on-fix] diff --git a/CHANGELOG.md b/CHANGELOG.md index a2196f8..03c9400 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,12 @@ # Changelog +## 2023-10-07 + +- v23.10.0 +- 허깅페이스 모델을 다운로드 실패했을 때, 계속 다운로드를 시도하지 않음 +- img2img에서 img2img단계를 건너뛰는 기능 추가 +- live preview에서 감지 단계를 보여줌 (PR #352) + ## 2023-09-20 - v23.9.3 diff --git a/README.md b/README.md index d6fe5e4..3a00116 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ -# !After Detailer +# ADetailer -!After Detailer is a extension for stable diffusion webui, similar to Detection Detailer, except it uses ultralytics instead of the mmdet. +ADetailer is a extension for stable diffusion webui, similar to Detection Detailer, except it uses ultralytics instead of the mmdet. ## Install @@ -22,15 +22,17 @@ You **DON'T** need to download any model from huggingface. ## Options -| Model, Prompts | | | -| --------------------------------- | ------------------------------------- | ------------------------------------------------- | -| ADetailer model | Determine what to detect. | `None` = disable | -| ADetailer prompt, negative prompt | Prompts and negative prompts to apply | If left blank, it will use the same as the input. | +| Model, Prompts | | | +| --------------------------------- | --------------------------------------------------------------------------------- | ------------------------------------------------- | +| ADetailer model | Determine what to detect. | `None` = disable | +| ADetailer prompt, negative prompt | Prompts and negative prompts to apply | If left blank, it will use the same as the input. | +| Skip img2img | Skip img2img. In practice, this works by changing the step count of img2img to 1. | img2img only | -| Detection | | | -| ------------------------------------ | -------------------------------------------------------------------------------------------- | --- | -| Detection model confidence threshold | Only objects with a detection model confidence above this threshold are used for inpainting. | | -| Mask min/max ratio | Only use masks whose area is between those ratios for the area of the entire image. | | +| Detection | | | +| ------------------------------------ | -------------------------------------------------------------------------------------------- | ------------ | +| Detection model confidence threshold | Only objects with a detection model confidence above this threshold are used for inpainting. | | +| Mask min/max ratio | Only use masks whose area is between those ratios for the area of the entire image. | | +| Mask only the top k largest | Only use the k objects with the largest area of the bbox. | 0 to disable | If you want to exclude objects in the background, try setting the min ratio to around `0.01`. @@ -86,9 +88,10 @@ Put your [ultralytics](https://github.com/ultralytics/ultralytics) yolo model in It must be a bbox detection or segment model and use all label. -## Example +## How it works -![image](https://i.imgur.com/38RSxSO.png) -![image](https://i.imgur.com/2CYgjLx.png) +ADetailer works in three simple steps. -[![ko-fi](https://ko-fi.com/img/githubbutton_sm.svg)](https://ko-fi.com/F1F1L7V2N) +1. Create an image. +2. Detect object with a detection model and create a mask image. +3. Inpaint using the image from 1 and the mask from 2. diff --git a/adetailer/__init__.py b/adetailer/__init__.py index dae6181..ce38959 100644 --- a/adetailer/__init__.py +++ b/adetailer/__init__.py @@ -1,5 +1,5 @@ from .__version__ import __version__ -from .args import AD_ENABLE, ALL_ARGS, ADetailerArgs, EnableChecker +from .args import ALL_ARGS, ADetailerArgs from .common import PredictOutput, get_models from .mediapipe import mediapipe_predict from .ultralytics import ultralytics_predict @@ -8,11 +8,9 @@ AFTER_DETAILER = "ADetailer" __all__ = [ "__version__", - "AD_ENABLE", "ADetailerArgs", "AFTER_DETAILER", "ALL_ARGS", - "EnableChecker", "PredictOutput", "get_models", "mediapipe_predict", diff --git a/adetailer/__version__.py b/adetailer/__version__.py index 88729b7..32f93ce 100644 --- a/adetailer/__version__.py +++ b/adetailer/__version__.py @@ -1 +1 @@ -__version__ = "23.9.3" +__version__ = "23.10.0" diff --git a/adetailer/args.py b/adetailer/args.py index ea22b99..503eebb 100644 --- a/adetailer/args.py +++ b/adetailer/args.py @@ -2,7 +2,7 @@ from __future__ import annotations from collections import UserList from functools import cached_property, partial -from typing import Any, Literal, NamedTuple, Optional, Union +from typing import Any, Literal, NamedTuple, Optional import pydantic from pydantic import ( @@ -185,19 +185,7 @@ class ADetailerArgs(BaseModel, extra=Extra.forbid): return p -class EnableChecker(BaseModel): - enable: bool - arg_list: list - - def is_enabled(self) -> bool: - ad_model = ALL_ARGS[0].attr - if not self.enable: - return False - return any(arg.get(ad_model, "None") != "None" for arg in self.arg_list) - - _all_args = [ - ("ad_enable", "ADetailer enable"), ("ad_model", "ADetailer model"), ("ad_prompt", "ADetailer prompt"), ("ad_negative_prompt", "ADetailer negative prompt"), @@ -238,8 +226,7 @@ _all_args = [ ("ad_controlnet_guidance_end", "ADetailer ControlNet guidance end"), ] -AD_ENABLE = Arg(*_all_args[0]) -_args = [Arg(*args) for args in _all_args[1:]] +_args = [Arg(*args) for args in _all_args] ALL_ARGS = ArgsList(_args) BBOX_SORTBY = [ diff --git a/adetailer/common.py b/adetailer/common.py index a29bdd5..3bf81cb 100644 --- a/adetailer/common.py +++ b/adetailer/common.py @@ -10,6 +10,7 @@ from PIL import Image, ImageDraw from rich import print repo_id = "Bingsu/adetailer" +_download_failed = False @dataclass @@ -20,12 +21,18 @@ class PredictOutput: def hf_download(file: str): + global _download_failed + + if _download_failed: + return "INVALID" + try: path = hf_hub_download(repo_id, file) except Exception: msg = f"[-] ADetailer: Failed to load model {file!r} from huggingface" print(msg) path = "INVALID" + _download_failed = True return path diff --git a/adetailer/ui.py b/adetailer/ui.py index b05e3aa..8a218e1 100644 --- a/adetailer/ui.py +++ b/adetailer/ui.py @@ -3,12 +3,12 @@ from __future__ import annotations from dataclasses import dataclass from functools import partial from types import SimpleNamespace -from typing import Any, Callable +from typing import Any import gradio as gr from adetailer import AFTER_DETAILER, __version__ -from adetailer.args import AD_ENABLE, ALL_ARGS, MASK_MERGE_INVERT +from adetailer.args import ALL_ARGS, MASK_MERGE_INVERT from controlnet_ext import controlnet_exists, get_cn_models cn_module_choices = [ @@ -91,13 +91,22 @@ def adui( elem_id=eid("ad_enable"), ) + with gr.Column(scale=6): + ad_skip_img2img = gr.Checkbox( + label="Skip img2img", + value=False, + visible=is_img2img, + elem_id=eid("ad_skip_img2img"), + ) + with gr.Column(scale=1, min_width=180): gr.Markdown( f"v{__version__}", elem_id=eid("ad_version"), ) - infotext_fields.append((ad_enable, AD_ENABLE.name)) + infotext_fields.append((ad_enable, "ADetailer enable")) + infotext_fields.append((ad_skip_img2img, "ADetailer skip img2img")) with gr.Group(), gr.Tabs(): for n in range(num_models): @@ -112,7 +121,7 @@ def adui( infotext_fields.extend(infofields) # components: [bool, dict, dict, ...] - components = [ad_enable, *states] + components = [ad_enable, ad_skip_img2img, *states] return components, infotext_fields @@ -421,7 +430,7 @@ def inpainting(w: Widgets, n: int, is_img2img: bool, webui_info: WebuiInfo): with gr.Row(): with gr.Column(variant="compact"): w.ad_use_checkpoint = gr.Checkbox( - label="Use separate checkpoint (experimental)" + suffix(n), + label="Use separate checkpoint" + suffix(n), value=False, visible=True, elem_id=eid("ad_use_checkpoint"), @@ -439,7 +448,7 @@ def inpainting(w: Widgets, n: int, is_img2img: bool, webui_info: WebuiInfo): with gr.Column(variant="compact"): w.ad_use_vae = gr.Checkbox( - label="Use separate VAE (experimental)" + suffix(n), + label="Use separate VAE" + suffix(n), value=False, visible=True, elem_id=eid("ad_use_vae"), diff --git a/install.py b/install.py index ceab808..12cc6a6 100644 --- a/install.py +++ b/install.py @@ -44,7 +44,7 @@ def run_pip(*args): def install(): deps = [ # requirements - ("ultralytics", "8.0.181", None), + ("ultralytics", "8.0.194", None), ("mediapipe", "0.10.5", None), ("rich", "13.0.0", None), # mediapipe diff --git a/scripts/!adetailer.py b/scripts/!adetailer.py index 544e7b2..c507068 100644 --- a/scripts/!adetailer.py +++ b/scripts/!adetailer.py @@ -5,7 +5,7 @@ import platform import re import sys import traceback -from contextlib import contextmanager +from contextlib import contextmanager, suppress from copy import copy, deepcopy from functools import partial from pathlib import Path @@ -24,7 +24,7 @@ from adetailer import ( mediapipe_predict, ultralytics_predict, ) -from adetailer.args import ALL_ARGS, BBOX_SORTBY, ADetailerArgs, EnableChecker +from adetailer.args import ALL_ARGS, BBOX_SORTBY, ADetailerArgs from adetailer.common import PredictOutput from adetailer.mask import ( filter_by_ratio, @@ -176,7 +176,7 @@ class AfterDetailerScript(scripts.Script): def is_ad_enabled(self, *args_) -> bool: arg_list = [arg for arg in args_ if isinstance(arg, dict)] - if not args_ or not arg_list or not isinstance(args_[0], (bool, dict)): + if not args_ or not arg_list: message = f""" [-] ADetailer: Invalid arguments passed to ADetailer. input: {args_!r} @@ -184,9 +184,26 @@ class AfterDetailerScript(scripts.Script): """ print(dedent(message), file=sys.stderr) return False - enable = args_[0] if isinstance(args_[0], bool) else True - checker = EnableChecker(enable=enable, arg_list=arg_list) - return checker.is_enabled() + + ad_enabled = args_[0] if isinstance(args_[0], bool) else True + not_none = any(arg.get("ad_model", "None") != "None" for arg in arg_list) + return ad_enabled and not_none + + def check_skip_img2img(self, p, *args_) -> None: + if ( + hasattr(p, "_ad_skip_img2img") + or not hasattr(p, "init_images") + or not p.init_images + ): + return + + if len(args_) >= 2 and isinstance(args_[1], bool): + p._ad_skip_img2img = args_[1] + if args_[1]: + p._ad_orig_steps = p.steps + p.steps = 1 + else: + p._ad_skip_img2img = False def get_args(self, p, *args_) -> list[ADetailerArgs]: """ @@ -198,8 +215,8 @@ class AfterDetailerScript(scripts.Script): message = f"[-] ADetailer: Invalid arguments passed to ADetailer: {args_!r}" raise ValueError(message) - if hasattr(p, "adetailer_xyz"): - args[0] = {**args[0], **p.adetailer_xyz} + if hasattr(p, "_ad_xyz"): + args[0] = {**args[0], **p._ad_xyz} all_inputs = [] @@ -306,7 +323,11 @@ class AfterDetailerScript(scripts.Script): return width, height def get_steps(self, p, args: ADetailerArgs) -> int: - return args.ad_steps if args.ad_use_steps else p.steps + if args.ad_use_steps: + return args.ad_steps + if hasattr(p, "_ad_orig_steps"): + return p._ad_orig_steps + return p.steps def get_cfg_scale(self, p, args: ADetailerArgs) -> float: return args.ad_cfg_scale if args.ad_use_cfg_scale else p.cfg_scale @@ -345,9 +366,23 @@ class AfterDetailerScript(scripts.Script): ) def write_params_txt(self, p) -> None: + i = p._ad_idx + lenp = len(p.all_prompts) + if i % lenp != lenp - 1: + return + + prev = None + if hasattr(p, "_ad_orig_steps"): + prev = p.steps + p.steps = p._ad_orig_steps + infotext = self.infotext(p) params_txt = Path(data_path, "params.txt") - params_txt.write_text(infotext, encoding="utf-8") + with suppress(Exception): + params_txt.write_text(infotext, encoding="utf-8") + + if hasattr(p, "_ad_orig_steps"): + p.steps = prev def script_filter(self, p, args: ADetailerArgs): script_runner = copy(p.scripts) @@ -525,16 +560,26 @@ class AfterDetailerScript(scripts.Script): @staticmethod def need_call_process(p) -> bool: + if p.scripts is None: + return False i = p._ad_idx bs = p.batch_size return i % bs == bs - 1 @staticmethod def need_call_postprocess(p) -> bool: + if p.scripts is None: + return False i = p._ad_idx bs = p.batch_size return i % bs == 0 + @staticmethod + def get_i2i_init_image(p, pp): + if getattr(p, "_ad_skip_img2img", False): + return p.init_images[0] + return pp.image + @rich_traceback def process(self, p, *args_): if getattr(p, "_ad_disabled", False): @@ -542,8 +587,11 @@ class AfterDetailerScript(scripts.Script): if self.is_ad_enabled(*args_): arg_list = self.get_args(p, *args_) + self.check_skip_img2img(p, *args_) extra_params = self.extra_params(arg_list) p.extra_generation_params.update(extra_params) + else: + p._ad_disabled = True def _postprocess_image_inner( self, p, pp, args: ADetailerArgs, *, n: int = 0 @@ -560,6 +608,7 @@ class AfterDetailerScript(scripts.Script): i = p._ad_idx + pp.image = self.get_i2i_init_image(p, pp) i2i = self.get_i2i_p(p, args, pp.image) seed, subseed = self.get_seed(p) ad_prompts, ad_negatives = self.get_prompt(p, args) @@ -579,6 +628,7 @@ class AfterDetailerScript(scripts.Script): pred = predictor(ad_model, pp.image, args.ad_confidence, **kwargs) masks = self.pred_preprocessing(pred, args) + shared.state.assign_current_image(pred.preview) if not masks: print( @@ -633,17 +683,14 @@ class AfterDetailerScript(scripts.Script): @rich_traceback def postprocess_image(self, p, pp, *args_): - if getattr(p, "_ad_disabled", False): - return - - if not self.is_ad_enabled(*args_): + if getattr(p, "_ad_disabled", False) or not self.is_ad_enabled(*args_): return p._ad_idx = getattr(p, "_ad_idx", -1) + 1 init_image = copy(pp.image) arg_list = self.get_args(p, *args_) - if p.scripts is not None and self.need_call_postprocess(p): + if self.need_call_postprocess(p): dummy = Processed(p, [], p.seed, "") with preseve_prompts(p): p.scripts.postprocess(copy(p), dummy) @@ -655,22 +702,16 @@ class AfterDetailerScript(scripts.Script): continue is_processed |= self._postprocess_image_inner(p, pp, args, n=n) - if is_processed: + if is_processed and not getattr(p, "_ad_skip_img2img", False): self.save_image( p, init_image, condition="ad_save_images_before", suffix="-ad-before" ) - if p.scripts is not None and self.need_call_process(p): + if self.need_call_process(p): with preseve_prompts(p): p.scripts.process(copy(p)) - try: - ia = p._ad_idx - lenp = len(p.all_prompts) - if ia % lenp == lenp - 1: - self.write_params_txt(p) - except Exception: - pass + self.write_params_txt(p) def on_after_component(component, **_kwargs): @@ -758,9 +799,9 @@ def make_axis_on_xyz_grid(): samplers = [sampler.name for sampler in all_samplers] def set_value(p, x, xs, *, field: str): - if not hasattr(p, "adetailer_xyz"): - p.adetailer_xyz = {} - p.adetailer_xyz[field] = x + if not hasattr(p, "_ad_xyz"): + p._ad_xyz = {} + p._ad_xyz[field] = x axis = [ xyz_grid.AxisOption(