models/__init__.py

from .prob_unet import ProbUNet
from .gen_prob_unet import GenProbUNet
from .phiseg import PhiSeg

MODEL_REGISTRY = {
    "prob_unet": ProbUNet,
    "gen_prob_unet": GenProbUNet,
    "phiseg": PhiSeg,
}