Create scheduler.py

scheduler.py

@@ -0,0 +1,50 @@
+import torch
+from torch import nn
+
+
+class DDPMSchedule(nn.Module):
+    def __init__(self, beta_start=1e-4, beta_end=0.02, steps=1000):
+        super().__init__()
+        beta = torch.linspace(beta_start, beta_end, steps)
+        alpha = 1 - beta
+        alpha_cumprod = torch.cumprod(alpha, dim=-1)
+        self.steps = steps
+        self.register_buffer("beta", beta, persistent=False)
+        self.register_buffer("alpha", alpha, persistent=False)
+        self.register_buffer("alpha_cumprod", alpha_cumprod, persistent=False)
+
+    def add_noise(
+        self, image: torch.Tensor, time: torch.LongTensor, noise: torch.Tensor
+    ) -> torch.Tensor:
+        alpha_cumprod_t = self.alpha_cumprod[time]
+
+        while len(alpha_cumprod_t.shape) < len(image.shape):
+            alpha_cumprod_t = alpha_cumprod_t.unsqueeze(-1)
+
+        noise_image = torch.sqrt(alpha_cumprod_t) * image + torch.sqrt(1-alpha_cumprod_t) * noise
+        return noise_image
+    
+    def step(self, x_t:torch.Tensor, pred_noise:torch.Tensor, time:torch.LongTensor):
+        alpha_t = self.alpha[time]
+        alpha_cumprod_t = self.alpha_cumprod[time]
+        alpha_cumprod_t_minus_1 = self.alpha_cumprod[time-1]
+        beta_t = self.beta[time]
+        zero_mask = (time > 0).float()
+
+        while len(alpha_t.shape) < len(x_t.shape):
+            alpha_t = alpha_t.unsqueeze(-1)
+            alpha_cumprod_t = alpha_cumprod_t.unsqueeze(-1)
+            beta_t = beta_t.unsqueeze(-1)
+            zero_mask = zero_mask.unsqueeze(-1)
+            alpha_cumprod_t_minus_1 = alpha_cumprod_t_minus_1.unsqueeze(-1)
+
+        noise_weight = (1-alpha_t)/torch.sqrt(1-alpha_cumprod_t)
+        noise = torch.randn_like(x_t)
+
+        
+        eta = (1-alpha_cumprod_t_minus_1) / (1-alpha_cumprod_t)
+        eta = torch.sqrt(eta * beta_t)
+
+
+        x_t_minus_1 = (x_t - noise_weight * pred_noise) / torch.sqrt(alpha_t) + eta * noise * zero_mask
+        return x_t_minus_1