Fix model download directory alias resolution
This commit is contained in:
@@ -87,11 +87,35 @@ def get_download_directory(
|
||||
|
||||
return paths[0]
|
||||
|
||||
def _get_folder_paths(folder_name: str) -> List[str]:
|
||||
"""Return paths for one key without rejecting the other aliases."""
|
||||
try:
|
||||
return list(folder_paths.get_folder_paths(folder_name) or [])
|
||||
except Exception as exc:
|
||||
facade.log.debug(
|
||||
f"Could not get folder path for {folder_name}: {exc}"
|
||||
)
|
||||
return []
|
||||
|
||||
def _get_folder_names() -> List[str]:
|
||||
"""Read folder names across ComfyUI versions."""
|
||||
get_folder_names = getattr(folder_paths, "get_folder_names", None)
|
||||
if callable(get_folder_names):
|
||||
try:
|
||||
return list(get_folder_names() or [])
|
||||
except Exception as exc:
|
||||
facade.log.debug(f"Could not get folder names: {exc}")
|
||||
|
||||
folder_registry = getattr(folder_paths, "folder_names_and_paths", {})
|
||||
if isinstance(folder_registry, dict):
|
||||
return list(folder_registry.keys())
|
||||
return []
|
||||
|
||||
try:
|
||||
paths = []
|
||||
seen_paths = set()
|
||||
for candidate_key in folder_keys:
|
||||
for path in folder_paths.get_folder_paths(candidate_key) or []:
|
||||
for path in _get_folder_paths(candidate_key):
|
||||
path_key = _normalize(path)
|
||||
if path_key in seen_paths:
|
||||
continue
|
||||
@@ -106,11 +130,11 @@ def get_download_directory(
|
||||
return _choose_preferred_path(paths, folder_key)
|
||||
|
||||
# If category not found, try to get any models directory as fallback.
|
||||
all_names = folder_paths.get_folder_names()
|
||||
if all_names:
|
||||
fallback_paths = folder_paths.get_folder_paths(all_names[0])
|
||||
all_names = _get_folder_names()
|
||||
for fallback_key in all_names:
|
||||
fallback_paths = _get_folder_paths(fallback_key)
|
||||
if fallback_paths:
|
||||
return _choose_preferred_path(fallback_paths, all_names[0])
|
||||
return _choose_preferred_path(fallback_paths, fallback_key)
|
||||
except Exception as e:
|
||||
facade.log.debug(f"Could not get folder path for {folder_key}: {e}")
|
||||
|
||||
|
||||
@@ -0,0 +1,39 @@
|
||||
import os
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from core.download.directories import get_download_directory
|
||||
from core.type_utils import get_category_folder_keys
|
||||
|
||||
|
||||
class _FolderPaths:
|
||||
def __init__(self, category, path):
|
||||
self.folder_names_and_paths = {category: ([path], set())}
|
||||
|
||||
def get_folder_paths(self, category):
|
||||
if category not in self.folder_names_and_paths:
|
||||
raise KeyError(category)
|
||||
return self.folder_names_and_paths[category][0]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("registered_category", ["diffusion_models", "unet"])
|
||||
def test_download_directory_skips_unavailable_category_aliases(
|
||||
tmp_path, registered_category
|
||||
):
|
||||
model_directory = str(tmp_path / registered_category)
|
||||
folder_paths = _FolderPaths(registered_category, model_directory)
|
||||
dependencies = SimpleNamespace(
|
||||
folder_paths=folder_paths,
|
||||
get_category_folder_keys=get_category_folder_keys,
|
||||
get_path_identity=lambda value: os.path.normcase(os.path.abspath(value)),
|
||||
is_path_within=lambda _path, _root: False,
|
||||
get_comfy_root_path=lambda _folder_paths: "",
|
||||
get_filename_from_path=os.path.basename,
|
||||
log=MagicMock(),
|
||||
)
|
||||
|
||||
assert get_download_directory("diffusion_models", dependencies=dependencies) == (
|
||||
model_directory
|
||||
)
|
||||
Reference in New Issue
Block a user