models.py

"""Model loading and caching."""

import sys
import types

from caliby import CalibyModel

MODELS: dict[str, CalibyModel] = {}


def get_model(variant: str, device: str) -> CalibyModel:
    """Load and cache a CalibyModel by variant name."""
    if variant not in MODELS:
        # ZeroGPU's @spaces.GPU decorator may remove sys.modules["__main__"].
        # Lightning's load_from_checkpoint calls inspect.stack() which
        # requires it, so ensure a placeholder exists.
        if "__main__" not in sys.modules:
            sys.modules["__main__"] = types.ModuleType("__main__")

        from caliby import load_model

        MODELS[variant] = load_model(variant, device=device)
    return MODELS[variant]