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:
2026-07-03 02:50:19 +00:00
parent 7a87a00ca1
commit da4164e95a
5 changed files with 51 additions and 4 deletions
+9 -2
View File
@@ -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"]