From b3ceeebf9417b521cc24eb732417a0727e1a0a77 Mon Sep 17 00:00:00 2001 From: doctorpangloss <@hiddenswitch.com> Date: Tue, 29 Oct 2024 19:22:51 -0700 Subject: [PATCH] Fix bugs in folder paths - Adding the output paths now correctly registers a relative path, i.e., outputs/loras and models/lorals will now be searched on all your base paths - Adding absolute paths with models/ works better - All the base paths and directories are queried better --- comfy/cmd/folder_paths.py | 17 ++++---- comfy/component_model/folder_path_types.py | 50 ++++++++++++++-------- tests/unit/comfy_test/folder_path_test.py | 15 +++++++ 3 files changed, 56 insertions(+), 26 deletions(-) diff --git a/comfy/cmd/folder_paths.py b/comfy/cmd/folder_paths.py index 88df4e164..275b00a74 100644 --- a/comfy/cmd/folder_paths.py +++ b/comfy/cmd/folder_paths.py @@ -46,6 +46,8 @@ def _resolve_path_with_compatibility(path: Path | str) -> PurePosixPath | Path: """ if isinstance(path, PurePosixPath) and path.is_absolute(): return path + if not path.is_absolute(): + return Path.resolve(_base_path() / path) return Path(path).resolve() @@ -67,16 +69,16 @@ def init_default_paths(folder_names_and_paths: FolderNames, configuration: Optio folder_names_and_paths.add_base_path(base_path) folder_names_and_paths.add(ModelPaths(["checkpoints"], supported_extensions=set(supported_pt_extensions))) folder_names_and_paths.add(ModelPaths(["configs"], additional_absolute_directory_paths={get_package_as_path("comfy.configs")}, supported_extensions={".yaml"})) - folder_names_and_paths.add(ModelPaths(["vae"], supported_extensions={".yaml"})) - folder_names_and_paths.add(ModelPaths(["clip"], supported_extensions={".yaml"})) - folder_names_and_paths.add(ModelPaths(["loras"], supported_extensions={".yaml"})) - folder_names_and_paths.add(ModelPaths(["diffusion_models", "unet"], supported_extensions=set(supported_pt_extensions))) + folder_names_and_paths.add(ModelPaths(["vae"], supported_extensions=set(supported_pt_extensions))) + folder_names_and_paths.add(ModelPaths(["clip"], supported_extensions=set(supported_pt_extensions))) + folder_names_and_paths.add(ModelPaths(["loras"], supported_extensions=set(supported_pt_extensions))) + folder_names_and_paths.add(ModelPaths(folder_names=["diffusion_models", "unet"], supported_extensions=set(supported_pt_extensions), folder_names_are_relative_directory_paths_too=True)) folder_names_and_paths.add(ModelPaths(["clip_vision"], supported_extensions=set(supported_pt_extensions))) folder_names_and_paths.add(ModelPaths(["style_models"], supported_extensions=set(supported_pt_extensions))) folder_names_and_paths.add(ModelPaths(["embeddings"], supported_extensions=set(supported_pt_extensions))) folder_names_and_paths.add(ModelPaths(["diffusers"], supported_extensions=set())) folder_names_and_paths.add(ModelPaths(["vae_approx"], supported_extensions=set(supported_pt_extensions))) - folder_names_and_paths.add(ModelPaths(["controlnet", "t2i_adapter"], supported_extensions=set(supported_pt_extensions))) + folder_names_and_paths.add(ModelPaths(folder_names=["controlnet", "t2i_adapter"], supported_extensions=set(supported_pt_extensions), folder_names_are_relative_directory_paths_too=True)) folder_names_and_paths.add(ModelPaths(["gligen"], supported_extensions=set(supported_pt_extensions))) folder_names_and_paths.add(ModelPaths(["upscale_models"], supported_extensions=set(supported_pt_extensions))) folder_names_and_paths.add(ModelPaths(["custom_nodes"], folder_name_base_path_subdir=construct_path(""), supported_extensions=set())) @@ -228,14 +230,13 @@ def add_model_folder_path(folder_name, full_folder_path: Optional[str] = None, e """ Registers a model path for the given canonical name. :param folder_name: the folder name - :param full_folder_path: When none, defaults to os.path.join(models_dir, folder_name) aka the folder as - a subpath to the default models directory + :param full_folder_path: When none, defaults to os.path.join(models_dir, folder_name) aka the folder as a subpath to the default models directory :param extensions: supported file extensions :return: the folder path """ folder_names_and_paths = _folder_names_and_paths() if full_folder_path is None: - # todo: this should use the subdir patter + # todo: this should use the subdir pattern full_folder_path = os.path.join(_models_dir(), folder_name) diff --git a/comfy/component_model/folder_path_types.py b/comfy/component_model/folder_path_types.py index ffa311236..c3d14e3d8 100644 --- a/comfy/component_model/folder_path_types.py +++ b/comfy/component_model/folder_path_types.py @@ -240,20 +240,25 @@ class FolderNames: yield directory_path def file_paths(self, folder_name: str, relative=False) -> typing.Generator[Path]: - for file_path in itertools.chain.from_iterable([candidate.file_paths(self.base_paths, relative=relative) - for candidate in self.contents if candidate.has_folder_name(folder_name)]): - yield file_path + for candidate in self.contents: + if not candidate.has_folder_name(folder_name): + continue + + for file_path in candidate.file_paths(self.base_paths, relative=relative): + yield file_path def first_existing_or_none(self, folder_name: str, relative_file_path: Path) -> Optional[Path]: - for directory_path in itertools.chain.from_iterable([candidate.directory_paths(self.base_paths) - for candidate in self.contents if candidate.has_folder_name(folder_name)]): - candidate_file_path: Path = construct_path(directory_path / relative_file_path) - try: - # todo: this should follow the symlink - if Path.exists(candidate_file_path): - return candidate_file_path - except OSError: + for candidate in self.contents: + if not candidate.has_folder_name(folder_name): continue + for directory_path in candidate.directory_paths(self.base_paths): + candidate_file_path: Path = construct_path(directory_path / relative_file_path) + try: + # todo: this should follow the symlink + if Path.exists(candidate_file_path): + return candidate_file_path + except OSError: + continue return None def add_supported_extension(self, folder_name: str, *supported_extensions: str | None): @@ -324,6 +329,11 @@ class FolderNames: if candidate.has_folder_name(folder_name): self._modify_model_paths(folder_name, paths, set(), candidate, index=index) + def get_paths(self, folder_name: str) -> typing.Generator[AbstractPaths]: + for candidate in self.contents: + if candidate.has_folder_name(folder_name): + yield candidate + def _modify_model_paths(self, key: str, paths: Iterable[Path | str], supported_extensions: set[str], model_paths: AbstractPaths = None, index: Optional[int] = None) -> AbstractPaths: model_paths = model_paths or ModelPaths([key], supported_extensions=set(supported_extensions), @@ -337,34 +347,37 @@ class FolderNames: # when given absolute paths, try to formulate them as relative paths anyway if path.is_absolute(): for base_path in self.base_paths: + did_add = False try: relative_to_basepath = path.relative_to(base_path) potential_folder_name = relative_to_basepath.stem potential_subdir = relative_to_basepath.parent # does the folder_name so far match the key? - # or have we not seen this folder before? + # and have we not seen this folder before? folder_name_not_already_in_contents = all(not candidate.has_folder_name(potential_folder_name) for candidate in self.contents) - if potential_folder_name == key or folder_name_not_already_in_contents: + if potential_folder_name == key and folder_name_not_already_in_contents: # fix the subdir model_paths.folder_name_base_path_subdir = potential_subdir model_paths.folder_names_are_relative_directory_paths_too = True if folder_name_not_already_in_contents: do_add(model_paths.folder_names, index, potential_folder_name) + did_add = True else: # if the folder name doesn't match the key, check if we have ever seen a folder name that matches the key: if model_paths.folder_names_are_relative_directory_paths_too: if potential_subdir == model_paths.folder_name_base_path_subdir: # then we want to add this to the folder name, because it's probably compatibility do_add(model_paths.folder_names, index, potential_folder_name) + did_add = True else: - # not this case - model_paths.folder_names_are_relative_directory_paths_too = False - - if not model_paths.folder_names_are_relative_directory_paths_too: + 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) 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) + did_add = True # since this was an absolute path that was a subdirectory of one of the base paths, # we are done @@ -375,7 +388,8 @@ class FolderNames: # if we got this far, none of the absolute paths were subdirectories of any base paths # add it to our absolute paths - model_paths.additional_absolute_directory_paths.add(path) + if not did_add: + model_paths.additional_absolute_directory_paths.add(path) else: # since this is a relative path, peacefully add it to model_paths potential_folder_name = path.stem diff --git a/tests/unit/comfy_test/folder_path_test.py b/tests/unit/comfy_test/folder_path_test.py index 51f6bc31c..f5b284b99 100644 --- a/tests/unit/comfy_test/folder_path_test.py +++ b/tests/unit/comfy_test/folder_path_test.py @@ -6,7 +6,9 @@ from pathlib import Path import pytest +from comfy.cli_args_types import Configuration from comfy.cmd import folder_paths +from comfy.cmd.folder_paths import init_default_paths from comfy.component_model.folder_path_types import FolderNames, ModelPaths from comfy.execution_context import context_folder_names_and_paths @@ -73,3 +75,16 @@ def test_get_save_image_path(temp_dir): assert counter == 1 assert subfolder == "" assert filename_prefix == "test" + + +def test_add_output_path_absolute(temp_dir): + names = FolderNames() + config = Configuration() + config.cwd = str(temp_dir) + init_default_paths(names, config) + with context_folder_names_and_paths(names): + folder_paths.add_model_folder_path("diffusion_models", os.path.join(folder_paths.get_output_directory(), "diffusion_models")) + mp: ModelPaths = next(names.get_paths("diffusion_models")) + assert len(mp.additional_absolute_directory_paths) == 0 + assert len(mp.additional_relative_directory_paths) == 1 + assert list(mp.additional_relative_directory_paths)[0] == (Path("output") / "diffusion_models")