common/models.py

"""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