← Back to list

Building a Diffusion Model from Scratch: Generating Faces with DDPM in PyTorch

How I implemented a Denoising Diffusion Probabilistic Model on CelebA-HQ, FFHQ, and WikiArt — without using any pretrained pipelines

Ghayas · 2026-05-02 14:14 · 1 claps · 7.6 min read
#ddpm #diffusion-models #generate-faces-with-ddpm
Open on Medium ↗
Wiki topics: MM · Multimodal & Generative Media ML · Machine Learning

Building a Diffusion Model from Scratch: Generating Faces with DDPM in PyTorch

How I implemented a Denoising Diffusion Probabilistic Model on CelebA-HQ, FFHQ, and WikiArt — without using any pretrained pipelines

“What if you could teach a neural network to sculpt faces out of pure noise?” That’s exactly what diffusion models do — and this post walks you through building one from scratch.

Introduction

Generative AI has been one of the most explosive areas of machine learning in recent years. Among the many generative approaches — GANs, VAEs, flow-based models — Diffusion Models have quietly taken the crown for image quality. They power tools like DALL·E 2, Stable Diffusion, and Imagen.

But most tutorials just hand you a pretrained pipeline and say “call .generate()." That's not learning. That's copy-paste.

In this post, I’ll walk you through how I built a Denoising Diffusion Probabilistic Model (DDPM) completely from scratch using base PyTorch — no HuggingFace Diffusers, no pretrained weights, no shortcuts. Just math, code, and a lot of GPUs.

This was part of my Generative AI (AI4009) Assignment at FAST-NUCES, and I’m sharing everything: the architecture decisions, the training tricks, the results, and the lessons learned.

What Is a Diffusion Model?

Before we dive into code, let’s understand the core idea intuitively.

Imagine you have a beautiful photograph. Now imagine gradually adding tiny amounts of random noise to it — step by step — until after 300 steps, it’s just white noise. That’s the forward process.

Now flip it. What if a neural network could reverse that process? Given a noisy image at step t, it predicts what the noise looked like and subtracts it. After 300 reverse steps, starting from pure Gaussian noise, it produces a realistic image. That’s the reverse process — and that’s what we’re training.

Pure Noise → [Step 299] → [Step 250] → ... → [Step 1] → Generated Image

Mathematically, the forward process is defined as:

q(x_t | x_{t-1}) = N(x_t; √(1-β_t) * x_{t-1}, β_t * I)

And the full noising in one shot (no need to iterate) uses the closed form:

x_t = √(ᾱ_t) * x_0 + √(1 - ᾱ_t) * ε,   where ε ~ N(0, I)

The model learns to predict ε given x_t and t. That's it. That's the whole idea.

The Dataset

I worked with three datasets for this assignment:

1. CelebA-HQ 256

High-resolution celebrity faces at 256×256, split into train/ and valid/ folders. Perfect for learning facial structure — eyes, nose, hair, skin tones.

2. FFHQ (Flickr-Faces-HQ) Thumbnails

70,000 face images at 128×128. A flat folder of sequentially numbered PNGs (00000.png, 00001.png, ...). More diverse than CelebA — includes people of all ages, ethnicities, and expressions.

3. WikiArt

Over 80,000 artwork images organized by art style (Impressionism, Cubism, Baroque, etc.). The most challenging dataset — highly varied textures, colors, and compositions.

All images were normalized to [-1, 1] for stable training:

python

transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])

The Architecture: A Simplified U-Net

The denoising network is a U-Net — an encoder-decoder with skip connections. The key insight is that the network needs to understand both the local pixel details and the global structure to predict noise accurately.

Channel Progression: 64 → 128 → 256

Input (3, 128, 128)
    ↓ init_conv
(64, 128, 128) ── enc1a, enc1b ──────────────────────────────┐ skip
    ↓ Downsample                                              │
(64, 64, 64) → enc2a, enc2b, SelfAttn ──────────────────────┐│ skip
    ↓ Downsample                                             ││
(128, 32, 32) → enc3a, enc3b, SelfAttn ────────────────────┐││ skip
    ↓ Downsample                                            │││
(256, 16, 16) → Bottleneck (mid1, SelfAttn, mid2)          │││
    ↑ Upsample                                              │││
(256+256, 32, 32) → dec3a, dec3b ←─────────────────────────┘││
    ↑ Upsample                                               ││
(256+128, 64, 64) → dec2a, dec2b, SelfAttn ←────────────────┘│
    ↑ Upsample                                                │
(128+64, 128, 128) → dec1a, dec1b ←──────────────────────────┘
    ↓ out_norm, out_conv
Output (3, 128, 128) — predicted noise

