From 0cc4aecfd5cc1583436165223177436453bf790b Mon Sep 17 00:00:00 2001
From: Phil Wang <lucidrains@gmail.com>
Date: Sun, 6 Sep 2020 01:40:57 -0700
Subject: [PATCH] first alpha release

---
 README.md                                     | 34 +++++++++++++++++++
 .../denoising_diffusion_pytorch.py            | 18 ++++------
 setup.py                                      |  3 +-
 3 files changed, 43 insertions(+), 12 deletions(-)

diff --git a/README.md b/README.md
index 0e2910e..f8f3b88 100644
--- a/README.md
+++ b/README.md
@@ -2,6 +2,40 @@
 
 Implementation of <a href="https://arxiv.org/abs/2006.11239">Denoising Diffusion Probabilistic Model</a> in Pytorch
 
+## Install
+
+```bash
+$ pip install denoising_diffusion_pytorch
+```
+
+## Usage
+
+```python
+import torch
+from denoising_diffusion_pytorch import Unet, GaussianDiffusion
+
+model = Unet(
+    dim = 64,
+    dim_mults = (1, 2, 4, 8)
+)
+
+diffusion = GaussianDiffusion(
+    model,
+    beta_start = 0.0001,
+    beta_end = 0.02,
+    num_diffusion_timesteps = 1000,   # number of steps
+    loss_type = 'l1'                  # L1 or L2
+)
+
+training_images = torch.randn(8, 3, 128, 128)
+loss = diffusion(training_images)
+loss.backward()
+# after a lot of training
+
+sampled_images = diffusion.p_sample_loop((1, 3, 128, 128))
+sampled_images.shape # (1, 3, 128, 128)
+```
+
 ## Citations
 
 ```bibtex
diff --git a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
index 3c0e45b..06c30cf 100644
--- a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
+++ b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
@@ -6,6 +6,7 @@ from torch import nn, einsum
 import torch.nn.functional as F
 
 import numpy as np
+from tqdm import tqdm
 from einops import rearrange
 
 # helpers functions
@@ -291,7 +292,7 @@ class GaussianDiffusion(nn.Module):
         b = shape[0]
         img = torch.randn(shape, device=device)
 
-        for i in reversed(range(0, self.num_timesteps)):
+        for i in tqdm(reversed(range(0, self.num_timesteps)), desc='sampling loop time step', total=self.num_timesteps):
             img = self.p_sample(img, torch.full((b,), i, device=device, dtype=torch.long))
         return img
 
@@ -303,7 +304,7 @@ class GaussianDiffusion(nn.Module):
             extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise
         )
 
-    def p_losses(self, x_start, t, denoise_fn, noise = None):
+    def p_losses(self, x_start, t, noise = None):
         b, c, h, w = x_start.shape
         noise = default(noise, lambda: torch.randn_like(x_start))
 
@@ -319,12 +320,7 @@ class GaussianDiffusion(nn.Module):
 
         return loss
 
-    def forward(self, *args, **kwargs):
-        return self.p_losses(*args, **kwargs)
-
-class Trainer(nn.Module):
-    def __init__(self):
-        super().__init__()
-
-    def forward(self, x):
-        return x
+    def forward(self, x, *args, **kwargs):
+        b, *_, device = *x.shape, x.device
+        t = torch.randint(0, 1000, (b,), device=device).long()
+        return self.p_losses(x, t, *args, **kwargs)
diff --git a/setup.py b/setup.py
index 3d86987..97b3895 100644
--- a/setup.py
+++ b/setup.py
@@ -16,7 +16,8 @@ setup(
   install_requires=[
     'einops',
     'numpy',
-    'torch'
+    'torch',
+    'tqdm'
   ],
   classifiers=[
     'Development Status :: 4 - Beta',