diff --git a/k_diffusion/models/__init__.py b/k_diffusion/models/__init__.py
new file mode 100644
index 0000000..2ee2244
--- /dev/null
+++ b/k_diffusion/models/__init__.py
@@ -0,0 +1,5 @@
+from . import flops
+from .flags import checkpointing, get_checkpointing
+from .image_v1 import ImageDenoiserModelV1
+from .image_transformer_v1 import ImageTransformerDenoiserModelV1
+from .image_transformer_v2 import ImageTransformerDenoiserModelV2