This project implements a Gen AI image generator running on a $1 Microcontroller, the RP2350, as used in the Raspberry Pi Pico 2.
It can generate 128×128 RGB images of human faces in 5-20s each and display them on a VGA monitor or stream them over USB. The model implements a latent flow diffusion transformer (DiT), similar to what is used in models like Flux. There are two variants at 2.9 and 1.7 million parameters, 5000x times less than even a typical local diffusion model. It supports conditional generation in 5 classes (gender × smile + neutral).
It is more than astonishing that a model this small is able to generate complex images at all. Many MNIST toy diffusion projects use far more parameters and are barely able to generate anything coherent. Interestingly, a lot of the optimizations that helped large models were also necessary for this micro model.
Article with more details here
pico_faces_monitor.mp4 #
This indicates that the model does indeed behave like a diffusion model. Smaller diffusion models often tend to collapse to an initial bias without actual refinement in subsequent steps. This is not the case here.
This demonstrates that the model scales as expected with number of steps k and classifier free guidance w (CFG).
Hold BOOTSEL while resetting the RP2350 board, copy
uf2/pico_faces_m3_decD_deep_full.uf2onto theRPI-RP2
drive. -
pip install pyserial matplotlib numpy pillow
, then:
python viewer/view_serial.py --port com10 --seed 3 --steps 8 --class 4 --cfg 6 --show
Check the viewer/README.md for the full parameters descriptions.
Optional: on a Pimoroni VGA Demo Base the firmware displays the images on a connected VGA monitor.
High quality model m3_decD_deep_full |
Fast model m3_long_cfg |
|
|---|---|---|
| DiT | dim 128 × depth 12, 2.37M params | dim 128 × depth 8, 1.59M params |
| VAE decoder | ~493K params | small, ~116K params |
| Blob size | 4.02MB | 2.57MB |
| Gen-FID (device, N=5000, K=8 w=4) | 53.8 (fp reference: 52.4) |
— (speed build) |
| Time / image | ~10 s @ K=4 w=4 (≈20 s @ K=8 w=8) | ~4.3 s @ K=4 w=4 (5.4 s @ K=8 plain) |
Path A — from the released checkpoints (no GPU, ~minutes).
pip install -r requirements.txt # plus torch (CPU is fine for folding)
bash scripts/finalize.sh m3_decD_deep_full
bash scripts/finalize.sh m3_long_cfg
This folds the released QAT checkpoint with its frozen calibration into
model.bin
, verifies the desktop C engine byte-exact against the released
goldens, confirms the blob is byte-identical to checkpoints/<model>/model.bin
, and builds the UF2 (needs the Pico SDK).
Path B — full retrain (CUDA GPU, ~a day). Dataset download → VAE → decoder D → latents → DiT → calibration → distillation-QAT → path A. The stage-by-stage commands live in train/README.md and quant/README.md; training is seeded but GPU nondeterminism means your checkpoints (and CRCs) will differ.
| dir | contents |
|---|---|
vae.yaml
, dit.yaml
, export.yaml
)checkpoints/data/train/quant/engine/firmware/scripts/finalize.sh
, verify_model.sh
, build_firmware.sh
, train_model.sh