azula.plugins.jit

Just Image Transformer (JIT) plugin.

from azula.plugins import jit

References

Back to Basics: Let Denoising Generative Models Denoise (Li et al., 2025)

https://arxiv.org/abs/2511.13720

Classes

JITDenoiser

Creates a JIT denoiser.

Functions

load_model

Loads a pre-trained JIT denoiser.

Descriptions

class azula.plugins.jit.JITDenoiser(backbone, schedule=None, num_classes=1000)

Creates a JIT denoiser.

Parameters:

• backbone (Module) – A time conditional network.

• schedule (Schedule) – A noise schedule. If None, use azula.noise.RectifiedSchedule instead.

• num_classes (int) – The number of classes.

forward(x_t, t, label=None, **kwargs)

Parameters:

• x_t (Tensor) – A noisy tensor \(x_t\), with shape \((B, 3, H, W)\).

• t (Tensor) – The time \(t\), with shape \(()\) or \((B)\).

• label (Tensor | None) – The class label \(c\) as an integer, with shape \((B)\).

• kwargs – Optional keyword arguments.

Returns:

The Dirac delta \(\delta(X - \mu_\phi(x_t \mid c))\).

Return type:

DiracPosterior

azula.plugins.jit.load_model(name, ema=True, **kwargs)

Loads a pre-trained JIT denoiser.

Parameters:

• name (str) – The pre-trained model name.

• ema (bool) – Whether to load EMA weights or not.

• kwargs – Keyword arguments passed to torch.load.

Returns:

A pre-trained denoiser.

Return type:

Denoiser