from .slotformer import SlotFormer
from .single_step_slotformer import SingleStepSlotFormer
from .steve_slotformer import STEVESlotFormer
from .lslotformer import LSlotFormer
from .steve_lslotformer import STEVELSlotFormer
from .savi import StoSAVi
from .e2e import E2E

from .savi import StoSAVi
from .savi_tl import StoSAViTL
from .dVAE import dVAE
from .steve import STEVE
from .utils import get_lr, to_rgb_from_tensor, assert_shape, SoftPositionEmbed
from .steve_utils import cosine_anneal, gumbel_softmax, make_one_hot
from .steve_transformer import STEVETransformerDecoder

def build_model(params):
    if params.model == "E2E":
    
        return E2E(
                    resolution=params.resolution,
                    clip_len=params.input_frames,
                    slot_dict=params.slot_dict,
                    enc_dict=params.enc_dict,
                    dec_dict=params.dec_dict,
                    pred_dict=params.pred_dict,
                    rollout_dict=params.rollout_dict,
                    loss_dict=params.loss_dict,
                )