From da4164e95aa36b4efbe1afff3db9da2ff0b2c6fe Mon Sep 17 00:00:00 2001 From: Linghong Date: Fri, 3 Jul 2026 02:50:19 +0000 Subject: [PATCH] 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 /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 --- .gitignore | 2 ++ README.md | 7 +++++++ bgfilter/segmentation.py | 11 +++++++++-- bgfilter/vitmatte_infer.py | 6 ++++-- bgfilter/weights.py | 29 +++++++++++++++++++++++++++++ 5 files changed, 51 insertions(+), 4 deletions(-) create mode 100644 bgfilter/weights.py diff --git a/.gitignore b/.gitignore index a6f8414..0ceb88f 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,6 @@ Samples/ Outputs/ +models/ +model-cache/ __pycache__/ *.py[cod] diff --git a/README.md b/README.md index 8ebbbd5..bc5ba34 100644 --- a/README.md +++ b/README.md @@ -205,6 +205,13 @@ D:\MiniConda\envs\lightML\python.exe scripts\smoke_samples.py ` `FileMetadataError`; (2) with `hf-xet` installed the Xet download path fails instantly. Once the weights are cached, run offline with `HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1` (as the service does in production). +- Bundling weights in the project — instead of the HF cache, drop each model into a + plain folder named after the repo basename under `models/`: + `models/vitmatte-base-composition-1k`, `models/BiRefNet`, `models/anime-seg`. The loader + prefers a matching local folder and falls back to the HF repo id / cache when absent, so + it is opt-in. Populate them with e.g. `hf download ZhengPeng7/BiRefNet --local-dir + models/BiRefNet`. Override the base directory with `BGFILTER_WEIGHTS_DIR`. `models/` is + gitignored. - Foreground colour estimation uses pymatting's `estimate_foreground_ml` to propagate clean foreground colour into semi-transparent edges before de-spill. Set `foreground.method: unmix` to fall back to the legacy heuristic. diff --git a/bgfilter/segmentation.py b/bgfilter/segmentation.py index e6269bf..f7fa882 100644 --- a/bgfilter/segmentation.py +++ b/bgfilter/segmentation.py @@ -1,9 +1,12 @@ from __future__ import annotations +import os + import numpy as np from PIL import Image from .settings import SegmentationSettings +from .weights import resolve_model_source class BiRefNetSegmenter: @@ -35,7 +38,7 @@ class BiRefNetSegmenter: self.torch = torch self.model = AutoModelForImageSegmentation.from_pretrained( - settings.model_name, trust_remote_code=True + resolve_model_source(settings.model_name), trust_remote_code=True ) self.model.eval() self.model.float() # checkpoint ships as fp16; force fp32 to match inputs @@ -91,7 +94,11 @@ class AnimeSegSegmenter: "pip install onnxruntime huggingface_hub" ) from exc - model_file = hf_hub_download(settings.model_name, "isnetis.onnx") + source = resolve_model_source(settings.model_name) + if os.path.isdir(source): + model_file = os.path.join(source, "isnetis.onnx") + else: + model_file = hf_hub_download(source, "isnetis.onnx") providers = ["CPUExecutionProvider"] if settings.device == "cuda" and "CUDAExecutionProvider" in ort.get_available_providers(): providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] diff --git a/bgfilter/vitmatte_infer.py b/bgfilter/vitmatte_infer.py index d33a265..9884f4f 100644 --- a/bgfilter/vitmatte_infer.py +++ b/bgfilter/vitmatte_infer.py @@ -4,6 +4,7 @@ import numpy as np from PIL import Image from .settings import ModelSettings +from .weights import resolve_model_source class ViTMatteRunner: @@ -18,8 +19,9 @@ class ViTMatteRunner: ) from exc self.torch = torch - self.processor = VitMatteImageProcessor.from_pretrained(settings.model_name) - self.model = VitMatteForImageMatting.from_pretrained(settings.model_name) + source = resolve_model_source(settings.model_name) + self.processor = VitMatteImageProcessor.from_pretrained(source) + self.model = VitMatteForImageMatting.from_pretrained(source) self.device = self._resolve_device(settings.device) self.model.to(self.device) self.model.eval() diff --git a/bgfilter/weights.py b/bgfilter/weights.py new file mode 100644 index 0000000..82bf2de --- /dev/null +++ b/bgfilter/weights.py @@ -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 ``/`` 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