nodes/Azornes-Comfyui-Model-Resolver
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Refactor: centralize metadata sidecar helpers and custom-result builders

Browse Source
This commit is contained in:
Dariusz L
2026-07-18 09:43:27 +02:00
parent be12494953
commit 719b3b3f43
15 changed files with 477 additions and 459 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+10 -352
View File
@@ -28,7 +28,7 @@ if not __package__ or __package__ == "":
from .core.log_system import LogLevel, create_module_logger
from .core.log_system import logger as backend_log_controller
from .core.log_system.config import LOG_LEVEL as BACKEND_DEFAULT_LOG_LEVEL
from .core.path_utils import get_filename_from_path
from .core.path_utils import get_filename_from_path, get_metadata_sidecar_path, get_safe_metadata_sidecar_path
# Web directory for JavaScript interface
WEB_DIRECTORY = "./web"
@@ -373,6 +373,7 @@ class ModelResolverExtension:
looks_like_model_file,
normalize_category_to_model_type,
normalize_sha256,
select_primary_model_file,
to_bool,
to_int,
)
@@ -389,9 +390,7 @@ class ModelResolverExtension:
get_all_progress,
get_aria2_status,
get_download_directory,
get_metadata_sidecar_path,
get_progress,
get_safe_metadata_sidecar_path,
is_allowed_model_download_filename,
normalize_download_category,
pause_download,
@@ -405,6 +404,7 @@ class ModelResolverExtension:
from .core.sources import clear_all_search_caches
from .core.sources.civarchive import (
CivArchiveSearchError,
build_civarchive_custom_result,
get_civarchive_model_details,
is_civarchive_available,
parse_civarchive_url,
@@ -416,11 +416,13 @@ class ModelResolverExtension:
clear_search_cache as clear_civarchive_search_cache,
)
from .core.sources.civitai import (
build_civitai_custom_result,
check_civitai_api_key,
check_civitai_session_token,
get_civitai_download_url,
get_civitai_model_details,
parse_civitai_url,
resolve_civitai_version_custom_result,
resolve_urn,
search_civitai,
search_civitai_for_file,
@@ -429,6 +431,7 @@ class ModelResolverExtension:
clear_search_cache as clear_civitai_search_cache,
)
from .core.sources.huggingface import (
build_huggingface_custom_result,
check_brave_search_api_key,
check_huggingface_token,
get_author_fallback_index_status,
@@ -2796,50 +2799,7 @@ class ModelResolverExtension:
)
return web.json_response(response)
def _select_custom_url_file(files, expected_filename=""):
"""Pick the most likely downloadable model file from provider details."""
if not isinstance(files, list):
files = [files] if files else []
candidates = [item for item in files if isinstance(item, dict)]
if not candidates:
return {}
expected = get_filename_from_path(expected_filename).lower()
def file_name(file_info):
return str(
file_info.get("name")
or file_info.get("filename")
or file_info.get("path")
or ""
).strip()
def has_download(file_info):
return bool(
file_info.get("download_url")
or file_info.get("downloadUrl")
or file_info.get("url")
)
if expected:
for file_info in candidates:
if get_filename_from_path(file_name(file_info)).lower() == expected:
return file_info
for file_info in candidates:
if file_info.get("primary") and has_download(file_info):
return file_info
for file_info in candidates:
if str(file_info.get("type") or "").lower() == "model" and has_download(file_info):
return file_info
for file_info in candidates:
if has_download(file_info):
return file_info
return candidates[0]
def _custom_result_timestamp():
return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
@@ -2849,313 +2809,11 @@ class ModelResolverExtension:
return _quote(str(value or "").replace("\\", "/"), safe="/")
def _build_civitai_custom_result(details, expected_filename="", api_key=None):
if not isinstance(details, dict):
return None
selected_version = (
details.get("selected_version")
if isinstance(details.get("selected_version"), dict)
else {}
)
version_id = details.get("version_id") or selected_version.get("id")
model_id = details.get("model_id")
file_info = _select_custom_url_file(
selected_version.get("files") or [],
expected_filename,
)
filename = (
file_info.get("name")
or file_info.get("filename")
or get_filename_from_path(expected_filename)
or details.get("name")
or f"civitai-{version_id or model_id}"
)
download_url = (
file_info.get("download_url")
or file_info.get("downloadUrl")
or (get_civitai_download_url(version_id, api_key) if version_id else "")
)
if not download_url:
return None
hashes = (
file_info.get("hashes")
if isinstance(file_info.get("hashes"), dict)
else {}
)
sha256 = (
file_info.get("sha256")
or file_info.get("hash")
or hashes.get("SHA256")
or hashes.get("sha256")
)
version_url = (
details.get("version_url")
or selected_version.get("url")
or (
f"https://civitai.com/models/{model_id}?modelVersionId={version_id}"
if model_id and version_id
else details.get("url")
)
)
return {
"source": "civitai",
"details_source": "civitai",
"model_id": model_id,
"version_id": version_id,
"name": details.get("name") or filename,
"version_name": selected_version.get("name") or "",
"type": details.get("type") or file_info.get("type") or "",
"filename": filename,
"url": version_url or details.get("url"),
"version_url": version_url or details.get("url"),
"download_url": download_url,
"size": file_info.get("size"),
"base_model": selected_version.get("base_model"),
"tags": details.get("tags") or [],
"trained_words": selected_version.get("trained_words") or [],
"images": details.get("images") or selected_version.get("images") or [],
"description": selected_version.get("description")
or details.get("description")
or "",
"model_description": details.get("description") or "",
"sha256": sha256,
"hashes": hashes,
"match_type": "custom_url",
"custom_url": True,
}
def _resolve_civitai_version_custom_result(version_id, expected_filename="", api_key=None):
try:
import requests
except Exception:
requests = None
if not requests or not version_id:
return None
headers = {}
if api_key:
headers["Authorization"] = f"Bearer {api_key}"
try:
response = requests.get(
f"https://civitai.com/api/v1/model-versions/{version_id}",
headers=headers,
timeout=20,
)
try:
if response.status_code != 200:
return None
data = response.json()
finally:
response.close()
except Exception as e:
self.logger.warning(
f"CivitAI custom URL version lookup failed: version_id={version_id}, error={e}"
)
return None
model_block = data.get("model") if isinstance(data.get("model"), dict) else {}
model_id = data.get("modelId") or data.get("model_id") or model_block.get("id")
if model_id:
try:
details = get_civitai_model_details(model_id, version_id, api_key)
result = _build_civitai_custom_result(
details,
expected_filename=expected_filename,
api_key=api_key,
)
if result:
return result
except Exception as e:
self.logger.warning(
f"CivitAI custom URL details lookup failed: model_id={model_id}, version_id={version_id}, error={e}"
)
file_info = _select_custom_url_file(data.get("files") or [], expected_filename)
hashes = (
file_info.get("hashes")
if isinstance(file_info.get("hashes"), dict)
else {}
)
sha256 = (
file_info.get("sha256")
or file_info.get("hash")
or hashes.get("SHA256")
or hashes.get("sha256")
)
filename = (
file_info.get("name")
or file_info.get("filename")
or get_filename_from_path(expected_filename)
or f"civitai-{version_id}"
)
download_url = (
file_info.get("downloadUrl")
or file_info.get("download_url")
or get_civitai_download_url(version_id, api_key)
)
model_url = (
f"https://civitai.com/models/{model_id}?modelVersionId={version_id}"
if model_id
else f"https://civitai.com/api/download/models/{version_id}"
)
return {
"source": "civitai",
"details_source": "civitai",
"model_id": model_id,
"version_id": version_id,
"name": model_block.get("name") or data.get("modelName") or filename,
"version_name": data.get("name") or "",
"type": model_block.get("type") or data.get("modelType") or file_info.get("type") or "",
"filename": filename,
"url": model_url,
"version_url": model_url,
"download_url": download_url,
"size": file_info.get("size")
or ((file_info.get("sizeKB") or 0) * 1024 if file_info.get("sizeKB") else None),
"base_model": data.get("baseModel"),
"tags": model_block.get("tags") or [],
"trained_words": data.get("trainedWords") or data.get("trained_words") or [],
"images": data.get("images") or [],
"description": data.get("description") or "",
"sha256": sha256,
"hashes": hashes,
"match_type": "custom_url",
"custom_url": True,
}
def _build_civarchive_custom_result(details, expected_filename=""):
if not isinstance(details, dict):
return None
selected_version = (
details.get("selected_version")
if isinstance(details.get("selected_version"), dict)
else {}
)
file_info = _select_custom_url_file(
selected_version.get("files") or [],
expected_filename,
)
download_urls = file_info.get("download_urls") or []
if not isinstance(download_urls, list):
download_urls = [download_urls]
download_url = (
file_info.get("download_url")
or file_info.get("downloadUrl")
or (download_urls[0] if download_urls else "")
)
if not download_url:
return None
filename = (
file_info.get("name")
or file_info.get("filename")
or get_filename_from_path(expected_filename)
or details.get("name")
or "model"
)
hashes = (
file_info.get("hashes")
if isinstance(file_info.get("hashes"), dict)
else {}
)
sha256 = (
file_info.get("sha256")
or file_info.get("hash")
or hashes.get("SHA256")
or hashes.get("sha256")
)
version_id = details.get("version_id") or selected_version.get("id")
model_id = details.get("model_id")
version_url = (
details.get("version_url")
or selected_version.get("url")
or (
f"https://civarchive.com/models/{model_id}?modelVersionId={version_id}"
if model_id and version_id
else details.get("url")
)
)
return {
"source": "civarchive",
"details_source": "civarchive",
"model_id": model_id,
"version_id": version_id,
"name": details.get("name") or filename,
"version_name": selected_version.get("name") or "",
"type": details.get("type") or file_info.get("type") or "",
"filename": filename,
"url": version_url or details.get("url"),
"version_url": version_url or details.get("url"),
"platform_url": details.get("platform_url")
or selected_version.get("platform_url"),
"download_url": download_url,
"download_urls": [url for url in download_urls if url],
"size": file_info.get("size"),
"base_model": selected_version.get("base_model"),
"tags": details.get("tags") or [],
"trained_words": selected_version.get("trained_words") or [],
"images": details.get("images") or selected_version.get("images") or [],
"description": selected_version.get("description")
or details.get("description")
or "",
"sha256": sha256,
"hash": sha256,
"hashes": hashes,
"platform": details.get("platform"),
"match_type": "custom_url",
"custom_url": True,
}
def _build_huggingface_custom_result(url, expected_filename="", token=None):
from urllib.parse import unquote
parsed = parse_huggingface_url(url)
if not parsed:
return None
repo_id = parsed.get("repo") or ""
branch = parsed.get("branch") or "main"
file_path = unquote(parsed.get("filename") or "")
if not repo_id or not file_path:
return None
filename = get_filename_from_path(file_path)
download_url = get_huggingface_download_url(repo_id, file_path, branch)
if not looks_like_model_file(download_url, expected_filename or filename):
return None
headers = {}
if token:
headers["Authorization"] = f"Bearer {token}"
size = fetch_remote_file_size_cached(
download_url,
headers=headers,
timeout=10,
)
page_url = (
f"https://huggingface.co/{repo_id}/blob/{branch}/{_quote_url_path(file_path)}"
)
return {
"source": "huggingface",
"details_source": "huggingface",
"repo_id": repo_id,
"repo": repo_id,
"path": file_path,
"filename": filename,
"name": repo_id,
"url": download_url,
"download_url": download_url,
"page_url": page_url,
"version_url": page_url,
"size": size,
"match_type": "custom_url",
"custom_url": True,
}
def _build_direct_custom_download_result(url, source, expected_filename=""):
from urllib.parse import urlparse
@@ -3275,14 +2933,14 @@ class ModelResolverExtension:
version_id,
civitai_key or None,
)
result = _build_civitai_custom_result(
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,
resolve_civitai_version_custom_result,
version_id,
expected_filename,
civitai_key or None,
@@ -3332,7 +2990,7 @@ class ModelResolverExtension:
version_id,
True,
)
result = _build_civarchive_custom_result(
result = build_civarchive_custom_result(
details,
expected_filename=expected_filename,
)
@@ -3359,7 +3017,7 @@ class ModelResolverExtension:
)
else:
result = await asyncio.to_thread(
_build_huggingface_custom_result,
build_huggingface_custom_result,
normalized_url,
expected_filename,
hf_token or None,
core/downloader.py
+1 -24
View File
@@ -32,6 +32,7 @@ from .path_utils import (
calculate_file_sha256,
get_comfy_root_path,
get_filename_from_path,
get_metadata_sidecar_path,
get_path_identity,
is_path_within,
read_json_safe,
@@ -346,30 +347,6 @@ def _normalise_metadata_file_path(path_value: str) -> str:
return str(path_value or "").replace(os.sep, "/")
def get_metadata_sidecar_path(file_path: str) -> str:
"""Return the LoRA Manager-compatible sidecar path for a model file."""
base_path, _extension = os.path.splitext(file_path)
return f"{base_path}.metadata.json"
def get_safe_metadata_sidecar_path(file_path: str) -> str:
"""Return the canonical sidecar path without allowing caller-selected targets."""
raw_model_path = str(file_path or "").strip()
if not raw_model_path:
raise ValueError("A model path is required")
model_path = os.path.realpath(os.path.abspath(raw_model_path))
model_dir = os.path.realpath(os.path.dirname(model_path))
metadata_path = os.path.realpath(get_metadata_sidecar_path(model_path))
if (
not model_dir
or os.path.dirname(metadata_path) != model_dir
or metadata_path == model_path
or not is_path_within(metadata_path, model_dir)
):
raise ValueError("Metadata path is outside the model directory")
return metadata_path
def _resolve_lora_manager_model_type(category: str, source_type: Any = "") -> str:
res = normalize_category_to_model_type(category)
if res in ("checkpoint", "diffusion_model", "embedding"):
core/metadata_audit.py
+1 -23
View File
@@ -10,7 +10,7 @@ from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Any, Dict, Iterable, List, Optional, Tuple
from .log_system import create_module_logger
from .path_utils import get_filename_from_path, get_path_identity, read_json_safe
from .path_utils import _metadata_sidecar_paths, get_filename_from_path, get_path_identity, read_json_safe
from .type_utils import MODEL_EXTENSIONS, format_size_bytes
log = create_module_logger(__name__)
@@ -161,28 +161,6 @@ def _is_model_file_path(path: str) -> bool:
return os.path.splitext(filename)[1].lower() in MODEL_EXTENSIONS
def _metadata_sidecar_paths(model_path: str) -> List[str]:
directory = os.path.dirname(model_path)
filename = get_filename_from_path(model_path)
base_name = os.path.splitext(filename)[0]
candidates = [
os.path.join(directory, f"{base_name}.metadata.json"),
os.path.join(directory, f"{filename}.metadata.json"),
]
result: List[str] = []
seen = set()
for candidate in candidates:
try:
key = os.path.normcase(os.path.abspath(candidate))
except (OSError, ValueError):
key = os.path.normcase(candidate)
if key in seen:
continue
seen.add(key)
result.append(candidate)
return result
def _format_signed_size(bytes_value: int) -> str:
if bytes_value == 0:
core/metadata_builder.py
+2 -28
View File
@@ -15,9 +15,11 @@ from typing import Any, Callable, Dict, List, Optional, Tuple
from .log_system import create_module_logger
from .path_utils import (
HashCalculationCancelled,
_metadata_sidecar_paths,
calculate_file_sha256,
extract_safetensors_header_metadata,
get_filename_from_path,
get_metadata_sidecar_path,
get_path_identity,
read_json_safe,
read_safetensors_header,
@@ -44,34 +46,6 @@ MAX_METADATA_BUILD_WORKERS = 64
DEFAULT_METADATA_BUILD_WORKER_LIMIT = 4
def _metadata_sidecar_paths(model_path: str) -> List[str]:
directory = os.path.dirname(model_path)
filename = get_filename_from_path(model_path)
base_name = os.path.splitext(filename)[0]
candidates = [
os.path.join(directory, f"{base_name}.metadata.json"),
os.path.join(directory, f"{filename}.metadata.json"),
]
result: List[str] = []
seen = set()
for candidate in candidates:
try:
key = os.path.normcase(os.path.abspath(candidate))
except (OSError, ValueError):
key = os.path.normcase(candidate)
if key in seen:
continue
seen.add(key)
result.append(candidate)
return result
def get_metadata_sidecar_path(model_path: str) -> str:
"""Return the canonical local metadata sidecar path for a model file."""
return _metadata_sidecar_paths(model_path)[0]
def _select_metadata_path(model_path: str) -> Tuple[str, bool]:
for candidate in _metadata_sidecar_paths(model_path):
if os.path.isfile(candidate):
core/path_utils.py
+47
View File
@@ -1076,3 +1076,50 @@ def split_path_segments(path_value: Any, filter_dots: bool = True) -> list[str]:
return parts
def _metadata_sidecar_paths(model_path: str) -> List[str]:
directory = os.path.dirname(model_path)
filename = get_filename_from_path(model_path)
base_name = os.path.splitext(filename)[0]
candidates = [
os.path.join(directory, f"{base_name}.metadata.json"),
os.path.join(directory, f"{filename}.metadata.json"),
]
result: List[str] = []
seen = set()
for candidate in candidates:
try:
key = os.path.normcase(os.path.abspath(candidate))
except (OSError, ValueError):
key = os.path.normcase(candidate)
if key in seen:
continue
seen.add(key)
result.append(candidate)
return result
def get_metadata_sidecar_path(model_path: str) -> str:
"""Return the canonical local metadata sidecar path for a model file."""
return _metadata_sidecar_paths(model_path)[0]
def get_safe_metadata_sidecar_path(file_path: str) -> str:
"""Return the canonical sidecar path without allowing caller-selected targets."""
raw_model_path = str(file_path or "").strip()
if not raw_model_path:
raise ValueError("A model path is required")
model_path = os.path.realpath(os.path.abspath(raw_model_path))
model_dir = os.path.realpath(os.path.dirname(model_path))
metadata_path = os.path.realpath(get_metadata_sidecar_path(model_path))
if (
not model_dir
or os.path.dirname(metadata_path) != model_dir
or metadata_path == model_path
or not is_path_within(metadata_path, model_dir)
):
raise ValueError("Metadata path is outside the model directory")
return metadata_path
core/sources/civarchive.py
+89
View File
@@ -2284,3 +2284,92 @@ def search_civarchive_for_file(
candidate_limit=detail_limit,
)
return best_match
def build_civarchive_custom_result(
details: Dict[str, Any],
expected_filename: str = "",
) -> Optional[Dict[str, Any]]:
if not isinstance(details, dict):
return None
selected_version = (
details.get("selected_version")
if isinstance(details.get("selected_version"), dict)
else {}
)
file_info = select_primary_model_file(
selected_version.get("files") or [],
expected_filename, require_download=True
) or {}
download_urls = file_info.get("download_urls") or []
if not isinstance(download_urls, list):
download_urls = [download_urls]
download_url = (
file_info.get("download_url")
or file_info.get("downloadUrl")
or (download_urls[0] if download_urls else "")
)
if not download_url:
return None
filename = (
file_info.get("name")
or file_info.get("filename")
or get_filename_from_path(expected_filename)
or details.get("name")
or "model"
)
hashes = (
file_info.get("hashes")
if isinstance(file_info.get("hashes"), dict)
else {}
)
sha256 = (
file_info.get("sha256")
or file_info.get("hash")
or hashes.get("SHA256")
or hashes.get("sha256")
)
version_id = details.get("version_id") or selected_version.get("id")
model_id = details.get("model_id")
version_url = (
details.get("version_url")
or selected_version.get("url")
or (
f"https://civarchive.com/models/{model_id}?modelVersionId={version_id}"
if model_id and version_id
else details.get("url")
)
)
return {
"source": "civarchive",
"details_source": "civarchive",
"model_id": model_id,
"version_id": version_id,
"name": details.get("name") or filename,
"version_name": selected_version.get("name") or "",
"type": details.get("type") or file_info.get("type") or "",
"filename": filename,
"url": version_url or details.get("url"),
"version_url": version_url or details.get("url"),
"platform_url": details.get("platform_url")
or selected_version.get("platform_url"),
"download_url": download_url,
"download_urls": [url for url in download_urls if url],
"size": file_info.get("size"),
"base_model": selected_version.get("base_model"),
"tags": details.get("tags") or [],
"trained_words": selected_version.get("trained_words") or [],
"images": details.get("images") or selected_version.get("images") or [],
"description": selected_version.get("description")
or details.get("description")
or "",
"sha256": sha256,
"hash": sha256,
"hashes": hashes,
"platform": details.get("platform"),
"match_type": "custom_url",
"custom_url": True,
}
core/sources/civitai.py
+133 -1
View File
@@ -10,7 +10,7 @@ import re
from typing import Any, Callable, Dict, List, Optional
from urllib.parse import quote, urlencode
import requests # noqa: F401 - used indirectly via unittest.mock patch targets
import requests
from ..log_system import create_module_logger
from ..matcher import (
@@ -1972,3 +1972,135 @@ def get_model_info_for_file(
return result
return _build_local_model_info(file_path, require_inferred=True)
def build_civitai_custom_result(
details: Dict[str, Any],
expected_filename: str = "",
api_key: Optional[str] = None,
) -> Optional[Dict[str, Any]]:
if not isinstance(details, dict):
return None
selected_version = (
details.get("selected_version")
if isinstance(details.get("selected_version"), dict)
else {}
)
version_id = details.get("version_id") or selected_version.get("id")
model_id = details.get("model_id")
file_info = select_primary_model_file(
selected_version.get("files") or [],
expected_filename, require_download=True
) or {}
filename = (
file_info.get("name")
or file_info.get("filename")
or get_filename_from_path(expected_filename)
or details.get("name")
or f"civitai-{version_id or model_id}"
)
download_url = (
file_info.get("download_url")
or file_info.get("downloadUrl")
or (get_civitai_download_url(version_id, api_key) if version_id else "")
)
if not download_url:
return None
hashes = (
file_info.get("hashes")
if isinstance(file_info.get("hashes"), dict)
else {}
)
sha256 = (
file_info.get("sha256")
or file_info.get("hash")
or hashes.get("SHA256")
or hashes.get("sha256")
)
version_url = (
details.get("version_url")
or selected_version.get("url")
or (
f"https://civitai.com/models/{model_id}?modelVersionId={version_id}"
if model_id and version_id
else details.get("url")
)
)
return {
"source": "civitai",
"details_source": "civitai",
"model_id": model_id,
"version_id": version_id,
"name": details.get("name") or filename,
"version_name": selected_version.get("name") or "",
"type": details.get("type") or file_info.get("type") or "",
"filename": filename,
"url": version_url or details.get("url"),
"version_url": version_url or details.get("url"),
"download_url": download_url,
"size": file_info.get("size"),
"base_model": selected_version.get("base_model"),
"tags": details.get("tags") or [],
"trained_words": selected_version.get("trained_words") or [],
"images": details.get("images") or selected_version.get("images") or [],
"description": selected_version.get("description")
or details.get("description")
or "",
"model_description": details.get("description") or "",
"sha256": sha256,
"hashes": hashes,
"match_type": "custom_url",
"custom_url": True,
}
def resolve_civitai_version_custom_result(
version_id: int,
expected_filename: str = "",
api_key: Optional[str] = None,
) -> Optional[Dict[str, Any]]:
if not version_id:
return None
headers = {}
if api_key:
headers["Authorization"] = f"Bearer {api_key}"
try:
response = requests.get(
f"https://civitai.com/api/v1/model-versions/{version_id}",
headers=headers,
timeout=20,
)
try:
if response.status_code != 200:
return None
data = response.json()
finally:
response.close()
except Exception as e:
log.warning(
f"CivitAI custom URL version lookup failed: version_id={version_id}, error={e}"
)
return None
model_block = data.get("model") if isinstance(data.get("model"), dict) else {}
model_id = data.get("modelId") or data.get("model_id") or model_block.get("id")
if model_id:
try:
details = get_civitai_model_details(model_id, version_id, api_key)
result = build_civitai_custom_result(
details,
expected_filename=expected_filename,
api_key=api_key,
)
if result:
return result
except Exception as e:
log.warning(
f"CivitAI custom URL details lookup failed: model_id={model_id}, version_id={version_id}, error={e}"
)
return None
core/sources/huggingface.py
+57
View File
@@ -24,6 +24,7 @@ from ..type_utils import (
clear_remote_size_cache,
extract_file_size,
fetch_remote_file_size_cached,
looks_like_model_file,
prepare_remote_size_probe_url,
)
@@ -1030,3 +1031,59 @@ def get_repo_files(repo: str, token: Optional[str] = None) -> List[Dict[str, Any
log.error(f"Error getting repo files: {e}")
return files
def build_huggingface_custom_result(
url: str,
expected_filename: str = "",
token: Optional[str] = None,
) -> Optional[Dict[str, Any]]:
from urllib.parse import quote, unquote
parsed = parse_huggingface_url(url)
if not parsed:
return None
repo_id = parsed.get("repo") or ""
branch = parsed.get("branch") or "main"
file_path = unquote(parsed.get("filename") or "")
if not repo_id or not file_path:
return None
filename = get_filename_from_path(file_path)
download_url = get_huggingface_download_url(repo_id, file_path, branch)
if not looks_like_model_file(download_url, expected_filename or filename):
return None
headers = {}
if token:
headers["Authorization"] = f"Bearer {token}"
size = fetch_remote_file_size_cached(
download_url,
headers=headers,
timeout=10,
)
def quote_url_path(val):
return quote(str(val or "").replace("\\", "/"), safe="/")
page_url = (
f"https://huggingface.co/{repo_id}/blob/{branch}/{quote_url_path(file_path)}"
)
return {
"source": "huggingface",
"details_source": "huggingface",
"repo_id": repo_id,
"repo": repo_id,
"path": file_path,
"filename": filename,
"name": repo_id,
"url": download_url,
"download_url": download_url,
"page_url": page_url,
"version_url": page_url,
"size": size,
"match_type": "custom_url",
"custom_url": True,
}
core/type_utils.py
+37 -8
View File
@@ -283,31 +283,60 @@ def normalize_download_category(category: str) -> str:
return CATEGORY_MAP.get(token, token or "checkpoints")
def select_primary_model_file(files: List[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
def select_primary_model_file(
files: List[Dict[str, Any]],
expected_filename: Optional[str] = None,
require_download: bool = False,
) -> Optional[Dict[str, Any]]:
"""
Selects the primary file from a list of model version files.
Selects the primary file from a list of model version files,
optionally filtering by expected filename or requiring a download URL.
"""
valid_files = [f for f in files if isinstance(f, dict) and (f.get("name") or f.get("filename"))]
if not isinstance(files, list):
files = [files] if files else []
valid_files = [f for f in files if isinstance(f, dict)]
if not valid_files:
return None
def file_name(item):
return str(item.get("name") or item.get("filename") or "")
def has_download(item):
return bool(item.get("downloadUrl") or item.get("download_url") or item.get("url"))
if require_download:
candidates = [f for f in valid_files if has_download(f)]
else:
candidates = [f for f in valid_files if file_name(f)]
if not candidates:
candidates = valid_files
expected = str(expected_filename or "").strip().lower()
if expected:
from .path_utils import get_filename_from_path
for file_info in candidates:
if get_filename_from_path(file_name(file_info)).lower() == expected:
return file_info
# 1. Look for file marked primary and of type Model
for file_info in valid_files:
for file_info in candidates:
if file_info.get("primary") and str(file_info.get("type") or "").lower() == "model":
return file_info
# 2. Look for any file marked primary
for file_info in valid_files:
for file_info in candidates:
if file_info.get("primary"):
return file_info
# 3. Look for any file of type Model
for file_info in valid_files:
for file_info in candidates:
if str(file_info.get("type") or "").lower() == "model":
return file_info
# 4. Fallback to the first file
return valid_files[0]
# 4. Fallback to the first candidate
return candidates[0]
def parse_size_to_bytes(value: Any) -> Optional[int]:
tests/test_path_security.py
+3 -3
View File
@@ -1,5 +1,5 @@
import os
import io
import os
import tarfile
import tempfile
import unittest
@@ -8,6 +8,7 @@ from types import SimpleNamespace
from unittest.mock import MagicMock, patch
import requests
import core.downloader as downloader_module
from core.aria2_installer import Aria2InstallError, _safe_extract_tar
from core.downloader import (
@@ -18,7 +19,6 @@ from core.downloader import (
download_file,
download_file_with_aria2,
download_model,
get_safe_metadata_sidecar_path,
is_allowed_model_download_filename,
sanitize_download_filename,
)
@@ -28,7 +28,7 @@ from core.network_utils import (
request_public_url,
validate_public_http_url,
)
from core.path_utils import is_path_in_configured_model_roots
from core.path_utils import get_safe_metadata_sidecar_path, is_path_in_configured_model_roots
class DummyFolderPaths:
tests/test_refactoring_targets.py
+77 -11
View File
@@ -1,8 +1,8 @@
import sys
import os
import unittest
from unittest.mock import MagicMock, patch, AsyncMock
import importlib
import os
import sys
import unittest
from unittest.mock import AsyncMock, MagicMock, patch
# Make sure parent package directory is in sys.path
parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
@@ -52,9 +52,13 @@ mock_routes.post = post_decorator
# Import the module so that routes register
node_mod = importlib.import_module("comfyui-model-resolver")
# Instantiating the extension and initializing routes to register them
extension = node_mod.ModelResolverExtension()
extension.setup_routes()
# Force setup_routes to run to populate routes_registered mock registry
node_mod.extension.routes_setup = False
node_mod.extension.setup_routes()
extension = node_mod.extension
class TestRefactoringTargets(unittest.IsolatedAsyncioTestCase):
@@ -64,8 +68,9 @@ class TestRefactoringTargets(unittest.IsolatedAsyncioTestCase):
patch_get_aria2_status.stop()
def test_downloader_calculate_file_sha256(self):
from core.downloader import calculate_file_sha256
import tempfile
from core.downloader import calculate_file_sha256
with tempfile.TemporaryDirectory() as tmpdir:
file_path = os.path.join(tmpdir, "test.txt")
with open(file_path, "wb") as f:
@@ -143,11 +148,21 @@ class TestRefactoringTargets(unittest.IsolatedAsyncioTestCase):
async def test_search_progress_route(self):
progress_id = "test_search_job"
extension.search_tracker.update(progress_id, status="running", stage="civitai", message="Searching Civitai", percent=50)
get_handler = routes_registered.get(("GET", "/model_resolver/search-progress/{progress_id}"))
self.assertIsNotNone(get_handler)
bound_extension = None
if hasattr(get_handler, "__closure__") and get_handler.__closure__:
for cell in get_handler.__closure__:
val = cell.cell_contents
if hasattr(val, "search_tracker"):
bound_extension = val
break
target_extension = bound_extension or extension
target_extension.search_tracker.update(progress_id, status="running", stage="civitai", message="Searching Civitai", percent=50)
mock_request = MagicMock()
mock_request.match_info = {"progress_id": progress_id}
@@ -161,11 +176,21 @@ class TestRefactoringTargets(unittest.IsolatedAsyncioTestCase):
async def test_search_cancel_route(self):
progress_id = "test_cancel_job"
extension.search_tracker.update(progress_id, status="running")
cancel_handler = routes_registered.get(("POST", "/model_resolver/search-cancel/{progress_id}"))
self.assertIsNotNone(cancel_handler)
bound_extension = None
if hasattr(cancel_handler, "__closure__") and cancel_handler.__closure__:
for cell in cancel_handler.__closure__:
val = cell.cell_contents
if hasattr(val, "search_tracker"):
bound_extension = val
break
target_extension = bound_extension or extension
target_extension.search_tracker.update(progress_id, status="running")
mock_request = MagicMock()
mock_request.match_info = {"progress_id": progress_id}
@@ -175,7 +200,8 @@ class TestRefactoringTargets(unittest.IsolatedAsyncioTestCase):
response_data = mock_json_res.call_args[0][0]
self.assertTrue(response_data.get("success"))
self.assertTrue(response_data.get("cancelled"))
self.assertTrue(extension.search_tracker.is_cancelled(progress_id))
self.assertTrue(target_extension.search_tracker.is_cancelled(progress_id))
async def test_root_directories_route_skips_non_model_categories(self):
import tempfile
@@ -234,3 +260,43 @@ class TestRefactoringTargets(unittest.IsolatedAsyncioTestCase):
self.assertIn("my_new_cat", categories)
self.assertNotIn("custom_nodes", categories)
self.assertNotIn("unet", categories)
def test_metadata_sidecar_paths_behavior(self):
from core.metadata_audit import _metadata_sidecar_paths as ma_sidecars
from core.metadata_builder import _metadata_sidecar_paths as mb_sidecars
from core.path_utils import get_metadata_sidecar_path, get_safe_metadata_sidecar_path
model_path = os.path.abspath("e:/models/checkpoints/sd15.safetensors")
expected_sidecar = os.path.abspath("e:/models/checkpoints/sd15.metadata.json")
self.assertEqual(get_metadata_sidecar_path(model_path), expected_sidecar)
# Safe sidecar path checks
safe_path = get_safe_metadata_sidecar_path(model_path)
self.assertEqual(os.path.normcase(safe_path), os.path.normcase(expected_sidecar))
# Test MB and MA sidecar candidates
mb_candidates = mb_sidecars(model_path)
ma_candidates = ma_sidecars(model_path)
self.assertEqual(mb_candidates, ma_candidates)
self.assertIn(expected_sidecar, mb_candidates)
def test_select_primary_model_file_extended(self):
from core.type_utils import select_primary_model_file
files = [
{"id": 1, "name": "other.safetensors", "primary": False, "type": "model", "url": "http://x"},
{"id": 2, "name": "main.safetensors", "primary": True, "type": "model"},
{"id": 3, "name": "config.json", "primary": False, "type": "config"}
]
# Original logic matches primary: True
self.assertEqual(select_primary_model_file(files)["id"], 2)
# Expected filename matches by name
self.assertEqual(select_primary_model_file(files, expected_filename="other.safetensors")["id"], 1)
# Require download excludes main.safetensors because it has no download url
self.assertEqual(select_primary_model_file(files, require_download=True)["id"], 1)
web/resolver/utils/html_utils.js
+12
View File
@@ -264,3 +264,15 @@ export async function copyTextWithFeedback(text, button, options = {}) {
}, duration);
}
/**
* Extracts a filename from a path (handling both windows and unix slash styles).
* @param {string} path
* @returns {string}
*/
export function getFilenameFromPath(path) {
if (!path) return '';
return path.split(/[\/\\]/).pop() || path;
}
web/resolver/utils/render_format_methods.js
+2 -3
View File
@@ -2,7 +2,7 @@ import { app } from "../../../../../scripts/app.js";
import { api } from "../../../../../scripts/api.js";
import { $el } from "../../../../../scripts/ui.js";
import { getSvgIcon } from "../../utils/icon_utils.js";
import { escapeHtml, escapeJsString, pollBackgroundTask } from "./html_utils.js";
import { escapeHtml, escapeJsString, getFilenameFromPath, pollBackgroundTask } from "./html_utils.js";
export const renderFormatMethods = {
encodeContextMenuModel(context = null) {
@@ -80,8 +80,7 @@ export const renderFormatMethods = {
},
getFilenameFromPath(path) {
if (!path) return '';
return path.split(/[\/\\]/).pop() || path;
return getFilenameFromPath(path);
},
/**
web/resolver/views/model_info_methods.js
+5 -5
View File
@@ -386,8 +386,8 @@ export const modelInfoMethods = {
getContextMenuSourceLookupFilename(model = {}) {
return model?.filename
|| model?.name
|| model?.original_path?.split(/[\/\\]/).pop()
|| String(this.getContextMenuSourceLookupPath(model) || '').split(/[\/\\]/).pop()
|| this.getFilenameFromPath(model?.original_path)
|| this.getFilenameFromPath(this.getContextMenuSourceLookupPath(model))
|| '';
},
@@ -643,7 +643,7 @@ export const modelInfoMethods = {
getHashCompareFilename(model = {}) {
return model.filename
|| model.name
|| String(this.getLocalHashComparePath(model) || '').split(/[\/\\]/).pop()
|| this.getFilenameFromPath(this.getLocalHashComparePath(model))
|| 'Selected local model';
},
@@ -1071,7 +1071,7 @@ export const modelInfoMethods = {
return;
}
const name = model.name || model.original_path?.split(/[\/\\]/).pop() || '';
const name = model.name || this.getFilenameFromPath(model.original_path) || '';
if (!name) return;
try {
@@ -1125,7 +1125,7 @@ export const modelInfoMethods = {
async showModelInfo(model) {
if (!model) return;
const name = model.name || model.original_path?.split(/[\/\\]/).pop() || '';
const name = model.name || this.getFilenameFromPath(model.original_path) || '';
if (!name) return;
// Create and show the info dialog
web/resolver/views/tabs_loaded_methods.js
+1 -1
View File
@@ -158,7 +158,7 @@ export const tabsLoadedMethods = {
},
getModelNameAndStrength(model) {
const fullName = model.name || model.original_path?.split(/[\/\\]/).pop() || 'Unknown';
const fullName = model.name || this.getFilenameFromPath(model.original_path) || 'Unknown';
const name = this.stripModelExtension(fullName);
const strength = model.strength !== null && model.strength !== undefined
? model.strength.toFixed(2)