NCSNModel

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. Divergence, deliberate: this matches none of the three published recipes exactly. NCSN v1 also conditions on the index, but through CondInstanceNorm++ per-$\sigma$ gain/bias rather than an additive MLP; NCSNv2's Technique 3 drops noise conditioning entirely and divides the output by $\sigma$; NCSN++ conditions on continuous $\sigma$ via Fourier features of $\log \sigma$. Feeding the index through the timestep MLP keeps _DDPMUNet reusable unchanged, which is what makes one backbone serve both families here. The geometric $\sigma$ table itself is held as a non-persistent buffer so it travels with .to(device=...).