an extra assert to make sure images are not integer type

imagen_pytorch/elucidated_imagen.py

@@ -28,6 +28,7 @@
     default,
     cast_tuple,
     cast_uint8_images_to_float,
+    is_float_dtype,
     eval_decorator,
     check_shape,
     pad_tuple_to_length,
@@ -684,6 +685,8 @@ def forward(
         images = cast_uint8_images_to_float(images)
         cond_images = maybe(cast_uint8_images_to_float)(cond_images)
 
+        assert is_float_dtype(images.dtype), f'images tensor needs to be floats but {images.dtype} dtype found instead'
+
         unet_index = unet_number - 1
 
         unet = default(unet, lambda: self.get_unet(unet_number))

imagen_pytorch/imagen_pytorch.py

@@ -76,6 +76,9 @@ def cast_tuple(val, length = None):
 
     return output
 
+def is_float_dtype(dtype):
+    return any([dtype == float_dtype for float_dtype in (torch.float64, torch.float32, torch.float16, torch.bfloat16)])
+
 def cast_uint8_images_to_float(images):
     if not images.dtype == torch.uint8:
         return images
@@ -2411,6 +2414,8 @@ def forward(
         images = cast_uint8_images_to_float(images)
         cond_images = maybe(cast_uint8_images_to_float)(cond_images)
 
+        assert is_float_dtype(images.dtype), f'images tensor needs to be floats but {images.dtype} dtype found instead'
+
         unet_index = unet_number - 1
 
         unet = default(unet, lambda: self.get_unet(unet_number))

imagen_pytorch/version.py

@@ -1 +1 @@
-__version__ = '1.11.0'
+__version__ = '1.11.1'