Dev.to AI πŸ€– Ai πŸ‘ 0 πŸ“– 9 min read

I Built a Latent Diffusion Model That Fits in 5.5MB and Runs Without a GPU

Stable Diffusion needs an A100 to feel fast. It ships with 860 million parameters, a UNet backbone, and a multi-gigabyte footprint that assumes you have real GPU infrastructure behind it. I wanted to know how small a di

Stable Diffusion needs an A100 to feel fast. It ships with 860 million parameters, a UNet backbone, and a multi-gigabyte footprint that assumes you have real GPU infrastructure behind it.

I wanted to know how small a diffusion model could get before it stopped working at all. The answer, after a lot of ablation studies I'll show you below, is TinyDiffusion 128: a 1.45 million parameter Diffusion Transformer that fits in 5.55MB of FP32 weights, runs in under 200ms on a plain CPU in the project's fastest benchmark configuration (more on how that number relates to the others below), and needs less than 25MB of RAM to do it.

This isn't a toy that only draws noise. It generates coherent 128x128 images from text prompts, and every TinyDiffusion number in this post is pulled directly from the project's own evaluation suite and ablation reports, not estimated or rounded up for effect. The figures for other models are reference values from published or documented sources, and they're labelled that way where they appear.

What is TinyDiffusion 128?

A quality-first, CPU-native, micro-scale latent diffusion system. It replaces the heavy UNet architecture most diffusion models use with a Diffusion Transformer (DiT) operating on an 8x compressed latent space, which is most of where the size reduction comes from.

The core numbers, as documented in the repository:

Attribute Value
Architecture Diffusion Transformer (DiT)
Parameters 1.45 million
FP32 weight size 5.55 MB
Peak RAM < 25 MB
CPU inference sub-200 ms for a full 12-step generation (fastest benchmark configuration)

Nothing here needs a GPU. It runs on Windows, macOS, and Linux with pure CPU inference, and there's a zero-dependency C++17 runtime for embedded or edge deployment on top of the Python package.

How It Actually Works

The pipeline has three stages: turn the text prompt into a conditioning signal, run that signal through a DiT backbone across a handful of denoising steps, then decode the result into a real image.

1. Conditioning. A prompt like "a ruby red circle on a white background" goes through a TinyTokenizer and TinyTextEncoder, producing a single context embedding (a 1x64 vector) that steers everything that follows.

2. Latent diffusion sampling. Gaussian noise starts in a compressed 1x4x16x16 latent space. The TinyDiT backbone predicts the noise to remove at each step, guided by both the timestep and the text embedding, and a DDIM scheduler runs this for 10 to 15 steps rather than the hundreds a naive diffusion process would need. Classifier-free guidance (CFG) blends a conditioned and unconditioned prediction at a scale of 3.5 by default, which the project's own ablation study found to be the best trade-off between boundary sharpness and saturation, without introducing clipping artifacts.

3. Reconstruction. Once the latent is denoised, a TinyVAE decoder upscales it 8x back into a full 128x128 RGB image.

The model can also be exported to a single-file binary format (.tdiff) and run through the standalone C++ runtime, which uses raw pointer arithmetic and arena allocators with no external dependencies.

How It Was Evaluated

A model this small is easy to dismiss as not really working. So before anything else, it's worth looking at how the project measured whether it actually does, rather than just trusting the claim.

Metric Score Target What it checks
FrΓ©chet Distance (FID) 301.64 < 350 Similarity to 792 reference samples
Inception Score (IS) 2.17 Β± 0.06 > 2.0 Class salience, batch variety
Mean brightness 166.9 / 255 > 140 No dark-image collapse
Spatial variance 36.2 > 25 Edge contrast, not flat output
Shannon entropy 6.65 bits > 4.0 Information density, pixel balance
Dark image bias (< 80 RGB) 0.1% (1/792) < 5.0% Color variety across the gallery

These are the project's own self-reported targets and results, evaluated against 792 procedurally generated reference samples, not an independent third-party benchmark. Worth keeping in mind as you read the rest of this: a small FID-style check against your own reference set tells you the model is internally consistent and not collapsing to noise, but it isn't the same claim as matching a large, diverse, human-curated dataset.

A note on the FID number itself: 301.64 will look high next to scores in published papers, and it shouldn't be compared with them. FID depends on the feature extractor and on the reference set, and it is biased upward when the sample count is small (here, 792 images). A score computed on 128x128 procedural images against a small custom reference set is not on the same scale as one computed on a large benchmark such as ImageNet or COCO. That's why the target is a loose threshold (< 350): it works as a sanity check that the model hasn't collapsed, not as a quality ranking.

