dew.objectives.diffusion

Name	Summary
VALIDATION_SAMPLES
BlockDiffusionObjective	Fine-tune DiffusionGemma on clean text rows split into prompt and canvases.
DiffusionObjective	Denoising diffusion: sample a noise level, corrupt, predict, weight.
DiffusionRunConfig	Describe a run, plus the diffusion objective’s own knobs.
MaskedDiffusionObjective	Train a masked diffusion model on the MDLM negative ELBO.
StableDiffusionAutoencoder	Run latent diffusion behind the vendored Stable Diffusion VAE.
TextCondition	Condition the model on text: which registered encoder reads the batch’s tokens, and from which checkpoint.

VALIDATION_SAMPLES

attribute source

VALIDATION_SAMPLES = 4

BlockDiffusionObjective

class source

class BlockDiffusionObjective(
    model: DiffusionGemma,
    *,
    prompt_length: int,
    num_canvases: int = 1,
    canvas_size: int | None = None,
    pretrained: Variables | None = None,
    pad_token_id: int = 0,
    self_cond_prob: float = 0.5,
    safety_epsilon: float = 0.0001,
    stop_gradient_from_denoiser_to_encoder: bool = False,
    encoder_loss_weight: float = 1.0,
    decoder_loss_weight: float = 1.0,
    ema_decay: float | None = None,
    trainable: PathFilter | None = None,
    head_chunks: int = 4,
)

Fine-tune DiffusionGemma on clean text rows split into prompt and canvases.

A row has prompt_length + canvas_size * num_canvases tokens. text accepts token arrays or ModelInputs; media conditions only the clean encoder. Supplied attention validity controls cache occupancy, otherwise the pad ID and canvas mask do. Optional canvas_mask and encoder_target_mask select text targets; media placeholders are never labels. Default encoder targets require adjacent valid slots, matching Google’s SequenceTargetShift. Every response token is corrupted, but only a uniformly selected valid canvas contributes diffusion CE.

trainable selects the parameter leaves the optimizer moves, by their full path (dew.objectives.base.PathFilter), the way LMObjective takes it; the rest of the tree is kept under frozen, which init returns and a checkpoint stores. An adapter’s own filter (dew.lora.LoRA.trainable) goes here. None trains every leaf.

Both cross-entropies score the final states through the bounded head (dew.objectives.lm.chunked.chunked_cross_entropy), head_chunks vocabulary tiles at a time, so no vocabulary-sized fp32 logits or softmax of a whole row is held for the backward pass. The first denoising pass, whose logits condition the second and carry no gradient, is the one place a full row of logits exists.

bank_sites: tuple[DecoderBank, ...]

Name the shared text stack, as the training model declares it.

BlockDiffusionObjective.pipeline

def pipeline(
    state: TrainState,
    *,
    ema: bool = True,
    processor: Processor | None = None,
) -> BlockGeneration

Publish the state’s weights as a BlockGeneration task.

The sampler keeps the published defaults, and the tokenizer’s EOS ids are the caller’s to set.

BlockDiffusionObjective.held_variables

def held_variables() -> Variables | None

Return the SFT source this objective starts from.

BlockDiffusionObjective.init

def init(key: jax.Array, variables: Variables | None = None) -> Variables

BlockDiffusionObjective.loss

def loss(params: Variables, batch: Batch, step: Step)

BlockDiffusionObjective.evaluate

def evaluate(params: Variables, batch: Batch, step: Step) -> TokenScores

Score the denoiser’s cross entropy on every canvas target of the batch.

One noise level and one canvas per row are drawn from the pass’s key, as training draws them, with dropout off and the averaged weights when the run keeps them, so perplexity over a validation pass is exp of the denoising loss per target.

BlockDiffusionObjective.reduce_loss

def reduce_loss(stats: BlockSFTStatistics)

DiffusionObjective

class source

class DiffusionObjective(
    model: nn.Module,
    process: Process,
    inputs: InputSpec,
    *,
    autoencoder: AutoEncoder | None = None,
    unconditional_prob: float = 0.12,
    ema_decay: float | None = 0.999,
    sampler: Solver = DDIM(),
    guidance: CFG | None = CFG(3.0),
    steps: int = 200,
    pretrained: Variables | None = None,
)

