diff --git a/comfy/cmd/folder_paths.py b/comfy/cmd/folder_paths.py index 7e22a041e..1b9c2b972 100644 --- a/comfy/cmd/folder_paths.py +++ b/comfy/cmd/folder_paths.py @@ -77,11 +77,11 @@ def init_default_paths(folder_names_and_paths: FolderNames, configuration: Optio hf_cache_paths = ModelPaths(["huggingface_cache"], supported_extensions=set()) # TODO: explore if there is a better way to do this if "HF_HUB_CACHE" in os.environ: - hf_cache_paths.additional_absolute_directory_paths.add(os.environ.get("HF_HUB_CACHE")) + hf_cache_paths.additional_absolute_directory_paths.append(os.environ.get("HF_HUB_CACHE")) model_paths_to_add = [ ModelPaths(["checkpoints"], supported_extensions=set(supported_pt_extensions)), - ModelPaths(["configs"], additional_absolute_directory_paths={get_package_as_path("comfy.configs")}, supported_extensions={".yaml"}), + ModelPaths(["configs"], additional_absolute_directory_paths=[get_package_as_path("comfy.configs")], supported_extensions={".yaml"}), ModelPaths(["vae"], supported_extensions=set(supported_pt_extensions)), ModelPaths(folder_names=["clip", "text_encoders"], supported_extensions=set(supported_pt_extensions)), ModelPaths(["loras"], supported_extensions=set(supported_pt_extensions)), @@ -271,6 +271,17 @@ def add_model_folder_path(folder_name, full_folder_path: Optional[str] = None, e folder_path.paths.insert(0, full_folder_path) else: folder_path.paths.append(full_folder_path) + else: + try: + current_default = folder_path.paths.index(full_folder_path) == 0 + except ValueError: + current_default = False + if current_default != is_default: + folder_path.paths.remove(full_folder_path) + if is_default: + folder_path.paths.insert(0, full_folder_path) + else: + folder_path.paths.append(full_folder_path) if extensions is not None: folder_path.supported_extensions |= extensions diff --git a/comfy/component_model/folder_path_types.py b/comfy/component_model/folder_path_types.py index 49cb7d55d..ffdb08097 100644 --- a/comfy/component_model/folder_path_types.py +++ b/comfy/component_model/folder_path_types.py @@ -1,7 +1,9 @@ from __future__ import annotations +import copy import dataclasses import itertools +import logging import os import typing import weakref @@ -16,6 +18,8 @@ extension_mimetypes_cache = { "webp": "image", } +logger = logging.getLogger(__name__) + def do_add(collection: list | set, index: int | None, item: Any): if isinstance(collection, list) and index == 0: @@ -101,10 +105,18 @@ class PathsList: p: FolderNames = self.parent() p.add_paths(self.folder_name, [path_str]) - def insert(self, path_str: str, index: int): + def insert(self, index: int, path_str: str | Path): p: FolderNames = self.parent() p.add_paths(self.folder_name, [path_str], index=index) + def index(self, value: str | Path): + value = construct_path(value) + p: FolderNames = self.parent() + return [path for path in p.directory_paths(self.folder_name)].index(value) + + def remove(self, value: str | Path): + p: FolderNames = self.parent() + p.remove_paths(self.folder_name, [value]) @dataclasses.dataclass class SupportedExtensions: @@ -168,12 +180,21 @@ class AbstractPaths(ABC): """Check if the given folder name is in folder_names.""" pass + @abstractmethod + def remove_path(self, path: str | Path) -> int: + """ + removes a path + :param path: the path + :return: the number of paths removed + """ + pass + @dataclasses.dataclass class ModelPaths(AbstractPaths): folder_name_base_path_subdir: Path = dataclasses.field(default_factory=lambda: construct_path("models")) - additional_relative_directory_paths: set[Path] = dataclasses.field(default_factory=set) - additional_absolute_directory_paths: set[str | Path] = dataclasses.field(default_factory=set) + additional_relative_directory_paths: list[Path] = dataclasses.field(default_factory=list) + additional_absolute_directory_paths: list[str | Path] = dataclasses.field(default_factory=list) folder_names_are_relative_directory_paths_too: bool = dataclasses.field(default_factory=lambda: True) def directory_paths(self, base_paths: Iterable[Path]) -> typing.Generator[Path]: @@ -220,12 +241,26 @@ class ModelPaths(AbstractPaths): def has_folder_name(self, folder_name: str) -> bool: return folder_name in self.folder_names + def remove_path(self, path: str | Path) -> int: + total = 0 + path = construct_path(path) + for paths_list in (self.additional_absolute_directory_paths, self.additional_relative_directory_paths): + try: + while True: + paths_list.remove(path) + total += 1 + except ValueError: + pass + + return total + @dataclasses.dataclass class FolderNames: application_paths: typing.Optional[ApplicationPaths] = dataclasses.field(default_factory=ApplicationPaths) contents: list[AbstractPaths] = dataclasses.field(default_factory=list) base_paths: list[Path] = dataclasses.field(default_factory=list) + is_root: bool = dataclasses.field(default=lambda: False) def supported_extensions(self, folder_name: str) -> typing.Generator[str]: for candidate in self.contents: @@ -298,8 +333,8 @@ class FolderNames: fn.add( ModelPaths(folder_names=[folder_name], supported_extensions=set(extensions), - additional_relative_directory_paths=set(path for path in paths if not Path(path).is_absolute()), - additional_absolute_directory_paths=set(path for path in paths if Path(path).is_absolute()), folder_names_are_relative_directory_paths_too=False + additional_relative_directory_paths=[path for path in paths if not Path(path).is_absolute()], + additional_absolute_directory_paths=[path for path in paths if Path(path).is_absolute()], folder_names_are_relative_directory_paths_too=False )) return fn @@ -329,6 +364,12 @@ class FolderNames: if candidate.has_folder_name(folder_name): self._modify_model_paths(folder_name, paths, set(), candidate, index=index) + def remove_paths(self, folder_name: str, paths: list[Path | str]): + for candidate in self.contents: + if candidate.has_folder_name(folder_name): + for path in paths: + candidate.remove_path(path) + def get_paths(self, folder_name: str) -> typing.Generator[AbstractPaths]: for candidate in self.contents: if candidate.has_folder_name(folder_name): @@ -341,6 +382,7 @@ class FolderNames: if index is not None and index != 0: raise ValueError(f"index was {index} but only 0 or None is supported") + did_add = False for path in paths: if isinstance(path, str): path = construct_path(path) @@ -374,9 +416,9 @@ class FolderNames: do_add(model_paths.additional_relative_directory_paths, index, relative_to_basepath) did_add = True else: - model_paths.additional_relative_directory_paths.add(relative_to_basepath) + do_add(model_paths.additional_relative_directory_paths, index, relative_to_basepath) for resolve_folder_name in model_paths.folder_names: - model_paths.additional_relative_directory_paths.add(model_paths.folder_name_base_path_subdir / resolve_folder_name) + do_add(model_paths.additional_relative_directory_paths, index, model_paths.folder_name_base_path_subdir / resolve_folder_name) did_add = True # since this was an absolute path that was a subdirectory of one of the base paths, @@ -389,7 +431,7 @@ class FolderNames: # if we got this far, none of the absolute paths were subdirectories of any base paths # add it to our absolute paths if not did_add: - model_paths.additional_absolute_directory_paths.add(path) + do_add(model_paths.additional_absolute_directory_paths, index, path) else: # since this is a relative path, peacefully add it to model_paths potential_folder_name = path.stem @@ -405,7 +447,7 @@ class FolderNames: # if there already exists a folder_name by this name, do not add it, and switch to all relative paths if any(candidate.has_folder_name(potential_folder_name) for candidate in self.contents): model_paths.folder_names_are_relative_directory_paths_too = False - model_paths.additional_relative_directory_paths.add(path) + do_add(model_paths.additional_relative_directory_paths, index, path) model_paths.folder_name_base_path_subdir = construct_path() else: do_add(model_paths.folder_names, index, potential_folder_name) @@ -418,7 +460,7 @@ class FolderNames: else: if any(candidate.has_folder_name(potential_folder_name) for candidate in self.contents): model_paths.folder_names_are_relative_directory_paths_too = False - model_paths.additional_relative_directory_paths.add(path) + do_add(model_paths.additional_relative_directory_paths, index, path) model_paths.folder_name_base_path_subdir = construct_path() else: do_add(model_paths.folder_names, index, potential_folder_name) @@ -488,6 +530,14 @@ class FolderNames: if __default is not None: raise ValueError("get with default is not supported") + def copy(self): + return copy.deepcopy(self) + + def clear(self): + if self.is_root: + logger.warning(f"trying to clear the root folder names and paths instance, this will cause unexpected behavior") + self.contents = [] + class SaveImagePathTuple(NamedTuple): full_output_folder: str diff --git a/comfy/execution_context.py b/comfy/execution_context.py index e63bf6e6a..6bf66c2f6 100644 --- a/comfy/execution_context.py +++ b/comfy/execution_context.py @@ -21,7 +21,7 @@ class ExecutionContext: inference_mode: bool = True -_current_context.set(ExecutionContext(server=ServerStub(), folder_names_and_paths=FolderNames())) +_current_context.set(ExecutionContext(server=ServerStub(), folder_names_and_paths=FolderNames(is_root=True))) def current_execution_context() -> ExecutionContext: diff --git a/tests/unit/comfy_test/folder_path_test.py b/tests/unit/comfy_test/folder_path_test.py index 038d740d7..442f55433 100644 --- a/tests/unit/comfy_test/folder_path_test.py +++ b/tests/unit/comfy_test/folder_path_test.py @@ -16,10 +16,8 @@ from comfy.execution_context import context_folder_names_and_paths @pytest.fixture() def clear_folder_paths(): # Clear the global dictionary before each test to ensure isolation - original = folder_paths.folder_names_and_paths.copy() - folder_paths.folder_names_and_paths.clear() - yield - folder_paths.folder_names_and_paths = original + with context_folder_names_and_paths(FolderNames()): + yield @pytest.fixture def temp_dir():