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:
@@ -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"]
|
||||
|
||||
Reference in New Issue
Block a user