mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-10-14 13:30:55 +02:00
Victor Kuznetsov
5b52782f01
The gate head is fitted on the freeze CLIP-L-ft embedding space, so it belongs with the model: _load_receipt_gate now prefers receipt-gate-2026-09-02.npz from the Hub snapshot or the RAIW_CLASSIFY_WEIGHTS directory, with the package asset kept as the fallback for weights frozen before the gate existed. Card mode of the publish workflow now uploads the gate alongside the README and the operating point, and the snapshot allow-pattern includes it.
130 lines
4.5 KiB
Python
130 lines
4.5 KiB
Python
#!/usr/bin/env python3
|
|
"""Publish the photo-classify freeze to Hugging Face wiltodelta/raiw-photo-classify.
|
|
|
|
Manual. Not a library release step. Weights are not in git: pass --src, a
|
|
directory that holds clip-l-ft.pt and the probe, with detector.pt and
|
|
provider.pt either at the root or in run1/.
|
|
|
|
uv run python scripts/publish_photo_classify_hf.py --src DIR
|
|
uv run python scripts/publish_photo_classify_hf.py --card-only
|
|
|
|
Needs HF_TOKEN with write role. The GitHub Action
|
|
publish-photo-classify-hf.yml injects secrets.HF_TOKEN.
|
|
|
|
Card mode uploads the README, operating-point.json, and the receipt-gate
|
|
head (staged from the package assets) without touching the freeze weights.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import logging
|
|
import os
|
|
import shutil
|
|
import tempfile
|
|
from pathlib import Path
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
HUB_REPO = "wiltodelta/raiw-photo-classify"
|
|
CLIP_FILE = "clip-l-ft.pt"
|
|
PROBE_FILE = "probe-weights-clip-l-ft.npz"
|
|
DETECTOR_FILE = "detector.pt"
|
|
PROVIDER_FILE = "provider.pt"
|
|
WEIGHT_FILES = (CLIP_FILE, PROBE_FILE, DETECTOR_FILE, PROVIDER_FILE)
|
|
GATE_FILE = "receipt-gate-2026-09-02.npz"
|
|
REPO_ROOT = Path(__file__).resolve().parents[1]
|
|
GATE_SOURCE = REPO_ROOT / "src" / "remove_ai_watermarks" / "assets" / GATE_FILE
|
|
CARD_DIR = REPO_ROOT / "docs" / "photo-classify-hf"
|
|
|
|
|
|
def weight_paths(src: Path) -> dict[str, Path]:
|
|
"""Map Hub filenames to files under ``src`` or ``src/run1``."""
|
|
src = src.expanduser().resolve()
|
|
if not src.is_dir():
|
|
raise SystemExit(f"missing weights directory: {src}")
|
|
found: dict[str, Path] = {}
|
|
for name in WEIGHT_FILES:
|
|
direct = src / name
|
|
nested = src / "run1" / name
|
|
if direct.is_file():
|
|
found[name] = direct
|
|
elif nested.is_file():
|
|
found[name] = nested
|
|
else:
|
|
raise SystemExit(f"missing {name} in {src} or {src / 'run1'}")
|
|
return found
|
|
|
|
|
|
def stage_card(dest: Path) -> None:
|
|
readme = CARD_DIR / "README.md"
|
|
operating = CARD_DIR / "operating-point.json"
|
|
if not readme.is_file() or not operating.is_file():
|
|
raise SystemExit(f"Hub card missing under {CARD_DIR}")
|
|
if not GATE_SOURCE.is_file():
|
|
raise SystemExit(f"receipt-gate asset missing: {GATE_SOURCE}")
|
|
dest.mkdir(parents=True, exist_ok=True)
|
|
shutil.copy2(readme, dest / "README.md")
|
|
shutil.copy2(operating, dest / "operating-point.json")
|
|
# The gate head versions with the model: card mode carries it even when
|
|
# the four freeze weight files are not being re-uploaded.
|
|
shutil.copy2(GATE_SOURCE, dest / GATE_FILE)
|
|
|
|
|
|
def _place(src: Path, dest: Path) -> None:
|
|
dest.parent.mkdir(parents=True, exist_ok=True)
|
|
try:
|
|
os.link(src, dest)
|
|
except OSError:
|
|
shutil.copy2(src, dest)
|
|
|
|
|
|
def stage_weights(dest: Path, src: Path) -> None:
|
|
for name, path in weight_paths(src).items():
|
|
_place(path, dest / name)
|
|
log.info("stage %s (%s bytes)", name, path.stat().st_size)
|
|
|
|
|
|
def publish(stage: Path, *, token: str, message: str) -> str:
|
|
from huggingface_hub import HfApi, create_repo
|
|
|
|
create_repo(HUB_REPO, repo_type="model", exist_ok=True, private=False, token=token)
|
|
api = HfApi(token=token)
|
|
commit = api.upload_folder(
|
|
folder_path=str(stage),
|
|
repo_id=HUB_REPO,
|
|
repo_type="model",
|
|
commit_message=message,
|
|
allow_patterns=["README.md", "operating-point.json", *WEIGHT_FILES, GATE_FILE],
|
|
)
|
|
return getattr(commit, "oid", "") or ""
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--src", type=Path, help="Freeze or weights directory")
|
|
parser.add_argument("--card-only", action="store_true")
|
|
parser.add_argument("--message", default="")
|
|
args = parser.parse_args()
|
|
logging.basicConfig(level=logging.INFO, format="%(message)s")
|
|
token = os.environ.get("HF_TOKEN", "").strip()
|
|
if not token:
|
|
raise SystemExit("HF_TOKEN is not set (write role required)")
|
|
if args.card_only == bool(args.src):
|
|
raise SystemExit("pass exactly one of --src or --card-only")
|
|
message = args.message.strip() or (
|
|
"Update photo-classify card" if args.card_only else "Publish photo-classify freeze"
|
|
)
|
|
with tempfile.TemporaryDirectory(prefix="raiw-hf-") as tmp:
|
|
stage = Path(tmp)
|
|
stage_card(stage)
|
|
if args.src is not None:
|
|
stage_weights(stage, args.src)
|
|
oid = publish(stage, token=token, message=message)
|
|
log.info("published %s revision=%s", HUB_REPO, oid or "(unknown)")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|