from datasets.video import (
    MinecraftVideoDataset,
    MinecraftVideoLatentDataset,
)

from algorithms.worldmem import DeMemWMMinecraft
from .exp_base import BaseLightningExperiment


class VideoPredictionExperiment(BaseLightningExperiment):
    """
    A video prediction experiment
    """

    compatible_algorithms = dict(
        dememwm_memory_dit=DeMemWMMinecraft,
    )

    compatible_datasets = dict(
        # video datasets
        video_minecraft=MinecraftVideoDataset,
        video_minecraft_latent=MinecraftVideoLatentDataset,
    )