NCSNpp.py¶

Source: LION/models/diffusion/NCSNpp.py

NCSN++ diffusion denoiser adapted for LION and PaDIS.

The network originates from the score-SDE NCSN++ implementation by Song et al. LION adds geometry-aware channel selection, PaDIS position channels, paper-compatible embedding/attention behaviour, and model presets used by the reproduction pipeline.

LION.models.diffusion.NCSNpp.score_from_denoiser(noisy_image_patch, denoised_patch, sigma)[source]¶

Convert a denoised prediction into a VE score estimate.

Parameters:

noisy_image_patch (Tensor)
denoised_patch (Tensor)
sigma (Tensor)

Return type:

Tensor

class LION.models.diffusion.NCSNpp.NCSNpp(params, geometry)[source]¶

Bases: LIONmodel

Geometry-aware NCSN++ denoiser used by PaDIS.

Parameters:

params (LIONModelParameter) – Architecture and PaDIS compatibility settings. Use default_parameters() to construct a supported preset.
geometry (LION.CTtools.ct_geometry.Geometry) – Geometry providing image channels and spatial dimensions.

Notes

The network predicts a denoised image. Score conversion is kept separate so training and reconstruction can use the same output convention.

static default_parameters(mode='padis-paper-ct-256')[source]¶

Return model parameters for a supported PaDIS prior preset.

Parameters:: mode (str, optional) – Patch, whole-image, resolution, or patch-size-ablation preset. A -no-position suffix disables absolute position channels.
Returns:: Fully populated NCSN++ architecture settings.
Return type:: LIONModelParameter

forward(x, time_cond, class_labels=None, augment_labels=None)[source]¶

Evaluate the denoiser at one or more continuous noise levels.

Parameters:

x (torch.Tensor) – Noisy NCHW inputs, including position channels when configured.
time_cond (torch.Tensor) – Per-sample sigma values or discrete scale indices.
class_labels (optional) – Reserved compatibility arguments; PaDIS does not condition on class labels.
augment_labels (optional) – Reserved compatibility arguments; PaDIS does not condition on class labels.

Returns:

Denoised image channels in the same spatial layout as x.

Return type:

torch.Tensor