画像生成AIを手元で動かすには、相応のメモリーを積んだPCが要る。待ち時間を嫌えばGPUも要る。RAMが520KBしかない1個170円のマイコンは、まず候補に挙がらない。Tim氏はそこに拡散モデルを載せた。Raspberry Piが設計したRP2350の上で、128×128ピクセルの人の顔を1枚10〜20秒で描く。学習から書き込み用のファイルまで一式をGitHubに置いている。
画素をそのまま作るのではなく、圧縮した表現の側で作る。128×128×3の画像を24分の1の16×16×8まで縮めるVAEを別に学習させ、マイコンには小さいデコーダーだけを載せた。学習のときはエンコーダーを大きく取り、小さなデコーダーでも戻しやすい表現を覚えさせている。本体は潜在フロー拡散トランスフォーマーで、Fluxのような大規模モデルと原理は変わらない。パラメーターは290万と170万の2種類あり、Fluxの約4000分の1にあたる。
小さくする工夫は積み重ねだ。重みは8ビットの整数に落とし、量子化で崩れた品質は量子化後の自己蒸留で取り戻した。活性化関数にReLUの2乗を使って計算を省ける箇所を増やし、推論時間を約15％削っている。時刻とクラスの条件づけは組み合わせが8×5しかないため、学習させた小さなネットワークを捨てて表に置き換えた。重みはフラッシュからDMAで少しずつ読み込み、前の層を計算している間に次の層を用意する。2つのコアは標準の150MHzから300MHzへ引き上げた。最初の実装から約15倍速くなって1枚5秒まで来たが、品質を上げた大きいモデルと誘導の処理を入れて20秒に戻っている。
できた画像はUSBの仮想COMポートから取り出せる。PimoroniのVGAボードを足せばモニターにも映る。色は12ビットしか出せないのでFloyd-Steinbergのディザリングで縞を抑え、生成中は途中経過を低い解像度で映す。画面用に確保したSRAMは、推論中は計算の作業領域に回している。
Tim氏は、コードの下働きの多くをClaude CodeとFable 5に任せたうえで、2週間近くRTX 5090で夜ごと学習を回したと振り返る。リポジトリにライセンスの記載はない。書き込み済みのUF2ファイルが置いてあるので、RP2350のボードが手元にあれば、まず動かしてから中身を追いたい人には入りやすい構成だろう。