denoise

app.py

@@ -122,6 +122,22 @@ ab_t[0] = 1
 # construct model
 nn_model = ContextUnet(in_channels=3, n_feat=n_feat, n_cfeat=n_cfeat, height=height).to(device)
 
+# define sampling function for DDIM
+# removes the noise using ddim
+def denoise_ddim(x, t, t_prev, pred_noise):
+    ab = ab_t[t]
+    ab_prev = ab_t[t_prev]
+
+    x0_pred = ab_prev.sqrt() / ab.sqrt() * (x - (1 - ab).sqrt() * pred_noise)
+    dir_xt = (1 - ab_prev).sqrt() * pred_noise
+
+    return x0_pred + dir_xt
+
+# load in model weights and set to eval mode
+nn_model.load_state_dict(torch.load(f"{save_dir}/model_31.pth", map_location=device))
+nn_model.eval()
+print("Loaded in Model without context")
+
 # sample quickly using DDIM
 @torch.no_grad()
 def sample_ddim(n_sample, n=20):