image is in range of 0 to 1

denoising_diffusion_pytorch/denoising_diffusion_pytorch.py

@@ -22,6 +22,15 @@ try:
 except:
     APEX_AVAILABLE = False
 
+# constants
+
+SAVE_AND_SAMPLE_EVERY = 1000
+UPDATE_EMA_EVERY = 10
+EXTS = ['jpg', 'jpeg', 'png']
+
+RESULTS_FOLDER = Path('./results')
+RESULTS_FOLDER.mkdir(exist_ok = True)
+
 # helpers functions
 
 def exists(x):
@@ -111,15 +120,14 @@ class Downsample(nn.Module):
     def forward(self, x):
         return self.conv(x)
 
-class PreNorm(nn.Module):
-    def __init__(self, dim, fn):
+class Rezero(nn.Module):
+    def __init__(self, fn):
         super().__init__()
         self.fn = fn
-        self.norm = nn.InstanceNorm2d(dim, affine = True)
+        self.g = nn.Parameter(torch.zeros(1))
 
     def forward(self, x):
-        x = self.norm(x)
-        return self.fn(x)
+        return self.fn(x) * self.g
 
 # building block modules
 
@@ -135,12 +143,12 @@ class Block(nn.Module):
         return self.block(x)
 
 class ResnetBlock(nn.Module):
-    def __init__(self, dim, dim_out, *, time_emb_dim = None, groups = 8):
+    def __init__(self, dim, dim_out, *, time_emb_dim, groups = 8):
         super().__init__()
         self.mlp = nn.Sequential(
             Mish(),
             nn.Linear(time_emb_dim, dim_out)
-        ) if exists(time_emb_dim) else None
+        )
 
         self.block1 = Block(dim, dim_out)
         self.block2 = Block(dim_out, dim_out)
@@ -148,10 +156,7 @@ class ResnetBlock(nn.Module):
 
     def forward(self, x, time_emb):
         h = self.block1(x)
-
-        if exists(self.mlp):
-            h += self.mlp(time_emb)[:, :, None, None]
-
+        h += self.mlp(time_emb)[:, :, None, None]
         h = self.block2(h)
         return h + self.res_conv(x)
 
@@ -182,8 +187,7 @@ class Unet(nn.Module):
         out_dim = None,
         dim_mults=(1, 2, 4, 8),
         groups = 8,
-        channels = 3,
-        with_time_emb = True
+        channels = 3
     ):
         super().__init__()
         self.channels = channels
@@ -191,17 +195,12 @@ class Unet(nn.Module):
         dims = [channels, *map(lambda m: dim * m, dim_mults)]
         in_out = list(zip(dims[:-1], dims[1:]))
 
-        if with_time_emb:
-            time_dim = dim
-            self.time_mlp = nn.Sequential(
-                SinusoidalPosEmb(dim),
-                nn.Linear(dim, dim * 4),
-                Mish(),
-                nn.Linear(dim * 4, dim)
-            )
-        else:
-            time_dim = None
-            self.time_mlp = None
+        self.time_pos_emb = SinusoidalPosEmb(dim)
+        self.mlp = nn.Sequential(
+            nn.Linear(dim, dim * 4),
+            Mish(),
+            nn.Linear(dim * 4, dim)
+        )
 
         self.downs = nn.ModuleList([])
         self.ups = nn.ModuleList([])
@@ -211,24 +210,24 @@ class Unet(nn.Module):
             is_last = ind >= (num_resolutions - 1)
 
             self.downs.append(nn.ModuleList([
-                ResnetBlock(dim_in, dim_out, time_emb_dim = time_dim),
-                ResnetBlock(dim_out, dim_out, time_emb_dim = time_dim),
-                Residual(PreNorm(dim_out, LinearAttention(dim_out))),
+                ResnetBlock(dim_in, dim_out, time_emb_dim = dim),
+                ResnetBlock(dim_out, dim_out, time_emb_dim = dim),
+                Residual(Rezero(LinearAttention(dim_out))),
                 Downsample(dim_out) if not is_last else nn.Identity()
             ]))
 
         mid_dim = dims[-1]
-        self.mid_block1 = ResnetBlock(mid_dim, mid_dim, time_emb_dim = time_dim)
-        self.mid_attn = Residual(PreNorm(mid_dim, LinearAttention(mid_dim)))
-        self.mid_block2 = ResnetBlock(mid_dim, mid_dim, time_emb_dim = time_dim)
+        self.mid_block1 = ResnetBlock(mid_dim, mid_dim, time_emb_dim = dim)
+        self.mid_attn = Residual(Rezero(LinearAttention(mid_dim)))
+        self.mid_block2 = ResnetBlock(mid_dim, mid_dim, time_emb_dim = dim)
 
         for ind, (dim_in, dim_out) in enumerate(reversed(in_out[1:])):
             is_last = ind >= (num_resolutions - 1)
 
             self.ups.append(nn.ModuleList([
-                ResnetBlock(dim_out * 2, dim_in, time_emb_dim = time_dim),
-                ResnetBlock(dim_in, dim_in, time_emb_dim = time_dim),
-                Residual(PreNorm(dim_in, LinearAttention(dim_in))),
+                ResnetBlock(dim_out * 2, dim_in, time_emb_dim = dim),
+                ResnetBlock(dim_in, dim_in, time_emb_dim = dim),
+                Residual(Rezero(LinearAttention(dim_in))),
                 Upsample(dim_in) if not is_last else nn.Identity()
             ]))
 
