Skip to content

dew.diffusion.schedules.source

Native grids and policies for published diffusion checkpoint scheduler files.

One class of a published scheduler_config.json is interpreted once, here, into the native process, solver and time grid that reproduce its set_timesteps and step.

Each pinned Diffusers 0.34.0 class declares its own constructor controls with its own defaults and builds its grid its own way. _SOURCES therefore names, per class, exactly the keys that class reads and the grid family it belongs to. A key another class declares is not read here, the way the source ignores it, and whatever a class does declare and this file does not reconstruct is refused rather than dropped.

The five families are the shapes those set_timesteps take:

  • tabulated: DDIM, PNDM, DDPM, LCM and TCD step between integer indices of the training beta table, so the schedule is that table and the grid is the indices. DDIM and PNDM transfer over a fixed training stride whatever their evaluation spacing. DDPM, LCM and TCD step to the grid’s own next point.
  • lambda: DPM-Solver multistep and singlestep, DEIS and UniPC integrate in log-SNR over paired sigma and model-time tables, normalized so that alpha^2 + sigma^2 is 1, and truncate their model times to integers.
  • sigma: LMS, Euler, Euler ancestral and Heun integrate the variance-exploding sigma directly and scale the model input by 1 / sqrt(sigma^2 + 1).
  • stage: KDPM2, KDPM2 ancestral and DPMSolverSDE evaluate the model twice per interval. Their grids carry the interpolated stage rows the source places between grid points. The solver’s second evaluation reads the source’s sigma and model time there, while the outer walk still visits one point per interval.
  • edm: EDMDPMSolverMultistep is EDM’s own convention, sigma_min to sigma_max at rho with c_noise = log(sigma) / 4 and a signed c_out. It has no beta table and no VP training law, so its training process is EDM’s log-normal sigma draw over the same preconditioning.
NameSummary
SourceScheduleSource-file policy interpreted once into native solver and grid fields.
published_betasThe class’s beta table, rescaled for zero terminal SNR when it asks.

dataclass source

class SourceSchedule(
config: Mapping[str, object],
betas: np.ndarray,
prediction: PredictionTransform,
policy: _Policy,
sampler: Solver,
_grids: dict[tuple[int, int | None, Origin], tuple[Process, jax.Array]] = dict(),
)

Source-file policy interpreted once into native solver and grid fields.

train_steps: int

The training step count the class declares, which is the beta table’s length wherever the class tabulates one.

def from_config(config: Mapping[str, object]) -> SourceSchedule
def training_process() -> Process

The process the checkpoint was trained under.

Every class but the EDM one tabulates a VP beta table and trains on it. EDM’s convention has no beta table and no VP law. Its training process is EDM’s own log-normal sigma draw, over the preconditioning the sampler reads.

def solver() -> Solver

The native solver this file’s class and controls name, resolved once when the file was read.

def sampling(
steps: int,
*,
tokens: int | None = None,
origin: Origin = 'scheduler',
) -> tuple[Process, jax.Array]

The process and the explicit descending grid a steps walk takes.

tokens is the latent token count a resolution-dependent flow shift reads, and origin is where a flow file’s sigmas start. Both belong to the calling pipeline, bound through the task’s grid callable.

function source

def published_betas(
*,
count: JSON,
start: JSON,
end: JSON,
schedule: JSON,
trained: JSON | np.ndarray,
zero_snr: bool,
schedules: tuple[str, ...],
) -> np.ndarray

The class’s beta table, rescaled for zero terminal SNR when it asks.

Every control arrives already resolved against the class’s own declared default. A file that omits one gets that class’s value, and a control the class does not declare never reaches the table. schedules are the beta_schedule tables the class implements: all of them accept the three common ones, DDPM adds GeoDiff’s sigmoid and Heun the exponential alpha-bar.