Automate Hugging Face asset discovery
This commit is contained in:
@@ -13,6 +13,7 @@
|
||||
!/execution_environments.py
|
||||
!/execution_history.py
|
||||
!/huggingface_assets.py
|
||||
!/huggingface_discovery.py
|
||||
!/instance_identity.py
|
||||
!/llm_profiles.json
|
||||
!/llm_profiles.py
|
||||
|
||||
@@ -301,18 +301,13 @@ The sync engine looks for node inputs that resolve to model files ending in `.sa
|
||||
|
||||
Assets are content-addressed and uploaded once: repeated references across nodes and components share one hash and upload decision, and unchanged files are skipped on later runs. In local mode, a local mirror stands in for remote storage.
|
||||
|
||||
For Vast.ai, a single-file model can instead be downloaded directly from Hugging Face over the rented instance's data-center connection. Register the association once from the source checkout:
|
||||
For Vast.ai, a single-file model can instead be downloaded directly from Hugging Face over the rented instance's data-center connection. This is automatic for models installed through ComfyUI Manager, files or symlinks in the standard Hugging Face cache, Hugging Face Git checkouts, browser downloads that retain their source URL, and safetensors that embed an official Hub source URL. These sources are treated only as candidates: Modal-Sync resolves the revision to an exact 40-character commit and requires the Hub's LFS/Xet SHA-256 and size to match the local bytes before persisting the learned mapping below `<ComfyUI user directory>/comfyui-modal/huggingface-assets.json`.
|
||||
|
||||
```bash
|
||||
COMFYUI_ROOT=/path/to/ComfyUI \
|
||||
uv run python scripts/register_huggingface_asset.py \
|
||||
/path/to/model.safetensors owner/model weights/model.safetensors \
|
||||
--revision main
|
||||
```
|
||||
Future Vast cache misses use `huggingface_hub` with `hf-xet` and high-performance mode inside the worker, verify the downloaded size and SHA-256 again, then atomically publish the file under the same content-addressed `/storage/assets/...` path used by ordinary sync. The configurator panel reports both automatic source identification and the Hugging Face download while preparation is active. Because Hugging Face does not provide a global reverse lookup from an arbitrary digest to every repository path, an originless file simply falls back to the normal streamed SSH upload. No user action is required.
|
||||
|
||||
Registration resolves the branch or tag to an exact 40-character commit, requires the Hub's LFS/Xet SHA-256 and size to match the local file, and stores only that immutable provenance below `<ComfyUI user directory>/comfyui-modal/huggingface-assets.json`. Hugging Face does not provide a global reverse lookup from an arbitrary digest to every repository path, so existing files need this one-time repository association. Future Vast cache misses use `huggingface_hub` with `hf-xet` and high-performance mode inside the worker, verify the downloaded size and SHA-256 again, then atomically publish the file under the same content-addressed `/storage/assets/...` path used by ordinary sync. The configurator panel reports the Hugging Face download while preparation is active.
|
||||
The `scripts/register_huggingface_asset.py` utility remains available as a diagnostic override for a file whose installation metadata was removed. It is not part of normal model installation or workflow use.
|
||||
|
||||
Public files need no token. For private or gated files, set `HF_TOKEN` in the local ComfyUI process. Modal-Sync sends it only through the protected SSH standard input of the transient materializer; it is not stored in the workflow, registry, Vast launch environment, command arguments, or logs. If Hugging Face is unavailable, authorization fails, or verification rejects the downloaded bytes, preparation reports the fallback and streams the local file through SSH without buffering the complete model in ComfyUI memory.
|
||||
Public files need no token. For private or gated files, set `HF_TOKEN` in the local ComfyUI process. Modal-Sync uses it for metadata verification and sends it only through the protected SSH standard input of the transient materializer; it is not stored in the workflow, registry, Vast launch environment, command arguments, or logs. If Hugging Face is unavailable, authorization fails, or verification rejects the downloaded bytes, preparation reports the fallback and streams the local file through SSH without buffering the complete model in ComfyUI memory.
|
||||
|
||||
The worker-side materializer is part of the versioned runtime source. Rebuild and publish the SSH/Vast worker image after upgrading before expecting an already configured Vast deployment to use this path.
|
||||
|
||||
|
||||
@@ -4,6 +4,7 @@
|
||||
|
||||
## 2026-08-24
|
||||
|
||||
- A custom node cannot require per-model registration as its ordinary asset workflow. Treat ComfyUI Manager catalogs, standard Hub cache paths, Hugging Face Git checkouts, browser download provenance, and embedded safetensors URLs as untrusted location hints; require the Hub's exact content digest and size to match the local file, then persist the verified mapping automatically. Preserve provenance-bearing symlink paths until discovery, and use streamed SSH for genuinely originless files.
|
||||
- Hugging Face Xet is an origin transport, not a global reverse-SHA search service. Preserve a verified mapping from local content SHA-256 to repository, exact commit, filename, and size; let the Vast worker download through `huggingface_hub`, verify the local content contract again, and atomically publish at the existing remote asset path. Keep unauthorized or unavailable Hub sources recoverable through a streamed SSH fallback, and never place `HF_TOKEN` in a workflow, registry, launch environment, command line, or log.
|
||||
- A storage backend can appear to stream over SSH while still buffering the entire input before `subprocess.run`. Passing a regular file as OpenSSH stdin gives bounded-memory kernel streaming; the remote atomic writer must also copy stdin in chunks rather than calling `read()` with no bound.
|
||||
- A configuration node's ordinary output is the wrong transport for queue-time planning UX because ComfyUI does not execute the node until its normal scheduler reaches it. Keep planning authoritative in the queue route, carry the serialized configurator node id in replayable websocket status, and render a non-serialized DOM widget directly on that node. Populate the final assignment table from the accepted queue response and retain it after prompt cleanup so setup progress and placement remain inspectable out of band.
|
||||
|
||||
@@ -0,0 +1,596 @@
|
||||
"""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",
|
||||
]
|
||||
@@ -5,7 +5,7 @@ from __future__ import annotations
|
||||
import argparse
|
||||
import logging
|
||||
import sys
|
||||
from collections.abc import Mapping, Sequence
|
||||
from collections.abc import Sequence
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
@@ -18,6 +18,10 @@ from huggingface_assets import ( # noqa: E402
|
||||
HuggingFaceAssetSource,
|
||||
sha256_file,
|
||||
)
|
||||
from huggingface_discovery import ( # noqa: E402
|
||||
HuggingFaceAssetHint,
|
||||
resolve_huggingface_asset_hint,
|
||||
)
|
||||
from settings import discover_comfyui_user_directory, get_settings # noqa: E402
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -45,23 +49,6 @@ def _parser() -> argparse.ArgumentParser:
|
||||
return parser
|
||||
|
||||
|
||||
def _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_lfs_sha256(sibling: object) -> str | None:
|
||||
"""Return the raw content SHA-256 advertised for one Hub LFS/Xet file."""
|
||||
lfs = _field(sibling, "lfs")
|
||||
if lfs is None:
|
||||
return None
|
||||
value = _field(lfs, "sha256") or _field(lfs, "oid")
|
||||
normalized = str(value or "").strip().lower()
|
||||
return normalized or None
|
||||
|
||||
|
||||
def resolve_huggingface_asset_source(
|
||||
local_path: Path,
|
||||
*,
|
||||
@@ -74,55 +61,24 @@ def resolve_huggingface_asset_source(
|
||||
resolved_path = local_path.expanduser().resolve()
|
||||
if not resolved_path.is_file():
|
||||
raise FileNotFoundError(f"Local model asset not found: {resolved_path}")
|
||||
model_info = api.model_info(
|
||||
repo_id,
|
||||
revision=revision,
|
||||
files_metadata=True,
|
||||
)
|
||||
exact_revision = str(_field(model_info, "sha") or "").strip().lower()
|
||||
siblings = _field(model_info, "siblings")
|
||||
if not isinstance(siblings, Sequence):
|
||||
raise ValueError(f"Hugging Face repository {repo_id!r} returned no file metadata.")
|
||||
sibling = next(
|
||||
(
|
||||
item
|
||||
for item in siblings
|
||||
if str(_field(item, "rfilename") or "") == filename
|
||||
),
|
||||
None,
|
||||
)
|
||||
if sibling is None:
|
||||
raise ValueError(
|
||||
f"Hugging Face repository {repo_id!r} has no file {filename!r} at "
|
||||
f"revision {revision!r}."
|
||||
)
|
||||
local_sha256 = sha256_file(resolved_path)
|
||||
remote_sha256 = _sibling_lfs_sha256(sibling)
|
||||
if remote_sha256 is None:
|
||||
raise ValueError(
|
||||
f"Hugging Face file {repo_id}/{filename} does not expose an LFS/Xet "
|
||||
"content SHA-256 and cannot be registered safely."
|
||||
)
|
||||
if remote_sha256 != local_sha256:
|
||||
raise ValueError(
|
||||
"Local file SHA-256 does not match Hugging Face metadata: "
|
||||
f"local={local_sha256} remote={remote_sha256}."
|
||||
)
|
||||
remote_size_value = _field(sibling, "size") or _field(_field(sibling, "lfs"), "size")
|
||||
remote_size = int(remote_size_value or 0)
|
||||
local_size = resolved_path.stat().st_size
|
||||
if remote_size != local_size:
|
||||
raise ValueError(
|
||||
"Local file size does not match Hugging Face metadata: "
|
||||
f"local={local_size} remote={remote_size}."
|
||||
)
|
||||
return HuggingFaceAssetSource(
|
||||
repo_id=repo_id,
|
||||
revision=exact_revision,
|
||||
filename=filename,
|
||||
source = resolve_huggingface_asset_hint(
|
||||
HuggingFaceAssetHint(
|
||||
repo_id=repo_id,
|
||||
revision=revision,
|
||||
filename=filename,
|
||||
evidence="manual diagnostic override",
|
||||
),
|
||||
sha256=local_sha256,
|
||||
size_bytes=local_size,
|
||||
size_bytes=resolved_path.stat().st_size,
|
||||
api=api,
|
||||
)
|
||||
if source is None:
|
||||
raise ValueError(
|
||||
"Local file does not match Hugging Face metadata for "
|
||||
f"{repo_id}/{filename} at revision {revision!r}."
|
||||
)
|
||||
return source
|
||||
|
||||
|
||||
def _registry(explicit_path: Path | None) -> HuggingFaceAssetRegistry:
|
||||
|
||||
+71
-26
@@ -17,6 +17,7 @@ from pathlib import Path
|
||||
from typing import Any, Callable, Iterable, Protocol, runtime_checkable
|
||||
|
||||
from .huggingface_assets import HuggingFaceAssetRegistry, HuggingFaceAssetSource
|
||||
from .huggingface_discovery import HuggingFaceAssetDiscovery
|
||||
from .settings import ModalSyncSettings, get_settings
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -76,7 +77,7 @@ def resolve_model_path(
|
||||
if path.suffix.lower() not in extensions:
|
||||
return None
|
||||
if path.is_file():
|
||||
return path.resolve()
|
||||
return path.absolute()
|
||||
if os.path.isabs(value):
|
||||
return None
|
||||
|
||||
@@ -89,12 +90,12 @@ def resolve_model_path(
|
||||
for folder_name in folder_paths.folder_names_and_paths:
|
||||
full_path = folder_paths.get_full_path(folder_name, value)
|
||||
if full_path is not None:
|
||||
return Path(full_path).resolve()
|
||||
return Path(full_path).absolute()
|
||||
|
||||
if comfyui_root is not None:
|
||||
candidate = comfyui_root / value
|
||||
if candidate.is_file():
|
||||
return candidate.resolve()
|
||||
return candidate.absolute()
|
||||
return None
|
||||
|
||||
|
||||
@@ -138,6 +139,21 @@ def _format_huggingface_download_status(
|
||||
return f"Downloading asset from Hugging Face on Vast.ai: {asset_name}"
|
||||
|
||||
|
||||
def _format_huggingface_discovery_status(
|
||||
asset_name: str,
|
||||
*,
|
||||
item_index: int | None,
|
||||
total_items: int | None,
|
||||
) -> str:
|
||||
"""Return one queue-time status for automatic Hugging Face source discovery."""
|
||||
if item_index is not None and total_items is not None and total_items > 1:
|
||||
return (
|
||||
f"Identifying Hugging Face source {item_index}/{total_items} for "
|
||||
f"Vast.ai: {asset_name}"
|
||||
)
|
||||
return f"Identifying Hugging Face source for Vast.ai: {asset_name}"
|
||||
|
||||
|
||||
class VolumeBackend(Protocol):
|
||||
"""Minimal storage interface needed by the sync engine."""
|
||||
|
||||
@@ -216,9 +232,9 @@ class AssetSyncRequestCache:
|
||||
def synced_assets(self) -> tuple[SyncedAsset, ...]:
|
||||
"""Return unique sync results in request planning order."""
|
||||
return tuple(
|
||||
self._assets_by_path[path]
|
||||
self._assets_by_path[path.resolve()]
|
||||
for path in self.planned_paths
|
||||
if path in self._assets_by_path
|
||||
if path.resolve() in self._assets_by_path
|
||||
)
|
||||
|
||||
|
||||
@@ -544,6 +560,7 @@ class ModalAssetSyncEngine:
|
||||
settings: ModalSyncSettings
|
||||
sync_index: SyncIndexBackend | None = None
|
||||
huggingface_asset_registry: HuggingFaceAssetRegistry | None = None
|
||||
huggingface_asset_discovery: HuggingFaceAssetDiscovery | None = None
|
||||
_hash_cache: dict[str, dict[str, Any]] = field(init=False, default_factory=dict)
|
||||
_path_resolution_cache: dict[str, str | None] = field(init=False, default_factory=dict)
|
||||
_hash_cache_dirty: bool = field(init=False, default=False)
|
||||
@@ -600,24 +617,33 @@ class ModalAssetSyncEngine:
|
||||
total_items: int | None = None,
|
||||
) -> SyncedAsset:
|
||||
"""Sync a file into content-addressable remote storage."""
|
||||
resolved_path = local_path.expanduser().resolve()
|
||||
if not resolved_path.is_file():
|
||||
raise FileNotFoundError(f"Asset not found: {resolved_path}")
|
||||
source_path = local_path.expanduser().absolute()
|
||||
if not source_path.is_file():
|
||||
raise FileNotFoundError(f"Asset not found: {source_path}")
|
||||
resolved_path = source_path.resolve()
|
||||
|
||||
sha256 = self._hash_file(resolved_path)
|
||||
proposed_remote_path = f"{remote_folder.rstrip('/')}/{sha256}_{resolved_path.name}"
|
||||
huggingface_source = self._huggingface_source_for_asset(
|
||||
resolved_path,
|
||||
sha256=sha256,
|
||||
proposed_remote_path = f"{remote_folder.rstrip('/')}/{sha256}_{source_path.name}"
|
||||
sync_key = self._asset_sync_index_key(sha256)
|
||||
huggingface_source = (
|
||||
self._huggingface_source_for_asset(
|
||||
source_path,
|
||||
sha256=sha256,
|
||||
status_callback=status_callback,
|
||||
item_index=item_index,
|
||||
total_items=total_items,
|
||||
)
|
||||
if self._lookup_sync_record(sync_key) is None
|
||||
else None
|
||||
)
|
||||
sync_result = self._sync_content_addressed_file(
|
||||
local_path=resolved_path,
|
||||
local_path=source_path,
|
||||
remote_path=proposed_remote_path,
|
||||
sync_key=self._asset_sync_index_key(sha256),
|
||||
source_description=str(resolved_path),
|
||||
sync_key=sync_key,
|
||||
source_description=str(source_path),
|
||||
status_callback=status_callback,
|
||||
upload_status_message=_format_asset_upload_status(
|
||||
resolved_path.name,
|
||||
source_path.name,
|
||||
item_index=item_index,
|
||||
total_items=total_items,
|
||||
destination=self._destination_label(),
|
||||
@@ -628,7 +654,7 @@ class ModalAssetSyncEngine:
|
||||
)
|
||||
|
||||
return SyncedAsset(
|
||||
local_path=resolved_path,
|
||||
local_path=source_path,
|
||||
remote_path=sync_result.remote_path,
|
||||
sha256=sha256,
|
||||
uploaded=sync_result.uploaded,
|
||||
@@ -693,11 +719,11 @@ class ModalAssetSyncEngine:
|
||||
prompt_input_values: Iterable[Any],
|
||||
) -> AssetSyncRequestCache:
|
||||
"""Plan unique syncable assets for one queued prompt in stable encounter order."""
|
||||
unique_paths: dict[Path, None] = {}
|
||||
unique_paths: dict[Path, Path] = {}
|
||||
for prompt_input_value in prompt_input_values:
|
||||
for local_path in self._collect_syncable_asset_paths(prompt_input_value):
|
||||
unique_paths.setdefault(local_path.resolve(), None)
|
||||
return AssetSyncRequestCache(planned_paths=tuple(unique_paths))
|
||||
unique_paths.setdefault(local_path.resolve(), local_path.absolute())
|
||||
return AssetSyncRequestCache(planned_paths=tuple(unique_paths.values()))
|
||||
|
||||
def sync_custom_nodes_directory(
|
||||
self,
|
||||
@@ -942,14 +968,33 @@ class ModalAssetSyncEngine:
|
||||
local_path: Path,
|
||||
*,
|
||||
sha256: str,
|
||||
status_callback: SyncStatusCallback | None,
|
||||
item_index: int | None,
|
||||
total_items: int | None,
|
||||
) -> HuggingFaceAssetSource | None:
|
||||
"""Return matching provenance when this engine can materialize it remotely."""
|
||||
if self.huggingface_asset_registry is None or not isinstance(
|
||||
self.volume,
|
||||
HuggingFaceMaterializingBackend,
|
||||
):
|
||||
"""Return registered or automatically discovered immutable provenance."""
|
||||
if not isinstance(self.volume, HuggingFaceMaterializingBackend):
|
||||
return None
|
||||
source = self.huggingface_asset_registry.get(sha256)
|
||||
source = (
|
||||
self.huggingface_asset_registry.get(sha256)
|
||||
if self.huggingface_asset_registry is not None
|
||||
else None
|
||||
)
|
||||
if source is None and self.huggingface_asset_discovery is not None:
|
||||
_emit_sync_status(
|
||||
status_callback,
|
||||
_format_huggingface_discovery_status(
|
||||
local_path.name,
|
||||
item_index=item_index,
|
||||
total_items=total_items,
|
||||
),
|
||||
item_index,
|
||||
total_items,
|
||||
)
|
||||
source = self.huggingface_asset_discovery.discover(
|
||||
local_path,
|
||||
sha256=sha256,
|
||||
)
|
||||
if source is None:
|
||||
return None
|
||||
actual_size = local_path.stat().st_size
|
||||
|
||||
@@ -214,6 +214,12 @@ def huggingface_assets_module(extension_package: object) -> object:
|
||||
return importlib.import_module(f"{PACKAGE_NAME}.huggingface_assets")
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def huggingface_discovery_module(extension_package: object) -> object:
|
||||
"""Return automatic Hugging Face asset provenance discovery helpers."""
|
||||
return importlib.import_module(f"{PACKAGE_NAME}.huggingface_discovery")
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def huggingface_materializer_module(extension_package: object) -> object:
|
||||
"""Return the remote Hugging Face file materializer."""
|
||||
|
||||
@@ -0,0 +1,414 @@
|
||||
"""Tests for zero-touch Hugging Face model provenance discovery."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import struct
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
|
||||
class FakeHubApi:
|
||||
"""Return deterministic file metadata for candidate verification tests."""
|
||||
|
||||
def __init__(self, responses: dict[tuple[str, str], object]) -> None:
|
||||
"""Store responses indexed by repository and requested revision."""
|
||||
self.responses = responses
|
||||
self.calls: list[tuple[str, str]] = []
|
||||
|
||||
def model_info(self, repo_id: str, **kwargs: object) -> object:
|
||||
"""Return one configured repository response and record the request."""
|
||||
revision = str(kwargs["revision"])
|
||||
assert kwargs["files_metadata"] is True
|
||||
self.calls.append((repo_id, revision))
|
||||
return self.responses[(repo_id, revision)]
|
||||
|
||||
|
||||
def _hub_response(
|
||||
*,
|
||||
filename: str,
|
||||
sha256: str,
|
||||
size_bytes: int,
|
||||
commit: str,
|
||||
) -> dict[str, object]:
|
||||
"""Build one Hub model-info payload containing exact file metadata."""
|
||||
return {
|
||||
"sha": commit,
|
||||
"siblings": [
|
||||
{
|
||||
"rfilename": filename,
|
||||
"size": size_bytes,
|
||||
"lfs": {"sha256": sha256},
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def _discovery(
|
||||
module: Any,
|
||||
assets_module: Any,
|
||||
*,
|
||||
tmp_path: Path,
|
||||
user_directory: Path,
|
||||
api: FakeHubApi,
|
||||
where_from_urls: tuple[str, ...] = (),
|
||||
) -> Any:
|
||||
"""Create an isolated discovery service with deterministic local evidence."""
|
||||
registry = assets_module.HuggingFaceAssetRegistry(
|
||||
tmp_path / "registry" / "huggingface-assets.json"
|
||||
)
|
||||
return module.HuggingFaceAssetDiscovery(
|
||||
registry=registry,
|
||||
user_directory=user_directory,
|
||||
comfyui_root=tmp_path / "comfyui",
|
||||
api=api,
|
||||
where_from_reader=lambda path: where_from_urls,
|
||||
)
|
||||
|
||||
|
||||
def test_parses_official_resolve_and_cache_urls(
|
||||
huggingface_discovery_module: Any,
|
||||
) -> None:
|
||||
"""Normal and CDN-backed Hub URLs should yield the same immutable hint shape."""
|
||||
resolve_hint = huggingface_discovery_module.huggingface_hint_from_url(
|
||||
"https://huggingface.co/owner/model/resolve/main/weights/model.safetensors?download=true",
|
||||
evidence="browser",
|
||||
)
|
||||
cache_hint = huggingface_discovery_module.huggingface_hint_from_url(
|
||||
"https://huggingface.co/api/resolve-cache/models/owner/model/"
|
||||
f"{'a' * 40}/weights/model.safetensors",
|
||||
evidence="browser",
|
||||
)
|
||||
|
||||
assert resolve_hint.repo_id == "owner/model"
|
||||
assert resolve_hint.revision == "main"
|
||||
assert resolve_hint.filename == "weights/model.safetensors"
|
||||
assert cache_hint.repo_id == "owner/model"
|
||||
assert cache_hint.revision == "a" * 40
|
||||
assert cache_hint.filename == "weights/model.safetensors"
|
||||
|
||||
|
||||
def test_discovers_manager_installed_model_and_persists_verified_mapping(
|
||||
huggingface_assets_module: Any,
|
||||
huggingface_discovery_module: Any,
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""Manager's catalog should make ordinary installs zero-touch after hash validation."""
|
||||
user_directory = tmp_path / "user"
|
||||
catalog = user_directory / "__manager" / "cache" / "123_model-list.json"
|
||||
catalog.parent.mkdir(parents=True)
|
||||
catalog.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"models": [
|
||||
{
|
||||
"filename": "manager-model.safetensors",
|
||||
"save_path": "vae",
|
||||
"url": (
|
||||
"https://huggingface.co/wrong/model/resolve/main/"
|
||||
"manager-model.safetensors"
|
||||
),
|
||||
},
|
||||
{
|
||||
"filename": "manager-model.safetensors",
|
||||
"save_path": "checkpoints",
|
||||
"url": (
|
||||
"https://huggingface.co/owner/model/resolve/main/"
|
||||
"weights/manager-model.safetensors"
|
||||
),
|
||||
}
|
||||
]
|
||||
}
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
asset_path = tmp_path / "models" / "checkpoints" / "manager-model.safetensors"
|
||||
asset_path.parent.mkdir(parents=True)
|
||||
asset_path.write_bytes(b"manager-model")
|
||||
sha256 = huggingface_assets_module.sha256_file(asset_path)
|
||||
commit = "b" * 40
|
||||
api = FakeHubApi(
|
||||
{
|
||||
("owner/model", "main"): _hub_response(
|
||||
filename="weights/manager-model.safetensors",
|
||||
sha256=sha256,
|
||||
size_bytes=asset_path.stat().st_size,
|
||||
commit=commit,
|
||||
)
|
||||
}
|
||||
)
|
||||
discovery = _discovery(
|
||||
huggingface_discovery_module,
|
||||
huggingface_assets_module,
|
||||
tmp_path=tmp_path,
|
||||
user_directory=user_directory,
|
||||
api=api,
|
||||
)
|
||||
|
||||
source = discovery.discover(asset_path, sha256=sha256)
|
||||
|
||||
assert source.repo_id == "owner/model"
|
||||
assert source.revision == commit
|
||||
assert source.filename == "weights/manager-model.safetensors"
|
||||
assert discovery.registry.get(sha256) == source
|
||||
assert discovery.discover(asset_path, sha256=sha256) == source
|
||||
assert api.calls == [("owner/model", "main")]
|
||||
|
||||
|
||||
def test_discovers_model_inside_huggingface_git_checkout(
|
||||
huggingface_assets_module: Any,
|
||||
huggingface_discovery_module: Any,
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""An ordinary Hub Git clone should supply its repository, commit, and file path."""
|
||||
commit = "8" * 40
|
||||
checkout = tmp_path / "checkout"
|
||||
git_directory = checkout / ".git"
|
||||
git_directory.mkdir(parents=True)
|
||||
(git_directory / "config").write_text(
|
||||
'[remote "origin"]\n\turl = https://huggingface.co/git/model.git\n',
|
||||
encoding="utf-8",
|
||||
)
|
||||
(git_directory / "HEAD").write_text(f"{commit}\n", encoding="utf-8")
|
||||
asset_path = checkout / "weights" / "git.safetensors"
|
||||
asset_path.parent.mkdir()
|
||||
asset_path.write_bytes(b"git-model")
|
||||
sha256 = huggingface_assets_module.sha256_file(asset_path)
|
||||
api = FakeHubApi(
|
||||
{
|
||||
("git/model", commit): _hub_response(
|
||||
filename="weights/git.safetensors",
|
||||
sha256=sha256,
|
||||
size_bytes=asset_path.stat().st_size,
|
||||
commit=commit,
|
||||
)
|
||||
}
|
||||
)
|
||||
discovery = _discovery(
|
||||
huggingface_discovery_module,
|
||||
huggingface_assets_module,
|
||||
tmp_path=tmp_path,
|
||||
user_directory=tmp_path / "user",
|
||||
api=api,
|
||||
)
|
||||
|
||||
source = discovery.discover(asset_path, sha256=sha256)
|
||||
|
||||
assert source.repo_id == "git/model"
|
||||
assert source.revision == commit
|
||||
assert source.filename == "weights/git.safetensors"
|
||||
|
||||
|
||||
def test_discovers_browser_download_source_metadata(
|
||||
huggingface_assets_module: Any,
|
||||
huggingface_discovery_module: Any,
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""Browser WhereFrom metadata should remove the need for manual registration."""
|
||||
asset_path = tmp_path / "browser-model.safetensors"
|
||||
asset_path.write_bytes(b"browser-model")
|
||||
sha256 = huggingface_assets_module.sha256_file(asset_path)
|
||||
api = FakeHubApi(
|
||||
{
|
||||
("browser/source", "main"): _hub_response(
|
||||
filename="browser-model.safetensors",
|
||||
sha256=sha256,
|
||||
size_bytes=asset_path.stat().st_size,
|
||||
commit="c" * 40,
|
||||
)
|
||||
}
|
||||
)
|
||||
discovery = _discovery(
|
||||
huggingface_discovery_module,
|
||||
huggingface_assets_module,
|
||||
tmp_path=tmp_path,
|
||||
user_directory=tmp_path / "user",
|
||||
api=api,
|
||||
where_from_urls=(
|
||||
"https://huggingface.co/browser/source/resolve/main/browser-model.safetensors",
|
||||
),
|
||||
)
|
||||
|
||||
source = discovery.discover(asset_path, sha256=sha256)
|
||||
|
||||
assert source.repo_id == "browser/source"
|
||||
assert source.revision == "c" * 40
|
||||
|
||||
|
||||
def test_discovers_standard_huggingface_cache_path(
|
||||
huggingface_assets_module: Any,
|
||||
huggingface_discovery_module: Any,
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""A file or symlink into the standard Hub cache should carry its own source."""
|
||||
commit = "d" * 40
|
||||
asset_path = (
|
||||
tmp_path
|
||||
/ "hub"
|
||||
/ "models--cache--model"
|
||||
/ "snapshots"
|
||||
/ commit
|
||||
/ "weights"
|
||||
/ "cached.safetensors"
|
||||
)
|
||||
asset_path.parent.mkdir(parents=True)
|
||||
asset_path.write_bytes(b"cached-model")
|
||||
sha256 = huggingface_assets_module.sha256_file(asset_path)
|
||||
api = FakeHubApi(
|
||||
{
|
||||
("cache/model", commit): _hub_response(
|
||||
filename="weights/cached.safetensors",
|
||||
sha256=sha256,
|
||||
size_bytes=asset_path.stat().st_size,
|
||||
commit=commit,
|
||||
)
|
||||
}
|
||||
)
|
||||
discovery = _discovery(
|
||||
huggingface_discovery_module,
|
||||
huggingface_assets_module,
|
||||
tmp_path=tmp_path,
|
||||
user_directory=tmp_path / "user",
|
||||
api=api,
|
||||
)
|
||||
|
||||
source = discovery.discover(asset_path, sha256=sha256)
|
||||
|
||||
assert source.repo_id == "cache/model"
|
||||
assert source.filename == "weights/cached.safetensors"
|
||||
|
||||
|
||||
def test_discovers_hub_url_embedded_in_safetensors_metadata(
|
||||
huggingface_assets_module: Any,
|
||||
huggingface_discovery_module: Any,
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""Model metadata containing a Hub file URL should be useful but still verified."""
|
||||
header = json.dumps(
|
||||
{
|
||||
"__metadata__": {
|
||||
"modelspec.source": (
|
||||
"https://huggingface.co/embedded/model/resolve/main/"
|
||||
"weights/embedded.safetensors"
|
||||
)
|
||||
}
|
||||
}
|
||||
).encode("utf-8")
|
||||
asset_path = tmp_path / "embedded.safetensors"
|
||||
asset_path.write_bytes(struct.pack("<Q", len(header)) + header + b"model-bytes")
|
||||
sha256 = huggingface_assets_module.sha256_file(asset_path)
|
||||
api = FakeHubApi(
|
||||
{
|
||||
("embedded/model", "main"): _hub_response(
|
||||
filename="weights/embedded.safetensors",
|
||||
sha256=sha256,
|
||||
size_bytes=asset_path.stat().st_size,
|
||||
commit="e" * 40,
|
||||
)
|
||||
}
|
||||
)
|
||||
discovery = _discovery(
|
||||
huggingface_discovery_module,
|
||||
huggingface_assets_module,
|
||||
tmp_path=tmp_path,
|
||||
user_directory=tmp_path / "user",
|
||||
api=api,
|
||||
)
|
||||
|
||||
source = discovery.discover(asset_path, sha256=sha256)
|
||||
|
||||
assert source.repo_id == "embedded/model"
|
||||
assert source.filename == "weights/embedded.safetensors"
|
||||
|
||||
|
||||
def test_rejects_unverified_hint_and_negative_caches_result(
|
||||
huggingface_assets_module: Any,
|
||||
huggingface_discovery_module: Any,
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""Install hints must never bypass exact content verification or repeat per prompt."""
|
||||
asset_path = tmp_path / "mismatch.safetensors"
|
||||
asset_path.write_bytes(b"local-content")
|
||||
sha256 = huggingface_assets_module.sha256_file(asset_path)
|
||||
api = FakeHubApi(
|
||||
{
|
||||
("wrong/model", "main"): _hub_response(
|
||||
filename="mismatch.safetensors",
|
||||
sha256="f" * 64,
|
||||
size_bytes=asset_path.stat().st_size,
|
||||
commit="f" * 40,
|
||||
)
|
||||
}
|
||||
)
|
||||
discovery = _discovery(
|
||||
huggingface_discovery_module,
|
||||
huggingface_assets_module,
|
||||
tmp_path=tmp_path,
|
||||
user_directory=tmp_path / "user",
|
||||
api=api,
|
||||
where_from_urls=(
|
||||
"https://huggingface.co/wrong/model/resolve/main/mismatch.safetensors",
|
||||
),
|
||||
)
|
||||
|
||||
assert discovery.discover(asset_path, sha256=sha256) is None
|
||||
assert discovery.discover(asset_path, sha256=sha256) is None
|
||||
assert discovery.registry.get(sha256) is None
|
||||
assert api.calls == [("wrong/model", "main")]
|
||||
|
||||
|
||||
def test_retries_unresolved_asset_after_manager_catalog_changes(
|
||||
huggingface_assets_module: Any,
|
||||
huggingface_discovery_module: Any,
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""Installing a model while ComfyUI is open should invalidate negative evidence."""
|
||||
user_directory = tmp_path / "user"
|
||||
asset_path = tmp_path / "models" / "late.safetensors"
|
||||
asset_path.parent.mkdir()
|
||||
asset_path.write_bytes(b"late-manager-install")
|
||||
sha256 = huggingface_assets_module.sha256_file(asset_path)
|
||||
api = FakeHubApi(
|
||||
{
|
||||
("late/model", "main"): _hub_response(
|
||||
filename="late.safetensors",
|
||||
sha256=sha256,
|
||||
size_bytes=asset_path.stat().st_size,
|
||||
commit="7" * 40,
|
||||
)
|
||||
}
|
||||
)
|
||||
discovery = _discovery(
|
||||
huggingface_discovery_module,
|
||||
huggingface_assets_module,
|
||||
tmp_path=tmp_path,
|
||||
user_directory=user_directory,
|
||||
api=api,
|
||||
)
|
||||
|
||||
assert discovery.discover(asset_path, sha256=sha256) is None
|
||||
|
||||
catalog = user_directory / "__manager" / "cache" / "new_model-list.json"
|
||||
catalog.parent.mkdir(parents=True)
|
||||
catalog.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"models": [
|
||||
{
|
||||
"filename": "late.safetensors",
|
||||
"url": (
|
||||
"https://huggingface.co/late/model/resolve/main/"
|
||||
"late.safetensors"
|
||||
),
|
||||
}
|
||||
]
|
||||
}
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
source = discovery.discover(asset_path, sha256=sha256)
|
||||
|
||||
assert source.repo_id == "late/model"
|
||||
assert source.revision == "7" * 40
|
||||
assert api.calls == [("late/model", "main")]
|
||||
@@ -294,6 +294,117 @@ def test_vast_sync_materializes_registered_huggingface_asset_before_upload(
|
||||
]
|
||||
|
||||
|
||||
def test_vast_sync_automatically_discovers_huggingface_asset(
|
||||
settings_module: Any,
|
||||
sync_engine_module: Any,
|
||||
huggingface_assets_module: Any,
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""Vast sync should discover provenance without a user-run registration command."""
|
||||
asset_path = tmp_path / "automatic.safetensors"
|
||||
asset_path.write_bytes(b"automatically-discovered")
|
||||
source = huggingface_assets_module.HuggingFaceAssetSource(
|
||||
repo_id="owner/automatic",
|
||||
revision="9" * 40,
|
||||
filename="automatic.safetensors",
|
||||
sha256=huggingface_assets_module.sha256_file(asset_path),
|
||||
size_bytes=asset_path.stat().st_size,
|
||||
)
|
||||
statuses: list[str] = []
|
||||
discoveries: list[tuple[Path, str]] = []
|
||||
|
||||
class DiscoveringVolume:
|
||||
"""Accept direct Hub materialization and reject local upload use."""
|
||||
|
||||
def exists(self, remote_path: str) -> bool:
|
||||
"""Report no preexisting remote asset."""
|
||||
del remote_path
|
||||
return False
|
||||
|
||||
def put_file(self, local_path: Path, remote_path: str) -> None:
|
||||
"""Fail if automatic discovery unexpectedly falls back to upload."""
|
||||
raise AssertionError(f"Unexpected upload of {local_path} to {remote_path}")
|
||||
|
||||
def put_bytes(self, payload: bytes, remote_path: str) -> None:
|
||||
"""Accept protocol metadata writes."""
|
||||
del payload, remote_path
|
||||
|
||||
def materialize_huggingface_file(
|
||||
self,
|
||||
registered_source: Any,
|
||||
remote_path: str,
|
||||
*,
|
||||
token: str | None,
|
||||
) -> bool:
|
||||
"""Accept the automatically discovered immutable source."""
|
||||
del remote_path, token
|
||||
assert registered_source == source
|
||||
return True
|
||||
|
||||
class FakeDiscovery:
|
||||
"""Return an automatically verified source for the local asset."""
|
||||
|
||||
def discover(self, local_path: Path, *, sha256: str) -> Any:
|
||||
"""Record discovery and return the expected source."""
|
||||
discoveries.append((local_path, sha256))
|
||||
return source
|
||||
|
||||
settings = settings_module.ModalSyncSettings(
|
||||
app_name="app",
|
||||
auto_deploy=True,
|
||||
allow_ephemeral_fallback=False,
|
||||
enable_memory_snapshot=True,
|
||||
enable_gpu_memory_snapshot=False,
|
||||
execution_mode="vast",
|
||||
sync_custom_nodes=False,
|
||||
volume_name="volume",
|
||||
route_path="/modal/queue_prompt",
|
||||
marker_property="is_modal_remote",
|
||||
local_storage_root=tmp_path / "storage",
|
||||
remote_storage_root="/storage",
|
||||
custom_nodes_archive_name="custom_nodes_bundle.zip",
|
||||
comfyui_root=None,
|
||||
custom_nodes_dir=None,
|
||||
)
|
||||
engine = sync_engine_module.ModalAssetSyncEngine(
|
||||
volume=DiscoveringVolume(),
|
||||
settings=settings,
|
||||
huggingface_asset_registry=huggingface_assets_module.HuggingFaceAssetRegistry(
|
||||
tmp_path / "empty-registry.json"
|
||||
),
|
||||
huggingface_asset_discovery=FakeDiscovery(),
|
||||
)
|
||||
|
||||
engine.sync_file(
|
||||
asset_path,
|
||||
status_callback=lambda message, current, total: statuses.append(message),
|
||||
)
|
||||
|
||||
assert discoveries == [(asset_path, source.sha256)]
|
||||
assert statuses == [
|
||||
"Identifying Hugging Face source for Vast.ai: automatic.safetensors",
|
||||
"Downloading asset from Hugging Face on Vast.ai: automatic.safetensors",
|
||||
]
|
||||
|
||||
|
||||
def test_resolve_model_path_preserves_huggingface_cache_symlink(
|
||||
sync_engine_module: Any,
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""Model resolution must retain a cache symlink long enough for provenance discovery."""
|
||||
cached_file = tmp_path / "cache" / "model.safetensors"
|
||||
cached_file.parent.mkdir()
|
||||
cached_file.write_bytes(b"cached")
|
||||
linked_file = tmp_path / "models" / "model.safetensors"
|
||||
linked_file.parent.mkdir()
|
||||
linked_file.symlink_to(cached_file)
|
||||
|
||||
resolved = sync_engine_module.resolve_model_path(str(linked_file))
|
||||
|
||||
assert resolved == linked_file
|
||||
assert resolved.resolve() == cached_file
|
||||
|
||||
|
||||
def test_vast_sync_falls_back_to_upload_after_huggingface_failure(
|
||||
settings_module: Any,
|
||||
sync_engine_module: Any,
|
||||
|
||||
+12
-3
@@ -20,6 +20,7 @@ if __package__:
|
||||
GpuCapability,
|
||||
)
|
||||
from .huggingface_assets import HuggingFaceAssetRegistry
|
||||
from .huggingface_discovery import HuggingFaceAssetDiscovery
|
||||
from .runtime_environment import build_remote_runtime_identity
|
||||
from .settings import ModalSyncSettings, discover_comfyui_user_directory
|
||||
from .sync_engine import ModalAssetSyncEngine
|
||||
@@ -43,6 +44,7 @@ else: # pragma: no cover - direct debugging imports.
|
||||
GpuCapability,
|
||||
)
|
||||
from huggingface_assets import HuggingFaceAssetRegistry
|
||||
from huggingface_discovery import HuggingFaceAssetDiscovery
|
||||
from runtime_environment import build_remote_runtime_identity
|
||||
from settings import ModalSyncSettings, discover_comfyui_user_directory
|
||||
from sync_engine import ModalAssetSyncEngine
|
||||
@@ -136,6 +138,14 @@ class VastService:
|
||||
self.runtime_configuration = runtime_configuration
|
||||
self.registry = registry
|
||||
self.identity_file = identity_file
|
||||
self.huggingface_asset_registry = (
|
||||
HuggingFaceAssetRegistry.for_user_directory(self.user_directory)
|
||||
)
|
||||
self.huggingface_asset_discovery = HuggingFaceAssetDiscovery(
|
||||
registry=self.huggingface_asset_registry,
|
||||
user_directory=self.user_directory,
|
||||
comfyui_root=getattr(settings, "comfyui_root", None),
|
||||
)
|
||||
self.lease_manager = VastLeaseManager(
|
||||
api_client=api_client,
|
||||
registry=registry,
|
||||
@@ -497,9 +507,8 @@ class VastService:
|
||||
storage_root=self.runtime_configuration.remote_storage_root,
|
||||
),
|
||||
settings=vast_settings,
|
||||
huggingface_asset_registry=(
|
||||
HuggingFaceAssetRegistry.for_user_directory(self.user_directory)
|
||||
),
|
||||
huggingface_asset_registry=self.huggingface_asset_registry,
|
||||
huggingface_asset_discovery=self.huggingface_asset_discovery,
|
||||
)
|
||||
|
||||
def executor(self) -> VastExecutorClient:
|
||||
|
||||
Reference in New Issue
Block a user