File size: 3,951 Bytes
34097e9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 |
from __future__ import annotations
import importlib
import re
from functools import lru_cache
from pathlib import Path
from modules import extensions, 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")
controlnet_exists = False
controlnet_path = None
cn_base_path = ""
for extension in extensions.active():
if not extension.enabled:
continue
# For cases like sd-webui-controlnet-master
if "sd-webui-controlnet" in extension.name:
controlnet_exists = True
controlnet_path = Path(extension.path)
cn_base_path = ".".join(controlnet_path.parts[-2:])
break
cn_model_module = {
"inpaint": "inpaint_global_harmonious",
"scribble": "t2ia_sketch_pidi",
"lineart": "lineart_coarse",
"openpose": "openpose_full",
"tile": None,
}
cn_model_regex = re.compile("|".join(cn_model_module.keys()))
class ControlNetExt:
def __init__(self):
self.cn_models = ["None"]
self.cn_available = False
self.external_cn = None
def init_controlnet(self):
import_path = cn_base_path + ".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 cn_model_regex.search(m))
def update_scripts_args(
self, p, model: str, weight: float, guidance_start: float, guidance_end: float
):
if (not self.cn_available) or model == "None":
return
module = None
for m, v in cn_model_module.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_start=guidance_start,
guidance_end=guidance_end,
pixel_perfect=True,
)
]
self.external_cn.update_cn_script_in_processing(p, cn_units)
def get_cn_model_dirs() -> list[Path]:
cn_model_dir = Path(models_path, "ControlNet")
if controlnet_path is not None:
cn_model_dir_old = controlnet_path.joinpath("models")
else:
cn_model_dir_old = None
ext_dir1 = shared.opts.data.get("control_net_models_path", "")
ext_dir2 = shared.opts.data.get("controlnet_dir", "")
dirs = [cn_model_dir]
for ext_dir in [cn_model_dir_old, ext_dir1, ext_dir2]:
if ext_dir:
dirs.append(Path(ext_dir))
return dirs
@lru_cache
def _get_cn_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 cn_model_regex.search(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_models() -> list[str]:
if controlnet_exists:
return _get_cn_models()
return []
|