|
|
from __future__ import annotations |
|
|
|
|
|
import importlib |
|
|
from functools import lru_cache |
|
|
from pathlib import Path |
|
|
import re |
|
|
|
|
|
from modules import sd_models, shared |
|
|
from modules.paths import data_path, models_path, script_path |
|
|
|
|
|
ext_path = Path(data_path, "extensions") |
|
|
ext_builtin_path = Path(script_path, "extensions-builtin") |
|
|
is_in_builtin = False |
|
|
controlnet_exists = False |
|
|
controlnet_enabled_models = { |
|
|
'inpaint': 'inpaint_global_harmonious', |
|
|
'scribble': 't2ia_sketch_pidi', |
|
|
'lineart': 'lineart_coarse', |
|
|
'openpose': 'openpose_full', |
|
|
'tile': None, |
|
|
} |
|
|
controlnet_model_regex = re.compile(r'.*('+('|'.join(controlnet_enabled_models.keys()))+').*') |
|
|
|
|
|
if ext_path.exists(): |
|
|
controlnet_exists = any( |
|
|
p.name == "sd-webui-controlnet" for p in ext_path.iterdir() if p.is_dir() |
|
|
) |
|
|
|
|
|
if not controlnet_exists and ext_builtin_path.exists(): |
|
|
controlnet_exists = any( |
|
|
p.name == "sd-webui-controlnet" |
|
|
for p in ext_builtin_path.iterdir() |
|
|
if p.is_dir() |
|
|
) |
|
|
|
|
|
if controlnet_exists: |
|
|
is_in_builtin = True |
|
|
|
|
|
|
|
|
class ControlNetExt: |
|
|
def __init__(self): |
|
|
self.cn_models = ["None"] |
|
|
self.cn_available = False |
|
|
self.external_cn = None |
|
|
|
|
|
def init_controlnet(self): |
|
|
if is_in_builtin: |
|
|
import_path = "extensions-builtin.sd-webui-controlnet.scripts.external_code" |
|
|
else: |
|
|
import_path = "extensions.sd-webui-controlnet.scripts.external_code" |
|
|
|
|
|
self.external_cn = importlib.import_module(import_path, "external_code") |
|
|
self.cn_available = True |
|
|
models = self.external_cn.get_models() |
|
|
self.cn_models.extend(m for m in models if controlnet_model_regex.match(m)) |
|
|
|
|
|
def _update_scripts_args(self, p, model: str, weight: float, guidance_end: float): |
|
|
module = None |
|
|
for m, v in controlnet_enabled_models.items(): |
|
|
if m in model: |
|
|
module = v |
|
|
break |
|
|
|
|
|
cn_units = [ |
|
|
self.external_cn.ControlNetUnit( |
|
|
model=model, |
|
|
weight=weight, |
|
|
control_mode=self.external_cn.ControlMode.BALANCED, |
|
|
module=module, |
|
|
guidance_end=guidance_end, |
|
|
pixel_perfect=True, |
|
|
) |
|
|
] |
|
|
|
|
|
self.external_cn.update_cn_script_in_processing(p, cn_units) |
|
|
|
|
|
def update_scripts_args(self, p, model: str, weight: float, guidance_end: float): |
|
|
if self.cn_available and model != "None": |
|
|
self._update_scripts_args(p, model, weight, guidance_end) |
|
|
|
|
|
|
|
|
def get_cn_model_dirs() -> list[Path]: |
|
|
cn_model_dir = Path(models_path, "ControlNet") |
|
|
if is_in_builtin: |
|
|
cn_model_dir_old = Path(ext_builtin_path, "sd-webui-controlnet", "models") |
|
|
else: |
|
|
cn_model_dir_old = Path(ext_path, "sd-webui-controlnet", "models") |
|
|
ext_dir1 = shared.opts.data.get("control_net_models_path", "") |
|
|
ext_dir2 = shared.opts.data.get("controlnet_dir", "") |
|
|
|
|
|
dirs = [cn_model_dir, cn_model_dir_old] |
|
|
for ext_dir in [ext_dir1, ext_dir2]: |
|
|
if ext_dir: |
|
|
dirs.append(Path(ext_dir)) |
|
|
|
|
|
return dirs |
|
|
|
|
|
|
|
|
@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") |
|
|
dirs = get_cn_model_dirs() |
|
|
name_filter = shared.opts.data.get("control_net_models_name_filter", "") |
|
|
name_filter = name_filter.strip(" ").lower() |
|
|
|
|
|
model_paths = [] |
|
|
|
|
|
for base in dirs: |
|
|
if not base.exists(): |
|
|
continue |
|
|
|
|
|
for p in base.rglob("*"): |
|
|
if p.is_file() and p.suffix in cn_model_exts and controlnet_model_regex.match(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 [] |
|
|
|
