From 748924de1251c0ec4d55867833e452a3010ae38a Mon Sep 17 00:00:00 2001 From: doctorpangloss <2229300+doctorpangloss@users.noreply.github.com> Date: Mon, 6 Oct 2025 15:21:10 -0700 Subject: [PATCH] Improve model downloading - adding known controlnet models now works better - comfyui_controlnet_aux sometimes wants torchhub paths with files that are in the form directory/filename.safetensors. This is now supported - save_with_filename now correctly matches again --- comfy/cmd/folder_paths.py | 2 +- comfy/distributed/server_stub.py | 3 +++ comfy/model_downloader.py | 12 ++++++++---- comfy/model_downloader_types.py | 9 ++++++--- comfy/nodes/vanilla_node_importing.py | 6 +++++- 5 files changed, 23 insertions(+), 9 deletions(-) diff --git a/comfy/cmd/folder_paths.py b/comfy/cmd/folder_paths.py index da0986a2f..3544ac202 100644 --- a/comfy/cmd/folder_paths.py +++ b/comfy/cmd/folder_paths.py @@ -101,7 +101,7 @@ def init_default_paths(folder_names_and_paths: FolderNames, configuration: Optio ModelPaths(["embeddings"], supported_extensions=set(supported_pt_extensions)), ModelPaths(["diffusers"], supported_extensions=set()), ModelPaths(["vae_approx"], supported_extensions=set(supported_pt_extensions)), - ModelPaths(folder_names=["controlnet", "t2i_adapter"], supported_extensions=set(supported_pt_extensions), folder_names_are_relative_directory_paths_too=True), + ModelPaths(folder_names=["controlnet", "t2i_adapter", "diff_controlnet"], supported_extensions=set(supported_pt_extensions), folder_names_are_relative_directory_paths_too=True), ModelPaths(["gligen"], supported_extensions=set(supported_pt_extensions)), ModelPaths(["upscale_models"], supported_extensions=set(supported_pt_extensions)), ModelPaths(["custom_nodes"], folder_name_base_path_subdir=construct_path(""), supported_extensions=set()), diff --git a/comfy/distributed/server_stub.py b/comfy/distributed/server_stub.py index 3a24a1333..912cf65d9 100644 --- a/comfy/distributed/server_stub.py +++ b/comfy/distributed/server_stub.py @@ -31,3 +31,6 @@ class ServerStub(ExecutorToClientProgress): @property def receive_all_progress_notifications(self) -> bool: return False + + def add_on_prompt_handler(self, handler): + pass \ No newline at end of file diff --git a/comfy/model_downloader.py b/comfy/model_downloader.py index 1fd929f6a..3231bfaa8 100644 --- a/comfy/model_downloader.py +++ b/comfy/model_downloader.py @@ -185,12 +185,16 @@ def get_or_download(folder_name: str, filename: str, known_files: Optional[List[ link_successful = True exc_info_link = {} if path is not None: - destination_link = os.path.join(this_model_directory, linked_filename) - if Path(destination_link).is_file(): + if Path(linked_filename).is_absolute(): + raise ValueError(f"{known_file.repo_id}/{known_file.filename} surprisingly was trying to link to an absolute path {linked_filename}, failing") + + destination_link = Path(this_model_directory) / linked_filename + if destination_link.is_file(): logger.warning(f"{known_file.repo_id}/{known_file.filename} could not link to {destination_link} because the path already exists, which is unexpected") else: try: - os.makedirs(this_model_directory, exist_ok=True) + # sometimes, linked filename has a path in it, on purpose, such as with controlnet_aux nodes + Path(destination_link).parent.mkdir(parents=True, exist_ok=True) os.symlink(path, destination_link) except FileExistsError: # the download was resumed @@ -520,7 +524,7 @@ KNOWN_DIFF_CONTROLNETS: Final[KnownDownloadables] = KnownDownloadables([ 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"), -], folder_name="controlnet") +], folder_name="diff_controlnet") KNOWN_APPROX_VAES: Final[KnownDownloadables] = KnownDownloadables([ HuggingFile("madebyollin/taesd", "taesd_decoder.safetensors", show_in_ui=False), diff --git a/comfy/model_downloader_types.py b/comfy/model_downloader_types.py index 749156c85..818ff53cf 100644 --- a/comfy/model_downloader_types.py +++ b/comfy/model_downloader_types.py @@ -5,15 +5,17 @@ import dataclasses import functools from os.path import split from pathlib import PurePosixPath -from typing import Optional, List, Sequence, Union, Iterable +from typing import Optional, List, Sequence, Union, Iterable, Protocol from can_ada import parse, URL # pylint: disable=no-name-in-module -from typing_extensions import TypedDict, NotRequired +from typing_extensions import TypedDict, NotRequired, runtime_checkable from .component_model.executor_types import ComboOptions from .component_model.files import canonicalize_path + + @dataclasses.dataclass(frozen=True) class UrlFile: _url: str @@ -112,7 +114,8 @@ class DownloadableFileList(ComboOptions, list[str]): main_name = str(f) self._validation_view.add(canonicalize_path(main_name)) self._validation_view.update(map(canonicalize_path, f.alternate_filenames)) - + if f.save_with_filename is not None: + self._validation_view.add(canonicalize_path(f.save_with_filename)) if getattr(f, 'show_in_ui', True): ui_view.add(main_name) diff --git a/comfy/nodes/vanilla_node_importing.py b/comfy/nodes/vanilla_node_importing.py index 971455fa7..166481155 100644 --- a/comfy/nodes/vanilla_node_importing.py +++ b/comfy/nodes/vanilla_node_importing.py @@ -43,10 +43,14 @@ class StreamToLogger: pass -class _PromptServerStub(): +class _PromptServerStub: def __init__(self): self.routes = prompt_server_instance_routes + self.on_prompt_handlers = [] + def add_on_prompt_handler(self, handler): + # todo: these need to be added to a real prompt server if the loading order is behaving in a complex way + self.on_prompt_handlers.append(handler) def _vanilla_load_importing_execute_prestartup_script(node_paths: Iterable[str]) -> None: def execute_script(script_path):