| 3 | |
| 4 | |
| 5 | class Schedule: |
| 6 | def __init__(self, schedule, timesteps): |
| 7 | self.timesteps = timesteps |
| 8 | self.schedule = schedule |
| 9 | |
| 10 | def cosine_beta_schedule(self, s=0.001): |
| 11 | timesteps = self.timesteps |
| 12 | steps = timesteps + 1 |
| 13 | x = torch.linspace(0, timesteps, steps) |
| 14 | alphas_cumprod = torch.cos(((x / timesteps) + s) / (1 + s) * np.pi * 0.5) ** 2 |
| 15 | alphas_cumprod = alphas_cumprod / alphas_cumprod[0] |
| 16 | betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1]) |
| 17 | return torch.clip(betas, 0.0001, 0.9999) |
| 18 | |
| 19 | def linear_beta_schedule(self): |
| 20 | timesteps = self.timesteps |
| 21 | scale = 1000 / timesteps |
| 22 | beta_start = 1e-6 * scale |
| 23 | beta_end = 0.02 * scale |
| 24 | return torch.linspace(beta_start, beta_end, timesteps) |
| 25 | |
| 26 | def quadratic_beta_schedule(self): |
| 27 | timesteps = self.timesteps |
| 28 | scale = 1000 / timesteps |
| 29 | beta_start = 1e-6 * scale |
| 30 | beta_end = 0.02 * scale |
| 31 | return torch.linspace(beta_start ** 0.5, beta_end ** 0.5, timesteps) ** 2 |
| 32 | |
| 33 | def sigmoid_beta_schedule(self): |
| 34 | timesteps = self.timesteps |
| 35 | scale = 1000 / timesteps |
| 36 | beta_start = 1e-6 * scale |
| 37 | beta_end = 0.02 * scale |
| 38 | betas = torch.linspace(-6, 6, timesteps) |
| 39 | return torch.sigmoid(betas) * (beta_end - beta_start) + beta_start |
| 40 | |
| 41 | def get_betas(self): |
| 42 | if self.schedule == "linear": |
| 43 | return self.linear_beta_schedule() |
| 44 | elif self.schedule == 'cosine': |
| 45 | return self.cosine_beta_schedule() |
| 46 | else: |
| 47 | raise NotImplementedError |
| 48 | |
| 49 | |
| 50 | if __name__ == "__main__": |
no outgoing calls
no test coverage detected