diff --git a/README.md b/README.md index e502b80..c264e61 100644 --- a/README.md +++ b/README.md @@ -10,10 +10,12 @@ model. The reference RepViT-M0.9 artifact has 4.76 million parameters and is 19.45 MiB in FP32. > [!IMPORTANT] -> This repository contains research and training code. It does not distribute a -> pretrained checkpoint. The CPC and GAICD archives reviewed during development -> do not state explicit image or annotation licenses; confirm training and model -> redistribution rights before publishing a derived checkpoint. +> Pretrained ONNX and PyTorch checkpoints are on +> [Hugging Face](https://huggingface.co/appwrite/focalnet) and the +> [2026-09-14-rc1 GitHub release](https://github.com/appwrite/focalnet/releases/tag/2026-09-14-rc1). +> The CPC and GAICD archives reviewed during development do not state explicit +> image or annotation licenses; the published weights do not grant rights to +> those datasets. ## How it works @@ -93,9 +95,15 @@ ONNX Runtime, so a runtime-only environment can use `pip install .`. ## Inference -Predict an importance map and an aspect-aware crop: +Download the published checkpoints, then predict an importance map and an +aspect-aware crop: ```sh +mkdir -p artifacts +curl -fsSL \ + "https://huggingface.co/appwrite/focalnet/resolve/main/focalnet.onnx" \ + -o artifacts/focalnet.onnx +# same file: https://github.com/appwrite/focalnet/releases/download/2026-09-14-rc1/focalnet.onnx uv run focalnet predict artifacts/focalnet.onnx photo.jpg \ --aspect-ratio 16:9 \ --heatmap artifacts/photo-importance.npy @@ -120,26 +128,27 @@ the same image; it is not a calibrated confidence value. ### Browser demo A static crop lab under `web/` runs the same format-v2 ONNX in the browser with -ONNX Runtime Web. Photos stay on-device. This repository does not publish a -checkpoint; the 19.45 MiB weights stay gitignored. Export a model, then point -the lab at that file. Bun is the package manager, bundler, test runner, and -local static server: +ONNX Runtime Web. Photos stay on-device. The 19.45 MiB weights stay gitignored; +`bun run fetch-model` downloads `focalnet-human.onnx` from +[Hugging Face](https://huggingface.co/appwrite/focalnet) and verifies the +published SHA-256. Bun is the package manager, bundler, test runner, and local +static server: ```sh cd web bun install -FOCALNET_ONNX=../artifacts/focalnet-human.onnx bun run fetch-model +bun run fetch-model bun test bun run dev ``` -`FOCALNET_ALLOW_DUMMY=1 bun run fetch-model` writes a tiny placeholder for UI -layout only. Set `FOCALNET_ONNX_SHA256` to verify a known file, or -`FOCALNET_ONNX_URL` to fetch a checkpoint you already host. Source order is -existing `web/public/focalnet-human.onnx`, then `FOCALNET_ONNX`, then -`FOCALNET_ONNX_URL`; dummy mode still runs if a configured URL fails. The demo -compares the human-ranked crop with the importance-retention crop from the same -map. +Override the source with `FOCALNET_ONNX` (local file) or `FOCALNET_ONNX_URL`. +`FOCALNET_ONNX_SHA256` defaults to the published ranking-model hash, including +when the variable is set but empty. Source order is existing +`web/public/focalnet-human.onnx`, then `FOCALNET_ONNX`, then the Hub URL. +`FOCALNET_ALLOW_DUMMY=1` writes a tiny placeholder for UI layout only if a +configured source is missing or fails. The demo compares the human-ranked crop +with the importance-retention crop from the same map. ## Training @@ -300,8 +309,9 @@ crop optimization, and ONNX parity. They use synthetic inputs and do not require teacher models or pretrained weights. `modal_app.py` and the scripts in `scripts/` provide reference workflows for -remote data preparation and GPU training. Review their resource names, storage -paths, and expected dataset hashes before running them in another environment. +remote data preparation, GPU training, and Hugging Face publishing. Review +their resource names, storage paths, and expected dataset hashes before running +them in another environment. ## Limitations @@ -311,12 +321,16 @@ paths, and expected dataset hashes before running them in another environment. - The GAICD test split was inspected during several development iterations, so it should be treated as an engineering benchmark rather than an untouched final test. -- Dataset rights must be resolved before distributing a trained checkpoint. +- CPC and GAICD archives do not state explicit licenses; published checkpoints + do not grant rights to those datasets. - Model latency depends on the ONNX Runtime build, processor, thread count, and surrounding service workload. ## License -FocalNet's source code is available under the [MIT License](LICENSE). This -license does not grant rights to third-party datasets, teacher model weights, or -derived model artifacts; review their respective terms before redistribution. +FocalNet's source code and the published checkpoints on +[Hugging Face](https://huggingface.co/appwrite/focalnet) and the +[GitHub release](https://github.com/appwrite/focalnet/releases/tag/2026-09-14-rc1) +are available under the [MIT License](LICENSE). This license does not grant +rights to third-party datasets or teacher model weights; review their +respective terms before redistribution. diff --git a/docs/hub-images/comparison-faces.png b/docs/hub-images/comparison-faces.png new file mode 100644 index 0000000..25cce4b Binary files /dev/null and b/docs/hub-images/comparison-faces.png differ diff --git a/docs/hub-images/comparison-retriever.png b/docs/hub-images/comparison-retriever.png new file mode 100644 index 0000000..d08b70a Binary files /dev/null and b/docs/hub-images/comparison-retriever.png differ diff --git a/docs/huggingface-model-card.md b/docs/huggingface-model-card.md new file mode 100644 index 0000000..33b9969 --- /dev/null +++ b/docs/huggingface-model-card.md @@ -0,0 +1,151 @@ +--- +license: mit +library_name: onnxruntime +pipeline_tag: image-to-image +tags: + - onnx + - computer-vision + - image-cropping + - saliency +base_model: timm/repvit_m0_9 +--- + +# FocalNet + +Appwrite FocalNet is a compact vision model for content-aware image cropping. It +predicts a 64×64 importance map, then ranks crop candidates with a small +composition head trained on human preferences. Training code lives in +[appwrite/focalnet](https://github.com/appwrite/focalnet). + +The reference RepViT-M0.9 ranker has 4.76 million parameters. The published +`focalnet-human.onnx` artifact is 19.45 MiB in FP32. + +## Examples + +These photographs come from [Autogravity](https://github.com/appwrite/autogravity). +A centered 1:1 crop of the park scene keeps grass. FocalNet keeps the dog. + +![Center 1:1 crop of empty grass versus FocalNet 1:1 crop of the golden retriever](images/comparison-retriever.png) + +On a two-person portrait, FocalNet holds the nearer subject instead of splitting +the frame. + +![Center 1:1 crop versus FocalNet 1:1 crop of a two-person portrait](images/comparison-faces.png) + +## Files + +| File | Role | SHA-256 | +| --- | --- | --- | +| `focalnet-human.onnx` | Importance map + crop ranking (format v2) | `59164c601c98cea3f62b25166710831dac63e1a872fc64767c65316ad5385439` | +| `focalnet.onnx` | Importance map only (format v1) | `87becceb269a2973c359df789783be49a9f840f47170a015d3776d7c4145a2ce` | +| `focalnet-human.pt` | PyTorch ranking checkpoint | `82c4d695310c220e5f3bcfbf5e64d315a5e826cfd77e945e61e032460014480a` | +| `focalnet.pt` | PyTorch importance checkpoint | `d1942f0652f8ea85f75ffc0cb1bf40d7e70b38e7ad102ab0e6df5c2e07ce52cf` | + +JSON sidecars next to the ONNX files record the preprocessing contract, parameter +counts, and export verification. + +Use `focalnet-human.onnx` for inference. The PyTorch files are for continued +training, not for the ONNX Runtime path. + +## Usage + +Download the ranking model and crop a photo: + +```sh +mkdir -p artifacts +curl -fsSL \ + "https://huggingface.co/appwrite/focalnet/resolve/main/focalnet-human.onnx" \ + -o artifacts/focalnet-human.onnx +uv run focalnet predict-human artifacts/focalnet-human.onnx photo.jpg \ + --aspect-ratio 16:9 +``` + +From Python: + +```python +from pathlib import Path +from urllib.request import urlretrieve + +from focalnet.human_runtime import HumanCropPredictor + +url = "https://huggingface.co/appwrite/focalnet/resolve/main/focalnet-human.onnx" +path = Path("artifacts/focalnet-human.onnx") +path.parent.mkdir(parents=True, exist_ok=True) +urlretrieve(url, path) +crop = HumanCropPredictor(path).predict("photo.jpg", aspect_ratio=16 / 9) +``` + +Importance-only inference: + +```sh +mkdir -p artifacts +curl -fsSL \ + "https://huggingface.co/appwrite/focalnet/resolve/main/focalnet.onnx" \ + -o artifacts/focalnet.onnx +uv run focalnet predict artifacts/focalnet.onnx photo.jpg --aspect-ratio 16:9 +``` + +`predict-human` generates up to 125 crops across five positions and five zoom +levels, scores a padded maximum of 128 candidates in one ONNX call, and selects +the highest preference score among crops whose retained importance is within +0.05 of the best candidate. The returned human score is a relative logit. +Compare it only among crops for the same image. + +## ONNX contract + +| Field | Contract | +| --- | --- | +| Image input | `image`, float32 `[1,3,256,256]`, RGB NCHW | +| Normalization | `(pixel/255 - [0.485,0.456,0.406]) / [0.229,0.224,0.225]` | +| Resize | Apply EXIF orientation, then aspect-fit with bilinear interpolation | +| Alpha and padding | Composite alpha on black; normalized padding values are zero | +| Importance output | `importance`, float32 `[1,1,64,64]`, sigmoid applied | +| Ranking inputs | `boxes` `[1,128,4]` and letterbox `content` `[1,4]` | +| Ranking output | `crop_scores`, float32 `[1,128]` | + +## Results + +The importance model was trained on 500,000 teacher-labeled Open Images V7 +samples (RepViT-M0.9 encoder, 48-channel decoder, 10 epochs, best epoch 7). +Teacher agreement on the group-disjoint 10,000-image validation split: + +| Metric | Result | +| --- | ---: | +| Map MAE | 0.09975 | +| Normalized centroid error | 0.06430 | +| 1:1 teacher importance retained | 91.09% | +| 16:9 teacher importance retained | 90.86% | +| 4:5 teacher importance retained | 85.56% | + +The ranking head (7,069 parameters) was pretrained on CPC and fine-tuned on +GAICD. The importance network stays frozen. On GAICD's 500-image test split: + +| Metric | GAICD only | CPC → GAICD | +| --- | ---: | ---: | +| Spearman correlation | 0.75696 | **0.76450** | +| Pearson correlation | 0.78211 | **0.79192** | +| Top-5 accuracy | 42.48% | **43.70%** | +| Top-10 accuracy | 64.45% | **64.54%** | +| Pairwise accuracy | 86.00% | **86.44%** | + +The GAICD test split was inspected during several development iterations. Treat +these figures as engineering benchmarks, not an untouched final test. INT8 +post-training quantization did not pass the crop-quality gate; FP32 is the +published artifact. + +## Training data + +- Open Images V7 (teacher-labeled importance maps from U²-Net + YuNet) +- [Comparative Photo Composition (CPC)](https://www3.cs.stonybrook.edu/~cvl/projects/wei2018goods/VPN_CVPR2018s.html) +- [GAICD](https://github.com/HuiZeng/Grid-Anchor-based-Image-Cropping-Pytorch) + +The reviewed CPC and GAICD archives do not state explicit image or annotation +licenses. This checkpoint is published under MIT for the model files; that +license does not grant rights to the third-party datasets or teacher weights +used during training. + +## License + +Model files are available under the [MIT License](https://github.com/appwrite/focalnet/blob/main/LICENSE), +the same license as the training code. Review dataset and teacher-model terms +before redistribution of derived artifacts. diff --git a/docs/results-500k.md b/docs/results-500k.md index 4305b3f..0b2d226 100644 --- a/docs/results-500k.md +++ b/docs/results-500k.md @@ -78,4 +78,8 @@ focalnet.onnx 87becceb269a2973c359df789783be49a9f840f47170a015d3776d7c4145a2ce best.pt d1942f0652f8ea85f75ffc0cb1bf40d7e70b38e7ad102ab0e6df5c2e07ce52cf ``` -The files identified by these hashes are not distributed in this repository. +These files are published at +[huggingface.co/appwrite/focalnet](https://huggingface.co/appwrite/focalnet) +and the +[2026-09-14-rc1 GitHub release](https://github.com/appwrite/focalnet/releases/tag/2026-09-14-rc1) +as `focalnet.onnx` and `focalnet.pt`. They stay gitignored in this repository. diff --git a/docs/results-cpc-gaic-v2.md b/docs/results-cpc-gaic-v2.md index 5d599a6..6ce8b18 100644 --- a/docs/results-cpc-gaic-v2.md +++ b/docs/results-cpc-gaic-v2.md @@ -100,7 +100,11 @@ CPCDataset.tar.gz dfa4ec73c9d9b4b525a8f79aee5670fac4797bad2eb1bd0e1f26f051ac GAIC.zip b895a3f9e03c8f70c37194370441dc2f17bb60e0ab437447ee240b003cd6550b ``` -The files identified by these hashes are not distributed in this repository. -Neither reviewed dataset archive states an explicit license for its images or -annotations. Confirm training and model redistribution rights before publishing -a derived checkpoint. +The ranking checkpoint identified by these hashes is published at +[huggingface.co/appwrite/focalnet](https://huggingface.co/appwrite/focalnet) +and the +[2026-09-14-rc1 GitHub release](https://github.com/appwrite/focalnet/releases/tag/2026-09-14-rc1) +as `focalnet-human.onnx` and `focalnet-human.pt`. Dataset archives stay off the +Hub and the release. Neither reviewed CPC nor GAICD archive states an explicit +license for its images or annotations; the published weights do not grant rights +to those datasets. diff --git a/scripts/publish-huggingface.py b/scripts/publish-huggingface.py new file mode 100644 index 0000000..e0addc7 --- /dev/null +++ b/scripts/publish-huggingface.py @@ -0,0 +1,222 @@ +"""Stage and upload verified FocalNet checkpoints to Hugging Face Hub.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import re +import shutil +from collections.abc import Callable +from pathlib import Path + +REPO_ID = "appwrite/focalnet" +DEFAULT_CARD = Path("docs/huggingface-model-card.md") +DEFAULT_IMAGES = Path("docs/hub-images") +IMAGE_SUFFIXES = {".gif", ".jpeg", ".jpg", ".png", ".webp"} +IMAGE_REFERENCE = re.compile(r"images/[\w.-]+") +EXPECTED_SHA256 = { + "focalnet.onnx": "87becceb269a2973c359df789783be49a9f840f47170a015d3776d7c4145a2ce", + "focalnet.pt": "d1942f0652f8ea85f75ffc0cb1bf40d7e70b38e7ad102ab0e6df5c2e07ce52cf", + "focalnet-human.onnx": "59164c601c98cea3f62b25166710831dac63e1a872fc64767c65316ad5385439", + "focalnet-human.pt": "82c4d695310c220e5f3bcfbf5e64d315a5e826cfd77e945e61e032460014480a", +} +HUB_WEIGHTS = ( + "focalnet.onnx", + "focalnet.onnx.json", + "focalnet.pt", + "focalnet-human.onnx", + "focalnet-human.onnx.json", + "focalnet-human.pt", +) +ONNX_SIDECARS = { + "focalnet.onnx": "focalnet.onnx.json", + "focalnet-human.onnx": "focalnet-human.onnx.json", +} + + +def sha256_file(path: Path) -> str: + with path.open("rb") as handle: + return hashlib.file_digest(handle, "sha256").hexdigest() + + +def require_hash(path: Path, expected: str) -> str: + actual = sha256_file(path) + if actual != expected: + raise ValueError(f"{path} hash {actual} does not match {expected}") + return actual + + +def copy_file(source: Path, destination: Path) -> None: + if not source.is_file(): + raise FileNotFoundError(f"Missing artifact: {source}") + destination.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(source, destination) + + +def image_files(images_dir: Path) -> list[Path]: + if not images_dir.is_dir(): + raise FileNotFoundError(f"Missing image directory: {images_dir}") + return sorted( + path + for path in images_dir.iterdir() + if path.is_file() and path.suffix.lower() in IMAGE_SUFFIXES + ) + + +def staged_paths(staging: Path) -> list[str]: + return sorted(str(path.relative_to(staging)) for path in staging.rglob("*") if path.is_file()) + + +def read_onnx_metadata(path: Path) -> dict: + from onnxruntime import InferenceSession + + session = InferenceSession(str(path), providers=["CPUExecutionProvider"]) + raw = session.get_modelmeta().custom_metadata_map.get("focalnet") + if not raw: + raise ValueError(f"{path} is missing embedded FocalNet metadata") + return json.loads(raw) + + +def require_paired_hashes(card_text: str, expected: dict[str, str]) -> None: + for name, digest in expected.items(): + paired = any( + f"`{name}`" in line and f"`{digest}`" in line for line in card_text.splitlines() + ) + if not paired: + raise ValueError(f"Model card does not pair {name} with SHA-256 {digest}") + + +def require_card_images(card_text: str, pictures: list[Path]) -> None: + referenced = set(IMAGE_REFERENCE.findall(card_text)) + present = {f"images/{picture.name}" for picture in pictures} + missing = sorted(referenced - present) + extra = sorted(present - referenced) + if missing: + raise ValueError(f"Model card references missing images: {missing}") + if extra: + raise ValueError(f"Image directory has files not referenced by the model card: {extra}") + if not present: + raise ValueError("Model card and image directory contain no example images") + + +def require_sidecar( + onnx_path: Path, + sidecar: Path, + read_metadata: Callable[[Path], dict], +) -> None: + if not sidecar.is_file(): + raise FileNotFoundError(f"Missing ONNX sidecar: {sidecar}") + embedded = read_metadata(onnx_path) + recorded = json.loads(sidecar.read_text()) + if recorded != embedded: + raise ValueError(f"{sidecar} does not match embedded metadata in {onnx_path}") + + +def prepare_staging( + importance_dir: Path, + human_dir: Path, + card: Path, + staging: Path, + expected: dict[str, str] | None = None, + images_dir: Path | None = None, + read_metadata: Callable[[Path], dict] | None = None, +) -> dict[str, str]: + expected = EXPECTED_SHA256 if expected is None else expected + images_dir = DEFAULT_IMAGES if images_dir is None else images_dir + read_metadata = read_onnx_metadata if read_metadata is None else read_metadata + if not card.is_file(): + raise FileNotFoundError(f"Missing model card: {card}") + card_text = card.read_text() + require_paired_hashes(card_text, expected) + pictures = image_files(images_dir) + require_card_images(card_text, pictures) + + sources = { + "focalnet.onnx": importance_dir / "focalnet.onnx", + "focalnet.onnx.json": importance_dir / "focalnet.onnx.json", + "focalnet.pt": importance_dir / "best.pt", + "focalnet-human.onnx": human_dir / "focalnet-human.onnx", + "focalnet-human.onnx.json": human_dir / "focalnet-human.onnx.json", + "focalnet-human.pt": human_dir / "best.pt", + } + for name, digest in expected.items(): + require_hash(sources[name], digest) + for onnx_name, sidecar_name in ONNX_SIDECARS.items(): + require_sidecar(sources[onnx_name], sources[sidecar_name], read_metadata) + + if staging.exists(): + shutil.rmtree(staging) + staging.mkdir(parents=True) + copy_file(card, staging / "README.md") + for name, source in sources.items(): + copy_file(source, staging / name) + for name, digest in expected.items(): + require_hash(staging / name, digest) + for onnx_name, sidecar_name in ONNX_SIDECARS.items(): + require_sidecar(staging / onnx_name, staging / sidecar_name, read_metadata) + for picture in pictures: + copy_file(picture, staging / "images" / picture.name) + + allowed = {"README.md", *HUB_WEIGHTS, *(f"images/{picture.name}" for picture in pictures)} + missing = [name for name in allowed if not (staging / name).is_file()] + if missing: + raise FileNotFoundError(f"Staging is missing {missing}") + extra = [name for name in staged_paths(staging) if name not in allowed] + if extra: + raise ValueError(f"Unexpected staged files: {extra}") + return {name: sha256_file(staging / name) for name in allowed} + + +def upload_staging(staging: Path, repo_id: str = REPO_ID, *, token: str | None = None) -> str: + from huggingface_hub import HfApi + + missing = [name for name in ("README.md", *HUB_WEIGHTS) if not (staging / name).is_file()] + if missing: + raise FileNotFoundError(f"Staging is missing {missing}") + for name, digest in EXPECTED_SHA256.items(): + require_hash(staging / name, digest) + for onnx_name, sidecar_name in ONNX_SIDECARS.items(): + require_sidecar(staging / onnx_name, staging / sidecar_name, read_onnx_metadata) + api = HfApi(token=token) + commit = api.upload_folder( + folder_path=str(staging), + repo_id=repo_id, + repo_type="model", + commit_message="Publish FocalNet checkpoints, model card, and examples", + allow_patterns=["README.md", "images/*", *HUB_WEIGHTS], + ) + return getattr(commit, "oid", str(commit)) + + +def parser() -> argparse.ArgumentParser: + root = argparse.ArgumentParser(description="Publish verified FocalNet weights to Hugging Face") + root.add_argument("--importance-dir", type=Path, required=True) + root.add_argument("--human-dir", type=Path, required=True) + root.add_argument("--card", type=Path, default=DEFAULT_CARD) + root.add_argument("--images-dir", type=Path, default=DEFAULT_IMAGES) + root.add_argument("--staging", type=Path, required=True) + root.add_argument("--repo-id", default=REPO_ID) + root.add_argument("--upload", action="store_true") + return root + + +def main() -> None: + args = parser().parse_args() + staged = prepare_staging( + args.importance_dir, + args.human_dir, + args.card, + args.staging, + images_dir=args.images_dir, + ) + for name, digest in staged.items(): + print(f"{digest} {name}") + if not args.upload: + return + commit = upload_staging(args.staging, args.repo_id) + print(f"Uploaded to https://huggingface.co/{args.repo_id} ({commit})") + + +if __name__ == "__main__": + main() diff --git a/scripts/render-hub-examples.py b/scripts/render-hub-examples.py new file mode 100644 index 0000000..d6e9589 --- /dev/null +++ b/scripts/render-hub-examples.py @@ -0,0 +1,265 @@ +"""Render labeled FocalNet comparison figures for the Hub model card.""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import numpy as np +from PIL import Image, ImageDraw, ImageFont + +from focalnet.human_runtime import HumanCropPredictor + +REPO = Path(__file__).resolve().parents[1] +REGULAR_FONTS = ( + Path("/usr/share/fonts/truetype/macos/Inter-Medium.ttf"), + Path("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf"), + Path("/usr/share/fonts/truetype/liberation/LiberationSans-Regular.ttf"), +) +BOLD_FONTS = ( + Path("/usr/share/fonts/truetype/macos/Inter-SemiBold.ttf"), + Path("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf"), + Path("/usr/share/fonts/truetype/liberation/LiberationSans-Bold.ttf"), +) +COMPARISONS = ( + ("golden-retriever-original.jpg", 1.0, "1:1", "comparison-retriever.png", False), + ("multiple-faces.jpg", 1.0, "1:1", "comparison-faces.png", False), +) + +FONT_MEDIUM = REGULAR_FONTS[0] +FONT_SEMIBOLD = BOLD_FONTS[0] + +BG = (246, 247, 249, 255) +INK = (17, 18, 23, 255) +WHITE = (255, 255, 255, 255) +CENTER_COLOR = (232, 168, 12, 255) +FOCAL_COLOR = (240, 46, 101, 255) +PAD = 24 +GAP = 16 +LABEL_H = 40 +RADIUS = 20 +CROP_SIZE = 460 + + +def resolve_font(*candidates: Path) -> Path: + for path in candidates: + if path.is_file(): + return path + names = ", ".join(str(path) for path in candidates) + raise FileNotFoundError(f"No usable TTF font among: {names}") + + +def font(path: Path, size: int) -> ImageFont.FreeTypeFont: + return ImageFont.truetype(str(path), size) + + +def center_crop(image: Image.Image, aspect: float) -> tuple[int, int, int, int]: + width, height = image.size + if width / height >= aspect: + crop_h = height + crop_w = max(1, min(width, int(round(height * aspect)))) + else: + crop_w = width + crop_h = max(1, min(height, int(round(width / aspect)))) + left = (width - crop_w) // 2 + top = (height - crop_h) // 2 + return left, top, crop_w, crop_h + + +def take_crop(image: Image.Image, box: tuple[int, int, int, int]) -> Image.Image: + left, top, width, height = box + return image.crop((left, top, left + width, top + height)) + + +def scale_to_width(image: Image.Image, width: int) -> Image.Image: + height = max(1, round(image.height * width / image.width)) + return image.resize((width, height), Image.Resampling.LANCZOS) + + +def scale_to_fit(image: Image.Image, size: int) -> Image.Image: + fitted = image.copy() + fitted.thumbnail((size, size), Image.Resampling.LANCZOS) + canvas = Image.new("RGB", (size, size), (246, 247, 249)) + canvas.paste(fitted, ((size - fitted.width) // 2, (size - fitted.height) // 2)) + return canvas + + +def round_panel(image: Image.Image) -> Image.Image: + mask = Image.new("L", image.size, 0) + ImageDraw.Draw(mask).rounded_rectangle((0, 0, *image.size), RADIUS, fill=255) + out = Image.new("RGBA", image.size, (0, 0, 0, 0)) + out.paste(image.convert("RGBA"), mask=mask) + return out + + +def draw_box( + draw: ImageDraw.ImageDraw, + box: tuple[int, int, int, int], + color: tuple[int, ...], + width: int, +) -> None: + left, top, crop_w, crop_h = box + for inset in range(width): + draw.rounded_rectangle( + ( + left + inset, + top + inset, + left + crop_w - 1 - inset, + top + crop_h - 1 - inset, + ), + radius=max(4, 14 - inset), + outline=color, + ) + + +def draw_pill( + draw: ImageDraw.ImageDraw, + xy: tuple[int, int], + text: str, + fill: tuple[int, ...], + text_font: ImageFont.FreeTypeFont, + text_fill: tuple[int, ...] = WHITE, +) -> None: + x, y = xy + left, top, right, bottom = draw.textbbox((x, y), text, font=text_font) + pad_x, pad_y = 10, 5 + draw.rounded_rectangle( + (left - pad_x, top - pad_y, right + pad_x, bottom + pad_y), + 8, + fill=fill, + ) + draw.text((x, y), text, font=text_font, fill=text_fill) + + +def annotate_source( + image: Image.Image, + center: tuple[int, int, int, int], + focal: tuple[int, int, int, int], +) -> Image.Image: + canvas = image.convert("RGBA") + shade = Image.new("RGBA", canvas.size, (10, 12, 18, 96)) + hole = Image.new("L", canvas.size, 0) + left, top, width, height = focal + ImageDraw.Draw(hole).rounded_rectangle((left, top, left + width, top + height), 16, fill=255) + canvas = Image.composite(canvas, Image.alpha_composite(canvas, shade), hole) + overlay = Image.new("RGBA", canvas.size, (0, 0, 0, 0)) + draw = ImageDraw.Draw(overlay) + draw_box(draw, center, CENTER_COLOR, 5) + draw_box(draw, focal, FOCAL_COLOR, 7) + caption = font(FONT_SEMIBOLD, max(20, image.width // 36)) + draw_pill(draw, (focal[0] + 18, focal[1] + 16), "FocalNet", FOCAL_COLOR, caption) + draw_pill(draw, (center[0] + 18, center[1] + 16), "Center", CENTER_COLOR, caption, INK) + return Image.alpha_composite(canvas, overlay) + + +def colorize_heatmap(heatmap: np.ndarray, size: tuple[int, int]) -> Image.Image: + values = np.clip(heatmap.astype(np.float32), 0, 1) + values = (values - values.min()) / (float(np.ptp(values)) + 1e-6) + heat = ( + np.asarray( + Image.fromarray((values * 255).astype(np.uint8), "L").resize( + size, Image.Resampling.BILINEAR + ), + dtype=np.float32, + ) + / 255.0 + ) + rgb = np.zeros((*heat.shape, 4), dtype=np.uint8) + rgb[..., 0] = np.clip(80 + 175 * heat + 80 * np.power(heat, 4), 0, 255).astype(np.uint8) + rgb[..., 1] = np.clip(20 + 40 * heat + 180 * np.power(heat, 3), 0, 255).astype(np.uint8) + rgb[..., 2] = np.clip(60 + 40 * (1 - heat) + 80 * heat, 0, 255).astype(np.uint8) + rgb[..., 3] = np.clip(50 + 160 * heat, 0, 255).astype(np.uint8) + return Image.fromarray(rgb, "RGBA") + + +def labeled_panel(image: Image.Image, title: str) -> Image.Image: + photo = round_panel(image) + canvas = Image.new("RGBA", (photo.width, photo.height + LABEL_H), (0, 0, 0, 0)) + canvas.paste(photo, (0, 0), photo) + draw = ImageDraw.Draw(canvas) + draw.text((6, photo.height + 10), title, font=font(FONT_MEDIUM, 22), fill=INK) + return canvas + + +def compose_story(source: Image.Image, left: Image.Image, right: Image.Image) -> Image.Image: + width = PAD * 2 + max(source.width, left.width + GAP + right.width) + height = PAD * 2 + source.height + GAP + max(left.height, right.height) + canvas = Image.new("RGBA", (width, height), BG) + canvas.paste(source, ((width - source.width) // 2, PAD), source) + y = PAD + source.height + GAP + x = (width - (left.width + GAP + right.width)) // 2 + canvas.paste(left, (x, y), left) + canvas.paste(right, (x + left.width + GAP, y), right) + return canvas.convert("RGB") + + +def render_comparison( + source: Path, + human: HumanCropPredictor, + aspect: float, + aspect_label: str, + output: Path, + *, + heatmap: bool = False, +) -> None: + image = Image.open(source).convert("RGB") + result, importance = human.predict(source, aspect) + crop = result["crop"] + focal = (crop["left"], crop["top"], crop["width"], crop["height"]) + center = center_crop(image, aspect) + annotated = annotate_source(image, center, focal) + if heatmap: + annotated = Image.alpha_composite( + annotated.convert("RGBA"), colorize_heatmap(importance, image.size) + ) + source_w = CROP_SIZE * 2 + GAP + source_panel = labeled_panel(scale_to_width(annotated, source_w), "Source") + center_panel = labeled_panel( + scale_to_fit(take_crop(image, center), CROP_SIZE), f"Center {aspect_label}" + ) + focal_panel = labeled_panel( + scale_to_fit(take_crop(image, focal), CROP_SIZE), + f"FocalNet {aspect_label}", + ) + figure = compose_story(source_panel, center_panel, focal_panel) + output.parent.mkdir(parents=True, exist_ok=True) + figure.save(output, "PNG", optimize=True) + print(f"wrote {output} {figure.size} {output.stat().st_size / 1024:.0f} KiB") + + +def parser() -> argparse.ArgumentParser: + root = argparse.ArgumentParser(description="Render Hub model-card comparison figures") + root.add_argument("--model", type=Path, required=True, help="Path to focalnet-human.onnx") + root.add_argument("--photos-dir", type=Path, required=True, help="Directory of source JPEGs") + root.add_argument("--output-dir", type=Path, default=REPO / "docs" / "hub-images") + root.add_argument("--font", type=Path, help="Regular TTF for panel labels") + root.add_argument("--font-bold", type=Path, help="Bold TTF for crop captions") + return root + + +def main() -> None: + global FONT_MEDIUM, FONT_SEMIBOLD + args = parser().parse_args() + FONT_MEDIUM = args.font if args.font is not None else resolve_font(*REGULAR_FONTS) + FONT_SEMIBOLD = args.font_bold if args.font_bold is not None else resolve_font(*BOLD_FONTS) + if not FONT_MEDIUM.is_file(): + raise FileNotFoundError(f"Missing regular font: {FONT_MEDIUM}") + if not FONT_SEMIBOLD.is_file(): + raise FileNotFoundError(f"Missing bold font: {FONT_SEMIBOLD}") + human = HumanCropPredictor(args.model) + for name, aspect, label, output_name, heat in COMPARISONS: + source = args.photos_dir / name + if not source.is_file(): + raise FileNotFoundError(f"Missing source photograph: {source}") + render_comparison( + source, + human, + aspect, + label, + args.output_dir / output_name, + heatmap=heat, + ) + + +if __name__ == "__main__": + main() diff --git a/tests/test_publish_huggingface.py b/tests/test_publish_huggingface.py new file mode 100644 index 0000000..d4f4b7a --- /dev/null +++ b/tests/test_publish_huggingface.py @@ -0,0 +1,210 @@ +import importlib.util +import json +from pathlib import Path + +import pytest + +REPO = Path(__file__).resolve().parents[1] +SCRIPT = REPO / "scripts" / "publish-huggingface.py" +RENDER = REPO / "scripts" / "render-hub-examples.py" + + +def load_script(path: Path, name: str): + spec = importlib.util.spec_from_file_location(name, path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def load_publisher(): + return load_script(SCRIPT, "publish_huggingface") + + +def write_release_inputs(tmp_path: Path) -> dict: + publisher = load_publisher() + importance = tmp_path / "importance" + human = tmp_path / "human" + pictures = tmp_path / "images" + importance.mkdir() + human.mkdir() + pictures.mkdir() + files = { + importance / "focalnet.onnx": b"onnx-v1", + importance / "best.pt": b"importance-pt", + human / "focalnet-human.onnx": b"onnx-v2", + human / "best.pt": b"human-pt", + } + for path, payload in files.items(): + path.write_bytes(payload) + embedded = { + "focalnet.onnx": {"format_version": 1, "task": "importance"}, + "focalnet-human.onnx": {"format_version": 2, "task": "human-crop-ranking"}, + } + (importance / "focalnet.onnx.json").write_text(json.dumps(embedded["focalnet.onnx"])) + (human / "focalnet-human.onnx.json").write_text(json.dumps(embedded["focalnet-human.onnx"])) + (pictures / "example.jpg").write_bytes(b"jpeg-bytes") + expected = { + "focalnet.onnx": publisher.sha256_file(importance / "focalnet.onnx"), + "focalnet.pt": publisher.sha256_file(importance / "best.pt"), + "focalnet-human.onnx": publisher.sha256_file(human / "focalnet-human.onnx"), + "focalnet-human.pt": publisher.sha256_file(human / "best.pt"), + } + card = tmp_path / "card.md" + card.write_text( + "\n".join( + [ + *(f"`{name}` `{digest}`" for name, digest in expected.items()), + "![example](images/example.jpg)", + ] + ) + ) + return { + "publisher": publisher, + "importance": importance, + "human": human, + "pictures": pictures, + "card": card, + "expected": expected, + "embedded": embedded, + "read_metadata": lambda path: json.loads(json.dumps(embedded[path.name])), + } + + +def test_committed_card_stays_aligned_with_publisher_and_images(): + publisher = load_publisher() + card = (REPO / "docs" / "huggingface-model-card.md").read_text() + assert card.startswith("---\n") + for field in ( + "license: mit", + "library_name: onnxruntime", + "pipeline_tag: image-to-image", + "base_model: timm/repvit_m0_9", + ): + assert field in card + assert "Microsoft FocalNet" not in card + publisher.require_paired_hashes(card, publisher.EXPECTED_SHA256) + publisher.require_card_images(card, publisher.image_files(REPO / "docs" / "hub-images")) + + +def test_docs_and_fetch_script_keep_the_publisher_hashes(): + publisher = load_publisher() + docs = (REPO / "docs" / "results-500k.md").read_text() + ( + REPO / "docs" / "results-cpc-gaic-v2.md" + ).read_text() + fetch = (REPO / "web" / "scripts" / "fetch-model.ts").read_text() + assert publisher.EXPECTED_SHA256["focalnet.onnx"] in docs + assert publisher.EXPECTED_SHA256["focalnet.pt"] in docs + assert publisher.EXPECTED_SHA256["focalnet-human.onnx"] in docs + assert publisher.EXPECTED_SHA256["focalnet-human.pt"] in docs + assert publisher.EXPECTED_SHA256["focalnet-human.onnx"] in fetch + assert "huggingface.co/appwrite/focalnet" in fetch + + +def test_prepare_staging_renames_checkpoints_and_rejects_hash_mismatch(tmp_path): + setup = write_release_inputs(tmp_path) + publisher = setup["publisher"] + staging = tmp_path / "hub" + staged = publisher.prepare_staging( + setup["importance"], + setup["human"], + setup["card"], + staging, + expected=setup["expected"], + images_dir=setup["pictures"], + read_metadata=setup["read_metadata"], + ) + assert (staging / "focalnet.pt").read_bytes() == b"importance-pt" + assert (staging / "focalnet-human.pt").read_bytes() == b"human-pt" + assert (staging / "images" / "example.jpg").read_bytes() == b"jpeg-bytes" + assert staged["focalnet.onnx"] == setup["expected"]["focalnet.onnx"] + assert staged["focalnet-human.pt"] == setup["expected"]["focalnet-human.pt"] + assert staged["images/example.jpg"] == publisher.sha256_file(setup["pictures"] / "example.jpg") + + (setup["human"] / "focalnet-human.onnx").write_bytes(b"tampered") + with pytest.raises(ValueError, match="does not match"): + publisher.prepare_staging( + setup["importance"], + setup["human"], + setup["card"], + staging, + expected=setup["expected"], + images_dir=setup["pictures"], + read_metadata=setup["read_metadata"], + ) + + +def test_prepare_staging_rejects_swapped_card_hashes(tmp_path): + setup = write_release_inputs(tmp_path) + names = list(setup["expected"]) + swapped = dict(setup["expected"]) + swapped[names[0]], swapped[names[2]] = swapped[names[2]], swapped[names[0]] + setup["card"].write_text( + "\n".join( + [ + *(f"`{name}` `{digest}`" for name, digest in swapped.items()), + "![example](images/example.jpg)", + ] + ) + ) + with pytest.raises(ValueError, match="does not pair"): + setup["publisher"].prepare_staging( + setup["importance"], + setup["human"], + setup["card"], + tmp_path / "hub", + expected=setup["expected"], + images_dir=setup["pictures"], + read_metadata=setup["read_metadata"], + ) + + +def test_prepare_staging_rejects_missing_images_and_stale_sidecars(tmp_path): + setup = write_release_inputs(tmp_path) + publisher = setup["publisher"] + staging = tmp_path / "hub" + with pytest.raises(FileNotFoundError, match="Missing image directory"): + publisher.prepare_staging( + setup["importance"], + setup["human"], + setup["card"], + staging, + expected=setup["expected"], + images_dir=tmp_path / "absent-images", + read_metadata=setup["read_metadata"], + ) + + (setup["pictures"] / "orphan.png").write_bytes(b"png") + with pytest.raises(ValueError, match="not referenced"): + publisher.prepare_staging( + setup["importance"], + setup["human"], + setup["card"], + staging, + expected=setup["expected"], + images_dir=setup["pictures"], + read_metadata=setup["read_metadata"], + ) + (setup["pictures"] / "orphan.png").unlink() + + (setup["human"] / "focalnet-human.onnx.json").write_text('{"format_version": 1}') + with pytest.raises(ValueError, match="does not match embedded metadata"): + publisher.prepare_staging( + setup["importance"], + setup["human"], + setup["card"], + staging, + expected=setup["expected"], + images_dir=setup["pictures"], + read_metadata=setup["read_metadata"], + ) + + +def test_renderer_uses_a_font_from_the_environment(): + renderer = load_script(RENDER, "render_hub_examples") + regular = renderer.resolve_font(*renderer.REGULAR_FONTS) + bold = renderer.resolve_font(*renderer.BOLD_FONTS) + assert regular.is_file() + assert bold.is_file() + with pytest.raises(FileNotFoundError, match="No usable TTF"): + renderer.resolve_font(Path("/tmp/missing-focalnet-font.ttf")) diff --git a/web/scripts/fetch-model.ts b/web/scripts/fetch-model.ts index 627b837..8a8bdfe 100644 --- a/web/scripts/fetch-model.ts +++ b/web/scripts/fetch-model.ts @@ -3,83 +3,120 @@ import { mkdir, unlink } from "node:fs/promises"; import { tmpdir } from "node:os"; import { join } from "node:path"; +export const HUB_REPO = "appwrite/focalnet"; +export const DEFAULT_ONNX_URL = + "https://huggingface.co/appwrite/focalnet/resolve/main/focalnet-human.onnx"; +export const DEFAULT_ONNX_SHA256 = + "59164c601c98cea3f62b25166710831dac63e1a872fc64767c65316ad5385439"; + const root = join(import.meta.dir, "..", ".."); const dest = join(root, "web", "public", "focalnet-human.onnx"); -const expected = Bun.env.FOCALNET_ONNX_SHA256; -const modelUrl = Bun.env.FOCALNET_ONNX_URL; -const allowDummy = Bun.env.FOCALNET_ALLOW_DUMMY === "1"; -await mkdir(join(dest, ".."), { recursive: true }); +export type FetchModelEnv = { + FOCALNET_ONNX_SHA256?: string; + FOCALNET_ONNX_URL?: string; + FOCALNET_ONNX?: string; + FOCALNET_ALLOW_DUMMY?: string; +}; -async function sha256(path: string): Promise { - const hasher = new Bun.CryptoHasher("sha256"); - hasher.update(await Bun.file(path).arrayBuffer()); - return hasher.digest("hex"); +export function resolvedHash(value: string | undefined): string { + const trimmed = value?.trim() ?? ""; + return trimmed || DEFAULT_ONNX_SHA256; } -async function installModel(source: string): Promise { - await Bun.write(dest, Bun.file(source)); - const actual = await sha256(dest); - if (expected && actual !== expected) { - await unlink(dest); - throw new Error(`Model hash ${actual} does not match ${expected}`); - } - console.log(`Wrote ${dest} (${actual})`); -} +export async function fetchModel( + env: FetchModelEnv = Bun.env, + destPath = dest, +): Promise { + const expected = resolvedHash(env.FOCALNET_ONNX_SHA256); + const modelUrl = env.FOCALNET_ONNX_URL?.trim() || DEFAULT_ONNX_URL; + const allowDummy = env.FOCALNET_ALLOW_DUMMY === "1"; + + await mkdir(join(destPath, ".."), { recursive: true }); -if (await Bun.file(dest).exists()) { - const actual = await sha256(dest); - if (!expected || actual === expected || allowDummy) { - console.log(`Using existing model at ${dest} (${actual})`); - process.exit(0); + async function sha256(path: string): Promise { + const hasher = new Bun.CryptoHasher("sha256"); + hasher.update(await Bun.file(path).arrayBuffer()); + return hasher.digest("hex"); } - console.warn(`Existing model hash ${actual} does not match ${expected}; replacing`); - await unlink(dest); -} -if (Bun.env.FOCALNET_ONNX) { - await installModel(Bun.env.FOCALNET_ONNX); - process.exit(0); -} + async function installModel(source: string): Promise { + await Bun.write(destPath, Bun.file(source)); + const actual = await sha256(destPath); + if (actual !== expected) { + await unlink(destPath); + throw new Error(`Model hash ${actual} does not match ${expected}`); + } + console.log(`Wrote ${destPath} (${actual})`); + return actual; + } -if (modelUrl) { - const tmp = join(tmpdir(), "focalnet-human.onnx.download"); - let installed = false; - try { - const response = await fetch(modelUrl); - if (!response.ok) { - throw new Error(`GET ${modelUrl} → ${response.status}`); + if (await Bun.file(destPath).exists()) { + const actual = await sha256(destPath); + if (actual === expected) { + console.log(`Using existing model at ${destPath} (${actual})`); + return actual; } - await Bun.write(tmp, response); - await installModel(tmp); - installed = true; - } catch (error) { - const detail = error instanceof Error ? error.message : String(error); - if (!allowDummy) { - throw new Error(detail); + if (allowDummy) { + console.log(`Using existing model at ${destPath} (${actual})`); + return actual; } - console.warn(detail); + console.warn(`Existing model hash ${actual} does not match ${expected}; replacing`); + await unlink(destPath); } - if (installed) { - await unlink(tmp).catch((error) => { - console.warn(error instanceof Error ? error.message : error); - }); - process.exit(0); + + if (env.FOCALNET_ONNX) { + return installModel(env.FOCALNET_ONNX); + } + + if (modelUrl) { + const tmp = join(tmpdir(), `focalnet-human-${crypto.randomUUID()}.onnx.download`); + let installed = false; + try { + const response = await fetch(modelUrl); + if (!response.ok) { + throw new Error(`GET ${modelUrl} → ${response.status}`); + } + await Bun.write(tmp, response); + const actual = await installModel(tmp); + installed = true; + await unlink(tmp).catch((error) => { + console.warn(error instanceof Error ? error.message : error); + }); + return actual; + } catch (error) { + const detail = error instanceof Error ? error.message : String(error); + if (installed) { + throw new Error(detail); + } + if (!allowDummy) { + throw new Error(detail); + } + console.warn(detail); + } finally { + await unlink(tmp).catch(() => undefined); + } + } + + if (!allowDummy) { + throw new Error( + "No local checkpoint. bun run fetch-model downloads the published Hub " + + "weights, or set FOCALNET_ONNX to a local file. FOCALNET_ALLOW_DUMMY=1 " + + "writes a UI placeholder.", + ); } -} -if (!allowDummy) { - throw new Error( - "No local checkpoint. Export a model, then set FOCALNET_ONNX to that file " + - "(optional FOCALNET_ONNX_SHA256 to verify). FOCALNET_ALLOW_DUMMY=1 writes a UI placeholder.", + const dummy = Bun.spawn( + ["uv", "run", "--with", "onnx", "python", join(root, "web/scripts/export-dummy-onnx.py")], + { cwd: root, stdout: "inherit", stderr: "inherit" }, ); + if ((await dummy.exited) !== 0) { + throw new Error("Dummy ONNX export failed"); + } + console.log("Wrote dummy ONNX for local UI testing; do not publish this file"); + return sha256(destPath); } -const dummy = Bun.spawn( - ["uv", "run", "--with", "onnx", "python", join(root, "web/scripts/export-dummy-onnx.py")], - { cwd: root, stdout: "inherit", stderr: "inherit" }, -); -if ((await dummy.exited) !== 0) { - throw new Error("Dummy ONNX export failed"); +if (import.meta.main) { + await fetchModel(); } -console.log("Wrote dummy ONNX for local UI testing; do not publish this file"); diff --git a/web/src/fetch-model.test.ts b/web/src/fetch-model.test.ts new file mode 100644 index 0000000..c9c8643 --- /dev/null +++ b/web/src/fetch-model.test.ts @@ -0,0 +1,118 @@ +import { mkdtemp, rm, writeFile } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { afterEach, expect, test } from "bun:test"; +import { DEFAULT_ONNX_SHA256, fetchModel, resolvedHash } from "../scripts/fetch-model"; + +const dirs: string[] = []; +const originalFetch = globalThis.fetch; + +afterEach(async () => { + globalThis.fetch = originalFetch; + await Promise.all(dirs.splice(0).map((dir) => rm(dir, { recursive: true, force: true }))); +}); + +async function tempDest(): Promise { + const dir = await mkdtemp(join(tmpdir(), "focalnet-fetch-")); + dirs.push(dir); + return join(dir, "focalnet-human.onnx"); +} + +async function digest(payload: string | Uint8Array): Promise { + const hasher = new Bun.CryptoHasher("sha256"); + hasher.update(payload); + return hasher.digest("hex"); +} + +test("empty SHA256 still verifies against the published default", () => { + expect(resolvedHash(undefined)).toBe(DEFAULT_ONNX_SHA256); + expect(resolvedHash("")).toBe(DEFAULT_ONNX_SHA256); + expect(resolvedHash(" ")).toBe(DEFAULT_ONNX_SHA256); + const override = "a".repeat(64); + expect(resolvedHash(` ${override} `)).toBe(override); +}); + +test("installs a local checkpoint whose hash matches", async () => { + const dest = await tempDest(); + const source = join(dest, "..", "source.onnx"); + const payload = "matching-checkpoint"; + await writeFile(source, payload); + const expected = await digest(payload); + const actual = await fetchModel( + { FOCALNET_ONNX: source, FOCALNET_ONNX_SHA256: expected }, + dest, + ); + expect(actual).toBe(expected); + expect(await Bun.file(dest).text()).toBe(payload); +}); + +test("rejects a local checkpoint when SHA256 is empty and the default does not match", async () => { + const dest = await tempDest(); + const source = join(dest, "..", "source.onnx"); + await writeFile(source, "not-the-published-model"); + await expect(fetchModel({ FOCALNET_ONNX: source, FOCALNET_ONNX_SHA256: "" }, dest)).rejects.toThrow( + /does not match/, + ); + expect(await Bun.file(dest).exists()).toBe(false); +}); + +test("downloads a matching checkpoint and rejects a mismatched one", async () => { + const dest = await tempDest(); + const payload = "downloaded-checkpoint"; + const expected = await digest(payload); + globalThis.fetch = async () => new Response(payload, { status: 200 }); + const actual = await fetchModel( + { + FOCALNET_ONNX_URL: "https://example.test/focalnet-human.onnx", + FOCALNET_ONNX_SHA256: expected, + }, + dest, + ); + expect(actual).toBe(expected); + expect(await Bun.file(dest).text()).toBe(payload); + + const other = await tempDest(); + globalThis.fetch = async () => new Response("tampered-checkpoint", { status: 200 }); + await expect( + fetchModel( + { + FOCALNET_ONNX_URL: "https://example.test/focalnet-human.onnx", + FOCALNET_ONNX_SHA256: expected, + }, + other, + ), + ).rejects.toThrow(/does not match/); + expect(await Bun.file(other).exists()).toBe(false); +}); + +test("reuses an existing matching file and keeps a dummy when allowed", async () => { + const dest = await tempDest(); + const payload = "already-fetched"; + await writeFile(dest, payload); + const expected = await digest(payload); + globalThis.fetch = async () => { + throw new Error("should not download when the existing hash matches"); + }; + expect( + await fetchModel( + { + FOCALNET_ONNX_URL: "https://example.test/focalnet-human.onnx", + FOCALNET_ONNX_SHA256: expected, + }, + dest, + ), + ).toBe(expected); + + const dummy = await tempDest(); + await writeFile(dummy, "ui-placeholder"); + expect( + await fetchModel( + { + FOCALNET_ALLOW_DUMMY: "1", + FOCALNET_ONNX_SHA256: expected, + }, + dummy, + ), + ).not.toBe(expected); + expect(await Bun.file(dummy).text()).toBe("ui-placeholder"); +}); diff --git a/web/src/main.ts b/web/src/main.ts index 0df84cd..c6db6bc 100644 --- a/web/src/main.ts +++ b/web/src/main.ts @@ -65,8 +65,8 @@ app.innerHTML = `

- The ranking head is trained on CPC and GAICD judgments. This lab does not - ship model weights; load a local ONNX with bun run fetch-model. + The ranking head is trained on CPC and GAICD judgments. Weights download + from Hugging Face via bun run fetch-model; photos stay in this tab.

`;