dew.diffusion.process

Holds the one convention a model is trained and sampled with.

Name	Summary
Denoiser	Denoises with one model, its parameters and its conditions.
DenoisingCondition	Carries text conditioning as the published families read it.
Process	Pairs a schedule with what the model predicts on it and how the loss is weighted.
aligned_conditions	unconditional with each row’s own model inputs taken from conditions.

Denoiser

dataclass source

class Denoiser(
    process: Process,
    model: nn.Module,
    params: Variables,
    conditions: dict[str, Conditioning],
    unconditional: dict[str, Conditioning] | None = None,
)

Denoises with one model, its parameters and its conditions.

A call is the model’s raw output at (x_t, t) followed by the process’s conversion of it into (x_0, epsilon). The two are separate because a conversion is not always linear in the output: dynamic thresholding and sample clipping limit x_0, and a consistency boundary reads a function of it, so combining two raw outputs after conversion is not combining them before it. Guidance therefore reads raw_both, combines the raw outputs and converts once, the order a source pipeline runs its scheduler in.

Denoiser.convert

def convert(x_t, t, output) -> tuple[jax.Array, jax.Array]

(x_0, epsilon) read out of a raw model output at (x_t, t).

Denoiser.raw_both

def raw_both(x_t, t) -> tuple[jax.Array, jax.Array]

The conditional and the unconditional raw outputs, in one model call over the doubled batch.

DenoisingCondition

dataclass source

class DenoisingCondition(
    pooled: jax.Array | None = None,
    time_ids: jax.Array | None = None,
    guidance: jax.Array | None = None,
    mask: jax.Array | None = None,
)

Carries text conditioning as the published families read it.

The fields are the token states, a pooled vector, the size and crop ids the XL towers add, the distilled guidance value a guidance-embedded transformer takes as a model input rather than as two guided branches, and the [B, tokens] mask of the real token states in a right-padded context, which a family that excludes padded keys reads.

DenoisingCondition.aligned

def aligned(given: DenoisingCondition) -> DenoisingCondition

This conditioning with given’s own model inputs.

A distilled guidance value belongs to the row rather than to its caption. Dropping the caption, or guiding against an unconditional one, changes what the model reads about the text and not the scale the checkpoint was distilled to walk at. The two seams that pair a conditional record with an unconditional one align them here first, so both keep each row’s own scalar.

Process

dataclass source

class Process(
    schedule: NoiseScheduler,
    prediction: PredictionTransform,
    weighting: Weighting = ScheduleWeighting(),
    sampling: NoiseScheduler | None = None,
)

Pairs a schedule with what the model predicts on it and how the loss is weighted.

sampling is the schedule inference integrates when it is not the training one, as EDM trains on log-normal sigmas and samples on the Karras grid. None means the same schedule.

Process.weight

def weight(t) -> jax.Array

Process.times

def times(steps: int) -> jax.Array

The descending time grid of steps points a sampler walks, from T to 0.

A tabulated schedule cannot take more steps than it has entries, so steps is capped at T there.

Process.noise

def noise(key, shape) -> jax.Array

Draws the sampling schedule’s Gaussian prior at shape.

Its default scale is the unit-data marginal at T; a schedule may declare a different prior normalization.

Process.denoiser

def denoiser(
    model,
    params,
    conditions: Mapping[str, Conditioning],
    unconditional: Mapping[str, Conditioning] | None = None,
) -> Denoiser

(x_t, t) -> (x_0, epsilon) for model under params with the given conditions, on the sampling schedule.

unconditional carries the same keys with the unconditional values and is what classifier-free guidance interpolates against.

aligned_conditions

function source

def aligned_conditions(
    conditions: Mapping[str, Conditioning],
    unconditional: Mapping[str, Conditioning],
) -> dict[str, Conditioning]

unconditional with each row’s own model inputs taken from conditions.

A keyword names a conditioning record on both sides or on neither, so the spatial keywords an inpainting source adds are arrays and pass through. A keyword that is a record on one side only is a caller pairing two different conditionings, and raises rather than aligning nothing.