File size: 894 Bytes
e8b8483 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 | """Backbone loading."""
import sys
from pathlib import Path
from typing import Optional
from .paths import BACKBONE, BACKBONE_SRC
_argus = None
def argus_module():
"""Import argus.py from the environment, BACKBONE_SRC, or the backbone repo."""
global _argus
if _argus is not None:
return _argus
try:
import argus
except ImportError:
if BACKBONE_SRC:
sys.path.insert(0, str(BACKBONE_SRC))
else:
from huggingface_hub import hf_hub_download
sys.path.insert(0, str(Path(hf_hub_download(BACKBONE, 'argus.py')).parent))
import argus
_argus = argus
return argus
def load_backbone(repo: Optional[str] = None):
"""Load the stock backbone in eval mode."""
from transformers import AutoModel
return AutoModel.from_pretrained(repo or BACKBONE, trust_remote_code=True).eval().backbone
|