debugged flexible noising

liopeer · Nov 9, 2023 · e4b29da · e4b29da
1 parent c1f2e7d
commit e4b29da
Show file tree

Hide file tree

Showing 4 changed files with 96 additions and 6 deletions.
diff --git a/diffusion_models/models/diffusion.py b/diffusion_models/models/diffusion.py
@@ -121,8 +121,8 @@ def forward_flexible(
             alphas_dash_interval = torch.cumprod(alphas_interval, axis=0)
             sqrt_alphas_dash_interval = torch.sqrt(alphas_dash_interval)
             sqrt_one_minus_alphas_dash_interval = torch.sqrt(1. - alphas_dash_interval)
-            batch_sqrt_alphas_dash[sample] = sqrt_alphas_dash_interval
-            batch_sqrt_one_minus_alpha_dash[sample] = sqrt_one_minus_alphas_dash_interval
+            batch_sqrt_alphas_dash[sample] = sqrt_alphas_dash_interval[-1]
+            batch_sqrt_one_minus_alpha_dash[sample] = sqrt_one_minus_alphas_dash_interval[-1]
         mean = batch_sqrt_alphas_dash.view(-1, 1, 1, 1) * x_t1
         out = mean + batch_sqrt_one_minus_alpha_dash.view(-1, 1, 1, 1) * noise_normal
         return out, noise_normal

diff --git a/tests/sampler_tests/flexible_denoising.ipynb b/tests/sampler_tests/flexible_denoising.ipynb
diff --git a/tests/sampler_tests/flexible_noising.ipynb b/tests/sampler_tests/flexible_noising.ipynb
diff --git a/tests/sampler_tests/sampler_dutifulpond10.py b/tests/sampler_tests/sampler_dutifulpond10.py
@@ -13,7 +13,7 @@ def run():
     config = None
     with open("config_dutifulpond10.json", "r") as f:
         config = json.load(f)
-    ckp_num = 40
+    ckp_num = 32
     ckpt_path = f"/itet-stor/peerli/net_scratch/dutiful-pond-10/checkpoint{ckp_num}.pt"
     device = torch.device("cuda")
     backbone = UNet(