The CFG Ablation: Why 3.5 Is the Default

Classifier-free guidance scale controls the tradeoff between literal prompt-following and how saturated and sharp the output looks. The project ran a sweep across five values:

CFG scale Spatial variance Notes
1.0 41.03 Low contrast, under-conditioned
2.0 42.98 Balanced, mild boundary softness
3.5 48.05 Strong boundary sharpness, no clipping
7.0 48.70 Highest variance, but minor saturation artifacts
10.0 43.10 Over-guided, edge clipping appears

The curve isn't monotonic. Pushing guidance higher doesn't keep helping past a point, it starts clipping and losing fidelity. Spatial variance is technically highest at 7.0 (48.70), but 3.5 reaches within about 1.3% of that (48.05) without the saturation artifacts that appear at 7.0 and the edge clipping at 10.0. That trade-off, rather than the single highest number, is why 3.5 is the default instead of the maximum the sweep tested.

TinyDiT vs TinyUNet: Why Transformer Won

The project also benchmarked its own DiT backbone against a ConvNet-based UNet alternative, on the same hardware:

Metric TinyDiT TinyUNet Advantage
Parameters 1.45M 1.54M DiT, 5.8% smaller
FP32 size 5.55 MB 5.89 MB DiT
Single-step latency 31.81 ms 57.99 ms DiT, 1.82x faster
Sampling throughput 31.4 it/s 17.2 it/s DiT, +82.5%
12-step CFG total 763.5 ms 1391.8 ms DiT

Despite being the newer, less battle-tested architecture choice for a project this size, the transformer backbone won on every axis measured here: fewer parameters, smaller file, and roughly 1.8x faster per step.

How It Compares to Larger Diffusion Models

This is the table most people will want to see, so it deserves the most caution. These are the project's documented comparison figures. The TinyDiffusion row is this project's own measurement; the other rows are reference values:

Model Parameters FP32 size CPU latency/step Total (steps) RAM
TinyDiffusion 128 (ours) 1.45 M 5.55 MB 13.13 ms ~178 ms (12 steps) < 25 MB
SnapFusion 280 M 1,120 MB ~350 ms ~1,800 ms (8 steps) ~1.5 GB
MobileDiffusion (Google) 520 M 2,080 MB ~280 ms ~1,200 ms (6 steps) ~2.2 GB
Mini-SD 390 M 1,560 MB ~680 ms ~5,400 ms (15 steps) ~2.8 GB
Stable Diffusion v1.5 860 M 3,440 MB ~2,800 ms ~28,000 ms (20 steps) ~4.5 GB

A few things worth being precise about here, since it's easy for a table like this to imply more than it shows:

  • The rows for SnapFusion, MobileDiffusion, Mini-SD, and Stable Diffusion v1.5 are reference figures, not benchmarks re-run on the same machine as TinyDiffusion. Only the TinyDiffusion row reflects this project's own measurement.
  • The step counts differ across models (12 steps for TinyDiffusion vs 20 for SD v1.5), so the total-time comparison isn't a perfectly controlled apples-to-apples test. The per-step latency column (13.13ms vs ~2,800ms) is the fairer number to anchor on, and it still puts TinyDiffusion roughly two orders of magnitude faster per step on CPU.

None of this changes the underlying result, which is real: a 1.45M parameter model fits in CPU cache in a way an 860M parameter model simply cannot. But "213x faster per step" is a more honest framing than "150x faster overall," since the latter quietly compares different step budgets.

A note on the two sets of timing numbers

The latency figures in this post come from two different benchmark setups, and they don't match each other:

Source Per-step latency 12-step total
Comparison table above 13.13 ms ~178 ms
TinyDiT vs TinyUNet benchmark 31.81 ms 763.5 ms (with CFG)

The second row is consistent with classifier-free guidance doubling the work: 763.5 ms over 12 steps is about 63.6 ms per step, roughly twice the 31.81 ms single-step figure, because each step runs a conditioned and an unconditioned prediction. The first row is the faster configuration, and it is where the "under 200 ms" headline comes from. I haven't reconciled the two on identical hardware and settings, so treat the headline as the project's best measured case rather than a guarantee for every configuration.

The per-step comparison with Stable Diffusion holds up either way: using the slower 31.81 ms figure instead of 13.13 ms still gives roughly 88x faster per step (about 2,800 ms vs 31.81 ms), instead of 213x.

Try It

