(schedule, n_timestep, linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3)
| 11 | |
| 12 | |
| 13 | def make_beta_schedule(schedule, n_timestep, linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3): |
| 14 | if schedule == "linear": |
| 15 | betas = ( |
| 16 | torch.linspace(linear_start ** 0.5, linear_end ** 0.5, n_timestep, dtype=torch.float64) ** 2 |
| 17 | ) |
| 18 | elif schedule == "cosine": |
| 19 | timesteps = ( |
| 20 | torch.arange(n_timestep + 1, dtype=torch.float64) / n_timestep + cosine_s |
| 21 | ) |
| 22 | alphas = timesteps / (1 + cosine_s) * np.pi / 2 |
| 23 | alphas = torch.cos(alphas).pow(2) |
| 24 | alphas = alphas / alphas[0] |
| 25 | betas = 1 - alphas[1:] / alphas[:-1] |
| 26 | betas = np.clip(betas, a_min=0, a_max=0.999) |
| 27 | elif schedule == "squaredcos_cap_v2": # used for karlo prior |
| 28 | # return early |
| 29 | return betas_for_alpha_bar( |
| 30 | n_timestep, |
| 31 | lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2, |
| 32 | ) |
| 33 | elif schedule == "sqrt_linear": |
| 34 | betas = torch.linspace(linear_start, linear_end, n_timestep, dtype=torch.float64) |
| 35 | elif schedule == "sqrt": |
| 36 | betas = torch.linspace(linear_start, linear_end, n_timestep, dtype=torch.float64) ** 0.5 |
| 37 | else: |
| 38 | raise ValueError(f"schedule '{schedule}' unknown.") |
| 39 | return betas.numpy() |
| 40 | |
| 41 | |
| 42 | def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999): |
no test coverage detected