wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-01-11 06:40:48 +08:00
Code Issues Actions 12 Packages Projects Releases Wiki Activity

Fix bugs in folder paths

Browse Source 
- 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
This commit is contained in:
doctorpangloss 2024-10-29 19:22:51 -07:00
parent a8d8bff548
commit b3ceeebf94
3 changed files with 56 additions and 26 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
comfy/cmd/folder_paths.py
View File

@ -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)

50
comfy/component_model/folder_path_types.py
View File

@ -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

15
tests/unit/comfy_test/folder_path_test.py
View File

@ -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")