Finished the Sampler

diffusion.py

@@ -59,10 +59,55 @@ class GaussianDiffusionSampler(nn.Module):
         )
         self.alphas = 1 - self.betas
         self.beta_alphas = torch.cumprod(self.alphas,dim=0)
-        
+
         """
          This line of code pads the tensor self.beta_alphas by adding a single element with the value 1 to the beginning of the tensor. 
          The resulting tensor is stored in self.beta_alphas_prev.
         """
         self.beta_alphas_prev = F.pad(self.beta_alphas,[1,0],value=1)[:T]
 
+        self.register_buffer(
+            "coeff1",
+            (1 / torch.sqrt(self.alphas))
+        )
+
+        self.register_buffer(
+            "coeff2",
+            self.coeff1 * ((1- self.alphas) / (torch.sqrt(1-self.beta_alphas)))
+        )
+
+        self.register_buffer(
+            "posterior_coeff",
+            (1 - self.beta_alphas_prev) / (1-self.beta_alphas) * self.betas
+        )
+
+    def pred_xt_prev_mean_from_eps(self,x_t,t,eps):
+        return (
+            extract(self.coeff1,t,x_t.shape) * x_t - 
+            extract(self.coeff2,t,x_t.shape) * eps
+        )
+    
+    def p_mean_variance(self,x_t,t):
+        var = torch.cat([self.posterior_coeff[1:2],self.betas[1:]])
+        var = extract(var,t,x_t.shape)
+
+        eps = self.model(x_t,t)
+        xt_prev_mean = self.pred_xt_prev_mean_from_eps(x_t,t,eps)
+        return xt_prev_mean,var
+    
+    def forward(self,x_T):
+        x_t=x_T
+        for timestep in reversed(range(self.T)):
+            print(f"Sampling timestep: {timestep}")
+
+            t = x_t.new_ones([x_t.shape[0],], dtype=torch.long) * timestep
+            mean, var = self.p_mean_variance(x_t,t)
+            if timestep > 0:
+                noise = torch.randn_like(x_t)
+            else:
+                noise = 0
+            x_t = mean + torch.sqrt(var) * noise
+        
+        x_0 = x_t
+        return torch.clip(x_0,-1,1)
+