Update sampler.py

sampler.py

@@ -1,14 +1,13 @@
 import torch
-import torch.nn.functional as F
 
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
+
 def extract(a, t, x_shape):
-   batch_size = t.shape[0]
-   t = t.to(a.device)
-   out = a.gather(-1, t)
-   out = out.to(t.device)
-   return out.reshape(batch_size, *((1,) * (len(x_shape) - 1)))
+    batch_size = t.shape[0]
+    out = a.gather(0, t.to(a.device))
+    return out.reshape(batch_size, *((1,) * (len(x_shape) - 1)))
+
 
 beta_start = 1e-4
 beta_end = 0.02
@@ -21,28 +20,20 @@ betas = torch.linspace(
     dtype=torch.float32,
     device=device,
 )
+
 alphas = 1.0 - betas
 alphas_cumprod = torch.cumprod(alphas, dim=0)
 
-alphas_cumprod_prev = F.pad(alphas_cumprod[:-1], (1, 0), value=1.0)
-sqrt_alphas_cumprod = torch.sqrt(alphas_cumprod)
-sqrt_one_minus_alphas_cumprod = torch.sqrt(1.0 - alphas_cumprod)
-
-sqrt_recip_alphas = torch.sqrt(1.0 / alphas)
-
-posterior_variance = (betas * (1.0 - alphas_cumprod_prev) / (1.0 - alphas_cumprod))
 
 @torch.no_grad()
 def ddim_sample(model, x, t, t_prev, labels=None):
-
     alpha_bar_t = extract(alphas_cumprod, t, x.shape)
     alpha_bar_prev = extract(alphas_cumprod, t_prev, x.shape)
 
     pred_noise = model(x, t, labels)
 
     x0 = (
-        x
-        - torch.sqrt(1 - alpha_bar_t) * pred_noise
+        x - torch.sqrt(1 - alpha_bar_t) * pred_noise
     ) / torch.sqrt(alpha_bar_t)
 
     x_prev = (
@@ -50,5 +41,4 @@ def ddim_sample(model, x, t, t_prev, labels=None):
         + torch.sqrt(1 - alpha_bar_prev) * pred_noise
     )
 
-    return x_prev
-
+    return x_prev