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)
| 227 | |
| 228 | |
| 229 | def get_named_beta_schedule(schedule_name, num_diffusion_timesteps): |
| 230 | """ |
| 231 | Get a pre-defined beta schedule for the given name. |
| 232 | |
| 233 | The beta schedule library consists of beta schedules which remain similar |
| 234 | in the limit of num_diffusion_timesteps. |
| 235 | Beta schedules may be added, but should not be removed or changed once |
| 236 | they are committed to maintain backwards compatibility. |
| 237 | """ |
| 238 | if schedule_name == "linear": |
| 239 | # Linear schedule from Ho et al, extended to work for any number of |
| 240 | # diffusion steps. |
| 241 | scale = 1000 / num_diffusion_timesteps |
| 242 | beta_start = scale * 0.0001 |
| 243 | beta_end = scale * 0.02 |
| 244 | return np.linspace( |
| 245 | beta_start, beta_end, num_diffusion_timesteps, dtype=np.float64 |
| 246 | ) |
| 247 | elif schedule_name == "cosine": |
| 248 | return betas_for_alpha_bar( |
| 249 | num_diffusion_timesteps, |
| 250 | lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2, |
| 251 | ) |
| 252 | else: |
| 253 | raise NotImplementedError(f"unknown beta schedule: {schedule_name}") |
| 254 | |
| 255 | |
| 256 | def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999): |
no test coverage detected