From a892411cf82e000df8cd76008824fd4c43aa9a91 Mon Sep 17 00:00:00 2001 From: doctorpangloss <@hiddenswitch.com> Date: Wed, 13 Mar 2024 18:11:16 -0700 Subject: [PATCH] Add known controlnet models and add --disable-known-models to prevent it from appearing or downloading --- comfy/cli_args.py | 1 + comfy/cli_args_types.py | 2 + comfy/model_downloader.py | 107 ++++++++++++++++++++++++++++++++++---- comfy/nodes/base_nodes.py | 12 ++--- 4 files changed, 107 insertions(+), 15 deletions(-) diff --git a/comfy/cli_args.py b/comfy/cli_args.py index d52563e69..0497478af 100644 --- a/comfy/cli_args.py +++ b/comfy/cli_args.py @@ -177,6 +177,7 @@ def create_parser() -> argparse.ArgumentParser: parser.add_argument("--external-address", required=False, help="Specifies a base URL for external addresses reported by the API, such as for image paths.") parser.add_argument("--verbose", action="store_true", help="Enables more debug prints.") + parser.add_argument("--disable-known-models", action="store_true", help="Disables automatic downloads of known models and prevents them from appearing in the UI.") # now give plugins a chance to add configuration for entry_point in entry_points().select(group='comfyui.custom_config'): diff --git a/comfy/cli_args_types.py b/comfy/cli_args_types.py index ef467ee7d..19ddf3b9f 100644 --- a/comfy/cli_args_types.py +++ b/comfy/cli_args_types.py @@ -79,6 +79,7 @@ class Configuration(dict): distributed_queue_name (str): This name will be used by the frontends and workers to exchange prompt requests and replies. Progress updates will be prefixed by the queue name, followed by a '.', then the user ID. external_address (str): Specifies a base URL for external addresses reported by the API, such as for image paths. verbose (bool): Shows extra output for debugging purposes such as import errors of custom nodes. + disable_known_models (bool): Disables automatic downloads of known models and prevents them from appearing in the UI. """ def __init__(self, **kwargs): @@ -142,6 +143,7 @@ class Configuration(dict): self.distributed_queue_frontend: bool = False self.distributed_queue_name: str = "comfyui" self.external_address: Optional[str] = None + self.disable_known_models: bool = False for key, value in kwargs.items(): self[key] = value diff --git a/comfy/model_downloader.py b/comfy/model_downloader.py index c294c3728..7e9c66308 100644 --- a/comfy/model_downloader.py +++ b/comfy/model_downloader.py @@ -13,20 +13,21 @@ from .cmd import folder_paths from .model_downloader_types import CivitFile, HuggingFile, CivitModelsGetResponse, CivitFile_ from .interruption import InterruptProcessingException from .utils import ProgressBar, comfy_tqdm +from .cli_args import args -session = Session() +_session = Session() def get_filename_list_with_downloadable(folder_name: str, known_files: List[Any]) -> List[str]: existing = frozenset(folder_paths.get_filename_list(folder_name)) - downloadable = frozenset(str(f) for f in known_files) + downloadable = frozenset() if args.disable_known_models else frozenset(str(f) for f in known_files) return sorted(list(existing | downloadable)) def get_or_download(folder_name: str, filename: str, known_files: List[HuggingFile | CivitFile]) -> Optional[str]: path = folder_paths.get_full_path(folder_name, filename) - if path is None: + if path is None and not args.disable_known_models: try: # todo: should this be the first or last path? destination = folder_paths.get_folder_paths(folder_name)[0] @@ -49,14 +50,14 @@ def get_or_download(folder_name: str, filename: str, known_files: List[HuggingFi save_filename = known_file.save_with_filename or known_file.filename if isinstance(known_file, CivitFile): - model_info_res = session.get( + model_info_res = _session.get( f"https://civitai.com/api/v1/models/{known_file.model_id}?modelVersionId={known_file.model_version_id}") model_info: CivitModelsGetResponse = model_info_res.json() - file: CivitFile_ - for file in chain.from_iterable(version['files'] for version in model_info['modelVersions']): - if file['name'] == filename: - url = file['downloadUrl'] + civit_file: CivitFile_ + for civit_file in chain.from_iterable(version['files'] for version in model_info['modelVersions']): + if civit_file['name'] == filename: + url = civit_file['downloadUrl'] break else: raise RuntimeError("unknown file type") @@ -67,7 +68,7 @@ def get_or_download(folder_name: str, filename: str, known_files: List[HuggingFi destination_with_filename = join(destination, save_filename) try: - with session.get(url, stream=True, allow_redirects=True) as response: + with _session.get(url, stream=True, allow_redirects=True) as response: total_size = int(response.headers.get("content-length", 0)) progress_bar = ProgressBar(total=total_size) with open(destination_with_filename, "wb") as file: @@ -139,10 +140,98 @@ KNOWN_CLIP_VISION_MODELS = [ KNOWN_LORAS = [ CivitFile(model_id=211577, model_version_id=238349, filename="openxl_handsfix.safetensors"), + CivitFile(model_id=324815, model_version_id=364137, filename="blur_control_xl_v1.safetensors"), +] + +KNOWN_CONTROLNETS = [ + HuggingFile("thibaud/controlnet-openpose-sdxl-1.0", "OpenPoseXL2.safetensors"), + HuggingFile("thibaud/controlnet-openpose-sdxl-1.0", "control-lora-openposeXL2-rank256.safetensors"), + HuggingFile("comfyanonymous/ControlNet-v1-1_fp16_safetensors", "control_lora_rank128_v11e_sd15_ip2p_fp16.safetensors"), + HuggingFile("comfyanonymous/ControlNet-v1-1_fp16_safetensors", "control_lora_rank128_v11e_sd15_shuffle_fp16.safetensors"), + HuggingFile("comfyanonymous/ControlNet-v1-1_fp16_safetensors", "control_lora_rank128_v11f1e_sd15_tile_fp16.safetensors"), + HuggingFile("comfyanonymous/ControlNet-v1-1_fp16_safetensors", "control_lora_rank128_v11f1p_sd15_depth_fp16.safetensors"), + HuggingFile("comfyanonymous/ControlNet-v1-1_fp16_safetensors", "control_lora_rank128_v11p_sd15_canny_fp16.safetensors"), + HuggingFile("comfyanonymous/ControlNet-v1-1_fp16_safetensors", "control_lora_rank128_v11p_sd15_inpaint_fp16.safetensors"), + HuggingFile("comfyanonymous/ControlNet-v1-1_fp16_safetensors", "control_lora_rank128_v11p_sd15_lineart_fp16.safetensors"), + HuggingFile("comfyanonymous/ControlNet-v1-1_fp16_safetensors", "control_lora_rank128_v11p_sd15_mlsd_fp16.safetensors"), + HuggingFile("comfyanonymous/ControlNet-v1-1_fp16_safetensors", "control_lora_rank128_v11p_sd15_normalbae_fp16.safetensors"), + HuggingFile("comfyanonymous/ControlNet-v1-1_fp16_safetensors", "control_lora_rank128_v11p_sd15_openpose_fp16.safetensors"), + HuggingFile("comfyanonymous/ControlNet-v1-1_fp16_safetensors", "control_lora_rank128_v11p_sd15_scribble_fp16.safetensors"), + HuggingFile("comfyanonymous/ControlNet-v1-1_fp16_safetensors", "control_lora_rank128_v11p_sd15_seg_fp16.safetensors"), + HuggingFile("comfyanonymous/ControlNet-v1-1_fp16_safetensors", "control_lora_rank128_v11p_sd15_softedge_fp16.safetensors"), + HuggingFile("comfyanonymous/ControlNet-v1-1_fp16_safetensors", "control_lora_rank128_v11p_sd15s2_lineart_anime_fp16.safetensors"), + HuggingFile("comfyanonymous/ControlNet-v1-1_fp16_safetensors", "control_v11e_sd15_ip2p_fp16.safetensors"), + HuggingFile("comfyanonymous/ControlNet-v1-1_fp16_safetensors", "control_v11e_sd15_shuffle_fp16.safetensors"), + HuggingFile("comfyanonymous/ControlNet-v1-1_fp16_safetensors", "control_v11f1e_sd15_tile_fp16.safetensors"), + HuggingFile("comfyanonymous/ControlNet-v1-1_fp16_safetensors", "control_v11f1p_sd15_depth_fp16.safetensors"), + HuggingFile("comfyanonymous/ControlNet-v1-1_fp16_safetensors", "control_v11p_sd15_canny_fp16.safetensors"), + HuggingFile("comfyanonymous/ControlNet-v1-1_fp16_safetensors", "control_v11p_sd15_inpaint_fp16.safetensors"), + HuggingFile("comfyanonymous/ControlNet-v1-1_fp16_safetensors", "control_v11p_sd15_lineart_fp16.safetensors"), + HuggingFile("comfyanonymous/ControlNet-v1-1_fp16_safetensors", "control_v11p_sd15_mlsd_fp16.safetensors"), + HuggingFile("comfyanonymous/ControlNet-v1-1_fp16_safetensors", "control_v11p_sd15_normalbae_fp16.safetensors"), + HuggingFile("comfyanonymous/ControlNet-v1-1_fp16_safetensors", "control_v11p_sd15_openpose_fp16.safetensors"), + HuggingFile("comfyanonymous/ControlNet-v1-1_fp16_safetensors", "control_v11p_sd15_scribble_fp16.safetensors"), + HuggingFile("comfyanonymous/ControlNet-v1-1_fp16_safetensors", "control_v11p_sd15_seg_fp16.safetensors"), + HuggingFile("comfyanonymous/ControlNet-v1-1_fp16_safetensors", "control_v11p_sd15_softedge_fp16.safetensors"), + HuggingFile("comfyanonymous/ControlNet-v1-1_fp16_safetensors", "control_v11p_sd15s2_lineart_anime_fp16.safetensors"), + HuggingFile("comfyanonymous/ControlNet-v1-1_fp16_safetensors", "control_v11u_sd15_tile_fp16.safetensors"), + HuggingFile("lllyasviel/sd_control_collection", "diffusers_xl_canny_full.safetensors"), + HuggingFile("lllyasviel/sd_control_collection", "diffusers_xl_canny_mid.safetensors"), + HuggingFile("lllyasviel/sd_control_collection", "diffusers_xl_canny_small.safetensors"), + HuggingFile("lllyasviel/sd_control_collection", "diffusers_xl_depth_full.safetensors"), + HuggingFile("lllyasviel/sd_control_collection", "diffusers_xl_depth_mid.safetensors"), + HuggingFile("lllyasviel/sd_control_collection", "diffusers_xl_depth_small.safetensors"), + HuggingFile("lllyasviel/sd_control_collection", "ioclab_sd15_recolor.safetensors"), + HuggingFile("lllyasviel/sd_control_collection", "kohya_controllllite_xl_blur.safetensors"), + HuggingFile("lllyasviel/sd_control_collection", "kohya_controllllite_xl_blur_anime.safetensors"), + HuggingFile("lllyasviel/sd_control_collection", "kohya_controllllite_xl_blur_anime_beta.safetensors"), + HuggingFile("lllyasviel/sd_control_collection", "kohya_controllllite_xl_canny.safetensors"), + HuggingFile("lllyasviel/sd_control_collection", "kohya_controllllite_xl_canny_anime.safetensors"), + HuggingFile("lllyasviel/sd_control_collection", "kohya_controllllite_xl_depth.safetensors"), + HuggingFile("lllyasviel/sd_control_collection", "kohya_controllllite_xl_depth_anime.safetensors"), + HuggingFile("lllyasviel/sd_control_collection", "kohya_controllllite_xl_openpose_anime.safetensors"), + HuggingFile("lllyasviel/sd_control_collection", "kohya_controllllite_xl_openpose_anime_v2.safetensors"), + HuggingFile("lllyasviel/sd_control_collection", "kohya_controllllite_xl_scribble_anime.safetensors"), + HuggingFile("lllyasviel/sd_control_collection", "sai_xl_canny_128lora.safetensors"), + HuggingFile("lllyasviel/sd_control_collection", "sai_xl_canny_256lora.safetensors"), + HuggingFile("lllyasviel/sd_control_collection", "sai_xl_depth_128lora.safetensors"), + HuggingFile("lllyasviel/sd_control_collection", "sai_xl_depth_256lora.safetensors"), + HuggingFile("lllyasviel/sd_control_collection", "sai_xl_recolor_128lora.safetensors"), + HuggingFile("lllyasviel/sd_control_collection", "sai_xl_recolor_256lora.safetensors"), + HuggingFile("lllyasviel/sd_control_collection", "sai_xl_sketch_128lora.safetensors"), + HuggingFile("lllyasviel/sd_control_collection", "sai_xl_sketch_256lora.safetensors"), + HuggingFile("lllyasviel/sd_control_collection", "sargezt_xl_depth.safetensors"), + HuggingFile("lllyasviel/sd_control_collection", "sargezt_xl_depth_faid_vidit.safetensors"), + HuggingFile("lllyasviel/sd_control_collection", "sargezt_xl_depth_zeed.safetensors"), + HuggingFile("lllyasviel/sd_control_collection", "sargezt_xl_softedge.safetensors"), + HuggingFile("lllyasviel/sd_control_collection", "t2i-adapter_diffusers_xl_canny.safetensors"), + HuggingFile("lllyasviel/sd_control_collection", "t2i-adapter_diffusers_xl_depth_midas.safetensors"), + HuggingFile("lllyasviel/sd_control_collection", "t2i-adapter_diffusers_xl_depth_zoe.safetensors"), + HuggingFile("lllyasviel/sd_control_collection", "t2i-adapter_diffusers_xl_lineart.safetensors"), + HuggingFile("lllyasviel/sd_control_collection", "t2i-adapter_diffusers_xl_openpose.safetensors"), + HuggingFile("lllyasviel/sd_control_collection", "t2i-adapter_diffusers_xl_sketch.safetensors"), + HuggingFile("lllyasviel/sd_control_collection", "t2i-adapter_xl_canny.safetensors"), + HuggingFile("lllyasviel/sd_control_collection", "t2i-adapter_xl_openpose.safetensors"), + HuggingFile("lllyasviel/sd_control_collection", "t2i-adapter_xl_sketch.safetensors"), + HuggingFile("lllyasviel/sd_control_collection", "thibaud_xl_openpose.safetensors"), + HuggingFile("lllyasviel/sd_control_collection", "thibaud_xl_openpose_256lora.safetensors"), +] + +KNOWN_DIFF_CONTROLNETS = [ + HuggingFile("kohya-ss/ControlNet-diff-modules", "diff_control_sd15_canny_fp16.safetensors"), + HuggingFile("kohya-ss/ControlNet-diff-modules", "diff_control_sd15_depth_fp16.safetensors"), + HuggingFile("kohya-ss/ControlNet-diff-modules", "diff_control_sd15_hed_fp16.safetensors"), + HuggingFile("kohya-ss/ControlNet-diff-modules", "diff_control_sd15_mlsd_fp16.safetensors"), + HuggingFile("kohya-ss/ControlNet-diff-modules", "diff_control_sd15_normal_fp16.safetensors"), + HuggingFile("kohya-ss/ControlNet-diff-modules", "diff_control_sd15_openpose_fp16.safetensors"), + HuggingFile("kohya-ss/ControlNet-diff-modules", "diff_control_sd15_scribble_fp16.safetensors"), + HuggingFile("kohya-ss/ControlNet-diff-modules", "diff_control_sd15_seg_fp16.safetensors"), ] def add_known_models(folder_name: str, symbol: List[Union[CivitFile, HuggingFile]], *models: Union[CivitFile, HuggingFile]) -> List[Union[CivitFile, HuggingFile]]: + if args.disable_known_models: + logging.warning(f"Known models have been disabled in the options (while adding {folder_name}/{','.join(map(str, models))})") symbol += models folder_paths.invalidate_cache(folder_name) return symbol diff --git a/comfy/nodes/base_nodes.py b/comfy/nodes/base_nodes.py index 6c70b803f..af4e6bda2 100644 --- a/comfy/nodes/base_nodes.py +++ b/comfy/nodes/base_nodes.py @@ -24,7 +24,7 @@ from ..cli_args import args from ..cmd import folder_paths, latent_preview from ..model_downloader import get_filename_list_with_downloadable, get_or_download, KNOWN_CHECKPOINTS, \ - KNOWN_CLIP_VISION_MODELS, KNOWN_GLIGEN_MODELS, KNOWN_UNCLIP_CHECKPOINTS, KNOWN_LORAS + KNOWN_CLIP_VISION_MODELS, KNOWN_GLIGEN_MODELS, KNOWN_UNCLIP_CHECKPOINTS, KNOWN_LORAS, KNOWN_CONTROLNETS, KNOWN_DIFF_CONTROLNETS from ..nodes.common import MAX_RESOLUTION from .. import controlnet @@ -703,7 +703,7 @@ class VAELoader: class ControlNetLoader: @classmethod def INPUT_TYPES(s): - return {"required": { "control_net_name": (folder_paths.get_filename_list("controlnet"),)}} + return {"required": { "control_net_name": (get_filename_list_with_downloadable("controlnet", KNOWN_CONTROLNETS),)}} RETURN_TYPES = ("CONTROL_NET",) FUNCTION = "load_controlnet" @@ -711,7 +711,7 @@ class ControlNetLoader: CATEGORY = "loaders" def load_controlnet(self, control_net_name): - controlnet_path = folder_paths.get_full_path("controlnet", control_net_name) + controlnet_path = get_or_download("controlnet", control_net_name, KNOWN_CONTROLNETS) controlnet_ = controlnet.load_controlnet(controlnet_path) return (controlnet_,) @@ -719,7 +719,7 @@ class DiffControlNetLoader: @classmethod def INPUT_TYPES(s): return {"required": { "model": ("MODEL",), - "control_net_name": (folder_paths.get_filename_list("controlnet"),)}} + "control_net_name": (get_filename_list_with_downloadable("controlnet", KNOWN_DIFF_CONTROLNETS),)}} RETURN_TYPES = ("CONTROL_NET",) FUNCTION = "load_controlnet" @@ -727,9 +727,9 @@ class DiffControlNetLoader: CATEGORY = "loaders" def load_controlnet(self, model, control_net_name): - controlnet_path = folder_paths.get_full_path("controlnet", control_net_name) + controlnet_path = get_or_download("controlnet", control_net_name, KNOWN_DIFF_CONTROLNETS) controlnet_ = controlnet.load_controlnet(controlnet_path, model) - return (controlnet,) + return (controlnet_,) class ControlNetApply: