Files
wildminder-ComfyUI-KEEP/tools/fetch_weight_hashes.py
T
WildAi ec9640f34b chore(docs): point the tracked references at the local .dev/docs tree
The audit artifacts and the plan moved to the local .dev/docs tree, so the
references in the code, the tools and the tests follow them there. The
license file names the tool that produces the record instead of a path.
2026-10-24 22:44:06 +03:00

489 lines
18 KiB
Python

"""P5.0 - the provenance of the weight manifest (one verified download each).
This host has **no** KEEP weights (P0.2: 0 of 8 present), so the sha256 values
of the manifest cannot come from a local inventory. They come from a **one
time verified download**: every upstream URL is downloaded exactly once, the
sha256 of the downloaded bytes is computed and cross-checked against the
upstream ``-<8 hex>`` file name convention, and the result is recorded in
``.dev/docs/audit/2026-10-22-weight-manifest-sources.md``.
The script is
* **env-gated** - the network is only touched when ``KEEP_FETCH_WEIGHT_HASHES=1``
(or ``--fetch``) is given; without it the script runs **offline** and renders
the document from the committed cache. An offline run without a cache entry
fails with an actionable message instead of silently inventing a hash;
* **idempotent** - a fetch that finds a verified cache entry re-uses it (no
second download), and the document is rendered from the cache, so a re-run
produces byte-identical output;
* **404-safe** - a URL that answers ``404`` is recorded as *unavailable*
(``sha256: null``, ``unverified: true``) and **no download is attempted**.
The transport is injectable (:class:`HttpTransport`), so the behaviour above is
unit-testable without a network.
"""
from __future__ import annotations
import argparse
import hashlib
import json
import os
import re
import shutil
import sys
import tempfile
import time
import urllib.error
import urllib.request
from dataclasses import dataclass
from pathlib import Path
REPO_ROOT = Path(__file__).resolve().parent.parent
#: The machine-readable cache (the state of the one time fetch).
DEFAULT_CACHE = REPO_ROOT / "tools" / "weight_hashes.json"
#: The provenance document the script writes.
DEFAULT_DOC = REPO_ROOT / ".dev" / "docs" / "audit" / "2026-10-22-weight-manifest-sources.md"
#: The environment variable that allows the script to touch the network.
FETCH_ENV = "KEEP_FETCH_WEIGHT_HASHES"
#: The upstream release the weights come from.
RELEASE = "v1.0.0"
RELEASE_BASE = f"https://github.com/jnjaby/KEEP/releases/download/{RELEASE}"
#: The truthy values of :data:`FETCH_ENV`.
TRUTHY = frozenset({"1", "true", "yes", "on"})
#: The upstream file name convention: ``<name>-<8 hex chars>.pth``.
NAME_PREFIX_PATTERN = re.compile(r"[-_]([0-9a-f]{8})\.[A-Za-z0-9]+$")
@dataclass(frozen=True)
class WeightSource:
"""One upstream weight file.
``filename_prefix`` is the ``-<8 hex>`` part of the upstream file name - the
upstream convention is that it is the **first 8 hex chars of the sha256** of
the file. It is the cross-check of the provenance, not a hash.
"""
name: str
dest_folder: str
subdir: str
required: bool
filename_prefix: str | None = None
@property
def url(self) -> str:
return f"{RELEASE_BASE}/{self.name}"
#: The eight manifest entries (the legacy on-disk layout is kept exactly:
#: ``keep_models`` + the ``KEEP`` sub-directory, ``facedetection``, ``gmflow``).
WEIGHT_SOURCES: tuple[WeightSource, ...] = (
WeightSource("KEEP-b76feb75.pth", "keep_models", "KEEP", True, "b76feb75"),
# Not published upstream (HTTP 404 - P0.2 §3 / P0.4 §4.1). The name prefix
# is the only provenance the file name carries.
WeightSource("KEEP_Asian-4765ebe0.pth", "keep_models", "KEEP", False, "4765ebe0"),
WeightSource("detection_Resnet50_Final.pth", "facedetection", "", True),
WeightSource("detection_mobilenet0.25_Final.pth", "facedetection", "", True),
WeightSource("yolov5n-face.pth", "facedetection", "", True),
WeightSource("yolov5l-face.pth", "facedetection", "", True),
WeightSource("parsing_parsenet.pth", "facedetection", "", True),
# The GMFlow checkpoint is defensive (P0.4 §2): the KEEP checkpoint carries
# the trained flownet weights.
WeightSource("gmflow_sintel-0c07dcb3.pth", "gmflow", "", False, "0c07dcb3"),
)
class FetchError(RuntimeError):
"""Raised when a weight cannot be fetched or verified."""
class HttpTransport:
"""The network transport (stdlib only - no torch, no requests)."""
def __init__(self, timeout: int = 120) -> None:
self.timeout = timeout
self.user_agent = "ComfyUI-KEEP/2.0 (weight provenance fetch)"
def head(self, url: str) -> tuple[int, int | None]:
"""``(status, content_length)`` of the URL (404 is not an error)."""
request = urllib.request.Request(
url, method="HEAD", headers={"User-Agent": self.user_agent}
)
try:
with urllib.request.urlopen(request, timeout=self.timeout) as response:
length = response.headers.get("Content-Length")
return response.status, int(length) if length is not None else None
except urllib.error.HTTPError as exc:
return exc.code, None
def download(
self,
url: str,
dest: Path,
attempts: int = 3,
expected_size: int | None = None,
) -> tuple[str, int]:
"""Download ``url`` to ``dest``; returns ``(sha256, size)``.
A partial download is kept between the attempts and continued with an
HTTP ``Range`` request (the KEEP checkpoint is ~634 MB - a restart from
zero would be expensive).
"""
part = dest.with_name(dest.name + ".part")
last_error: Exception | None = None
for attempt in range(1, attempts + 1):
try:
return self._download_once(url, part, dest, expected_size)
except Exception as exc: # noqa: BLE001 - retried below
last_error = exc
print(
f" attempt {attempt}/{attempts} failed: {exc}",
file=sys.stderr,
flush=True,
)
if attempt < attempts:
time.sleep(2 * attempt)
raise FetchError(f"the download of {url} failed: {last_error}")
def _download_once(
self,
url: str,
part: Path,
dest: Path,
expected_size: int | None,
) -> tuple[str, int]:
already = part.stat().st_size if part.exists() else 0
if expected_size is not None and already == expected_size:
# the partial file is already complete - hash it, do not fetch it again
return self._finish(part, dest)
headers = {"User-Agent": self.user_agent}
if already:
headers["Range"] = f"bytes={already}-"
request = urllib.request.Request(url, headers=headers)
try:
response = urllib.request.urlopen(request, timeout=self.timeout)
except urllib.error.HTTPError as exc:
if exc.code == 416 and already:
# the server refuses the range: the partial file is complete
return self._finish(part, dest)
raise
with response:
append = already > 0 and response.status == 206
digest = hashlib.sha256()
if append:
with part.open("rb") as existing:
for chunk in iter(lambda: existing.read(1 << 20), b""):
digest.update(chunk)
else:
already = 0
with part.open("ab" if append else "wb") as handle:
while True:
chunk = response.read(1 << 20)
if not chunk:
break
handle.write(chunk)
digest.update(chunk)
size = part.stat().st_size
part.replace(dest)
return digest.hexdigest(), size
@staticmethod
def _finish(part: Path, dest: Path) -> tuple[str, int]:
"""Hash a complete partial file and move it to its destination."""
digest = hashlib.sha256()
with part.open("rb") as handle:
for chunk in iter(lambda: handle.read(1 << 20), b""):
digest.update(chunk)
size = part.stat().st_size
part.replace(dest)
return digest.hexdigest(), size
def _utc_now() -> str:
return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
def fetch_entry(
source: WeightSource, transport: HttpTransport, workdir: Path
) -> dict:
"""Fetch one weight and return its record.
A URL that answers ``404`` is recorded as *unavailable* and **no download
is attempted** - the file does not exist upstream, so there is nothing to
hash.
"""
status, length = transport.head(source.url)
if status == 404:
return {
"name": source.name,
"url": source.url,
"dest_folder": source.dest_folder,
"subdir": source.subdir,
"required": source.required,
"available": False,
"unverified": True,
"sha256": None,
"size": None,
"http_status": 404,
"filename_prefix": source.filename_prefix,
"prefix_match": None,
"fetched_at": _utc_now(),
"note": "not published upstream (HTTP 404) - no hash is obtainable",
}
if status != 200:
raise FetchError(
f"the URL {source.url} answered HTTP {status} (expected 200 or 404)"
)
target = workdir / source.name
sha256, size = transport.download(source.url, target, expected_size=length)
if length is not None and size != length:
raise FetchError(
f"the download of {source.url} is incomplete: {size} of {length} bytes"
)
if target.exists():
target.unlink()
prefix_match: bool | None = None
if source.filename_prefix:
prefix_match = sha256.startswith(source.filename_prefix)
if not prefix_match:
# Not fatal - the recorded sha256 is the truth - but it is a
# finding that must be visible.
print(
f"WARNING: the sha256 of {source.name} does not start with the "
f"file name prefix {source.filename_prefix!r} - the upstream "
"renaming convention does not hold for this asset",
file=sys.stderr,
)
return {
"name": source.name,
"url": source.url,
"dest_folder": source.dest_folder,
"subdir": source.subdir,
"required": source.required,
"available": True,
"unverified": False,
"sha256": sha256,
"size": size,
"http_status": 200,
"filename_prefix": source.filename_prefix,
"prefix_match": prefix_match,
"fetched_at": _utc_now(),
"note": "",
}
def load_cache(path: Path) -> dict:
"""The committed cache (an empty cache when the file does not exist)."""
if not path.is_file():
return {"release": RELEASE, "entries": {}}
data = json.loads(path.read_text(encoding="utf-8"))
data.setdefault("entries", {})
return data
def write_cache(path: Path, cache: dict) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
path.write_text(
json.dumps(cache, indent=2, sort_keys=True, ensure_ascii=False) + "\n",
encoding="utf-8",
)
def build_cache(
fetch: bool,
cache: dict,
transport: HttpTransport,
workdir: Path,
cache_path: Path | None = None,
) -> dict:
"""The cache with every entry fetched (or taken from the cache).
The cache is written after **every** entry, so an interrupted run (the
KEEP checkpoint is ~634 MB) keeps the hashes it already verified.
"""
entries = dict(cache.get("entries", {}))
for source in WEIGHT_SOURCES:
recorded = entries.get(source.name)
if recorded and recorded.get("url") == source.url:
if recorded.get("sha256") or recorded.get("http_status") == 404:
print(f" {source.name}: cached - no second download")
continue
if not fetch:
raise FetchError(
f"no cached record for {source.name} and the network is disabled. "
f"Re-run with {FETCH_ENV}=1 (or --fetch) to download it once."
)
print(f" {source.name}: downloading from {source.url}", flush=True)
entries[source.name] = fetch_entry(source, transport, workdir)
if cache_path is not None:
write_cache(cache_path, {"release": RELEASE, "entries": entries})
return {"release": RELEASE, "entries": entries}
def render_doc(cache: dict) -> str:
"""The provenance document (deterministic - a function of the cache)."""
entries = [cache["entries"][source.name] for source in WEIGHT_SOURCES]
verified = [e for e in entries if e.get("sha256")]
unavailable = [e for e in entries if not e.get("available")]
lines: list[str] = []
lines.append("# P5.0 - Weight manifest provenance: one verified download per weight")
lines.append("")
lines.append(
"- Date of record: 2026-10-22 (plan date) - executed 2026-10-23\n"
"- Script: `tools/fetch_weight_hashes.py` (env-gated, idempotent)\n"
"- Machine-readable cache: `tools/weight_hashes.json`\n"
f"- Upstream release: `jnjaby/KEEP {RELEASE}`"
)
lines.append("")
lines.append("## 1. Method")
lines.append("")
lines.append(
"This host has **no** KEEP weights (P0.2: 0 of 8 present), so the sha256\n"
"values of the manifest cannot come from a local inventory. Every URL was\n"
"therefore downloaded **once**; the sha256 of the downloaded bytes was\n"
"computed and cross-checked against the upstream file name convention\n"
"(``-<8 hex chars>`` = the first 8 hex chars of the sha256)."
)
lines.append("")
lines.append(
"The script is **env-gated** (`KEEP_FETCH_WEIGHT_HASHES=1`), **idempotent**\n"
"(a verified cache entry is never downloaded twice) and **404-safe** (a URL\n"
"that answers 404 is recorded as unavailable and *no download is attempted*)."
)
lines.append("")
lines.append(
f"**Result: {len(verified)} of {len(entries)} weights verified, "
f"{len(unavailable)} unavailable upstream.**"
)
lines.append("")
lines.append("## 2. The eight entries")
lines.append("")
lines.append(
"| # | File | Destination | Required | HTTP | Size (bytes) | sha256 | "
"Name prefix | Prefix matches |"
)
lines.append("|---|---|---|---|---|---|---|---|---|")
for index, entry in enumerate(entries, start=1):
destination = entry["dest_folder"] + (
"/" + entry["subdir"] if entry["subdir"] else ""
)
sha = entry.get("sha256") or "`null`"
prefix = entry.get("filename_prefix") or "-"
match = entry.get("prefix_match")
if match is None:
match_text = "-"
else:
match_text = "yes" if match else "**no**"
lines.append(
f"| {index} | `{entry['name']}` | `{destination}` | "
f"{'yes' if entry['required'] else 'no'} | {entry.get('http_status')} | "
f"{entry.get('size') if entry.get('size') is not None else '-'} | "
f"`{sha}` | `{prefix}` | {match_text} |"
)
lines.append("")
lines.append("## 3. The machine-readable record")
lines.append("")
lines.append("```json")
lines.append(
json.dumps(
{
"release": cache.get("release", RELEASE),
"entries": entries,
"verified_count": len(verified),
"unavailable_count": len(unavailable),
"total_count": len(entries),
},
indent=2,
sort_keys=True,
ensure_ascii=False,
)
)
lines.append("```")
lines.append("")
lines.append("## 4. The unavailable asset")
lines.append("")
for entry in unavailable:
lines.append(
f"`{entry['name']}` is **not published upstream** (HTTP "
f"{entry.get('http_status')}; the release API of `jnjaby/KEEP "
f"{RELEASE}` does not list it - P0.2 section 3 / P0.4 section 4.1).\n"
"No hash is obtainable, so the manifest entry carries `sha256: null`,\n"
"`unverified: true` and `available: false`. The entry is **not**\n"
"required: a required entry would break the *default* KEEP model too\n"
"(both checkpoints share the `keep_models/KEEP` destination). Selecting\n"
"the `Asian` option raises an actionable error instead (P5.3)."
)
lines.append("")
lines.append("## 5. How the manifest uses this record")
lines.append("")
lines.append(
"* the **7** verified entries carry their `sha256` in `WEIGHT_MANIFEST`\n"
" (`src/keep_model_configs.py`) and the downloader verifies every download\n"
" against it;\n"
"* the **Asian** entry carries `sha256: null` + `unverified: true` +\n"
" `available: false` + `sha256_prefix: \"4765ebe0\"` (the only provenance the\n"
" file name carries) - the downloader **skips the verification** for it and\n"
" warns loudly."
)
lines.append("")
return "\n".join(lines)
def write_doc(path: Path, text: str) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
path.write_text(text, encoding="utf-8")
def main(argv: list[str] | None = None) -> int:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--cache", type=Path, default=DEFAULT_CACHE)
parser.add_argument("--doc", type=Path, default=DEFAULT_DOC)
parser.add_argument(
"--fetch",
action="store_true",
help=f"allow the network (same as {FETCH_ENV}=1)",
)
parser.add_argument(
"--offline",
action="store_true",
help="never touch the network (the default without the env variable)",
)
args = parser.parse_args(argv)
env_fetch = (os.environ.get(FETCH_ENV) or "").strip().lower() in TRUTHY
fetch = env_fetch or (args.fetch and not args.offline)
if args.offline:
fetch = False
cache = load_cache(args.cache)
transport = HttpTransport()
workdir = Path(tempfile.mkdtemp(prefix="keep-weight-hashes-"))
try:
cache = build_cache(fetch, cache, transport, workdir, args.cache)
except FetchError as exc:
print(f"error: {exc}", file=sys.stderr)
return 1
finally:
shutil.rmtree(workdir, ignore_errors=True)
write_cache(args.cache, cache)
write_doc(args.doc, render_doc(cache))
print(f"wrote {args.cache}")
print(f"wrote {args.doc}")
return 0
if __name__ == "__main__":
raise SystemExit(main())