feat: new args, params refactor

adetailer/args.py

@@ -1,8 +1,8 @@
 from __future__ import annotations
 
 from collections import UserList
-from functools import cached_property
-from typing import Any, NamedTuple, Optional, Union
+from functools import cached_property, partial
+from typing import Any, NamedTuple, Union
 
 import pydantic
 from pydantic import (
@@ -12,6 +12,7 @@ from pydantic import (
     NonNegativeInt,
     PositiveInt,
     confloat,
+    conint,
     validator,
 )
 
@@ -36,9 +37,12 @@ class ADetailerArgs(BaseModel, extra=Extra.forbid):
     ad_prompt: str = ""
     ad_negative_prompt: str = ""
     ad_conf: confloat(ge=0.0, le=1.0) = 0.3
+    ad_mask_min_scale: confloat(ge=0.0, le=1.0) = 0.0
+    ad_mask_max_scale: confloat(ge=0.0, le=1.0) = 1.0
     ad_dilate_erode: int = 32
     ad_x_offset: int = 0
     ad_y_offset: int = 0
+    ad_mask_merge_invert: conint(ge=0, le=2) = 0
     ad_mask_blur: NonNegativeInt = 4
     ad_denoising_strength: confloat(ge=0.0, le=1.0) = 0.4
     ad_inpaint_full_res: bool = True
@@ -49,6 +53,7 @@ class ADetailerArgs(BaseModel, extra=Extra.forbid):
     ad_use_steps: bool = False
     ad_steps: PositiveInt = 28
     ad_use_cfg_scale: bool = False
+    ad_restore_face: bool = False
     ad_cfg_scale: NonNegativeFloat = 7.0
     ad_controlnet_model: str = "None"
     ad_controlnet_weight: confloat(ge=0.0, le=1.0) = 1.0
@@ -65,47 +70,65 @@ class ADetailerArgs(BaseModel, extra=Extra.forbid):
             v /= 100.0
         return v
 
+    @staticmethod
+    def ppop(
+        p: dict[str, Any],
+        key: str,
+        pops: list[str] | None = None,
+        cond: Any = None,
+    ):
+        if pops is None:
+            pops = [key]
+        value = p[key]
+        cond = (not bool(value)) if cond is None else value == cond
+
+        if cond:
+            for k in pops:
+                p.pop(k)
+
     def extra_params(self, suffix: str = ""):
         if self.ad_model == "None":
             return {}
 
-        params = {name: getattr(self, attr) for attr, name in ALL_ARGS}
-        params["ADetailer conf"] = int(params["ADetailer conf"] * 100)
+        p = {name: getattr(self, attr) for attr, name in ALL_ARGS}
+        p["ADetailer conf"] = int(p["ADetailer conf"] * 100)
+        ppop = partial(self.ppop, p)
 
-        if not params["ADetailer prompt"]:
-            params.pop("ADetailer prompt")
-        if not params["ADetailer negative prompt"]:
-            params.pop("ADetailer negative prompt")
-
-        if params["ADetailer x offset"] == 0:
-            params.pop("ADetailer x offset")
-        if params["ADetailer y offset"] == 0:
-            params.pop("ADetailer y offset")
-
-        if not params["ADetailer inpaint full"]:
-            params.pop("ADetailer inpaint padding")
-
-        if not params["ADetailer use inpaint width/height"]:
-            params.pop("ADetailer use inpaint width/height")
-            params.pop("ADetailer inpaint width")
-            params.pop("ADetailer inpaint height")
-
-        if not params["ADetailer use separate steps"]:
-            params.pop("ADetailer use separate steps")
-            params.pop("ADetailer steps")
-
-        if not params["ADetailer use separate CFG scale"]:
-            params.pop("ADetailer use separate CFG scale")
-            params.pop("ADetailer CFG scale")
-
-        if params["ADetailer ControlNet model"] == "None":
-            params.pop("ADetailer ControlNet model")
-            params.pop("ADetailer ControlNet weight")
+        ppop("ADetailer prompt")
+        ppop("ADetailer negative prompt")
+        ppop("ADetailer mask min ratio", cond=0.0)
+        ppop("ADetailer mask max ratio", cond=1.0)
+        ppop("ADetailer x offset", cond=0)
+        ppop("ADetailer y offset", cond=0)
+        ppop("ad_mask_merge_invert", cond=0)
+        ppop("ADetailer inpaint full", ["ADetailer inpaint padding"])
+        ppop(
+            "ADetailer use inpaint width/height",
+            [
+                "ADetailer use inpaint width/height",
+                "ADetailer inpaint width",
+                "ADetailer inpaint height",
+            ],
+        )
+        ppop(
+            "ADetailer use separate steps",
+            ["ADetailer use separate steps", "ADetailer steps"],
+        )
+        ppop(
+            "ADetailer use separate CFG scale",
+            ["ADetailer use separate CFG scale", "ADetailer CFG scale"],
+        )
+        ppop("ADetailer restore face")
+        ppop(
+            "ADetailer ControlNet model",
+            ["ADetailer ControlNet model", "ADetailer ControlNet weight"],
+            cond="None",
+        )
 
         if suffix:
-            params = {k + suffix: v for k, v in params.items()}
+            p = {k + suffix: v for k, v in p.items()}
 
-        return params
+        return p
 
 
 class EnableChecker(BaseModel):
@@ -127,9 +150,12 @@ _all_args = [
     ("ad_prompt", "ADetailer prompt"),
     ("ad_negative_prompt", "ADetailer negative prompt"),
     ("ad_conf", "ADetailer conf"),
-    ("ad_dilate_erode", "ADetailer dilate/erode"),
+    ("ad_mask_min_ratio", "ADetailer mask min ratio"),
+    ("ad_mask_max_ratio", "ADetailer mask max ratio"),
     ("ad_x_offset", "ADetailer x offset"),
     ("ad_y_offset", "ADetailer y offset"),
+    ("ad_dilate_erode", "ADetailer dilate/erode"),
+    ("ad_mask_merge_invert", "ADetailer mask merge/invert"),
     ("ad_mask_blur", "ADetailer mask blur"),
     ("ad_denoising_strength", "ADetailer denoising strength"),
     ("ad_inpaint_full_res", "ADetailer inpaint full"),
@@ -141,6 +167,7 @@ _all_args = [
     ("ad_steps", "ADetailer steps"),
     ("ad_use_cfg_scale", "ADetailer use separate CFG scale"),
     ("ad_cfg_scale", "ADetailer CFG scale"),
+    ("ad_restore_face", "ADetailer restore face"),
     ("ad_controlnet_model", "ADetailer ControlNet model"),
     ("ad_controlnet_weight", "ADetailer ControlNet weight"),
 ]