mirror of
https://github.com/lllyasviel/stable-diffusion-webui-forge.git
synced 2026-07-21 21:01:24 +08:00
63 lines
1.9 KiB
Python
63 lines
1.9 KiB
Python
import os
|
|
from typing import TYPE_CHECKING
|
|
|
|
if TYPE_CHECKING:
|
|
from modules_forge.supported_controlnet import ControlModelPatcher
|
|
from modules_forge.supported_preprocessor import Preprocessor
|
|
|
|
from backend import utils
|
|
from modules.paths_internal import models_path, normalized_filepath, parser
|
|
|
|
parser.add_argument(
|
|
"--controlnet-dir",
|
|
type=normalized_filepath,
|
|
help="Path to directory with ControlNet models",
|
|
default=os.path.join(models_path, "ControlNet"),
|
|
)
|
|
parser.add_argument(
|
|
"--controlnet-dirs",
|
|
type=normalized_filepath,
|
|
action="append",
|
|
help="Directories for ControlNet model(s)",
|
|
default=[],
|
|
)
|
|
parser.add_argument(
|
|
"--controlnet-preprocessor-models-dir",
|
|
type=normalized_filepath,
|
|
help="Path to directory with Annotator models",
|
|
default=os.path.join(models_path, "ControlNetPreprocessor"),
|
|
)
|
|
|
|
cmd_opts, _ = parser.parse_known_args()
|
|
|
|
controlnet_dir: str = cmd_opts.controlnet_dir
|
|
os.makedirs(controlnet_dir, exist_ok=True)
|
|
|
|
preprocessor_dir: str = cmd_opts.controlnet_preprocessor_models_dir
|
|
os.makedirs(preprocessor_dir, exist_ok=True)
|
|
|
|
diffusers_dir: str = os.path.join(models_path, "diffusers")
|
|
os.makedirs(diffusers_dir, exist_ok=True)
|
|
|
|
supported_preprocessors: dict[str, "Preprocessor"] = {}
|
|
supported_control_models: list["ControlModelPatcher"] = []
|
|
|
|
|
|
def add_supported_preprocessor(preprocessor: "Preprocessor"):
|
|
supported_preprocessors[preprocessor.name] = preprocessor
|
|
|
|
|
|
def add_supported_control_model(control_model: "ControlModelPatcher"):
|
|
supported_control_models.append(control_model)
|
|
|
|
|
|
def try_load_supported_control_model(ckpt_path: os.PathLike):
|
|
state_dict = utils.load_torch_file(ckpt_path, safe_load=True)
|
|
for supported_type in supported_control_models:
|
|
state_dict_copy = {k: v for k, v in state_dict.items()}
|
|
model = supported_type.try_build_from_state_dict(state_dict_copy, ckpt_path)
|
|
if model is not None:
|
|
return model
|
|
|
|
return None
|