handle tuple / list for resolutions

audiodiffusion/__init__.py

@@ -213,11 +213,12 @@ class AudioDiffusionPipeline(DiffusionPipeline):
         step_generator = step_generator or generator
         # For backwards compatibility
         if type(self.unet.sample_size) == int:
-            self.unet.sample_size = [self.unet.sample_size,
-                                     self.unet.sample_size]
+            self.unet.sample_size = (self.unet.sample_size,
+                                     self.unet.sample_size)
         if noise is None:
             noise = torch.randn(
-                [batch_size, self.unet.in_channels] + self.unet.sample_size,
+                (batch_size, self.unet.in_channels, self.unet.sample_size[0],
+                 self.unet.sample_size[1]),
                 generator=generator)
         images = noise
         mask = None