Automate Hugging Face asset discovery

This commit is contained in:
Ken Simpson
2026-08-24 12:49:02 -07:00
parent 96f594a6ec
commit f3dc4c2fad
10 changed files with 1236 additions and 102 deletions
+1
View File
@@ -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
+4 -9
View File
@@ -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.
+1
View File
@@ -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.
+596
View File
@@ -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",
]
+20 -64
View File
@@ -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
View File
@@ -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
+6
View File
@@ -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."""
+414
View File
@@ -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")]
+111
View File
@@ -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
View File
@@ -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: