debugged flexible noising

Author: liopeer

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

tests/sampler_tests/flexible_denoising.ipynb

@@ -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(