git clone https://github.com/bharqav/tiny-diffusion.git
cd tiny-diffusion
python -m venv venv
source venv/bin/activate    # or venv\Scripts\activate on Windows
pip install -e ".[dev]"

The fastest way to see it work is the interactive prompt studio:

tiny-diffusion interactive
Enter Prompt: a glowing gold star on a deep space background
TinyDiffusion Sampling: 100%|##########| 10/10 [00:00<00:00, 26.2it/s]
Success! Image saved to: output/result_1.png

Or generate directly from the CLI with explicit sampling parameters:

tiny-diffusion generate \
  --prompt "a ruby red circle on a white background" \
  --steps 12 \
  --guidance 3.5 \
  --seed 42 \
  --output output/ruby_circle.png

Or from Python:

from tiny_diffusion import TinyDiffusionPipeline

pipe = TinyDiffusionPipeline.from_pretrained(checkpoint_dir="checkpoints", model_type="dit")

image = pipe(
    prompt="a vibrant magenta hexagon on a golden glow background",
    num_inference_steps=10,
    guidance_scale=3.5,
    seed=513,
)
image.save("output/magenta_hexagon.png")

For the standalone C++ runtime:

cmake -B build
cmake --build build --config Release
./build/test_runtime
./build/tiny-diffusion-cpp

Repository Structure

tiny-diffusion/
β”œβ”€β”€ tiny_diffusion/           # Canonical Python package
β”‚   β”œβ”€β”€ cli.py                # Unified CLI entrypoint
β”‚   β”œβ”€β”€ models/               # DiT, UNet, VAE, text encoder
β”‚   β”œβ”€β”€ sampling/             # DDIM scheduler, sampling pipeline
β”‚   β”œβ”€β”€ data/                 # Tokenizer and dataset abstractions
β”‚   β”œβ”€β”€ data_factory/         # Procedural synthetic data generation
β”‚   β”œβ”€β”€ training/             # VAE, DiT, and distributed training loops
β”‚   β”œβ”€β”€ evaluation/           # FID, IS, Shannon entropy metrics
β”‚   └── compression/          # INT8 quantization, .tdiff export
β”œβ”€β”€ runtime/                  # Pure C++17 zero-dependency inference engine
β”‚   β”œβ”€β”€ include/              # arena.h, tensor.h, model.h, kernels.h
β”‚   β”œβ”€β”€ src/                  # C++ implementation
β”‚   └── tests/                # C++ unit tests
β”œβ”€β”€ output/                   # 1,000-image gallery, manifests, benchmark reports
β”œβ”€β”€ benchmarks/               # Ablation and comparative benchmark runners
β”œβ”€β”€ scripts/                  # Gallery generator, ops profiler, training pipeline
└── tests/                    # 27 passing pytest unit tests

Lessons Learned

Parameter count isn't the only thing that matters, shape does too. TinyDiT beat TinyUNet with fewer parameters and a smaller file, but the real win was architecture: the transformer's structure mapped better to CPU cache behavior than the convolutional alternative did, which is why the latency gap (1.82x) is bigger than the parameter gap (5.8%).

Guidance scale has a peak, not a ceiling. It would be easy to assume more CFG guidance is strictly better. The ablation shows clearly that it isn't: quality rises, peaks around 3.5 to 7.0, then degrades into clipping artifacts past that.

Self-evaluation needs an honest frame. An FID score against your own 792-sample reference set is a real, useful signal that the model isn't collapsing or stuck on noise. But it's a different claim from matching a benchmark built on a large external dataset, and it's worth saying so directly rather than letting the number imply more than it supports.

Compare the thing that's actually controlled. The per-step latency comparison against larger models is a clean, apples-to-apples number. The total-time comparison isn't, because the step counts differ. Both are useful, but only one of them is a fair test, and it matters which one you lead with.

The Real Takeaway

A 5.55MB model that runs in under 200ms on a CPU in its fastest configuration isn't trying to replace Stable Diffusion. It's answering a narrower question: how much can you strip away from a diffusion model's architecture, parameter count, and latent space before image generation stops working, and how do you prove it's still working once you have?

The ablations and evaluation suite are the actual contribution here, as much as the model itself. Anyone can claim a small model works. Showing the CFG sweep, the architecture comparison, and the FID numbers against a stated reference set is what makes that claim checkable.

The code, the 1,000-image gallery, and the full benchmark reports are on GitHub: github.com/bharqav/tiny-diffusion

πŸ“° Read the original article on Dev.to AI

Originally published by Dev.to AI. Aggregated on AIWithGhost for educational purposes β€” full credit and traffic to the original publisher.