diff --git a/aaaaaa/ui.py b/aaaaaa/ui.py index f637d79..212095c 100644 --- a/aaaaaa/ui.py +++ b/aaaaaa/ui.py @@ -12,6 +12,7 @@ from adetailer import ADETAILER, __version__ from adetailer.args import ALL_ARGS, MASK_MERGE_INVERT from controlnet_ext import controlnet_exists, controlnet_type, get_cn_models +from lib_controlnet.external_code import ControlMode if controlnet_type == "forge": from lib_controlnet import global_state @@ -718,3 +719,12 @@ def controlnet(w: Widgets, n: int, is_img2img: bool): interactive=controlnet_exists, elem_id=eid("ad_controlnet_guidance_end"), ) + + w.ad_controlnet_control_mode = gr.Radio( + choices=[e.value for e in ControlMode], + value=ControlMode.BALANCED.value, + label="ControlNet control mode" + suffix(n), + visible=True, + interactive=controlnet_exists, + elem_id=eid("ad_controlnet_control_mode"), + ) diff --git a/adetailer/args.py b/adetailer/args.py index e1b8751..d9d925c 100644 --- a/adetailer/args.py +++ b/adetailer/args.py @@ -4,7 +4,8 @@ from dataclasses import dataclass from enum import Enum from functools import cached_property, partial -from typing import Any, Literal, NamedTuple, Optional +from typing import Any, Literal, NamedTuple, Optional, Union +from lib_controlnet.external_code import ControlMode try: from pydantic.v1 import ( @@ -96,6 +97,7 @@ class ADetailerArgs(BaseModel, extra=Extra.forbid): ad_controlnet_weight: confloat(ge=0.0, le=1.0) = 1.0 ad_controlnet_guidance_start: confloat(ge=0.0, le=1.0) = 0.0 ad_controlnet_guidance_end: confloat(ge=0.0, le=1.0) = 1.0 + ad_controlnet_control_mode: Union[ControlMode, int, str] = ControlMode.BALANCED is_api: bool = True @validator("is_api", pre=True) @@ -195,6 +197,7 @@ def extra_params(self, suffix: str = "") -> dict[str, Any]: "ADetailer ControlNet weight", "ADetailer ControlNet guidance start", "ADetailer ControlNet guidance end", + "ADetailer ControlNet control mode", ], cond="None", ) @@ -258,6 +261,7 @@ def need_skip(self) -> bool: ("ad_controlnet_weight", "ADetailer ControlNet weight"), ("ad_controlnet_guidance_start", "ADetailer ControlNet guidance start"), ("ad_controlnet_guidance_end", "ADetailer ControlNet guidance end"), + ("ad_controlnet_control_mode", "ADetailer ControlNet control mode"), ] _args = [Arg(*args) for args in _all_args] diff --git a/controlnet_ext/controlnet_ext_forge.py b/controlnet_ext/controlnet_ext_forge.py index 5fba119..5d9652b 100644 --- a/controlnet_ext/controlnet_ext_forge.py +++ b/controlnet_ext/controlnet_ext_forge.py @@ -4,7 +4,7 @@ import numpy as np from lib_controlnet import external_code, global_state -from lib_controlnet.external_code import ControlNetUnit +from lib_controlnet.external_code import ControlNetUnit, ControlMode from modules import scripts from modules.processing import StableDiffusionProcessing @@ -53,6 +53,7 @@ def update_scripts_args( # noqa: PLR0913 weight: float, guidance_start: float, guidance_end: float, + control_mode: ControlMode, ): if (not self.cn_available) or model == "None": return @@ -81,6 +82,7 @@ def update_scripts_args( # noqa: PLR0913 weight=weight, guidance_start=guidance_start, guidance_end=guidance_end, + control_mode=control_mode, processor_res=pres, ) ], diff --git a/scripts/!adetailer.py b/scripts/!adetailer.py index 9ea743e..42470e2 100644 --- a/scripts/!adetailer.py +++ b/scripts/!adetailer.py @@ -174,6 +174,7 @@ def update_controlnet_args(self, p, args: ADetailerArgs) -> None: weight=args.ad_controlnet_weight, guidance_start=args.ad_controlnet_guidance_start, guidance_end=args.ad_controlnet_guidance_end, + control_mode=args.ad_controlnet_control_mode ) def is_ad_enabled(self, *args) -> bool: