diffusion_models.gaussian_diffusion.gaussian_diffuser

Module Contents

class GaussianDiffuser(beta_scheduler)

Initializes the class instance.

Parameters:

beta_scheduler (BaseBetaScheduler) – The beta scheduler instance to be used.

device: str = 'cpu'

The device to use. Defaults to cpu.

classmethod from_checkpoint(checkpoint)

Instantiate a Gaussian Diffuser from a training checkpoint.

Parameters:

checkpoint (Checkpoint) – The training checkpoint object containing the trained model’s parameters and configuration.

Returns:

An instance of the GaussianDiffuser class initialized with the parameters loaded from the given checkpoint.

Return type:

GaussianDiffuser

to(device='cpu')

Moves the data to the specified device.

This performs a similar behaviour to the to method of PyTorch. moving the GaussianDiffuser and the BetaScheduler to the specified device.

Parameters:

device (str) – The device to which the method should move the object. Default is “cpu”.

Example

>>> gaussian_diffuser = GaussianDiffuser()
>>> gaussian_diffuser = gaussian_diffuser.to(device="cuda")

diffuse_batch(images)

Diffuse a batch of images.

Diffuse the given batch of images by adding noise based on the beta scheduler.

Parameters:

images (Tensor) –

Batch of images to diffuse.

Shape should be (B, C, H, W).

Returns:

A tuple containing three tensors

• images: Diffused batch of images.

• noise: Noise added to the images.

• timesteps: Timesteps used for diffusion.

Return type:

Tuple[Tensor, Tensor, Tensor]

denoise_batch(images, model)

Denoise a batch of images.

This denoises a batch images. This is the image generation process.

Parameters:

• images (Tensor) – A batch of noisy images.

• model (diffusion_models.models.base_diffusion_model.BaseDiffusionModel) – The model to be used for denoising.

Returns:

A list of tensors containing a batch of denoised images.

Return type:

List[Tensor]