(schedule, n_timestep, linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3)
| 19 | |
| 20 | |
| 21 | def make_beta_schedule(schedule, n_timestep, linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3): |
| 22 | if schedule == "linear": |
| 23 | betas = ( |
| 24 | torch.linspace(linear_start ** 0.5, linear_end ** 0.5, n_timestep, dtype=torch.float64) ** 2 |
| 25 | ) |
| 26 | |
| 27 | elif schedule == "cosine": |
| 28 | timesteps = ( |
| 29 | torch.arange(n_timestep + 1, dtype=torch.float64) / n_timestep + cosine_s |
| 30 | ) |
| 31 | alphas = timesteps / (1 + cosine_s) * np.pi / 2 |
| 32 | alphas = torch.cos(alphas).pow(2) |
| 33 | alphas = alphas / alphas[0] |
| 34 | betas = 1 - alphas[1:] / alphas[:-1] |
| 35 | betas = np.clip(betas, a_min=0, a_max=0.999) |
| 36 | |
| 37 | elif schedule == "sqrt_linear": |
| 38 | betas = torch.linspace(linear_start, linear_end, n_timestep, dtype=torch.float64) |
| 39 | elif schedule == "sqrt": |
| 40 | betas = torch.linspace(linear_start, linear_end, n_timestep, dtype=torch.float64) ** 0.5 |
| 41 | else: |
| 42 | raise ValueError(f"schedule '{schedule}' unknown.") |
| 43 | return betas.numpy() |
| 44 | |
| 45 | |
| 46 | def make_ddim_timesteps(ddim_discr_method, num_ddim_timesteps, num_ddpm_timesteps, verbose=True): |
no outgoing calls
no test coverage detected