A $1 Chip Now Runs a Diffusion Model, 4,000x Smaller Than Flux
TL;DR
Tim (cpldcpu), the engineer behind BitNetMCU, has a real latent flow diffusion transformer generating 128x128 face images on an RP2350, the roughly one-dollar chip in a Raspberry Pi Pico 2. Two variants, 2.9 million and 1.7 million parameters, about 4,000x fewer than the 12 billion in FLUX.1-dev. Images take 4.3 to 20 seconds, fully on-device, inside 520 KB of SRAM, with weights streamed from flash. Pico-Faces ships flashable UF2 images, int8 checkpoints, a C99 inference engine, and the full training pipeline.
What the thing actually does
You hold BOOTSEL, drag a UF2 onto the RPI-RP2 drive, and the board starts generating faces. Output goes over USB as a raw RGB frame, or straight to a VGA monitor if you have a Pimoroni VGA demo base under the Waveshare RP2350-Plus. Conditioning is five classes: four gender-by-smile combinations plus an unconditional class.
You are not going to mistake these for photographs. You are going to mistake them for faces, which is the entire trick. The quality model scores a Gen-FID of 53.8 on-device over 5,000 samples at 8 steps with guidance 4, against a 52.4 floating-point reference. That gap of 1.4 is the total damage from int8 quantization, which is a much better result than the number itself suggests.
The critical detail is not the FID, it is that the model refines. Tim shows the step-by-step emergence grid, and features actually resolve across the 8 diffusion steps instead of snapping to a blurry average on step one and sitting there. Most toy MNIST diffusion demos with more parameters collapse exactly that way.
Why 2.9 million parameters is enough
The first move is the same one Stable Diffusion made in 2022: do not generate pixels. A variational autoencoder compresses the 128x128x3 image into a 16x16x8 latent, a factor of 24, and the diffusion transformer only ever works in that smaller space. Chasing 49,152 numbers is hopeless at this scale. Chasing 2,048 is not.
The second move is more interesting and is the one worth stealing. Tim trained the VAE deliberately lopsided: a large encoder paired with a tiny decoder, because only the decoder has to fit on the microcontroller. The big encoder gets to do the hard thinking about what a good latent representation looks like, under the explicit constraint that a small decoder has to be able to unpack it. It is the difference between a professor writing notes for themselves and a professor writing notes a student can actually revise from. Two decoders were trained, 115K and 494K parameters, and the bigger one is only marginally better (slightly cleaner teeth), which tells you the latent dimension is the binding constraint, not decoder capacity.
The tricks that work on 12B also work on 2.9M
This is the part that should interest anyone doing small-model work. Almost every technique that made large diffusion models good turned out to be necessary here too.
- Flow matching instead of the DDPM formalism. The model predicts a velocity toward the clean latent and the sampler walks it in fractional steps (Lipman et al.).
- Classifier-free guidance (Ho and Salimans), which Tim calls one very significant improvement. It doubles inference cost because you evaluate conditioned and unconditioned, and it is worth every cycle.
- A transformer, not a U-Net. He tried a convnet U-Net first, expecting it to suit a tiny device better. The DiT was both cleaner and better performing. The latent gets patchified into 8x8 = 64 tokens of dimension 128, with sinusoidal position embeddings (learned ones did not work).
- ReLU squared in the dense blocks, borrowed from Primer. It raises activation sparsity, which the C engine exploits for roughly 15% less inference time. A math trick that turns directly into wall-clock on a Cortex-M33.
The AdaLN answer key
Conditioning uses AdaLN-Zero, which normally runs a small MLP to produce per-step scale and bias values. Here there are only 5 classes and 8 timesteps, so Tim threw the MLP away and precomputed every result into a lookup table. It is the times table printed inside the cover of the exercise book: you never do the arithmetic, because there are only so many questions and you already wrote down all the answers.
It works, and it is also the single largest avoidable cost in the build. Those tables occupy 983,040 bytes, nearly a quarter of the 4.02 MB blob, more than the entire VAE decoder and positional embeddings combined. Tim notes he found out about AdaLN-single too late to shrink them.
Getting 4 MB through a 520 KB hole
The RP2350 has 520 KB of SRAM and two Cortex-M33 cores at a stock 150 MHz. The model blob is 4.02 MB. Those numbers do not fit, so the engine never tries to make them fit.
Weights live in flash and are streamed by DMA into an SRAM ping-pong buffer while the previous layer is still computing. The DiT makes this cheap in a way most architectures do not: all 64 tokens go through a layer together, so every weight is reused 64 times once it lands. Streaming throughput is almost never the bottleneck.
The rest is deliberate embedded discipline. Integer-only datapath, no floating point, no malloc, static arenas only, Cortex-M33 SMLAD intrinsics for the int8 matrix multiplies, and core 1 dispatched over disjoint row halves through the inter-core FIFO with no locks. The firmware raises VREG to 1.30 V and overclocks to 300 MHz. The SRAM that holds the VGA framebuffer gets reused as activation and weight scratch during inference, because there is nowhere else to put it.
Iterative optimization cut generation time by about 15x from the first working implementation, down to roughly 5 seconds. Then the bigger model and CFG pushed it back to around 20. The engine is verified byte-exact against a numpy simulator, an x86 golden harness, and the device itself, which is the kind of rigor that separates a demo from a thing you can build on.
What it costs you to reproduce
Two paths, and the repo is honest about both.
Path A, minutes, no GPU. Fold the released quantization-aware checkpoint with its frozen calibration into model.bin, verify the desktop C engine byte-exact against the released goldens, confirm the blob matches the reference, and build the UF2 with the Pico SDK 2.2.0. CPU torch is fine.
Path B, about a day on a CUDA GPU. Download FFHQ, label it by gender and smile, train the VAE, retrain the decoder, extract latents, train the DiT, calibrate, run distillation-QAT, then path A. Training is seeded, but GPU nondeterminism means your checkpoints will not match the released CRCs.
The a-day figure covers one clean run. The actual project took the better part of two weeks of ablations and nightly runs on an RTX 5090, which Tim notes was also heating his office. The energy accounting on generating a 128x128 face for a dollar chip is not flattering, but nobody built this for the energy accounting.
One more detail worth flagging for anyone tracking how this work gets done now: Tim says Claude Code with Fable 5 did a lot of the grunt work of implementing the code, while the two weeks went into ablations and training. The judgment stayed human, the typing did not.
The caveats, stated plainly
- Faces only, five classes. This is FFHQ with gender-by-smile conditioning. There is no text encoder and no prompt. It is not a tiny Flux, it is a proof that the architecture survives the shrink.
- No license file in the repo as of this writing. The code and weights are public, the terms are not. If you plan to build on it, ask first.
- Quantization hurt the bigger model. Tim says the larger variant suffered notable quantization damage, partly healed by self-distillation after the fact, and that training in a more quantization-friendly way is obvious remaining headroom.
- The FID is self-reported and measured on-device by the author. It is reproducible from the repo, which is better than most claims, but nobody independent has run it yet.
Why this matters beyond the novelty
The useful lesson is not that a diffusion model runs on a dollar chip. It is that the constraint stack for a 2.9M-parameter model and a 12B-parameter model turns out to be nearly identical: latent space, flow matching, guidance, a transformer backbone, adaptive normalization. If you are shipping a small model to an edge device, the small-model literature is not a separate discipline. The frontier recipe is the recipe, and the hard part is the plumbing.
The second lesson is architectural. The DiT is unusually friendly to flash streaming precisely because it batches 64 tokens through one layer at a time, and that reuse ratio is what makes a 4 MB model viable in 520 KB. Pick your architecture for the memory hierarchy you actually have, not the one you wish you had.
Key Takeaways
- Pico-Faces runs a latent flow diffusion transformer on an RP2350, generating 128x128 conditioned face images in 4.3 to 20 seconds entirely on-device.
- Two variants at 2.9M and 1.7M parameters, roughly 4,000x smaller than FLUX.1-dev's 12 billion, with int8 weights and an integer-only C99 inference engine.
- A 4.02 MB blob runs inside 520 KB of SRAM by DMA-streaming weights from flash into a ping-pong buffer, exploiting the DiT's 64x weight reuse per layer.
- Nearly every optimization that helps large diffusion models (latent space, flow matching, CFG, AdaLN-Zero, ReLU squared) was necessary at this scale too.
- The asymmetric VAE, a large encoder training a 115K or 494K parameter decoder, is the most portable idea here for anyone targeting constrained hardware.
- Everything is reproducible: UF2 images, int8 checkpoints, byte-exact goldens, and a full retraining pipeline. There is no license file yet, so check before you build on it.
Sources: Pico-Faces repository, AI Image Generation on an RP2350 Microcontroller, Raspberry Pi RP2350 product page, FLUX.1-dev model card, Scalable Diffusion Models with Transformers, Flow Matching for Generative Modeling, Classifier-Free Diffusion Guidance