refactor: split route registry into dedicated route modules
This commit is contained in:
Dariusz L
@@ -0,0 +1,318 @@
|
||||
"""Directories route registration."""
|
||||
|
||||
def register_directory_routes(context):
|
||||
TEMPLATE_KEY_ALIASES = context.get('TEMPLATE_KEY_ALIASES')
|
||||
asyncio = context.get('asyncio')
|
||||
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')
|
||||
get_model_files = context.get('get_model_files')
|
||||
infer_download_path_templates = context.get('infer_download_path_templates')
|
||||
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')
|
||||
|
||||
@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)
|
||||
|
||||
@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)
|
||||
|
||||
@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."""
|
||||
from ..sources.popular import get_base_models_config
|
||||
|
||||
force_rescan = request.query.get("force") == "1"
|
||||
models = await asyncio.to_thread(get_model_files, force_rescan)
|
||||
base_models_config = get_base_models_config()
|
||||
suggestions = await asyncio.to_thread(
|
||||
infer_download_path_templates,
|
||||
models,
|
||||
base_models_config,
|
||||
)
|
||||
return web.json_response(suggestions)
|
||||
|
||||
@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
|
||||
}
|
||||
)
|
||||
|
||||
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(),
|
||||
),
|
||||
)
|
||||
)
|
||||
@@ -0,0 +1,349 @@
|
||||
"""Downloads route registration."""
|
||||
|
||||
def register_download_routes(context):
|
||||
Aria2InstallError = context.get('Aria2InstallError')
|
||||
UnsafeUrlError = context.get('UnsafeUrlError')
|
||||
asyncio = context.get('asyncio')
|
||||
cancel_download = context.get('cancel_download')
|
||||
clear_completed_downloads = context.get('clear_completed_downloads')
|
||||
first_non_empty = context.get('first_non_empty')
|
||||
get_all_progress = context.get('get_all_progress')
|
||||
get_aria2_status = context.get('get_aria2_status')
|
||||
get_civarchive_model_details = context.get('get_civarchive_model_details')
|
||||
get_civitai_model_details = context.get('get_civitai_model_details')
|
||||
get_default_root_for_category = context.get('get_default_root_for_category')
|
||||
get_download_directory = context.get('get_download_directory')
|
||||
get_override_settings_from_request = context.get('get_override_settings_from_request')
|
||||
get_progress = context.get('get_progress')
|
||||
host_matches_domain = context.get('host_matches_domain')
|
||||
install_aria2_engine = context.get('install_aria2_engine')
|
||||
is_allowed_model_download_filename = context.get('is_allowed_model_download_filename')
|
||||
json_api_endpoint = context.get('json_api_endpoint')
|
||||
load_resolver_settings = context.get('load_resolver_settings')
|
||||
normalize_download_category = context.get('normalize_download_category')
|
||||
pause_download = context.get('pause_download')
|
||||
resolve_download_subfolder = context.get('resolve_download_subfolder')
|
||||
resume_download = context.get('resume_download')
|
||||
routes = context.get('routes')
|
||||
sanitize_download_filename = context.get('sanitize_download_filename')
|
||||
save_resolver_settings = context.get('save_resolver_settings')
|
||||
self = context.get('self')
|
||||
split_path_segments = context.get('split_path_segments')
|
||||
start_aria2_daemon = context.get('start_aria2_daemon')
|
||||
start_background_download = context.get('start_background_download')
|
||||
stop_aria2_daemon = context.get('stop_aria2_daemon')
|
||||
to_bool = context.get('to_bool')
|
||||
to_int = context.get('to_int')
|
||||
validate_public_http_url = context.get('validate_public_http_url')
|
||||
web = context.get('web')
|
||||
|
||||
@routes.post("/model_resolver/download")
|
||||
@json_api_endpoint("download", return_success_on_error=True)
|
||||
async def download_model(request):
|
||||
"""Start downloading a model."""
|
||||
data = await request.json()
|
||||
url = data.get("url", "")
|
||||
filename = data.get("filename", "")
|
||||
category = data.get("category", "checkpoints")
|
||||
category = normalize_download_category(category)
|
||||
subfolder = data.get("subfolder", "")
|
||||
base_directory = data.get("base_directory", "")
|
||||
path_metadata = data.get("path_metadata", {})
|
||||
if not isinstance(path_metadata, dict):
|
||||
path_metadata = {}
|
||||
download_metadata = data.get("download_metadata") or data.get(
|
||||
"metadata", {}
|
||||
)
|
||||
if not isinstance(download_metadata, dict):
|
||||
download_metadata = {}
|
||||
download_metadata = dict(download_metadata)
|
||||
settings = load_resolver_settings()
|
||||
if not base_directory:
|
||||
base_directory = get_default_root_for_category(category, settings)
|
||||
subfolder = resolve_download_subfolder(
|
||||
category,
|
||||
subfolder,
|
||||
path_metadata,
|
||||
settings,
|
||||
)
|
||||
|
||||
if not url:
|
||||
return web.json_response(
|
||||
{"error": "URL is required"}, status=400
|
||||
)
|
||||
try:
|
||||
url = await asyncio.to_thread(validate_public_http_url, url)
|
||||
except UnsafeUrlError as exc:
|
||||
return web.json_response({"error": str(exc)}, status=400)
|
||||
|
||||
from urllib.parse import urlparse as _download_urlparse
|
||||
|
||||
download_host = _download_urlparse(url).hostname
|
||||
|
||||
if not filename:
|
||||
# Extract filename from URL
|
||||
from urllib.parse import unquote, urlparse
|
||||
|
||||
parsed = urlparse(url)
|
||||
filename = unquote(parsed.path.split("/")[-1])
|
||||
filename = sanitize_download_filename(filename)
|
||||
|
||||
if not filename:
|
||||
return web.json_response(
|
||||
{"error": "Could not determine filename"}, status=400
|
||||
)
|
||||
if not is_allowed_model_download_filename(filename):
|
||||
return web.json_response(
|
||||
{"error": "Unsupported model file extension"}, status=400
|
||||
)
|
||||
|
||||
# Build headers if needed
|
||||
headers = {}
|
||||
if host_matches_domain(download_host, "huggingface.co"):
|
||||
hf_token = data.get("hf_token", "")
|
||||
if hf_token:
|
||||
headers["Authorization"] = f"Bearer {hf_token}"
|
||||
elif host_matches_domain(
|
||||
download_host,
|
||||
"civitai.com",
|
||||
"civitai.red",
|
||||
):
|
||||
civitai_key = data.get("civitai_key", "")
|
||||
if civitai_key and "token=" not in url:
|
||||
url += (
|
||||
f"{'&' if '?' in url else '?'}token={civitai_key}"
|
||||
)
|
||||
civitai_session_token = str(
|
||||
data.get("civitai_session_token", "") or ""
|
||||
).strip()
|
||||
if civitai_session_token:
|
||||
headers["Cookie"] = (
|
||||
f"__Secure-civitai-token={civitai_session_token}"
|
||||
)
|
||||
|
||||
def _first_metadata_value(*values):
|
||||
return first_non_empty(*values)
|
||||
|
||||
def _metadata_int(value):
|
||||
return to_int(value)
|
||||
|
||||
inferred_source = ""
|
||||
if host_matches_domain(download_host, "civitai.com", "civitai.red"):
|
||||
inferred_source = "civitai"
|
||||
elif host_matches_domain(download_host, "huggingface.co"):
|
||||
inferred_source = "huggingface"
|
||||
|
||||
download_metadata.setdefault("filename", filename)
|
||||
download_metadata.setdefault("category", category)
|
||||
download_metadata.setdefault("download_url", url)
|
||||
download_metadata.setdefault("source_url", url)
|
||||
download_metadata.setdefault("path_metadata", path_metadata)
|
||||
download_metadata.setdefault(
|
||||
"source",
|
||||
_first_metadata_value(
|
||||
download_metadata.get("details_source"),
|
||||
path_metadata.get("source"),
|
||||
inferred_source,
|
||||
),
|
||||
)
|
||||
|
||||
model_id = _metadata_int(
|
||||
_first_metadata_value(
|
||||
download_metadata.get("model_id"),
|
||||
download_metadata.get("modelId"),
|
||||
path_metadata.get("model_id"),
|
||||
)
|
||||
)
|
||||
version_id = _metadata_int(
|
||||
_first_metadata_value(
|
||||
download_metadata.get("version_id"),
|
||||
download_metadata.get("versionId"),
|
||||
path_metadata.get("version_id"),
|
||||
)
|
||||
)
|
||||
source_name = str(
|
||||
_first_metadata_value(
|
||||
download_metadata.get("details_source"),
|
||||
download_metadata.get("source"),
|
||||
)
|
||||
).lower()
|
||||
try:
|
||||
if (
|
||||
source_name == "civitai"
|
||||
and model_id
|
||||
and not download_metadata.get("civitai_details")
|
||||
):
|
||||
details = await asyncio.to_thread(
|
||||
get_civitai_model_details,
|
||||
model_id,
|
||||
version_id,
|
||||
data.get("civitai_key", ""),
|
||||
)
|
||||
if details:
|
||||
download_metadata["civitai_details"] = details
|
||||
elif (
|
||||
source_name == "civarchive"
|
||||
and model_id
|
||||
and not download_metadata.get("civitai_details")
|
||||
):
|
||||
details = await asyncio.to_thread(
|
||||
get_civarchive_model_details,
|
||||
model_id,
|
||||
version_id,
|
||||
)
|
||||
if details:
|
||||
download_metadata["civitai_details"] = details
|
||||
except Exception as metadata_error:
|
||||
self.logger.warning(
|
||||
f"Model metadata lookup failed: {metadata_error}"
|
||||
)
|
||||
|
||||
target_directory = ""
|
||||
target_path = ""
|
||||
try:
|
||||
import os as _download_os
|
||||
|
||||
target_directory = (
|
||||
get_download_directory(category, base_directory) or ""
|
||||
)
|
||||
if target_directory and subfolder:
|
||||
target_directory = _download_os.path.join(
|
||||
target_directory, *split_path_segments(subfolder)
|
||||
)
|
||||
if target_directory:
|
||||
target_path = _download_os.path.join(
|
||||
target_directory, filename
|
||||
)
|
||||
except Exception:
|
||||
target_directory = ""
|
||||
target_path = ""
|
||||
|
||||
# Start background download
|
||||
download_id = start_background_download(
|
||||
url=url,
|
||||
filename=filename,
|
||||
category=category,
|
||||
headers=headers if headers else None,
|
||||
subfolder=subfolder,
|
||||
base_directory=base_directory,
|
||||
metadata=download_metadata,
|
||||
)
|
||||
|
||||
return web.json_response(
|
||||
{
|
||||
"success": True,
|
||||
"download_id": download_id,
|
||||
"filename": filename,
|
||||
"category": category,
|
||||
"path": target_path,
|
||||
"directory": target_directory,
|
||||
"download_backend": settings.get("download_backend", "python"),
|
||||
}
|
||||
)
|
||||
|
||||
@routes.get("/model_resolver/progress/{download_id}")
|
||||
@json_api_endpoint("progress")
|
||||
async def get_download_progress(request):
|
||||
"""Get progress for a specific download."""
|
||||
download_id = request.match_info["download_id"]
|
||||
progress = get_progress(download_id)
|
||||
|
||||
if progress:
|
||||
return web.json_response(progress)
|
||||
else:
|
||||
return web.json_response(
|
||||
{"error": "Download not found"}, status=404
|
||||
)
|
||||
|
||||
@routes.get("/model_resolver/progress")
|
||||
@json_api_endpoint("progress")
|
||||
async def get_all_downloads_progress(request):
|
||||
"""Get progress for all downloads."""
|
||||
progress = get_all_progress()
|
||||
return web.json_response(progress)
|
||||
|
||||
@routes.post("/model_resolver/cancel/{download_id}")
|
||||
@json_api_endpoint("cancel", return_success_on_error=True)
|
||||
async def cancel_download_route(request):
|
||||
"""Cancel a download in progress."""
|
||||
download_id = request.match_info["download_id"]
|
||||
cancel_download(download_id)
|
||||
return web.json_response({"success": True})
|
||||
|
||||
@routes.post("/model_resolver/pause/{download_id}")
|
||||
@json_api_endpoint("pause", return_success_on_error=True)
|
||||
async def pause_download_route(request):
|
||||
"""Pause an aria2 download."""
|
||||
download_id = request.match_info["download_id"]
|
||||
result = pause_download(download_id)
|
||||
status = 200 if result.get("success") else 400
|
||||
return web.json_response(result, status=status)
|
||||
|
||||
@routes.post("/model_resolver/resume/{download_id}")
|
||||
@json_api_endpoint("resume", return_success_on_error=True)
|
||||
async def resume_download_route(request):
|
||||
"""Resume an aria2 download."""
|
||||
download_id = request.match_info["download_id"]
|
||||
result = resume_download(download_id)
|
||||
status = 200 if result.get("success") else 400
|
||||
return web.json_response(result, status=status)
|
||||
|
||||
@routes.post("/model_resolver/clear_completed_downloads")
|
||||
@json_api_endpoint("clear_completed_downloads", return_success_on_error=True)
|
||||
async def clear_completed_downloads_route(request):
|
||||
"""Clear completed and failed downloads from progress memory."""
|
||||
clear_completed_downloads()
|
||||
return web.json_response({"success": True})
|
||||
|
||||
@routes.get("/model_resolver/aria2/status")
|
||||
@routes.post("/model_resolver/aria2/status")
|
||||
@json_api_endpoint("aria2 status")
|
||||
async def aria2_status_route(request):
|
||||
"""Report aria2 availability using saved settings, with optional override via POST."""
|
||||
settings = await get_override_settings_from_request(request)
|
||||
return web.json_response(await asyncio.to_thread(get_aria2_status, settings))
|
||||
|
||||
|
||||
@routes.post("/model_resolver/aria2/start")
|
||||
@json_api_endpoint("aria2 start", return_success_on_error=True)
|
||||
async def aria2_start_route(request):
|
||||
"""Start the aria2 daemon without starting a download."""
|
||||
settings = await get_override_settings_from_request(request)
|
||||
result = await asyncio.to_thread(start_aria2_daemon, settings)
|
||||
status = 200 if result.get("success") else 400
|
||||
return web.json_response(result, status=status)
|
||||
|
||||
@routes.get("/model_resolver/aria2/stop")
|
||||
@routes.post("/model_resolver/aria2/stop")
|
||||
@json_api_endpoint("aria2 stop", return_success_on_error=True)
|
||||
async def aria2_stop_route(request):
|
||||
"""Stop the aria2 daemon started by Model Resolver."""
|
||||
result = await asyncio.to_thread(stop_aria2_daemon)
|
||||
status = 200 if result.get("success") else 400
|
||||
return web.json_response(result, status=status)
|
||||
|
||||
@routes.post("/model_resolver/aria2/install")
|
||||
@json_api_endpoint("aria2 install")
|
||||
async def aria2_install_route(request):
|
||||
"""Download official aria2c binary and save its path."""
|
||||
try:
|
||||
payload = await request.json()
|
||||
except Exception:
|
||||
payload = {}
|
||||
force = False
|
||||
if isinstance(payload, dict):
|
||||
force = to_bool(payload.get("force"), False)
|
||||
try:
|
||||
install_result = await asyncio.to_thread(install_aria2_engine, force)
|
||||
except Aria2InstallError as exc:
|
||||
self.logger.warning(f"Model Resolver aria2 install failed: {exc}")
|
||||
return web.json_response({"success": False, "error": str(exc)}, status=500)
|
||||
settings = await asyncio.to_thread(
|
||||
save_resolver_settings,
|
||||
{
|
||||
"aria2c_path": install_result.get("aria2c_path", ""),
|
||||
"download_backend": "aria2",
|
||||
},
|
||||
)
|
||||
install_result["settings"] = settings
|
||||
return web.json_response(install_result)
|
||||
@@ -0,0 +1,199 @@
|
||||
"""Metadata route registration."""
|
||||
|
||||
def register_metadata_routes(context):
|
||||
asyncio = context.get('asyncio')
|
||||
audit_metadata_sizes = context.get('audit_metadata_sizes')
|
||||
build_missing_local_metadata = context.get('build_missing_local_metadata')
|
||||
cancel_progress_response = context.get('cancel_progress_response')
|
||||
get_metadata_build_capabilities = context.get('get_metadata_build_capabilities')
|
||||
get_model_files = context.get('get_model_files')
|
||||
get_progress_response = context.get('get_progress_response')
|
||||
invalidate_local_hash_match_cache = context.get('invalidate_local_hash_match_cache')
|
||||
json_api_endpoint = context.get('json_api_endpoint')
|
||||
normalize_metadata_build_mode = context.get('normalize_metadata_build_mode')
|
||||
routes = context.get('routes')
|
||||
run_in_background_thread = context.get('run_in_background_thread')
|
||||
self = context.get('self')
|
||||
to_bool = context.get('to_bool')
|
||||
to_int = context.get('to_int')
|
||||
web = context.get('web')
|
||||
|
||||
@routes.get("/model_resolver/models")
|
||||
@json_api_endpoint("get_models")
|
||||
async def get_models(request):
|
||||
"""Get list of all available models."""
|
||||
force_rescan = to_bool(
|
||||
request.query.get("force") or request.query.get("force_rescan"),
|
||||
False
|
||||
)
|
||||
if force_rescan:
|
||||
invalidate_local_hash_match_cache()
|
||||
models = get_model_files(force_rescan=force_rescan)
|
||||
return web.json_response(models)
|
||||
|
||||
@routes.post("/model_resolver/metadata-size-audit")
|
||||
@json_api_endpoint("metadata-size-audit")
|
||||
async def metadata_size_audit_route(request):
|
||||
"""Check local model .metadata.json sidecars for stale size values."""
|
||||
payload = {}
|
||||
try:
|
||||
if request.can_read_body:
|
||||
payload = await request.json()
|
||||
except Exception:
|
||||
payload = {}
|
||||
if not isinstance(payload, dict):
|
||||
payload = {}
|
||||
|
||||
force_rescan = to_bool(payload.get("force_rescan"), True)
|
||||
result = await asyncio.to_thread(
|
||||
audit_metadata_sizes,
|
||||
force_rescan=force_rescan,
|
||||
worker_count=payload.get("worker_count"),
|
||||
batch_size=payload.get("batch_size"),
|
||||
)
|
||||
return web.json_response(result)
|
||||
|
||||
@routes.get("/model_resolver/metadata-build/capabilities")
|
||||
@json_api_endpoint("metadata-build-capabilities")
|
||||
async def metadata_build_capabilities_route(request):
|
||||
"""Return local CPU and worker limits for the metadata builder."""
|
||||
return web.json_response(get_metadata_build_capabilities())
|
||||
|
||||
@routes.post("/model_resolver/metadata-build/start")
|
||||
@json_api_endpoint("metadata-build-start")
|
||||
async def metadata_build_start_route(request):
|
||||
"""Create missing local metadata sidecars and fill missing SHA256 values."""
|
||||
import uuid
|
||||
|
||||
payload = {}
|
||||
try:
|
||||
if request.can_read_body:
|
||||
payload = await request.json()
|
||||
except Exception:
|
||||
payload = {}
|
||||
if not isinstance(payload, dict):
|
||||
payload = {}
|
||||
|
||||
force_rescan = to_bool(payload.get("force_rescan"), True)
|
||||
worker_count = to_int(payload.get("worker_count"), 0)
|
||||
metadata_mode = normalize_metadata_build_mode(
|
||||
payload.get("metadata_mode")
|
||||
)
|
||||
self.metadata_builder_progress.cleanup()
|
||||
progress_id = f"metadata_build_{uuid.uuid4().hex}"
|
||||
self.metadata_builder_progress.update(
|
||||
progress_id,
|
||||
status="queued",
|
||||
stage="queued",
|
||||
message="Preparing local metadata build...",
|
||||
percent=0,
|
||||
current=0,
|
||||
total=0,
|
||||
requested_worker_count=worker_count,
|
||||
metadata_mode=metadata_mode,
|
||||
)
|
||||
|
||||
def update_metadata_build_progress(progress_payload):
|
||||
self._update_metadata_build_progress(
|
||||
progress_id, progress_payload
|
||||
)
|
||||
|
||||
def is_metadata_build_cancelled():
|
||||
return self.metadata_builder_progress.is_cancelled(progress_id)
|
||||
|
||||
def run_metadata_build_task():
|
||||
result = build_missing_local_metadata(
|
||||
force_rescan=force_rescan,
|
||||
worker_count=worker_count,
|
||||
metadata_mode=metadata_mode,
|
||||
progress_callback=update_metadata_build_progress,
|
||||
is_cancelled=is_metadata_build_cancelled,
|
||||
)
|
||||
return result
|
||||
|
||||
def on_success(result):
|
||||
self.metadata_builder_progress.update(
|
||||
progress_id,
|
||||
status="done",
|
||||
stage="done",
|
||||
message="Local metadata build completed.",
|
||||
percent=100,
|
||||
active_models=[],
|
||||
active_worker_count=0,
|
||||
current_model="",
|
||||
current_path="",
|
||||
bytes_read=0,
|
||||
total_bytes=0,
|
||||
result=result,
|
||||
**result,
|
||||
)
|
||||
|
||||
def on_cancel(result=None):
|
||||
self.metadata_builder_progress.update(
|
||||
progress_id,
|
||||
status="cancelled",
|
||||
stage="cancelled",
|
||||
message="Metadata build cancelled",
|
||||
percent=100,
|
||||
active_models=[],
|
||||
active_worker_count=0,
|
||||
current_model="",
|
||||
current_path="",
|
||||
bytes_read=0,
|
||||
total_bytes=0,
|
||||
result=result or {},
|
||||
**(result or {}),
|
||||
)
|
||||
|
||||
def on_error(exc):
|
||||
self.metadata_builder_progress.update(
|
||||
progress_id,
|
||||
status="error",
|
||||
stage="error",
|
||||
message=str(exc) or "Metadata build failed",
|
||||
percent=100,
|
||||
active_models=[],
|
||||
active_worker_count=0,
|
||||
current_model="",
|
||||
current_path="",
|
||||
bytes_read=0,
|
||||
total_bytes=0,
|
||||
error=str(exc) or "Metadata build failed",
|
||||
)
|
||||
|
||||
run_in_background_thread(
|
||||
self.metadata_builder_progress,
|
||||
progress_id,
|
||||
run_metadata_build_task,
|
||||
on_success,
|
||||
on_cancel,
|
||||
on_error,
|
||||
error_log_msg="Metadata build failed",
|
||||
)
|
||||
return web.json_response(
|
||||
{
|
||||
"success": True,
|
||||
"progress_id": progress_id,
|
||||
"metadata_mode": metadata_mode,
|
||||
}
|
||||
)
|
||||
|
||||
@routes.get("/model_resolver/metadata-build/progress/{progress_id}")
|
||||
@json_api_endpoint("metadata-build-progress")
|
||||
async def metadata_build_progress_route(request):
|
||||
"""Return progress for local metadata sidecar creation."""
|
||||
return get_progress_response(
|
||||
self.metadata_builder_progress,
|
||||
request,
|
||||
not_found_status=404,
|
||||
)
|
||||
|
||||
@routes.post("/model_resolver/metadata-build/cancel/{progress_id}")
|
||||
@json_api_endpoint("metadata-build-cancel")
|
||||
async def metadata_build_cancel_route(request):
|
||||
"""Cancel local metadata sidecar creation."""
|
||||
return cancel_progress_response(
|
||||
self.metadata_builder_progress,
|
||||
request,
|
||||
cancel_message="Stopping metadata build...",
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
+23
-4495
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -79,6 +79,77 @@ def _find_bound_extension(handler):
|
||||
|
||||
class TestRefactoringTargets(unittest.IsolatedAsyncioTestCase):
|
||||
|
||||
ROUTE_GROUPS = {
|
||||
"workflow": {
|
||||
("POST", "/model_resolver/analyze"),
|
||||
("GET", "/model_resolver/analyze-progress/{analysis_id}"),
|
||||
("POST", "/model_resolver/resolve"),
|
||||
("POST", "/model_resolver/local-matches"),
|
||||
("POST", "/model_resolver/local-model-hashes"),
|
||||
("GET", "/model_resolver/model-preview"),
|
||||
("POST", "/model_resolver/workflow-model-hashes"),
|
||||
("POST", "/model_resolver/local-matches-by-hash"),
|
||||
("POST", "/model_resolver/open-containing-folder"),
|
||||
("POST", "/model_resolver/calculate-file-hash"),
|
||||
("POST", "/model_resolver/calculate-file-hash/start"),
|
||||
("GET", "/model_resolver/calculate-file-hash/progress/{progress_id}"),
|
||||
("POST", "/model_resolver/calculate-file-hash/cancel/{progress_id}"),
|
||||
},
|
||||
"metadata": {
|
||||
("GET", "/model_resolver/models"),
|
||||
("POST", "/model_resolver/metadata-size-audit"),
|
||||
("GET", "/model_resolver/metadata-build/capabilities"),
|
||||
("POST", "/model_resolver/metadata-build/start"),
|
||||
("GET", "/model_resolver/metadata-build/progress/{progress_id}"),
|
||||
("POST", "/model_resolver/metadata-build/cancel/{progress_id}"),
|
||||
},
|
||||
"loaded": {
|
||||
("POST", "/model_resolver/loaded"),
|
||||
("GET", "/model_resolver/loaded-progress/{loaded_id}"),
|
||||
},
|
||||
"model_info": {
|
||||
("POST", "/model_resolver/civitai-search"),
|
||||
("POST", "/model_resolver/custom-url"),
|
||||
("POST", "/model_resolver/model-details"),
|
||||
},
|
||||
"search": {
|
||||
("GET", "/model_resolver/search-progress/{progress_id}"),
|
||||
("POST", "/model_resolver/search-cancel/{progress_id}"),
|
||||
("POST", "/model_resolver/search"),
|
||||
("POST", "/model_resolver/clear-search-cache"),
|
||||
("POST", "/model_resolver/civitai/session-token/check"),
|
||||
("POST", "/model_resolver/civitai/api-key/check"),
|
||||
("POST", "/model_resolver/huggingface/token/check"),
|
||||
("POST", "/model_resolver/brave/api-key/check"),
|
||||
("GET", "/model_resolver/huggingface/author-index/status"),
|
||||
("POST", "/model_resolver/huggingface/author-index/refresh"),
|
||||
("GET", "/model_resolver/model-list/status"),
|
||||
("POST", "/model_resolver/model-list/update"),
|
||||
},
|
||||
"downloads": {
|
||||
("POST", "/model_resolver/download"),
|
||||
("GET", "/model_resolver/progress/{download_id}"),
|
||||
("GET", "/model_resolver/progress"),
|
||||
("POST", "/model_resolver/cancel/{download_id}"),
|
||||
("POST", "/model_resolver/pause/{download_id}"),
|
||||
("POST", "/model_resolver/resume/{download_id}"),
|
||||
("POST", "/model_resolver/clear_completed_downloads"),
|
||||
("GET", "/model_resolver/aria2/status"),
|
||||
("POST", "/model_resolver/aria2/status"),
|
||||
("POST", "/model_resolver/aria2/start"),
|
||||
("GET", "/model_resolver/aria2/stop"),
|
||||
("POST", "/model_resolver/aria2/stop"),
|
||||
("POST", "/model_resolver/aria2/install"),
|
||||
},
|
||||
"directories": {
|
||||
("GET", "/model_resolver/directories"),
|
||||
("GET", "/model_resolver/root-directories"),
|
||||
("GET", "/model_resolver/path-template-suggestions"),
|
||||
("GET", "/model_resolver/capabilities"),
|
||||
("GET", "/model_resolver/subfolders/{category}"),
|
||||
},
|
||||
}
|
||||
|
||||
def test_all_api_routes_are_registered(self):
|
||||
expected_routes = {
|
||||
("GET", "/model_resolver/base-models"),
|
||||
@@ -145,6 +216,26 @@ class TestRefactoringTargets(unittest.IsolatedAsyncioTestCase):
|
||||
}
|
||||
self.assertEqual(expected_routes, set(routes_registered))
|
||||
|
||||
def test_route_groups_cover_the_registered_feature_routes(self):
|
||||
grouped_routes = set().union(*self.ROUTE_GROUPS.values())
|
||||
self.assertEqual(
|
||||
grouped_routes,
|
||||
set(routes_registered)
|
||||
- {
|
||||
("GET", "/model_resolver/base-models"),
|
||||
("GET", "/model_resolver/base-models/status"),
|
||||
("POST", "/model_resolver/base-models/update"),
|
||||
("GET", "/model_resolver/version"),
|
||||
("GET", "/model_resolver/logs/backend/export"),
|
||||
("GET", "/model_resolver/settings"),
|
||||
("POST", "/model_resolver/settings"),
|
||||
},
|
||||
)
|
||||
self.assertEqual(
|
||||
sum(len(routes) for routes in self.ROUTE_GROUPS.values()),
|
||||
len(grouped_routes),
|
||||
)
|
||||
|
||||
def test_project_version_helpers_preserve_version_comparison_behavior(self):
|
||||
self.assertEqual(
|
||||
node_mod._extract_project_version('version = "1.2.3"'),
|
||||
@@ -198,6 +289,28 @@ class TestRefactoringTargets(unittest.IsolatedAsyncioTestCase):
|
||||
payload = mock_json_response.call_args.args[0]
|
||||
self.assertIn("settings", payload)
|
||||
self.assertIn("schema", payload)
|
||||
|
||||
async def test_metadata_capabilities_route_returns_capabilities(self):
|
||||
get_handler = routes_registered[("GET", "/model_resolver/metadata-build/capabilities")]
|
||||
with patch("aiohttp.web.json_response") as mock_json_response:
|
||||
await get_handler(MagicMock())
|
||||
|
||||
payload = mock_json_response.call_args.args[0]
|
||||
self.assertIn("metadata_modes", payload)
|
||||
self.assertIn("default_metadata_mode", payload)
|
||||
|
||||
async def test_custom_url_route_rejects_missing_url(self):
|
||||
post_handler = routes_registered[("POST", "/model_resolver/custom-url")]
|
||||
request = AsyncMock()
|
||||
request.json.return_value = {}
|
||||
|
||||
with patch("aiohttp.web.json_response") as mock_json_response:
|
||||
await post_handler(request)
|
||||
|
||||
mock_json_response.assert_called_once_with(
|
||||
{"error": "URL is required"},
|
||||
status=400,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
|
||||
Reference in New Issue
Block a user