Building a Diffusion Model From Scratch

A lot of "AI image generator" projects are a wrapper around someone else's model. I wanted to write the model myself, starting from an empty file, so I would understand every line of it. So I built thisnaturedoesnotexist - a denoising diffusion model, hand-written in PyTorch, with no pretrained weights and no hosted API anywhere in the stack. It learns to turn pure static into landscapes.

This is the deep version of that story: what a diffusion model actually is, how the network is put together, and - the part I found most valuable - every gotcha that cost me time or money. If you are learning this yourself, the mistakes are the useful part.

Teaching a network to un-destroy an image

The idea behind a diffusion model is odd. You take a real image and add a little random noise, then add a little more, and keep going for a thousand steps until nothing is left but television static. That direction - image to noise - is called the forward process, and it is dead simple: it is just addition of Gaussian noise on a fixed schedule. No learning required.

The useful part is running it backwards. If a network could look at a slightly noisy image and predict the noise that was added, you could subtract that estimate and get a cleaner image. Do that a thousand times, starting from nothing but static, and you walk backwards from noise into a brand-new image that never existed. That reverse walk is the only thing the network learns.

I built the mental model on the cheap first: a "toy" diffusion model on 2D points, where the whole dataset is a scatter of dots in a spiral and the network is a tiny multi-layer perceptron. No images and no U-Net, so there was nowhere for a bug to hide. When a cloud of random points learned to reassemble itself into the spiral, I knew the loop was right and could scale it up to pixels with confidence.

Why the noise predictor has to be a U-Net

The toy model used a plain MLP because its input was two numbers. An image is a grid, and an MLP would throw away the single most important fact about it: that pixel (5,5) sits next to pixel (5,6). So the real noise predictor is a U-Net, and I wrote it by hand - roughly 157 lines - rather than importing one, because the whole point was to understand it.

A U-Net has three parts:

• Down path - repeatedly halve the image (64→32→16...) while growing the channel count. Deeper layers see a bigger slice of the picture at once, so they see more context but less detail.
• Bottleneck - the smallest, most abstract representation. Here I added a single-head self-attention block so distant regions can talk to each other. Convolutions only see nearby pixels. Attention lets the sky on the left match the sky on the right.
• Up path - mirror the down path back to full resolution. At every step it concatenates a skip connection from the matching down-path layer, so the fine detail that got squeezed out on the way down is handed back on the way up.

A few deliberate choices inside the blocks, each of which matters more than it looks:

• Residual blocks with GroupNorm + SiLU. GroupNorm (not BatchNorm) because diffusion trains on small, noisy batches where BatchNorm's per-batch statistics get erratic. SiLU because it is the smooth activation these models empirically like.
• Sinusoidal timestep embeddings fed into every block through a small linear projection - the same trick transformers use for positions, here encoding "how far along the noise schedule are we."
• Attention only at the bottleneck. Attention is quadratic in the number of pixels, so it is ruinous at full resolution and cheap on the tiny bottleneck grid - which is exactly where global consistency matters most.

The training objective is very simple

This part surprised me. After all that architecture, the training step is four lines of intent: pick a random image, pick a random timestep, add that much noise, and ask the network to guess the noise back. The loss is plain mean-squared error between the real noise and the predicted noise.

# the entire training objective
t     = randint(0, 1000)              # random noise level
noise = randn_like(x0)                # the target
x_t   = schedule.add_noise(x0, t, noise)
pred  = model(x_t, t)                 # the U-Net's guess
loss  = mse(pred, noise)             # that's it

That is all there is to it. There is no adversarial discriminator to babysit, no mode collapse to fight - two of the things that made GANs miserable to train. Diffusion trades that instability for a longer, more predictable climb. The details that make it actually work are around the edges of that loop:

