tsgm.models.ddpm

The implementation is based on Keras DDPM implementation: https://keras.io/examples/generative/ddpm/

Module Contents

class GaussianDiffusion(beta_start: float = 0.0001, beta_end: float = 0.02, timesteps: int = 1000)[source]

Gaussian diffusion utility for generating samples using a diffusion process.

This class implements a Gaussian diffusion process, where a sample is gradually perturbed by adding Gaussian noise over a series of timesteps. It also includes methods to reverse the diffusion process, predicting the original data from the noisy samples.

Args:

beta_start (float): Start value of the scheduled variance for the diffusion process. beta_end (float): End value of the scheduled variance for the diffusion process. timesteps (int): Number of timesteps in the forward process.

_extract(a: tensorflow.python.types.core.TensorLike, t: int, x_shape: tensorflow.TensorShape) tensorflow.python.types.core.TensorLike[source]

Extracts coefficients for a specific timestep and reshapes them for broadcasting.

Args:

a: Tensor to extract from. t: Timestep for which the coefficients are to be extracted. x_shape: Shape of the current batched samples.

Returns:

Tensor reshaped to [batch_size, 1, 1] for broadcasting.

q_mean_variance(x_start: tensorflow.python.types.core.TensorLike, t: float) Tuple[source]

Extracts the mean and variance at a specific timestep in the forward diffusion process.

Args:

x_start: Initial sample (before the first diffusion step). t: A timestep.

Returns:

mean, variance, log_variance: Tensors representing the mean, variance, and log variance of the distribution at t.

q_sample(x_start: tensorflow.python.types.core.TensorLike, t: float, noise: float) Tuple[source]

Performs the forward diffusion step by adding Gaussian noise to the sample.

Args:

x_start: Initial sample (before the first diffusion step) t: Current timestep noise: Gaussian noise to be added at timestep t

Returns:

Diffused samples at timestep t

predict_start_from_noise(x_t: tensorflow.python.types.core.TensorLike, t, noise)[source]

Predicts the initial sample from the noisy sample at timestep t.

Args:

x_t: Noisy sample at timestep t. t: Current timestep. noise: Gaussian noise added at timestep t.

Returns:

Predicted initial sample.

q_posterior(x_start, x_t, t)[source]

Computes the mean and variance of the posterior distribution q(x_{t-1} | x_t, x_0).

Args:

x_start: Initial sample (x_0) for the posterior computation. x_t: Sample at timestep t. t: Current timestep.

Returns:

Posterior mean, variance, and clipped log variance at the current timestep.

p_mean_variance(pred_noise, x, t)[source]

Predicts the mean and variance for the reverse diffusion step.

Args:

pred_noise: Noise predicted by the diffusion model. x: Samples at a given timestep for which the noise was predicted. t: Current timestep.

Returns:

model_mean, posterior_variance, posterior_log_variance: Tensors representing the mean and variance of the model at the current timestep.

p_sample(pred_noise, x, t)[source]

Generates a sample from the diffusion model by reversing the diffusion process.

Args:

pred_noise: Noise predicted by the diffusion model. x: Samples at a given timestep for which the noise was predicted. t: Current timestep.

Returns:

Sample generated by reversing the diffusion process at timestep t.

class DDPM(network: tensorflow.keras.Model, ema_network: tensorflow.keras.Model, timesteps: int, ema: float = 0.999)[source]

Bases: tensorflow.keras.Model

Denoising Diffusion Probabilistic Model

Args:

network (keras.Model): A Keras model that predicts the noise added to the images. ema_network (keras.Model): EMA model, a clone of network timesteps (int): The number of timesteps in the diffusion process. ema (float): The decay factor for the EMA, default is 0.999.

train_step(images: tensorflow.python.types.core.TensorLike) Dict[source]

Performs a single training step on a batch of images.

Args:

images: A batch of images to train on.

Returns:

A dictionary containing the loss value for the training step.

generate(n_samples: int = 16) tensorflow.python.types.core.TensorLike[source]

Generates new samples by running the reverse diffusion process.

Args:

n_samples: The number of samples to generate.

Returns:

Generated samples after running the reverse diffusion process.

call(n_samples: int) tensorflow.python.types.core.TensorLike[source]

Calls the generate method to produce samples.

Args:

n_samples: The number of samples to generate.

Returns:

Generated samples.