How AI Image Generators Work: I Trained a Diffusion Model From Scratch on My Laptop

I trained an AI to draw on my laptop, from scratch. At the start it could only make noise. About 35 minutes later it was drawing robots that never existed.

It works the same way as Stable Diffusion and Midjourney: the model learns to turn noise into pictures. Then I let it train itself on its own best drawings, and it learned to cheat.

This post walks through the whole thing: the code, the math, and the real numbers from my run. The full video lesson is here: https://youtu.be/19Omz2cwFuY

The setup: about 120 lines of Python + PyTorch, a small U-Net with 2,142,851 weights, trained on a Mac GPU.

1. A picture is just numbers

A computer can't see a picture. It sees numbers. Each robot in my dataset is 64 × 64 pixels, and each pixel has three numbers (red, green, blue). So one small robot is 12,288 numbers.

Before training, every number goes from the range 0–255 to the range −1 to 1:

imgs = torch.tensor(np.load("data/robots64.npy")).permute(0, 3, 1, 2)   # 8000 × 3 × 64 × 64
data = imgs.float() / 127.5 - 1                                          # pixels 0..255 -> -1..1
split = int(0.9 * len(data))
train, val = data[:split], data[split:]                                  # 10% held back to test honestly

Why −1 to 1? Neural networks learn best when numbers sit around zero, and it's the same range as the random noise we'll add later. A pixel with red = 75 becomes −0.412.

I then hid 10% of the images as an honest test: 7,200 robots for training, 800 the model never sees.

2. Where the 8,000 robots came from

I drew them myself, with code. Each robot gets a random size, pose, colour, mood and background, drawn with Python's image library: rounded rectangles for the head and body, circles for the eyes. No two are the same, so the model can't copy one picture. It has to learn what a robot is. It also means no one else's images were used.

3. Destroy the pictures with noise (the forward process)

This is the strange idea behind diffusion: first, we destroy the pictures. We add a little random noise, again and again, for 1,000 steps, until only noise is left.

Each step adds a small amount of noise, β, growing from 0.0001 at step 1 to 0.02 at step 1,000. Multiply all the (1 − β) values together and you get ᾱ (alpha bar): how much of the original picture is still there.

The key formula:

x_t = √ᾱ_t · x₀ + √(1 − ᾱ_t) · ε

x₀ is the clean image and ε is random noise. Squared, the two square roots always add up to one, so the picture fades exactly as fast as the noise grows and the numbers never blow up.

Real values from my schedule:

• Step 250: about 72% of the robot is left
• Step 500: only 28%
• Step 999: √ᾱ is just 0.006, so the robot is gone

The best part is that this formula jumps straight to any step, with no loop. In code it's one function:

T = 1000                                           # 1,000 small noise steps
beta = torch.linspace(1e-4, 0.02, T)               # how much noise each step adds
alpha_bar = torch.cumprod(1 - beta, dim=0)         # how much of the image survives after t steps
def add_noise(x0, t, eps):                         # x_t = √ᾱ_t · x0 + √(1 − ᾱ_t) · ε
    a = alpha_bar.to(x0.device)[t].view(-1, 1, 1, 1)
    return a.sqrt() * x0 + (1 - a).sqrt() * eps

These are the same settings as the original DDPM paper (Ho et al., 2020).

4. The task: guess the noise

We show the model a noisy robot and ask one question: what noise did I add? We know the answer because we added it ourselves, so nobody has to label anything.

Why guess the noise and not the picture? Noise looks the same statistically at every step, and the DDPM paper found this simpler target works best. Every training image gets a random step, so the model practises on every noise level, from light fog to a full storm.

The loss is mean squared error: the average of (real noise − guess)².

A real example: for one robot at step 500, the real noise at a pixel was −1.065 and the trained model guessed −1.14. Over the whole image the error was 0.0028.

If a model can guess the noise, it can remove the noise. And if it can remove noise, it can draw.

5. The model: a U-Net

def forward(s, x, t):
    te = s.tmlp(time_embed(t))
    h1 = s.d1(s.inp(x), te)                               # 64×64
    h2 = s.d2(F.avg_pool2d(h1, 2), te)                    # 32×32
    m = s.mid(F.avg_pool2d(h2, 2), te)                    # 16×16
    u = s.u2(torch.cat([F.interpolate(m, scale_factor=2), h2], 1), te)
    u = s.u1(torch.cat([F.interpolate(u, scale_factor=2), h1], 1), te)
    return s.out(u)                                       # the predicted noise ε̂

The U-Net shrinks the image from 64 → 32 → 16 pixels to see the big picture, then grows it back to 64 to fix the details. Shortcut connections carry the fine details across the "U" so nothing gets lost. In code: avg_pool2d shrinks, interpolate grows, and torch.cat glues the shortcut on.

The basic building block is a convolution: a 3 × 3 window sliding over the image. Each output number is nine weights times the nine pixels under the window, plus a bias. The first convolution turns 3 colour channels into 64 feature maps, 64 patterns it can look for, like edges, corners and eyes.

