597 lines
22 KiB
Python
597 lines
22 KiB
Python
"""Zero-touch Hugging Face provenance discovery for local ComfyUI model files."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import configparser
|
|
import json
|
|
import logging
|
|
import plistlib
|
|
import re
|
|
import struct
|
|
import subprocess
|
|
import sys
|
|
import threading
|
|
from collections.abc import Callable, Mapping, Sequence
|
|
from dataclasses import dataclass, field, replace
|
|
from pathlib import Path, PurePosixPath
|
|
from typing import Any
|
|
from urllib.parse import unquote, urlparse
|
|
|
|
import httpx
|
|
from huggingface_hub import HfApi
|
|
from huggingface_hub.errors import HfHubHTTPError
|
|
|
|
if __package__:
|
|
from .huggingface_assets import (
|
|
HuggingFaceAssetRegistry,
|
|
HuggingFaceAssetSource,
|
|
)
|
|
else: # pragma: no cover - direct debugging imports.
|
|
from huggingface_assets import (
|
|
HuggingFaceAssetRegistry,
|
|
HuggingFaceAssetSource,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_COMMIT_PATTERN = re.compile(r"^[0-9a-f]{40}$")
|
|
_SHA256_PATTERN = re.compile(r"^[0-9a-f]{64}$")
|
|
_SAFE_REPO_COMPONENT = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*$")
|
|
_HUGGINGFACE_URL_PATTERN = re.compile(r"https://huggingface\.co/[^\s\"'<>]+")
|
|
_MAX_SAFETENSORS_HEADER_BYTES = 16 * 1024 * 1024
|
|
_MAX_XATTR_BYTES = 1024 * 1024
|
|
_MAX_MANAGER_FILENAME_FALLBACKS = 8
|
|
WhereFromReader = Callable[[Path], Sequence[str]]
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class HuggingFaceAssetHint:
|
|
"""Describe a candidate Hub location learned from trustworthy local metadata."""
|
|
|
|
repo_id: str
|
|
revision: str
|
|
filename: str
|
|
evidence: str
|
|
|
|
def __post_init__(self) -> None:
|
|
"""Reject malformed candidates before making a Hub API request."""
|
|
repo_parts = self.repo_id.split("/")
|
|
if len(repo_parts) != 2 or not all(
|
|
_SAFE_REPO_COMPONENT.fullmatch(part) for part in repo_parts
|
|
):
|
|
raise ValueError("Hugging Face hint repo_id is invalid.")
|
|
revision = self.revision.strip()
|
|
if (
|
|
not revision
|
|
or len(revision) > 200
|
|
or any(character in revision for character in "\x00\n\r")
|
|
):
|
|
raise ValueError("Hugging Face hint revision is invalid.")
|
|
path = PurePosixPath(self.filename)
|
|
if (
|
|
not self.filename
|
|
or path.is_absolute()
|
|
or any(part in {"", ".", ".."} for part in path.parts)
|
|
):
|
|
raise ValueError("Hugging Face hint filename is unsafe.")
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class LocalHuggingFaceDownloadMetadata:
|
|
"""Hold immutable metadata written by a Hugging Face local-dir download."""
|
|
|
|
commit_hash: str
|
|
etag: str
|
|
filename: str
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class _ManagerHuggingFaceHint:
|
|
"""Associate one Manager catalog hint with its expected local model folder."""
|
|
|
|
hint: HuggingFaceAssetHint
|
|
save_path: str
|
|
|
|
def matches(self, local_path: Path) -> bool:
|
|
"""Return whether the local path ends in Manager's declared save location."""
|
|
normalized_save_path = self.save_path.strip().replace("\\", "/")
|
|
save_parts = tuple(
|
|
part
|
|
for part in PurePosixPath(normalized_save_path).parts
|
|
if part not in {"", ".", "default"}
|
|
)
|
|
if not save_parts:
|
|
return False
|
|
expected_parts = (*save_parts, local_path.name)
|
|
return local_path.absolute().parts[-len(expected_parts) :] == expected_parts
|
|
|
|
|
|
def huggingface_hint_from_url(
|
|
url: str,
|
|
*,
|
|
evidence: str,
|
|
) -> HuggingFaceAssetHint | None:
|
|
"""Parse one official Hugging Face model file URL into a candidate hint."""
|
|
try:
|
|
parsed = urlparse(url.strip())
|
|
except ValueError:
|
|
return None
|
|
if parsed.scheme != "https" or parsed.hostname not in {
|
|
"huggingface.co",
|
|
"www.huggingface.co",
|
|
}:
|
|
return None
|
|
parts = [unquote(part) for part in parsed.path.split("/") if part]
|
|
if len(parts) >= 7 and parts[:3] == ["api", "resolve-cache", "models"]:
|
|
owner, repository, revision = parts[3:6]
|
|
filename_parts = parts[6:]
|
|
elif len(parts) >= 5 and parts[2] in {"blob", "resolve"}:
|
|
owner, repository = parts[:2]
|
|
revision = parts[3]
|
|
filename_parts = parts[4:]
|
|
else:
|
|
return None
|
|
try:
|
|
return HuggingFaceAssetHint(
|
|
repo_id=f"{owner}/{repository.removesuffix('.git')}",
|
|
revision=revision,
|
|
filename=PurePosixPath(*filename_parts).as_posix(),
|
|
evidence=evidence,
|
|
)
|
|
except ValueError:
|
|
return None
|
|
|
|
|
|
def _huggingface_repo_id_from_url(url: str) -> str | None:
|
|
"""Return an owner/repository pair from one Hugging Face Git remote URL."""
|
|
normalized = url.strip()
|
|
if normalized.startswith("git@hf.co:"):
|
|
repository_path = normalized.removeprefix("git@hf.co:")
|
|
else:
|
|
try:
|
|
parsed = urlparse(normalized)
|
|
except ValueError:
|
|
return None
|
|
if parsed.scheme != "https" or parsed.hostname not in {
|
|
"huggingface.co",
|
|
"www.huggingface.co",
|
|
}:
|
|
return None
|
|
repository_path = parsed.path.lstrip("/")
|
|
parts = repository_path.removesuffix(".git").split("/")
|
|
if len(parts) != 2 or not all(_SAFE_REPO_COMPONENT.fullmatch(part) for part in parts):
|
|
return None
|
|
return "/".join(parts)
|
|
|
|
|
|
def _metadata_field(value: object, name: str) -> object | None:
|
|
"""Read one field from a Hugging Face object or test mapping."""
|
|
if isinstance(value, Mapping):
|
|
return value.get(name)
|
|
return getattr(value, name, None)
|
|
|
|
|
|
def _sibling_sha256(sibling: object) -> str | None:
|
|
"""Return the raw LFS/Xet-compatible content digest for one Hub file."""
|
|
lfs = _metadata_field(sibling, "lfs")
|
|
if lfs is None:
|
|
return None
|
|
value = _metadata_field(lfs, "sha256") or _metadata_field(lfs, "oid")
|
|
normalized = str(value or "").strip().lower()
|
|
return normalized if _SHA256_PATTERN.fullmatch(normalized) else None
|
|
|
|
|
|
def resolve_huggingface_asset_hint(
|
|
hint: HuggingFaceAssetHint,
|
|
*,
|
|
sha256: str,
|
|
size_bytes: int,
|
|
api: Any,
|
|
) -> HuggingFaceAssetSource | None:
|
|
"""Resolve a candidate and return it only when Hub metadata matches local bytes."""
|
|
model_info = api.model_info(
|
|
hint.repo_id,
|
|
revision=hint.revision,
|
|
files_metadata=True,
|
|
)
|
|
exact_revision = str(_metadata_field(model_info, "sha") or "").strip().lower()
|
|
if not _COMMIT_PATTERN.fullmatch(exact_revision):
|
|
raise ValueError(
|
|
f"Hugging Face did not resolve {hint.repo_id!r} to an exact commit."
|
|
)
|
|
siblings = _metadata_field(model_info, "siblings")
|
|
if not isinstance(siblings, Sequence) or isinstance(siblings, (str, bytes)):
|
|
return None
|
|
sibling = next(
|
|
(
|
|
item
|
|
for item in siblings
|
|
if str(_metadata_field(item, "rfilename") or "") == hint.filename
|
|
),
|
|
None,
|
|
)
|
|
if sibling is None or _sibling_sha256(sibling) != sha256:
|
|
return None
|
|
lfs = _metadata_field(sibling, "lfs")
|
|
remote_size_value = _metadata_field(sibling, "size") or _metadata_field(lfs, "size")
|
|
if int(remote_size_value or 0) != size_bytes:
|
|
return None
|
|
return HuggingFaceAssetSource(
|
|
repo_id=hint.repo_id,
|
|
revision=exact_revision,
|
|
filename=hint.filename,
|
|
sha256=sha256,
|
|
size_bytes=size_bytes,
|
|
)
|
|
|
|
|
|
def macos_where_from_urls(path: Path) -> tuple[str, ...]:
|
|
"""Read browser source URLs from the macOS download provenance attribute."""
|
|
if sys.platform != "darwin":
|
|
return ()
|
|
try:
|
|
result = subprocess.run(
|
|
(
|
|
"/usr/bin/xattr",
|
|
"-px",
|
|
"com.apple.metadata:kMDItemWhereFroms",
|
|
str(path),
|
|
),
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.DEVNULL,
|
|
check=False,
|
|
timeout=2.0,
|
|
)
|
|
except (FileNotFoundError, subprocess.TimeoutExpired):
|
|
return ()
|
|
if result.returncode != 0 or len(result.stdout) > _MAX_XATTR_BYTES * 3:
|
|
return ()
|
|
try:
|
|
payload = bytes.fromhex(result.stdout.decode("ascii"))
|
|
urls = plistlib.loads(payload)
|
|
except (UnicodeDecodeError, ValueError, plistlib.InvalidFileException):
|
|
return ()
|
|
if not isinstance(urls, Sequence) or isinstance(urls, (str, bytes)):
|
|
return ()
|
|
return tuple(str(url) for url in urls if isinstance(url, str))
|
|
|
|
|
|
def _safetensors_source_urls(path: Path) -> tuple[str, ...]:
|
|
"""Read official Hub URLs embedded in a bounded safetensors metadata header."""
|
|
if path.suffix.casefold() != ".safetensors":
|
|
return ()
|
|
try:
|
|
with path.open("rb") as input_file:
|
|
header_size_payload = input_file.read(8)
|
|
if len(header_size_payload) != 8:
|
|
return ()
|
|
header_size = struct.unpack("<Q", header_size_payload)[0]
|
|
if not (0 < header_size <= _MAX_SAFETENSORS_HEADER_BYTES):
|
|
return ()
|
|
header = json.loads(input_file.read(header_size))
|
|
except (OSError, UnicodeDecodeError, json.JSONDecodeError, struct.error):
|
|
return ()
|
|
metadata = header.get("__metadata__") if isinstance(header, Mapping) else None
|
|
if not isinstance(metadata, Mapping):
|
|
return ()
|
|
urls: list[str] = []
|
|
for value in metadata.values():
|
|
if not isinstance(value, str):
|
|
continue
|
|
urls.extend(_HUGGINGFACE_URL_PATTERN.findall(value))
|
|
return tuple(urls)
|
|
|
|
|
|
def _hub_cache_hint(path: Path) -> HuggingFaceAssetHint | None:
|
|
"""Recognize the documented models--owner--repo/snapshots/commit cache layout."""
|
|
parts = path.absolute().parts
|
|
for index, part in enumerate(parts):
|
|
if not part.startswith("models--") or index + 2 >= len(parts):
|
|
continue
|
|
repository_parts = part.removeprefix("models--").split("--", maxsplit=1)
|
|
if len(repository_parts) != 2 or parts[index + 1] != "snapshots":
|
|
continue
|
|
revision = parts[index + 2]
|
|
filename_parts = parts[index + 3 :]
|
|
if not filename_parts:
|
|
continue
|
|
try:
|
|
return HuggingFaceAssetHint(
|
|
repo_id="/".join(repository_parts),
|
|
revision=revision,
|
|
filename=PurePosixPath(*filename_parts).as_posix(),
|
|
evidence="Hugging Face cache",
|
|
)
|
|
except ValueError:
|
|
return None
|
|
return None
|
|
|
|
|
|
def _git_head_commit(git_directory: Path) -> str | None:
|
|
"""Resolve one ordinary Git HEAD without executing a subprocess."""
|
|
try:
|
|
head = (git_directory / "HEAD").read_text(encoding="utf-8").strip()
|
|
except OSError:
|
|
return None
|
|
if _COMMIT_PATTERN.fullmatch(head):
|
|
return head
|
|
if not head.startswith("ref: "):
|
|
return None
|
|
reference = head.removeprefix("ref: ").strip()
|
|
reference_path = git_directory / reference
|
|
try:
|
|
value = reference_path.read_text(encoding="utf-8").strip()
|
|
except OSError:
|
|
value = ""
|
|
if _COMMIT_PATTERN.fullmatch(value):
|
|
return value
|
|
try:
|
|
packed_refs = (git_directory / "packed-refs").read_text(encoding="utf-8")
|
|
except OSError:
|
|
return None
|
|
for line in packed_refs.splitlines():
|
|
fields = line.split()
|
|
if len(fields) == 2 and fields[1] == reference and _COMMIT_PATTERN.fullmatch(fields[0]):
|
|
return fields[0]
|
|
return None
|
|
|
|
|
|
def _git_checkout_hint(path: Path) -> HuggingFaceAssetHint | None:
|
|
"""Recognize a model file inside an ordinary Hugging Face Git checkout."""
|
|
for repository_root in path.parents:
|
|
git_directory = repository_root / ".git"
|
|
if not git_directory.is_dir():
|
|
continue
|
|
parser = configparser.ConfigParser(interpolation=None)
|
|
try:
|
|
parser.read(git_directory / "config", encoding="utf-8")
|
|
except (OSError, configparser.Error):
|
|
return None
|
|
repo_id = next(
|
|
(
|
|
repo_id
|
|
for section in parser.sections()
|
|
if section.startswith("remote ")
|
|
for repo_id in [_huggingface_repo_id_from_url(parser.get(section, "url", fallback=""))]
|
|
if repo_id is not None
|
|
),
|
|
None,
|
|
)
|
|
revision = _git_head_commit(git_directory)
|
|
if repo_id is None or revision is None:
|
|
return None
|
|
try:
|
|
filename = path.absolute().relative_to(repository_root.absolute()).as_posix()
|
|
return HuggingFaceAssetHint(
|
|
repo_id=repo_id,
|
|
revision=revision,
|
|
filename=filename,
|
|
evidence="Hugging Face Git checkout",
|
|
)
|
|
except ValueError:
|
|
return None
|
|
return None
|
|
|
|
|
|
def _local_download_metadata(path: Path, sha256: str) -> LocalHuggingFaceDownloadMetadata | None:
|
|
"""Find matching local-dir metadata and its original repository-relative filename."""
|
|
absolute_path = path.absolute()
|
|
for local_directory in absolute_path.parents:
|
|
try:
|
|
filename = absolute_path.relative_to(local_directory).as_posix()
|
|
except ValueError:
|
|
continue
|
|
metadata_path = (
|
|
local_directory
|
|
/ ".cache"
|
|
/ "huggingface"
|
|
/ "download"
|
|
/ f"{filename}.metadata"
|
|
)
|
|
if not metadata_path.is_file():
|
|
continue
|
|
try:
|
|
lines = metadata_path.read_text(encoding="utf-8").splitlines()
|
|
except OSError:
|
|
return None
|
|
if len(lines) >= 2 and _COMMIT_PATTERN.fullmatch(lines[0]) and lines[1] == sha256:
|
|
return LocalHuggingFaceDownloadMetadata(
|
|
commit_hash=lines[0],
|
|
etag=lines[1],
|
|
filename=filename,
|
|
)
|
|
return None
|
|
|
|
|
|
@dataclass
|
|
class HuggingFaceAssetDiscovery:
|
|
"""Learn and persist immutable Hub provenance from normal model installation state."""
|
|
|
|
registry: HuggingFaceAssetRegistry
|
|
user_directory: Path
|
|
comfyui_root: Path | None
|
|
api: Any = field(default_factory=HfApi)
|
|
where_from_reader: WhereFromReader = macos_where_from_urls
|
|
_manager_hints_by_filename: (
|
|
dict[str, tuple[_ManagerHuggingFaceHint, ...]] | None
|
|
) = field(init=False, default=None)
|
|
_manager_catalog_signature: tuple[tuple[str, int, int], ...] | None = field(
|
|
init=False,
|
|
default=None,
|
|
)
|
|
_unresolved_evidence_by_sha256: dict[str, tuple[tuple[str, str, str], ...]] = field(
|
|
init=False,
|
|
default_factory=dict,
|
|
)
|
|
_cache_lock: threading.Lock = field(init=False, default_factory=threading.Lock)
|
|
|
|
def discover(self, local_path: Path, *, sha256: str) -> HuggingFaceAssetSource | None:
|
|
"""Return cached or automatically verified provenance for one local file."""
|
|
registered = self.registry.get(sha256)
|
|
if registered is not None:
|
|
return registered
|
|
source_path = local_path.expanduser().absolute()
|
|
size_bytes = source_path.stat().st_size
|
|
metadata = _local_download_metadata(source_path, sha256)
|
|
hints = self._candidate_hints(source_path)
|
|
if metadata is not None:
|
|
hints = tuple(
|
|
replace(
|
|
hint,
|
|
revision=metadata.commit_hash,
|
|
filename=(
|
|
metadata.filename
|
|
if hint.filename == source_path.name
|
|
else hint.filename
|
|
),
|
|
)
|
|
for hint in hints
|
|
)
|
|
evidence_fingerprint = tuple(
|
|
(hint.repo_id, hint.revision, hint.filename) for hint in hints
|
|
)
|
|
with self._cache_lock:
|
|
if self._unresolved_evidence_by_sha256.get(sha256) == evidence_fingerprint:
|
|
return None
|
|
for hint in hints:
|
|
try:
|
|
source = resolve_huggingface_asset_hint(
|
|
hint,
|
|
sha256=sha256,
|
|
size_bytes=size_bytes,
|
|
api=self.api,
|
|
)
|
|
except (HfHubHTTPError, httpx.HTTPError, OSError, ValueError) as exc:
|
|
logger.debug(
|
|
"Hugging Face provenance candidate failed evidence=%s repo=%s file=%s: %s",
|
|
hint.evidence,
|
|
hint.repo_id,
|
|
hint.filename,
|
|
exc,
|
|
)
|
|
continue
|
|
if source is not None:
|
|
logger.info(
|
|
"Automatically discovered Hugging Face asset source=%s evidence=%s.",
|
|
source.display_reference,
|
|
hint.evidence,
|
|
)
|
|
return self.registry.upsert(source)
|
|
with self._cache_lock:
|
|
self._unresolved_evidence_by_sha256[sha256] = evidence_fingerprint
|
|
return None
|
|
|
|
def _candidate_hints(self, local_path: Path) -> tuple[HuggingFaceAssetHint, ...]:
|
|
"""Collect unique candidate locations from local install provenance."""
|
|
candidates: list[HuggingFaceAssetHint] = []
|
|
for path_candidate in (local_path, local_path.resolve()):
|
|
for hint in (_hub_cache_hint(path_candidate), _git_checkout_hint(path_candidate)):
|
|
if hint is not None:
|
|
candidates.append(hint)
|
|
for url in (*self.where_from_reader(local_path), *_safetensors_source_urls(local_path)):
|
|
hint = huggingface_hint_from_url(url, evidence="local download metadata")
|
|
if hint is not None:
|
|
candidates.append(hint)
|
|
candidates.extend(self._manager_candidates(local_path))
|
|
unique: dict[tuple[str, str, str], HuggingFaceAssetHint] = {}
|
|
for hint in candidates:
|
|
unique.setdefault((hint.repo_id, hint.revision, hint.filename), hint)
|
|
return tuple(unique.values())
|
|
|
|
def _manager_candidates(self, local_path: Path) -> tuple[HuggingFaceAssetHint, ...]:
|
|
"""Prefer Manager entries whose configured save path matches this model."""
|
|
records = self._manager_hints().get(local_path.name, ())
|
|
matching = tuple(record.hint for record in records if record.matches(local_path))
|
|
if matching:
|
|
return matching
|
|
return tuple(
|
|
record.hint for record in records[:_MAX_MANAGER_FILENAME_FALLBACKS]
|
|
)
|
|
|
|
def _manager_hints(self) -> dict[str, tuple[_ManagerHuggingFaceHint, ...]]:
|
|
"""Load Manager's model catalog and refresh it when installation state changes."""
|
|
catalog_paths = self._manager_catalog_paths()
|
|
signature = self._catalog_signature(catalog_paths)
|
|
with self._cache_lock:
|
|
if (
|
|
self._manager_hints_by_filename is not None
|
|
and self._manager_catalog_signature == signature
|
|
):
|
|
return self._manager_hints_by_filename
|
|
loaded = self._load_manager_hints(catalog_paths)
|
|
self._manager_hints_by_filename = loaded
|
|
self._manager_catalog_signature = signature
|
|
return loaded
|
|
|
|
def _load_manager_hints(
|
|
self,
|
|
catalog_paths: tuple[Path, ...],
|
|
) -> dict[str, tuple[_ManagerHuggingFaceHint, ...]]:
|
|
"""Parse every available Manager model-list cache without trusting its bytes."""
|
|
mutable: dict[str, list[_ManagerHuggingFaceHint]] = {}
|
|
for catalog_path in catalog_paths:
|
|
try:
|
|
payload = json.loads(catalog_path.read_text(encoding="utf-8"))
|
|
except (OSError, json.JSONDecodeError):
|
|
continue
|
|
models = payload.get("models") if isinstance(payload, Mapping) else None
|
|
if not isinstance(models, Sequence) or isinstance(models, (str, bytes)):
|
|
continue
|
|
for model in models:
|
|
if not isinstance(model, Mapping):
|
|
continue
|
|
local_filename = str(model.get("filename") or "").strip()
|
|
url = str(model.get("url") or "").strip()
|
|
hint = huggingface_hint_from_url(url, evidence="ComfyUI Manager catalog")
|
|
if local_filename and hint is not None:
|
|
mutable.setdefault(local_filename, []).append(
|
|
_ManagerHuggingFaceHint(
|
|
hint=hint,
|
|
save_path=str(model.get("save_path") or ""),
|
|
)
|
|
)
|
|
return {
|
|
filename: tuple(hints)
|
|
for filename, hints in mutable.items()
|
|
}
|
|
|
|
@staticmethod
|
|
def _catalog_signature(
|
|
catalog_paths: tuple[Path, ...],
|
|
) -> tuple[tuple[str, int, int], ...]:
|
|
"""Return a stable signature that changes when Manager rewrites its catalog."""
|
|
signature: list[tuple[str, int, int]] = []
|
|
for path in catalog_paths:
|
|
try:
|
|
stat_result = path.stat()
|
|
except OSError:
|
|
continue
|
|
signature.append((str(path), stat_result.st_mtime_ns, stat_result.st_size))
|
|
return tuple(signature)
|
|
|
|
def _manager_catalog_paths(self) -> tuple[Path, ...]:
|
|
"""Return current and legacy ComfyUI Manager model-list cache paths."""
|
|
candidates = list(
|
|
(self.user_directory / "__manager" / "cache").glob("*_model-list.json")
|
|
)
|
|
if self.comfyui_root is not None:
|
|
candidates.extend(
|
|
(
|
|
self.comfyui_root
|
|
/ "custom_nodes"
|
|
/ "ComfyUI-Manager"
|
|
/ "cache"
|
|
).glob("*_model-list.json")
|
|
)
|
|
return tuple(sorted(set(path.resolve() for path in candidates if path.is_file())))
|
|
|
|
|
|
__all__ = [
|
|
"HuggingFaceAssetDiscovery",
|
|
"HuggingFaceAssetHint",
|
|
"LocalHuggingFaceDownloadMetadata",
|
|
"huggingface_hint_from_url",
|
|
"macos_where_from_urls",
|
|
"resolve_huggingface_asset_hint",
|
|
]
|