2024-01-28 17:08:33 +00:00
|
|
|
import os
|
2024-01-29 22:25:03 +00:00
|
|
|
import ldm_patched.modules.utils
|
2024-02-17 16:03:17 +00:00
|
|
|
import argparse
|
2024-01-28 17:08:33 +00:00
|
|
|
|
2024-01-31 23:22:06 +00:00
|
|
|
from modules.paths_internal import models_path
|
2024-02-17 16:03:17 +00:00
|
|
|
from pathlib import Path
|
|
|
|
|
|
|
|
|
|
|
|
parser = argparse.ArgumentParser()
|
|
|
|
|
|
|
|
parser.add_argument(
|
|
|
|
"--controlnet-dir",
|
|
|
|
type=Path,
|
|
|
|
help="Path to directory with ControlNet models",
|
|
|
|
default=None,
|
|
|
|
)
|
|
|
|
parser.add_argument(
|
|
|
|
"--controlnet-preprocessor-models-dir",
|
|
|
|
type=Path,
|
|
|
|
help="Path to directory with annotator model directories",
|
|
|
|
default=None,
|
|
|
|
)
|
|
|
|
|
|
|
|
cmd_opts = parser.parse_known_args()[0]
|
2024-01-28 17:08:33 +00:00
|
|
|
|
2024-02-12 02:22:31 +00:00
|
|
|
if cmd_opts.controlnet_dir:
|
|
|
|
controlnet_dir = str(cmd_opts.controlnet_dir)
|
|
|
|
else:
|
|
|
|
controlnet_dir = os.path.join(models_path, 'ControlNet')
|
2024-01-28 17:08:33 +00:00
|
|
|
os.makedirs(controlnet_dir, exist_ok=True)
|
|
|
|
|
2024-02-12 02:22:31 +00:00
|
|
|
if cmd_opts.controlnet_preprocessor_models_dir:
|
|
|
|
preprocessor_dir = str(cmd_opts.controlnet_preprocessor_models_dir)
|
|
|
|
else:
|
|
|
|
preprocessor_dir = os.path.join(models_path, 'ControlNetPreprocessor')
|
2024-01-28 17:08:33 +00:00
|
|
|
os.makedirs(preprocessor_dir, exist_ok=True)
|
2024-01-28 14:44:41 +00:00
|
|
|
|
2024-02-04 01:16:21 +00:00
|
|
|
diffusers_dir = os.path.join(models_path, 'diffusers')
|
|
|
|
os.makedirs(diffusers_dir, exist_ok=True)
|
|
|
|
|
2024-01-29 22:25:03 +00:00
|
|
|
supported_preprocessors = {}
|
|
|
|
supported_control_models = []
|
2024-01-28 14:44:41 +00:00
|
|
|
|
|
|
|
|
2024-01-29 22:25:03 +00:00
|
|
|
def add_supported_preprocessor(preprocessor):
|
|
|
|
global supported_preprocessors
|
2024-01-29 16:03:03 +00:00
|
|
|
p = preprocessor
|
2024-01-29 22:25:03 +00:00
|
|
|
supported_preprocessors[p.name] = p
|
2024-01-28 18:27:43 +00:00
|
|
|
return
|
|
|
|
|
|
|
|
|
2024-01-29 22:25:03 +00:00
|
|
|
def add_supported_control_model(control_model):
|
|
|
|
global supported_control_models
|
|
|
|
supported_control_models.append(control_model)
|
|
|
|
return
|
2024-01-28 14:44:41 +00:00
|
|
|
|
|
|
|
|
2024-01-29 22:25:03 +00:00
|
|
|
def try_load_supported_control_model(ckpt_path):
|
|
|
|
global supported_control_models
|
|
|
|
state_dict = ldm_patched.modules.utils.load_torch_file(ckpt_path, safe_load=True)
|
|
|
|
for supported_type in supported_control_models:
|
2024-02-09 02:24:04 +00:00
|
|
|
state_dict_copy = {k: v for k, v in state_dict.items()}
|
2024-01-29 22:25:03 +00:00
|
|
|
model = supported_type.try_build_from_state_dict(state_dict_copy, ckpt_path)
|
|
|
|
if model is not None:
|
|
|
|
return model
|
|
|
|
return None
|