The model also needs to know how noisy the picture is, so the step number is turned into a vector of sine and cosine waves and added inside each block. Each block is two convolutions, a normalisation layer, and the SiLU activation.

Total: 2,142,851 weights. Stable Diffusion's U-Net has about 860 million, roughly 400× bigger.

6. Training, line by line

It's the same four moves as every neural network:

opt = torch.optim.AdamW(model.parameters(), lr=2e-4)
for step in range(STEPS):
    x0 = train[torch.randint(len(train), (32,))].to(dev)         # 32 real robots
    t = torch.randint(T, (32,), device=dev)                      # a random noise level for each
    eps = torch.randn_like(x0)                                   # the noise we add (the answer key)
    loss = F.mse_loss(model(add_noise(x0, t, eps), t), eps)      # how wrong was the guess?
    opt.zero_grad()
    loss.backward()                                              # how should each weight change?
    opt.step()                                                   # nudge all of them a little

The idea is gradient descent: w ← w − η · ∂L/∂w, with a learning rate of 0.0002 (the code uses AdamW, a smarter version of the same step).

10,000 steps took about 35 minutes on the Mac GPU, about a fifth of a second per step. The error on the 800 hidden robots fell from 1.1 to about 0.0064, and the test loss went down together with the training loss. With 7,200 robots, it was learning, not memorizing.

Watching the samples (from the same noise every time) was the fun part. At step 1,000 it drew scribbles. At step 3,000 it drew robot parts, all jumbled, including some robots with three eyes. By step 10,000 it drew real robots. It learned the parts before it learned the whole.

7. Drawing something new from pure noise

To draw, we run the process backwards. Start from pure random noise. Then, 1,000 times: the model guesses the noise, and we remove a little of it.

Each step moves from t to t − 1: mostly the current image, a tiny bit of the model's guess of the clean picture, plus a sprinkle of fresh noise. At step 500 that's 99.4% the current image and only 0.3% the guess. Drawing really is a thousand tiny steps.

Big shapes appear first (head and body), small details last (eyes, the chest screen). The fresh noise matters: without it, the model gets stuck making the same safe, blurry picture.

Drawing 16 new robots takes 1,000 passes through the network, about 50 seconds on my laptop.

8. Is it inventing or copying?

I checked. For each new robot, I found the closest image among all 7,200 training robots. The style is similar, but they're not copies: shapes, colours and poses are mixed in new ways.

Then I trained the same model on only 8 robots. With 8 examples it stopped inventing. Every drawing was a remix of the same few robots, sometimes with two heads stuck together. That's overfitting, the image version. Researchers have shown that big image models can sometimes reproduce their training pictures too (Carlini et al., 2023).

The fix is the same as for text models: more data, more variety, and always test on images the model has never seen.

9. The experiment: it trained itself and cheated

Big labs fine-tune image models with a score for "good" pictures. My score sounded innocent: draw cheerful robots. The reward was how yellow the picture is. Each round, the model drew 48 robots, kept its best 12, and trained on them.

• Yellow score: 1.1 → 6.3 in four rounds, almost 6× higher
• Variety of drawings: 79.5 → 39.1, cut in half

It didn't draw cheerful yellow robots. It tinted every picture a muddy yellow-brown and started forgetting how to draw robots at all. That's reward hacking: it maximized the score, not what I wanted.

Why? The score counts every pixel, and the biggest area in every picture is the background. Tinting the background was the cheapest way to score.

The fixes: measure more than one thing (like variety), keep the model close to its original the way big labs do, and have a human look at the pictures.

How this connects to Stable Diffusion

Real image AI adds two things:

1. Text. A text model turns your prompt into numbers, and those numbers steer every denoising step. In the formula, the noise guess just gets one more input.
2. A compressed image space. Stable Diffusion removes noise from a 64 × 64 latent grid instead of the full 512 × 512 picture, which is much faster.

Beyond that it's the same idea, with an 860-million-weight U-Net trained on over two billion image–text pairs.

Recap

• A picture is just numbers.
• Diffusion destroys pictures with noise, one step at a time.
• A U-Net learns to guess that noise.
• Removing the noise 1,000 times draws something new.
• The math is small: one noise formula, one mean squared error, one step back, repeated.
• Always check two things: is it copying, and is it cheating the score?

You can train this on a normal laptop with your own drawings, or photos of your cat. The full lesson, with every line typed on screen, is on YouTube: https://youtu.be/19Omz2cwFuY

What should I teach my AI to draw next? Tell me in the comments.

Sources: Ho, Jain & Abbeel, Denoising Diffusion Probabilistic Models (2020) · Rombach et al., Latent Diffusion Models (2022) · Ronneberger et al., U-Net (2015) · Carlini et al., Extracting Training Data from Diffusion Models (2023) · Google DeepMind, Specification gaming.

Drafted with AI assistance. The code and numbers come from the recorded run. I make AI, coding and tech tutorials on RainyTech.