haybale-models 0.3.1__py3-none-any.whl

haybale_models/OVERVIEW.md

@@ -0,0 +1,30 @@
+# Models
+
+Model references, local model store, downloads and graph needs-list
+
+## Nodes
+### Models
+- **Resolve Model** — Gives the file of a model on this machine, with its sha256 and size
+
+## Types
+- **Model** — A model, named by a catalog entry or by a file (url and sha256). Resolve Model turns it into a path on disk.
+
+## Widgets
+- **ModelRefWidget** — Pick a model
+
+## Settings
+- **Models** — Model download behaviour
+
+## Farmhands
+- **Fetch models** — Queue downloads: everything an open graph needs, or catalog entries by key (admin).
+- **List models** — Catalog entries with per-variant status, indexed local files and download jobs.
+- **Models a graph needs** — The models an open graph references and does not have locally.
+
+## States
+- **Models** — Model volumes, file index, download queue and the pool of loaded models
+
+## Panels
+- **Models** — Model download behaviour: how many files download in parallel.
+
+## Editors
+- **Models** — Model files on this machine and where they come from

haybale_models/QUICKREF.md

@@ -0,0 +1,27 @@
+# haybale-models — component index (v0.3.1)
+
+## node
+- `haybale-models:node:ResolveModelNode` — Resolve Model — Gives the file of a model on this machine, with its sha256 and size  _tags: model, resolve, file, path, weights, catalog, sha256_
+
+## type
+- `haybale-models:type:MODEL_REF` — Model — A model, named by a catalog entry or by a file (url and sha256). Resolve Model turns it into a path on disk.
+
+## widget
+- `haybale-models:widget:ModelRefWidget` — ModelRefWidget — Pick a model
+
+## setting
+- `haybale-models:setting:ModelsSettings` — Models — Model download behaviour
+
+## farmhand
+- `haybale-models:farmhand:fetch` — Fetch models — Queue downloads: everything an open graph needs, or catalog entries by key (admin).
+- `haybale-models:farmhand:list` — List models — Catalog entries with per-variant status, indexed local files and download jobs.
+- `haybale-models:farmhand:needs` — Models a graph needs — The models an open graph references and does not have locally.
+
+## state
+- `haybale-models:state:ModelsState` — Models — Model volumes, file index, download queue and the pool of loaded models
+
+## panel
+- `haybale-models:panel:ModelsSettingsPanel` — Models — Model download behaviour: how many files download in parallel.
+
+## editor
+- `haybale-models:editor:ModelsEditor` — Models — Model files on this machine and where they come from

haybale_models/__init__.py

@@ -0,0 +1,44 @@
+"""Haybale-Models: model references, local model store, downloads and needs-list.
+
+Component folders (registered in ``register_components``):
+  - ``state``: ``ModelsState`` with the two model volumes.
+  - ``settings``, ``types``, ``widgets``, ``nodes``, ``editors``, ``farmhands``, ``panels``.
+
+Never imports ``torch``: variant choice comes from the consuming node's settings.
+"""
+
+from pathlib import Path
+
+from haywire.core.farmhand import FarmhandRegistry
+from haywire.core.library.base import BaseLibrary
+from haywire.core.library.decorator import library
+from haywire.core.node.registry import NodeRegistry
+from haywire.core.settings.registry import SettingsRegistry
+from haywire.core.state import LibraryStateRegistry
+from haywire.core.types.registry import TypeRegistry
+from haywire.ui.editor.registry import EditorTypeRegistry
+from haywire.ui.panel.registry import PanelRegistry
+from haywire.ui.widget.registry import WidgetRegistry
+
+
+@library(
+    file_watcher=True,
+)
+class Library(BaseLibrary):
+    """Models library — model references, store and downloads."""
+
+    def register_components(self):
+        """Register the component folders with their registries."""
+        base_path = Path(__file__).parent
+
+        self.add_folder_to_registry(folder_path=str(base_path / "state"), registry_cls=LibraryStateRegistry)
+        self.add_folder_to_registry(folder_path=str(base_path / "settings"), registry_cls=SettingsRegistry)
+        self.add_folder_to_registry(folder_path=str(base_path / "types"), registry_cls=TypeRegistry)
+        self.add_folder_to_registry(folder_path=str(base_path / "widgets"), registry_cls=WidgetRegistry)
+        self.add_folder_to_registry(folder_path=str(base_path / "nodes"), registry_cls=NodeRegistry)
+        self.add_folder_to_registry(folder_path=str(base_path / "editors"), registry_cls=EditorTypeRegistry)
+        self.add_folder_to_registry(folder_path=str(base_path / "farmhands"), registry_cls=FarmhandRegistry)
+        self.add_folder_to_registry(folder_path=str(base_path / "panels"), registry_cls=PanelRegistry)
+
+    def validate(self) -> bool:
+        return True

haybale_models/consumer.py

@@ -0,0 +1,115 @@
+"""Helpers for node authors: resolve a MODEL_REF port and hold its loaded model per run."""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Hashable
+from typing import TYPE_CHECKING, TypeVar
+
+from haywire.core.di.context import get_library_state_container
+from haywire.core.node.node_warning import NodeWarning
+
+from haybale_models.matcher import Preference
+from haybale_models.resolver import ModelMissing, ResolvedModel, resolve
+from haybale_models.state.models_state import ModelsState
+from haybale_models.value.model_ref import ModelRef
+
+if TYPE_CHECKING:
+    from haywire.core.node.base import BaseNode
+
+__all__ = [
+    "KIND",
+    "MISSING_WARNING",
+    "ModelMissing",
+    "acquire_port",
+    "models_state",
+    "preference_for",
+    "release_node",
+    "resolve_port",
+]
+
+T = TypeVar("T")
+
+KIND = "model"
+MISSING_WARNING = NodeWarning(message="Model missing — Show in Models", source_version=None, kind=KIND)
+
+
+def models_state() -> ModelsState:
+    """The app-wide :class:`ModelsState`.
+
+    Raises:
+        RuntimeError: If haybale-models is not enabled.
+    """
+    state = get_library_state_container().get(ModelsState)
+    if state is None:
+        raise RuntimeError("haybale-models is not enabled")
+    return state
+
+
+def preference_for(node: BaseNode, port_id: str) -> Preference:
+    """The variant preference for *port_id*.
+
+    The node's ``model_preference(port_id)`` if it defines one, else the port's ``prefer``
+    property, else ``[{}]``.
+    """
+    hook = getattr(node, "model_preference", None)
+    if callable(hook):
+        return hook(port_id)  # type: ignore[no-any-return]
+    prefer = node.ports[port_id].widget_config.get("properties", {}).get("prefer")
+    return list(prefer) if prefer else [{}]
+
+
+def resolve_port(node: BaseNode, port_id: str) -> ResolvedModel:
+    """Resolve the model on *port_id* and set the node's ``"model"`` warnings.
+
+    All ``MODEL_REF`` ports of one node share a single ``"model"`` warning list, so with
+    several ports the last call wins. Key the list per port when a two-port consumer lands.
+
+    Raises:
+        ModelMissing: If no local file matches the reference.
+    """
+    state = models_state()
+    value = node.ports[port_id].get_value()
+    ref = value if isinstance(value, ModelRef) else ModelRef.from_dict(value or {})
+    try:
+        resolved = resolve(
+            ref,
+            preference_for(node, port_id),
+            catalog=state.catalog,
+            index=state.index,
+            sources=state.sources,
+        )
+    except ModelMissing:
+        node.wrapper.replace_warnings(KIND, [] if ref.is_empty else [MISSING_WARNING])
+        raise
+    node.wrapper.replace_warnings(KIND, _warnings(resolved))
+    return resolved
+
+
+def _warnings(resolved: ResolvedModel) -> list[NodeWarning]:
+    messages: list[str] = []
+    if resolved.better_unfetched is not None:
+        variant = resolved.better_unfetched.variant_name()
+        messages.append(f"Preferred variant {variant} is not downloaded — Show in Models")
+    if resolved.ref.is_file and not resolved.ref.sha256:
+        messages.append("Not pinned — use the pin on the card")
+    return [NodeWarning(message=m, source_version=None, kind=KIND) for m in messages]
+
+
+def acquire_port(
+    node: BaseNode,
+    port_id: str,
+    *,
+    variant_key: Hashable,
+    load: Callable[[ResolvedModel], T],
+    release: Callable[[T], None],
+) -> T:
+    """Resolve *port_id* and hold its loaded model for *node*; nodes share one load per key."""
+    resolved = resolve_port(node, port_id)
+    return models_state().pool.acquire(
+        (resolved.key, variant_key), node, load=lambda: load(resolved), release=release
+    )
+
+
+def release_node(node: BaseNode) -> None:
+    """Drop *node*'s holds. Call from ``on_shutdown`` and, as a backstop, ``on_teardown``."""
+    models_state().pool.release(node)

