diff --git a/.github/workflows/pypi.yml b/.github/workflows/pypi.yml index a8d10cf..b98eee8 100644 --- a/.github/workflows/pypi.yml +++ b/.github/workflows/pypi.yml @@ -7,30 +7,7 @@ on: jobs: test: name: test - runs-on: macos-latest - strategy: - matrix: - python-version: - - "3.10" - - "3.11" - - "3.12" - - steps: - - uses: actions/checkout@v4 - - - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v5 - with: - python-version: ${{ matrix.python-version }} - - - uses: yezz123/setup-uv@v4 - - - name: Install dependencies - run: | - uv pip install --system ".[test]" - - - name: Run tests - run: pytest -v + uses: ./.github/workflows/test.yml build: name: build diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index f345788..b8be5e4 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -1,9 +1,12 @@ -name: Test on PR +name: Test on: pull_request: paths: - "adetailer/**.py" + workflow_call: + schedule: + - cron: "0 0 * * 0" jobs: test: @@ -24,7 +27,7 @@ jobs: with: python-version: ${{ matrix.python-version }} - - uses: yezz123/setup-uv@v4 + - uses: astral-sh/setup-uv@v3 - name: Install dependencies run: | diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 6d11532..3343ccb 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -5,7 +5,7 @@ exclude: ^modules/ repos: - repo: https://github.com/pre-commit/pre-commit-hooks - rev: v4.6.0 + rev: v5.0.0 hooks: - id: check-added-large-files args: [--maxkb=100] @@ -24,7 +24,7 @@ repos: - id: prettier - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.5.7 + rev: v0.7.2 hooks: - id: ruff args: [--fix, --exit-non-zero-on-fix] diff --git a/CHANGELOG.md b/CHANGELOG.md index b70e4e4..83d0eac 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,15 @@ # Changelog +## 2024-11-10 + +- v24.11.0 +- `disable_controlnet_units` 함수가 `script_args`의 상태를 변경된 상태로 저장하는 문제 수정 +- XYZ Grid에 CFG Scale, scheduler, noise multiplier 추가 +- Area 또는 Confidence를 기준으로 마스크 최대 갯수를 지정할 수 있도록 함 (PR #720) + +- `ADetailer detector classes`의 element id를 `ad_classes`에서 `ad_model_classes`로 변경 +- `mediapipe` 최대 버전을 0.10.15로 제한 + ## 2024-09-02 - v24.9.0 diff --git a/Taskfile.yml b/Taskfile.yml index 834d82e..bd3c8f1 100644 --- a/Taskfile.yml +++ b/Taskfile.yml @@ -25,8 +25,8 @@ tasks: update: cmds: - - "{{.PYTHON}} -m uv pip install -U ultralytics mediapipe ruff pre-commit black devtools pytest hypothesis" + - "{{.PYTHON}} -m uv pip install -U ultralytics mediapipe ruff pre-commit-uv black devtools pytest hypothesis" update-torch: cmds: - - "{{.PYTHON}} -m uv pip install -U torch torchvision torchaudio -f https://download.pytorch.org/whl/torch_stable.html" + - "{{.PYTHON}} -m uv pip install -U torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu124" diff --git a/aaaaaa/ui.py b/aaaaaa/ui.py index 4b9b484..f637d79 100644 --- a/aaaaaa/ui.py +++ b/aaaaaa/ui.py @@ -204,7 +204,7 @@ def one_ui_group(n: int, is_img2img: bool, webui_info: WebuiInfo): label="ADetailer detector classes" + suffix(n), value="", visible=False, - elem_id=eid("ad_classes"), + elem_id=eid("ad_model_classes"), ) w.ad_model.change( @@ -294,14 +294,22 @@ def detection(w: Widgets, n: int, is_img2img: bool): visible=True, elem_id=eid("ad_confidence"), ) - w.ad_mask_k_largest = gr.Slider( - label="Mask only the top k largest (0 to disable)" + suffix(n), + w.ad_mask_filter_method = gr.Radio( + choices=["Area", "Confidence"], + value="Area", + label="Method to filter top k masks by (confidence or area)" + + suffix(n), + visible=True, + elem_id=eid("ad_mask_filter_method"), + ) + w.ad_mask_k = gr.Slider( + label="Mask only the top k (0 to disable)" + suffix(n), minimum=0, maximum=10, step=1, value=0, visible=True, - elem_id=eid("ad_mask_k_largest"), + elem_id=eid("ad_mask_k"), ) with gr.Column(variant="compact"): diff --git a/adetailer/__version__.py b/adetailer/__version__.py index ca0cfd9..66b04ec 100644 --- a/adetailer/__version__.py +++ b/adetailer/__version__.py @@ -1 +1 @@ -__version__ = "24.9.0" +__version__ = "24.11.0" diff --git a/adetailer/args.py b/adetailer/args.py index 4efcdd7..e1b8751 100644 --- a/adetailer/args.py +++ b/adetailer/args.py @@ -60,7 +60,8 @@ class ADetailerArgs(BaseModel, extra=Extra.forbid): ad_prompt: str = "" ad_negative_prompt: str = "" ad_confidence: confloat(ge=0.0, le=1.0) = 0.3 - ad_mask_k_largest: NonNegativeInt = 0 + ad_mask_filter_method: Literal["Area", "Confidence"] = "Area" + ad_mask_k: NonNegativeInt = 0 ad_mask_min_ratio: confloat(ge=0.0, le=1.0) = 0.0 ad_mask_max_ratio: confloat(ge=0.0, le=1.0) = 1.0 ad_dilate_erode: int = 4 @@ -131,7 +132,11 @@ class ADetailerArgs(BaseModel, extra=Extra.forbid): ppop("ADetailer prompt") ppop("ADetailer negative prompt") p.pop("ADetailer tab enable", None) # always pop - ppop("ADetailer mask only top k largest", cond=0) + ppop( + "ADetailer mask only top k", + ["ADetailer mask only top k", "ADetailer method to decide top k masks"], + cond=0, + ) ppop("ADetailer mask min ratio", cond=0.0) ppop("ADetailer mask max ratio", cond=1.0) ppop("ADetailer x offset", cond=0) @@ -217,7 +222,8 @@ _all_args = [ ("ad_prompt", "ADetailer prompt"), ("ad_negative_prompt", "ADetailer negative prompt"), ("ad_confidence", "ADetailer confidence"), - ("ad_mask_k_largest", "ADetailer mask only top k largest"), + ("ad_mask_filter_method", "ADetailer method to decide top k masks"), + ("ad_mask_k", "ADetailer mask only top k"), ("ad_mask_min_ratio", "ADetailer mask min ratio"), ("ad_mask_max_ratio", "ADetailer mask max ratio"), ("ad_x_offset", "ADetailer x offset"), diff --git a/adetailer/common.py b/adetailer/common.py index 0a4fb7a..ca6415a 100644 --- a/adetailer/common.py +++ b/adetailer/common.py @@ -22,6 +22,7 @@ T = TypeVar("T", int, float) class PredictOutput(Generic[T]): bboxes: list[list[T]] = field(default_factory=list) masks: list[Image.Image] = field(default_factory=list) + confidences: list[float] = field(default_factory=list) preview: Optional[Image.Image] = None diff --git a/adetailer/mask.py b/adetailer/mask.py index 9496aa4..65388c5 100644 --- a/adetailer/mask.py +++ b/adetailer/mask.py @@ -225,6 +225,7 @@ def filter_by_ratio( idx = [i for i in range(items) if is_in_ratio(pred.bboxes[i], low, high, orig_area)] pred.bboxes = [pred.bboxes[i] for i in idx] pred.masks = [pred.masks[i] for i in idx] + pred.confidences = [pred.confidences[i] for i in idx] return pred @@ -236,9 +237,31 @@ def filter_k_largest(pred: PredictOutput[T], k: int = 0) -> PredictOutput[T]: idx = idx[::-1] pred.bboxes = [pred.bboxes[i] for i in idx] pred.masks = [pred.masks[i] for i in idx] + pred.confidences = [pred.confidences[i] for i in idx] return pred +def filter_k_most_confident(pred: PredictOutput[T], k: int = 0) -> PredictOutput[T]: + if not pred.bboxes or not pred.confidences or k == 0: + return pred + idx = np.argsort(pred.confidences)[-k:] + idx = idx[::-1] + pred.bboxes = [pred.bboxes[i] for i in idx] + pred.masks = [pred.masks[i] for i in idx] + pred.confidences = [pred.confidences[i] for i in idx] + return pred + + +def filter_k_by( + pred: PredictOutput[T], k: int = 0, by: str = "Area" +) -> PredictOutput[T]: + if by == "Area": + return filter_k_largest(pred, k) + if by == "Confidence": + return filter_k_most_confident(pred, k) + raise RuntimeError + + # Merge / Invert def mask_merge(masks: list[Image.Image]) -> list[Image.Image]: arrs = [np.array(m) for m in masks] diff --git a/adetailer/mediapipe.py b/adetailer/mediapipe.py index b05fa00..067aa53 100644 --- a/adetailer/mediapipe.py +++ b/adetailer/mediapipe.py @@ -52,6 +52,7 @@ def mediapipe_face_detection( preview_array = img_array.copy() bboxes = [] + confidences = [] for detection in pred.detections: draw_util.draw_detection(preview_array, detection) @@ -63,12 +64,15 @@ def mediapipe_face_detection( x2 = x1 + w y2 = y1 + h + confidences.append(detection.score) bboxes.append([x1, y1, x2, y2]) masks = create_mask_from_bbox(bboxes, image.size) preview = Image.fromarray(preview_array) - return PredictOutput(bboxes=bboxes, masks=masks, preview=preview) + return PredictOutput( + bboxes=bboxes, masks=masks, confidences=confidences, preview=preview + ) def mediapipe_face_mesh( @@ -141,7 +145,6 @@ def mediapipe_face_mesh_eyes_only( preview = image.copy() masks = [] - for landmarks in pred.multi_face_landmarks: points = np.array( [[land.x * w, land.y * h] for land in landmarks.landmark], dtype=int diff --git a/adetailer/ultralytics.py b/adetailer/ultralytics.py index dc93482..7c7a1a7 100644 --- a/adetailer/ultralytics.py +++ b/adetailer/ultralytics.py @@ -37,11 +37,16 @@ def ultralytics_predict( masks = create_mask_from_bbox(bboxes, image.size) else: masks = mask_to_pil(pred[0].masks.data, image.size) + + confidences = pred[0].boxes.conf.cpu().numpy().tolist() + preview = pred[0].plot() preview = cv2.cvtColor(preview, cv2.COLOR_BGR2RGB) preview = Image.fromarray(preview) - return PredictOutput(bboxes=bboxes, masks=masks, preview=preview) + return PredictOutput( + bboxes=bboxes, masks=masks, confidences=confidences, preview=preview + ) def apply_classes(model: YOLO | YOLOWorld, model_path: str | Path, classes: str): diff --git a/install.py b/install.py index 6afd916..e7242cd 100644 --- a/install.py +++ b/install.py @@ -45,7 +45,7 @@ def install(): deps = [ # requirements ("ultralytics", "8.2.0", None), - ("mediapipe", "0.10.13", None), + ("mediapipe", "0.10.13", "0.10.15"), ("rich", "13.0.0", None), ] diff --git a/scripts/!adetailer.py b/scripts/!adetailer.py index 52baf15..9ea743e 100644 --- a/scripts/!adetailer.py +++ b/scripts/!adetailer.py @@ -52,7 +52,7 @@ from adetailer.args import ( from adetailer.common import PredictOutput, ensure_pil_image, safe_mkdir from adetailer.mask import ( filter_by_ratio, - filter_k_largest, + filter_k_by, has_intersection, is_all_black, mask_preprocess, @@ -392,7 +392,7 @@ class AfterDetailerScript(scripts.Script): value = args.ad_scheduler return {"scheduler": value} - def get_override_settings(self, p, args: ADetailerArgs) -> dict[str, Any]: + def get_override_settings(self, _p, args: ADetailerArgs) -> dict[str, Any]: d = {} if args.ad_use_clip_skip: @@ -413,7 +413,7 @@ class AfterDetailerScript(scripts.Script): d["sd_vae"] = args.ad_vae return d - def get_initial_noise_multiplier(self, p, args: ADetailerArgs) -> float | None: + def get_initial_noise_multiplier(self, _p, args: ADetailerArgs) -> float | None: return args.ad_noise_multiplier if args.ad_use_noise_multiplier else None @staticmethod @@ -474,20 +474,30 @@ class AfterDetailerScript(scripts.Script): script_runner.alwayson_scripts = filtered_alwayson return script_runner, script_args - def disable_controlnet_units( - self, script_args: list[Any] | tuple[Any, ...] - ) -> None: - for obj in script_args: - if "controlnet" in obj.__class__.__name__.lower(): - if hasattr(obj, "enabled"): - obj.enabled = False - if hasattr(obj, "input_mode"): - obj.input_mode = getattr(obj.input_mode, "SIMPLE", "simple") + def disable_controlnet_units(self, script_args: Sequence[Any]) -> list[Any]: + new_args = [] + for arg in script_args: + if "controlnet" in arg.__class__.__name__.lower(): + new = copy(arg) + if hasattr(new, "enabled"): + new.enabled = False + if hasattr(new, "input_mode"): + new.input_mode = getattr(new.input_mode, "SIMPLE", "simple") + new_args.append(new) - elif isinstance(obj, dict) and "module" in obj: - obj["enabled"] = False + elif isinstance(arg, dict) and "module" in arg: + new = copy(arg) + new["enabled"] = False + new_args.append(new) - def get_i2i_p(self, p, args: ADetailerArgs, image): + else: + new_args.append(arg) + + return new_args + + def get_i2i_p( + self, p, args: ADetailerArgs, image: Image.Image + ) -> StableDiffusionProcessingImg2Img: seed, subseed = self.get_seed(p) width, height = self.get_width_height(p, args) steps = self.get_steps(p, args) @@ -545,7 +555,7 @@ class AfterDetailerScript(scripts.Script): i2i._ad_inner = True if args.ad_controlnet_model != "Passthrough" and controlnet_type != "forge": - self.disable_controlnet_units(i2i.script_args) + i2i.script_args = self.disable_controlnet_units(i2i.script_args) if args.ad_controlnet_model not in ["None", "Passthrough"]: self.update_controlnet_args(i2i, args) @@ -555,6 +565,9 @@ class AfterDetailerScript(scripts.Script): return i2i def save_image(self, p, image, *, condition: str, suffix: str) -> None: + if not opts.data.get(condition, False): + return + i = get_i(p) if p.all_prompts: i %= len(p.all_prompts) @@ -563,23 +576,22 @@ class AfterDetailerScript(scripts.Script): save_prompt = p.prompt seed, _ = self.get_seed(p) - if opts.data.get(condition, False): - ad_save_images_dir: str = opts.data.get("ad_save_images_dir", "") + ad_save_images_dir: str = opts.data.get("ad_save_images_dir", "") - if not ad_save_images_dir.strip(): - ad_save_images_dir = p.outpath_samples + if not ad_save_images_dir.strip(): + ad_save_images_dir = p.outpath_samples - images.save_image( - image=image, - path=ad_save_images_dir, - basename="", - seed=seed, - prompt=save_prompt, - extension=opts.samples_format, - info=self.infotext(p), - p=p, - suffix=suffix, - ) + images.save_image( + image=image, + path=ad_save_images_dir, + basename="", + seed=seed, + prompt=save_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: @@ -596,7 +608,7 @@ class AfterDetailerScript(scripts.Script): pred = filter_by_ratio( pred, low=args.ad_mask_min_ratio, high=args.ad_mask_max_ratio ) - pred = filter_k_largest(pred, k=args.ad_mask_k_largest) + pred = filter_k_by(pred, k=args.ad_mask_k, by=args.ad_mask_filter_method) pred = self.sort_bboxes(pred) masks = mask_preprocess( pred.masks, @@ -652,7 +664,7 @@ class AfterDetailerScript(scripts.Script): img2img_mask: Image.Image, ad_mask: list[Image.Image] ) -> list[Image.Image]: if ad_mask and img2img_mask.size != ad_mask[0].size: - img2img_mask = img2img_mask.resize(ad_mask[0].size, resample=images.LANCZOS) + img2img_mask = img2img_mask.resize(ad_mask[0].size, resample=Image.LANCZOS) return [mask for mask in ad_mask if has_intersection(img2img_mask, mask)] @staticmethod @@ -663,14 +675,9 @@ class AfterDetailerScript(scripts.Script): mask = ImageChops.invert(mask) mask = create_binary_mask(mask) - if is_skip_img2img(p): - if hasattr(p, "init_images") and p.init_images: - width, height = p.init_images[0].size - else: - msg = "[-] ADetailer: no init_images." - raise RuntimeError(msg) - else: - width, height = p.width, p.height + width, height = p.width, p.height + if is_skip_img2img(p) and hasattr(p, "init_images") and p.init_images: + width, height = p.init_images[0].size return images.resize_image(p.resize_mode, mask, width, height) @staticmethod @@ -969,18 +976,22 @@ def on_ui_settings(): shared.opts.add_option( "ad_save_previews", - shared.OptionInfo(False, "Save mask previews", section=section), + shared.OptionInfo(default=False, label="Save mask previews", section=section), ) shared.opts.add_option( "ad_save_images_before", - shared.OptionInfo(False, "Save images before ADetailer", section=section), + shared.OptionInfo( + default=False, label="Save images before ADetailer", section=section + ), ) shared.opts.add_option( "ad_only_selected_scripts", shared.OptionInfo( - True, "Apply only selected scripts to ADetailer", section=section + default=True, + label="Apply only selected scripts to ADetailer", + section=section, ), ) @@ -1014,7 +1025,9 @@ def on_ui_settings(): shared.opts.add_option( "ad_same_seed_for_each_tab", shared.OptionInfo( - False, "Use same seed for each tab in adetailer", section=section + default=False, + label="Use same seed for each tab in adetailer", + section=section, ), ) @@ -1080,7 +1093,8 @@ def make_axis_on_xyz_grid(): return model_list = ["None", *model_mapping.keys()] - samplers = [sampler.name for sampler in all_samplers] + xyz_samplers = [sampler.name for sampler in all_samplers] + xyz_schedulers = [scheduler.label for scheduler in schedulers] axis = [ xyz_grid.AxisOption( @@ -1119,6 +1133,11 @@ def make_axis_on_xyz_grid(): float, partial(set_value, field="ad_denoising_strength"), ), + xyz_grid.AxisOption( + "[ADetailer] CFG scale 1st", + float, + partial(set_value, field="ad_cfg_scale"), + ), xyz_grid.AxisOption( "[ADetailer] Inpaint only masked 1st", str, @@ -1134,7 +1153,18 @@ def make_axis_on_xyz_grid(): "[ADetailer] ADetailer sampler 1st", str, partial(set_value, field="ad_sampler"), - choices=lambda: samplers, + choices=lambda: xyz_samplers, + ), + xyz_grid.AxisOption( + "[ADetailer] ADetailer scheduler 1st", + str, + partial(set_value, field="ad_scheduler"), + choices=lambda: xyz_schedulers, + ), + xyz_grid.AxisOption( + "[ADetailer] noise multiplier 1st", + float, + partial(set_value, field="ad_noise_multiplier"), ), xyz_grid.AxisOption( "[ADetailer] ControlNet model 1st",