| 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 [] |
|
|
