Vedhathiri
Author
Generative AI image models usually require powerful processors and large amounts of memory, but a new project demonstrates that image generation can be pushed onto a small microcontroller. Tim, from Tim’s Blog, has implemented a generative AI image model on the RP2350, the microcontroller used in the Raspberry Pi Pico 2. Running on a Waveshare RP2350-Plus board, the system can generate 128×128 RGB images of human faces in around 10-20 seconds and display them on a VGA monitor or transfer them through USB. Most importantly, the image generation is performed directly on the microcontroller, without a GPU or cloud-based AI service doing the inference. The model is a latent flow diffusion transformer (DiT) with two variants containing 1.7 million and 2.9 million parameters. It was trained using the FFHQ face dataset and supports five generation classes based on gender, smiling, and an unconditional class. Unlike general-purpose text-to-image systems, this implementation is specifically designed to generate small human-face images.
The key to making generative AI work on the RP2350 is reducing both the model's computational and memory requirements. Instead of generating the image directly as individual pixels, the system uses a Variational Auto-Encoder (VAE) to compress a 128×128×3 image into a much smaller 16×16×8 latent representation. The diffusion model then starts from noise and uses flow matching to predict the direction toward the clean image, taking eight steps during inference so that the image gradually emerges from the noise. A Diffusion Transformer processes the latent data as 64 tokens, while Classifier-Free Guidance (CFG) is used to strengthen the selected generation condition. To fit the model into the microcontroller, its weights are stored using INT8 quantization, with the larger model and inference engine fitting within about 4 MB of flash. The inference engine also streams weights from flash using DMA, uses both Cortex-M33 cores, and runs them at 300 MHz, compared with the RP2350's standard 150 MHz clock. These optimizations allow the Pico 2 to perform the required AI calculations despite having only 520 KB of RAM, although generating each image still takes several seconds.
The image-generation process is controlled through the RP2350's USB virtual COM port, using a custom protocol to send commands to the device. Once generation starts, the Pico creates the initial noise, processes it through the diffusion model, reconstructs the resulting latent representation using the VAE, and produces the final 128×128 RGB image. A Pimoroni VGA board is used to output the result to a VGA monitor, while the image can also be transferred over USB. During generation, a low-resolution preview of the latent representation can be displayed, allowing the image to be seen developing as the model processes it. The project does have significant limitations: the output is restricted to 128×128 resolution, the model is trained specifically for human-face generation rather than arbitrary images, it does not accept normal text prompts like larger image-generation models, and generation takes roughly 10–20 seconds per image.
