refactor: extract model route handlers into ModelService
This commit is contained in:
Dariusz L
@@ -1,899 +1,16 @@
|
||||
"""CivitAI and exact model metadata route registration."""
|
||||
|
||||
from ..services.model_service import ModelService
|
||||
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')
|
||||
json_api_endpoint = context.get("json_api_endpoint")
|
||||
routes = context.get("routes")
|
||||
model_service = ModelService(context)
|
||||
|
||||
@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)
|
||||
"""Delegate model operation to the model service."""
|
||||
return await model_service.civitai_search(request)
|
||||
|
||||
+6
-282
@@ -1,292 +1,16 @@
|
||||
"""Custom model URL route registration."""
|
||||
|
||||
from ..services.model_service import ModelService
|
||||
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
|
||||
]
|
||||
json_api_endpoint = context.get("json_api_endpoint")
|
||||
routes = context.get("routes")
|
||||
model_service = ModelService(context)
|
||||
|
||||
@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)
|
||||
"""Delegate model operation to the model service."""
|
||||
return await model_service.custom_url(request)
|
||||
|
||||
@@ -1,97 +1,16 @@
|
||||
"""Model details route registration."""
|
||||
|
||||
from ..services.model_service import ModelService
|
||||
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')
|
||||
json_api_endpoint = context.get("json_api_endpoint")
|
||||
routes = context.get("routes")
|
||||
model_service = ModelService(context)
|
||||
|
||||
@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)
|
||||
"""Delegate model operation to the model service."""
|
||||
return await model_service.model_details(request)
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user