from utils.interpolation_models.EDEN import EDEN
from utils.interpolation_models.EDEN_VAE import VAE
from utils.interpolation_models.EDEN_DiT import DiT
from utils.interpolation_models.Discriminator import NLayerDiscriminator


def load_model(model_name, **model_args):
    if model_name == "EDEN":
        return EDEN(**model_args)
    elif model_name == "EDEN_VAE":
        return VAE(**model_args)
    elif model_name == "EDEN_DiT":
        return DiT(**model_args)
    elif model_name == "Discriminator":
        return NLayerDiscriminator(**model_args)
    else:
        raise f"No model named {model_name} in models!"