refactor(routes): split large route modules into focused submodules
Break monolithic model_info, search, and workflow route files into smaller modules by responsibility: civitai_search, custom_url, model_details, search_support, source_search, workflow_analysis, hashes, and loaded_models. Update registry imports to preserve the existing route registrations and API surface. Add partition coverage tests to ensure no route endpoints are dropped or duplicated.
This commit is contained in:
Dariusz L
@@ -0,0 +1,899 @@
|
||||
"""CivitAI and exact model metadata route registration."""
|
||||
|
||||
from .context import RouteContext
|
||||
|
||||
|
||||
def register_civitai_search_routes(context: RouteContext):
|
||||
download_available = context.get('download_available')
|
||||
extract_sha256_from_metadata = context.get('extract_sha256_from_metadata')
|
||||
find_external_metadata_sidecar_path = context.get('find_external_metadata_sidecar_path')
|
||||
find_local_file_path = context.get('find_local_file_path')
|
||||
get_existing_model_preview_path = context.get('get_existing_model_preview_path')
|
||||
get_filename_from_path = context.get('get_filename_from_path')
|
||||
get_model_resolver_sidecar_path = context.get('get_model_resolver_sidecar_path')
|
||||
is_path_in_configured_model_roots = context.get('is_path_in_configured_model_roots')
|
||||
json_api_endpoint = context.get('json_api_endpoint')
|
||||
looks_like_model_file = context.get('looks_like_model_file')
|
||||
normalize_category_to_model_type = context.get('normalize_category_to_model_type')
|
||||
normalize_sha256 = context.get('normalize_sha256')
|
||||
read_json_safe = context.get('read_json_safe')
|
||||
request_public_url = context.get('request_public_url')
|
||||
resolve_civarchive_by_hash = context.get('resolve_civarchive_by_hash')
|
||||
routes = context.get('routes')
|
||||
search_huggingface_for_file = context.get('search_huggingface_for_file')
|
||||
self = context.get('self')
|
||||
to_bool = context.get('to_bool')
|
||||
web = context.get('web')
|
||||
write_model_resolver_metadata = context.get('write_model_resolver_metadata')
|
||||
|
||||
@routes.post("/model_resolver/civitai-search")
|
||||
@json_api_endpoint("civitai-search")
|
||||
async def civitai_search(request):
|
||||
"""Fetch model metadata from trusted exact-match sources."""
|
||||
data = await request.json()
|
||||
filename = data.get("filename", "")
|
||||
category = data.get("category", "")
|
||||
resolved_path = data.get("resolved_path", "")
|
||||
local_only = to_bool(data.get("local_only"), False)
|
||||
force_refresh = to_bool(
|
||||
data.get("force_refresh") or data.get("force"), False
|
||||
)
|
||||
provided_hash = normalize_sha256(
|
||||
data.get("sha256")
|
||||
or data.get("hash")
|
||||
or data.get("file_hash")
|
||||
or ""
|
||||
)
|
||||
hf_token = data.get("hf_token", "")
|
||||
brave_search_api_key = data.get("brave_search_api_key", "")
|
||||
hf_use_brave_fallback = to_bool(
|
||||
data.get("hf_use_brave_fallback", True),
|
||||
True,
|
||||
)
|
||||
|
||||
if not filename:
|
||||
return web.json_response(
|
||||
{"error": "Filename is required"}, status=400
|
||||
)
|
||||
|
||||
# Clean filename for display
|
||||
import os as _os
|
||||
|
||||
clean_name = _os.path.splitext(filename)[0]
|
||||
|
||||
local_metadata_fields = (
|
||||
"author",
|
||||
"creator",
|
||||
"license",
|
||||
"usage_hint",
|
||||
"usage_tips",
|
||||
"preview_url",
|
||||
"metadata_source",
|
||||
"from_safetensors_header",
|
||||
"local_metadata_available",
|
||||
"header_metadata_keys",
|
||||
"metadata_summary",
|
||||
"base_model_raw",
|
||||
"sha256_source",
|
||||
"hashes",
|
||||
"hash_status",
|
||||
)
|
||||
|
||||
def add_local_metadata_fields(payload, result):
|
||||
if not isinstance(payload, dict) or not isinstance(result, dict):
|
||||
return payload
|
||||
for key in local_metadata_fields:
|
||||
value = result.get(key)
|
||||
if value in (None, "", [], {}):
|
||||
continue
|
||||
payload[key] = value
|
||||
return payload
|
||||
|
||||
# Get the file path to hash
|
||||
file_path = resolved_path if resolved_path else None
|
||||
file_location = ""
|
||||
|
||||
if not file_path:
|
||||
file_path = find_local_file_path(filename, category)
|
||||
|
||||
if file_path:
|
||||
file_path = _os.path.realpath(
|
||||
_os.path.abspath(_os.path.normpath(str(file_path)))
|
||||
)
|
||||
if not is_path_in_configured_model_roots(file_path):
|
||||
if resolved_path:
|
||||
return web.json_response(
|
||||
{
|
||||
"error": (
|
||||
"resolved_path is outside configured model directories"
|
||||
)
|
||||
},
|
||||
status=403,
|
||||
)
|
||||
file_path = None
|
||||
elif not _os.path.isfile(file_path):
|
||||
if resolved_path:
|
||||
return web.json_response(
|
||||
{"error": "resolved_path is not an existing model file"},
|
||||
status=404,
|
||||
)
|
||||
file_path = None
|
||||
|
||||
if file_path and _os.path.exists(file_path):
|
||||
file_location = _os.path.dirname(file_path).replace("\\", "/")
|
||||
if file_location and not file_location.endswith("/"):
|
||||
file_location += "/"
|
||||
|
||||
external_model_description = ""
|
||||
if file_path:
|
||||
external_metadata_path = find_external_metadata_sidecar_path(
|
||||
file_path
|
||||
)
|
||||
external_metadata = (
|
||||
read_json_safe(external_metadata_path, {})
|
||||
if external_metadata_path
|
||||
else {}
|
||||
)
|
||||
if isinstance(external_metadata, dict):
|
||||
external_model_description = str(
|
||||
external_metadata.get("modelDescription")
|
||||
or external_metadata.get("model_description")
|
||||
or external_metadata.get("description")
|
||||
or ""
|
||||
).strip()
|
||||
|
||||
def infer_model_type_from_category(value):
|
||||
return normalize_category_to_model_type(value)
|
||||
|
||||
def split_result_descriptions(result, preferred_model_description=""):
|
||||
result = result if isinstance(result, dict) else {}
|
||||
explicit_model_description = (
|
||||
result.get("model_description")
|
||||
or result.get("modelDescription")
|
||||
or ""
|
||||
)
|
||||
generic_description = result.get("description") or ""
|
||||
model_description = (
|
||||
preferred_model_description
|
||||
or explicit_model_description
|
||||
or generic_description
|
||||
)
|
||||
version_description = (
|
||||
result.get("version_description")
|
||||
or result.get("versionDescription")
|
||||
or ""
|
||||
)
|
||||
if (
|
||||
not version_description
|
||||
and explicit_model_description
|
||||
and generic_description
|
||||
and generic_description != explicit_model_description
|
||||
):
|
||||
version_description = generic_description
|
||||
if version_description == model_description:
|
||||
version_description = ""
|
||||
return model_description, version_description
|
||||
|
||||
def build_info_response(
|
||||
result=None,
|
||||
*,
|
||||
metadata_path="",
|
||||
metadata_saved=False,
|
||||
civitai_checked=False,
|
||||
local_payload=False,
|
||||
):
|
||||
result = result or {}
|
||||
model_description, version_description = (
|
||||
split_result_descriptions(
|
||||
result,
|
||||
external_model_description,
|
||||
)
|
||||
)
|
||||
size_value = result.get("size")
|
||||
if not size_value and file_path and _os.path.exists(file_path):
|
||||
try:
|
||||
size_value = _os.path.getsize(file_path)
|
||||
except Exception:
|
||||
size_value = None
|
||||
|
||||
model_type = (
|
||||
result.get("model_type")
|
||||
or result.get("type")
|
||||
or infer_model_type_from_category(category)
|
||||
)
|
||||
|
||||
response = {
|
||||
"filename": result.get("filename") or filename,
|
||||
"category": category,
|
||||
"file_path": result.get("file_path") or file_path or "",
|
||||
"resolved_path": result.get("resolved_path") or file_path or "",
|
||||
"metadata_path": metadata_path
|
||||
or result.get("metadata_path")
|
||||
or "",
|
||||
"metadata_saved": bool(metadata_saved),
|
||||
"location": result.get("location") or file_location,
|
||||
"source": result.get("source") or "",
|
||||
"details_source": result.get("details_source")
|
||||
or result.get("source")
|
||||
or "",
|
||||
"url": result.get("url"),
|
||||
"version_url": result.get("version_url"),
|
||||
"download_url": result.get("download_url"),
|
||||
"platform_url": result.get("platform_url"),
|
||||
"repo_id": result.get("repo_id"),
|
||||
"path": result.get("path"),
|
||||
"model_id": result.get("model_id"),
|
||||
"model_name": result.get("model_name") or clean_name,
|
||||
"model_type": model_type,
|
||||
"version_id": result.get("version_id"),
|
||||
"version_name": result.get("version_name", ""),
|
||||
"sha256": result.get("sha256") or provided_hash,
|
||||
"size": size_value,
|
||||
"base_model": result.get("base_model"),
|
||||
"base_model_source": result.get("base_model_source"),
|
||||
"base_model_inferred": bool(result.get("base_model_inferred")),
|
||||
"tags": result.get("tags", []),
|
||||
"trained_words": result.get("trained_words", []),
|
||||
"images": result.get("images", []),
|
||||
"clip_skip": result.get("clip_skip"),
|
||||
"description": model_description,
|
||||
"model_description": model_description,
|
||||
"version_description": version_description,
|
||||
"from_metadata": bool(result.get("from_metadata")),
|
||||
"local_only": bool(local_payload),
|
||||
"metadata_checked": bool(civitai_checked),
|
||||
"civitai_checked": bool(civitai_checked),
|
||||
}
|
||||
return add_local_metadata_fields(response, result)
|
||||
def extract_result_sha256(result):
|
||||
return extract_sha256_from_metadata(result)
|
||||
|
||||
def result_filename_matches(result):
|
||||
if not isinstance(result, dict):
|
||||
return False
|
||||
expected = get_filename_from_path(filename).lower()
|
||||
candidates = [
|
||||
result.get("filename"),
|
||||
result.get("path"),
|
||||
result.get("file_path"),
|
||||
]
|
||||
for candidate in candidates:
|
||||
basename = get_filename_from_path(candidate).lower()
|
||||
if basename and basename == expected:
|
||||
return True
|
||||
return False
|
||||
|
||||
def result_hash_matches(result, *, require_filename=False):
|
||||
if not provided_hash or not isinstance(result, dict):
|
||||
return False
|
||||
result_hash = extract_result_sha256(result)
|
||||
if result_hash != provided_hash:
|
||||
return False
|
||||
return not require_filename or result_filename_matches(result)
|
||||
|
||||
def huggingface_page_url(result):
|
||||
try:
|
||||
from urllib.parse import quote as _quote
|
||||
except Exception:
|
||||
_quote = None
|
||||
repo_id = str(
|
||||
result.get("repo_id") or result.get("repo") or ""
|
||||
).strip()
|
||||
hf_path = str(
|
||||
result.get("path") or result.get("filename") or ""
|
||||
).strip()
|
||||
if not repo_id or not hf_path:
|
||||
return ""
|
||||
if _quote:
|
||||
hf_path = _quote(hf_path.replace("\\", "/"), safe="/")
|
||||
return f"https://huggingface.co/{repo_id}/blob/main/{hf_path}"
|
||||
|
||||
def prepare_remote_result(result, source_name):
|
||||
result = dict(result or {})
|
||||
source_name = str(source_name or result.get("source") or "").lower()
|
||||
if source_name == "huggingface":
|
||||
download_url = result.get("download_url") or result.get("url")
|
||||
page_url = result.get("page_url") or huggingface_page_url(result)
|
||||
if page_url:
|
||||
result["url"] = page_url
|
||||
result["version_url"] = page_url
|
||||
result["download_url"] = download_url
|
||||
result["model_name"] = (
|
||||
result.get("model_name")
|
||||
or result.get("name")
|
||||
or result.get("repo_id")
|
||||
or clean_name
|
||||
)
|
||||
result["model_type"] = (
|
||||
result.get("model_type")
|
||||
or infer_model_type_from_category(category)
|
||||
)
|
||||
else:
|
||||
result["model_name"] = (
|
||||
result.get("model_name")
|
||||
or result.get("name")
|
||||
or clean_name
|
||||
)
|
||||
if result.get("url") and not result.get("version_url"):
|
||||
result["version_url"] = result.get("url")
|
||||
|
||||
result["source"] = source_name or result.get("source") or ""
|
||||
result["details_source"] = (
|
||||
result.get("details_source")
|
||||
or source_name
|
||||
or result.get("source")
|
||||
or ""
|
||||
)
|
||||
result["file_path"] = file_path or result.get("file_path") or ""
|
||||
result["resolved_path"] = (
|
||||
file_path or result.get("resolved_path") or ""
|
||||
)
|
||||
result["location"] = result.get("location") or file_location
|
||||
result["filename"] = result.get("filename") or filename
|
||||
result["sha256"] = (
|
||||
extract_result_sha256(result)
|
||||
or provided_hash
|
||||
or result.get("sha256")
|
||||
)
|
||||
if not result.get("size") and file_path and _os.path.exists(file_path):
|
||||
try:
|
||||
result["size"] = _os.path.getsize(file_path)
|
||||
except Exception:
|
||||
pass
|
||||
return result
|
||||
|
||||
def result_has_remote_identity(result):
|
||||
if not isinstance(result, dict):
|
||||
return False
|
||||
return bool(
|
||||
result.get("url")
|
||||
or result.get("version_url")
|
||||
or result.get("download_url")
|
||||
or result.get("platform_url")
|
||||
or result.get("repo_id")
|
||||
or result.get("model_id")
|
||||
or result.get("version_id")
|
||||
or result.get("from_metadata")
|
||||
)
|
||||
|
||||
def remote_link_is_marked_dead(item):
|
||||
if not isinstance(item, dict):
|
||||
return False
|
||||
status = str(item.get("status") or "").lower()
|
||||
return bool(
|
||||
item.get("deletedAt")
|
||||
or item.get("deleted_at")
|
||||
or item.get("is_dead")
|
||||
or item.get("isDead")
|
||||
or item.get("likelyDead")
|
||||
or item.get("likely_dead")
|
||||
or item.get("dead")
|
||||
or status in {"dead", "deleted", "unavailable", "missing"}
|
||||
)
|
||||
|
||||
def collect_result_download_urls(result):
|
||||
urls = []
|
||||
if remote_link_is_marked_dead(result):
|
||||
return urls
|
||||
expected_filename = result.get("filename") or filename
|
||||
dead_urls = set()
|
||||
mirrors = result.get("mirrors") or []
|
||||
if not isinstance(mirrors, list):
|
||||
mirrors = [mirrors]
|
||||
for mirror in mirrors:
|
||||
if not isinstance(mirror, dict):
|
||||
continue
|
||||
if remote_link_is_marked_dead(mirror):
|
||||
dead_url = str(mirror.get("url") or "").strip()
|
||||
if dead_url.startswith(("http://", "https://")):
|
||||
dead_urls.add(dead_url)
|
||||
continue
|
||||
url = str(mirror.get("url") or "").strip()
|
||||
mirror_filename = (
|
||||
mirror.get("filename")
|
||||
or mirror.get("name")
|
||||
or expected_filename
|
||||
)
|
||||
if (
|
||||
looks_like_model_file(url, mirror_filename)
|
||||
and url not in urls
|
||||
):
|
||||
urls.append(url)
|
||||
raw_urls = result.get("download_urls") or []
|
||||
if not isinstance(raw_urls, list):
|
||||
raw_urls = [raw_urls]
|
||||
for raw_url in raw_urls:
|
||||
url = str(raw_url or "").strip()
|
||||
if (
|
||||
looks_like_model_file(url, expected_filename)
|
||||
and url not in dead_urls
|
||||
and url not in urls
|
||||
):
|
||||
urls.append(url)
|
||||
for key in ("download_url", "downloadUrl"):
|
||||
url = str(result.get(key) or "").strip()
|
||||
if (
|
||||
looks_like_model_file(url, expected_filename)
|
||||
and url not in dead_urls
|
||||
and url not in urls
|
||||
):
|
||||
urls.append(url)
|
||||
return urls
|
||||
|
||||
def remote_download_url_is_alive(url):
|
||||
headers = {
|
||||
"User-Agent": "ComfyUI-Model-Resolver/1.0",
|
||||
"Accept": "*/*",
|
||||
}
|
||||
try:
|
||||
response, _final_url, _final_headers = request_public_url(
|
||||
"HEAD",
|
||||
url,
|
||||
headers=headers,
|
||||
timeout=8,
|
||||
)
|
||||
try:
|
||||
if response.status_code < 400:
|
||||
return True
|
||||
if response.status_code in {401, 403, 404, 410}:
|
||||
return False
|
||||
finally:
|
||||
response.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
try:
|
||||
response, _final_url, _final_headers = request_public_url(
|
||||
"GET",
|
||||
url,
|
||||
headers={**headers, "Range": "bytes=0-0"},
|
||||
stream=True,
|
||||
timeout=8,
|
||||
)
|
||||
try:
|
||||
return response.status_code < 400
|
||||
finally:
|
||||
response.close()
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def civarchive_result_has_live_download(result):
|
||||
urls = collect_result_download_urls(result)
|
||||
if not urls:
|
||||
return False
|
||||
for url in urls[:3]:
|
||||
if remote_download_url_is_alive(url):
|
||||
result["download_url"] = url
|
||||
result["download_urls"] = [
|
||||
url,
|
||||
*[other for other in urls if other != url],
|
||||
]
|
||||
return True
|
||||
return False
|
||||
|
||||
def save_remote_metadata(result, source_name):
|
||||
metadata_path = result.get("metadata_path") or ""
|
||||
metadata_saved = False
|
||||
if (
|
||||
result.get("from_metadata")
|
||||
or not file_path
|
||||
or not _os.path.exists(file_path)
|
||||
):
|
||||
return metadata_path, metadata_saved
|
||||
try:
|
||||
model_description, version_description = (
|
||||
split_result_descriptions(result)
|
||||
)
|
||||
metadata_payload = {
|
||||
"source": source_name,
|
||||
"details_source": result.get("details_source")
|
||||
or source_name,
|
||||
"filename": filename,
|
||||
"category": category,
|
||||
"model_name": result.get("model_name", clean_name),
|
||||
"name": result.get("model_name", clean_name),
|
||||
"model_type": result.get("model_type", "")
|
||||
or result.get("type", ""),
|
||||
"type": result.get("model_type", "")
|
||||
or result.get("type", ""),
|
||||
"model_id": result.get("model_id"),
|
||||
"version_id": result.get("version_id"),
|
||||
"version_name": result.get("version_name", ""),
|
||||
"sha256": result.get("sha256") or provided_hash,
|
||||
"size": result.get("size"),
|
||||
"base_model": result.get("base_model"),
|
||||
"base_model_source": result.get("base_model_source"),
|
||||
"base_model_inferred": bool(result.get("base_model_inferred")),
|
||||
"tags": result.get("tags", []),
|
||||
"trained_words": result.get("trained_words", []),
|
||||
"images": result.get("images", []),
|
||||
"clip_skip": result.get("clip_skip"),
|
||||
"description": model_description,
|
||||
"model_description": model_description,
|
||||
"version_description": version_description,
|
||||
"download_url": result.get("download_url"),
|
||||
"source_url": result.get("version_url") or result.get("url"),
|
||||
"version_url": result.get("version_url") or result.get("url"),
|
||||
"model_url": result.get("url"),
|
||||
"url": result.get("version_url") or result.get("url"),
|
||||
"platform_url": result.get("platform_url"),
|
||||
"repo_id": result.get("repo_id"),
|
||||
"path": result.get("path"),
|
||||
"path_metadata": {
|
||||
"filename": filename,
|
||||
"category": category,
|
||||
"source": source_name,
|
||||
"model_id": result.get("model_id"),
|
||||
"version_id": result.get("version_id"),
|
||||
"repo_id": result.get("repo_id"),
|
||||
"path": result.get("path"),
|
||||
},
|
||||
}
|
||||
add_local_metadata_fields(metadata_payload, result)
|
||||
metadata_path = write_model_resolver_metadata(
|
||||
file_path,
|
||||
metadata_payload,
|
||||
category,
|
||||
result.get("version_url")
|
||||
or result.get("url")
|
||||
or result.get("platform_url")
|
||||
or result.get("download_url")
|
||||
or "",
|
||||
create_preview=True,
|
||||
) or ""
|
||||
metadata_saved = bool(metadata_path)
|
||||
if metadata_path:
|
||||
result["metadata_path"] = metadata_path
|
||||
except Exception as metadata_error:
|
||||
self.logger.warning(
|
||||
f"{source_name} metadata sidecar save failed: {metadata_error}"
|
||||
)
|
||||
return metadata_path, metadata_saved
|
||||
|
||||
if local_only:
|
||||
result = None
|
||||
metadata_path = ""
|
||||
metadata_saved = False
|
||||
if download_available and file_path and _os.path.exists(file_path):
|
||||
try:
|
||||
from ..sources.civitai import (
|
||||
get_model_info_for_file,
|
||||
)
|
||||
|
||||
result = get_model_info_for_file(
|
||||
file_path,
|
||||
local_only=True,
|
||||
)
|
||||
source_metadata_path = (
|
||||
result.get("metadata_path")
|
||||
if isinstance(result, dict)
|
||||
else ""
|
||||
)
|
||||
canonical_metadata_path = get_model_resolver_sidecar_path(
|
||||
file_path
|
||||
)
|
||||
if (
|
||||
isinstance(result, dict)
|
||||
and result.get("from_metadata")
|
||||
and source_metadata_path
|
||||
and canonical_metadata_path
|
||||
and _os.path.abspath(source_metadata_path)
|
||||
!= _os.path.abspath(canonical_metadata_path)
|
||||
and not _os.path.exists(canonical_metadata_path)
|
||||
):
|
||||
preview_url = result.get("preview_url") or ""
|
||||
if not preview_url:
|
||||
preview_path = get_existing_model_preview_path(file_path)
|
||||
if preview_path:
|
||||
preview_url = preview_path.replace("\\", "/")
|
||||
model_description, version_description = (
|
||||
split_result_descriptions(result)
|
||||
)
|
||||
metadata_payload = {
|
||||
"source": "metadata_import",
|
||||
"details_source": result.get("source") or "metadata",
|
||||
"filename": filename,
|
||||
"category": category,
|
||||
"model_name": result.get("model_name", clean_name),
|
||||
"name": result.get("model_name", clean_name),
|
||||
"model_type": result.get("model_type", ""),
|
||||
"type": result.get("model_type", ""),
|
||||
"model_id": result.get("model_id"),
|
||||
"version_id": result.get("version_id"),
|
||||
"version_name": result.get("version_name", ""),
|
||||
"sha256": result.get("sha256"),
|
||||
"size": result.get("size"),
|
||||
"base_model": result.get("base_model"),
|
||||
"base_model_source": result.get("base_model_source"),
|
||||
"base_model_inferred": bool(result.get("base_model_inferred")),
|
||||
"tags": result.get("tags", []),
|
||||
"trained_words": result.get("trained_words", []),
|
||||
"images": result.get("images", []),
|
||||
"clip_skip": result.get("clip_skip"),
|
||||
"description": model_description,
|
||||
"model_description": model_description,
|
||||
"version_description": version_description,
|
||||
"download_url": result.get("download_url"),
|
||||
"preview_url": preview_url,
|
||||
"source_url": result.get("version_url")
|
||||
or result.get("url"),
|
||||
"version_url": result.get("version_url")
|
||||
or result.get("url"),
|
||||
"model_url": result.get("url"),
|
||||
"url": result.get("version_url") or result.get("url"),
|
||||
"path_metadata": {
|
||||
"filename": filename,
|
||||
"category": category,
|
||||
"source": "metadata_import",
|
||||
"model_id": result.get("model_id"),
|
||||
"version_id": result.get("version_id"),
|
||||
"imported_from": source_metadata_path,
|
||||
},
|
||||
}
|
||||
add_local_metadata_fields(metadata_payload, result)
|
||||
metadata_path = write_model_resolver_metadata(
|
||||
file_path,
|
||||
metadata_payload,
|
||||
category,
|
||||
result.get("version_url")
|
||||
or result.get("url")
|
||||
or result.get("download_url")
|
||||
or "",
|
||||
create_preview=True,
|
||||
) or ""
|
||||
metadata_saved = bool(metadata_path)
|
||||
if metadata_path:
|
||||
result["metadata_path"] = metadata_path
|
||||
result["metadata_imported_from"] = source_metadata_path
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Local model info error: {e}")
|
||||
return web.json_response(
|
||||
build_info_response(
|
||||
result,
|
||||
metadata_path=metadata_path,
|
||||
metadata_saved=metadata_saved,
|
||||
local_payload=True,
|
||||
civitai_checked=False,
|
||||
)
|
||||
)
|
||||
|
||||
# Search remote metadata. CivitAI and CivArchive are hash-only;
|
||||
# HuggingFace name results must still confirm the exact SHA.
|
||||
if download_available and file_path and _os.path.exists(file_path):
|
||||
try:
|
||||
from ..sources.civitai import (
|
||||
get_model_info_by_hash,
|
||||
get_model_info_for_file,
|
||||
)
|
||||
|
||||
if provided_hash:
|
||||
result = get_model_info_by_hash(
|
||||
provided_hash,
|
||||
use_cache=not force_refresh,
|
||||
)
|
||||
if result:
|
||||
result["file_path"] = file_path
|
||||
result["resolved_path"] = file_path
|
||||
result["location"] = file_location
|
||||
if not result.get("size"):
|
||||
try:
|
||||
result["size"] = _os.path.getsize(file_path)
|
||||
except Exception:
|
||||
pass
|
||||
if not extract_result_sha256(result):
|
||||
result["sha256"] = provided_hash
|
||||
elif not force_refresh:
|
||||
result = get_model_info_for_file(file_path)
|
||||
else:
|
||||
result = None
|
||||
|
||||
if result and (
|
||||
result.get("url")
|
||||
or result.get("version_url")
|
||||
or result.get("from_metadata")
|
||||
or result.get("trained_words")
|
||||
or result.get("base_model_inferred")
|
||||
or result.get("local_metadata_available")
|
||||
or result.get("from_safetensors_header")
|
||||
) and (
|
||||
result.get("from_metadata")
|
||||
or not provided_hash
|
||||
or result_hash_matches(result)
|
||||
):
|
||||
result_source = (
|
||||
"local"
|
||||
if result.get("source") == "local"
|
||||
and not result_has_remote_identity(result)
|
||||
else "civitai"
|
||||
)
|
||||
result = prepare_remote_result(result, result_source)
|
||||
metadata_path, metadata_saved = save_remote_metadata(
|
||||
result,
|
||||
result_source,
|
||||
)
|
||||
|
||||
return web.json_response(
|
||||
build_info_response(
|
||||
result,
|
||||
metadata_path=metadata_path,
|
||||
metadata_saved=metadata_saved,
|
||||
civitai_checked=True,
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.warning(f"CivitAI search error: {e}")
|
||||
|
||||
if provided_hash:
|
||||
try:
|
||||
result = resolve_civarchive_by_hash(
|
||||
provided_hash,
|
||||
query=filename,
|
||||
exact_only=False,
|
||||
model_type=infer_model_type_from_category(category),
|
||||
)
|
||||
if result:
|
||||
if not extract_result_sha256(result):
|
||||
result["sha256"] = provided_hash
|
||||
result = prepare_remote_result(result, "civarchive")
|
||||
if result_hash_matches(result) and civarchive_result_has_live_download(result):
|
||||
metadata_path, metadata_saved = save_remote_metadata(
|
||||
result,
|
||||
"civarchive",
|
||||
)
|
||||
return web.json_response(
|
||||
build_info_response(
|
||||
result,
|
||||
metadata_path=metadata_path,
|
||||
metadata_saved=metadata_saved,
|
||||
civitai_checked=True,
|
||||
)
|
||||
)
|
||||
if result_hash_matches(result):
|
||||
self.logger.info(
|
||||
f"CivArchive metadata candidate rejected: "
|
||||
f"no live download link for {filename}"
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.warning(f"CivArchive hash lookup error: {e}")
|
||||
|
||||
if provided_hash:
|
||||
hf_attempts = [
|
||||
{
|
||||
"label": "HuggingFace",
|
||||
"use_api_search": True,
|
||||
"use_comfy_org_fallback": True,
|
||||
"use_brave_fallback": False,
|
||||
}
|
||||
]
|
||||
if hf_use_brave_fallback and brave_search_api_key:
|
||||
hf_attempts.append(
|
||||
{
|
||||
"label": "HuggingFace Brave",
|
||||
"use_api_search": False,
|
||||
"use_comfy_org_fallback": False,
|
||||
"use_brave_fallback": True,
|
||||
}
|
||||
)
|
||||
|
||||
for hf_attempt in hf_attempts:
|
||||
try:
|
||||
result = search_huggingface_for_file(
|
||||
filename,
|
||||
token=hf_token or None,
|
||||
exact_only=True,
|
||||
brave_api_key=brave_search_api_key or None,
|
||||
use_api_search=hf_attempt["use_api_search"],
|
||||
use_comfy_org_fallback=hf_attempt[
|
||||
"use_comfy_org_fallback"
|
||||
],
|
||||
use_brave_fallback=hf_attempt[
|
||||
"use_brave_fallback"
|
||||
],
|
||||
force_refresh=force_refresh,
|
||||
)
|
||||
if result and result_hash_matches(
|
||||
result,
|
||||
require_filename=True,
|
||||
):
|
||||
result = prepare_remote_result(
|
||||
result,
|
||||
"huggingface",
|
||||
)
|
||||
metadata_path, metadata_saved = save_remote_metadata(
|
||||
result,
|
||||
"huggingface",
|
||||
)
|
||||
return web.json_response(
|
||||
build_info_response(
|
||||
result,
|
||||
metadata_path=metadata_path,
|
||||
metadata_saved=metadata_saved,
|
||||
civitai_checked=True,
|
||||
)
|
||||
)
|
||||
if result:
|
||||
self.logger.info(
|
||||
f"{hf_attempt['label']} metadata candidate rejected: "
|
||||
f"filename/hash mismatch for {filename}"
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.warning(
|
||||
f"{hf_attempt['label']} metadata lookup error: {e}"
|
||||
)
|
||||
|
||||
# No remote result found. Keep any useful local safetensors
|
||||
# header metadata, including base-model fingerprints, model
|
||||
# titles, descriptions, tags, trigger words, and author data.
|
||||
local_fallback_result = None
|
||||
if file_path and _os.path.exists(file_path):
|
||||
try:
|
||||
from ..sources.civitai import get_model_info_for_file
|
||||
|
||||
local_candidate = get_model_info_for_file(
|
||||
file_path,
|
||||
local_only=True,
|
||||
)
|
||||
if (
|
||||
isinstance(local_candidate, dict)
|
||||
and (
|
||||
local_candidate.get("base_model")
|
||||
or local_candidate.get("local_metadata_available")
|
||||
or local_candidate.get("from_safetensors_header")
|
||||
)
|
||||
):
|
||||
local_fallback_result = local_candidate
|
||||
except Exception as local_fallback_error:
|
||||
self.logger.debug(
|
||||
f"Local base model inference fallback failed: {local_fallback_error}"
|
||||
)
|
||||
|
||||
response = build_info_response(
|
||||
local_fallback_result,
|
||||
civitai_checked=True,
|
||||
local_payload=bool(local_fallback_result),
|
||||
)
|
||||
response["url"] = None
|
||||
if force_refresh and file_path and _os.path.exists(file_path):
|
||||
try:
|
||||
metadata_payload = {
|
||||
"filename": filename,
|
||||
"category": category,
|
||||
"model_name": response.get("model_name") or clean_name,
|
||||
"name": response.get("model_name") or clean_name,
|
||||
"model_type": response.get("model_type")
|
||||
or infer_model_type_from_category(category),
|
||||
"sha256": response.get("sha256") or provided_hash,
|
||||
"size": response.get("size") or _os.path.getsize(file_path),
|
||||
"base_model": response.get("base_model"),
|
||||
"base_model_source": response.get("base_model_source"),
|
||||
"base_model_inferred": bool(response.get("base_model_inferred")),
|
||||
"description": response.get("model_description")
|
||||
or response.get("description")
|
||||
or "",
|
||||
"model_description": response.get("model_description")
|
||||
or response.get("description")
|
||||
or "",
|
||||
"version_description": response.get(
|
||||
"version_description"
|
||||
)
|
||||
or "",
|
||||
"civitai_deleted": True,
|
||||
"civitai_checked": True,
|
||||
"remote_metadata_missing": True,
|
||||
"source": "local",
|
||||
"details_source": response.get("details_source") or "",
|
||||
}
|
||||
add_local_metadata_fields(metadata_payload, response)
|
||||
metadata_path = write_model_resolver_metadata(
|
||||
file_path,
|
||||
metadata_payload,
|
||||
category,
|
||||
"",
|
||||
) or ""
|
||||
response["metadata_path"] = metadata_path
|
||||
response["metadata_saved"] = bool(metadata_path)
|
||||
except Exception as metadata_error:
|
||||
self.logger.warning(
|
||||
f"Remote metadata no-match sidecar save failed: {metadata_error}"
|
||||
)
|
||||
return web.json_response(response)
|
||||
@@ -0,0 +1,292 @@
|
||||
"""Custom model URL route registration."""
|
||||
|
||||
from .context import RouteContext
|
||||
|
||||
|
||||
def register_custom_url_routes(context: RouteContext):
|
||||
UnsafeUrlError = context.get('UnsafeUrlError')
|
||||
asyncio = context.get('asyncio')
|
||||
build_civarchive_custom_result = context.get('build_civarchive_custom_result')
|
||||
build_civitai_custom_result = context.get('build_civitai_custom_result')
|
||||
build_huggingface_custom_result = context.get('build_huggingface_custom_result')
|
||||
extract_sha256_from_metadata = context.get('extract_sha256_from_metadata')
|
||||
get_civarchive_model_details = context.get('get_civarchive_model_details')
|
||||
get_civitai_download_url = context.get('get_civitai_download_url')
|
||||
get_civitai_model_details = context.get('get_civitai_model_details')
|
||||
get_filename_from_path = context.get('get_filename_from_path')
|
||||
host_matches_domain = context.get('host_matches_domain')
|
||||
json_api_endpoint = context.get('json_api_endpoint')
|
||||
looks_like_model_file = context.get('looks_like_model_file')
|
||||
normalize_category_to_model_type = context.get('normalize_category_to_model_type')
|
||||
normalize_sha256 = context.get('normalize_sha256')
|
||||
parse_civarchive_url = context.get('parse_civarchive_url')
|
||||
parse_civitai_url = context.get('parse_civitai_url')
|
||||
resolve_civarchive_by_hash = context.get('resolve_civarchive_by_hash')
|
||||
resolve_civarchive_model_version = context.get('resolve_civarchive_model_version')
|
||||
resolve_civitai_version_custom_result = context.get('resolve_civitai_version_custom_result')
|
||||
routes = context.get('routes')
|
||||
search_local_matches_by_hash = context.get('search_local_matches_by_hash')
|
||||
self = context.get('self')
|
||||
time = context.get('time')
|
||||
validate_public_http_url = context.get('validate_public_http_url')
|
||||
web = context.get('web')
|
||||
|
||||
def _custom_result_timestamp():
|
||||
return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
def _build_direct_custom_download_result(url, source, expected_filename=""):
|
||||
from urllib.parse import urlparse
|
||||
|
||||
if not looks_like_model_file(url, expected_filename):
|
||||
return None
|
||||
|
||||
parsed = urlparse(url)
|
||||
filename = (
|
||||
get_filename_from_path(parsed.path)
|
||||
or get_filename_from_path(expected_filename)
|
||||
or "model"
|
||||
)
|
||||
if expected_filename and "." in get_filename_from_path(expected_filename):
|
||||
filename = get_filename_from_path(expected_filename)
|
||||
source = str(source or "custom").strip().lower()
|
||||
source_label = {
|
||||
"civarchive": "CivArchive",
|
||||
"civitai": "CivitAI",
|
||||
"huggingface": "HuggingFace",
|
||||
}.get(source, "Custom URL")
|
||||
return {
|
||||
"source": source,
|
||||
"details_source": source,
|
||||
"name": source_label,
|
||||
"filename": filename,
|
||||
"url": url,
|
||||
"version_url": url,
|
||||
"download_url": url,
|
||||
"match_type": "custom_url",
|
||||
"custom_url": True,
|
||||
}
|
||||
|
||||
def _collect_custom_url_local_hash_matches(result, category):
|
||||
source_key = str(result.get("source") or "custom").strip().lower()
|
||||
sha256 = normalize_sha256(extract_sha256_from_metadata(result))
|
||||
if not sha256:
|
||||
return []
|
||||
|
||||
try:
|
||||
matches = search_local_matches_by_hash(
|
||||
sha256,
|
||||
category=category or None,
|
||||
max_matches=20,
|
||||
)
|
||||
except Exception as hash_error:
|
||||
self.logger.warning(
|
||||
f"Custom URL local metadata hash lookup failed for {source_key}:{sha256}: {hash_error}"
|
||||
)
|
||||
return []
|
||||
|
||||
return [
|
||||
{
|
||||
**match,
|
||||
"hash_lookup_source": source_key,
|
||||
"hash_lookup_filename": result.get("filename")
|
||||
or result.get("path")
|
||||
or "",
|
||||
"hash_lookup_sha256": sha256,
|
||||
}
|
||||
for match in matches
|
||||
]
|
||||
|
||||
@routes.post("/model_resolver/custom-url")
|
||||
@json_api_endpoint("custom-url")
|
||||
async def custom_url(request):
|
||||
"""Resolve a user-provided provider URL into a normal search result."""
|
||||
data = await request.json()
|
||||
raw_url = str(data.get("url") or data.get("custom_url") or "").strip()
|
||||
if not raw_url:
|
||||
return web.json_response(
|
||||
{"error": "URL is required"}, status=400
|
||||
)
|
||||
|
||||
if raw_url.startswith("hf://"):
|
||||
normalized_url = raw_url
|
||||
else:
|
||||
try:
|
||||
normalized_url = await asyncio.to_thread(
|
||||
validate_public_http_url,
|
||||
raw_url,
|
||||
)
|
||||
except UnsafeUrlError as exc:
|
||||
return web.json_response(
|
||||
{"error": str(exc)},
|
||||
status=400,
|
||||
)
|
||||
|
||||
category = data.get("category") or ""
|
||||
expected_filename = (
|
||||
data.get("filename")
|
||||
or get_filename_from_path(data.get("original_path") or "")
|
||||
or ""
|
||||
)
|
||||
civitai_key = data.get("civitai_key", "")
|
||||
hf_token = data.get("hf_token", "")
|
||||
|
||||
result = None
|
||||
source = ""
|
||||
try:
|
||||
civitai_parsed = parse_civitai_url(normalized_url)
|
||||
except Exception:
|
||||
civitai_parsed = None
|
||||
try:
|
||||
civarchive_parsed = parse_civarchive_url(normalized_url)
|
||||
except Exception:
|
||||
civarchive_parsed = None
|
||||
|
||||
if civitai_parsed:
|
||||
source = "civitai"
|
||||
model_id = civitai_parsed.get("model_id")
|
||||
version_id = civitai_parsed.get("version_id")
|
||||
if model_id:
|
||||
details = await asyncio.to_thread(
|
||||
get_civitai_model_details,
|
||||
model_id,
|
||||
version_id,
|
||||
civitai_key or None,
|
||||
)
|
||||
result = build_civitai_custom_result(
|
||||
details,
|
||||
expected_filename=expected_filename,
|
||||
api_key=civitai_key or None,
|
||||
)
|
||||
elif version_id:
|
||||
result = await asyncio.to_thread(
|
||||
resolve_civitai_version_custom_result,
|
||||
version_id,
|
||||
expected_filename,
|
||||
civitai_key or None,
|
||||
)
|
||||
if not result and version_id:
|
||||
result = {
|
||||
"source": "civitai",
|
||||
"details_source": "civitai",
|
||||
"version_id": version_id,
|
||||
"name": expected_filename or f"CivitAI version {version_id}",
|
||||
"filename": expected_filename or f"civitai-{version_id}",
|
||||
"url": normalized_url,
|
||||
"version_url": normalized_url,
|
||||
"download_url": get_civitai_download_url(
|
||||
version_id,
|
||||
civitai_key or None,
|
||||
),
|
||||
"match_type": "custom_url",
|
||||
"custom_url": True,
|
||||
}
|
||||
elif civarchive_parsed:
|
||||
source = "civarchive"
|
||||
if civarchive_parsed.get("sha256"):
|
||||
result = await asyncio.to_thread(
|
||||
resolve_civarchive_by_hash,
|
||||
civarchive_parsed.get("sha256"),
|
||||
expected_filename,
|
||||
False,
|
||||
normalize_category_to_model_type(category),
|
||||
)
|
||||
else:
|
||||
model_id = civarchive_parsed.get("model_id")
|
||||
version_id = civarchive_parsed.get("version_id")
|
||||
if model_id:
|
||||
result = await asyncio.to_thread(
|
||||
resolve_civarchive_model_version,
|
||||
model_id,
|
||||
version_id,
|
||||
expected_filename or str(model_id),
|
||||
False,
|
||||
True,
|
||||
)
|
||||
if not result:
|
||||
details = await asyncio.to_thread(
|
||||
get_civarchive_model_details,
|
||||
model_id,
|
||||
version_id,
|
||||
True,
|
||||
)
|
||||
result = build_civarchive_custom_result(
|
||||
details,
|
||||
expected_filename=expected_filename,
|
||||
)
|
||||
if result:
|
||||
result = dict(result)
|
||||
result["source"] = "civarchive"
|
||||
result["details_source"] = "civarchive"
|
||||
result["match_type"] = "custom_url"
|
||||
result["custom_url"] = True
|
||||
else:
|
||||
from urllib.parse import urlparse
|
||||
|
||||
parsed_direct = urlparse(normalized_url)
|
||||
direct_host = parsed_direct.hostname
|
||||
if (
|
||||
host_matches_domain(direct_host, "civarchive.com")
|
||||
and "/api/download/" in parsed_direct.path
|
||||
):
|
||||
source = "civarchive"
|
||||
result = _build_direct_custom_download_result(
|
||||
normalized_url,
|
||||
"civarchive",
|
||||
expected_filename,
|
||||
)
|
||||
else:
|
||||
result = await asyncio.to_thread(
|
||||
build_huggingface_custom_result,
|
||||
normalized_url,
|
||||
expected_filename,
|
||||
hf_token or None,
|
||||
)
|
||||
if result:
|
||||
source = "huggingface"
|
||||
|
||||
if not result:
|
||||
return web.json_response(
|
||||
{
|
||||
"error": (
|
||||
"Unsupported or unresolved URL. Use a HuggingFace file URL, "
|
||||
"CivitAI model/download URL, or CivArchive model/hash URL."
|
||||
)
|
||||
},
|
||||
status=400,
|
||||
)
|
||||
|
||||
result = dict(result)
|
||||
result["provided_url"] = normalized_url
|
||||
result["url_source"] = "custom"
|
||||
result["searched_at"] = result.get("searched_at") or _custom_result_timestamp()
|
||||
result.setdefault("category", category)
|
||||
if expected_filename and not result.get("filename"):
|
||||
result["filename"] = expected_filename
|
||||
|
||||
if not (result.get("download_url") or result.get("url")):
|
||||
return web.json_response(
|
||||
{"error": "The URL resolved, but no download URL was found"},
|
||||
status=400,
|
||||
)
|
||||
|
||||
source = source or result.get("source") or "custom"
|
||||
local_hash_matches = _collect_custom_url_local_hash_matches(
|
||||
result,
|
||||
category,
|
||||
)
|
||||
response = {
|
||||
"success": True,
|
||||
"source": source,
|
||||
"result": result,
|
||||
"custom": [result],
|
||||
"searched_sources": ["custom"],
|
||||
"local_hash_matches": local_hash_matches,
|
||||
}
|
||||
response[source] = result
|
||||
return web.json_response(response)
|
||||
@@ -0,0 +1,663 @@
|
||||
"""Local model and hash route registration."""
|
||||
|
||||
from .context import RouteContext
|
||||
|
||||
|
||||
def register_hash_routes(context: RouteContext):
|
||||
FileManagerError = context.get('FileManagerError')
|
||||
FileManagerUnavailableError = context.get('FileManagerUnavailableError')
|
||||
MODEL_RESOLVER_METADATA_SCHEMA = context.get('MODEL_RESOLVER_METADATA_SCHEMA')
|
||||
MODEL_RESOLVER_METADATA_SCHEMA_VERSION = context.get('MODEL_RESOLVER_METADATA_SCHEMA_VERSION')
|
||||
UnsupportedFileManagerPlatformError = context.get('UnsupportedFileManagerPlatformError')
|
||||
asyncio = context.get('asyncio')
|
||||
cancel_progress_response = context.get('cancel_progress_response')
|
||||
get_existing_model_preview_path = context.get('get_existing_model_preview_path')
|
||||
get_filename_from_path = context.get('get_filename_from_path')
|
||||
get_local_model_hash_metadata = context.get('get_local_model_hash_metadata')
|
||||
get_progress_response = context.get('get_progress_response')
|
||||
get_safe_model_resolver_sidecar_path = context.get('get_safe_model_resolver_sidecar_path')
|
||||
get_workflow_model_inventory = context.get('get_workflow_model_inventory')
|
||||
is_path_in_configured_model_roots = context.get('is_path_in_configured_model_roots')
|
||||
json_api_endpoint = context.get('json_api_endpoint')
|
||||
load_resolver_settings = context.get('load_resolver_settings')
|
||||
normalize_file_manager_path = context.get('normalize_file_manager_path')
|
||||
normalize_sha256 = context.get('normalize_sha256')
|
||||
open_in_file_manager = context.get('open_in_file_manager')
|
||||
os = context.get('os')
|
||||
read_json_safe = context.get('read_json_safe')
|
||||
resolver_bool_setting = context.get('resolver_bool_setting')
|
||||
routes = context.get('routes')
|
||||
run_in_background_thread = context.get('run_in_background_thread')
|
||||
search_local_matches_by_hash = context.get('search_local_matches_by_hash')
|
||||
self = context.get('self')
|
||||
time = context.get('time')
|
||||
to_bool = context.get('to_bool')
|
||||
to_int = context.get('to_int')
|
||||
web = context.get('web')
|
||||
write_json_atomic = context.get('write_json_atomic')
|
||||
|
||||
@routes.post("/model_resolver/local-model-hashes")
|
||||
@json_api_endpoint("local-model-hashes")
|
||||
async def local_model_hashes(request):
|
||||
"""Return SHA256 hashes already stored in local sidecar metadata."""
|
||||
data = await request.json()
|
||||
model = data.get("model") if isinstance(data.get("model"), dict) else {}
|
||||
path = (
|
||||
data.get("path")
|
||||
or data.get("file_path")
|
||||
or data.get("resolved_path")
|
||||
or model.get("path")
|
||||
or model.get("resolved_path")
|
||||
or ""
|
||||
)
|
||||
|
||||
if not path:
|
||||
return web.json_response(
|
||||
{"error": "path is required"}, status=400
|
||||
)
|
||||
|
||||
normalized_path = os.path.realpath(
|
||||
os.path.abspath(os.path.normpath(str(path)))
|
||||
)
|
||||
if not is_path_in_configured_model_roots(normalized_path):
|
||||
return web.json_response(
|
||||
{"error": "path is outside configured model directories"},
|
||||
status=403,
|
||||
)
|
||||
|
||||
return web.json_response(
|
||||
get_local_model_hash_metadata(normalized_path, model=model)
|
||||
)
|
||||
|
||||
@routes.get("/model_resolver/model-preview")
|
||||
async def get_model_preview(request):
|
||||
"""Serve an adjacent model preview from a configured model directory."""
|
||||
model_path = str(request.query.get("path") or "").strip()
|
||||
if not model_path:
|
||||
return web.Response(text="path is required", status=400)
|
||||
|
||||
try:
|
||||
normalized_path = os.path.realpath(
|
||||
os.path.abspath(os.path.normpath(model_path))
|
||||
)
|
||||
except (OSError, TypeError, ValueError):
|
||||
return web.Response(text="invalid model path", status=400)
|
||||
if not os.path.isfile(normalized_path):
|
||||
return web.Response(text="model file does not exist", status=404)
|
||||
if not is_path_in_configured_model_roots(normalized_path):
|
||||
return web.Response(
|
||||
text="path is outside configured model directories",
|
||||
status=403,
|
||||
)
|
||||
|
||||
preview_path = get_existing_model_preview_path(normalized_path)
|
||||
if not preview_path:
|
||||
return web.Response(text="preview not found", status=404)
|
||||
if not is_path_in_configured_model_roots(preview_path):
|
||||
return web.Response(
|
||||
text="preview is outside configured model directories",
|
||||
status=403,
|
||||
)
|
||||
|
||||
return web.FileResponse(
|
||||
preview_path,
|
||||
headers={"Cache-Control": "private, no-cache"},
|
||||
)
|
||||
|
||||
@routes.post("/model_resolver/workflow-model-hashes")
|
||||
@json_api_endpoint("workflow-model-hashes")
|
||||
async def workflow_model_hashes(request):
|
||||
"""Return hash metadata for existing local models used by a workflow."""
|
||||
data = await request.json()
|
||||
workflow_json = data.get("workflow")
|
||||
if not isinstance(workflow_json, dict):
|
||||
return web.json_response(
|
||||
{"error": "Workflow JSON must be an object"}, status=400
|
||||
)
|
||||
|
||||
settings = await asyncio.to_thread(load_resolver_settings)
|
||||
if not resolver_bool_setting(
|
||||
settings.get("workflow_hash_metadata_enabled"), True
|
||||
):
|
||||
return web.json_response(
|
||||
{
|
||||
"success": True,
|
||||
"enabled": False,
|
||||
"models": [],
|
||||
"by_node": {},
|
||||
"by_path": {},
|
||||
"count": 0,
|
||||
}
|
||||
)
|
||||
|
||||
inventory = await asyncio.to_thread(
|
||||
get_workflow_model_inventory,
|
||||
workflow_json,
|
||||
)
|
||||
refs = inventory["model_refs"]
|
||||
|
||||
by_node = {}
|
||||
by_path = {}
|
||||
models = []
|
||||
seen = set()
|
||||
|
||||
for ref in refs:
|
||||
if not isinstance(ref, dict) or not ref.get("exists"):
|
||||
continue
|
||||
full_path = str(ref.get("full_path") or "").strip()
|
||||
if not full_path:
|
||||
continue
|
||||
model_info = {
|
||||
"path": full_path,
|
||||
"filename": get_filename_from_path(full_path),
|
||||
"relative_path": ref.get("original_path") or "",
|
||||
"category": ref.get("category") or "",
|
||||
}
|
||||
metadata = get_local_model_hash_metadata(
|
||||
full_path, model=model_info
|
||||
)
|
||||
sha256 = normalize_sha256(metadata.get("sha256"))
|
||||
if not sha256:
|
||||
continue
|
||||
|
||||
entry = {
|
||||
"node_id": ref.get("node_id"),
|
||||
"node_type": ref.get("node_type") or "",
|
||||
"widget_index": ref.get("widget_index"),
|
||||
"widget_name": ref.get("widget_name") or "",
|
||||
"path": ref.get("original_path") or "",
|
||||
"filename": get_filename_from_path(
|
||||
ref.get("original_path") or full_path
|
||||
),
|
||||
"category": ref.get("category") or "",
|
||||
"sha256": sha256,
|
||||
"size": metadata.get("size") or 0,
|
||||
}
|
||||
entry_key = (
|
||||
str(entry.get("node_id")),
|
||||
str(entry.get("widget_index")),
|
||||
entry.get("path"),
|
||||
sha256,
|
||||
)
|
||||
if entry_key in seen:
|
||||
continue
|
||||
seen.add(entry_key)
|
||||
models.append(entry)
|
||||
|
||||
path_key = str(entry.get("path") or entry.get("filename") or "")
|
||||
if path_key:
|
||||
by_path[path_key] = entry
|
||||
by_path[get_filename_from_path(path_key)] = entry
|
||||
node_key = f"{entry.get('node_id')}:{entry.get('widget_index')}"
|
||||
by_node[node_key] = entry
|
||||
|
||||
return web.json_response(
|
||||
{
|
||||
"success": True,
|
||||
"enabled": True,
|
||||
"models": models,
|
||||
"by_node": by_node,
|
||||
"by_path": by_path,
|
||||
"count": len(models),
|
||||
}
|
||||
)
|
||||
|
||||
@routes.post("/model_resolver/local-matches-by-hash")
|
||||
@json_api_endpoint("local-matches-by-hash")
|
||||
async def local_matches_by_hash(request):
|
||||
"""Search local model metadata sidecars for a remote SHA256."""
|
||||
data = await request.json()
|
||||
sha256 = normalize_sha256(
|
||||
data.get("sha256")
|
||||
or data.get("hash")
|
||||
or data.get("SHA256")
|
||||
or ""
|
||||
)
|
||||
if not sha256:
|
||||
return web.json_response(
|
||||
{"error": "sha256 is required"}, status=400
|
||||
)
|
||||
|
||||
category = data.get("category", "")
|
||||
source = str(
|
||||
data.get("source")
|
||||
or data.get("hash_lookup_source")
|
||||
or "download_source"
|
||||
).strip().lower().replace("-", "_")
|
||||
filename = (
|
||||
data.get("filename")
|
||||
or data.get("path")
|
||||
or data.get("model_name")
|
||||
or ""
|
||||
)
|
||||
max_matches = to_int(data.get("max_matches"), 20)
|
||||
force_rescan = to_bool(data.get("force_rescan"), False)
|
||||
|
||||
matches = search_local_matches_by_hash(
|
||||
sha256,
|
||||
category=category or None,
|
||||
max_matches=max_matches,
|
||||
force_rescan=force_rescan,
|
||||
)
|
||||
enriched_matches = [
|
||||
{
|
||||
**match,
|
||||
"hash_lookup_source": source or "download_source",
|
||||
"hash_lookup_filename": filename,
|
||||
"hash_lookup_sha256": sha256,
|
||||
}
|
||||
for match in matches
|
||||
]
|
||||
return web.json_response(
|
||||
{
|
||||
"sha256": sha256,
|
||||
"local_hash_matches": enriched_matches,
|
||||
"matches": enriched_matches,
|
||||
}
|
||||
)
|
||||
|
||||
@routes.post("/model_resolver/open-containing-folder")
|
||||
@json_api_endpoint("open-containing-folder")
|
||||
async def open_containing_folder(request):
|
||||
"""Reveal a file or open a directory in the host file manager."""
|
||||
try:
|
||||
data = await request.json()
|
||||
except Exception as exc:
|
||||
return web.json_response(
|
||||
{"success": False, "error": f"invalid JSON body: {exc}"},
|
||||
status=400,
|
||||
)
|
||||
if not isinstance(data, dict):
|
||||
return web.json_response(
|
||||
{"success": False, "error": "JSON body must be an object"},
|
||||
status=400,
|
||||
)
|
||||
target_path = data.get("path", "")
|
||||
|
||||
try:
|
||||
normalized_path = normalize_file_manager_path(target_path)
|
||||
except ValueError as exc:
|
||||
return web.json_response(
|
||||
{"success": False, "error": str(exc)}, status=400
|
||||
)
|
||||
|
||||
if not os.path.exists(normalized_path):
|
||||
return web.json_response(
|
||||
{"success": False, "error": "path does not exist"}, status=404
|
||||
)
|
||||
if not is_path_in_configured_model_roots(normalized_path):
|
||||
return web.json_response(
|
||||
{
|
||||
"success": False,
|
||||
"error": "path is outside configured model directories",
|
||||
},
|
||||
status=403,
|
||||
)
|
||||
|
||||
try:
|
||||
result = await asyncio.to_thread(
|
||||
open_in_file_manager, normalized_path
|
||||
)
|
||||
except FileNotFoundError:
|
||||
return web.json_response(
|
||||
{"success": False, "error": "path no longer exists"},
|
||||
status=404,
|
||||
)
|
||||
except UnsupportedFileManagerPlatformError as exc:
|
||||
return web.json_response(
|
||||
{"success": False, "error": str(exc)}, status=501
|
||||
)
|
||||
except FileManagerUnavailableError as exc:
|
||||
return web.json_response(
|
||||
{
|
||||
"success": False,
|
||||
"error": (
|
||||
f"{exc} The ComfyUI host may be running without "
|
||||
"a graphical desktop session."
|
||||
),
|
||||
},
|
||||
status=503,
|
||||
)
|
||||
except FileManagerError as exc:
|
||||
return web.json_response(
|
||||
{"success": False, "error": str(exc)}, status=500
|
||||
)
|
||||
|
||||
return web.json_response({"success": True, **result})
|
||||
|
||||
# cleanup_hash_progress, update_hash_progress, is_hash_progress_cancelled, and mark_hash_progress_cancelled removed
|
||||
|
||||
from ..path_utils import HashCalculationCancelled
|
||||
|
||||
def resolve_hash_file_request(data):
|
||||
import os as _os
|
||||
|
||||
file_path = (
|
||||
data.get("file_path")
|
||||
or data.get("resolved_path")
|
||||
or data.get("path")
|
||||
or ""
|
||||
)
|
||||
if not file_path:
|
||||
return "", "file_path is required"
|
||||
|
||||
normalized_path = _os.path.realpath(
|
||||
_os.path.abspath(_os.path.normpath(file_path))
|
||||
)
|
||||
if not _os.path.exists(normalized_path) or not _os.path.isfile(normalized_path):
|
||||
return "", "file does not exist"
|
||||
if not is_path_in_configured_model_roots(normalized_path):
|
||||
return "", "file is outside configured model directories"
|
||||
return normalized_path, ""
|
||||
|
||||
def write_calculated_hash_metadata(
|
||||
normalized_path,
|
||||
sha256,
|
||||
sha256_source="file",
|
||||
):
|
||||
import os as _os
|
||||
|
||||
resolved_metadata_path = get_safe_model_resolver_sidecar_path(
|
||||
normalized_path
|
||||
)
|
||||
|
||||
metadata_updated = False
|
||||
try:
|
||||
metadata = read_json_safe(resolved_metadata_path, {})
|
||||
if not isinstance(metadata, dict):
|
||||
metadata = {}
|
||||
|
||||
filename = get_filename_from_path(normalized_path)
|
||||
stem, _ext = _os.path.splitext(filename)
|
||||
hashes = metadata.get("hashes")
|
||||
if not isinstance(hashes, dict):
|
||||
hashes = {}
|
||||
hashes["SHA256"] = sha256
|
||||
|
||||
metadata["schema"] = MODEL_RESOLVER_METADATA_SCHEMA
|
||||
metadata["schema_version"] = (
|
||||
MODEL_RESOLVER_METADATA_SCHEMA_VERSION
|
||||
)
|
||||
metadata["managed_by"] = MODEL_RESOLVER_METADATA_SCHEMA
|
||||
metadata["sha256"] = sha256
|
||||
metadata["hashes"] = hashes
|
||||
metadata["hash_status"] = "completed"
|
||||
metadata["sha256_source"] = sha256_source or "file"
|
||||
metadata["last_checked_at"] = time.time()
|
||||
metadata.setdefault("file_name", stem)
|
||||
metadata.setdefault("model_name", stem)
|
||||
metadata.setdefault("file_path", normalized_path.replace("\\", "/"))
|
||||
try:
|
||||
metadata.setdefault("size", _os.path.getsize(normalized_path))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
write_json_atomic(resolved_metadata_path, metadata, indent=2)
|
||||
metadata_updated = True
|
||||
self.logger.info(
|
||||
f"Stored SHA256 and updated metadata: {resolved_metadata_path}"
|
||||
)
|
||||
except Exception as metadata_error:
|
||||
self.logger.warning(
|
||||
f"Could not update metadata with calculated SHA256 for {normalized_path}: {metadata_error}"
|
||||
)
|
||||
|
||||
return resolved_metadata_path, metadata_updated
|
||||
|
||||
def calculate_sha256_with_progress(normalized_path, progress_id=""):
|
||||
import os as _os
|
||||
|
||||
from ..path_utils import calculate_file_sha256 as _calculate_file_sha256_core
|
||||
|
||||
total_bytes = max(0, _os.path.getsize(normalized_path))
|
||||
|
||||
self.hash_tracker.update(
|
||||
progress_id,
|
||||
status="running",
|
||||
stage="header",
|
||||
message="Checking safetensors header...",
|
||||
percent=0,
|
||||
bytes_read=0,
|
||||
total_bytes=total_bytes,
|
||||
)
|
||||
|
||||
if self.hash_tracker.is_cancelled(progress_id):
|
||||
raise HashCalculationCancelled()
|
||||
|
||||
hash_source = ["file"]
|
||||
stage_transitioned = [False]
|
||||
_bytes_read = [0]
|
||||
_last_update = [0.0]
|
||||
|
||||
def on_progress(bytes_read, total_bytes):
|
||||
_bytes_read[0] = bytes_read
|
||||
if not stage_transitioned[0]:
|
||||
stage_transitioned[0] = True
|
||||
self.logger.info(
|
||||
"No SHA256 found in safetensors header; calculating full "
|
||||
f"file SHA256: {normalized_path}"
|
||||
)
|
||||
self.hash_tracker.update(
|
||||
progress_id,
|
||||
status="running",
|
||||
stage="hashing",
|
||||
message="Calculating SHA256...",
|
||||
percent=0,
|
||||
bytes_read=0,
|
||||
total_bytes=total_bytes,
|
||||
)
|
||||
|
||||
now = time.time()
|
||||
if now - _last_update[0] >= 0.15 or bytes_read >= total_bytes:
|
||||
percent = 98 if total_bytes <= 0 else min(
|
||||
98,
|
||||
(bytes_read / total_bytes) * 98,
|
||||
)
|
||||
self.hash_tracker.update(
|
||||
progress_id,
|
||||
status="running",
|
||||
stage="hashing",
|
||||
message="Calculating SHA256...",
|
||||
percent=percent,
|
||||
bytes_read=bytes_read,
|
||||
total_bytes=total_bytes,
|
||||
)
|
||||
_last_update[0] = now
|
||||
|
||||
def is_cancelled():
|
||||
if self.hash_tracker.is_cancelled(progress_id):
|
||||
percent = 0 if total_bytes <= 0 else min(
|
||||
98,
|
||||
(_bytes_read[0] / total_bytes) * 98,
|
||||
)
|
||||
self.hash_tracker.update(
|
||||
progress_id,
|
||||
status="cancelled",
|
||||
stage="cancelled",
|
||||
message="Hash calculation cancelled",
|
||||
percent=percent,
|
||||
bytes_read=_bytes_read[0],
|
||||
total_bytes=total_bytes,
|
||||
)
|
||||
return True
|
||||
return False
|
||||
|
||||
def on_hash_source(source_name):
|
||||
hash_source[0] = source_name
|
||||
if source_name == "safetensors_header":
|
||||
self.logger.info(
|
||||
f"Found SHA256 in safetensors header: {normalized_path}"
|
||||
)
|
||||
self.hash_tracker.update(
|
||||
progress_id,
|
||||
status="running",
|
||||
stage="header",
|
||||
message="Found SHA256 in safetensors header",
|
||||
percent=98,
|
||||
bytes_read=0,
|
||||
total_bytes=total_bytes,
|
||||
sha256_source="safetensors_header",
|
||||
)
|
||||
|
||||
sha256 = _calculate_file_sha256_core(
|
||||
normalized_path,
|
||||
chunk_size=1024 * 1024 * 4,
|
||||
on_progress=on_progress,
|
||||
is_cancelled=is_cancelled,
|
||||
use_safetensors_header=True,
|
||||
on_hash_source=on_hash_source,
|
||||
)
|
||||
|
||||
if is_cancelled():
|
||||
raise HashCalculationCancelled()
|
||||
|
||||
return sha256, hash_source[0]
|
||||
|
||||
@routes.post("/model_resolver/calculate-file-hash")
|
||||
@json_api_endpoint("calculate-file-hash")
|
||||
async def calculate_file_hash_route(request):
|
||||
"""Calculate SHA256 for a local model and persist it to sidecar metadata."""
|
||||
data = await request.json()
|
||||
|
||||
normalized_path, error = resolve_hash_file_request(data)
|
||||
if error == "file_path is required":
|
||||
return web.json_response(
|
||||
{"error": "file_path is required"}, status=400
|
||||
)
|
||||
if error:
|
||||
status = 403 if error == "file is outside configured model directories" else 404
|
||||
return web.json_response(
|
||||
{"error": error}, status=status
|
||||
)
|
||||
|
||||
sha256, sha256_source = calculate_sha256_with_progress(normalized_path)
|
||||
if not sha256:
|
||||
return web.json_response(
|
||||
{"error": "could not calculate hash"}, status=500
|
||||
)
|
||||
resolved_metadata_path, metadata_updated = write_calculated_hash_metadata(
|
||||
normalized_path,
|
||||
sha256,
|
||||
sha256_source,
|
||||
)
|
||||
|
||||
return web.json_response(
|
||||
{
|
||||
"success": True,
|
||||
"sha256": sha256,
|
||||
"hash": sha256,
|
||||
"sha256_source": sha256_source,
|
||||
"file_path": normalized_path,
|
||||
"metadata_path": resolved_metadata_path,
|
||||
"metadata_updated": metadata_updated,
|
||||
}
|
||||
)
|
||||
|
||||
@routes.post("/model_resolver/calculate-file-hash/start")
|
||||
@json_api_endpoint("calculate-file-hash-start")
|
||||
async def calculate_file_hash_start_route(request):
|
||||
"""Start SHA256 calculation in a background thread."""
|
||||
import uuid
|
||||
|
||||
data = await request.json()
|
||||
normalized_path, error = resolve_hash_file_request(data)
|
||||
if error == "file_path is required":
|
||||
return web.json_response(
|
||||
{"error": "file_path is required"}, status=400
|
||||
)
|
||||
if error:
|
||||
status = 403 if error == "file is outside configured model directories" else 404
|
||||
return web.json_response(
|
||||
{"error": error}, status=status
|
||||
)
|
||||
|
||||
self.hash_tracker.cleanup()
|
||||
progress_id = f"hash_{uuid.uuid4().hex}"
|
||||
self.hash_tracker.update(
|
||||
progress_id,
|
||||
status="queued",
|
||||
stage="queued",
|
||||
message="Preparing hash calculation...",
|
||||
percent=0,
|
||||
file_path=normalized_path,
|
||||
)
|
||||
|
||||
def run_hash_task():
|
||||
sha256, sha256_source = calculate_sha256_with_progress(
|
||||
normalized_path,
|
||||
progress_id=progress_id,
|
||||
)
|
||||
if self.hash_tracker.is_cancelled(progress_id):
|
||||
raise HashCalculationCancelled()
|
||||
self.hash_tracker.update(
|
||||
progress_id,
|
||||
status="running",
|
||||
stage="metadata",
|
||||
message="Saving metadata...",
|
||||
percent=99,
|
||||
)
|
||||
resolved_metadata_path, metadata_updated = write_calculated_hash_metadata(
|
||||
normalized_path,
|
||||
sha256,
|
||||
sha256_source,
|
||||
)
|
||||
return sha256, sha256_source, resolved_metadata_path, metadata_updated
|
||||
|
||||
def on_success(res):
|
||||
sha256, sha256_source, resolved_metadata_path, metadata_updated = res
|
||||
self.hash_tracker.update(
|
||||
progress_id,
|
||||
status="done",
|
||||
stage="done",
|
||||
message="Hash calculated",
|
||||
percent=100,
|
||||
sha256=sha256,
|
||||
hash=sha256,
|
||||
sha256_source=sha256_source,
|
||||
file_path=normalized_path,
|
||||
metadata_path=resolved_metadata_path,
|
||||
metadata_updated=metadata_updated,
|
||||
)
|
||||
|
||||
def on_cancel(res=None):
|
||||
self.hash_tracker.update(
|
||||
progress_id,
|
||||
status="cancelled",
|
||||
stage="cancelled",
|
||||
message="Hash calculation cancelled",
|
||||
)
|
||||
|
||||
run_in_background_thread(
|
||||
self.hash_tracker,
|
||||
progress_id,
|
||||
run_hash_task,
|
||||
on_success,
|
||||
on_cancel,
|
||||
error_log_msg=f"Hash calculation failed for {normalized_path}",
|
||||
)
|
||||
return web.json_response(
|
||||
{
|
||||
"success": True,
|
||||
"progress_id": progress_id,
|
||||
}
|
||||
)
|
||||
|
||||
@routes.get("/model_resolver/calculate-file-hash/progress/{progress_id}")
|
||||
@json_api_endpoint("calculate-file-hash-progress")
|
||||
async def calculate_file_hash_progress_route(request):
|
||||
"""Return progress for a background SHA256 calculation."""
|
||||
return get_progress_response(
|
||||
self.hash_tracker,
|
||||
request,
|
||||
not_found_status=404
|
||||
)
|
||||
|
||||
@routes.post("/model_resolver/calculate-file-hash/cancel/{progress_id}")
|
||||
@json_api_endpoint("calculate-file-hash-cancel")
|
||||
async def calculate_file_hash_cancel_route(request):
|
||||
"""Cancel a background SHA256 calculation."""
|
||||
return cancel_progress_response(
|
||||
self.hash_tracker,
|
||||
request,
|
||||
cancel_message="Stopping hash calculation..."
|
||||
)
|
||||
@@ -0,0 +1,479 @@
|
||||
"""Loaded model route registration."""
|
||||
|
||||
from .context import RouteContext
|
||||
|
||||
|
||||
def register_loaded_model_routes(context: RouteContext):
|
||||
adapt_custom_node_loaded_model = context.get('adapt_custom_node_loaded_model')
|
||||
asyncio = context.get('asyncio')
|
||||
get_filename_from_path = context.get('get_filename_from_path')
|
||||
get_progress_response = context.get('get_progress_response')
|
||||
get_workflow_model_inventory = context.get('get_workflow_model_inventory')
|
||||
json_api_endpoint = context.get('json_api_endpoint')
|
||||
routes = context.get('routes')
|
||||
self = context.get('self')
|
||||
web = context.get('web')
|
||||
|
||||
@routes.post("/model_resolver/loaded")
|
||||
@json_api_endpoint("get_loaded_models")
|
||||
async def get_loaded_models(request):
|
||||
"""Get all currently loaded models in the workflow."""
|
||||
data = await request.json()
|
||||
workflow_json = data.get("workflow")
|
||||
loaded_id = str(
|
||||
data.get("loaded_id") or data.get("progress_id") or ""
|
||||
).strip()
|
||||
|
||||
if not workflow_json:
|
||||
return web.json_response(
|
||||
{"error": "Workflow JSON is required"}, status=400
|
||||
)
|
||||
if not isinstance(workflow_json, dict):
|
||||
return web.json_response(
|
||||
{"error": "Workflow JSON must be an object"}, status=400
|
||||
)
|
||||
|
||||
def update_loaded_progress(*args, **kwargs):
|
||||
self._update_loaded_progress(loaded_id, *args, **kwargs)
|
||||
|
||||
def get_workflow_node_count():
|
||||
node_count = 0
|
||||
nodes = workflow_json.get("nodes", [])
|
||||
if isinstance(nodes, list):
|
||||
node_count += len(nodes)
|
||||
|
||||
definitions = workflow_json.get("definitions", {})
|
||||
subgraphs = (
|
||||
definitions.get("subgraphs", [])
|
||||
if isinstance(definitions, dict)
|
||||
else []
|
||||
)
|
||||
if isinstance(subgraphs, list):
|
||||
for subgraph in subgraphs:
|
||||
if not isinstance(subgraph, dict):
|
||||
continue
|
||||
subgraph_nodes = subgraph.get("nodes", [])
|
||||
if isinstance(subgraph_nodes, list):
|
||||
node_count += len(subgraph_nodes)
|
||||
return node_count
|
||||
|
||||
def interpolate_percent(start, end, current, total):
|
||||
if not total:
|
||||
return start
|
||||
try:
|
||||
ratio = max(0.0, min(1.0, float(current) / float(total)))
|
||||
except (TypeError, ValueError, ZeroDivisionError):
|
||||
ratio = 0.0
|
||||
return start + ((end - start) * ratio)
|
||||
|
||||
def build_loaded_models_response():
|
||||
from ..workflow_analyzer import (
|
||||
URN_TYPE_MAP,
|
||||
get_lora_model_strength,
|
||||
)
|
||||
|
||||
update_loaded_progress(
|
||||
"scanning",
|
||||
"Scanning local model index...",
|
||||
percent=5,
|
||||
)
|
||||
|
||||
workflow_node_count = get_workflow_node_count()
|
||||
|
||||
def update_workflow_analysis_progress(payload):
|
||||
self._update_workflow_analysis_progress(
|
||||
loaded_id,
|
||||
workflow_node_count,
|
||||
interpolate_percent,
|
||||
payload,
|
||||
start_percent=5,
|
||||
end_percent=45,
|
||||
)
|
||||
|
||||
inventory = get_workflow_model_inventory(
|
||||
workflow_json,
|
||||
progress_callback=update_workflow_analysis_progress,
|
||||
)
|
||||
available_models = inventory["available_models"]
|
||||
all_model_refs = inventory["model_refs"]
|
||||
|
||||
# Create lookup for full paths by filename (with and without extension)
|
||||
path_by_filename = {}
|
||||
total_local_models = len(available_models)
|
||||
update_loaded_progress(
|
||||
"indexing",
|
||||
f"Indexing {total_local_models} local models...",
|
||||
percent=45,
|
||||
current=0,
|
||||
total=total_local_models,
|
||||
)
|
||||
for index, model_info in enumerate(available_models, start=1):
|
||||
rel_path = model_info.get("relative_path", "")
|
||||
if rel_path:
|
||||
filename = get_filename_from_path(rel_path)
|
||||
path_by_filename[filename] = model_info.get("path")
|
||||
# Also add without extension for matching (simple approach)
|
||||
if "." in filename:
|
||||
filename_no_ext = filename.rsplit(".", 1)[0]
|
||||
if filename_no_ext not in path_by_filename:
|
||||
path_by_filename[filename_no_ext] = model_info.get("path")
|
||||
# Add the full relative path as key too
|
||||
path_by_filename[rel_path] = model_info.get("path")
|
||||
|
||||
if (
|
||||
total_local_models
|
||||
and (index == total_local_models or index % 250 == 0)
|
||||
):
|
||||
update_loaded_progress(
|
||||
"indexing",
|
||||
f"Indexing local model {index} of {total_local_models}",
|
||||
percent=interpolate_percent(
|
||||
45, 55, index, total_local_models
|
||||
),
|
||||
current=index,
|
||||
total=total_local_models,
|
||||
)
|
||||
|
||||
# Also use folder_paths.get_full_path() to get paths
|
||||
try:
|
||||
import folder_paths
|
||||
except Exception:
|
||||
folder_paths = None
|
||||
|
||||
folder_categories = [
|
||||
"loras",
|
||||
"checkpoints",
|
||||
"vae",
|
||||
"controlnet",
|
||||
"upscale_models",
|
||||
]
|
||||
if folder_paths:
|
||||
for index, category_name in enumerate(folder_categories, start=1):
|
||||
update_loaded_progress(
|
||||
"indexing",
|
||||
f"Reading {category_name} model list...",
|
||||
percent=interpolate_percent(
|
||||
55, 65, index - 1, len(folder_categories)
|
||||
),
|
||||
current=index - 1,
|
||||
total=len(folder_categories),
|
||||
)
|
||||
try:
|
||||
filenames = folder_paths.get_filename_list(category_name)
|
||||
for filename in filenames:
|
||||
full_path = folder_paths.get_full_path(
|
||||
category_name, filename
|
||||
)
|
||||
if (
|
||||
full_path
|
||||
and full_path not in path_by_filename.values()
|
||||
):
|
||||
path_by_filename[filename] = full_path
|
||||
filename_no_ext = (
|
||||
filename.rsplit(".", 1)[0]
|
||||
if "." in filename
|
||||
else filename
|
||||
)
|
||||
if filename_no_ext not in path_by_filename:
|
||||
path_by_filename[filename_no_ext] = full_path
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Also extract from node.properties.models
|
||||
workflow_nodes = workflow_json.get("nodes", [])
|
||||
nodes = list(workflow_nodes) if isinstance(workflow_nodes, list) else []
|
||||
node_scopes = {
|
||||
id(node): {
|
||||
"subgraph_id": "",
|
||||
"subgraph_name": "",
|
||||
"is_top_level": True,
|
||||
}
|
||||
for node in nodes
|
||||
if isinstance(node, dict)
|
||||
}
|
||||
definitions = workflow_json.get("definitions", {})
|
||||
subgraphs = (
|
||||
definitions.get("subgraphs", [])
|
||||
if isinstance(definitions, dict)
|
||||
else []
|
||||
)
|
||||
for subgraph in subgraphs:
|
||||
if isinstance(subgraph, dict):
|
||||
subgraph_nodes = subgraph.get("nodes", [])
|
||||
if isinstance(subgraph_nodes, list):
|
||||
nodes.extend(subgraph_nodes)
|
||||
for node in subgraph_nodes:
|
||||
if not isinstance(node, dict):
|
||||
continue
|
||||
node_scopes[id(node)] = {
|
||||
"subgraph_id": subgraph.get("id") or "",
|
||||
"subgraph_name": (
|
||||
subgraph.get("name")
|
||||
or subgraph.get("id")
|
||||
or ""
|
||||
),
|
||||
"is_top_level": False,
|
||||
}
|
||||
|
||||
def get_node_scope(node):
|
||||
return node_scopes.get(
|
||||
id(node),
|
||||
{
|
||||
"subgraph_id": "",
|
||||
"subgraph_name": "",
|
||||
"is_top_level": True,
|
||||
},
|
||||
)
|
||||
|
||||
def node_matches_ref(node, ref):
|
||||
if not isinstance(node, dict):
|
||||
return False
|
||||
if str(node.get("id")) != str(ref.get("node_id")):
|
||||
return False
|
||||
scope = get_node_scope(node)
|
||||
if scope["is_top_level"] != (
|
||||
ref.get("is_top_level") is not False
|
||||
):
|
||||
return False
|
||||
if not scope["is_top_level"]:
|
||||
return str(scope["subgraph_id"]) == str(
|
||||
ref.get("subgraph_id") or ""
|
||||
)
|
||||
return True
|
||||
|
||||
# Collect all loaded models with their values
|
||||
loaded_models = []
|
||||
total_refs = len(all_model_refs)
|
||||
update_loaded_progress(
|
||||
"building",
|
||||
"Building loaded models list...",
|
||||
percent=65,
|
||||
current=0,
|
||||
total=total_refs,
|
||||
)
|
||||
|
||||
# Process each model reference from analyze_workflow_models
|
||||
for index, ref in enumerate(all_model_refs, start=1):
|
||||
if total_refs and (index == total_refs or index % 50 == 0):
|
||||
update_loaded_progress(
|
||||
"building",
|
||||
f"Building loaded model {index} of {total_refs}",
|
||||
percent=interpolate_percent(65, 94, index, total_refs),
|
||||
current=index,
|
||||
total=total_refs,
|
||||
)
|
||||
|
||||
original_path = ref.get("original_path", "")
|
||||
node_id = ref.get("node_id")
|
||||
widget_index = ref.get("widget_index")
|
||||
node_type = ref.get("node_type", "")
|
||||
category = ref.get("category", "unknown")
|
||||
|
||||
# Determine model name and strength
|
||||
model_name = get_filename_from_path(original_path)
|
||||
strength = None
|
||||
|
||||
if (
|
||||
ref.get("strength") is None
|
||||
and category in {"lora", "loras"}
|
||||
):
|
||||
for node in nodes:
|
||||
if node_matches_ref(node, ref):
|
||||
strength = get_lora_model_strength(
|
||||
node,
|
||||
widget_index,
|
||||
)
|
||||
break
|
||||
|
||||
if ref.get("strength") is not None:
|
||||
strength = ref.get("strength")
|
||||
|
||||
model_name, strength = adapt_custom_node_loaded_model(
|
||||
ref,
|
||||
model_name,
|
||||
strength,
|
||||
)
|
||||
|
||||
# Check if model exists locally
|
||||
exists = ref.get("exists", False)
|
||||
|
||||
# If URN, resolve to display name
|
||||
if ref.get("is_urn"):
|
||||
urn = ref.get("urn", {})
|
||||
# Use model name from URN as display name
|
||||
model_name = (
|
||||
f"urn:{urn.get('type', 'model')}:{urn.get('model_id')}"
|
||||
)
|
||||
category = urn.get("type", category)
|
||||
if category in URN_TYPE_MAP:
|
||||
category = URN_TYPE_MAP[category]
|
||||
|
||||
loaded_models.append(
|
||||
{
|
||||
"name": model_name,
|
||||
"category": category,
|
||||
"node_id": node_id,
|
||||
"widget_index": widget_index,
|
||||
"node_type": node_type,
|
||||
"node_title": ref.get("node_title", ""),
|
||||
"subgraph_id": ref.get("subgraph_id") or "",
|
||||
"subgraph_name": ref.get("subgraph_name") or "",
|
||||
"is_top_level": ref.get("is_top_level") is not False,
|
||||
"locate_node_id": ref.get("locate_node_id"),
|
||||
"locate_node_type": ref.get("locate_node_type", ""),
|
||||
"locate_node_title": ref.get("locate_node_title", ""),
|
||||
"locate_subgraph_id": (
|
||||
ref.get("locate_subgraph_id") or ""
|
||||
),
|
||||
"locate_subgraph_name": (
|
||||
ref.get("locate_subgraph_name") or ""
|
||||
),
|
||||
"locate_is_top_level": (
|
||||
ref.get("locate_is_top_level") is not False
|
||||
),
|
||||
"locate_via_promoted_widget": ref.get(
|
||||
"locate_via_promoted_widget", False
|
||||
),
|
||||
"exists": exists,
|
||||
"strength": strength,
|
||||
"original_path": original_path,
|
||||
"is_urn": ref.get("is_urn", False),
|
||||
"custom_node_adapter": ref.get(
|
||||
"custom_node_adapter"
|
||||
),
|
||||
"active": ref.get("active"),
|
||||
"connected": ref.get("connected", True),
|
||||
"resolved_path": (
|
||||
path_by_filename.get(model_name)
|
||||
or path_by_filename.get(original_path)
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
# Also check node.properties.models for embedded models
|
||||
total_nodes = len(nodes)
|
||||
for node_index, node in enumerate(nodes, start=1):
|
||||
if total_nodes and (
|
||||
node_index == total_nodes or node_index % 50 == 0
|
||||
):
|
||||
update_loaded_progress(
|
||||
"embedded",
|
||||
"Checking embedded model metadata...",
|
||||
percent=interpolate_percent(
|
||||
94, 98, node_index, total_nodes
|
||||
),
|
||||
current=node_index,
|
||||
total=total_nodes,
|
||||
)
|
||||
|
||||
node_type = node.get("type", "")
|
||||
properties = node.get("properties", {})
|
||||
if not isinstance(properties, dict):
|
||||
properties = {}
|
||||
models_list = properties.get("models", [])
|
||||
if not isinstance(models_list, list):
|
||||
models_list = []
|
||||
|
||||
for model_info in models_list:
|
||||
if isinstance(model_info, dict):
|
||||
name = model_info.get("name", "")
|
||||
directory = model_info.get("directory", "")
|
||||
node_scope = get_node_scope(node)
|
||||
|
||||
if name:
|
||||
# Check if this model is already in loaded_models
|
||||
existing = next(
|
||||
(
|
||||
model
|
||||
for model in loaded_models
|
||||
if model.get("original_path") == name
|
||||
),
|
||||
None,
|
||||
)
|
||||
if not existing:
|
||||
loaded_models.append(
|
||||
{
|
||||
"name": get_filename_from_path(name),
|
||||
"category": directory or "checkpoints",
|
||||
"node_id": node.get("id"),
|
||||
"widget_index": None,
|
||||
"node_type": node_type,
|
||||
"node_title": node.get("title", ""),
|
||||
"subgraph_id": node_scope["subgraph_id"],
|
||||
"subgraph_name": node_scope[
|
||||
"subgraph_name"
|
||||
],
|
||||
"is_top_level": node_scope[
|
||||
"is_top_level"
|
||||
],
|
||||
"exists": True, # Embedded models are loaded
|
||||
"strength": None,
|
||||
"original_path": name,
|
||||
"is_urn": False,
|
||||
}
|
||||
)
|
||||
|
||||
update_loaded_progress(
|
||||
"finalizing",
|
||||
"Loaded models ready",
|
||||
percent=99,
|
||||
current=len(loaded_models),
|
||||
total=len(loaded_models),
|
||||
)
|
||||
return {
|
||||
"loaded_models": loaded_models,
|
||||
"total": len(loaded_models),
|
||||
}
|
||||
|
||||
update_loaded_progress(
|
||||
"starting",
|
||||
"Preparing loaded model scan...",
|
||||
percent=0,
|
||||
status="starting",
|
||||
current=0,
|
||||
total=0,
|
||||
)
|
||||
|
||||
try:
|
||||
result = await asyncio.to_thread(build_loaded_models_response)
|
||||
except Exception as e:
|
||||
update_loaded_progress(
|
||||
"error",
|
||||
str(e),
|
||||
percent=100,
|
||||
status="error",
|
||||
current=0,
|
||||
total=0,
|
||||
)
|
||||
self.logger.error(
|
||||
f"Model Resolver loaded models error: {e}", exc_info=True
|
||||
)
|
||||
return web.json_response({"error": str(e)}, status=500)
|
||||
|
||||
update_loaded_progress(
|
||||
"completed",
|
||||
"Loaded models ready",
|
||||
percent=100,
|
||||
status="completed",
|
||||
current=result.get("total", 0),
|
||||
total=result.get("total", 0),
|
||||
)
|
||||
return web.json_response(result)
|
||||
|
||||
@routes.get("/model_resolver/loaded-progress/{loaded_id}")
|
||||
@json_api_endpoint("loaded-progress")
|
||||
async def get_loaded_models_progress(request):
|
||||
"""Get loaded models inspection progress."""
|
||||
return get_progress_response(
|
||||
self.loaded_progress,
|
||||
request,
|
||||
param_name="loaded_id",
|
||||
not_found_payload={
|
||||
"status": "unknown",
|
||||
"stage": "unknown",
|
||||
"message": "No loaded models progress available",
|
||||
"percent": 0,
|
||||
"current": 0,
|
||||
"total": 0,
|
||||
}
|
||||
)
|
||||
@@ -0,0 +1,97 @@
|
||||
"""Model details route registration."""
|
||||
|
||||
from .context import RouteContext
|
||||
|
||||
|
||||
def register_model_details_routes(context: RouteContext):
|
||||
asyncio = context.get('asyncio')
|
||||
download_available = context.get('download_available')
|
||||
get_civarchive_model_details = context.get('get_civarchive_model_details')
|
||||
get_civitai_model_details = context.get('get_civitai_model_details')
|
||||
get_huggingface_model_details = context.get('get_huggingface_model_details')
|
||||
json_api_endpoint = context.get('json_api_endpoint')
|
||||
routes = context.get('routes')
|
||||
web = context.get('web')
|
||||
|
||||
@routes.post("/model_resolver/model-details")
|
||||
@json_api_endpoint("model-details")
|
||||
async def model_details(request):
|
||||
"""Return normalized full model details for sources that expose model pages."""
|
||||
data = await request.json()
|
||||
source = str(data.get("source", "")).strip().lower()
|
||||
model_id = data.get("model_id")
|
||||
version_id = data.get("version_id")
|
||||
civitai_key = data.get("civitai_key", "")
|
||||
hf_token = data.get("hf_token", "")
|
||||
file_path = data.get("file_path", "")
|
||||
branch = data.get("branch", "")
|
||||
|
||||
if not download_available:
|
||||
return web.json_response(
|
||||
{"error": "Download providers are not available"}, status=503
|
||||
)
|
||||
|
||||
if source == "lora_manager_archive":
|
||||
source = "civitai"
|
||||
|
||||
if source not in {"civitai", "civarchive", "huggingface"}:
|
||||
return web.json_response(
|
||||
{"error": "Unsupported model details source"}, status=400
|
||||
)
|
||||
|
||||
if source == "huggingface":
|
||||
model_id = str(model_id or "").strip()
|
||||
branch = str(branch or version_id or "main").strip() or "main"
|
||||
version_id = branch
|
||||
else:
|
||||
try:
|
||||
model_id = (
|
||||
int(model_id)
|
||||
if model_id is not None and str(model_id).strip()
|
||||
else None
|
||||
)
|
||||
except (TypeError, ValueError):
|
||||
model_id = None
|
||||
|
||||
try:
|
||||
version_id = (
|
||||
int(version_id)
|
||||
if version_id is not None and str(version_id).strip()
|
||||
else None
|
||||
)
|
||||
except (TypeError, ValueError):
|
||||
version_id = None
|
||||
|
||||
if not model_id:
|
||||
return web.json_response(
|
||||
{"error": "model_id is required"}, status=400
|
||||
)
|
||||
|
||||
if source == "civitai":
|
||||
details = await asyncio.to_thread(
|
||||
get_civitai_model_details,
|
||||
model_id,
|
||||
version_id,
|
||||
civitai_key or None,
|
||||
)
|
||||
elif source == "civarchive":
|
||||
details = await asyncio.to_thread(
|
||||
get_civarchive_model_details,
|
||||
model_id,
|
||||
version_id,
|
||||
)
|
||||
else:
|
||||
details = await asyncio.to_thread(
|
||||
get_huggingface_model_details,
|
||||
model_id,
|
||||
file_path,
|
||||
branch,
|
||||
hf_token or None,
|
||||
)
|
||||
|
||||
if not details:
|
||||
return web.json_response(
|
||||
{"error": "Model details not found"}, status=404
|
||||
)
|
||||
|
||||
return web.json_response(details)
|
||||
+7
-1256
File diff suppressed because it is too large
Load Diff
+15
-9
@@ -122,18 +122,20 @@ def register_routes(self):
|
||||
get_workflow_model_inventory,
|
||||
)
|
||||
from .base_models import register_base_model_routes
|
||||
from .civitai_search import register_civitai_search_routes
|
||||
from .custom_url import register_custom_url_routes
|
||||
from .directories import register_directory_routes
|
||||
from .downloads import register_download_routes
|
||||
from .hashes import register_hash_routes
|
||||
from .helpers import create_route_helpers
|
||||
from .loaded_models import register_loaded_model_routes
|
||||
from .metadata import register_metadata_routes
|
||||
from .model_info import register_model_info_routes
|
||||
from .search import register_search_routes
|
||||
from .model_details import register_model_details_routes
|
||||
from .search_support import register_search_support_routes
|
||||
from .settings import register_settings_routes
|
||||
from .source_search import register_source_search_routes
|
||||
from .version import register_version_routes
|
||||
from .workflow import (
|
||||
register_loaded_model_routes,
|
||||
register_workflow_routes,
|
||||
)
|
||||
from .workflow_analysis import register_workflow_analysis_routes
|
||||
except ImportError as e:
|
||||
self.logger.error(f"Model Resolver: Could not import core modules: {e}")
|
||||
return False
|
||||
@@ -242,13 +244,17 @@ def register_routes(self):
|
||||
|
||||
register_base_model_routes(routes, web, json_api_endpoint)
|
||||
route_context = RouteContext.from_namespaces(globals(), locals())
|
||||
register_workflow_routes(route_context)
|
||||
register_workflow_analysis_routes(route_context)
|
||||
register_hash_routes(route_context)
|
||||
register_metadata_routes(route_context)
|
||||
register_loaded_model_routes(route_context)
|
||||
register_model_info_routes(route_context)
|
||||
register_civitai_search_routes(route_context)
|
||||
register_custom_url_routes(route_context)
|
||||
register_model_details_routes(route_context)
|
||||
|
||||
if download_available:
|
||||
register_search_routes(route_context)
|
||||
register_source_search_routes(route_context)
|
||||
register_search_support_routes(route_context)
|
||||
register_download_routes(route_context)
|
||||
register_directory_routes(route_context)
|
||||
|
||||
|
||||
+5
-1183
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,117 @@
|
||||
"""Search cache, credentials, and index route registration."""
|
||||
|
||||
from .context import RouteContext
|
||||
|
||||
|
||||
def register_search_support_routes(context: RouteContext):
|
||||
asyncio = context.get('asyncio')
|
||||
check_brave_search_api_key = context.get('check_brave_search_api_key')
|
||||
check_civitai_api_key = context.get('check_civitai_api_key')
|
||||
check_civitai_session_token = context.get('check_civitai_session_token')
|
||||
check_huggingface_token = context.get('check_huggingface_token')
|
||||
clear_all_search_caches = context.get('clear_all_search_caches')
|
||||
clear_huggingface_search_cache = context.get('clear_huggingface_search_cache')
|
||||
get_known_author_fallback_indexes_status = context.get('get_known_author_fallback_indexes_status')
|
||||
get_model_list_update_status = context.get('get_model_list_update_status')
|
||||
invalidate_local_hash_match_cache = context.get('invalidate_local_hash_match_cache')
|
||||
invalidate_model_files_cache = context.get('invalidate_model_files_cache')
|
||||
json_api_endpoint = context.get('json_api_endpoint')
|
||||
refresh_known_author_fallback_indexes = context.get('refresh_known_author_fallback_indexes')
|
||||
reload_model_list = context.get('reload_model_list')
|
||||
reload_popular_databases = context.get('reload_popular_databases')
|
||||
routes = context.get('routes')
|
||||
self = context.get('self')
|
||||
update_model_list_from_remote = context.get('update_model_list_from_remote')
|
||||
web = context.get('web')
|
||||
|
||||
@routes.post("/model_resolver/clear-search-cache")
|
||||
@json_api_endpoint("Clear search cache")
|
||||
async def clear_search_cache_route(request):
|
||||
"""Clear backend search caches after token/settings changes."""
|
||||
clear_all_search_caches()
|
||||
reload_popular_databases()
|
||||
reload_model_list()
|
||||
invalidate_model_files_cache()
|
||||
invalidate_local_hash_match_cache()
|
||||
self.search_result_timestamps.clear()
|
||||
self.logger.info("Cleared backend search caches")
|
||||
return web.json_response({"success": True, "cleared": "all"})
|
||||
|
||||
async def _check_credential_helper(request, payload_key, check_func, log_name):
|
||||
try:
|
||||
data = await request.json()
|
||||
val = data.get(payload_key, "")
|
||||
result = await asyncio.to_thread(check_func, val)
|
||||
return web.json_response(result)
|
||||
except Exception as e:
|
||||
self.logger.exception(f"{log_name} check error: {e}")
|
||||
return web.json_response({"error": str(e)}, status=500)
|
||||
|
||||
@routes.post("/model_resolver/civitai/session-token/check")
|
||||
async def civitai_session_token_check_route(request):
|
||||
"""Check whether a CivitAI browser session token is valid."""
|
||||
return await _check_credential_helper(
|
||||
request, "civitai_session_token", check_civitai_session_token, "CivitAI session token"
|
||||
)
|
||||
|
||||
@routes.post("/model_resolver/civitai/api-key/check")
|
||||
async def civitai_api_key_check_route(request):
|
||||
"""Check whether a CivitAI API key is valid."""
|
||||
return await _check_credential_helper(
|
||||
request, "civitai_key", check_civitai_api_key, "CivitAI API key"
|
||||
)
|
||||
|
||||
@routes.post("/model_resolver/huggingface/token/check")
|
||||
async def huggingface_token_check_route(request):
|
||||
"""Check whether a HuggingFace token is valid."""
|
||||
return await _check_credential_helper(
|
||||
request, "hf_token", check_huggingface_token, "HuggingFace token"
|
||||
)
|
||||
|
||||
@routes.post("/model_resolver/brave/api-key/check")
|
||||
async def brave_api_key_check_route(request):
|
||||
"""Check whether a Brave Search API key is valid."""
|
||||
return await _check_credential_helper(
|
||||
request, "brave_search_api_key", check_brave_search_api_key, "Brave Search API key"
|
||||
)
|
||||
|
||||
@routes.get("/model_resolver/huggingface/author-index/status")
|
||||
@json_api_endpoint("HuggingFace author index status")
|
||||
async def huggingface_author_index_status_route(request):
|
||||
"""Return local HuggingFace author fallback index status."""
|
||||
return web.json_response(
|
||||
get_known_author_fallback_indexes_status()
|
||||
)
|
||||
|
||||
@routes.post("/model_resolver/huggingface/author-index/refresh")
|
||||
@json_api_endpoint("HuggingFace author index refresh")
|
||||
async def huggingface_author_index_refresh_route(request):
|
||||
"""Refresh HuggingFace author fallback index."""
|
||||
data = await request.json()
|
||||
hf_token = data.get("hf_token", "")
|
||||
result = await asyncio.to_thread(
|
||||
refresh_known_author_fallback_indexes, hf_token or None
|
||||
)
|
||||
clear_huggingface_search_cache()
|
||||
return web.json_response(result)
|
||||
|
||||
@routes.get("/model_resolver/model-list/status")
|
||||
@json_api_endpoint("Model list status")
|
||||
async def model_list_status_route(request):
|
||||
"""Return local model-list status and optionally compare with GitHub."""
|
||||
check_remote = (
|
||||
str(request.query.get("check_remote", "")).lower()
|
||||
in {"1", "true", "yes"}
|
||||
)
|
||||
return web.json_response(
|
||||
get_model_list_update_status(check_remote=check_remote)
|
||||
)
|
||||
|
||||
@routes.post("/model_resolver/model-list/update")
|
||||
@json_api_endpoint("Model list update")
|
||||
async def model_list_update_route(request):
|
||||
"""Download latest ComfyUI-Manager model-list.json."""
|
||||
result = await asyncio.to_thread(update_model_list_from_remote)
|
||||
clear_all_search_caches()
|
||||
self.search_result_timestamps.clear()
|
||||
return web.json_response(result)
|
||||
File diff suppressed because it is too large
Load Diff
+7
-1390
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,269 @@
|
||||
"""Workflow analysis route registration."""
|
||||
|
||||
from .context import RouteContext
|
||||
|
||||
|
||||
def register_workflow_analysis_routes(context: RouteContext):
|
||||
analyze_and_find_matches = context.get('analyze_and_find_matches')
|
||||
apply_resolution = context.get('apply_resolution')
|
||||
asyncio = context.get('asyncio')
|
||||
download_available = context.get('download_available')
|
||||
fetch_remote_file_size_cached = context.get('fetch_remote_file_size_cached')
|
||||
get_filename_from_path = context.get('get_filename_from_path')
|
||||
get_popular_model_url = context.get('get_popular_model_url')
|
||||
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')
|
||||
routes = context.get('routes')
|
||||
search_local_matches = context.get('search_local_matches')
|
||||
search_model_list = context.get('search_model_list')
|
||||
self = context.get('self')
|
||||
should_skip_existing_custom_node_reference = context.get('should_skip_existing_custom_node_reference')
|
||||
to_bool = context.get('to_bool')
|
||||
web = context.get('web')
|
||||
|
||||
@routes.post("/model_resolver/analyze")
|
||||
async def analyze_workflow(request):
|
||||
"""Analyze workflow and return missing models with matches."""
|
||||
try:
|
||||
data = await request.json()
|
||||
workflow_json = data.get("workflow")
|
||||
analysis_id = str(data.get("analysis_id") or "").strip()
|
||||
force_rescan = to_bool(data.get("force_rescan"), False)
|
||||
if force_rescan:
|
||||
invalidate_local_hash_match_cache()
|
||||
|
||||
if workflow_json is None:
|
||||
return web.json_response(
|
||||
{"error": "Workflow JSON is required"}, status=400
|
||||
)
|
||||
if not isinstance(workflow_json, dict):
|
||||
return web.json_response(
|
||||
{"error": "Workflow JSON must be an object"}, status=400
|
||||
)
|
||||
|
||||
if analysis_id:
|
||||
self._update_analysis_progress(
|
||||
analysis_id,
|
||||
{
|
||||
"status": "starting",
|
||||
"stage": "starting",
|
||||
"message": "Starting analysis...",
|
||||
"current": 0,
|
||||
"total": 0,
|
||||
},
|
||||
)
|
||||
|
||||
def update_analysis_progress(payload):
|
||||
self._update_analysis_progress(analysis_id, payload)
|
||||
|
||||
# Analyze and find matches
|
||||
result = await asyncio.to_thread(
|
||||
analyze_and_find_matches,
|
||||
workflow_json,
|
||||
0.0,
|
||||
10,
|
||||
update_analysis_progress if analysis_id else None,
|
||||
force_rescan=force_rescan,
|
||||
)
|
||||
|
||||
missing_models = result.get("missing_models", [])
|
||||
filtered_missing = []
|
||||
for missing in missing_models:
|
||||
name = missing.get("name") or missing.get("original_path", "")
|
||||
if should_skip_existing_custom_node_reference(missing):
|
||||
self.logger.info(
|
||||
"Filtered existing custom-node model "
|
||||
f"reference: {name}"
|
||||
)
|
||||
continue
|
||||
filtered_missing.append(missing)
|
||||
result["missing_models"] = filtered_missing
|
||||
result["total_missing"] = len(filtered_missing)
|
||||
|
||||
# If download available, check for download sources only from LOCAL sources
|
||||
# (workflow_url, popular, model-list.json) - skip automatic online search
|
||||
# Online search is now only triggered on-demand via search button
|
||||
if download_available:
|
||||
for missing in result.get("missing_models", []):
|
||||
# Check if there's a 100% local match
|
||||
matches = missing.get("matches", [])
|
||||
has_perfect_match = any(
|
||||
m.get("confidence", 0) == 100 for m in matches
|
||||
)
|
||||
|
||||
if not has_perfect_match:
|
||||
filename = get_filename_from_path(
|
||||
missing.get("original_path", "")
|
||||
)
|
||||
|
||||
# 0. Check workflow URL first (highest priority - directly from workflow)
|
||||
workflow_url = missing.get("workflow_url", "")
|
||||
if workflow_url:
|
||||
# Determine source from URL
|
||||
if "huggingface.co" in workflow_url:
|
||||
source = "huggingface"
|
||||
elif "civitai.com" in workflow_url:
|
||||
source = "civitai"
|
||||
else:
|
||||
source = "workflow"
|
||||
|
||||
# Try to get file size using cached remote helper
|
||||
file_size = fetch_remote_file_size_cached(workflow_url, timeout=5)
|
||||
|
||||
missing["download_source"] = {
|
||||
"source": source,
|
||||
"url": workflow_url,
|
||||
"model_url": missing.get(
|
||||
"workflow_model_url", workflow_url
|
||||
),
|
||||
"filename": filename,
|
||||
"directory": missing.get(
|
||||
"workflow_directory", ""
|
||||
)
|
||||
or missing.get("category", "checkpoints"),
|
||||
"match_type": "exact",
|
||||
"url_source": "workflow",
|
||||
"size": file_size,
|
||||
}
|
||||
continue
|
||||
|
||||
# 1. Check popular models (always exact match)
|
||||
popular_info = get_popular_model_url(filename)
|
||||
if popular_info:
|
||||
popular_model_list_result = search_model_list(
|
||||
filename, exact_only=True
|
||||
)
|
||||
missing["download_source"] = {
|
||||
"source": "popular",
|
||||
"url": popular_info.get("url"),
|
||||
"filename": filename,
|
||||
"type": popular_info.get("type"),
|
||||
"directory": popular_info.get("directory"),
|
||||
"size": (
|
||||
popular_model_list_result.get("size")
|
||||
if popular_model_list_result
|
||||
else None
|
||||
)
|
||||
or popular_info.get("size"),
|
||||
"match_type": "exact",
|
||||
}
|
||||
continue
|
||||
|
||||
# 2. Check model list (ComfyUI Manager database)
|
||||
# Use exact_only=True to avoid confusing fuzzy matches for downloads
|
||||
model_list_result = search_model_list(
|
||||
filename, exact_only=True
|
||||
)
|
||||
if model_list_result:
|
||||
missing["download_source"] = {
|
||||
"source": "model_list",
|
||||
"url": model_list_result.get("url"),
|
||||
"filename": model_list_result.get("filename"),
|
||||
"name": model_list_result.get("name"),
|
||||
"type": model_list_result.get("type"),
|
||||
"directory": model_list_result.get("directory"),
|
||||
"size": model_list_result.get("size"),
|
||||
"match_type": model_list_result.get(
|
||||
"match_type"
|
||||
),
|
||||
"confidence": model_list_result.get(
|
||||
"confidence"
|
||||
),
|
||||
}
|
||||
continue
|
||||
|
||||
# NOTE: Search for online sources (HuggingFace, CivitAI) is
|
||||
# now done on-demand via /model_resolver/search endpoint
|
||||
# when user clicks "Search Online" button, not automatically
|
||||
|
||||
if analysis_id:
|
||||
self.analysis_progress.update(
|
||||
analysis_id,
|
||||
status="completed",
|
||||
stage="completed",
|
||||
message="Analysis complete",
|
||||
current=result.get("total_missing", 0),
|
||||
total=result.get("total_missing", 0),
|
||||
)
|
||||
|
||||
return web.json_response(result)
|
||||
except Exception as e:
|
||||
if "analysis_id" in locals() and analysis_id:
|
||||
self.analysis_progress.update(
|
||||
analysis_id,
|
||||
status="error",
|
||||
stage="error",
|
||||
message=str(e),
|
||||
current=0,
|
||||
total=0,
|
||||
)
|
||||
self.logger.error(f"Model Resolver analyze error: {e}", exc_info=True)
|
||||
return web.json_response({"error": str(e)}, status=500)
|
||||
|
||||
@routes.get("/model_resolver/analyze-progress/{analysis_id}")
|
||||
@json_api_endpoint("analyze-progress")
|
||||
async def get_analyze_progress(request):
|
||||
"""Get workflow analysis progress."""
|
||||
return get_progress_response(
|
||||
self.analysis_progress,
|
||||
request,
|
||||
param_name="analysis_id",
|
||||
not_found_payload={
|
||||
"status": "unknown",
|
||||
"stage": "unknown",
|
||||
"message": "No analysis progress available",
|
||||
"current": 0,
|
||||
"total": 0,
|
||||
}
|
||||
)
|
||||
|
||||
@routes.post("/model_resolver/resolve")
|
||||
@json_api_endpoint("resolve", return_success_on_error=True)
|
||||
async def resolve_models(request):
|
||||
"""Apply model resolution and return updated workflow."""
|
||||
data = await request.json()
|
||||
workflow_json = data.get("workflow")
|
||||
resolutions = data.get("resolutions", [])
|
||||
|
||||
if not workflow_json:
|
||||
return web.json_response(
|
||||
{"error": "Workflow JSON is required"}, status=400
|
||||
)
|
||||
|
||||
if not resolutions:
|
||||
return web.json_response(
|
||||
{"error": "Resolutions array is required"}, status=400
|
||||
)
|
||||
|
||||
# Apply resolutions
|
||||
updated_workflow = apply_resolution(workflow_json, resolutions)
|
||||
|
||||
return web.json_response(
|
||||
{"workflow": updated_workflow, "success": True}
|
||||
)
|
||||
|
||||
@routes.post("/model_resolver/local-matches")
|
||||
@json_api_endpoint("local-matches")
|
||||
async def local_matches(request):
|
||||
"""Search local model files by filename/path."""
|
||||
data = await request.json()
|
||||
filename = data.get("filename", "")
|
||||
category = data.get("category", "")
|
||||
force_rescan = to_bool(data.get("force_rescan"), False)
|
||||
if force_rescan:
|
||||
invalidate_local_hash_match_cache()
|
||||
|
||||
if not filename:
|
||||
return web.json_response(
|
||||
{"error": "filename is required"}, status=400
|
||||
)
|
||||
|
||||
matches = search_local_matches(
|
||||
filename,
|
||||
category=category or None,
|
||||
similarity_threshold=0.0,
|
||||
max_matches_per_model=10,
|
||||
force_rescan=force_rescan,
|
||||
)
|
||||
return web.json_response({"matches": matches})
|
||||
@@ -150,6 +150,55 @@ class TestRefactoringTargets(unittest.IsolatedAsyncioTestCase):
|
||||
},
|
||||
}
|
||||
|
||||
ROUTE_SUBGROUPS = {
|
||||
"workflow_analysis": {
|
||||
("POST", "/model_resolver/analyze"),
|
||||
("GET", "/model_resolver/analyze-progress/{analysis_id}"),
|
||||
("POST", "/model_resolver/resolve"),
|
||||
("POST", "/model_resolver/local-matches"),
|
||||
},
|
||||
"workflow_hashes": {
|
||||
("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}"),
|
||||
},
|
||||
"loaded_models": {
|
||||
("POST", "/model_resolver/loaded"),
|
||||
("GET", "/model_resolver/loaded-progress/{loaded_id}"),
|
||||
},
|
||||
"civitai_search": {
|
||||
("POST", "/model_resolver/civitai-search"),
|
||||
},
|
||||
"custom_url": {
|
||||
("POST", "/model_resolver/custom-url"),
|
||||
},
|
||||
"model_details": {
|
||||
("POST", "/model_resolver/model-details"),
|
||||
},
|
||||
"source_search": {
|
||||
("GET", "/model_resolver/search-progress/{progress_id}"),
|
||||
("POST", "/model_resolver/search-cancel/{progress_id}"),
|
||||
("POST", "/model_resolver/search"),
|
||||
},
|
||||
"search_support": {
|
||||
("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"),
|
||||
},
|
||||
}
|
||||
|
||||
def test_all_api_routes_are_registered(self):
|
||||
expected_routes = {
|
||||
("GET", "/model_resolver/base-models"),
|
||||
@@ -236,6 +285,20 @@ class TestRefactoringTargets(unittest.IsolatedAsyncioTestCase):
|
||||
len(grouped_routes),
|
||||
)
|
||||
|
||||
def test_large_route_groups_are_partitioned_without_gaps_or_overlaps(self):
|
||||
partitioned_routes = set().union(*self.ROUTE_SUBGROUPS.values())
|
||||
expected_routes = (
|
||||
self.ROUTE_GROUPS["workflow"]
|
||||
| self.ROUTE_GROUPS["loaded"]
|
||||
| self.ROUTE_GROUPS["model_info"]
|
||||
| self.ROUTE_GROUPS["search"]
|
||||
)
|
||||
self.assertEqual(partitioned_routes, expected_routes)
|
||||
self.assertEqual(
|
||||
sum(len(routes) for routes in self.ROUTE_SUBGROUPS.values()),
|
||||
len(partitioned_routes),
|
||||
)
|
||||
|
||||
def test_project_version_helpers_preserve_version_comparison_behavior(self):
|
||||
self.assertEqual(
|
||||
node_mod._extract_project_version('version = "1.2.3"'),
|
||||
@@ -332,6 +395,32 @@ class TestRefactoringTargets(unittest.IsolatedAsyncioTestCase):
|
||||
{"error": "URL is required"},
|
||||
status=400,
|
||||
)
|
||||
|
||||
async def test_local_model_hashes_route_rejects_missing_path(self):
|
||||
post_handler = routes_registered[("POST", "/model_resolver/local-model-hashes")]
|
||||
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": "path is required"},
|
||||
status=400,
|
||||
)
|
||||
|
||||
async def test_model_details_route_rejects_unsupported_source(self):
|
||||
post_handler = routes_registered[("POST", "/model_resolver/model-details")]
|
||||
request = AsyncMock()
|
||||
request.json.return_value = {"source": "unknown", "model_id": 1}
|
||||
|
||||
with patch("aiohttp.web.json_response") as mock_json_response:
|
||||
await post_handler(request)
|
||||
|
||||
mock_json_response.assert_called_once_with(
|
||||
{"error": "Unsupported model details source"},
|
||||
status=400,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
|
||||
Reference in New Issue
Block a user