| 30 | |
| 31 | |
| 32 | def make_beta_schedule(schedule, n_timestep, linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3): |
| 33 | if schedule == "linear": |
| 34 | betas = ( |
| 35 | torch.linspace(linear_start ** 0.5, linear_end ** 0.5, n_timestep, dtype=torch.float64) ** 2 |
| 36 | ) |
| 37 | |
| 38 | elif schedule == "cosine": |
| 39 | timesteps = ( |
| 40 | torch.arange(n_timestep + 1, dtype=torch.float64) / n_timestep + cosine_s |
| 41 | ) |
| 42 | alphas = timesteps / (1 + cosine_s) * np.pi / 2 |
| 43 | alphas = torch.cos(alphas).pow(2) |
| 44 | alphas = alphas / alphas[0] |
| 45 | betas = 1 - alphas[1:] / alphas[:-1] |
| 46 | betas = np.clip(betas, a_min=0, a_max=0.999) |
| 47 | |
| 48 | elif schedule == "sqrt_linear": |
| 49 | betas = torch.linspace(linear_start, linear_end, n_timestep, dtype=torch.float64) |
| 50 | elif schedule == "sqrt": |
| 51 | betas = torch.linspace(linear_start, linear_end, n_timestep, dtype=torch.float64) ** 0.5 |
| 52 | else: |
| 53 | raise ValueError(f"schedule '{schedule}' unknown.") |
| 54 | return betas.numpy() |
| 55 | |
| 56 | |
| 57 | def make_ddim_timesteps(ddim_discr_method, num_ddim_timesteps, num_ddpm_timesteps, verbose=True): |