move all constants into configurable class init parameters

denoising_diffusion_pytorch/denoising_diffusion_pytorch.py

@@ -22,15 +22,6 @@ 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):
@@ -190,6 +181,8 @@ class Unet(nn.Module):
         channels = 3
     ):
         super().__init__()
+        self.channels = channels
+
         dims = [channels, *map(lambda m: dim * m, dim_mults)]
         in_out = list(zip(dims[:-1], dims[1:]))
 
@@ -229,7 +222,7 @@ class Unet(nn.Module):
                 Upsample(dim_in) if not is_last else nn.Identity()
             ]))
 
-        out_dim = default(out_dim, 3)
+        out_dim = default(out_dim, channels)
         self.final_conv = nn.Sequential(
             Block(dim, dim),
             nn.Conv2d(dim, out_dim, 1)
@@ -291,11 +284,13 @@ class GaussianDiffusion(nn.Module):
         denoise_fn,
         *,
         image_size,
+        channels = 3,
         timesteps = 1000,
         loss_type = 'l1',
         betas = None
     ):
         super().__init__()
+        self.channels = channels
         self.image_size = image_size
         self.denoise_fn = denoise_fn
 
@@ -389,7 +384,8 @@ class GaussianDiffusion(nn.Module):
     @torch.no_grad()
     def sample(self, batch_size = 16):
         image_size = self.image_size
-        return self.p_sample_loop((batch_size, 3, image_size, image_size))
+        channels = self.channels
+        return self.p_sample_loop((batch_size, channels, image_size, image_size))
 
     @torch.no_grad()
     def interpolate(self, x1, x2, t = None, lam = 0.5):
@@ -440,17 +436,18 @@ class GaussianDiffusion(nn.Module):
 # dataset classes
 
 class Dataset(data.Dataset):
-    def __init__(self, folder, image_size):
+    def __init__(self, folder, image_size, exts = ['jpg', 'jpeg', 'png']):
         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.ToTensor(),
+            transforms.Lambda(lambda t: (t * 2) - 1)
         ])
 
     def __len__(self):
@@ -476,13 +473,19 @@ class Trainer(object):
         train_num_steps = 100000,
         gradient_accumulate_every = 2,
         fp16 = False,
-        step_start_ema = 2000
+        step_start_ema = 2000,
+        update_ema_every = 10,
+        save_and_sample_every = 1000,
+        results_folder = './results'
     ):
         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
@@ -501,6 +504,9 @@ 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):
@@ -518,10 +524,10 @@ class Trainer(object):
             'model': self.model.state_dict(),
             'ema': self.ema_model.state_dict()
         }
-        torch.save(data, str(RESULTS_FOLDER / f'model-{milestone}.pt'))
+        torch.save(data, str(self.results_folder / f'model-{milestone}.pt'))
 
     def load(self, milestone):
-        data = torch.load(str(RESULTS_FOLDER / f'model-{milestone}.pt'))
+        data = torch.load(str(self.results_folder / f'model-{milestone}.pt'))
 
         self.step = data['step']
         self.model.load_state_dict(data['model'])
@@ -540,15 +546,16 @@ class Trainer(object):
             self.opt.step()
             self.opt.zero_grad()
 
-            if self.step % UPDATE_EMA_EVERY == 0:
+            if self.step % self.update_ema_every == 0:
                 self.step_ema()
 
-            if self.step != 0 and self.step % SAVE_AND_SAMPLE_EVERY == 0:
+            if self.step != 0 and self.step % self.save_and_sample_every == 0:
-                milestone = self.step // SAVE_AND_SAMPLE_EVERY
+                milestone = self.step // self.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)
-                utils.save_image(all_images, str(RESULTS_FOLDER / f'sample-{milestone}.png'), nrow=6)
+                all_images = (all_images + 1) * 0.5
+                utils.save_image(all_images, str(self.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.0',
+  version = '0.6.6',
   license='MIT',
   description = 'Denoising Diffusion Probabilistic Models - Pytorch',
   author = 'Phil Wang',