haybale_models/describe.py

@@ -0,0 +1,111 @@
+"""What a catalog entry is, in words: a one-line summary, a plain-text description, and search text.
+
+Used wherever a person chooses between models (the model picker, the Models editor),
+so both describe an entry the same way.
+"""
+
+from __future__ import annotations
+
+import re
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from haywire.core.catalog import CatalogEntry
+
+from haybale_models.filters import COMMERCIAL, NON_COMMERCIAL
+
+_NOT_NON_COMMERCIAL = re.compile(r"(?<!non-)commercial")
+
+__all__ = ["entry_description", "entry_search_text", "entry_summary", "search_matches"]
+
+#: Tags shown in a summary line; the rest are counted (``+2``).
+SUMMARY_TAGS = 4
+
+#: Longest description shown in a tip, in characters; longer text ends in ``…``.
+DESCRIPTION_LIMIT = 400
+
+_LINK = re.compile(r"\[([^\]]+)\]\([^)]*\)")
+_MARKUP = re.compile(r"[*_`#>]+")
+
+
+def entry_summary(entry: CatalogEntry, *, scale: bool = True) -> str:
+    """Return what tells *entry* apart at a glance, such as ``4x · ESRGAN · anime, photo +2``.
+
+    Scale, architecture and the first tags, each left out when the entry has none.
+    ``scale=False`` leaves the scale out, for a place that already shows it.
+    """
+    parts: list[str] = []
+    if scale and entry.scale:
+        parts.append(f"{entry.scale}x")
+    if entry.architecture:
+        parts.append(entry.architecture)
+    if entry.tags:
+        shown = ", ".join(entry.tags[:SUMMARY_TAGS])
+        rest = len(entry.tags) - SUMMARY_TAGS
+        parts.append(f"{shown} +{rest}" if rest > 0 else shown)
+    return " · ".join(parts)
+
+
+def entry_description(entry: CatalogEntry) -> str:
+    """Return *entry*'s description as plain text for a tip, with purpose, author and licence below.
+
+    Markdown links keep their text, other markup is dropped, and a description longer
+    than ``DESCRIPTION_LIMIT`` characters is cut at a word and ends in ``…``.
+    """
+    text = _plain(entry.description)
+    if len(text) > DESCRIPTION_LIMIT:
+        text = text[:DESCRIPTION_LIMIT].rsplit(" ", 1)[0] + " …"
+    facts = []
+    if entry.purpose:
+        facts.append("Purpose: " + ", ".join(entry.purpose))
+    if entry.author:
+        facts.append("By " + entry.author)
+    if licence := entry_licence_text(entry):
+        facts.append("Licence: " + licence)
+    return "\n".join(part for part in (text, " · ".join(facts)) if part)
+
+
+def entry_licence_text(entry: CatalogEntry) -> str:
+    """Return *entry*'s licence, followed by ``non-commercial`` unless its author states otherwise.
+
+    Empty for a commercial entry without a licence.
+    """
+    parts = [entry.license, "" if entry.commercial else NON_COMMERCIAL]
+    return " · ".join(part for part in parts if part)
+
+
+def entry_search_text(entry: CatalogEntry) -> str:
+    """Return the lower-case text a search over *entry* matches; see :func:`search_matches`.
+
+    Its name, scale (``4x``), architecture, purpose, tags, licence, commercial use and description.
+    """
+    words = [
+        entry.name,
+        f"{entry.scale}x" if entry.scale else "",
+        entry.architecture,
+        *entry.purpose,
+        *entry.tags,
+        entry.license,
+        COMMERCIAL if entry.commercial else NON_COMMERCIAL,
+        _plain(entry.description),
+    ]
+    return " ".join(w for w in words if w).lower()
+
+
+def search_matches(query: str, text: str) -> bool:
+    """Return whether every word of *query* appears somewhere in *text*, ignoring case.
+
+    A word may be part of a longer one, in any order. An empty query matches everything. The word
+    ``commercial`` does not match ``non-commercial``.
+    """
+    text = text.lower()
+    return all(
+        _NOT_NON_COMMERCIAL.search(text) if word == COMMERCIAL else word in text
+        for word in query.lower().split()
+    )
+
+
+def _plain(markdown: str) -> str:
+    text = _LINK.sub(r"\1", markdown or "")
+    text = _MARKUP.sub("", text)
+    return " ".join(text.split())

