refactor(directories): extract directory route handlers into DirectoryService
This commit is contained in:
Dariusz L
+11
-279
@@ -1,310 +1,42 @@
|
||||
"""Directories route registration."""
|
||||
|
||||
from ..services.directory_service import DirectoryService
|
||||
from ..services.scanner_service import ScannerService
|
||||
from .context import RouteContext
|
||||
|
||||
|
||||
def register_directory_routes(context: RouteContext):
|
||||
TEMPLATE_KEY_ALIASES = context.get('TEMPLATE_KEY_ALIASES')
|
||||
dedupe_local_base_directories = context.get('dedupe_local_base_directories')
|
||||
get_category_folder_keys = context.get('get_category_folder_keys')
|
||||
get_comfy_root_path = context.get('get_comfy_root_path')
|
||||
get_default_root_for_category = context.get('get_default_root_for_category')
|
||||
get_download_directory = context.get('get_download_directory')
|
||||
get_enabled_download_categories = context.get('get_enabled_download_categories')
|
||||
get_local_path_identity = context.get('get_local_path_identity')
|
||||
is_civarchive_available = context.get('is_civarchive_available')
|
||||
is_lora_manager_archive_available = context.get('is_lora_manager_archive_available')
|
||||
json_api_endpoint = context.get('json_api_endpoint')
|
||||
load_resolver_settings = context.get('load_resolver_settings')
|
||||
normalize_download_category = context.get('normalize_download_category')
|
||||
prefer_local_base_directory = context.get('prefer_local_base_directory')
|
||||
register_version_routes = context.get('register_version_routes')
|
||||
routes = context.get('routes')
|
||||
self = context.get('self')
|
||||
split_path_segments = context.get('split_path_segments')
|
||||
web = context.get('web')
|
||||
"""Register directory routes and delegate behavior to services."""
|
||||
json_api_endpoint = context.get("json_api_endpoint")
|
||||
register_version_routes = context.get("register_version_routes")
|
||||
routes = context.get("routes")
|
||||
web = context.get("web")
|
||||
directory_service = DirectoryService(context)
|
||||
scanner_service = ScannerService(context)
|
||||
|
||||
@routes.get("/model_resolver/directories")
|
||||
@json_api_endpoint("directories")
|
||||
async def get_directories(request):
|
||||
"""Get available model directories."""
|
||||
import folder_paths
|
||||
|
||||
categories = get_enabled_download_categories(list(folder_paths.folder_names_and_paths.keys()))
|
||||
directories = {}
|
||||
for cat in categories:
|
||||
path = get_download_directory(cat)
|
||||
if path:
|
||||
directories[cat] = path
|
||||
|
||||
return web.json_response(directories)
|
||||
return await directory_service.get_directories(request)
|
||||
|
||||
@routes.get("/model_resolver/root-directories")
|
||||
@json_api_endpoint("root directories")
|
||||
async def get_root_directories(request):
|
||||
"""Get configured ComfyUI root directories for path settings."""
|
||||
import os
|
||||
|
||||
import folder_paths
|
||||
|
||||
known_categories = set(folder_paths.folder_names_and_paths.keys())
|
||||
skip_categories = {"custom_nodes", "configs"}
|
||||
categories = get_enabled_download_categories(list(known_categories))
|
||||
category_source_keys = {}
|
||||
for cat in [*known_categories]:
|
||||
folder_key = normalize_download_category(cat)
|
||||
if folder_key:
|
||||
category_source_keys.setdefault(folder_key, []).append(cat)
|
||||
|
||||
roots = {}
|
||||
settings = load_resolver_settings()
|
||||
comfy_root = get_comfy_root_path(folder_paths)
|
||||
for cat in categories:
|
||||
folder_key = normalize_download_category(cat)
|
||||
if folder_key in known_categories:
|
||||
raw_candidate_keys = [folder_key]
|
||||
else:
|
||||
raw_candidate_keys = [
|
||||
folder_key,
|
||||
*TEMPLATE_KEY_ALIASES.get(folder_key, ()),
|
||||
*category_source_keys.get(folder_key, []),
|
||||
]
|
||||
if folder_key == "ultralytics":
|
||||
raw_candidate_keys = [
|
||||
candidate_key
|
||||
for candidate_key in raw_candidate_keys
|
||||
if str(candidate_key or "").strip().lower() != "yolo"
|
||||
]
|
||||
|
||||
candidate_keys = []
|
||||
for candidate_key in raw_candidate_keys:
|
||||
if (
|
||||
candidate_key
|
||||
and candidate_key in known_categories
|
||||
and candidate_key not in candidate_keys
|
||||
):
|
||||
candidate_keys.append(candidate_key)
|
||||
paths = []
|
||||
for candidate_key in candidate_keys:
|
||||
paths.extend(folder_paths.get_folder_paths(candidate_key) or [])
|
||||
if folder_key == "ultralytics":
|
||||
normalized_ultralytics_paths = []
|
||||
for path in paths:
|
||||
normalized_path = os.path.normpath(str(path or ""))
|
||||
basename = os.path.basename(normalized_path).lower()
|
||||
if basename == "yolo":
|
||||
continue
|
||||
parent_dir = os.path.dirname(normalized_path)
|
||||
if (
|
||||
basename in {"bbox", "segm"}
|
||||
and os.path.basename(parent_dir).lower()
|
||||
== "ultralytics"
|
||||
):
|
||||
normalized_path = parent_dir
|
||||
normalized_ultralytics_paths.append(normalized_path)
|
||||
paths = normalized_ultralytics_paths
|
||||
preferred_directory = get_default_root_for_category(
|
||||
folder_key, settings
|
||||
)
|
||||
normalized_paths = dedupe_local_base_directories(
|
||||
paths,
|
||||
preferred_directory=preferred_directory,
|
||||
comfy_root=comfy_root,
|
||||
)
|
||||
roots[folder_key] = normalized_paths
|
||||
|
||||
for raw_key in known_categories:
|
||||
if (
|
||||
not raw_key
|
||||
or raw_key in skip_categories
|
||||
or raw_key in roots
|
||||
):
|
||||
continue
|
||||
raw_paths = folder_paths.get_folder_paths(raw_key) or []
|
||||
roots[raw_key] = dedupe_local_base_directories(
|
||||
raw_paths,
|
||||
comfy_root=comfy_root,
|
||||
)
|
||||
|
||||
return web.json_response(roots)
|
||||
return await directory_service.get_root_directories(request)
|
||||
|
||||
@routes.get("/model_resolver/path-template-suggestions")
|
||||
@json_api_endpoint("path template suggestions")
|
||||
async def get_path_template_suggestions(request):
|
||||
"""Infer path template presets from existing local model folders."""
|
||||
return await scanner_service.get_path_template_suggestions(request)
|
||||
|
||||
@routes.get("/model_resolver/capabilities")
|
||||
@json_api_endpoint("capabilities")
|
||||
async def get_capabilities(request):
|
||||
"""Get optional source capabilities available in this install."""
|
||||
from ..workflow_analyzer import NODE_TYPE_MODEL_WIDGET_CATEGORIES
|
||||
return web.json_response(
|
||||
{
|
||||
"sources": {
|
||||
"civarchive": is_civarchive_available(),
|
||||
"lora_manager_archive": is_lora_manager_archive_available()
|
||||
},
|
||||
"node_rules": NODE_TYPE_MODEL_WIDGET_CATEGORIES
|
||||
}
|
||||
)
|
||||
return await directory_service.get_capabilities(request)
|
||||
|
||||
register_version_routes(routes, web, json_api_endpoint)
|
||||
|
||||
@routes.get("/model_resolver/subfolders/{category}")
|
||||
@json_api_endpoint("subfolders")
|
||||
async def get_subfolders(request):
|
||||
"""Get known subfolders for a category using ComfyUI folder_paths."""
|
||||
import os
|
||||
|
||||
import folder_paths
|
||||
|
||||
raw_category = (request.match_info.get("category") or "").strip()
|
||||
category = normalize_download_category(raw_category)
|
||||
|
||||
if not category or category == "unknown":
|
||||
return web.json_response([])
|
||||
|
||||
known_categories = set(folder_paths.folder_names_and_paths.keys())
|
||||
folder_keys = get_category_folder_keys(category)
|
||||
available_folder_keys = [
|
||||
folder_key for folder_key in folder_keys if folder_key in known_categories
|
||||
]
|
||||
if not available_folder_keys:
|
||||
self.logger.debug(
|
||||
f"Model Resolver: skipping subfolder lookup for unknown category '{raw_category}' -> '{category}'"
|
||||
)
|
||||
return web.json_response([])
|
||||
|
||||
subfolders = {}
|
||||
settings = load_resolver_settings()
|
||||
comfy_root = get_comfy_root_path(folder_paths)
|
||||
preferred_directory = (
|
||||
get_default_root_for_category(category, settings)
|
||||
or get_download_directory(category)
|
||||
or ""
|
||||
)
|
||||
|
||||
def add_subfolder(rel_path, base_dir=""):
|
||||
rel_path = "/".join(split_path_segments(rel_path))
|
||||
if not rel_path or rel_path == ".":
|
||||
return
|
||||
base_dir = os.path.abspath(base_dir) if base_dir else ""
|
||||
base_identity = (
|
||||
get_local_path_identity(base_dir) if base_dir else ""
|
||||
)
|
||||
key = (rel_path.lower(), base_identity)
|
||||
base_label = (
|
||||
os.path.basename(os.path.normpath(base_dir))
|
||||
if base_dir
|
||||
else ""
|
||||
)
|
||||
current = subfolders.get(key)
|
||||
if current and not prefer_local_base_directory(
|
||||
base_dir,
|
||||
current.get("base_directory", ""),
|
||||
preferred_directory,
|
||||
comfy_root,
|
||||
):
|
||||
return
|
||||
subfolders[key] = {
|
||||
"value": rel_path,
|
||||
"label": rel_path,
|
||||
"base_label": base_label,
|
||||
"base_directory": base_dir,
|
||||
}
|
||||
|
||||
raw_base_dirs = []
|
||||
for folder_key in available_folder_keys:
|
||||
for base_dir in folder_paths.get_folder_paths(folder_key) or []:
|
||||
if not base_dir or not os.path.isdir(base_dir):
|
||||
continue
|
||||
raw_base_dirs.append(base_dir)
|
||||
base_dirs = dedupe_local_base_directories(
|
||||
raw_base_dirs,
|
||||
preferred_directory=preferred_directory,
|
||||
comfy_root=comfy_root,
|
||||
)
|
||||
|
||||
def find_base_dir(full_path):
|
||||
if not full_path:
|
||||
return ""
|
||||
full_path_identity = get_local_path_identity(full_path)
|
||||
for base_dir in base_dirs:
|
||||
base_identity = get_local_path_identity(base_dir)
|
||||
try:
|
||||
if (
|
||||
os.path.commonpath(
|
||||
[full_path_identity, base_identity]
|
||||
)
|
||||
== base_identity
|
||||
):
|
||||
return base_dir
|
||||
except Exception:
|
||||
continue
|
||||
return ""
|
||||
|
||||
for folder_key in available_folder_keys:
|
||||
filenames = folder_paths.get_filename_list(folder_key) or []
|
||||
for rel_path in filenames:
|
||||
if not isinstance(rel_path, str):
|
||||
continue
|
||||
base_dir = ""
|
||||
try:
|
||||
base_dir = find_base_dir(
|
||||
folder_paths.get_full_path(folder_key, rel_path)
|
||||
)
|
||||
except Exception:
|
||||
base_dir = ""
|
||||
parts = split_path_segments(rel_path)
|
||||
if len(parts) <= 1:
|
||||
continue
|
||||
current = ""
|
||||
for part in parts[:-1]:
|
||||
current = f"{current}/{part}" if current else part
|
||||
add_subfolder(current, base_dir)
|
||||
|
||||
for base_dir in base_dirs:
|
||||
for root, dirs, _files in os.walk(base_dir):
|
||||
rel_root = os.path.relpath(root, base_dir)
|
||||
for dirname in dirs:
|
||||
rel_path = (
|
||||
dirname
|
||||
if rel_root in ("", ".")
|
||||
else os.path.join(rel_root, dirname)
|
||||
)
|
||||
add_subfolder(rel_path, base_dir)
|
||||
|
||||
value_counts = {}
|
||||
for item in subfolders.values():
|
||||
value_key = item.get("value", "").lower()
|
||||
value_counts[value_key] = value_counts.get(value_key, 0) + 1
|
||||
|
||||
response_items = []
|
||||
for item in subfolders.values():
|
||||
value = item.get("value", "")
|
||||
base_label = item.get("base_label", "")
|
||||
label = (
|
||||
f"{value} ({base_label})"
|
||||
if base_label and value_counts.get(value.lower(), 0) > 1
|
||||
else value
|
||||
)
|
||||
response_items.append(
|
||||
{
|
||||
"value": value,
|
||||
"label": label,
|
||||
"base_directory": item.get("base_directory", ""),
|
||||
}
|
||||
)
|
||||
|
||||
return web.json_response(
|
||||
sorted(
|
||||
response_items,
|
||||
key=lambda item: (
|
||||
item.get("value", "").lower(),
|
||||
item.get("base_directory", "").lower(),
|
||||
),
|
||||
)
|
||||
)
|
||||
return await directory_service.get_subfolders(request)
|
||||
|
||||
@@ -0,0 +1,309 @@
|
||||
"""Directory and subfolder discovery used by HTTP route adapters."""
|
||||
|
||||
import os
|
||||
|
||||
from ..routes.context import RouteContext
|
||||
|
||||
|
||||
class DirectoryService:
|
||||
"""Coordinate configured model directory and subfolder discovery."""
|
||||
|
||||
def __init__(self, context: RouteContext):
|
||||
self.template_key_aliases = context.get("TEMPLATE_KEY_ALIASES")
|
||||
self.dedupe_local_base_directories = context.get(
|
||||
"dedupe_local_base_directories"
|
||||
)
|
||||
self.get_category_folder_keys = context.get("get_category_folder_keys")
|
||||
self.get_comfy_root_path = context.get("get_comfy_root_path")
|
||||
self.get_default_root_for_category = context.get(
|
||||
"get_default_root_for_category"
|
||||
)
|
||||
self.get_download_directory = context.get("get_download_directory")
|
||||
self.get_enabled_download_categories = context.get(
|
||||
"get_enabled_download_categories"
|
||||
)
|
||||
self.get_local_path_identity = context.get("get_local_path_identity")
|
||||
self.is_civarchive_available = context.get("is_civarchive_available")
|
||||
self.is_lora_manager_archive_available = context.get(
|
||||
"is_lora_manager_archive_available"
|
||||
)
|
||||
self.load_resolver_settings = context.get("load_resolver_settings")
|
||||
self.normalize_download_category = context.get(
|
||||
"normalize_download_category"
|
||||
)
|
||||
self.prefer_local_base_directory = context.get(
|
||||
"prefer_local_base_directory"
|
||||
)
|
||||
self.extension = context.get("self")
|
||||
self.split_path_segments = context.get("split_path_segments")
|
||||
self.web = context.get("web")
|
||||
|
||||
async def get_directories(self, request):
|
||||
"""Get available model directories."""
|
||||
import folder_paths
|
||||
|
||||
categories = self.get_enabled_download_categories(
|
||||
list(folder_paths.folder_names_and_paths.keys())
|
||||
)
|
||||
directories = {}
|
||||
for category in categories:
|
||||
path = self.get_download_directory(category)
|
||||
if path:
|
||||
directories[category] = path
|
||||
|
||||
return self.web.json_response(directories)
|
||||
|
||||
async def get_root_directories(self, request):
|
||||
"""Get configured ComfyUI root directories for path settings."""
|
||||
import folder_paths
|
||||
|
||||
known_categories = set(folder_paths.folder_names_and_paths.keys())
|
||||
skip_categories = {"custom_nodes", "configs"}
|
||||
categories = self.get_enabled_download_categories(
|
||||
list(known_categories)
|
||||
)
|
||||
category_source_keys = {}
|
||||
for category in [*known_categories]:
|
||||
folder_key = self.normalize_download_category(category)
|
||||
if folder_key:
|
||||
category_source_keys.setdefault(folder_key, []).append(category)
|
||||
|
||||
roots = {}
|
||||
settings = self.load_resolver_settings()
|
||||
comfy_root = self.get_comfy_root_path(folder_paths)
|
||||
for category in categories:
|
||||
folder_key = self.normalize_download_category(category)
|
||||
if folder_key in known_categories:
|
||||
raw_candidate_keys = [folder_key]
|
||||
else:
|
||||
raw_candidate_keys = [
|
||||
folder_key,
|
||||
*self.template_key_aliases.get(folder_key, ()),
|
||||
*category_source_keys.get(folder_key, []),
|
||||
]
|
||||
if folder_key == "ultralytics":
|
||||
raw_candidate_keys = [
|
||||
candidate_key
|
||||
for candidate_key in raw_candidate_keys
|
||||
if str(candidate_key or "").strip().lower() != "yolo"
|
||||
]
|
||||
|
||||
candidate_keys = []
|
||||
for candidate_key in raw_candidate_keys:
|
||||
if (
|
||||
candidate_key
|
||||
and candidate_key in known_categories
|
||||
and candidate_key not in candidate_keys
|
||||
):
|
||||
candidate_keys.append(candidate_key)
|
||||
paths = []
|
||||
for candidate_key in candidate_keys:
|
||||
paths.extend(folder_paths.get_folder_paths(candidate_key) or [])
|
||||
if folder_key == "ultralytics":
|
||||
normalized_ultralytics_paths = []
|
||||
for path in paths:
|
||||
normalized_path = os.path.normpath(str(path or ""))
|
||||
basename = os.path.basename(normalized_path).lower()
|
||||
if basename == "yolo":
|
||||
continue
|
||||
parent_dir = os.path.dirname(normalized_path)
|
||||
if (
|
||||
basename in {"bbox", "segm"}
|
||||
and os.path.basename(parent_dir).lower()
|
||||
== "ultralytics"
|
||||
):
|
||||
normalized_path = parent_dir
|
||||
normalized_ultralytics_paths.append(normalized_path)
|
||||
paths = normalized_ultralytics_paths
|
||||
preferred_directory = self.get_default_root_for_category(
|
||||
folder_key,
|
||||
settings,
|
||||
)
|
||||
normalized_paths = self.dedupe_local_base_directories(
|
||||
paths,
|
||||
preferred_directory=preferred_directory,
|
||||
comfy_root=comfy_root,
|
||||
)
|
||||
roots[folder_key] = normalized_paths
|
||||
|
||||
for raw_key in known_categories:
|
||||
if (
|
||||
not raw_key
|
||||
or raw_key in skip_categories
|
||||
or raw_key in roots
|
||||
):
|
||||
continue
|
||||
raw_paths = folder_paths.get_folder_paths(raw_key) or []
|
||||
roots[raw_key] = self.dedupe_local_base_directories(
|
||||
raw_paths,
|
||||
comfy_root=comfy_root,
|
||||
)
|
||||
|
||||
return self.web.json_response(roots)
|
||||
|
||||
async def get_capabilities(self, request):
|
||||
"""Get optional source capabilities available in this install."""
|
||||
from ..workflow_analyzer import NODE_TYPE_MODEL_WIDGET_CATEGORIES
|
||||
|
||||
return self.web.json_response(
|
||||
{
|
||||
"sources": {
|
||||
"civarchive": self.is_civarchive_available(),
|
||||
"lora_manager_archive": self.is_lora_manager_archive_available(),
|
||||
},
|
||||
"node_rules": NODE_TYPE_MODEL_WIDGET_CATEGORIES,
|
||||
}
|
||||
)
|
||||
|
||||
async def get_subfolders(self, request):
|
||||
"""Get known subfolders for a category using ComfyUI folder_paths."""
|
||||
import folder_paths
|
||||
|
||||
raw_category = (request.match_info.get("category") or "").strip()
|
||||
category = self.normalize_download_category(raw_category)
|
||||
|
||||
if not category or category == "unknown":
|
||||
return self.web.json_response([])
|
||||
|
||||
known_categories = set(folder_paths.folder_names_and_paths.keys())
|
||||
folder_keys = self.get_category_folder_keys(category)
|
||||
available_folder_keys = [
|
||||
folder_key for folder_key in folder_keys if folder_key in known_categories
|
||||
]
|
||||
if not available_folder_keys:
|
||||
self.extension.logger.debug(
|
||||
f"Model Resolver: skipping subfolder lookup for unknown category "
|
||||
f"'{raw_category}' -> '{category}'"
|
||||
)
|
||||
return self.web.json_response([])
|
||||
|
||||
subfolders = {}
|
||||
settings = self.load_resolver_settings()
|
||||
comfy_root = self.get_comfy_root_path(folder_paths)
|
||||
preferred_directory = (
|
||||
self.get_default_root_for_category(category, settings)
|
||||
or self.get_download_directory(category)
|
||||
or ""
|
||||
)
|
||||
|
||||
def add_subfolder(rel_path, base_dir=""):
|
||||
rel_path = "/".join(self.split_path_segments(rel_path))
|
||||
if not rel_path or rel_path == ".":
|
||||
return
|
||||
base_dir = os.path.abspath(base_dir) if base_dir else ""
|
||||
base_identity = (
|
||||
self.get_local_path_identity(base_dir) if base_dir else ""
|
||||
)
|
||||
key = (rel_path.lower(), base_identity)
|
||||
base_label = (
|
||||
os.path.basename(os.path.normpath(base_dir))
|
||||
if base_dir
|
||||
else ""
|
||||
)
|
||||
current = subfolders.get(key)
|
||||
if current and not self.prefer_local_base_directory(
|
||||
base_dir,
|
||||
current.get("base_directory", ""),
|
||||
preferred_directory,
|
||||
comfy_root,
|
||||
):
|
||||
return
|
||||
subfolders[key] = {
|
||||
"value": rel_path,
|
||||
"label": rel_path,
|
||||
"base_label": base_label,
|
||||
"base_directory": base_dir,
|
||||
}
|
||||
|
||||
raw_base_dirs = []
|
||||
for folder_key in available_folder_keys:
|
||||
for base_dir in folder_paths.get_folder_paths(folder_key) or []:
|
||||
if not base_dir or not os.path.isdir(base_dir):
|
||||
continue
|
||||
raw_base_dirs.append(base_dir)
|
||||
base_dirs = self.dedupe_local_base_directories(
|
||||
raw_base_dirs,
|
||||
preferred_directory=preferred_directory,
|
||||
comfy_root=comfy_root,
|
||||
)
|
||||
|
||||
def find_base_dir(full_path):
|
||||
if not full_path:
|
||||
return ""
|
||||
full_path_identity = self.get_local_path_identity(full_path)
|
||||
for base_dir in base_dirs:
|
||||
base_identity = self.get_local_path_identity(base_dir)
|
||||
try:
|
||||
if (
|
||||
os.path.commonpath(
|
||||
[full_path_identity, base_identity]
|
||||
)
|
||||
== base_identity
|
||||
):
|
||||
return base_dir
|
||||
except Exception:
|
||||
continue
|
||||
return ""
|
||||
|
||||
for folder_key in available_folder_keys:
|
||||
filenames = folder_paths.get_filename_list(folder_key) or []
|
||||
for rel_path in filenames:
|
||||
if not isinstance(rel_path, str):
|
||||
continue
|
||||
base_dir = ""
|
||||
try:
|
||||
base_dir = find_base_dir(
|
||||
folder_paths.get_full_path(folder_key, rel_path)
|
||||
)
|
||||
except Exception:
|
||||
base_dir = ""
|
||||
parts = self.split_path_segments(rel_path)
|
||||
if len(parts) <= 1:
|
||||
continue
|
||||
current = ""
|
||||
for part in parts[:-1]:
|
||||
current = f"{current}/{part}" if current else part
|
||||
add_subfolder(current, base_dir)
|
||||
|
||||
for base_dir in base_dirs:
|
||||
for root, dirs, _files in os.walk(base_dir):
|
||||
rel_root = os.path.relpath(root, base_dir)
|
||||
for dirname in dirs:
|
||||
rel_path = (
|
||||
dirname
|
||||
if rel_root in ("", ".")
|
||||
else os.path.join(rel_root, dirname)
|
||||
)
|
||||
add_subfolder(rel_path, base_dir)
|
||||
|
||||
value_counts = {}
|
||||
for item in subfolders.values():
|
||||
value_key = item.get("value", "").lower()
|
||||
value_counts[value_key] = value_counts.get(value_key, 0) + 1
|
||||
|
||||
response_items = []
|
||||
for item in subfolders.values():
|
||||
value = item.get("value", "")
|
||||
base_label = item.get("base_label", "")
|
||||
label = (
|
||||
f"{value} ({base_label})"
|
||||
if base_label and value_counts.get(value.lower(), 0) > 1
|
||||
else value
|
||||
)
|
||||
response_items.append(
|
||||
{
|
||||
"value": value,
|
||||
"label": label,
|
||||
"base_directory": item.get("base_directory", ""),
|
||||
}
|
||||
)
|
||||
|
||||
return self.web.json_response(
|
||||
sorted(
|
||||
response_items,
|
||||
key=lambda item: (
|
||||
item.get("value", "").lower(),
|
||||
item.get("base_directory", "").lower(),
|
||||
),
|
||||
)
|
||||
)
|
||||
@@ -0,0 +1,319 @@
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from aiohttp import web
|
||||
|
||||
from core.routes.context import RouteContext
|
||||
from core.routes.directories import register_directory_routes
|
||||
|
||||
|
||||
class _Routes:
|
||||
def __init__(self):
|
||||
self.handlers = {}
|
||||
|
||||
def get(self, path):
|
||||
def decorator(handler):
|
||||
self.handlers[("GET", path)] = handler
|
||||
return handler
|
||||
|
||||
return decorator
|
||||
|
||||
def post(self, path):
|
||||
def decorator(handler):
|
||||
self.handlers[("POST", path)] = handler
|
||||
return handler
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
def _json_api_endpoint(_name, **_kwargs):
|
||||
return lambda handler: handler
|
||||
|
||||
|
||||
def _dedupe(paths, **_kwargs):
|
||||
return list(dict.fromkeys(paths))
|
||||
|
||||
|
||||
def _build_routes(**overrides):
|
||||
routes = _Routes()
|
||||
values = {
|
||||
"TEMPLATE_KEY_ALIASES": {"diffusion_models": ("unet",)},
|
||||
"asyncio": asyncio,
|
||||
"dedupe_local_base_directories": _dedupe,
|
||||
"get_category_folder_keys": lambda category: [category],
|
||||
"get_comfy_root_path": MagicMock(return_value=r"C:\ComfyUI"),
|
||||
"get_default_root_for_category": MagicMock(return_value=None),
|
||||
"get_download_directory": MagicMock(return_value=r"C:\models"),
|
||||
"get_enabled_download_categories": MagicMock(
|
||||
side_effect=lambda categories: list(categories)
|
||||
),
|
||||
"get_local_path_identity": lambda path: os.path.normcase(
|
||||
os.path.abspath(path)
|
||||
),
|
||||
"get_model_files": MagicMock(return_value=[]),
|
||||
"infer_download_path_templates": MagicMock(return_value=[]),
|
||||
"is_civarchive_available": MagicMock(return_value=False),
|
||||
"is_lora_manager_archive_available": MagicMock(return_value=False),
|
||||
"json_api_endpoint": _json_api_endpoint,
|
||||
"load_resolver_settings": MagicMock(return_value={}),
|
||||
"normalize_download_category": lambda category: str(category).lower(),
|
||||
"prefer_local_base_directory": MagicMock(return_value=True),
|
||||
"register_version_routes": MagicMock(),
|
||||
"routes": routes,
|
||||
"self": SimpleNamespace(logger=MagicMock()),
|
||||
"split_path_segments": lambda value: [
|
||||
part for part in str(value).replace("\\", "/").split("/") if part
|
||||
],
|
||||
"web": web,
|
||||
}
|
||||
values.update(overrides)
|
||||
register_directory_routes(RouteContext(values))
|
||||
return routes.handlers, values
|
||||
|
||||
|
||||
def _request(category="checkpoints"):
|
||||
return SimpleNamespace(
|
||||
match_info={"category": category},
|
||||
query={},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_directories_route_returns_only_configured_download_directories():
|
||||
handlers, values = _build_routes(
|
||||
get_download_directory=MagicMock(
|
||||
side_effect=lambda category: {
|
||||
"checkpoints": r"C:\models\checkpoints",
|
||||
"loras": "",
|
||||
}.get(category)
|
||||
)
|
||||
)
|
||||
folder_paths = SimpleNamespace(
|
||||
folder_names_and_paths={
|
||||
"checkpoints": ([], set()),
|
||||
"loras": ([], set()),
|
||||
}
|
||||
)
|
||||
|
||||
with patch.dict(sys.modules, {"folder_paths": folder_paths}):
|
||||
response = await handlers[("GET", "/model_resolver/directories")](
|
||||
_request()
|
||||
)
|
||||
|
||||
assert json.loads(response.text) == {
|
||||
"checkpoints": r"C:\models\checkpoints"
|
||||
}
|
||||
values["get_enabled_download_categories"].assert_called_once_with(
|
||||
["checkpoints", "loras"]
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_root_directories_normalizes_ultralytics_paths_and_skips_yolo():
|
||||
handlers, values = _build_routes(
|
||||
get_enabled_download_categories=MagicMock(
|
||||
return_value=["ultralytics"]
|
||||
),
|
||||
get_default_root_for_category=MagicMock(
|
||||
return_value=r"C:\models\ultralytics"
|
||||
),
|
||||
)
|
||||
folder_paths = MagicMock()
|
||||
folder_paths.__file__ = r"C:\ComfyUI\folder_paths.py"
|
||||
folder_paths.folder_names_and_paths = {
|
||||
"ultralytics": (
|
||||
[
|
||||
r"C:\models\ultralytics\bbox",
|
||||
r"C:\models\ultralytics\segm",
|
||||
r"C:\models\ultralytics\yolo",
|
||||
],
|
||||
set(),
|
||||
)
|
||||
}
|
||||
folder_paths.get_folder_paths.side_effect = lambda category: (
|
||||
folder_paths.folder_names_and_paths.get(category, ([], set()))[0]
|
||||
)
|
||||
|
||||
with patch.dict(sys.modules, {"folder_paths": folder_paths}):
|
||||
response = await handlers[
|
||||
("GET", "/model_resolver/root-directories")
|
||||
](_request())
|
||||
|
||||
roots = json.loads(response.text)
|
||||
assert roots["ultralytics"] == [r"C:\models\ultralytics"]
|
||||
values["get_default_root_for_category"].assert_called_once_with(
|
||||
"ultralytics",
|
||||
{},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_root_directories_resolves_aliases_and_preserves_extra_categories():
|
||||
checkpoint_directory = r"C:\models\checkpoints"
|
||||
extra_directory = r"C:\models\extra"
|
||||
handlers, _ = _build_routes(
|
||||
TEMPLATE_KEY_ALIASES={"custom_category": ("checkpoints",)},
|
||||
get_enabled_download_categories=MagicMock(
|
||||
return_value=["custom_category"]
|
||||
),
|
||||
)
|
||||
folder_paths = MagicMock()
|
||||
folder_paths.folder_names_and_paths = {
|
||||
"checkpoints": ([checkpoint_directory], set()),
|
||||
"extra": ([extra_directory], set()),
|
||||
}
|
||||
folder_paths.get_folder_paths.side_effect = lambda category: (
|
||||
folder_paths.folder_names_and_paths.get(category, ([], set()))[0]
|
||||
)
|
||||
|
||||
with patch.dict(sys.modules, {"folder_paths": folder_paths}):
|
||||
response = await handlers[
|
||||
("GET", "/model_resolver/root-directories")
|
||||
](_request())
|
||||
|
||||
roots = json.loads(response.text)
|
||||
assert roots["custom_category"] == [checkpoint_directory]
|
||||
assert roots["extra"] == [extra_directory]
|
||||
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_capabilities_route_returns_optional_sources_and_node_rules():
|
||||
handlers, values = _build_routes()
|
||||
|
||||
with patch(
|
||||
"core.workflow_analyzer.NODE_TYPE_MODEL_WIDGET_CATEGORIES",
|
||||
{"CheckpointLoaderSimple": ["checkpoints"]},
|
||||
):
|
||||
response = await handlers[("GET", "/model_resolver/capabilities")](
|
||||
_request()
|
||||
)
|
||||
|
||||
body = json.loads(response.text)
|
||||
assert body == {
|
||||
"sources": {
|
||||
"civarchive": False,
|
||||
"lora_manager_archive": False,
|
||||
},
|
||||
"node_rules": {"CheckpointLoaderSimple": ["checkpoints"]},
|
||||
}
|
||||
values["is_civarchive_available"].assert_called_once_with()
|
||||
values["is_lora_manager_archive_available"].assert_called_once_with()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_subfolders_route_returns_empty_for_unknown_category(tmp_path):
|
||||
handlers, values = _build_routes(
|
||||
get_category_folder_keys=MagicMock(),
|
||||
normalize_download_category=lambda category: "unknown",
|
||||
)
|
||||
folder_paths = SimpleNamespace(folder_names_and_paths={})
|
||||
|
||||
with patch.dict(sys.modules, {"folder_paths": folder_paths}):
|
||||
response = await handlers[
|
||||
("GET", "/model_resolver/subfolders/{category}")
|
||||
](_request("not-a-category"))
|
||||
|
||||
assert json.loads(response.text) == []
|
||||
values["get_category_folder_keys"].assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_subfolders_route_logs_and_returns_empty_for_unavailable_category():
|
||||
get_category_folder_keys = MagicMock(return_value=["checkpoints"])
|
||||
handlers, values = _build_routes(
|
||||
get_category_folder_keys=get_category_folder_keys
|
||||
)
|
||||
folder_paths = SimpleNamespace(folder_names_and_paths={"loras": ([], set())})
|
||||
|
||||
with patch.dict(sys.modules, {"folder_paths": folder_paths}):
|
||||
response = await handlers[
|
||||
("GET", "/model_resolver/subfolders/{category}")
|
||||
](_request("checkpoints"))
|
||||
|
||||
assert json.loads(response.text) == []
|
||||
values["self"].logger.debug.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_subfolders_route_ignores_invalid_and_unresolvable_entries(tmp_path):
|
||||
base_directory = Path(tmp_path) / "checkpoints"
|
||||
base_directory.mkdir()
|
||||
missing_directory = Path(tmp_path) / "missing"
|
||||
folder_paths = MagicMock()
|
||||
folder_paths.folder_names_and_paths = {
|
||||
"checkpoints": (
|
||||
[str(base_directory), "", str(missing_directory)],
|
||||
set(),
|
||||
)
|
||||
}
|
||||
folder_paths.get_folder_paths.side_effect = lambda category: (
|
||||
folder_paths.folder_names_and_paths.get(category, ([], set()))[0]
|
||||
)
|
||||
folder_paths.get_filename_list.return_value = [
|
||||
None,
|
||||
"model.safetensors",
|
||||
"nested/model.safetensors",
|
||||
]
|
||||
folder_paths.get_full_path.side_effect = RuntimeError("not indexed")
|
||||
handlers, _ = _build_routes()
|
||||
|
||||
with patch.dict(sys.modules, {"folder_paths": folder_paths}):
|
||||
response = await handlers[
|
||||
("GET", "/model_resolver/subfolders/{category}")
|
||||
](_request())
|
||||
|
||||
assert json.loads(response.text) == [
|
||||
{
|
||||
"value": "nested",
|
||||
"label": "nested",
|
||||
"base_directory": "",
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_subfolders_route_discovers_nested_directories(tmp_path):
|
||||
base_directory = Path(tmp_path) / "checkpoints"
|
||||
nested_directory = base_directory / "character" / "pony"
|
||||
nested_directory.mkdir(parents=True)
|
||||
(base_directory / "character" / "model.safetensors").touch()
|
||||
|
||||
folder_paths = MagicMock()
|
||||
folder_paths.folder_names_and_paths = {
|
||||
"checkpoints": ([str(base_directory)], set())
|
||||
}
|
||||
folder_paths.get_folder_paths.side_effect = lambda category: (
|
||||
folder_paths.folder_names_and_paths.get(category, ([], set()))[0]
|
||||
)
|
||||
folder_paths.get_filename_list.return_value = [
|
||||
"character/model.safetensors"
|
||||
]
|
||||
folder_paths.get_full_path.side_effect = (
|
||||
lambda _category, relative_path: str(base_directory / relative_path)
|
||||
)
|
||||
values = {
|
||||
"get_category_folder_keys": MagicMock(return_value=["checkpoints"]),
|
||||
"get_comfy_root_path": MagicMock(return_value=str(tmp_path)),
|
||||
"get_default_root_for_category": MagicMock(return_value=None),
|
||||
"get_download_directory": MagicMock(return_value=str(base_directory)),
|
||||
}
|
||||
handlers, _ = _build_routes(**values)
|
||||
|
||||
with patch.dict(sys.modules, {"folder_paths": folder_paths}):
|
||||
response = await handlers[
|
||||
("GET", "/model_resolver/subfolders/{category}")
|
||||
](_request())
|
||||
|
||||
items = json.loads(response.text)
|
||||
assert [item["value"] for item in items] == [
|
||||
"character",
|
||||
"character/pony",
|
||||
]
|
||||
assert all(item["base_directory"] == str(base_directory) for item in items)
|
||||
Reference in New Issue
Block a user