st_mlbee/load_model_s.py

r"""
Load model_s from hf.

cf aslo align-model-pool\model_pool\load_model.py and ycco make-upload-model-s.ipynb.
"""
import joblib
from huggingface_hub import hf_hub_download
from loguru import logger

try:
    # loc = hf_hub_download("mikeee/model_s_512", "model-s", local_dir=".")  # v1 MODEL_ID2 dim512
    # loc2 = hf_hub_download("mikeee/model_s_512v2", "model-s-v2", local_dir=".")  # v1 MODEL_ID2 dim512
    loc3 = hf_hub_download("mikeee/model_s_384", "model-s", local_dir=".")  # v1 MODEL_ID2 dim512
except Exception as exc:
    logger.error(exc)
    raise SystemExit(1) from exc


def load_model_s(model_file=None):
    """Load a model from hf."""
    if model_file is None:
        model_file = loc3
    try:
        model = joblib.load(model_file)
    except Exception as exc:
        logger.error(exc)
        raise
    return model