# Copyright (c) CAIRI AI Lab. All rights reserved

from .convlstm import ConvLSTM
from .e3dlstm import E3DLSTM
from .mim import MIM
from .predrnn import PredRNN
from .predrnnpp import PredRNNpp
from .predrnnv2 import PredRNNv2
from .simvp import SimVP
from .tau import TAU
from .ustep import USTEP

method_maps = {
    'convlstm': ConvLSTM,
    'e3dlstm': E3DLSTM,
    'mim': MIM,
    'predrnn': PredRNN,
    'predrnnpp': PredRNNpp,
    'predrnnv2': PredRNNv2,
    'simvp': SimVP,
    'tau': TAU,
    'ustep': USTEP,
}

__all__ = [
    'method_maps', 'ConvLSTM', 'E3DLSTM', 'MIM',
    'PredRNN', 'PredRNNpp', 'PredRNNv2', 'SimVP', 
    'TAU', 'USTEP'
]