diff --git a/docs/module-internals.md b/docs/module-internals.md index 9f0d2df..98a796d 100644 --- a/docs/module-internals.md +++ b/docs/module-internals.md @@ -1035,7 +1035,10 @@ not one producer. The freeze file keys the Muse Image head downloads and pin that `identify` does not import this module. Weights stay out of git. The Hub snapshot is `wiltodelta/raiw-photo-classify`. -`RAIW_CLASSIFY_WEIGHTS` overrides it. The extra is `classify`. User guide: +`RAIW_CLASSIFY_WEIGHTS` overrides it. The extra is `classify`. The runtime +resolves the snapshot through `WEIGHTS_ALLOW_PATTERNS`, exported so an +offline deploy's pre-cache installs the exact runtime request as one name +(`tests/test_classify.py` pins the seam). User guide: [photo-classify.md](photo-classify.md). Hub card: [photo-classify-hf/README.md](photo-classify-hf/README.md). diff --git a/docs/photo-classify.md b/docs/photo-classify.md index 1af04db..864b032 100644 --- a/docs/photo-classify.md +++ b/docs/photo-classify.md @@ -45,7 +45,11 @@ Weights are not in git. First call downloads at the freeze revision pinned in `classify.py`, or reads `RAIW_CLASSIFY_WEIGHTS` if that directory already has `clip-l-ft.pt`, `probe-weights-clip-l-ft.npz`, `detector.pt`, and `provider.pt`. The training -catalog is not published with the package. +catalog is not published with the package. Deployments that pre-cache the +snapshot for offline resolution must request exactly +`classify.WEIGHTS_ALLOW_PATTERNS` (the runtime resolves that same list under +`HF_HUB_OFFLINE=1`; a partial pre-cache fails the whole resolution, even for +files the runtime would fall back on). ## What one call returns diff --git a/src/remove_ai_watermarks/classify.py b/src/remove_ai_watermarks/classify.py index e3a88d3..e1b5b8a 100644 --- a/src/remove_ai_watermarks/classify.py +++ b/src/remove_ai_watermarks/classify.py @@ -39,6 +39,18 @@ _WEIGHT_FILES = (CLIP_FILE, PROBE_FILE, DETECTOR_FILE, PROVIDER_FILE) RECEIPT_GATE_FILE = "receipt-gate-2026-09-02.npz" RECEIPT_GATE_THRESHOLD = 2.0432573877459697 +# The complete file set the runtime's snapshot_download may request. Deploy-time +# pre-caches in offline environments (HF_HUB_OFFLINE=1) must install exactly +# this list: a subset leaves the runtime resolution unsatisfiable and every +# classify call fails on weights that are four-fifths present. Exported so a +# pre-cache imports one name instead of re-listing the files and silently +# drifting the next time a release adds one. +WEIGHTS_ALLOW_PATTERNS: tuple[str, ...] = ( + *_WEIGHT_FILES, + OPERATING_POINT_FILE, + RECEIPT_GATE_FILE, +) + # 2026-08-31 freeze operating point. Ridge threshold is also in the probe file; # the runtime prefers the file when weights load. MLP_THRESHOLD = 5.9586493237495395 @@ -289,7 +301,7 @@ def _weights_dir() -> Path: snapshot_download( repo_id=WEIGHTS_REPO, revision=WEIGHTS_REVISION, - allow_patterns=[*_WEIGHT_FILES, OPERATING_POINT_FILE, RECEIPT_GATE_FILE], + allow_patterns=list(WEIGHTS_ALLOW_PATTERNS), ) ) except Exception as exc: diff --git a/tests/test_classify.py b/tests/test_classify.py index fcaa6d4..265d1ad 100644 --- a/tests/test_classify.py +++ b/tests/test_classify.py @@ -18,6 +18,7 @@ from remove_ai_watermarks.classify import ( RECEIPT_GATE_FILE, RECEIPT_GATE_THRESHOLD, RIDGE_THRESHOLD, + WEIGHTS_ALLOW_PATTERNS, WEIGHTS_REPO, WEIGHTS_REVISION, classify_from_scores, @@ -52,6 +53,32 @@ def test_hub_snapshot_is_the_freeze_revision() -> None: assert WEIGHTS_REVISION == "ffc46db5135ee3a83f51538f2c8f7483b9b8b40c" +def test_runtime_requests_the_exported_allow_patterns(monkeypatch: pytest.MonkeyPatch) -> None: + """The exported pre-cache list and the runtime request must stay one list. + + An offline deploy pre-caches WEIGHTS_ALLOW_PATTERNS from the Hub; the + runtime then resolves the same snapshot offline. If the runtime's request + grows a file the export misses, every offline deployment breaks on a + four-fifths-present cache, so pin the seam in both directions. + """ + from remove_ai_watermarks.classify import _WEIGHT_FILES, _weights_dir + + captured: dict[str, object] = {} + + def fake_snapshot_download(**kwargs: object) -> str: + captured.update(kwargs) + return "/cached/snapshot" + + monkeypatch.setattr("huggingface_hub.snapshot_download", fake_snapshot_download) + monkeypatch.delenv("RAIW_CLASSIFY_WEIGHTS", raising=False) + + assert _weights_dir() == Path("/cached/snapshot") + assert captured["repo_id"] == WEIGHTS_REPO + assert captured["revision"] == WEIGHTS_REVISION + assert captured["allow_patterns"] == list(WEIGHTS_ALLOW_PATTERNS) + assert set(_WEIGHT_FILES) <= set(WEIGHTS_ALLOW_PATTERNS) + + def test_shipped_operating_point_matches_the_runtime_defaults() -> None: payload = json.loads(Path("docs/photo-classify-hf/operating-point.json").read_text()) assert payload["model1"]["mlp_threshold"] == MLP_THRESHOLD