class

NCSNModel

extendsPretrainedModel

NCSNModel(config: NCSNConfig)

source edit

Score network $s_\theta(\tilde{x}, \sigma_i)$ over a geometric $\sigma$ schedule.

Implements the noise-conditional score network of Song and Ermon, 2019 (and the NCSNv2 / NCSN++ refinements of Song 2020, Song et al. 2021). Reuses DDPMUNet as the score backbone — modern NCSN variants converged on the same U-Net architecture as diffusion — but feeds the integer noise-level index $i \in [0, L)$ rather than a raw $\sigma$ value through the timestep embedding (matching the NCSNv2 / NCSN++ recipe). The geometric $\sigma$ table itself is held as a non-persistent buffer so it travels with .to(device=...).

Parameters

configNCSNConfig

Hyperparameters controlling the U-Net topology (base_channels, channel_mult, num_res_blocks, attention_resolutions, num_heads, dropout, resnet_groups) and the score-based extras (num_noise_levels, sigma_max, sigma_min).

Attributes

unetDDPMUNet

Shared U-Net backbone instantiated from a derived DDPMConfig.

sigmasTensor

Buffer of shape (num_noise_levels,) holding the geometric

\sigma_1 > \sigma_2 > \cdots > \sigma_L

schedule.

Notes

Reference: Song and Ermon, "Generative Modeling by Estimating Gradients of the Data Distribution", NeurIPS, 2019 (arXiv:1907.05600); NCSNv2 refinements in Song and Ermon, 2020 (arXiv:2006.09011).

The trained score satisfies

s_\theta(\tilde{x}, \sigma) \;\approx\; \nabla_{\tilde{x}} \log p_\sigma(\tilde{x}),

where $p_\sigma$ is the data distribution convolved with a Gaussian of standard deviation $\sigma$ .

Examples

>>> import lucid
>>> from lucid.models.generative.ncsn import NCSNConfig, NCSNModel
>>> cfg = NCSNConfig(sample_size=32, base_channels=32,
...                  channel_mult=(1, 2), num_res_blocks=1,
...                  resnet_groups=16, num_noise_levels=10)
>>> model = NCSNModel(cfg).eval()
>>> x_tilde = lucid.randn((1, 3, 32, 32))
>>> sigma_idx = lucid.tensor([3]).long()
>>> out = model(x_tilde, sigma_idx)
>>> out.sample.shape   # (1, 3, 32, 32) — raw score
(1, 3, 32, 32)

Used by 2

Constructors

dunder

init

→None

__init__(config: NCSNConfig)

source edit

Properties

prop

num_noise_levels

→int

num_noise_levels: int

source edit

Instance methods

forward

→DiffusionModelOutput

forward(sample: Tensor, sigma_idx: Tensor)

source edit

Predict the score at sample for the given noise-level indices.

Parameters

sampleTensor

(B, in_channels, H, W) perturbed image x̃.

sigma_idxTensor

(B,) integer indices into the σ schedule.

Returns

DiffusionModelOutput

class:DiffusionModelOutput whose sample field is the raw

>>> import lucid >>> from lucid.models.generative.ncsn import NCSNConfig, NCSNModel >>> cfg = NCSNConfig(sample_size=32, base_channels=32, ... channel_mult=(1, 2), num_res_blocks=1, ... resnet_groups=16, num_noise_levels=10) >>> model = NCSNModel(cfg).eval() >>> x_tilde = lucid.randn((1, 3, 32, 32)) >>> sigma_idx = lucid.tensor([3]).long() >>> out = model(x_tilde, sigma_idx) >>> out.sample.shape # (1, 3, 32, 32) — raw score (1, 3, 32, 32)