Dariusz L
dca6bc93fb
Use order-insensitive token comparison in calculate_filename_confidence so variant names with reordered tokens (e.g., int8_convrot vs convrot_int8) score higher. Replace first-match partial matching in HuggingFace source with best-match selection based on filename confidence, and surface confidence in search results.
1600 lines
53 KiB
Python
1600 lines
53 KiB
Python
"""
|
|
HuggingFace Source Module
|
|
|
|
Search and download models from HuggingFace Hub.
|
|
"""
|
|
|
|
import os
|
|
import posixpath
|
|
import re
|
|
import threading
|
|
import time
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
from typing import Any, Callable, Dict, List, Optional
|
|
from urllib.parse import quote, unquote, urlparse
|
|
|
|
from ..contracts import HuggingFaceFileReference, SearchResult
|
|
from ..log_system import create_module_logger
|
|
from ..progress import get_progress_reporter
|
|
|
|
log = create_module_logger(__name__)
|
|
|
|
from ..matcher import (
|
|
build_filename_search_queries,
|
|
calculate_filename_confidence,
|
|
)
|
|
from ..network_utils import execute_provider_json_request, host_matches_domain, request_source_response
|
|
from ..path_utils import METADATA_DIR, get_filename_from_path, read_json_safe, write_json_atomic
|
|
from ..type_utils import (
|
|
build_model_result,
|
|
check_credential_http,
|
|
clear_remote_size_cache,
|
|
extract_file_size,
|
|
looks_like_model_file,
|
|
normalize_sha256,
|
|
)
|
|
from .common import probe_remote_file_size, resolve_file_size
|
|
|
|
HF_API_URL = "https://huggingface.co/api"
|
|
HF_AUTHOR_FALLBACKS = ["Comfy-Org", "Kijai"]
|
|
BRAVE_SEARCH_API_URL = "https://api.search.brave.com/res/v1/web/search"
|
|
HF_AUTHOR_INDEX_CACHE_TTL_SECONDS = 24 * 60 * 60
|
|
HF_AUTHOR_INDEX_CACHE_VERSION = 2
|
|
|
|
HF_AUTHOR_INDEX_CACHE_PATH = os.path.join(
|
|
METADATA_DIR, "huggingface-author-index.json"
|
|
)
|
|
|
|
# Cache for search results
|
|
_search_cache: Dict[str, Any] = {}
|
|
_author_index_cache: Dict[str, Dict[str, Any]] = {}
|
|
_author_index_lock = threading.RLock()
|
|
|
|
|
|
def clear_search_cache():
|
|
"""Clear cached HuggingFace search results and in-memory indexes."""
|
|
global _search_cache, _author_index_cache
|
|
_search_cache.clear()
|
|
clear_remote_size_cache()
|
|
with _author_index_lock:
|
|
_author_index_cache.clear()
|
|
|
|
|
|
_report_progress = get_progress_reporter("HuggingFace progress callback")
|
|
|
|
|
|
def check_huggingface_token(token: Optional[str]) -> Dict[str, Any]:
|
|
"""Check whether a HuggingFace access token is accepted."""
|
|
from ..type_utils import check_credential_preconditions
|
|
precheck = check_credential_preconditions(token, "HuggingFace token")
|
|
if precheck:
|
|
return precheck
|
|
|
|
value = (token or "").strip()
|
|
headers = {"Authorization": f"Bearer {value}"}
|
|
|
|
def get_user(data):
|
|
return data.get("name") or data.get("user", {}).get("name") or ""
|
|
|
|
result = check_credential_http(
|
|
f"{HF_API_URL}/whoami-v2",
|
|
headers=headers,
|
|
success_message="HuggingFace token is valid.",
|
|
get_username=get_user,
|
|
error_msg_401_403="HuggingFace token is not accepted.",
|
|
)
|
|
if result.get("status_code") == 404:
|
|
result = check_credential_http(
|
|
f"{HF_API_URL}/whoami",
|
|
headers=headers,
|
|
success_message="HuggingFace token is valid.",
|
|
get_username=get_user,
|
|
error_msg_401_403="HuggingFace token is not accepted.",
|
|
)
|
|
return result
|
|
|
|
|
|
def check_brave_search_api_key(api_key: Optional[str]) -> Dict[str, Any]:
|
|
"""Check whether a Brave Search API key is accepted."""
|
|
from ..type_utils import check_credential_preconditions
|
|
precheck = check_credential_preconditions(api_key, "Brave Search API key")
|
|
if precheck:
|
|
return precheck
|
|
|
|
key = (api_key or "").strip()
|
|
headers = {"X-Subscription-Token": key}
|
|
params = {"q": "test", "count": 1}
|
|
|
|
def custom_429(response):
|
|
return {
|
|
"success": False,
|
|
"valid": False,
|
|
"status": "limited",
|
|
"message": "Brave Search rate limit was reached.",
|
|
"status_code": 429,
|
|
}
|
|
|
|
return check_credential_http(
|
|
BRAVE_SEARCH_API_URL,
|
|
headers=headers,
|
|
params=params,
|
|
success_message="Brave Search API key is valid.",
|
|
error_msg_401_403="Brave Search API key is not accepted.",
|
|
custom_429_handler=custom_429,
|
|
)
|
|
|
|
|
|
def _author_index_cache_key(author: str, headers: Dict[str, str]) -> str:
|
|
auth_key = "token" if headers.get("Authorization") else "public"
|
|
return f"{author}::{auth_key}"
|
|
|
|
|
|
def _is_author_index_fresh(index: Any) -> bool:
|
|
if not isinstance(index, dict) or not index:
|
|
return False
|
|
|
|
updated_at = index.get("updated_at")
|
|
if not isinstance(updated_at, (int, float)):
|
|
return False
|
|
|
|
return (time.time() - updated_at) < HF_AUTHOR_INDEX_CACHE_TTL_SECONDS
|
|
|
|
|
|
def _read_persistent_author_indexes() -> Dict[str, Any]:
|
|
default_val = {"version": HF_AUTHOR_INDEX_CACHE_VERSION, "authors": {}}
|
|
data = read_json_safe(HF_AUTHOR_INDEX_CACHE_PATH, default_val)
|
|
if not isinstance(data, dict) or data.get("version") != HF_AUTHOR_INDEX_CACHE_VERSION:
|
|
return {"version": HF_AUTHOR_INDEX_CACHE_VERSION, "authors": {}}
|
|
if not isinstance(data.get("authors"), dict):
|
|
data["authors"] = {}
|
|
else:
|
|
# A short-lived regression stored the raw HuggingFace response list in
|
|
# place of the built author index. Ignore malformed entries so they are
|
|
# fetched and rebuilt normally.
|
|
data["authors"] = {
|
|
author: index
|
|
for author, index in data["authors"].items()
|
|
if isinstance(index, dict)
|
|
}
|
|
return data
|
|
|
|
|
|
def _write_persistent_author_index(author: str, index: Dict[str, Any]) -> None:
|
|
try:
|
|
data = _read_persistent_author_indexes()
|
|
data["authors"][author] = index
|
|
write_json_atomic(
|
|
HF_AUTHOR_INDEX_CACHE_PATH,
|
|
data,
|
|
indent=2,
|
|
)
|
|
except Exception as e:
|
|
log.debug(f"Error writing HuggingFace author index cache: {e}")
|
|
|
|
|
|
def _build_author_index_from_models(
|
|
author: str, repos: List[Dict[str, Any]]
|
|
) -> Dict[str, Any]:
|
|
files = []
|
|
repo_ids = []
|
|
hash_count = 0
|
|
|
|
for repo in repos:
|
|
repo_id = repo.get("id", "")
|
|
if not repo_id:
|
|
continue
|
|
|
|
repo_ids.append(repo_id)
|
|
for sibling in repo.get("siblings") or []:
|
|
file_path = sibling.get("rfilename") or sibling.get("path") or ""
|
|
if not file_path:
|
|
continue
|
|
|
|
sha256 = _extract_huggingface_file_sha256(sibling)
|
|
if sha256:
|
|
hash_count += 1
|
|
|
|
files.append(
|
|
{
|
|
"repo_id": repo_id,
|
|
"path": file_path,
|
|
"filename": get_filename_from_path(file_path),
|
|
"size": sibling.get("size"),
|
|
"sha256": sha256,
|
|
}
|
|
)
|
|
|
|
return {
|
|
"author": author,
|
|
"updated_at": time.time(),
|
|
"repo_count": len(repo_ids),
|
|
"file_count": len(files),
|
|
"hash_count": hash_count,
|
|
"repos": repo_ids,
|
|
"files": files,
|
|
}
|
|
|
|
|
|
def _is_huggingface_hash_candidate_path(file_path: str) -> bool:
|
|
return str(file_path or "").lower().endswith(
|
|
(
|
|
".safetensors",
|
|
".ckpt",
|
|
".pt",
|
|
".pth",
|
|
".bin",
|
|
".gguf",
|
|
".onnx",
|
|
)
|
|
)
|
|
|
|
|
|
def _enrich_author_index_with_repo_trees(
|
|
index: Dict[str, Any], headers: Dict[str, str]
|
|
) -> Dict[str, Any]:
|
|
"""Add SHA-256 values from repository tree responses to an author index."""
|
|
files_by_repo_path = {}
|
|
for file_info in index.get("files") or []:
|
|
repo_id = file_info.get("repo_id", "")
|
|
file_path = file_info.get("path", "")
|
|
if (
|
|
repo_id
|
|
and file_path
|
|
and _is_huggingface_hash_candidate_path(file_path)
|
|
and not _extract_huggingface_file_sha256(file_info)
|
|
):
|
|
files_by_repo_path[(repo_id, file_path)] = file_info
|
|
|
|
repo_ids = sorted({repo_id for repo_id, _ in files_by_repo_path})
|
|
if not repo_ids:
|
|
index["hash_count"] = sum(
|
|
1
|
|
for file_info in index.get("files") or []
|
|
if _extract_huggingface_file_sha256(file_info)
|
|
)
|
|
return index
|
|
|
|
def fetch_repo_tree(repo_id: str):
|
|
return repo_id, _get_repo_tree(repo_id, headers=headers)
|
|
|
|
max_workers = min(8, len(repo_ids))
|
|
with ThreadPoolExecutor(max_workers=max_workers) as executor:
|
|
futures = {
|
|
executor.submit(fetch_repo_tree, repo_id): repo_id
|
|
for repo_id in repo_ids
|
|
}
|
|
for future in as_completed(futures):
|
|
repo_id = futures[future]
|
|
try:
|
|
_, tree = future.result()
|
|
except Exception as e:
|
|
log.debug(
|
|
f"HuggingFace author index tree failed repo={repo_id}: {e}"
|
|
)
|
|
continue
|
|
|
|
for tree_file in tree or []:
|
|
file_path = tree_file.get("path", "")
|
|
file_info = files_by_repo_path.get((repo_id, file_path))
|
|
if not file_info:
|
|
continue
|
|
|
|
sha256 = _extract_huggingface_file_sha256(tree_file)
|
|
if sha256:
|
|
file_info["sha256"] = sha256
|
|
if file_info.get("size") is None:
|
|
file_info["size"] = extract_file_size(tree_file)
|
|
|
|
index["hash_count"] = sum(
|
|
1
|
|
for file_info in index.get("files") or []
|
|
if _extract_huggingface_file_sha256(file_info)
|
|
)
|
|
return index
|
|
|
|
|
|
def _fetch_author_index(author: str, headers: Dict[str, str], limit: int = 200):
|
|
try:
|
|
repos = execute_provider_json_request(
|
|
"HuggingFace author index",
|
|
f"{HF_API_URL}/models",
|
|
params={"author": author, "limit": limit, "full": "true"},
|
|
headers=headers,
|
|
timeout=15,
|
|
)
|
|
if not isinstance(repos, list):
|
|
return None
|
|
|
|
index = _build_author_index_from_models(author, repos)
|
|
index = _enrich_author_index_with_repo_trees(index, headers=headers)
|
|
log.info(
|
|
f"HuggingFace author index refreshed for author={author}: "
|
|
f"repos={index['repo_count']}, files={index['file_count']}, "
|
|
f"hashes={index['hash_count']}"
|
|
)
|
|
return index
|
|
except Exception as e:
|
|
log.debug(f"Error refreshing HuggingFace author index for author={author}: {e}")
|
|
return None
|
|
|
|
|
|
def _get_author_index(
|
|
author: str,
|
|
headers: Dict[str, str],
|
|
force_refresh: bool = False,
|
|
persist: bool = True,
|
|
allow_stale: bool = False,
|
|
) -> Optional[Dict[str, Any]]:
|
|
cache_key = _author_index_cache_key(author, headers)
|
|
can_persist = not headers.get("Authorization")
|
|
|
|
with _author_index_lock:
|
|
memory_index = _author_index_cache.get(cache_key)
|
|
if not force_refresh and _is_author_index_fresh(memory_index):
|
|
return memory_index
|
|
|
|
persistent_index = None
|
|
if can_persist:
|
|
persistent_index = _read_persistent_author_indexes().get("authors", {}).get(
|
|
author
|
|
)
|
|
if not force_refresh and _is_author_index_fresh(persistent_index):
|
|
_author_index_cache[cache_key] = persistent_index
|
|
log.debug(f"HuggingFace author index cache hit for author={author}")
|
|
return persistent_index
|
|
|
|
if allow_stale and not force_refresh:
|
|
stale_index = persistent_index or memory_index
|
|
if isinstance(stale_index, dict):
|
|
_author_index_cache[cache_key] = stale_index
|
|
log.debug(f"HuggingFace stale author index cache hit for author={author}")
|
|
return stale_index
|
|
|
|
fresh_index = _fetch_author_index(author, headers=headers)
|
|
if fresh_index:
|
|
_author_index_cache[cache_key] = fresh_index
|
|
if can_persist and persist:
|
|
_write_persistent_author_index(author, fresh_index)
|
|
return fresh_index
|
|
|
|
if persistent_index:
|
|
log.warning(
|
|
f"Using stale HuggingFace author index for author={author}; refresh failed"
|
|
)
|
|
_author_index_cache[cache_key] = persistent_index
|
|
return persistent_index
|
|
|
|
return None
|
|
|
|
|
|
def get_author_fallback_index_status(author: str = "Comfy-Org") -> Dict[str, Any]:
|
|
"""Return persisted public author fallback index status for the options UI."""
|
|
with _author_index_lock:
|
|
index = _read_persistent_author_indexes().get("authors", {}).get(author)
|
|
|
|
updated_at = index.get("updated_at") if isinstance(index, dict) else None
|
|
age_seconds = (time.time() - updated_at) if isinstance(updated_at, (int, float)) else None
|
|
return {
|
|
"author": author,
|
|
"exists": bool(index),
|
|
"updated_at": updated_at,
|
|
"age_seconds": age_seconds,
|
|
"ttl_seconds": HF_AUTHOR_INDEX_CACHE_TTL_SECONDS,
|
|
"stale": bool(index) and not _is_author_index_fresh(index),
|
|
"repo_count": index.get("repo_count", 0) if isinstance(index, dict) else 0,
|
|
"file_count": index.get("file_count", 0) if isinstance(index, dict) else 0,
|
|
"hash_count": index.get("hash_count", 0) if isinstance(index, dict) else 0,
|
|
}
|
|
|
|
|
|
def get_known_author_fallback_indexes_status() -> Dict[str, Any]:
|
|
"""Return aggregate status for all known HuggingFace author indexes."""
|
|
author_statuses = [
|
|
get_author_fallback_index_status(author) for author in HF_AUTHOR_FALLBACKS
|
|
]
|
|
updated_times = [
|
|
status["updated_at"]
|
|
for status in author_statuses
|
|
if isinstance(status.get("updated_at"), (int, float))
|
|
]
|
|
cached_author_count = sum(bool(status["exists"]) for status in author_statuses)
|
|
fully_cached = cached_author_count == len(author_statuses)
|
|
|
|
return {
|
|
"authors": author_statuses,
|
|
"author_count": len(author_statuses),
|
|
"cached_author_count": cached_author_count,
|
|
"exists": cached_author_count > 0,
|
|
"fully_cached": fully_cached,
|
|
"updated_at": max(updated_times) if updated_times else None,
|
|
"stale": not fully_cached
|
|
or any(bool(status["stale"]) for status in author_statuses),
|
|
"repo_count": sum(int(status["repo_count"]) for status in author_statuses),
|
|
"file_count": sum(int(status["file_count"]) for status in author_statuses),
|
|
"hash_count": sum(int(status.get("hash_count") or 0) for status in author_statuses),
|
|
}
|
|
|
|
|
|
def refresh_author_fallback_index(
|
|
token: Optional[str] = None, author: str = "Comfy-Org"
|
|
) -> Dict[str, Any]:
|
|
"""Force refresh the public HuggingFace author fallback index."""
|
|
headers = {}
|
|
|
|
index = _get_author_index(author, headers=headers, force_refresh=True)
|
|
status = get_author_fallback_index_status(author)
|
|
status["success"] = bool(index)
|
|
status["persisted"] = bool(index)
|
|
return status
|
|
|
|
|
|
def refresh_known_author_fallback_indexes(
|
|
token: Optional[str] = None,
|
|
) -> Dict[str, Any]:
|
|
"""Force refresh all known public HuggingFace author indexes."""
|
|
refresh_results = [
|
|
refresh_author_fallback_index(token, author) for author in HF_AUTHOR_FALLBACKS
|
|
]
|
|
status = get_known_author_fallback_indexes_status()
|
|
failed_authors = [
|
|
result["author"] for result in refresh_results if not result.get("success")
|
|
]
|
|
status["success"] = not failed_authors
|
|
status["persisted"] = all(
|
|
bool(result.get("persisted")) for result in refresh_results
|
|
)
|
|
status["failed_authors"] = failed_authors
|
|
return status
|
|
|
|
|
|
def parse_huggingface_url(url: str) -> Optional[HuggingFaceFileReference]:
|
|
"""
|
|
Parse a HuggingFace URL to extract repo and filename.
|
|
|
|
Handles formats:
|
|
- https://huggingface.co/user/repo/resolve/main/file.safetensors
|
|
- https://huggingface.co/user/repo/blob/main/file.safetensors
|
|
- hf://user/repo/file.safetensors
|
|
|
|
Returns:
|
|
Typed repository-file identity, or None if not an HF URL
|
|
"""
|
|
if url.startswith("hf://"):
|
|
parts = url[5:].split("/", 2)
|
|
if len(parts) >= 3:
|
|
return HuggingFaceFileReference(
|
|
repo=f"{parts[0]}/{parts[1]}",
|
|
filename=parts[2],
|
|
)
|
|
return None
|
|
|
|
parsed = urlparse(url)
|
|
if not host_matches_domain(parsed.hostname, "huggingface.co"):
|
|
return None
|
|
|
|
match = re.match(r"^/([^/]+/[^/]+)/(resolve|blob)/([^/]+)/(.+)$", parsed.path)
|
|
if match:
|
|
return HuggingFaceFileReference(
|
|
repo=match.group(1),
|
|
branch=match.group(3),
|
|
filename=match.group(4),
|
|
)
|
|
|
|
return None
|
|
|
|
|
|
def get_huggingface_download_url(repo: str, filename: str, branch: str = "main") -> str:
|
|
"""Generate a direct download URL for a HuggingFace file."""
|
|
return (
|
|
f"https://huggingface.co/{repo}/resolve/"
|
|
f"{quote(branch, safe='')}/{quote(filename)}"
|
|
)
|
|
|
|
|
|
def _normalize_huggingface_path(value: Any) -> str:
|
|
return str(value or "").replace("\\", "/").strip("/")
|
|
|
|
|
|
def _get_huggingface_variant_label(filename: str) -> str:
|
|
"""Return a concise precision/quantization label inferred from a filename."""
|
|
stem = os.path.splitext(get_filename_from_path(filename))[0].lower()
|
|
patterns = (
|
|
("nvfp4_mixed", "NVFP4 mixed"),
|
|
("fp8_scaled", "FP8 scaled"),
|
|
("int8_convrot", "INT8 convrot"),
|
|
("nvfp4", "NVFP4"),
|
|
("fp8", "FP8"),
|
|
("bf16", "BF16"),
|
|
("fp16", "FP16"),
|
|
("int8", "INT8"),
|
|
("int4", "INT4"),
|
|
)
|
|
for token, label in patterns:
|
|
if re.search(rf"(?:^|[_.-]){re.escape(token)}(?:$|[_.-])", stem):
|
|
return label
|
|
return ""
|
|
|
|
|
|
def _normalize_huggingface_details_file(
|
|
repo_id: str,
|
|
branch: str,
|
|
file_info: Dict[str, Any],
|
|
target_path: str,
|
|
) -> Optional[Dict[str, Any]]:
|
|
file_path = _normalize_huggingface_path(file_info.get("path"))
|
|
if not file_path:
|
|
return None
|
|
|
|
download_url = get_huggingface_download_url(repo_id, file_path, branch)
|
|
if not looks_like_model_file(download_url):
|
|
return None
|
|
|
|
filename = get_filename_from_path(file_path)
|
|
sha256 = _extract_huggingface_file_sha256(file_info)
|
|
extension = os.path.splitext(filename)[1].lstrip(".")
|
|
metadata = {
|
|
"format": extension.upper() if extension else "",
|
|
"fp": _get_huggingface_variant_label(filename),
|
|
}
|
|
metadata = {key: value for key, value in metadata.items() if value}
|
|
page_url = (
|
|
f"https://huggingface.co/{repo_id}/blob/"
|
|
f"{quote(branch, safe='')}/{quote(file_path, safe='/')}"
|
|
)
|
|
|
|
return {
|
|
"name": filename,
|
|
"filename": filename,
|
|
"path": file_path,
|
|
"size": extract_file_size(file_info),
|
|
"sha256": sha256,
|
|
"hash": sha256,
|
|
"hashes": {"SHA256": sha256} if sha256 else {},
|
|
"download_url": download_url,
|
|
"url": page_url,
|
|
"primary": bool(target_path and file_path == target_path),
|
|
"metadata": metadata,
|
|
}
|
|
|
|
|
|
def get_huggingface_model_details(
|
|
repo_id: str,
|
|
file_path: str = "",
|
|
branch: str = "main",
|
|
token: Optional[str] = None,
|
|
) -> Optional[Dict[str, Any]]:
|
|
"""Fetch model-file variants from the folder containing a matched HF file."""
|
|
normalized_repo_id = str(repo_id or "").strip().strip("/")
|
|
if not re.fullmatch(r"[A-Za-z0-9_.-]+/[A-Za-z0-9_.-]+", normalized_repo_id):
|
|
return None
|
|
|
|
normalized_branch = str(branch or "main").strip() or "main"
|
|
if len(normalized_branch) > 200 or any(
|
|
character in normalized_branch for character in "\r\n?#"
|
|
):
|
|
return None
|
|
|
|
target_path = _normalize_huggingface_path(file_path)
|
|
target_directory = posixpath.dirname(target_path)
|
|
headers = {"Authorization": f"Bearer {token.strip()}"} if token and token.strip() else {}
|
|
repo_tree = _get_repo_tree(
|
|
normalized_repo_id,
|
|
headers=headers,
|
|
branch=normalized_branch,
|
|
)
|
|
if not repo_tree:
|
|
return None
|
|
|
|
normalized_files = []
|
|
for file_info in repo_tree:
|
|
candidate_path = _normalize_huggingface_path(file_info.get("path"))
|
|
if not candidate_path:
|
|
continue
|
|
if target_path and posixpath.dirname(candidate_path) != target_directory:
|
|
continue
|
|
|
|
normalized_file = _normalize_huggingface_details_file(
|
|
normalized_repo_id,
|
|
normalized_branch,
|
|
file_info,
|
|
target_path,
|
|
)
|
|
if normalized_file:
|
|
normalized_files.append(normalized_file)
|
|
|
|
if not normalized_files:
|
|
return None
|
|
|
|
normalized_files.sort(
|
|
key=lambda item: (
|
|
not item.get("primary", False),
|
|
str(item.get("name") or "").lower(),
|
|
)
|
|
)
|
|
|
|
model_data: Dict[str, Any] = {}
|
|
try:
|
|
response_data = execute_provider_json_request(
|
|
"HuggingFace model details",
|
|
f"{HF_API_URL}/models/{normalized_repo_id}",
|
|
headers=headers,
|
|
timeout=15,
|
|
)
|
|
if isinstance(response_data, dict):
|
|
model_data = response_data
|
|
except Exception as e:
|
|
log.debug(
|
|
f"Error getting HuggingFace model metadata for {normalized_repo_id}: {e}"
|
|
)
|
|
|
|
card_data = (
|
|
model_data.get("cardData")
|
|
if isinstance(model_data.get("cardData"), dict)
|
|
else {}
|
|
)
|
|
tags = model_data.get("tags") if isinstance(model_data.get("tags"), list) else []
|
|
repo_url = f"https://huggingface.co/{normalized_repo_id}"
|
|
tree_url = (
|
|
f"{repo_url}/tree/{quote(normalized_branch, safe='')}"
|
|
+ (
|
|
f"/{quote(target_directory, safe='/')}"
|
|
if target_directory
|
|
else ""
|
|
)
|
|
)
|
|
version = {
|
|
"id": normalized_branch,
|
|
"name": normalized_branch,
|
|
"published_at": model_data.get("createdAt") or "",
|
|
"updated_at": model_data.get("lastModified") or "",
|
|
"description": "",
|
|
"trained_words": [],
|
|
"stats": {
|
|
"downloads": model_data.get("downloads", 0),
|
|
"thumbsUpCount": model_data.get("likes", 0),
|
|
},
|
|
"files": normalized_files,
|
|
"images": [],
|
|
"url": tree_url,
|
|
}
|
|
owner = normalized_repo_id.split("/", 1)[0]
|
|
description = (
|
|
model_data.get("description")
|
|
or card_data.get("description")
|
|
or (
|
|
f"Select a model file from `{target_directory}` in this Hugging Face "
|
|
"repository."
|
|
if target_directory
|
|
else "Select a model file from this Hugging Face repository."
|
|
)
|
|
)
|
|
|
|
return {
|
|
"source": "huggingface",
|
|
"details_source": "huggingface",
|
|
"model_id": normalized_repo_id,
|
|
"version_id": normalized_branch,
|
|
"name": model_data.get("modelId") or normalized_repo_id,
|
|
"type": model_data.get("pipeline_tag")
|
|
or model_data.get("library_name")
|
|
or "Model repository",
|
|
"description": description,
|
|
"creator": {"username": model_data.get("author") or owner},
|
|
"stats": {
|
|
"downloads": model_data.get("downloads", 0),
|
|
"thumbsUpCount": model_data.get("likes", 0),
|
|
},
|
|
"tags": tags,
|
|
"url": repo_url,
|
|
"version_url": tree_url,
|
|
"folder": target_directory,
|
|
"branch": normalized_branch,
|
|
"versions": [version],
|
|
"selected_version": version,
|
|
"images": [],
|
|
}
|
|
|
|
|
|
|
|
|
|
# _extract_file_info_size removed in favor of extract_file_size from type_utils
|
|
|
|
|
|
|
|
def _build_huggingface_result(
|
|
repo_id: str,
|
|
file_path: str,
|
|
file_info: Dict[str, Any],
|
|
match_type: str,
|
|
headers: Optional[Dict[str, str]] = None,
|
|
confidence: float = 0.0,
|
|
) -> SearchResult:
|
|
sha256 = _extract_huggingface_file_sha256(file_info)
|
|
download_url = get_huggingface_download_url(repo_id, file_path)
|
|
size = resolve_file_size(
|
|
file_info,
|
|
[download_url],
|
|
probe=lambda url: probe_remote_file_size(url, headers=headers),
|
|
)
|
|
return build_model_result(
|
|
"huggingface",
|
|
model_id=repo_id,
|
|
version_id=None,
|
|
filename=get_filename_from_path(file_path),
|
|
url=download_url,
|
|
download_url=download_url,
|
|
size=size,
|
|
match_type=match_type,
|
|
confidence=confidence,
|
|
sha256=sha256,
|
|
normalize_hashes=True,
|
|
repo_id=repo_id,
|
|
path=file_path,
|
|
branch="main",
|
|
)
|
|
|
|
|
|
def _extract_huggingface_file_sha256(file_info: Dict[str, Any]) -> str:
|
|
lfs_info = file_info.get("lfs") if isinstance(file_info.get("lfs"), dict) else {}
|
|
return normalize_sha256(
|
|
file_info.get("sha256")
|
|
or file_info.get("hash")
|
|
or lfs_info.get("sha256")
|
|
or lfs_info.get("oid")
|
|
)
|
|
|
|
|
|
def _get_repo_tree(
|
|
repo_id: str, headers: Dict[str, str], branch: str = "main"
|
|
) -> Optional[List[Dict[str, Any]]]:
|
|
files_url = f"{HF_API_URL}/models/{repo_id}/tree/{quote(branch, safe='')}"
|
|
|
|
try:
|
|
data = execute_provider_json_request(
|
|
"HuggingFace tree",
|
|
files_url,
|
|
params={"recursive": "1"},
|
|
headers=headers,
|
|
timeout=15,
|
|
)
|
|
if isinstance(data, list):
|
|
return data
|
|
return None
|
|
except Exception as e:
|
|
log.debug(f"Error getting HuggingFace tree for {repo_id}@{branch}: {e}")
|
|
return None
|
|
|
|
|
|
def get_huggingface_file_sha256(
|
|
url: str,
|
|
headers: Optional[Dict[str, str]] = None,
|
|
) -> str:
|
|
"""Return the regular file SHA256 for a Hugging Face file URL."""
|
|
parsed = parse_huggingface_url(url)
|
|
if not parsed:
|
|
return ""
|
|
|
|
repo_id = parsed.repo
|
|
branch = parsed.branch
|
|
target_path = _normalize_huggingface_path(
|
|
unquote(parsed.filename)
|
|
)
|
|
if not repo_id or not target_path:
|
|
return ""
|
|
|
|
repo_tree = _get_repo_tree(repo_id, headers=headers or {}, branch=branch)
|
|
for file_info in repo_tree or []:
|
|
candidate_path = _normalize_huggingface_path(file_info.get("path"))
|
|
if candidate_path == target_path:
|
|
return _extract_huggingface_file_sha256(file_info)
|
|
return ""
|
|
|
|
|
|
def _find_matching_file_in_repo(
|
|
repo_id: str,
|
|
files: List[Dict[str, Any]],
|
|
filename: str,
|
|
exact_only: bool = False,
|
|
headers: Optional[Dict[str, str]] = None,
|
|
) -> Optional[SearchResult]:
|
|
filename_lower = filename.lower()
|
|
filename_base = os.path.splitext(filename_lower)[0]
|
|
|
|
for file_info in files:
|
|
file_path = file_info.get("path", "")
|
|
if not file_path:
|
|
continue
|
|
|
|
file_name = get_filename_from_path(file_path)
|
|
file_name_lower = file_name.lower()
|
|
|
|
if file_name_lower == filename_lower or file_path.lower().endswith(filename_lower):
|
|
return _build_huggingface_result(
|
|
repo_id,
|
|
file_path,
|
|
file_info,
|
|
"exact",
|
|
headers=headers,
|
|
confidence=100.0,
|
|
)
|
|
|
|
if exact_only:
|
|
return None
|
|
|
|
best_match: Optional[SearchResult] = None
|
|
for file_info in files:
|
|
file_path = file_info.get("path", "")
|
|
if not file_path:
|
|
continue
|
|
|
|
file_name = get_filename_from_path(file_path)
|
|
file_name_lower = file_name.lower()
|
|
file_base = os.path.splitext(file_name_lower)[0]
|
|
if not file_path.lower().endswith((".safetensors", ".ckpt")):
|
|
continue
|
|
|
|
confidence = calculate_filename_confidence(filename, file_name)
|
|
is_partial_filename = (
|
|
filename_base in file_base or file_base in filename_base
|
|
)
|
|
if not is_partial_filename and confidence < 70.0:
|
|
continue
|
|
|
|
candidate = _build_huggingface_result(
|
|
repo_id,
|
|
file_path,
|
|
file_info,
|
|
"partial",
|
|
headers=headers,
|
|
confidence=confidence,
|
|
)
|
|
if best_match is None or candidate.confidence > best_match.confidence:
|
|
best_match = candidate
|
|
|
|
return best_match
|
|
|
|
|
|
def _find_matching_file_in_author_index(
|
|
index: Dict[str, Any],
|
|
filename: str,
|
|
exact_only: bool = False,
|
|
headers: Optional[Dict[str, str]] = None,
|
|
sha256: Optional[str] = None,
|
|
) -> Optional[SearchResult]:
|
|
files = index.get("files") or []
|
|
requested_sha256 = normalize_sha256(sha256)
|
|
|
|
if requested_sha256:
|
|
for file_info in files:
|
|
repo_id = file_info.get("repo_id", "")
|
|
file_path = file_info.get("path", "")
|
|
if (
|
|
not repo_id
|
|
or not file_path
|
|
or _extract_huggingface_file_sha256(file_info) != requested_sha256
|
|
):
|
|
continue
|
|
|
|
return _build_huggingface_result(
|
|
repo_id,
|
|
file_path,
|
|
file_info,
|
|
"hash",
|
|
headers=headers,
|
|
confidence=100.0,
|
|
)
|
|
|
|
return None
|
|
|
|
filename_lower = filename.lower()
|
|
filename_base = os.path.splitext(filename_lower)[0]
|
|
|
|
for file_info in files:
|
|
repo_id = file_info.get("repo_id", "")
|
|
file_path = file_info.get("path", "")
|
|
if not repo_id or not file_path:
|
|
continue
|
|
|
|
file_name = get_filename_from_path(file_path)
|
|
file_name_lower = file_name.lower()
|
|
if file_name_lower == filename_lower or file_path.lower().endswith(
|
|
filename_lower
|
|
):
|
|
return _build_huggingface_result(
|
|
repo_id,
|
|
file_path,
|
|
file_info,
|
|
"exact",
|
|
headers=headers,
|
|
confidence=100.0,
|
|
)
|
|
|
|
if exact_only:
|
|
return None
|
|
|
|
best_match: Optional[SearchResult] = None
|
|
for file_info in files:
|
|
repo_id = file_info.get("repo_id", "")
|
|
file_path = file_info.get("path", "")
|
|
if not repo_id or not file_path:
|
|
continue
|
|
|
|
file_name = get_filename_from_path(file_path)
|
|
file_name_lower = file_name.lower()
|
|
file_base = os.path.splitext(file_name_lower)[0]
|
|
file_path_lower = file_path.lower()
|
|
|
|
if not file_path_lower.endswith((".safetensors", ".ckpt")):
|
|
continue
|
|
|
|
confidence = calculate_filename_confidence(filename, file_name)
|
|
is_partial_filename = (
|
|
filename_base in file_base or file_base in filename_base
|
|
)
|
|
if not is_partial_filename and confidence < 70.0:
|
|
continue
|
|
|
|
candidate = _build_huggingface_result(
|
|
repo_id,
|
|
file_path,
|
|
file_info,
|
|
"partial",
|
|
headers=headers,
|
|
confidence=confidence,
|
|
)
|
|
if best_match is None or candidate.confidence > best_match.confidence:
|
|
best_match = candidate
|
|
|
|
return best_match
|
|
|
|
|
|
|
|
|
|
def _get_repos_by_author(
|
|
author: str, headers: Dict[str, str], limit: int = 200
|
|
) -> List[str]:
|
|
try:
|
|
response = request_source_response(
|
|
f"{HF_API_URL}/models",
|
|
params={"author": author, "limit": limit},
|
|
headers=headers,
|
|
timeout=15,
|
|
log_name="HuggingFace repos by author"
|
|
)
|
|
if response.status_code != 200:
|
|
log.debug(
|
|
f"HuggingFace author search returned {response.status_code} for author={author}"
|
|
)
|
|
return []
|
|
|
|
repos = response.json()
|
|
return [repo.get("id", "") for repo in repos if repo.get("id")]
|
|
except Exception as e:
|
|
log.debug(f"Error getting HuggingFace repos for author={author}: {e}")
|
|
return []
|
|
|
|
|
|
def _search_brave_for_huggingface_candidates(
|
|
filename: str, brave_api_key: str
|
|
) -> List[Dict[str, str]]:
|
|
if not brave_api_key:
|
|
return []
|
|
|
|
query = f'"{filename}" site:huggingface.co'
|
|
headers = {
|
|
"Accept": "application/json",
|
|
"Accept-Encoding": "gzip",
|
|
"X-Subscription-Token": brave_api_key,
|
|
}
|
|
params = {
|
|
"q": query,
|
|
"count": 20,
|
|
"safesearch": "off",
|
|
"spellcheck": "false",
|
|
}
|
|
candidates = []
|
|
seen = set()
|
|
|
|
try:
|
|
response = request_source_response(
|
|
BRAVE_SEARCH_API_URL, headers=headers, params=params, timeout=15, log_name="Brave search"
|
|
)
|
|
if response.status_code != 200:
|
|
log.warning(
|
|
f"Brave search returned {response.status_code} for query={query}"
|
|
)
|
|
return []
|
|
|
|
payload = response.json()
|
|
results = ((payload or {}).get("web") or {}).get("results") or []
|
|
for result in results:
|
|
candidate_url = result.get("url", "")
|
|
if not candidate_url:
|
|
continue
|
|
|
|
parsed = parse_huggingface_url(candidate_url)
|
|
if parsed:
|
|
repo_id = parsed.repo
|
|
file_path = parsed.filename
|
|
branch = parsed.branch
|
|
else:
|
|
parsed_url = urlparse(candidate_url)
|
|
if "huggingface.co" not in parsed_url.netloc:
|
|
continue
|
|
path_parts = [part for part in parsed_url.path.split("/") if part]
|
|
if len(path_parts) < 2:
|
|
continue
|
|
repo_id = f"{path_parts[0]}/{path_parts[1]}"
|
|
file_path = filename
|
|
branch = "main"
|
|
|
|
key = (repo_id, branch, file_path)
|
|
if not repo_id or key in seen:
|
|
continue
|
|
|
|
seen.add(key)
|
|
candidates.append(
|
|
{
|
|
"repo": repo_id,
|
|
"branch": branch,
|
|
"filename": file_path,
|
|
"url": candidate_url,
|
|
}
|
|
)
|
|
except Exception as e:
|
|
log.debug(f"Brave search error for filename={filename}: {e}")
|
|
|
|
return candidates
|
|
|
|
|
|
def search_huggingface_for_file(
|
|
filename: str,
|
|
token: Optional[str] = None,
|
|
exact_only: bool = False,
|
|
brave_api_key: Optional[str] = None,
|
|
use_api_search: bool = True,
|
|
use_comfy_org_fallback: bool = True,
|
|
use_brave_fallback: bool = True,
|
|
force_refresh: bool = False,
|
|
progress_callback: Optional[Callable[[Dict[str, Any]], None]] = None,
|
|
sha256: Optional[str] = None,
|
|
) -> Optional[SearchResult]:
|
|
"""
|
|
Search HuggingFace for a specific model file.
|
|
Returns the first repo that contains a matching file.
|
|
|
|
Args:
|
|
filename: Filename to search for
|
|
token: Optional HF token
|
|
exact_only: If True, only return exact filename matches (for downloads).
|
|
If False, also try partial matching (for local file resolution).
|
|
|
|
Returns:
|
|
A typed SearchResult if a matching file is found, otherwise None.
|
|
"""
|
|
global _search_cache
|
|
|
|
requested_sha256 = normalize_sha256(sha256)
|
|
brave_key = "brave" if brave_api_key and use_brave_fallback else "nobrave"
|
|
token_key = "token" if token else "notoken"
|
|
methods_key = (
|
|
f"api{int(bool(use_api_search))}_"
|
|
f"comfy{int(bool(use_comfy_org_fallback))}_"
|
|
f"brave{int(bool(use_brave_fallback))}_"
|
|
f"force{int(bool(force_refresh))}"
|
|
)
|
|
cache_key = f"hf_{filename}_sha{requested_sha256}_exact{exact_only}_{token_key}_{brave_key}_{methods_key}"
|
|
if cache_key in _search_cache:
|
|
log.debug(f"HuggingFace cache hit file={filename} exact={exact_only}")
|
|
_report_progress(
|
|
progress_callback,
|
|
"cache",
|
|
"Using HuggingFace cache",
|
|
86,
|
|
)
|
|
return _search_cache[cache_key]
|
|
|
|
try:
|
|
def result_matches_requested_hash(result: Optional[SearchResult]) -> bool:
|
|
if not requested_sha256:
|
|
return True
|
|
return normalize_sha256(
|
|
result.sha256
|
|
or result.extra_value("hash")
|
|
or (result.hashes or {}).get("SHA256")
|
|
or (result.hashes or {}).get("sha256")
|
|
) == requested_sha256
|
|
|
|
headers = {}
|
|
if token:
|
|
headers["Authorization"] = f"Bearer {token}"
|
|
|
|
log.info(f"HuggingFace start file={filename} exact={exact_only}")
|
|
|
|
def find_author_index_result(
|
|
author: str,
|
|
) -> tuple[Optional[Dict[str, Any]], Optional[SearchResult]]:
|
|
index = _get_author_index(
|
|
author,
|
|
headers={},
|
|
allow_stale=force_refresh,
|
|
)
|
|
result = (
|
|
_find_matching_file_in_author_index(
|
|
index,
|
|
filename,
|
|
exact_only=exact_only,
|
|
headers={},
|
|
sha256=requested_sha256,
|
|
)
|
|
if index
|
|
else None
|
|
)
|
|
|
|
if (
|
|
force_refresh
|
|
and index
|
|
and not _is_author_index_fresh(index)
|
|
and result is None
|
|
):
|
|
index = _get_author_index(
|
|
author,
|
|
headers={},
|
|
force_refresh=True,
|
|
)
|
|
result = (
|
|
_find_matching_file_in_author_index(
|
|
index,
|
|
filename,
|
|
exact_only=exact_only,
|
|
headers={},
|
|
sha256=requested_sha256,
|
|
)
|
|
if index
|
|
else None
|
|
)
|
|
|
|
return index, result
|
|
|
|
cached_author_indexes: Dict[str, Optional[Dict[str, Any]]] = {}
|
|
if requested_sha256 and use_comfy_org_fallback:
|
|
total_authors = max(1, len(HF_AUTHOR_FALLBACKS))
|
|
for author_index, author in enumerate(HF_AUTHOR_FALLBACKS, start=1):
|
|
_report_progress(
|
|
progress_callback,
|
|
"author_hash_index",
|
|
f"Checking HF hash index {author_index}/{total_authors}",
|
|
28 + ((author_index - 1) / total_authors) * 6,
|
|
author=author,
|
|
author_index=author_index,
|
|
author_count=total_authors,
|
|
)
|
|
index, result = find_author_index_result(author)
|
|
cached_author_indexes[author] = index
|
|
if not index:
|
|
continue
|
|
|
|
if result:
|
|
_search_cache[cache_key] = result
|
|
_report_progress(
|
|
progress_callback,
|
|
"found",
|
|
"Found HuggingFace SHA-256 match",
|
|
92,
|
|
repo=result.extra_value("repo_id"),
|
|
match_type=result.match_type,
|
|
)
|
|
log.info(
|
|
f"HuggingFace found source=author_index "
|
|
f"match=hash file={result.filename} "
|
|
f"repo={result.extra_value('repo_id')} path={result.extra_value('path')}"
|
|
)
|
|
return result
|
|
|
|
repos_to_check = []
|
|
seen_repos = set()
|
|
search_queries = build_filename_search_queries(filename)
|
|
|
|
if use_api_search:
|
|
total_queries = max(1, len(search_queries))
|
|
for query_index, search_query in enumerate(search_queries, start=1):
|
|
_report_progress(
|
|
progress_callback,
|
|
"api_search",
|
|
f"HuggingFace API query {query_index}/{total_queries}",
|
|
32 + ((query_index - 1) / total_queries) * 10,
|
|
query=search_query,
|
|
query_index=query_index,
|
|
query_count=total_queries,
|
|
)
|
|
search_url = f"{HF_API_URL}/models?search={quote(search_query)}&limit=20"
|
|
repos = execute_provider_json_request(
|
|
"HuggingFace models search",
|
|
search_url,
|
|
api_key=token,
|
|
timeout=10,
|
|
)
|
|
if repos is None:
|
|
log.warning(
|
|
f"HuggingFace API query={search_query} failed"
|
|
)
|
|
continue
|
|
if not repos:
|
|
continue
|
|
log.info(f"HuggingFace API query={search_query} repos={len(repos)}")
|
|
|
|
for repo in repos:
|
|
repo_id = repo.get("id", "")
|
|
if repo_id and repo_id not in seen_repos:
|
|
seen_repos.add(repo_id)
|
|
repos_to_check.append(repo_id)
|
|
|
|
_report_progress(
|
|
progress_callback,
|
|
"api_candidates",
|
|
f"HuggingFace API candidates: {len(repos_to_check)} repos",
|
|
44,
|
|
candidate_count=len(repos_to_check),
|
|
)
|
|
else:
|
|
_report_progress(
|
|
progress_callback,
|
|
"api_skip",
|
|
"Skipping HuggingFace API search",
|
|
44,
|
|
)
|
|
|
|
repo_count = len(repos_to_check)
|
|
for repo_index, repo_id in enumerate(repos_to_check, start=1):
|
|
_report_progress(
|
|
progress_callback,
|
|
"repo_files",
|
|
f"Checking HF repo {repo_index}/{repo_count}",
|
|
48 + (repo_index / max(1, repo_count)) * 14,
|
|
repo=repo_id,
|
|
candidate_index=repo_index,
|
|
candidate_count=repo_count,
|
|
)
|
|
log.debug(f"HuggingFace check repo={repo_id} file={filename}")
|
|
files = _get_repo_tree(repo_id, headers=headers)
|
|
if not files:
|
|
continue
|
|
|
|
log.debug(f"HuggingFace tree repo={repo_id} files={len(files)}")
|
|
result = _find_matching_file_in_repo(
|
|
repo_id,
|
|
files,
|
|
filename,
|
|
exact_only=exact_only,
|
|
headers=headers,
|
|
)
|
|
if result and result_matches_requested_hash(result):
|
|
_search_cache[cache_key] = result
|
|
_report_progress(
|
|
progress_callback,
|
|
"found",
|
|
"Found HuggingFace match",
|
|
92,
|
|
repo=repo_id,
|
|
match_type=result.match_type,
|
|
)
|
|
log.info(
|
|
f"HuggingFace found match={result.match_type} file={filename} repo={repo_id} path={result.extra_value('path')}"
|
|
)
|
|
return result
|
|
|
|
author_fallback_repos = []
|
|
author_fallback_repo_count = 0
|
|
if use_comfy_org_fallback:
|
|
total_authors = max(1, len(HF_AUTHOR_FALLBACKS))
|
|
for author_index, author in enumerate(HF_AUTHOR_FALLBACKS, start=1):
|
|
_report_progress(
|
|
progress_callback,
|
|
"author_index",
|
|
f"Checking HF fallback index {author_index}/{total_authors}",
|
|
64 + ((author_index - 1) / total_authors) * 6,
|
|
author=author,
|
|
author_index=author_index,
|
|
author_count=total_authors,
|
|
)
|
|
if author in cached_author_indexes:
|
|
index = cached_author_indexes[author]
|
|
result = (
|
|
_find_matching_file_in_author_index(
|
|
index,
|
|
filename,
|
|
exact_only=exact_only,
|
|
headers={},
|
|
sha256=requested_sha256,
|
|
)
|
|
if index
|
|
else None
|
|
)
|
|
else:
|
|
index, result = find_author_index_result(author)
|
|
if index:
|
|
author_fallback_repo_count += int(index.get("repo_count") or 0)
|
|
if result and result_matches_requested_hash(result):
|
|
_search_cache[cache_key] = result
|
|
_report_progress(
|
|
progress_callback,
|
|
"found",
|
|
"Found HuggingFace fallback match",
|
|
92,
|
|
repo=result.extra_value("repo_id"),
|
|
match_type=result.match_type,
|
|
)
|
|
log.info(
|
|
f"HuggingFace found source=author_index match={result.match_type} file={filename} repo={result.extra_value('repo_id')} path={result.extra_value('path')}"
|
|
)
|
|
return result
|
|
|
|
log.debug(
|
|
f"HuggingFace author_index miss author={author} file={filename} repos={index.get('repo_count', 0)} files={index.get('file_count', 0)}"
|
|
)
|
|
continue
|
|
|
|
repos = _get_repos_by_author(author, headers=headers)
|
|
log.info(f"HuggingFace author_fallback author={author} repos={len(repos)}")
|
|
for repo_id in repos:
|
|
if repo_id not in seen_repos:
|
|
seen_repos.add(repo_id)
|
|
author_fallback_repos.append(repo_id)
|
|
author_fallback_repo_count += 1
|
|
else:
|
|
_report_progress(
|
|
progress_callback,
|
|
"author_skip",
|
|
"Skipping HuggingFace fallback index",
|
|
70,
|
|
)
|
|
|
|
fallback_repo_count = len(author_fallback_repos)
|
|
for repo_index, repo_id in enumerate(author_fallback_repos, start=1):
|
|
_report_progress(
|
|
progress_callback,
|
|
"author_repo",
|
|
f"Checking HF fallback repo {repo_index}/{fallback_repo_count}",
|
|
70 + (repo_index / max(1, fallback_repo_count)) * 8,
|
|
repo=repo_id,
|
|
candidate_index=repo_index,
|
|
candidate_count=fallback_repo_count,
|
|
)
|
|
log.debug(
|
|
f"HuggingFace author_fallback check repo={repo_id} file={filename}"
|
|
)
|
|
files = _get_repo_tree(repo_id, headers=headers)
|
|
if not files:
|
|
continue
|
|
|
|
result = _find_matching_file_in_repo(
|
|
repo_id,
|
|
files,
|
|
filename,
|
|
exact_only=exact_only,
|
|
headers=headers,
|
|
)
|
|
if result and result_matches_requested_hash(result):
|
|
_search_cache[cache_key] = result
|
|
_report_progress(
|
|
progress_callback,
|
|
"found",
|
|
"Found HuggingFace fallback match",
|
|
92,
|
|
repo=repo_id,
|
|
match_type=result.match_type,
|
|
)
|
|
log.info(
|
|
f"HuggingFace found source=author_fallback match={result.match_type} file={filename} repo={repo_id} path={result.extra_value('path')}"
|
|
)
|
|
return result
|
|
|
|
if use_brave_fallback:
|
|
_report_progress(
|
|
progress_callback,
|
|
"brave_search",
|
|
"Searching Brave fallback",
|
|
80,
|
|
)
|
|
brave_candidates = (
|
|
_search_brave_for_huggingface_candidates(filename, brave_api_key or "")
|
|
if use_brave_fallback
|
|
else []
|
|
)
|
|
log.info(f"HuggingFace brave_fallback file={filename} candidates={len(brave_candidates)}")
|
|
_report_progress(
|
|
progress_callback,
|
|
"brave_candidates",
|
|
f"Brave candidates: {len(brave_candidates)}",
|
|
84,
|
|
candidate_count=len(brave_candidates),
|
|
)
|
|
|
|
brave_count = len(brave_candidates)
|
|
for candidate_index, candidate in enumerate(brave_candidates, start=1):
|
|
repo_id = candidate.get("repo", "")
|
|
branch = candidate.get("branch", "main")
|
|
expected_file = candidate.get("filename", filename)
|
|
_report_progress(
|
|
progress_callback,
|
|
"brave_repo",
|
|
f"Checking Brave candidate {candidate_index}/{brave_count}",
|
|
84 + (candidate_index / max(1, brave_count)) * 6,
|
|
repo=repo_id,
|
|
candidate_index=candidate_index,
|
|
candidate_count=brave_count,
|
|
)
|
|
files = _get_repo_tree(repo_id, headers=headers, branch=branch)
|
|
if not files:
|
|
continue
|
|
|
|
result = _find_matching_file_in_repo(
|
|
repo_id,
|
|
files,
|
|
expected_file,
|
|
exact_only=True,
|
|
headers=headers,
|
|
)
|
|
if (
|
|
result
|
|
and result.filename.lower() == filename.lower()
|
|
and result_matches_requested_hash(result)
|
|
):
|
|
_search_cache[cache_key] = result
|
|
_report_progress(
|
|
progress_callback,
|
|
"found",
|
|
"Found HuggingFace Brave match",
|
|
92,
|
|
repo=repo_id,
|
|
match_type=result.match_type,
|
|
)
|
|
log.info(
|
|
f"HuggingFace found source=brave_fallback match=exact file={filename} repo={repo_id} path={result.extra_value('path')}"
|
|
)
|
|
return result
|
|
|
|
# Not found
|
|
_search_cache[cache_key] = None
|
|
_report_progress(
|
|
progress_callback,
|
|
"done",
|
|
"HuggingFace checked",
|
|
92,
|
|
candidate_count=len(repos_to_check),
|
|
author_repo_count=author_fallback_repo_count,
|
|
brave_candidate_count=len(brave_candidates),
|
|
)
|
|
log.info(
|
|
f"HuggingFace miss file={filename} repos={len(repos_to_check)} author_repos={author_fallback_repo_count} brave={len(brave_candidates)}"
|
|
)
|
|
return None
|
|
|
|
except Exception as e:
|
|
log.exception(f"HuggingFace search error for {filename}: {e}")
|
|
_report_progress(
|
|
progress_callback,
|
|
"error",
|
|
f"HuggingFace search error: {e}",
|
|
100,
|
|
status="error",
|
|
)
|
|
return None
|
|
|
|
|
|
def search_huggingface(
|
|
query: str,
|
|
model_type: Optional[str] = None,
|
|
limit: int = 10,
|
|
token: Optional[str] = None,
|
|
) -> List[SearchResult]:
|
|
"""
|
|
Search HuggingFace Hub for models (general search).
|
|
Returns repos that might be relevant, not guaranteed to have exact file.
|
|
"""
|
|
results = []
|
|
|
|
try:
|
|
params = {"search": query, "limit": limit, "full": "true"}
|
|
|
|
data = execute_provider_json_request(
|
|
"HuggingFace general search",
|
|
f"{HF_API_URL}/models",
|
|
params=params,
|
|
api_key=token,
|
|
timeout=15,
|
|
)
|
|
|
|
if data:
|
|
models = data
|
|
|
|
for model in models:
|
|
repo_id = model.get("id", "")
|
|
|
|
results.append(
|
|
build_model_result(
|
|
"huggingface",
|
|
name=model.get("modelId", repo_id),
|
|
url=f"https://huggingface.co/{repo_id}",
|
|
repo=repo_id,
|
|
downloads=model.get("downloads", 0),
|
|
likes=model.get("likes", 0),
|
|
)
|
|
)
|
|
|
|
except Exception as e:
|
|
log.error(f"HuggingFace search error: {e}")
|
|
|
|
return results
|
|
|
|
|
|
def build_huggingface_custom_result(
|
|
url: str,
|
|
expected_filename: str = "",
|
|
token: Optional[str] = None,
|
|
) -> Optional[SearchResult]:
|
|
from urllib.parse import quote, unquote
|
|
|
|
parsed = parse_huggingface_url(url)
|
|
if not parsed:
|
|
return None
|
|
|
|
repo_id = parsed.repo
|
|
branch = parsed.branch
|
|
file_path = unquote(parsed.filename)
|
|
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}"
|
|
file_info: Dict[str, Any] = {}
|
|
repo_tree = _get_repo_tree(repo_id, headers, branch)
|
|
normalized_file_path = file_path.replace("\\", "/").strip("/")
|
|
for candidate in repo_tree or []:
|
|
candidate_path = str(candidate.get("path") or "").replace("\\", "/").strip("/")
|
|
if candidate_path == normalized_file_path:
|
|
file_info = candidate
|
|
break
|
|
|
|
size = resolve_file_size(
|
|
file_info,
|
|
[download_url],
|
|
probe=lambda url: probe_remote_file_size(url, headers=headers, timeout=10),
|
|
)
|
|
sha256 = _extract_huggingface_file_sha256(file_info)
|
|
|
|
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 build_model_result(
|
|
"huggingface",
|
|
model_id=repo_id,
|
|
version_id=None,
|
|
name=repo_id,
|
|
filename=filename,
|
|
url=download_url,
|
|
download_url=download_url,
|
|
size=size,
|
|
sha256=sha256,
|
|
hashes={"SHA256": sha256} if sha256 else {},
|
|
match_type="custom_url",
|
|
details_source="huggingface",
|
|
repo_id=repo_id,
|
|
repo=repo_id,
|
|
path=file_path,
|
|
branch=branch,
|
|
page_url=page_url,
|
|
version_url=page_url,
|
|
custom_url=True,
|
|
)
|
|
|