Files
FabledCurator/backend/app/scripts/download_models.py
T
bvandeusen 3b3e7565fb fix(ml): align tagger + downloader with Camie v2 actual layout (model.onnx -> camie-tagger-v2.onnx + JSON metadata + ImageNet preprocessing + sigmoid on refined output)
The HF repo Camais03/camie-tagger-v2 has camie-tagger-v2.onnx (789 MB)
+ camie-tagger-v2-metadata.json (7.77 MB) at root, NOT model.onnx +
selected_tags.csv. Tags ship as nested JSON (dataset_info.tag_mapping)
not CSV. Per the published onnx_inference.py reference: input is NCHW
not NHWC, normalize with ImageNet mean/std, pad-square color (124,116,
104), sigmoid the second output (refined predictions) not the first.

Operator hit this during the IR migration ML backfill — download_models
silently fetched only 3 json files (allow_patterns matched nothing
useful), tagger.load() then raised RuntimeError. Fetched the actual
v2 layout via WebFetch, rewrote tagger to match published reference.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-05-25 02:25:30 -04:00

73 lines
2.2 KiB
Python

"""Self-heal model weights into /models. Idempotent — only missing files
are fetched. Called by the ml-worker entrypoint before Celery starts.
"""
import os
import sys
from pathlib import Path
MODEL_ROOT = Path(os.environ.get("ML_MODEL_DIR", "/models"))
CAMIE_REPO = os.environ.get("CAMIE_HF_REPO", "Camais03/camie-tagger-v2")
SIGLIP_REPO = os.environ.get(
"SIGLIP_HF_REPO", "google/siglip-so400m-patch14-384"
)
def _snapshot(repo_id: str, dest: Path, allow_patterns: list[str] | None) -> None:
from huggingface_hub import snapshot_download
dest.mkdir(parents=True, exist_ok=True)
snapshot_download(
repo_id=repo_id,
local_dir=str(dest),
allow_patterns=allow_patterns,
)
def ensure_camie() -> None:
"""Fetch Camie v2 weights + metadata.
v2 layout (HuggingFace Camais03/camie-tagger-v2): the ONNX file is
named camie-tagger-v2.onnx (not model.onnx) and tags ship inside
camie-tagger-v2-metadata.json (not selected_tags.csv). Both at root.
The repo also contains app/, game/, training/, images/ subdirs full
of setup/demo files we don't need — allow_patterns scopes the fetch
to just the inference essentials (~790 MB instead of ~2 GB).
"""
dest = MODEL_ROOT / "camie"
model_file = dest / "camie-tagger-v2.onnx"
meta_file = dest / "camie-tagger-v2-metadata.json"
if model_file.is_file() and meta_file.is_file():
print(f"[download_models] Camie present at {dest}")
return
print(f"[download_models] Fetching {CAMIE_REPO} -> {dest}")
_snapshot(
CAMIE_REPO, dest,
[
"camie-tagger-v2.onnx",
"camie-tagger-v2-metadata.json",
"config.json",
"config.yaml",
],
)
def ensure_siglip() -> None:
dest = MODEL_ROOT / "siglip"
if (dest / "config.json").is_file() and any(dest.glob("*.safetensors")):
print(f"[download_models] SigLIP present at {dest}")
return
print(f"[download_models] Fetching {SIGLIP_REPO} -> {dest}")
_snapshot(SIGLIP_REPO, dest, None)
def main() -> int:
ensure_camie()
ensure_siglip()
print("[download_models] Done.")
return 0
if __name__ == "__main__":
sys.exit(main())