fix something nonsensical when sampling random times, thanks to @jtawade

imagen_pytorch/imagen_pytorch.py

@@ -207,8 +207,8 @@ def __init__(self, *, noise_schedule, timesteps = 1000):
     def get_times(self, batch_size, noise_level, *, device):
         return torch.full((batch_size,), noise_level, device = device, dtype = torch.float32)
 
-    def sample_random_times(self, batch_size, max_thres = 0.999, *, device):
-        return torch.zeros((batch_size,), device = device).float().uniform_(0, max_thres)
+    def sample_random_times(self, batch_size, *, device):
+        return torch.zeros((batch_size,), device = device).float().uniform_(0, 1)
 
     def get_condition(self, times):
         return maybe(self.log_snr)(times)

imagen_pytorch/version.py

@@ -1 +1 @@
-__version__ = '1.15.0'
+__version__ = '1.15.1'

setup.py

@@ -28,6 +28,7 @@
   install_requires=[
     'accelerate',
     'click',
+    'datasets',
     'einops>=0.4',
     'einops-exts',
     'ema-pytorch>=0.0.3',