Sampled Image using DDIM in Half the number of timesteps

Checkpoints/ckpt_199.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e55c1db65805437bc0443bf33ad9749b03531709f563c88bccc828447460e5c1
+size 332549097

main.py

@@ -1,4 +1,4 @@
-from train import train, eval
+from train import eval_ddim, train, eval
 
 
 def main(model_config=None):
@@ -6,7 +6,7 @@ def main(model_config=None):
         "state": "eval", # or eval
         "epochs": 200,
         "batch_size": 80,
-        "T": 1000,
+        "T": 500,
         "channel": 128,
         "ch_mult": [1, 2, 3, 4],
         "attn": [2],
@@ -30,8 +30,10 @@ def main(model_config=None):
 
     if modelConfig['state'] == 'train':
         train(modelConfig)
-    else:
+    elif modelConfig['state'] == 'eval':
         eval(modelConfig)
+    elif modelConfig['state'] == 'eval_ddim':
+        eval_ddim(modelConfig)
     
 
 if __name__ == "__main__":

train.py

@@ -7,7 +7,7 @@ from torch.utils.data import DataLoader
 from torchvision import transforms
 from torchvision.datasets import CIFAR10
 from torchvision.utils import save_image
-from diffusion import GaussianDiffusionSampler, GaussianDiffusionTrainer
+from diffusion import GaussianDiffusionSampler, GaussianDiffusionTrainer, DDIMSampler
 from model import UNet
 from scheduler import GradualWarmupScheduler
 
@@ -84,3 +84,24 @@ def eval(modelConfig:Dict):
             nrow = modelConfig['nrow']
         )
 
+
+def eval_ddim(modelConfig:Dict):
+    with torch.no_grad():
+        device = torch.device(modelConfig['device'])
+        model = torch.load(os.path.join(modelConfig['checkpoint_dir'],modelConfig['test_load_weight']),device)
+        print("Model loaded")
+        model.eval()
+        sampler = DDIMSampler(
+            modelConfig['beta_1'], modelConfig['beta_T'],
+            model,modelConfig['T']
+        )
+        noisyImage = torch.randn(
+            size=[modelConfig['batch_size'],3,32,32],
+            device=device
+        )
+        sampledImgs = sampler(noisyImage)
+        sampledImgs = sampledImgs * 0.5 + 0.5
+        save_image(sampledImgs,
+            os.path.join(modelConfig['sample_dir'],modelConfig['sampledImgName']),
+            nrow = modelConfig['nrow']
+        )