haybale_models/docs/haybale-models.editor.ModelsEditor.md

@@ -0,0 +1,35 @@
+# Models
+
+`haybale-models:editor:ModelsEditor` · kind: editor
+
+Model files on this machine and where they come from
+
+## Details
+
+- **default_slot**: `action`
+- **opens**: `OpenBehavior.REQUIRED`
+- **order**: `100`
+
+## Notes
+
+Lists the model catalog, the files in the Models volume and the Hugging Face files in use.
+
+The **Filters** icon in the header opens the filter section. Its **Show:** row picks the sections
+(**Catalog**, **Models**, **Huggingface**). Its search toggle opens a field that narrows every
+section: each word must appear, in any order, and scale, licence and `commercial` count too. Its
+**+** adds a filter on library, purpose, tags, architecture, scale, commercial use or status, which
+narrows the catalog. Closing it switches them all off, so a closed section always means the full
+list, and reopening it brings them back.
+The **Needed** section is always shown: per open graph, the models it lacks; **Fetch all**
+downloads them. Each section header counts what the section lists; hover it for the whole
+volume's size.
+
+A catalog row reads `non-commercial` unless its author states the model may be used commercially.
+Each catalog resource shows its status: `local`, `HF cache`, `downloading`, `failed` or
+`not downloaded`. **Models** lists the files in the Models volume that no catalog entry
+occupies. **Huggingface** lists the cached files that catalog entries or open graphs use,
+and what uses each; the Volumes editor shows where the cache lives. Download, Cancel and
+Remove need admin access; a removal names the open graphs that use the file and the bytes
+it frees. Cancel is offered only while a download can still be stopped: a Hugging Face
+download, once started, runs to the end. Opening the editor and **Rescan** index the Models
+volume on a background thread.

haybale_models/docs/haybale-models.farmhand.fetch.md

@@ -0,0 +1,23 @@
+# Fetch models
+
+`haybale-models:farmhand:fetch` · kind: farmhand
+
+Queue downloads: everything an open graph needs, or catalog entries by key (admin).
+
+## Agent Instructions
+
+Queues the downloads and returns their job ids; progress is in models_list jobs. Pass binding_id to fetch everything that graph needs, or keys=[<catalog key>, ...] to fetch the preferred variant of each entry. Needs admin access.
+
+## Details
+
+- **input_schema**: `{'type': 'object', 'properties': {'binding_id': {'type': 'string', 'default': None}, 'keys': {'type': 'array', 'items': {'type': 'string'}, 'default': None}}, 'required': []}`
+- **annotations**: `{'read_only_hint': False, 'destructive_hint': False, 'idempotent_hint': False, 'open_world_hint': True}`
+
+## Notes
+
+Queues downloads of missing models. Needs ADMIN access.
+
+Pass `binding_id` to fetch everything that open graph needs, or `keys` to fetch the
+preferred variant of each catalog entry. Returns the job ids at once; the files
+arrive in the background, so call `models_list` to follow the jobs. Downloads use
+network and disk: call it only when the user wants the models.

