A U-Net is a convolutional network with a contracting encoder, a bottleneck, and an expansive decoder. For diffusion-based image generation, its job is to predict the noise in an image at a given time step. That prediction lets the generation process gradually turn random noise into an image. section 15.1.2 section 15.4
The encoder downsamples the image to capture features at different scales; the decoder upsamples those features. Skip connections join encoder and decoder feature maps of the same spatial size, helping retain edges and textures that downsampling might lose. Compared with a simple convolutional network, this design captures both local and global image features more effectively for denoising. section 15.1.2 section 15.6
Figure 15.2 The architecture of the denoising U-Net model. The U-Net architecture is characterized by its symmetric shape, with a contracting path (encoder) and an expansive path (decoder), connected by a bottleneck layer. The model is designed to remove noise from images while preserving important details. The input to the model is a noisy image, along with which time step the image is in, and the output is the predicted noise in the image.
Use the architecture diagram to trace the noisy image and time step through the encoder, bottleneck, and decoder to the predicted noise. section 15.1.2
To train a U-Net on flower images:
- Resize Oxford 102 Flower images to 64 × 64 pixels and create noisy versions at selected time steps. section 15.1.3 section 15.2.1
- Give the noisy images and their time steps to the model. Compare its predicted noise with the noise actually added, then update its weights to reduce pixel-wise L1 loss. section 15.1.3 section 15.4.1
Time-step input matters: the balance between clean image and noise changes as diffusion progresses. In the PyTorch UNet class, sine-and-cosine time embeddings are expanded to match image features and added within the network. Attention also helps the model focus on relevant image features; it treats flattened pixels as a sequence to learn dependencies between them. section 15.3.2 section 15.3
Listing 15.4 Defining the UNet() class
class UNet(nn.Module):
…
def forward(self, sample, timesteps): #1
if not torch.is_tensor(timesteps):
timesteps = torch.tensor([timesteps],
dtype=torch.long,
device=sample.device)
timesteps = torch.flatten(timesteps)
timesteps = timesteps.broadcast_to(sample.shape[0])
t_emb = sinusoidal_embedding(timesteps, self.hidden_dims[0])
t_emb = self.time_embedding(t_emb) #2
x = self.init_conv(sample)
r = x.clone()
skips = []
for block1, block2, attn, downsample in self.down_blocks: #3
x = block1(x, t_emb)
skips.append(x)
x = block2(x, t_emb)
x = attn(x)
skips.append(x)
x = downsample(x)
x = self.mid_block1(x, t_emb)
x = self.mid_attn(x)
x = self.mid_block2(x, t_emb) #4
for block1, block2, attn, upsample in self.up_blocks:
x = torch.cat((x, skips.pop()), dim=1) #5
x = block1(x, t_emb)
x = torch.cat((x, skips.pop()), dim=1)
x = block2(x, t_emb)
x = attn(x)
x = upsample(x)
x = self.out_block(torch.cat((x, r), dim=1), t_emb)
out = self.conv_out(x)
return {"sample": out} #6
The partial forward method shows where time steps become embeddings and where down blocks begin processing image features. section 15.3.2
Generate a flower from noise
A PyTorch setup instantiates UNet with three input channels, hidden dimensions [128, 256, 512, 1024], and 64-pixel images—about 133 million parameters. After training, start with random noise and repeatedly feed the evolving image to the model; the example uses 50 inference steps. This flower generator uses image and time-step inputs, not a text prompt. Text-to-image systems can additionally condition a U-Net denoiser on an encoded prompt. section 15.3.2 section 15.4 section 15.5.1
