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