diffusion_models.models.SimpleUnet

Module Contents

class Block(in_ch, out_ch, time_emb_dim, up=False)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

time_mlp

conv2

bnorm1

bnorm2

relu

forward(x, t)

class SinusoidalPositionEmbeddings(dim)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

dim

forward(time)

class SimpleUnet(diffuser, image_channels)

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

down_channels = (64, 128, 256, 512, 1024)

up_channels = (1024, 512, 256, 128, 64)

out_dim

time_emb_dim = 32

time_mlp

conv0

downs

ups

output

forward(x, timestep)

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.