71 lines
2.7 KiB
Python
71 lines
2.7 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):
|
|
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):
|
|
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 = [
|
|
xyz_grid.AxisOption("[ControlNet] Model", str, apply_field("control_net_model"), choices=choices_model, confirm=confirm_model, cost=0.9),
|
|
xyz_grid.AxisOption("[ControlNet] Weight", float, apply_field("control_net_weight")),
|
|
xyz_grid.AxisOption("[ControlNet] Guidance Strength", float, apply_field("control_net_guidance_strength")),
|
|
xyz_grid.AxisOption("[ControlNet] Resize Mode", str, apply_field("control_net_resize_mode"), choices=choices_resize_mode),
|
|
xyz_grid.AxisOption("[ControlNet] Preprocessor", str, apply_field("control_net_module"), choices=choices_preprocessor, confirm=confirm_preprocessor),
|
|
xyz_grid.AxisOption("[ControlNet] Pre Resolution", int, apply_field("control_net_pres")),
|
|
xyz_grid.AxisOption("[ControlNet] Pre Threshold A", float, apply_field("control_net_pthr_a")),
|
|
xyz_grid.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)
|