use beartype

imagen_pytorch/imagen_pytorch.py

@@ -2,6 +2,7 @@
 import copy
 from random import random
 from typing import List, Union
+from beartype import beartype
 from tqdm.auto import tqdm
 from functools import partial, wraps
 from contextlib import contextmanager, nullcontext
@@ -1747,6 +1748,7 @@ def __init__(self, *args, **kwargs):
 
 # main imagen ddpm class, which is a cascading DDPM from Ho et al.
 
+@beartype
 class Imagen(nn.Module):
     def __init__(
         self,

imagen_pytorch/version.py

@@ -1 +1 @@
-__version__ = '1.16.4'
+__version__ = '1.16.5'

setup.py

@@ -27,6 +27,7 @@
   ],
   install_requires=[
     'accelerate',
+    'beartype',
     'click',
     'datasets',
     'einops>=0.6',