Time Step Embedding

Every residual block is conditioned on the current timestep t. We use sinusoidal positional embeddings — the same idea as in transformers — to encode time:

python

class SinusoidalTimeEmbedding(nn.Module):
    def forward(self, t):
        half  = self.dim // 2
        freqs = torch.exp(-math.log(10000) * torch.arange(half) / (half - 1))
        args  = t.float().unsqueeze(1) * freqs.unsqueeze(0)
        return torch.cat([args.sin(), args.cos()], dim=-1)

This embedding is passed through a small MLP and added to every residual block’s feature map — allowing the network to behave differently at different noise levels.

Residual Blocks

Each block follows the pattern: GroupNorm → SiLU → Conv → add time embedding → GroupNorm → Dropout → Conv → skip connection

python

class ResBlock(nn.Module):
    def forward(self, x, t_emb):
        h = self.conv1(F.silu(self.norm1(x)))
        h = h + self.time_mlp(t_emb)[:, :, None, None]
        h = self.conv2(self.drop(F.silu(self.norm2(h))))
        return h + self.skip(x)

Self-Attention

At the 128-channel and 256-channel resolution levels, we add multi-head self-attention blocks. These allow the model to capture long-range dependencies — critical for generating globally coherent faces.

The Noise Schedule

We use a linear beta schedule:

python

betas = torch.linspace(beta_start=1e-4, beta_end=0.02, steps=T)

From betas, we precompute all the quantities needed for both forward and reverse processes:

  • alphas = 1 - betas
  • alphas_cumprod = ∏ alphas (cumulative product)
  • sqrt_alphas_cumprod and sqrt_one_minus_alphas_cumprod — for forward noising
  • posterior_variance — for reverse denoising

With 300 timesteps, the image goes from clean to near-pure noise smoothly, giving the model a good learning signal at every step.

Training

Loss Function

The training objective is beautifully simple — MSE between predicted and actual noise:

python

x_t, noise = schedule.q_sample(x0, t)          # add noise
pred_noise  = model(x_t, t)                    # predict it
loss        = F.mse_loss(pred_noise, noise)    # minimize difference

Optimizations for Kaggle T4×2

Training a diffusion model is computationally demanding. Here’s what made it feasible on Kaggle’s dual T4 setup:

1. Mixed Precision Training

python

with autocast(enabled=True):
    pred_noise = model(x_t, t)
    loss = F.mse_loss(pred_noise, noise)
scaler.scale(loss).backward()

This halves memory usage and roughly doubles throughput on T4 GPUs.

2. DataParallel for Dual GPU

python

if torch.cuda.device_count() > 1:
    model = nn.DataParallel(model)

Both T4s share the batch load — effectively doubling throughput.

3. Gradient Clipping

python

torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

Prevents occasional gradient explosions that can destabilize training.

4. Cosine LR Scheduler

python

scheduler = CosineAnnealingLR(optimizer, T_max=cfg.EPOCHS)

Smoothly decays learning rate, helping convergence in later epochs.

Training Dynamics

The loss curve tells the story of learning. In the first few epochs, the model learns the easy stuff — removing heavy noise at high timesteps. Gradually it learns the harder task: recovering fine details from slightly noisy images.

A healthy training run shows:

  • Rapid initial drop in the first 5 epochs
  • Steady, smooth decrease through the middle epochs
  • Plateau near convergence in the final epochs

If you see spikes or instability, gradient clipping and reducing the learning rate usually fix it.

Forward Diffusion Visualization

One of the most satisfying things to visualize is the forward process — watching a clean face dissolve into noise:

t=0    → Clear face, full detail
t=50   → Slight grain, structure intact
t=100  → Obvious noise, rough shapes visible
t=150  → Heavy noise, only faint outlines
t=200  → Dominated by noise, barely recognizable
t=250  → Near-pure noise
t=299  → Pure Gaussian noise

The beauty here is that the model never actually sees this sequence during training — it just sees random (x_t, t, noise) triples. Yet it learns the entire spectrum implicitly.

Image Generation: From Noise to Face

Sampling works by running the reverse process for all T steps:

python

x = torch.randn(n, 3, 128, 128)   # start from pure noise
for t in reversed(range(0, T)):
    pred_noise = model(x, t)
    mean = sqrt_recip_alpha * (x - beta / sqrt_one_minus_alpha_cumprod * pred_noise)
    x    = mean + posterior_std * torch.randn_like(x)   # add small noise (except t=0)

Each step refines the image slightly. After 300 steps, what started as pure static has become a realistic face.

Image Reconstruction: The Core Task

