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 torch
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=".")
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 = loc
    try:
        model = joblib.load(model_file)
    except Exception as exc:
        logger.error(exc)
        raise
    return model