Support loading model weights from local project folders
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>
This commit is contained in:
@@ -0,0 +1,29 @@
|
||||
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
|
||||
Reference in New Issue
Block a user