class

DiffusionMixin

DiffusionMixin()

source edit

Reverse-process sampling loop for diffusion-family generative models.

Concrete subclasses must:

Inherit from lucid.models.PretrainedModel.
Define forward(sample, timestep, ...) -> DiffusionModelOutput that returns the network's prediction at the supplied timestep.
Expose a config attribute carrying sample_size / in_channels (fields DiffusionModelConfig defines) — or explicitly pass generator_shape to generate.

The mixin then provides generate — a vanilla reverse-process loop parametrised by an externally-supplied DiffusionScheduler. Sampler choice (DDPM, DDIM, …) is thus a runtime knob rather than baked into the model class.

Notes

The mixin is stateless. All hyper-parameters (number of inference steps, image shape, batch size) flow through generate's keyword arguments; the model class itself stays sampler-agnostic.

Examples

>>> from lucid.models import DDPMScheduler, AutoModelForImageGeneration
>>> model = AutoModelForImageGeneration.from_pretrained("ddpm_cifar_gen")
>>> scheduler = DDPMScheduler(num_train_timesteps=1000)
>>> out = model.generate(scheduler, n_samples=4, num_inference_steps=50)
>>> out.samples.shape
(4, 3, 32, 32)

Used by 2

Instance methods

generate

→GenerationOutput

generate(scheduler: DiffusionScheduler, n_samples: int = 1, num_inference_steps: int | None = None, generator_shape: tuple[int, ...] | None = None, return_intermediates: bool = False, device: str = 'cpu')

source edit

Sample n_samples images via the reverse diffusion process.

Parameters

schedulerDiffusionScheduler

Already-constructed scheduler (DDPM, DDIM, etc.). Encapsulates the noise schedule and the per-step update rule.

n_samples(int, optional, keyword - only)= 1

Batch size of the generated output.

num_inference_steps(int or None, optional, keyword - only)= None

If None, the scheduler's timesteps are used as-is; otherwise scheduler.set_timesteps(num_inference_steps, ...) is called first to (re)build the timestep schedule.

generator_shape(tuple[int, ...] or None, optional, keyword - only)= None

Per-sample tensor shape, (C, H, W). When omitted, defaults to (in_channels, H, W) derived from self.config.

return_intermediates(bool, optional, keyword - only)= False

If True, every per-step sample is recorded and returned in GenerationOutput.intermediates.

device(str, optional, keyword - only)= "cpu"

Device to allocate the initial Gaussian noise tensor on.

Returns

GenerationOutput

.samples is the final (n_samples, C, H, W) batch; .intermediates is populated when return_intermediates is True.

Raises

RuntimeError

If generator_shape is None and the model lacks a config with sample_size / in_channels fields.

Notes

Decorated with lucid.no_grad — sampling never builds an autograd graph. Standard reverse-process loop:

Draw $x_T \sim \mathcal{N}(0, I)$ .
For each t in scheduler.timesteps (descending): run the model forward, then apply scheduler.step(model_pred, t, x_t) to obtain $x_{t-1}$ .

Examples

>>> model = AutoModelForImageGeneration.from_pretrained("ddpm_cifar_gen")
>>> scheduler = DDPMScheduler(num_train_timesteps=1000)
>>> out = model.generate(scheduler, n_samples=4, num_inference_steps=50,
...                      device="cpu")
>>> out.samples.shape
(4, 3, 32, 32)

>>> from lucid.models import DDPMScheduler, AutoModelForImageGeneration >>> model = AutoModelForImageGeneration.from_pretrained("ddpm_cifar_gen") >>> scheduler = DDPMScheduler(num_train_timesteps=1000) >>> out = model.generate(scheduler, n_samples=4, num_inference_steps=50) >>> out.samples.shape (4, 3, 32, 32)

generate(scheduler: DiffusionScheduler, n_samples: int = 1, num_inference_steps: int | None = None, generator_shape: tuple[int, ...] | None = None, return_intermediates: bool = False, device: str = 'cpu')

>>> model = AutoModelForImageGeneration.from_pretrained("ddpm_cifar_gen") >>> scheduler = DDPMScheduler(num_train_timesteps=1000) >>> out = model.generate(scheduler, n_samples=4, num_inference_steps=50, ... device="cpu") >>> out.samples.shape (4, 3, 32, 32)