diff --git a/DEPLOY.md b/DEPLOY.md index 2c2ef81..2d42c2a 100644 --- a/DEPLOY.md +++ b/DEPLOY.md @@ -81,7 +81,7 @@ print('torch', torch.__version__, 'cuda=', torch.version.cuda)" Three models are used by default: `hustvl/vitmatte-base-composition-1k` (matting), `ZhengPeng7/BiRefNet` (default segmenter), `skytnt/anime-seg` (optional segmenter). -A fourth model, `ZhengPeng7/BiRefNet_HR-matting` (~0.9 GB), is only needed if you +A fourth model, `ZhengPeng7/BiRefNet_HR-matting` (~425 MB), is only needed if you enable the **cross-check** veto — see the note at the end of this section. Pick **one** of the two layouts below. @@ -127,7 +127,7 @@ finds it (`ZhengPeng7/BiRefNet` → `models/BiRefNet`). Required files per model ├── config.json ├── BiRefNet_config.py # custom code — required by trust_remote_code ├── birefnet.py # custom code — required by trust_remote_code - └── model.safetensors # ~0.9 GB + └── model.safetensors # ~425 MB ``` Notes: @@ -155,7 +155,7 @@ Override the base directory with `BGFILTER_WEIGHTS_DIR`. `models/` is gitignored ### Cross-check veto (optional, off by default in the service) The pipeline can run a second, high-resolution matting model -(`ZhengPeng7/BiRefNet_HR-matting`, ~0.9 GB) as a "second opinion" that only clears +(`ZhengPeng7/BiRefNet_HR-matting`, ~425 MB) as a "second opinion" that only clears background residue trapped between hair strands. In the **HTTP service it is OFF by default** and lazy-loaded, so you don't need this model unless you turn it on: diff --git a/DEPLOY_ZH.md b/DEPLOY_ZH.md index ec43680..238c08f 100644 --- a/DEPLOY_ZH.md +++ b/DEPLOY_ZH.md @@ -75,7 +75,7 @@ print('torch', torch.__version__, 'cuda=', torch.version.cuda)" 默认用到三个模型：`hustvl/vitmatte-base-composition-1k`（抠图）、 `ZhengPeng7/BiRefNet`（默认分割）、`skytnt/anime-seg`（可选分割）。 -第四个模型 `ZhengPeng7/BiRefNet_HR-matting`（约 0.9 GB）只有在启用 **cross-check（交叉校验） +第四个模型 `ZhengPeng7/BiRefNet_HR-matting`（约 425 MB）只有在启用 **cross-check（交叉校验） 否决**时才需要——见本节末尾说明。下面两种放法**二选一**。 ### 方案 A —— HuggingFace 缓存 + 离线模式（推荐） @@ -119,7 +119,7 @@ hf download ZhengPeng7/BiRefNet_HR-matting # 仅在使用 cross-check 时（ ├── config.json ├── BiRefNet_config.py # 自定义代码——trust_remote_code 必需 ├── birefnet.py # 自定义代码——trust_remote_code 必需 - └── model.safetensors # 约 0.9 GB + └── model.safetensors # 约 425 MB ``` 说明： @@ -145,7 +145,7 @@ hf download ZhengPeng7/BiRefNet_HR-matting --local-dir models/BiRefNet_HR- ### cross-check 否决（可选，服务里默认关闭） -管线可以再跑一个高分辨率 matting 模型（`ZhengPeng7/BiRefNet_HR-matting`，约 0.9 GB） +管线可以再跑一个高分辨率 matting 模型（`ZhengPeng7/BiRefNet_HR-matting`，约 425 MB） 作为"第二意见"，专门清理发丝间残留的背景色。在 **HTTP 服务里默认是关的**、且懒加载， 所以不开启就不需要这个模型： diff --git a/scripts/fetch_weights.py b/scripts/fetch_weights.py new file mode 100644 index 0000000..50b44e7 --- /dev/null +++ b/scripts/fetch_weights.py @@ -0,0 +1,116 @@ +#!/usr/bin/env python3 +"""Download every model weight into the local ``models/`` folder via the HF mirror. + +Run once, no arguments. Afterwards the service and CLI have **full functionality** +offline: ViTMatte (matting), both segmenters (BiRefNet + anime-seg), and the +cross-check second-opinion model are all present. + + python scripts/fetch_weights.py + +Weights go to ``BGFILTER_WEIGHTS_DIR`` (default ``/models``), one plain folder +per model named after the repo basename -- exactly what the loader looks for. +""" + +from __future__ import annotations + +import os +import sys +import time +from pathlib import Path + +# Route downloads through the domestic mirror. The mirror is reached DIRECTLY, so +# any local HTTP proxy must be bypassed -- with a proxy the mirror redirects the +# blob back to huggingface.co and newer huggingface_hub rejects it (FileMetadataError). +# We therefore force the bypass (clear proxy vars + NO_PROXY=*) rather than setdefault, +# because a pre-existing NO_PROXY that omits the CDN would leave the proxy in play. +# Xet is disabled because these repos are Xet-backed and that path fails here. +for _proxy_var in ("HTTP_PROXY", "HTTPS_PROXY", "http_proxy", "https_proxy", "ALL_PROXY", "all_proxy"): + os.environ.pop(_proxy_var, None) +os.environ["NO_PROXY"] = "*" +os.environ["no_proxy"] = "*" +os.environ["HF_HUB_DISABLE_XET"] = "1" +os.environ.setdefault("HF_ENDPOINT", "https://hf-mirror.com") +os.environ.setdefault("HF_HUB_DISABLE_SYMLINKS_WARNING", "1") + +REPO_ROOT = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(REPO_ROOT)) + +from huggingface_hub import snapshot_download # noqa: E402 + +from bgfilter.config import _SEG_MODELS, load_settings # noqa: E402 +from bgfilter.weights import weights_dir # noqa: E402 + + +_WEIGHT_SUFFIXES = (".safetensors", ".bin", ".onnx") + + +def _dir_size_mb(path: Path) -> float: + return sum(f.stat().st_size for f in path.rglob("*") if f.is_file()) / 1e6 + + +def _already_present(dest: Path) -> bool: + """True once a weight blob is on disk -- treat the model as fully downloaded. + + The weight file is the large, slow payload; if it is present the small config / + code files are too. A folder without any weight file is re-fetched (and + snapshot_download resumes whatever is missing). Delete the folder to force a + fresh download. + """ + return dest.is_dir() and any( + f.is_file() and f.suffix.lower() in _WEIGHT_SUFFIXES for f in dest.iterdir() + ) + + +def _download(repo_id: str, dest: Path, attempts: int = 5) -> None: + for attempt in range(1, attempts + 1): + try: + snapshot_download(repo_id, local_dir=str(dest)) + return + except Exception as exc: # noqa: BLE001 -- retry any transient failure + if attempt == attempts: + raise + print(f" attempt {attempt} failed ({type(exc).__name__}); retrying...", flush=True) + time.sleep(3) + + +def main() -> int: + cfg = REPO_ROOT / "configs" / "default.yaml" + settings = load_settings(str(cfg) if cfg.exists() else None) + + # Every model the app can use, read from config/code so it stays in sync. + repo_ids: list[str] = [] + for rid in ( + settings.model.model_name, # ViTMatte matting model + *_SEG_MODELS.values(), # BiRefNet + anime-seg segmenters + settings.cross_check.model_name, # cross-check second-opinion model + ): + if rid and rid not in repo_ids: + repo_ids.append(rid) + + base = weights_dir() + if not base.is_absolute(): + base = REPO_ROOT / base + base.mkdir(parents=True, exist_ok=True) + + print(f"endpoint : {os.environ['HF_ENDPOINT']}") + print(f"target : {base}") + print(f"models : {len(repo_ids)}") + downloaded = skipped = 0 + for repo_id in repo_ids: + dest = base / repo_id.split("/")[-1] + if _already_present(dest): + print(f"\n=== {repo_id} -> already present ({_dir_size_mb(dest):.0f} MB), skip", flush=True) + skipped += 1 + continue + print(f"\n>>> {repo_id} -> {dest}", flush=True) + _download(repo_id, dest) + print(f" OK ({_dir_size_mb(dest):.0f} MB)", flush=True) + downloaded += 1 + + print(f"\nDone: {downloaded} downloaded, {skipped} already present.") + print("All weights ready -- the app now has full functionality offline.") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main())