haybale_models/docs/haybale-models.farmhand.list.md

@@ -0,0 +1,28 @@
+# List models
+
+`haybale-models:farmhand:list` · kind: farmhand
+
+Catalog entries with per-variant status, indexed local files and download jobs.
+
+## Agent Instructions
+
+Lists the model catalog (every entry with each variant's size and status: local, HF cache, queued, downloading N %, failed, not downloaded), the files indexed in the models volume, and the download jobs. detail=true adds license, commercial and description to catalog rows.
+
+## Details
+
+- **input_schema**: `{'type': 'object', 'properties': {'detail': {'type': 'boolean', 'default': False}}, 'required': []}`
+- **annotations**: `{'read_only_hint': True, 'destructive_hint': False, 'idempotent_hint': False, 'open_world_hint': False}`
+
+## Notes
+
+Lists the model catalog, the local model files and the download jobs.
+
+Returns every catalog entry with its variants (size and status: local, HF cache,
+queued, downloading, failed or not downloaded), every file indexed in the models
+volume, and every download job with its progress. Call it first, to learn the
+catalog keys that `models_fetch` takes and to follow a download. Pass
+`detail=true` to add license, commercial (whether the author states the model may be used
+commercially) and description to the catalog rows.
+
+A failed job with `assisted` true can be finished by the user: they get the file from one of its
+`page_urls` and hand it over with Download in browser in the Models editor.

haybale_models/docs/haybale-models.farmhand.needs.md

@@ -0,0 +1,24 @@
+# Models a graph needs
+
+`haybale-models:farmhand:needs` · kind: farmhand
+
+The models an open graph references and does not have locally.
+
+## Agent Instructions
+
+Returns the missing models of an open graph, interiors included: label, size (null when unknown) and the node_ids that ask for it. Get binding_id from haystack_list_graphs.
+
+## Details
+
+- **input_schema**: `{'type': 'object', 'properties': {'binding_id': {'type': 'string'}}, 'required': ['binding_id']}`
+- **annotations**: `{'read_only_hint': True, 'destructive_hint': False, 'idempotent_hint': False, 'open_world_hint': False}`
+
+## Notes
+
+Lists the models an open graph needs and this machine does not have.
+
+Returns, for the graph under `binding_id`, each missing model with its label, size
+(`null` when unknown) and the ids of the nodes that ask for it, interiors of
+subgraphs and macros included, plus the known total in bytes. Call it before
+starting a graph, or to decide what `models_fetch` should download. Get
+`binding_id` from `haystack_list_graphs`.

haybale_models/docs/haybale-models.node.ResolveModelNode.md

@@ -0,0 +1,33 @@
+# Resolve Model
+
+`haybale-models:node:ResolveModelNode` · kind: node
+
+Gives the file of a model on this machine, with its sha256 and size
+
+## Ports
+
+| id | direction | type | description |
+|---|---|---|---|
+| `model` | inlet | `haybale-models:type:MODEL_REF` | Model to look up |
+| `path` | outlet | `haywire-core:type:FILE_PATH` | The model's file on this machine |
+| `sha256` | outlet | `haywire-core:type:STRING` | The file's sha256 checksum |
+| `size` | outlet | `haywire-core:type:INT` | The file's size in bytes |
+
+## Notes
+
+Gives the file of a model on this machine, with its sha256 and size.
+
+Pick the model on the card: a catalog entry, or a file by its address
+(`hf://` or `https://`) or from the Models or Hugging Face cache folder.
+Nothing is downloaded here: a model that is not on this machine is an
+error until it is fetched in the Models editor.
+
+- `Path`: the file, as a path inside the `Models` or `Hugging Face cache`
+  volume. For a model made of several files, or a whole-repo `hf://`
+  address, the folder holding them.
+- `sha256`: the file's digest, empty when it is not known yet.
+- `Size`: the size in bytes, of all files together for a model made of
+  several files.
+
+For a catalog entry with several variants, the first one on disk in the
+port's preferred order is used.

haybale_models/docs/haybale-models.panel.ModelsSettingsPanel.md

@@ -0,0 +1,14 @@
+# Models
+
+`haybale-models:panel:ModelsSettingsPanel` · kind: panel
+
+Model download behaviour: how many files download in parallel.
+
+## Details
+
+- **surface**: `barns`
+- **order**: `10`
+
+## Notes
+
+Model download behaviour: how many files download in parallel.

haybale_models/docs/haybale-models.setting.ModelsSettings.md

@@ -0,0 +1,9 @@
+# Models
+
+`haybale-models:setting:ModelsSettings` · kind: setting
+
+Model download behaviour
+
+## Notes
+
+Settings for the model download queue.

haybale_models/docs/haybale-models.state.ModelsState.md

@@ -0,0 +1,9 @@
+# Models
+
+`haybale-models:state:ModelsState` · kind: state
+
+Model volumes, file index, download queue and the pool of loaded models
+
+## Notes
+
+Owns the model volumes, the file index and the pool of loaded models.

haybale_models/docs/haybale-models.type.MODEL_REF.md

@@ -0,0 +1,19 @@
+# Model
+
+`haybale-models:type:MODEL_REF` · kind: type
+
+A model, named by a catalog entry or by a file (url and sha256). Resolve Model turns it into a path on disk.
+
+## Details
+
+- **wire_type**: `haywire-core:wire:DataWire`
+- **default**: `{'value': {}}`
+- **widget_key**: `haybale-models:widget:ModelRefWidget`
+- **color**: `#b07ad9`
+
+## Notes
+
+A model, named by a catalog entry or by a file.
+
+The port's `filter` narrows the picker by catalog field; `prefer` is an ordered
+list of tag sets used to pick a variant.

haybale_models/docs/haybale-models.widget.ModelRefWidget.md

@@ -0,0 +1,53 @@
+# ModelRefWidget
+
+`haybale-models:widget:ModelRefWidget` · kind: widget
+
+Pick a model
+
+## Details
+
+- **min_width**: `220`
+
+## Notes
+
+Pick the model a node uses.
+
+The card shows one line: a state icon, the model's name and an edit button.
+
+- Green icon: the model is on this machine.
+- Amber `online_prediction` icon: an address that is not pinned yet; the pin
+  button next to it records the digest of the file on this machine.
+- Amber icon: the model is not on this machine; fetch it in the Models editor.
+- Red icon: no model chosen.
+
+Hover the name for the details. The edit button opens the picker, with the
+choice at the top and four tabs:
+
+- **Catalog**: the catalog entries this port accepts. A catalog entry follows
+  the catalog: each machine uses the variant that suits it best.
+- **Models Volume**: a file in the Models volume, named by its sha256: one you
+  put there or fetched by address. Catalog variants are in the Catalog tab.
+- **Hugging Face**: a model file in the Hugging Face cache, named by its `hf://`
+  address and its sha256, so it works offline.
+- **Address**: an `hf://` or `https://` address. An `hf://` address is pinned when
+  the Hub answers; an `https://` address once the file is on this machine.
+
+**Use** stores the choice (pinned when its digest is known), **Cancel** drops it.
+The picker never downloads: missing models are fetched in the Models editor.
+
+Options (in `properties`):
+
+- `filter`: catalog field to allowed values, such as `{"architecture": ["esrgan"]}`;
+  only matching entries are listed.
+- `unavailable`: when set, the edit button is unavailable and this text is its
+  tooltip, for a node that cannot take a model yet. A node that learns later which
+  models fit sets `filter` and drops `unavailable` by rebuilding the port in `rejig`.
+- `catalog_only`: when true, only catalog entries can be chosen: the other tabs are not shown,
+  for a node that needs what a catalog entry says about the model. A value already saved stays
+  shown until it is replaced.
+- `prefer`: ordered list of tag sets for choosing a variant, such as
+  `[{"precision": "fp16"}, {}]`.
+
+```python
+MODEL_REF.as_inlet("model", widget_config={"properties": {"filter": {"scale": [4]}}})
+```