Get a pre-defined beta schedule for the given name. The beta schedule library consists of beta schedules which remain similar in the limit of num_diffusion_timesteps. Beta schedules may be added, but should not be removed or changed once they are committed to maintain backwards
(schedule_name, num_diffusion_timesteps)
| 24 | |
| 25 | |
| 26 | def get_named_beta_schedule(schedule_name, num_diffusion_timesteps): |
| 27 | """ |
| 28 | Get a pre-defined beta schedule for the given name. |
| 29 | |
| 30 | The beta schedule library consists of beta schedules which remain similar |
| 31 | in the limit of num_diffusion_timesteps. |
| 32 | Beta schedules may be added, but should not be removed or changed once |
| 33 | they are committed to maintain backwards compatibility. |
| 34 | """ |
| 35 | if schedule_name == "linear": |
| 36 | # Linear schedule from Ho et al, extended to work for any number of |
| 37 | # diffusion steps. |
| 38 | scale = 1000 / num_diffusion_timesteps |
| 39 | return get_beta_schedule( |
| 40 | "linear", |
| 41 | beta_start=scale * 0.0001, |
| 42 | beta_end=scale * 0.02, |
| 43 | num_diffusion_timesteps=num_diffusion_timesteps, |
| 44 | ) |
| 45 | elif schedule_name == "cosine": |
| 46 | return betas_for_alpha_bar( |
| 47 | num_diffusion_timesteps, |
| 48 | lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2, |
| 49 | ) |
| 50 | else: |
| 51 | raise NotImplementedError(f"unknown beta schedule: {schedule_name}") |
| 52 | |
| 53 | |
| 54 | def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999): |
no test coverage detected