Denoising diffusion: sample a noise level, corrupt, predict, weight.

latent_shape: tuple[int, ...]

The per-example shape the model denoises: the sample field’s, or its latent when an autoencoder sits in front of the model.

DiffusionObjective.pipeline

def pipeline(state: TrainState, *, ema: bool = True) -> TextToImage

The model over the state’s published weights as a TextToImage task, sampling the way this objective’s evaluation does.

DiffusionObjective.encoder_params

def encoder_params() -> dict

DiffusionObjective.encode

def encode(encoders, tokens: dict | None = None) -> dict

Encode conditions under the supplied parameters; omitted tokens select each condition’s configured unconditional datum.

DiffusionObjective.blank_conditions

def blank_conditions(like: dict) -> dict

Cast the stored unconditional conditions to the conditional branch’s dtypes.

like is the conditional branch. The values themselves were encoded once at construction, so this only changes dtype.

DiffusionObjective.held_variables

def held_variables() -> Variables

Every array init starts from rather than draws: a whole pretrained tree, or the frozen towers.

A text tower and a VAE are released weights: hundreds of megabytes that a nullary trace would compile into the state executable as constants. One mapping, so an objective that starts from more than the towers extends this and init together.

DiffusionObjective.init

def init(key, variables: Variables | None = None) -> Variables

DiffusionObjective.trainable

def trainable(params) -> dict

Return the model’s own collections, without the frozen towers.

DiffusionObjective.encoded_conditions

def encoded_conditions(params, batch) -> dict

Encode each condition’s own batch field under the tree’s frozen towers.

DiffusionObjective.denoiser

def denoiser(params, given, unconditional)

Build the process’s denoiser over the model’s own collections.

The unconditional branch is passed only when this objective is guided; without guidance the sampler never evaluates it.

DiffusionObjective.loss

def loss(params, batch, step: Step)

DiffusionObjective.evaluate

def evaluate(params, batch, step: Step)

One generated sample for every real row, without display decoding.

DiffusionObjective.preview

def preview(params, batch, step: Step, *, scored=None)

A separate small draw for display, with root-only caption decoding.

DiffusionRunConfig

dataclass source

class DiffusionRunConfig(
    model: ModelConfig = (lambda: ModelConfig('unet', dict(DEFAULT_MODEL_CONFIG)))(),
    data: CaptionedSpec = OxfordFlowers(),
    optim: OptimConfig = OptimConfig(),
    trainer: TrainerConfig = TrainerConfig(),
    objective: str = 'diffusion',
    lora: LoRA | None = None,
    preset: PresetSpec = presets.EDM(),
    sampler: SamplerSpec = samplers.EulerAncestral(),
    guidance: CFG | None = (lambda: CFG(3.0))(),
    sampling_steps: int = 200,
    unconditional_prob: float = 0.12,
    ema_decay: float | None = 0.999,
    text: TextCondition | None = TextCondition(),
    autoencoder: StableDiffusionAutoencoder | None = None,
    val_metrics: tuple[str, ...] = ('clip',),
)

Describe a run, plus the diffusion objective’s own knobs.

preset: PresetSpec

The convention the model is trained and sampled with.

sampler: SamplerSpec

The solver validation samples with.

guidance: CFG | None

How validation samples are guided, scale and interval; None samples the conditional prediction alone.

unconditional_prob: float

Fraction of training examples whose condition is dropped.

ema_decay: float | None

None disables EMA; 1.0 retains a frozen copy.

text: TextCondition | None

The text condition, under the models’ textcontext keyword; None trains unconditionally.

autoencoder: StableDiffusionAutoencoder | None

Set for latent diffusion; None trains in pixel space.

val_metrics: tuple[str, ...]

Names in the metrics registry, scored on every validation pass. The registry is the list of what a run can name, so a metric registered elsewhere is spelled here without this class knowing it; __post_init__ refuses a name nothing is registered under.

parameter_roots: tuple[tuple[str, ...], ...]

Parameter ownership in the variables tree this config builds.

DiffusionRunConfig.sample_field

def sample_field() -> Field

The batch field the model generates, at the resolution the data comes in.

DiffusionRunConfig.model_fields

def model_fields(autoencoder: AutoEncoder | None) -> dict

The fields the registry builds the model from: the run’s precision settings and the channels the model denoises, over model.config.

