small helper fn to make sampling more clear

README.md

@@ -34,8 +34,8 @@ loss = diffusion(training_images)
 loss.backward()
 # after a lot of training
 
-sampled_images = diffusion.p_sample_loop((1, 3, 128, 128))
+sampled_images = diffusion.sample(128, batch_size = 4)
-sampled_images.shape # (1, 3, 128, 128)
+sampled_images.shape # (4, 3, 128, 128)
 ```
 
 Or, if you simply want to pass in a folder name and the desired image dimensions, you can use the `Trainer` class to easily train a model.

denoising_diffusion_pytorch/denoising_diffusion_pytorch.py

@@ -309,6 +309,10 @@ class GaussianDiffusion(nn.Module):
             img = self.p_sample(img, torch.full((b,), i, device=device, dtype=torch.long))
         return img
 
+    @torch.no_grad()
+    def sample(self, image_size, batch_size = 16):
+        return self.p_sample_loop((16, 3, image_size, image_size))
+
     def q_sample(self, x_start, t, noise=None):
         noise = default(noise, lambda: torch.randn_like(x_start))