diff --git a/src/nobg/__init__.py b/src/nobg/__init__.py index dbace71..7e23134 100644 --- a/src/nobg/__init__.py +++ b/src/nobg/__init__.py @@ -4,4 +4,4 @@ from .birefnet.image_processing_birefnet import BiRefNetImageProcessor from .birefnet.modeling_birefnet import BiRefNet -__version__ = "0.2.4" +__version__ = "0.2.5" diff --git a/src/nobg/auto.py b/src/nobg/auto.py index d813fd3..2d687a0 100644 --- a/src/nobg/auto.py +++ b/src/nobg/auto.py @@ -14,7 +14,7 @@ class AutoModel: @set_doc(PyTorchModelHubMixin.from_pretrained.__doc__) def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs): tags: list[str] = model_info(pretrained_model_name_or_path).tags or [] - if "birefnet" in tags: + if "biref" in tags or "birefnet" in tags: return BiRefNet.from_pretrained( pretrained_model_name_or_path, *model_args, **kwargs ) @@ -52,7 +52,7 @@ def from_pretrained(cls, pretrained_model_name_or_path, **kwargs): @staticmethod def _from_model_config(pretrained_model_name_or_path, **kwargs): tags: list[str] = model_info(pretrained_model_name_or_path).tags or [] - if "birefnet" not in tags: + if "biref" not in tags and "birefnet" not in tags: raise ValueError("this model is not part of nobg") config_file = hf_hub_download( pretrained_model_name_or_path, diff --git a/src/nobg/birefnet/modeling_birefnet.py b/src/nobg/birefnet/modeling_birefnet.py index ded038d..946cdaa 100644 --- a/src/nobg/birefnet/modeling_birefnet.py +++ b/src/nobg/birefnet/modeling_birefnet.py @@ -427,7 +427,7 @@ class BiRefNet( repo_url="https://github.com/feyninc/nobg", paper_url="https://arxiv.org/abs/2401.03407", license="apache-2.0", - tags=["nobg", "birefnet"], + tags=["nobg", "biref"], model_card_template=model_card_template( class_name="BiRefNet", default_repo="nobg/birefnet",