DiffusionRunConfig.build

def build(*, variables: Variables | None = None) -> DiffusionObjective

Build the configured compute owners around their parameters.

Supplied variables are the authoritative saved snapshot. Encoders and the VAE read only configuration/tokenizer metadata and bind their respective subtrees without a source weight load or storage cast.

DiffusionRunConfig.build_eval_metrics

def build_eval_metrics() -> list

Validation metrics for val_metrics, each pulling its own weights on construction. A video run scores VideoGrid against its video field, so psnr and ssim read that grid there, and an image-only metric raises a ValueError naming it here, before the trainer.

MaskedDiffusionObjective

class source

class MaskedDiffusionObjective(
    model: CausalTransformer,
    process: DiscreteProcess,
    seq_len: int,
    *,
    head_chunks: int = 4,
    ema_decay: float | None = 0.999,
    sampler: Unmask = Unmask(),
    steps: int = MDLM_STEPS,
    samples: int = 4,
    decode: Callable[[Sequence[int]], str] | None = None,
    pretrained: Variables | None = None,
)

Train a masked diffusion model on the MDLM negative ELBO.

The rows are [B, seq_len] token ids under batch["text"].

MaskedDiffusionObjective.pipeline

def pipeline(
    state: TrainState,
    *,
    ema: bool = True,
    processor: Processor | None = None,
) -> MaskedGeneration

Publish the state’s weights as a native full-response MDLM task.

MaskedDiffusionObjective.held_variables

def held_variables() -> Variables | None

Return the checkpoint this run continues from, or None for a fresh init.

MaskedDiffusionObjective.init

def init(key, variables: Variables | None = None)

MaskedDiffusionObjective.loss

def loss(params, batch, step: Step)

MaskedDiffusionObjective.evaluate

def evaluate(params, batch, step: Step) -> TokenScores

Score the negative ELBO of every token in the batch.

One noise level and one masking are drawn from the pass’s key, as training draws them, with dropout off and the averaged weights when the run keeps them. Every token counts and carries its weighted masked cross entropy, zero where it was left visible, so perplexity over a validation pass is exp of the ELBO bound per token, the number MDLM reports.

MaskedDiffusionObjective.preview

def preview(params, batch, step: Step, *, scored=None)

Generate the configured display count, then decode on process zero.

StableDiffusionAutoencoder

dataclass source

class StableDiffusionAutoencoder(
    modelname: str = 'pcuenq/sd-vae-ft-mse-flax',
    revision: str = 'bf16',
    dtype: DtypeName = 'bfloat16',
    latent_shift: float | None = None,
    latent_scale: float | None = None,
)

Run latent diffusion behind the vendored Stable Diffusion VAE.

latent_scale: float | None

Per-dataset latent statistics; None keeps the checkpoint’s.

StableDiffusionAutoencoder.build

def build(*, params: Variables | None = None) -> AutoEncoder

Bind supplied VAE params while reconstructing its model metadata.

TextCondition

dataclass source

class TextCondition(
    encoder: str = 'clip_text',
    checkpoint: str = DEFAULT_MODEL,
    dtype: DtypeName | None = None,
    param_dtype: DtypeName = 'float32',
    field: str = 'text',
    unconditional: str = '',
    max_length: int | None = None,
    revision: str | None = None,
)

Condition the model on text: which registered encoder reads the batch’s tokens, and from which checkpoint.

dtype: DtypeName | None

The encoder’s compute dtype; None follows the model’s compute dtype.

param_dtype: DtypeName

Storage precision when loading source weights; supplied params retain theirs.

field: str

The batch field holding the tokenized text.

unconditional: str

The prompt the unconditional branch is encoded from.

max_length: int | None

Tokens every prompt is padded to; None keeps the encoder’s own default, which for CLIP is the checkpoint’s context length.

revision: str | None

The checkpoint’s git revision. A rerun then conditions on the weights the run named, even after the branch has moved on.

TextCondition.build

def build(
    *,
    params: Variables | None = None,
    dtype: DtypeName | None = None,
) -> Condition

Bind supplied encoder params without a source weight load or storage cast.

dtype is the run’s own compute dtype, which an unset self.dtype follows: the tower runs beside the model it conditions, in every step, and a checkpoint stored in float32 is no reason to run it there.