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 timestept
.- 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.