• A cosine noise schedule instead of linear. The linear schedule destroys an image too aggressively early on; cosine is gentler and spends more of its budget in the range where the network can actually learn something.
• EMA weights. The images I judge quality on are never generated from the live training weights. They come from an exponential moving average (decay 0.999) of those weights - a slow, smoothed version that skips past the per-step jitter. This one line noticeably cleans up samples.
• Two samplers. The faithful one (DDPM) walks all 1000 steps. The fast one (DDIM) takes a shortcut and gets a comparable image in about 30. During training I sample with DDIM so I can eyeball progress without waiting. 30 steps gives you a preview. 1000 steps gives you the finished render.

The loss curve lies to you

This one cost me a genuine "wait, is it broken?" afternoon. Diffusion loss drops fast and then crawls - in my runs it fell from about 0.043 to 0.026 in a handful of epochs and then barely moved for dozens more. If you were watching the loss like you would on a classifier, you would conclude the model finished learning early and kill the run.

That would be a mistake. The samples kept getting better long after the loss went flat - soft coloured blobs slowly resolving into things with horizons and structure. MSE-on-noise is dominated by the easy, high-noise timesteps where any reasonable guess is close. The hard, low-noise timesteps - the ones that decide whether you get a real landscape or mush - contribute a tiny slice of the average. So the number stops moving while the thing you care about is still improving.

Data did more of the work than resolution did