@@ -239,7 +238,8 @@ class Unet(nn.Module):
         )
 
     def forward(self, x, time):
-        t = self.time_mlp(time) if exists(self.time_mlp) else None
+        t = self.time_pos_emb(time)
+        t = self.mlp(t)
 
         h = []
 
@@ -365,7 +365,7 @@ class GaussianDiffusion(nn.Module):
         x_recon = self.predict_start_from_noise(x, t=t, noise=self.denoise_fn(x, t))
 
         if clip_denoised:
-            x_recon.clamp_(-1., 1.)
+            x_recon.clamp_(0., 1.)
 
         model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t)
         return model_mean, posterior_variance, posterior_log_variance
@@ -445,18 +445,17 @@ class GaussianDiffusion(nn.Module):
 # dataset classes
 
 class Dataset(data.Dataset):
-    def __init__(self, folder, image_size, exts = ['jpg', 'jpeg', 'png']):
+    def __init__(self, folder, image_size):
         super().__init__()
         self.folder = folder
         self.image_size = image_size
-        self.paths = [p for ext in exts for p in Path(f'{folder}').glob(f'**/*.{ext}')]
+        self.paths = [p for ext in EXTS for p in Path(f'{folder}').glob(f'**/*.{ext}')]
 
         self.transform = transforms.Compose([
             transforms.Resize(image_size),
             transforms.RandomHorizontalFlip(),
             transforms.CenterCrop(image_size),
-            transforms.ToTensor(),
-            transforms.Lambda(lambda t: (t * 2) - 1)
+            transforms.ToTensor()
         ])
 
     def __len__(self):
@@ -482,19 +481,13 @@ class Trainer(object):
         train_num_steps = 100000,
         gradient_accumulate_every = 2,
         fp16 = False,
-        step_start_ema = 2000,
-        update_ema_every = 10,
-        save_and_sample_every = 1000,
-        results_folder = './results'
+        step_start_ema = 2000
     ):
         super().__init__()
         self.model = diffusion_model
         self.ema = EMA(ema_decay)
         self.ema_model = copy.deepcopy(self.model)
-        self.update_ema_every = update_ema_every
-
         self.step_start_ema = step_start_ema
-        self.save_and_sample_every = save_and_sample_every
 
         self.batch_size = train_batch_size
         self.image_size = diffusion_model.image_size
@@ -513,9 +506,6 @@ class Trainer(object):
         if fp16:
             (self.model, self.ema_model), self.opt = amp.initialize([self.model, self.ema_model], self.opt, opt_level='O1')
 
-        self.results_folder = Path(results_folder)
-        self.results_folder.mkdir(exist_ok = True)
-
         self.reset_parameters()
 
     def reset_parameters(self):
@@ -533,10 +523,10 @@ class Trainer(object):
             'model': self.model.state_dict(),
             'ema': self.ema_model.state_dict()
         }
-        torch.save(data, str(self.results_folder / f'model-{milestone}.pt'))
+        torch.save(data, str(RESULTS_FOLDER / f'model-{milestone}.pt'))
 
     def load(self, milestone):
-        data = torch.load(str(self.results_folder / f'model-{milestone}.pt'))
+        data = torch.load(str(RESULTS_FOLDER / f'model-{milestone}.pt'))
 
         self.step = data['step']
         self.model.load_state_dict(data['model'])
@@ -555,16 +545,15 @@ class Trainer(object):
             self.opt.step()
             self.opt.zero_grad()
 
-            if self.step % self.update_ema_every == 0:
+            if self.step % UPDATE_EMA_EVERY == 0:
                 self.step_ema()
 
-            if self.step != 0 and self.step % self.save_and_sample_every == 0:
-                milestone = self.step // self.save_and_sample_every
+            if self.step != 0 and self.step % SAVE_AND_SAMPLE_EVERY == 0:
+                milestone = self.step // SAVE_AND_SAMPLE_EVERY
                 batches = num_to_groups(36, self.batch_size)
                 all_images_list = list(map(lambda n: self.ema_model.sample(batch_size=n), batches))
                 all_images = torch.cat(all_images_list, dim=0)
-                all_images = (all_images + 1) * 0.5
-                utils.save_image(all_images, str(self.results_folder / f'sample-{milestone}.png'), nrow = 6)
+                utils.save_image(all_images, str(RESULTS_FOLDER / f'sample-{milestone}.png'), nrow=6)
                 self.save(milestone)
 
             self.step += 1

setup.py

@@ -3,7 +3,7 @@ from setuptools import setup, find_packages
 setup(
   name = 'denoising-diffusion-pytorch',
   packages = find_packages(),
-  version = '0.6.8',
+  version = '0.6.2',
   license='MIT',
   description = 'Denoising Diffusion Probabilistic Models - Pytorch',
   author = 'Phil Wang',