fluxdev/stable-diffusion-webui-forge

1import os2import ldm_patched.modules.utils3import argparse4 5from modules.paths_internal import models_path6from pathlib import Path7 8 9parser = argparse.ArgumentParser()10 11parser.add_argument(12    "--controlnet-dir",13    type=Path,14    help="Path to directory with ControlNet models",15    default=None,16)17parser.add_argument(18    "--controlnet-preprocessor-models-dir",19    type=Path,20    help="Path to directory with annotator model directories",21    default=None,22)23 24cmd_opts = parser.parse_known_args()[0]25 26if cmd_opts.controlnet_dir:27    controlnet_dir = str(cmd_opts.controlnet_dir)28else:29    controlnet_dir = os.path.join(models_path, 'ControlNet')30os.makedirs(controlnet_dir, exist_ok=True)31 32if cmd_opts.controlnet_preprocessor_models_dir:33    preprocessor_dir = str(cmd_opts.controlnet_preprocessor_models_dir)34else:35    preprocessor_dir = os.path.join(models_path, 'ControlNetPreprocessor')36os.makedirs(preprocessor_dir, exist_ok=True)37 38diffusers_dir = os.path.join(models_path, 'diffusers')39os.makedirs(diffusers_dir, exist_ok=True)40 41supported_preprocessors = {}42supported_control_models = []43 44 45def add_supported_preprocessor(preprocessor):46    global supported_preprocessors47    p = preprocessor48    supported_preprocessors[p.name] = p49    return50 51 52def add_supported_control_model(control_model):53    global supported_control_models54    supported_control_models.append(control_model)55    return56 57 58def try_load_supported_control_model(ckpt_path):59    global supported_control_models60    state_dict = ldm_patched.modules.utils.load_torch_file(ckpt_path, safe_load=True)61    for supported_type in supported_control_models:62        state_dict_copy = {k: v for k, v in state_dict.items()}63        model = supported_type.try_build_from_state_dict(state_dict_copy, ckpt_path)64        if model is not None:65            return model66    return None67