| |
|
|
| |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
| from typing import TYPE_CHECKING, Any |
|
|
| if TYPE_CHECKING: |
| from collections.abc import Sequence |
|
|
| _DEFAULT_NUM_TRAIN_TIMESTEPS = 1000 |
|
|
|
|
| @dataclass(frozen=True) |
| class Schedule: |
| """The discretization a solver steps along for one denoising run. |
| |
| Attributes: |
| sigmas: Noise schedule (length = number of steps + 1, terminating at 0). |
| alphas: Signal schedule (length = number of steps + 1). |
| timesteps: Model conditioning timesteps (length = number of steps). |
| num_train_timesteps: Training step count the timesteps are scaled to. |
| """ |
|
|
| sigmas: tuple[float, ...] |
| alphas: tuple[float, ...] |
| timesteps: tuple[float, ...] |
| num_train_timesteps: int = _DEFAULT_NUM_TRAIN_TIMESTEPS |
|
|
| def __len__(self) -> int: |
| return len(self.timesteps) |
|
|
| @classmethod |
| def from_sigmas( |
| cls, |
| sigmas: Sequence[float], |
| *, |
| num_train_timesteps: int = _DEFAULT_NUM_TRAIN_TIMESTEPS, |
| ) -> Schedule: |
| """Build a schedule from an explicit sigma schedule, with no scheduler involved. |
| |
| Args: |
| sigmas: Full noise schedule (length = number of steps + 1), taken as given. |
| num_train_timesteps: Training step count the timesteps are scaled to. |
| """ |
| sigmas = tuple(float(s) for s in sigmas) |
| return cls( |
| sigmas=sigmas, |
| alphas=tuple(1.0 - s for s in sigmas), |
| timesteps=tuple(s * num_train_timesteps for s in sigmas[:-1]), |
| num_train_timesteps=num_train_timesteps, |
| ) |
|
|
| @classmethod |
| def from_scheduler( |
| cls, |
| scheduler: Any, |
| num_inference_steps: int | None = None, |
| **set_timesteps_kwargs: Any, |
| ) -> Schedule: |
| """Build a schedule by discretizing a diffusers scheduler. |
| |
| Args: |
| scheduler: Scheduler to discretize. |
| num_inference_steps: Number of denoising steps. |
| set_timesteps_kwargs: Forwarded verbatim to `scheduler.set_timesteps`. |
| """ |
| scheduler.set_timesteps(num_inference_steps, **set_timesteps_kwargs) |
|
|
| sigmas = tuple(float(s) for s in scheduler.sigmas.tolist()) |
| return cls( |
| sigmas=sigmas, |
| alphas=tuple(1.0 - s for s in sigmas), |
| timesteps=tuple(float(t) for t in scheduler.timesteps.tolist()), |
| num_train_timesteps=scheduler.config.num_train_timesteps, |
| ) |
|
|
