119 lines
4.3 KiB
Python
119 lines
4.3 KiB
Python
from modules import scripts, shared
|
|
from scripts import controlnet
|
|
|
|
|
|
def find_xyz_grid():
|
|
for data in scripts.scripts_data:
|
|
if data.script_class.__module__ in ["xyz_grid.py", "xy_grid.py"] and hasattr(data, "module"):
|
|
return data.module
|
|
|
|
return None
|
|
|
|
|
|
def add_axis_options(xyz_grid):
|
|
class AxisOption(xyz_grid.AxisOption):
|
|
"""This class returns an instance of the xyz_grid.AxisOption class.
|
|
|
|
If clone is not specified, a single instance is returned.
|
|
If clone is True, the number of copies is
|
|
obtained from control_net_max_models_num and a list is returned.
|
|
If clone is an integer greater than 1,
|
|
that number of copies is made and returned as a list.
|
|
"""
|
|
def __new__(cls, *args, clone=None, **kwargs):
|
|
def init():
|
|
# Compatible with older WebUI
|
|
try:
|
|
cls.__init__(this, *args, **kwargs)
|
|
except:
|
|
kwargs.pop("choices", None)
|
|
cls.__init__(this, *args, **kwargs)
|
|
|
|
if clone:
|
|
this = super().__new__(cls)
|
|
init()
|
|
return this._clone(clone)
|
|
else:
|
|
this = super().__new__(xyz_grid.AxisOption)
|
|
init()
|
|
return this
|
|
|
|
def _clone(self, num=True):
|
|
if num is True:
|
|
num = shared.opts.data.get("control_net_max_models_num", 1)
|
|
|
|
def copy():
|
|
return self.__class__(**self.__dict__)
|
|
|
|
instance_list = [copy()]
|
|
|
|
for i in range(1, num):
|
|
instance = copy()
|
|
instance.label = f"{self.label} - {i}"
|
|
|
|
apply_func = self.apply.__kwdefaults__["enclosure"]
|
|
apply_arg = f"{self.apply.__kwdefaults__['field']}_{i}"
|
|
instance.apply = apply_func(apply_arg)
|
|
|
|
instance_list.append(instance)
|
|
|
|
return instance_list
|
|
|
|
def enable_control_net(p):
|
|
shared.opts.data["control_net_allow_script_control"] = True
|
|
setattr(p, "control_net_enabled", True)
|
|
|
|
def apply_field(field):
|
|
def core(p, x, xs, *, field=field, enclosure=apply_field):
|
|
enable_control_net(p)
|
|
setattr(p, field, x)
|
|
|
|
return core
|
|
|
|
def choices_model():
|
|
controlnet.update_cn_models()
|
|
return list(controlnet.cn_models_names.values())
|
|
|
|
def confirm_model(p, xs):
|
|
confirm_list = choices_model()
|
|
for x in xs:
|
|
if x not in confirm_list:
|
|
raise RuntimeError(f"Unknown ControlNet Model: {x}")
|
|
|
|
def choices_resize_mode():
|
|
return ["Envelope (Outer Fit)", "Scale to Fit (Inner Fit)", "Just Resize"]
|
|
|
|
def confirm_resize_mode(p, xs):
|
|
confirm_list = choices_resize_mode()
|
|
for x in xs:
|
|
if x not in confirm_list:
|
|
raise RuntimeError(f"Unknown Resize Mode: {x}")
|
|
|
|
def choices_preprocessor():
|
|
return list(controlnet.Script().preprocessor)
|
|
|
|
def confirm_preprocessor(p, xs):
|
|
confirm_list = choices_preprocessor()
|
|
for x in xs:
|
|
if x not in confirm_list:
|
|
raise RuntimeError(f"Unknown Preprocessor: {x}")
|
|
|
|
extra_axis_options = [
|
|
AxisOption("[ControlNet] Model", str, apply_field("control_net_model"), choices=choices_model, confirm=confirm_model, cost=0.9),
|
|
AxisOption("[ControlNet] Weight", float, apply_field("control_net_weight")),
|
|
AxisOption("[ControlNet] Guidance Strength", float, apply_field("control_net_guidance_strength")),
|
|
AxisOption("[ControlNet] Resize Mode", str, apply_field("control_net_resize_mode"), choices=choices_resize_mode, confirm=confirm_resize_mode),
|
|
AxisOption("[ControlNet] Preprocessor", str, apply_field("control_net_module"), choices=choices_preprocessor, confirm=confirm_preprocessor),
|
|
AxisOption("[ControlNet] Pre Resolution", int, apply_field("control_net_pres")),
|
|
AxisOption("[ControlNet] Pre Threshold A", float, apply_field("control_net_pthr_a")),
|
|
AxisOption("[ControlNet] Pre Threshold B", float, apply_field("control_net_pthr_b")),
|
|
]
|
|
|
|
xyz_grid.axis_options.extend(extra_axis_options)
|
|
|
|
|
|
xyz_grid = find_xyz_grid()
|
|
|
|
if xyz_grid:
|
|
add_axis_options(xyz_grid)
|