lhk229
da4164e95a
Add bgfilter/weights.py: resolve_model_source() maps a HuggingFace repo id to a local folder under the weights dir (models/ by default, override with BGFILTER_WEIGHTS_DIR) when one named after the repo basename exists; otherwise the repo id is returned unchanged. Opt-in and backward compatible. - vitmatte_infer.py / segmentation.py resolve model_name through it; anime-seg reads <dir>/isnetis.onnx directly instead of hf_hub_download when local. - .gitignore: models/, model-cache/ - README: document bundling weights as plain project folders. Verified: anime-seg loads from a local models/anime-seg/ folder offline. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
30 lines
1000 B
Python
30 lines
1000 B
Python
from __future__ import annotations
|
|
|
|
import os
|
|
from pathlib import Path
|
|
|
|
|
|
def weights_dir() -> Path:
|
|
"""Directory scanned for locally-bundled model weights.
|
|
|
|
Defaults to ``models`` (relative to the working directory); override with the
|
|
``BGFILTER_WEIGHTS_DIR`` environment variable.
|
|
"""
|
|
return Path(os.environ.get("BGFILTER_WEIGHTS_DIR", "models"))
|
|
|
|
|
|
def resolve_model_source(repo_id: str) -> str:
|
|
"""Map a HuggingFace repo id to a local weights folder when one is present.
|
|
|
|
If ``<weights_dir>/<basename>`` exists it is returned (so models can be shipped
|
|
as plain folders inside the project, e.g. ``ZhengPeng7/BiRefNet`` ->
|
|
``models/BiRefNet``); otherwise the repo id is returned unchanged so it loads
|
|
from the HF hub/cache as before. An absolute/existing path is passed through.
|
|
"""
|
|
if os.path.isdir(repo_id):
|
|
return repo_id
|
|
local = weights_dir() / repo_id.split("/")[-1]
|
|
if local.is_dir():
|
|
return str(local)
|
|
return repo_id
|