diffusion_models.models.SimpleUnet

Module Contents

class Block(in_ch, out_ch, time_emb_dim, up=False)[source]
time_mlp[source]
conv2[source]
bnorm1[source]
bnorm2[source]
relu[source]
forward(x, t)[source]
class SinusoidalPositionEmbeddings(dim)[source]
dim[source]
forward(time)[source]
class SimpleUnet(diffuser, image_channels)[source]

A Simplified variant of the Unet architecture used in DDPM.

Parameters:

  • diffuser (GaussianDiffuser) – A gaussian diffuser.

  • image_channels (int) – The number of image channels.

image_channels[source]
down_channels = (64, 128, 256, 512, 1024)[source]
up_channels = (1024, 512, 256, 128, 64)[source]
out_dim[source]
time_emb_dim = 32[source]
time_mlp[source]
conv0[source]
downs[source]
ups[source]
output[source]
forward(x, timestep)[source]

Forward pass of the diffusion model.

The forward pass of the diffusion model, predicting the noise at a single step.

Parameters:

  • x – A batch of noisy images.

  • timestep – The timesteps of each image in the batch.

Returns:

A tensor representing the noise predicted for each image.