Beyond generation, we also tested image reconstruction — given a target image, corrupt it to near-pure noise, then denoise back:

  1. Take a real face from the dataset
  2. Apply forward diffusion at t = T-1 (maximum corruption)
  3. Run the full reverse process
  4. Compare the result to the original

The reconstructed image won’t be pixel-perfect — diffusion models are stochastic — but it should capture the same general structure: similar facial proportions, similar colors, similar style.

Quantitative Evaluation

PSNR (Peak Signal-to-Noise Ratio)

Measures reconstruction fidelity in dB. Higher is better. For diffusion model reconstruction:

  • > 20 dB — good reconstruction
  • > 25 dB — excellent reconstruction

SSIM (Structural Similarity Index)

Ranges from 0 to 1. Measures perceptual similarity accounting for luminance, contrast, and structure.

  • > 0.7 — structurally similar
  • > 0.85 — high structural fidelity

These metrics are computed using torchmetrics:

python

psnr = PeakSignalNoiseRatio(data_range=1.0)
ssim = StructuralSimilarityIndexMeasure(data_range=1.0)
print(f"PSNR : {psnr(recon, target):.2f} dB")
print(f"SSIM : {ssim(recon, target):.4f}")

The Gradio App

To make the model interactive, we built a Gradio app that:

  1. Starts from pure Gaussian noise
  2. Runs the full reverse diffusion
  3. Returns the generated image + intermediate denoising steps

python

def generate_with_steps(num_images):
    imgs, steps = sample_images(model, num_images, ...)
    return final_image, denoising_grid, all_generated
demo = gr.Blocks()
# ... slider for num_images, button, output panels
demo.launch(share=True)

The share=True flag creates a public Gradio link — perfect for sharing your results without any deployment infrastructure.

Key Lessons Learned

1. Start small, scale up. Begin with T=100 timesteps and 128×128 images. Verify the loss decreases before committing to a full 300-step, 256×256 run.

2. GroupNorm over BatchNorm. With small batch sizes (16–32), BatchNorm statistics are noisy. GroupNorm is stable regardless of batch size.

3. SiLU (Swish) beats ReLU here. The smooth activation function helps gradients flow better through deep U-Nets.

4. The noise schedule matters. Linear schedules work, but cosine schedules (from Improved DDPM) give smoother transitions and often better final quality.

5. Mixed precision is non-negotiable. Without it, you either run out of memory or train 2× slower on T4s.

6. Patience. Diffusion models need more epochs than GANs to produce good samples. Don’t judge the model at epoch 5.

What’s Next?

This project implements the original DDPM paper (Ho et al., 2020). The field has moved fast since then:

  • DDIM (Song et al., 2020) — 10–50× faster sampling with the same model
  • Improved DDPM (Nichol & Dhariwal, 2021) — cosine schedule, learned variances
  • Latent Diffusion (Rombach et al., 2022) — the basis of Stable Diffusion, operates in compressed latent space
  • Classifier-Free Guidance — conditioning on text or class labels for controlled generation

Each of these builds directly on the foundation we built here. Understanding DDPM from scratch is the essential first step.

References

  1. Ho, J., Jain, A., & Abbeel, P. (2020). Denoising Diffusion Probabilistic Models. NeurIPS 2020. arxiv.org/abs/2006.11239
  2. Song, J., Meng, C., & Ermon, S. (2020). Denoising Diffusion Implicit Models. arxiv.org/abs/2010.02502
  3. Nichol, A., & Dhariwal, P. (2021). Improved Denoising Diffusion Probabilistic Models. arxiv.org/abs/2102.09672
  4. Ronneberger, O., Fischer, P., & Brox, T. (2015). U-Net: Convolutional Networks for Biomedical Image Segmentation. arxiv.org/abs/1505.04597

Code & Resources

This project was completed as part of the Generative AI (AI4009) course at the National University of Computer and Emerging Sciences (FAST-NUCES), Spring 2026.

If this post helped you understand diffusion models better, please leave a clap 👏 — it helps others find the article. Questions or suggestions? Drop them in the comments below.


메타데이터
post_id
baecb433cea9
slug
building-a-diffusion-model-from-scratch-generating-faces-with-ddpm-in-pytorch-baecb433cea9
url
https://medium.com/@ghayas.7214/building-a-diffusion-model-from-scratch-generating-faces-with-ddpm-in-pytorch-baecb433cea9
canonical_url
https://medium.com/@ghayas.7214/building-a-diffusion-model-from-scratch-generating-faces-with-ddpm-in-pytorch-baecb433cea9
author_url
https://medium.com/@ghayas.7214
status
ok
fetched_at
2026-06-15 20:49:13