multigeom.diffusion.model

Conditional diffusion model: a UNet2DConditionModel from diffusers augmented with a small response-to-token MLP.

The conditioning vector r ∈ ℝ^{50} is mapped through a 3-layer SiLU MLP to a
single 256-d token, which is supplied to UNet2DConditionModel as the
single-token `encoder_hidden_states` input. Classifier-free guidance dropout
on r is applied at the *training* level (`multigeom.diffusion.train`), not inside the
forward pass — this keeps inference / G(x) estimation independent of any
training-only randomness.

Functions

build_diffusion_model(cfg)	Construct a NeuralConditionedUNet from a diffusion config namespace.
build_scheduler(cfg)	Construct the DDPM scheduler with the paper's linear β schedule.

Classes

ConditioningMLP(*args, **kwargs)	Maps a [B, n_neurons] response vector to a [B, 1, embedding_dim] token, suitable as encoder_hidden_states for a single-token cross-attention block.
NeuralConditionedUNet(*args, **kwargs)	Conditional UNet for the multi-scale-Fisher diffusion estimator.