@@ -7,7 +7,7 @@
 __version__ = "0.0.4"
 
 from .modeling_utils import ModelMixin
-from .models import NCSNpp, TemporalUNet, UNetLDMModel, UNetModel
+from .models import AutoencoderKL, NCSNpp, TemporalUNet, UNetLDMModel, UNetModel, VQModel
 from .pipeline_utils import DiffusionPipeline
 from .pipelines import (
     BDDMPipeline,
 
@@ -22,3 +22,4 @@
 from .unet_ldm import UNetLDMModel
 from .unet_rl import TemporalUNet
 from .unet_sde_score_estimation import NCSNpp
+from .vae import AutoencoderKL, VQModel