horde-engine 3.6.2__tar.gz → 3.7.1__tar.gz

{horde_engine-3.6.2 → horde_engine-3.7.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: horde-engine
-Version: 3.6.2
+Version: 3.7.1
 Summary: A wrapper around ComfyUI to allow use by AI Horde.
 Project-URL: Homepage, https://github.com/Haidra-Org/hordelib
 Project-URL: Bug Tracker, https://github.com/Haidra-Org/hordelib/issues
@@ -690,7 +690,7 @@ Requires-Dist: filelock
 Requires-Dist: ftfy
 Requires-Dist: fuzzywuzzy>=0.18.0
 Requires-Dist: gitpython>=3.1.45
-Requires-Dist: horde-model-reference<6.0.0,>=5.1.1
+Requires-Dist: horde-model-reference<7.0.0,>=6.0.0
 Requires-Dist: horde-sdk>=0.20.1
 Requires-Dist: kornia>=0.7.1
 Requires-Dist: logfire[system-metrics]==4.14.2

{horde_engine-3.6.2 → horde_engine-3.7.1}/hordelib/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '3.6.2'
-__version_tuple__ = version_tuple = (3, 6, 2)
+__version__ = version = '3.7.1'
+__version_tuple__ = version_tuple = (3, 7, 1)
 
 __commit_id__ = commit_id = None

{horde_engine-3.6.2 → horde_engine-3.7.1}/hordelib/model_manager/base.py

@@ -2,6 +2,7 @@ import os
 import threading
 from abc import ABC
 from collections.abc import Callable, Iterable
+from functools import lru_cache
 from pathlib import Path
 from typing import Any, cast
 from urllib import parse
@@ -26,6 +27,59 @@ from hordelib.config_path import get_hordelib_path
 from hordelib.consts import CIVITAI_API_PATH
 from hordelib.settings import UserSettings
 
+_ANON_API_KEY = "0000000000"
+
+
+def _resolve_horde_api_key() -> str | None:
+    """Return the worker's horde API key from the environment, or None for anonymous/standalone use.
+
+    The key gates the content-addressed R2 mirror in the download engine. The anonymous key cannot be trusted,
+    and a standalone hordelib user without a key simply downloads from each model's origin host. The worker
+    exposes its configured key as ``AIHORDE_API_KEY`` (inherited by the download subprocess).
+    """
+    key = os.environ.get("AIHORDE_API_KEY") or os.environ.get("AI_HORDE_API_KEY")
+    if not key or key == _ANON_API_KEY:
+        return None
+    return key
+
+
+@lru_cache(maxsize=1)
+def _resolve_r2_gateway_url() -> str | None:
+    """Return the gated-R2 gateway base URL from horde_model_reference settings (process-static, cached).
+
+    None when no gateway is configured, in which case the engine downloads only from origin hosts. An older
+    horde_model_reference without the R2 settings block degrades to the same origin-only behaviour rather than
+    breaking the download path, so the worker keeps working ahead of a coordinated hmr release.
+
+    The gateway URL is owned by the official deployment (the horde_model_reference default); it is deliberately
+    not a per-worker config knob. A sophisticated operator can still override it via the
+    ``HORDE_MODEL_REFERENCE_R2__GATEWAY_URL`` environment variable, in which case the safety guard below applies:
+    because the worker's API key is sent to the gateway, a plaintext (non-``https``) gateway is refused outright
+    (the mirror is disabled and origins are used) so the key cannot travel in the clear.
+    """
+    try:
+        from horde_model_reference import HordeModelReferenceSettings
+        from horde_model_reference.download_engine import gateway_accepts_key
+
+        gateway_url = HordeModelReferenceSettings().r2.gateway_url
+    except (AttributeError, ImportError):
+        # Older horde_model_reference without the R2 settings block or the key-safety guard: origin-only.
+        return None
+
+    if not gateway_url:
+        return None
+    if not gateway_accepts_key(gateway_url):
+        logger.warning(
+            "Configured R2 gateway is not https; refusing to send the API key over plaintext and disabling the "
+            "mirror (downloads use origin hosts): gateway={}",
+            gateway_url,
+        )
+        return None
+    # Normal path logs at DEBUG only: the gateway is a deployment-owned technical detail, not something a typical
+    # operator configures or needs to reason about, so it stays out of the default console.
+    logger.debug("Gated R2 mirror enabled; the API key will be sent to gateway={}", parse.urlparse(gateway_url).netloc)
+    return gateway_url
+
 
 class BaseModelManager[RecordT: GenericModelRecord | dict[str, Any]](ABC):
     """Abstract base for the per-category model managers.
@@ -477,6 +531,7 @@ class BaseModelManager[RecordT: GenericModelRecord | dict[str, Any]](ABC):
         model_name: str,
         *,
         callback: Callable[[int, int], None] | None = None,
+        connections: int = 1,
     ) -> bool | None:
         """Ensure *model_name*'s declared files are present, fetching any that are missing.
 
@@ -487,6 +542,9 @@ class BaseModelManager[RecordT: GenericModelRecord | dict[str, Any]](ABC):
         Args:
             model_name: The reference key of the model to download.
             callback: Optional ``(downloaded_bytes, total_bytes)`` progress callback, per file.
+            connections: Maximum concurrent connections per large file, forwarded to the engine. 1 keeps the
+                single-stream path; a higher value lets the engine segment a large file across that many
+                ranged connections to raise single-file throughput (the caller's configured value).
 
         Returns:
             True if the model is present and validated, False on download/validation failure, and None
@@ -520,6 +578,9 @@ class BaseModelManager[RecordT: GenericModelRecord | dict[str, Any]](ABC):
             self._weights_root,
             progress_callback=callback,
             auth_query_token=self._civitai_token_for(record),
+            gateway_base_url=_resolve_r2_gateway_url(),
+            apikey=_resolve_horde_api_key(),
+            connections=connections,
         )
         if not download_succeeded:
             return False

{horde_engine-3.6.2 → horde_engine-3.7.1}/hordelib/preload.py

@@ -27,11 +27,37 @@ Two optimisations keep this cheap after the first time:
 import contextlib
 import os
 import threading
-from collections.abc import Iterator
+from collections.abc import Iterable, Iterator
 from pathlib import Path
+from typing import Protocol
 
 from loguru import logger
 
+
+class _AnnotatorFileLike(Protocol):
+    """The shape the prefetch reads from a ``horde_model_reference.annotator_catalog.AnnotatorFile``.
+
+    Described structurally so this module type-checks even against an older installed ``horde_model_reference``
+    that does not yet ship the catalog (the concrete type is imported lazily and may be absent; see
+    :func:`_prefetch_annotator_files`).
+    """
+
+    repo: str
+    subfolder: str
+    filename: str
+    sha256: str | None
+    preprocessors: tuple[str, ...]
+
+    @property
+    def relative_path(self) -> str:
+        """The file's path relative to the annotator checkpoints directory (``<repo>/<subfolder>/<filename>``)."""
+        ...
+
+    @property
+    def origin_url(self) -> str:
+        """The HuggingFace origin URL the file is fetched from when the gated mirror does not serve it."""
+        ...
+
 _preload_mutex = threading.Lock()
 _preload_completed = False
 
@@ -52,6 +78,83 @@ _PRELOAD_MARKER_NAME = ".hordelib_preload_complete"
 _ONNXRUNTIME_GATED_PREPROCESSORS = frozenset({"OpenposePreprocessor"})
 
 
+def _annotator_files_to_prefetch(all_files: Iterable[_AnnotatorFileLike]) -> list[_AnnotatorFileLike]:
+    """Filter the annotator catalog to the files whose preprocessor will actually run on this install.
+
+    On a lean base install (no ``controlnet`` extra) the onnxruntime-gated preprocessors never register and the
+    preload skips them (see :func:`_run_preload`), so prefetching their large weights would be wasted. A file is
+    kept when at least one preprocessor that loads it is not gated away. Best-effort: if the feature check itself
+    cannot be evaluated, everything is kept (the worst case is a redundant download, never a missing one).
+    """
+    try:
+        from hordelib.feature_impact import FEATURE_KIND
+        from hordelib.feature_requirements import feature_available
+
+        if feature_available(FEATURE_KIND.controlnet):
+            return list(all_files)
+    except Exception as e:
+        logger.debug("Could not evaluate controlnet feature for annotator prefetch; keeping all: error={}", e)
+        return list(all_files)
+
+    return [
+        entry
+        for entry in all_files
+        if any(preprocessor not in _ONNXRUNTIME_GATED_PREPROCESSORS for preprocessor in entry.preprocessors)
+    ]
+
+
+def _prefetch_annotator_files(ckpts_dir: Path) -> None:
+    """Pre-place known annotator checkpoints via the unified gated-mirror-then-origin download engine.
+
+    For each catalog file not already on disk, fetch it to the exact path ``custom_hf_download`` expects, so the
+    detector finds it present and skips its own HuggingFace download. The gated R2 mirror (when configured and
+    the file's hash is known) is tried first and falls back to the HuggingFace origin on any failure, so this
+    only ever adds a faster, free path; it never blocks a download. Best-effort throughout: any failure simply
+    leaves the file for the detector to fetch itself, exactly as before this hook existed.
+    """
+    try:
+        from horde_model_reference import annotator_catalog, download_engine
+    except Exception as e:
+        # An older installed horde_model_reference without the catalog (or its record-free fetch helper): the
+        # detectors download from HuggingFace themselves, exactly as before. Degrade quietly.
+        logger.debug("Annotator catalog unavailable (older horde_model_reference?); skipping prefetch: {}", e)
+        return
+
+    gateway_base_url: str | None = None
+    apikey: str | None = None
+    try:
+        from hordelib.model_manager.base import _resolve_horde_api_key, _resolve_r2_gateway_url
+
+        gateway_base_url = _resolve_r2_gateway_url()
+        apikey = _resolve_horde_api_key()
+    except Exception as e:
+        logger.debug("Could not resolve gated-mirror settings; annotator prefetch uses origin only: error={}", e)
+
+    for entry in _annotator_files_to_prefetch(annotator_catalog.ANNOTATOR_FILES):
+        destination = ckpts_dir / Path(entry.relative_path)
+        if destination.is_file():
+            continue
+        try:
+            outcome = download_engine.download_addressed_file(
+                entry.origin_url,
+                destination,
+                sha256=entry.sha256,
+                gateway_base_url=gateway_base_url,
+                apikey=apikey,
+            )
+            if not outcome.success:
+                logger.info(
+                    "Annotator prefetch did not complete; the detector will fetch it at first use: file={}",
+                    entry.filename,
+                )
+        except Exception as e:
+            logger.warning(
+                "Annotator prefetch errored; the detector will fetch it at first use: file={} error={}",
+                entry.filename,
+                e,
+            )
+
+
 def _midas_already_cached() -> bool:
     """Return whether the transformers-based MiDaS annotator is fully in the HF cache.
 
@@ -153,6 +256,66 @@ def controlnet_annotators_present() -> bool | None:
         return None
 
 
+def _annotator_file_resolvable(entry: _AnnotatorFileLike, ckpts_dir: Path | None) -> bool:
+    """Whether one annotator file is already where the engine will find it at first use (no download needed).
+
+    Existence-based, unlike the preload marker: a detector loads a file it finds on disk and skips its own
+    HuggingFace fetch, so a file present in the checkpoints dir (the flat ``<repo>/<subfolder>/<filename>``
+    layout ``custom_hf_download`` writes) -- or already in the HuggingFace hub cache, which a fetch resolves
+    without touching the network -- will not be re-downloaded. The cheap on-disk path is checked first.
+    """
+    if ckpts_dir is not None and (ckpts_dir / Path(entry.relative_path)).is_file():
+        return True
+    try:
+        from huggingface_hub import try_to_load_from_cache
+
+        in_repo = "/".join(part for part in (entry.subfolder, entry.filename) if part)
+        if isinstance(try_to_load_from_cache(repo_id=entry.repo, filename=in_repo), str):
+            return True
+        if ckpts_dir is not None and isinstance(
+            try_to_load_from_cache(repo_id=entry.repo, filename=in_repo, cache_dir=str(ckpts_dir)),
+            str,
+        ):
+            return True
+    except Exception as e:  # pragma: no cover - a hub-cache probe failure just means "not via the cache"
+        logger.debug("Hub-cache annotator lookup failed (treating this file as absent): error={}", e)
+    return False
+
+
+def annotators_resolvable(control_types: Iterable[str]) -> bool | None:
+    """Whether the controlnet annotators for *control_types* are on disk or cached (existence, not the marker).
+
+    The authoritative answer to "will a controlnet job for these control types need a slow annotator download
+    before it can run", as distinct from :func:`controlnet_annotators_present`, which reads the preload
+    *marker* (a stricter, pin-keyed fast-path that is ``False`` whenever a full verify has not run for the
+    current pin, even when every file is already present on disk). A surface deciding whether to prompt a
+    pre-download should use this; the marker stays only the preload's own skip optimisation.
+
+    Weightless control types (``canny``, ``scribble``) and types unknown to the catalog need no files and are
+    vacuously resolvable. Returns ``None`` only when the file catalog or the checkpoint directory cannot be
+    determined, so a caller treats the unknown as "do not claim missing" rather than nagging spuriously.
+    """
+    try:
+        from horde_model_reference import annotator_catalog
+    except Exception as e:
+        logger.debug("Annotator catalog unavailable; cannot resolve annotator presence: error={}", e)
+        return None
+    # Prefer the catalog's own control-type selector; fall back to filtering the file list directly so an
+    # older horde_model_reference that ships ANNOTATOR_FILES but not the helper still resolves correctly.
+    selector = getattr(annotator_catalog, "annotators_for_control_types", None)
+    if selector is not None:
+        needed = selector(control_types)
+    else:
+        wanted = set(control_types)
+        needed = tuple(entry for entry in annotator_catalog.ANNOTATOR_FILES if wanted.intersection(entry.control_types))
+    if not needed:
+        return True
+    ckpts_dir = _annotator_ckpts_dir()
+    if ckpts_dir is None:
+        return None
+    return all(_annotator_file_resolvable(entry, ckpts_dir) for entry in needed)
+
+
 def _record_annotators_verified(ref: str) -> None:
     """Persist that annotators for ``ref`` are downloaded and verified on this machine.
 
@@ -284,6 +447,13 @@ def download_all_controlnet_annotators() -> bool:
             _preload_completed = True
             return True
 
+        # Pre-place the annotator checkpoints through the unified (gated R2 -> HuggingFace) download engine
+        # *before* the detectors run, so custom_hf_download finds them present and skips its own HF fetch. Runs
+        # with the Hub online (the offline window is inside _run_preload). Best-effort; never blocks the preload.
+        ckpts_dir = _annotator_ckpts_dir()
+        if ckpts_dir is not None:
+            _prefetch_annotator_files(ckpts_dir)
+
         force_offline = _midas_already_cached()
         if force_offline:
             logger.debug(

{horde_engine-3.6.2 → horde_engine-3.7.1}/pyproject.toml

@@ -52,7 +52,7 @@ dependencies = [
     "ftfy",  # comfyui_controlnet_aux node requirement (previously satisfied transitively via open-clip-torch)
     "fuzzywuzzy>=0.18.0",
     "gitpython>=3.1.45",
-    "horde_model_reference>=5.1.1,<6.0.0",
+    "horde_model_reference>=6.0.0,<7.0.0",
     "horde-sdk>=0.20.1",
     "kornia>=0.7.1",
     "logfire[system-metrics]==4.14.2",