feat: controlnet inpaint support (#4)

Co-authored-by: @catboxanon

adetailer/__version__.py

@@ -1 +1 @@
-__version__ = "23.4.2"
+__version__ = "23.5.0.dev0"

controlnet_ext/__init__.py

@@ -0,0 +1,3 @@
+from .controlnet_ext import ControlNetExt, controlnet_exists, get_cn_inpaint_models
+
+__all__ = ["ControlNetExt", "controlnet_exists", "get_cn_inpaint_models"]

controlnet_ext/controlnet_ext.py

@@ -0,0 +1,93 @@
+from __future__ import annotations
+
+import importlib
+from functools import lru_cache
+from pathlib import Path
+
+from modules import sd_models, shared
+from modules.paths import data_path, models_path
+
+extensions_path = Path(data_path, "extensions")
+controlnet_exists = any(
+    p.name == "sd-webui-controlnet" for p in extensions_path.iterdir() if p.is_dir()
+)
+
+
+class ControlNetExt:
+    def __init__(self):
+        self.cn_models = ["None"]
+        self.cn_available = False
+        self.external_cn = None
+
+    def init_controlnet(self) -> bool:
+        try:
+            self.external_cn = importlib.import_module(
+                "extensions.sd-webui-controlnet.scripts.external_code", "external_code"
+            )
+            self.cn_available = True
+            models = self.external_cn.get_models()
+            self.cn_models.extend(m for m in models if "inpaint" in m)
+            return True
+        except ImportError:
+            return False
+
+    def _update_scripts_args(self, p, model: str, weight: float):
+        cn_units = [
+            self.external_cn.ControlNetUnit(
+                model=model,
+                weight=weight,
+                control_mode=self.external_cn.ControlMode.BALANCED,
+                module="inpaint_global_harmonious",
+                pixel_perfect=True,
+            )
+        ]
+
+        self.external_cn.update_cn_script_in_processing(p, cn_units)
+
+    def update_scripts_args(self, p, model: str, weight: float):
+        if self.cn_available and model != "None":
+            self._update_scripts_args(p, model, weight)
+
+
+@lru_cache
+def _get_cn_inpaint_models() -> list[str]:
+    """
+    Since we can't import ControlNet, we use a function that does something like
+    controlnet's `list(global_state.cn_models_names.values())`.
+    """
+    cn_model_exts = (".pt", ".pth", ".ckpt", ".safetensors")
+    cn_model_dir = Path(models_path, "ControlNet")
+    cn_model_dir_old = Path(extensions_path, "sd-webui-controlnet", "models")
+    ext_dir1 = shared.opts.data.get("control_net_models_path", "")
+    ext_dir2 = shared.opts.data.get("controlnet_dir", "")
+    name_filter = shared.opts.data.get("control_net_models_name_filter", "")
+    name_filter = name_filter.strip(" ").lower()
+
+    model_paths = []
+
+    for base in [cn_model_dir, cn_model_dir_old, ext_dir1, ext_dir2]:
+        if not base:
+            continue
+        base = Path(base)
+        if not base.exists():
+            continue
+
+        for p in base.rglob("*"):
+            if p.is_file() and p.suffix in cn_model_exts and "inpaint" in p.name:
+                if name_filter and name_filter not in p.name.lower():
+                    continue
+                model_paths.append(p)
+    model_paths.sort(key=lambda p: p.name)
+
+    models = []
+    for p in model_paths:
+        model_hash = sd_models.model_hash(p)
+        name = f"{p.stem} [{model_hash}]"
+        models.append(name)
+    return models
+
+
+def get_cn_inpaint_models() -> list[str]:
+    if controlnet_exists:
+        return _get_cn_inpaint_models()
+    return []

scripts/!adetailer.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import platform
+import sys
 from copy import copy
 from pathlib import Path
 
@@ -10,6 +11,7 @@ import torch
 import modules
 from adetailer import __version__, get_models, mediapipe_predict, ultralytics_predict
 from adetailer.common import dilate_erode, is_all_black, offset
+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 (
@@ -47,6 +49,8 @@ class ADetailerArgs:
         self.ad_inpaint_full_res: bool = args[9]
         self.ad_inpaint_full_res_padding: int = args[10]
         self.ad_cfg_scale: float = args[11]
+        self.ad_controlnet_model: str = args[12]
+        self.ad_controlnet_weight: float = args[13]
 
     def asdict(self):
         return {
@@ -62,6 +66,8 @@ class ADetailerArgs:
             "ad_inpaint_full_res": self.ad_inpaint_full_res,
             "ad_inpaint_full_res_padding": self.ad_inpaint_full_res_padding,
             "ad_cfg_scale": self.ad_cfg_scale,
+            "ad_controlnet_model": self.ad_controlnet_model,
+            "ad_controlnet_weight": self.ad_controlnet_weight,
         }
 
 
@@ -79,6 +85,10 @@ def gr_show(visible=True):
 
 
 class AfterDetailerScript(scripts.Script):
+    def __init__(self):
+        super().__init__()
+        self.controlnet_ext = None
+
     def title(self):
         return AFTER_DETAILER
 
@@ -196,6 +206,28 @@ class AfterDetailerScript(scripts.Script):
                         visible=True,
                     )
 
+                cn_inpaint_models = ["None"] + get_cn_inpaint_models()
+
+                with gr.Group():
+                    with gr.Row():
+                        ad_controlnet_model = gr.Dropdown(
+                            label="ControlNet model",
+                            choices=cn_inpaint_models,
+                            value="None",
+                            visible=controlnet_exists,
+                            type="value",
+                        )
+
+                    with gr.Row():
+                        ad_controlnet_weight = gr.Slider(
+                            label="ControlNet weight",
+                            minimum=0.0,
+                            maximum=1.0,
+                            step=0.05,
+                            value=1.0,
+                            visible=controlnet_exists,
+                        )
+
         all_widgets = [
             ad_model,
             ad_prompt,
@@ -209,14 +241,10 @@ class AfterDetailerScript(scripts.Script):
             ad_inpaint_full_res,
             ad_inpaint_full_res_padding,
             ad_cfg_scale,
+            ad_controlnet_model,
+            ad_controlnet_weight,
         ]
 
-        def on_ad_model_change(model_name):
-            visible = model_name != "None"
-            return {widget: gr_show(visible) for widget in all_widgets[1:]}
-
-        ad_model.change(on_ad_model_change, inputs=[ad_model], outputs=all_widgets[1:])
-
         self.infotext_fields = [
             (ad_model, "ADetailer model"),
             (ad_prompt, "ADetailer prompt"),
@@ -230,12 +258,21 @@ class AfterDetailerScript(scripts.Script):
             (ad_inpaint_full_res, "ADetailer inpaint full"),
             (ad_inpaint_full_res_padding, "ADetailer inpaint padding"),
             (ad_cfg_scale, "ADetailer CFG scale"),
+            (ad_controlnet_model, "ADetailer ControlNet model"),
+            (ad_controlnet_weight, "ADetailer ControlNet weight"),
         ]
 
         return all_widgets
 
-    @staticmethod
+    def init_controlnet_ext(self):
+        if self.controlnet_ext is None:
+            self.controlnet_ext = ControlNetExt()
+            success = self.controlnet_ext.init_controlnet()
+            if not success:
+                print("[-] ADetailer: ControlNetExt init failed.", file=sys.stderr)
+
     def extra_params(
+        self,
         ad_model,
         ad_prompt,
         ad_negative_prompt,
@@ -248,6 +285,8 @@ class AfterDetailerScript(scripts.Script):
         ad_inpaint_full_res,
         ad_inpaint_full_res_padding,
         ad_cfg_scale,
+        ad_controlnet_model,
+        ad_controlnet_weight,
     ):
         params = {
             "ADetailer model": ad_model,
@@ -262,6 +301,8 @@ class AfterDetailerScript(scripts.Script):
             "ADetailer inpaint full": ad_inpaint_full_res,
             "ADetailer inpaint padding": ad_inpaint_full_res_padding,
             "ADetailer CFG scale": ad_cfg_scale,
+            "ADetailer ControlNet model": ad_controlnet_model,
+            "ADetailer ControlNet weight": ad_controlnet_weight,
             "ADetailer version": __version__,
         }
 
@@ -269,6 +310,9 @@ class AfterDetailerScript(scripts.Script):
             params.pop("ADetailer prompt")
         if not ad_negative_prompt:
             params.pop("ADetailer negative prompt")
+        if ad_controlnet_model == "None":
+            params.pop("ADetailer ControlNet model")
+            params.pop("ADetailer ControlNet weight")
 
         return params
 
@@ -344,6 +388,16 @@ class AfterDetailerScript(scripts.Script):
         params_txt = Path(data_path, "params.txt")
         params_txt.write_text(infotext, encoding="utf-8")
 
+    def update_controlnet_args(self, p, args):
+        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 process(self, p, *args_):
         args = self.get_args(*args_)
         if args.ad_model != "None":
@@ -359,6 +413,8 @@ class AfterDetailerScript(scripts.Script):
         if args.ad_model == "None":
             return
 
+        self.init_controlnet_ext()
+
         p._idx = getattr(p, "_idx", -1) + 1
         i = p._idx
 
@@ -407,6 +463,8 @@ class AfterDetailerScript(scripts.Script):
         i2i.script_args = p.script_args
         i2i._disable_adetailer = True
 
+        self.update_controlnet_args(i2i, args)
+
         kwargs = {}
         if args.ad_model.lower().startswith("mediapipe"):
             predictor = mediapipe_predict
@@ -456,7 +514,7 @@ class AfterDetailerScript(scripts.Script):
                 processed = process_images(p2)
 
                 p2 = copy(i2i)
-                p2.init_images = processed.images
+                p2.init_images = [processed.images[0]]
 
             p2.seed = seed + j + 1
             p2.subseed = subseed + j + 1