I trained in two environments. On my laptop (an M1 Pro, using Apple's MPS backend) I trained small U-Nets at 32 and 64 pixels on about 6,300 real landscape photos - Flickr shots I pulled out of the huggan/monet2photo dataset. Then I rented a real GPU and trained a bigger U-Net at 128 pixels on the LHQ dataset - roughly 90,000 landscapes.

The jump in quality from that second run was huge, and it is tempting to credit the higher resolution. But most of the gain came from the data: 15× more images, and cleaner ones. Resolution sets how much detail the model can show. The dataset decides whether there is anything worth showing. If I could only fix one of the two, I would pick the data.

Scaling also surfaced a laptop-versus-cloud gotcha. Mixed precision (AMP) - training in half precision for roughly 2× speed and half the memory - simply does not run on the M1's MPS backend. I wrote the code for it, but the very first time I could confirm it actually worked was watching amp: True print on the rented NVIDIA pod. If you develop on Apple Silicon and deploy to CUDA, assume a class of features is untestable until you are on the real hardware, and leave a clean switch for them.

The cheapest card per hour is a trap

For the 128px run I rented a single RTX 4090 on RunPod's community cloud at $0.69/hour. Before committing I worked out what the alternatives would cost. The answer was not what I expected:

My model is small - 48M parameters, only a few GB of VRAM in use. It physically cannot fill a big data-center GPU, so paying for one means paying for silicon that sits idle. Measured baseline: the 4090 did 433 seconds per epoch. Everything else is extrapolated from that one real datapoint, sorted by estimated total for a 100-epoch run.

The A40, 3090 and L4 are all cheaper per hour than the 4090 and all cost more in total because they finish so much later. The giant cards are money on fire for a model this size. The part I did not figure out until afterward:

Two bugs that only bite once you are paying by the hour

The resume bug. When I handed weights from the 32px run to the 64px run, my --epochs flag was being read as an absolute target epoch, not "run this many more." So asking for "60 epochs" after a checkpoint at epoch 40 would have silently trained only 20. I caught it before it wasted a run, and the fix encodes a real distinction: same resolution as the checkpoint means "continue an interrupted run" (restore the optimizer and epoch counter too); a different resolution means "new training phase" (carry the weights, reset everything else). If the checkpoint and the CLI disagree about what "epochs" means, you pay for it in GPU time.

The storage blip. A cloud run died at epoch 32 - not a code or data bug. A DataLoader worker hit OSError: No such device or address reading a single image off the network-attached volume. Every one of the 90,000 files had already been read cleanly 31 times over; it was a transient I/O hiccup on shared cloud storage. One flaky read should never kill a multi-hour job, so the dataset now retries a failed read a few times and falls back to a neighbouring image if a file is genuinely unreadable:

# a long run must survive a momentary I/O hiccup
for attempt in range(5):
    try:
        return transform(Image.open(paths[idx]))
    except OSError:
        sleep(0.2 * (attempt + 1))       # back off
        if attempt == 3:
            idx = (idx + 1) % len(paths)  # fall back to a neighbour

Because I was already checkpointing every epoch, resuming from the epoch-31 checkpoint cost about seven minutes instead of the whole run. Both bugs come down to the same thing: long jobs need per-epoch checkpoints and resilient data loading. You find out the hard way if you skip either.

The images were broken but the model was fine

Once the 128px run finished, I ran the first-ever local inference pass on the checkpoint - not more training, just sampling from it - to see if more DDIM steps would clean up some speckle and flat-colour tiles I was seeing in the epoch-99 grid. Fifty steps instead of thirty did not fix it. So I reached for what should have been the reference: the full, no-shortcuts 1000-step ancestral sampler. It came back worse - grids full of solid flat tiles: pure red, pure white, pure black, pure cyan.

The root cause was a gap I had never noticed because I had never run that code path against a real trained model. The fast sampler (DDIM) self-corrects every single step: it reconstructs the model's current best guess at the final image and clamps it back into a valid range before continuing. The ancestral sampler never had that guard. During training I only ever used DDIM to render the periodic progress grids, so the ancestral path had sat untested since the day I wrote it. Small per-step prediction errors, with nothing correcting them, compounded over a thousand steps and diverged.

# before: nothing keeps x bounded across 1000 uncorrected steps
x = (x - (beta / sqrt(1 - alpha_bar)) * pred_noise) / sqrt(alpha)

# after: reconstruct x0, clamp it, then take the posterior mean from that
x0_pred = ((x - sqrt(1 - alpha_bar) * pred_noise) / sqrt(alpha_bar)).clamp(-1, 1)
mean    = coef_x0 * x0_pred + coef_xt * x   # same clamp-then-correct DDIM already had

I fixed the sampler and re-ran the exact same seed. What came back was a real coastal cliff with water reflections, a glacier lake, hillside sunsets, a golden-hour ocean horizon, a forest edge - no speckle, no flat tiles.

That reopens a different, more interesting problem than the one I thought I had. The fixed ancestral sampler looks great and takes about 30 seconds per image on an M1 - much too slow for a page where a visitor clicks "generate" and expects something in a couple of seconds. The fast sampler is quick but still shows some artifacts, likely because skipping across large timestep jumps is less forgiving of a small model's imperfections than many tiny, noise-smoothed ancestral steps are. The model is good enough. What I have not solved is how to sample it quickly without losing that quality.

The honest part: shipping a hand-rolled model is its own problem

The concept borrows from thispersondoesnotexist.com - a fresh generated image on every visit. That site does not run its model in your browser. It serves pre-generated images. My model is 48M parameters, which lands around 95-190MB as ONNX. That is the real design fork for shipping:

• Pre-generated gallery - render a few thousand images offline, drop them in object storage, serve one at random. Instant, free, but finite and not truly "live."
• Live in-browser via ONNX Runtime Web on WebGPU - genuinely generating on the visitor's device, zero server GPU cost, but a heavy model download and a device-dependent wait.
• Hybrid - gallery by default, live generation as an opt-in for capable machines.

I was ready to write here that a from-scratch 48M-parameter DDPM simply produces dreamy, painterly landscapes rather than photoreal ones - and then the sampler fix above proved that framing wrong. Sampled correctly, the model produces genuinely convincing landscapes. The real constraint is somewhere else: the sampler fast enough to run live on a visitor's device is not yet the same sampler that produces that quality, and closing that gap - not the model's raw capability - is the unresolved problem as I publish this. I never expected photoreal output on day one. The point was to write every layer myself and to be honest about what still does not work.

What I would tell myself on day one

I would start with the toy model again. I would watch the sample grid instead of the loss, and put the effort into more data before raising the resolution. Checkpoint every epoch. Make the data loader handle a bad read. And write the U-Net by hand at least once. Debugging my own at 2am taught me more than a pretrained checkpoint would have.

If you want the companion piece on the audio side - where I trained a real classifier, measured it honestly, and decided not to ship it yet - that story is how I taught a model to read dog barks. And the code, checkpoints, and sample grids for this one live back on my portfolio.