Что такое text-to-image на ESP32-S3 и почему это вообще работает
Генерация изображения по текстовому описанию уместилась в микроконтроллер. Автор проекта Nikuson запустил полный пайплайн text-to-image на ESP32-S3: диффузионный трансформер (DiT) на 2,9 млн параметров, собственный текстовый энкодер на 300 тыс. параметров и декодер на 179 тыс. параметров. Облако не участвует: текст входит, картинка выходит на самом чипе.
Цифры производительности: 5,4 секунды на лицо 64×64 при 8 шагах rectified flow и 2,9 секунды при 4 шагах. Итоговый model.bin весит 3,1 МБ и читается из flash через esp_partition_mmap, пик потребления SRAM около 250 КБ. Промпт уходит по USB, ESP32 отвечает логом шагов и картинкой, которую клиент рисует в терминале полублоками, так что отдельный дисплей не нужен. Автор подробно описал проект в статье на Хабре.
Границы видны сразу. Модель рисует лица 64×64, а не произвольные сцены: в бюджете микроконтроллера «что угодно» не помещается, и лица дают лучший компромисс между разрешением и узнаваемостью. Это демонстрация технологии на пределе железа, а не замена Stable Diffusion.
Ресурсы ESP32-S3: что есть в распоряжении
ESP32-S3 это микроконтроллер с двумя ядрами Xtensa LX7 по 240 МГц и без единого тензорного ядра. Вся арифметика ложится на общие ядра, поэтому вычисления приходится экономить вручную.
| Ресурс | Объём | Что это значит для нейросети |
|---|---|---|
| CPU | 2 × Xtensa LX7, 240 МГц | Тензорных ядер нет, всё считают обычные ядра |
| SRAM | 512 КБ, реально свободно около 320 КБ | Основная рабочая память под активации |
| PSRAM | 2 МБ | Медленная, подключена по quad SPI |
| Flash | 8 МБ | Отображается в адресное пространство через кэш |
| SIMD PIE | 128-битные регистры | 16 умножений int8 за одну инструкцию |
512 КБ SRAM выглядят внушительно ровно до момента, когда рядом появляются веса современных моделей. Текстовый энкодер CLIP из Stable Diffusion весит 250 МБ, то есть в 500 раз больше всей оперативной памяти чипа. Стандартный стек text-to-image на такое железо не портируется: Stable Diffusion требует видеокарту с 4 ГБ памяти и больше, Flux просит 12 ГБ, а лучше 24 ГБ. Отсюда вывод: переносить готовые веса бессмысленно, архитектуру нужно собирать заново под бюджет в сотни килобайт.
Архитектура: латентная диффузия, DiT и собственный текстовый энкодер
Диффузия идёт не по пикселям, а по сжатому латенту размером 16×16×4. Это в 12 раз меньше данных, чем картинка 64×64×3: каждый шаг генерации работает с 1024 числами вместо 12 288. Декодер в пиксели запускается один раз в самом конце, а не на каждом шаге.
Вместо U-Net работает DiT на 64 токенах. Такой трансформер сводится к набору матричных умножений, а значит всю сеть можно описать одним ядром int8 × int8 → int32 и оптимизировать только его. Декодер на 179 тыс. параметров превращает латент в готовое изображение 64×64.
Текстовую часть закрывает собственный крошечный трансформер на 300 тыс. параметров со словарём из 45 «слов». Промпт вроде «улыбающаяся девушка со светлыми волосами» превращается в эмбеддинги этим энкодером, дальше DiT разворачивает их в картинку.
Почему не CLIP и не U-Net
CLIP отпадает по размеру: 250 МБ против 512 КБ SRAM. U-Net отпадает по структуре: пропускные связи и работа на нескольких разрешениях дают разнородный набор слоёв, который трудно свести к одному оптимизированному ядру. DiT на 64 токенах даёт однородную последовательность операций. Компромисс по качеству виден сразу: словарь из 45 токенов вместо сотен тысяч у CLIP, зато модель целиком влезает в чип.
Rectified flow вместо DDPM: 8 шагов вместо 50-1000
Классический DDPM требует 50-1000 шагов на генерацию, и на микроконтроллере это превратилось бы в минуты ожидания. Автор использует rectified flow: приемлемое качество получается уже за 8 шагов, а при 4 шагах лицо 64×64 готово за 2,9 секунды. Сокращение числа шагов дало выигрыш, которого не выжать никакой оптимизацией арифметики.
Как ускорить нейросеть на микроконтроллере: int8, PIE и ассемблерный GEMM
Скорость сложилась из нескольких низкоуровневых решений, которые работают только вместе. Ниже разобраны приёмы, которые дали 5,4 секунды; автор описывает их с кодом в исходном разборе.
int8-квантование с динамическим масштабом
Веса хранятся в int8, активации считаются в int8, накопление идёт в int32. Масштаб пересчитывается динамически по токенам, что помогает удержать точность на модели из миллионов параметров. Квантование экономит память, а попутно открывает доступ к пропускной способности PIE: такие инструкции перемножают 16 пар int8 за одну операцию, и fp32 такой скорости не даёт.
PIE и ручной ассемблерный GEMM
PIE это набор SIMD-инструкций с 128-битными регистрами; инструкция EE.VMULAS.S8.ACCX выполняет умножение с накоплением. Матричные умножения написаны на ассемблере вручную: компилятор не расставляет такие инструкции сам и не выжимает из них максимум. Одно вылизанное ядро GEMM обслуживает всю сеть, от энкодера промпта до последнего блока DiT.
Два ядра и FreeRTOS
ESP32-S3 имеет два ядра Xtensa LX7 по 240 МГц, и FreeRTOS позволяет разнести вычисления между ними. Распараллеливание даёт дополнительный прирост на матричных операциях, где работу легко нарезать на независимые куски.
Поэлементные операции и чтение весов из flash
После матриц остаются поэлементные операции, и они тоже оптимизированы: округление rne (round to nearest even, приведение к ближайшему чётному), быстрые реализации exp и GELU вместо библиотечных. Отдельный приём касается памяти. На каждом шаге генерации сеть обычно прогоняется дважды, для условной и безусловной ветви; автор склеивает обе ветви в один батч, чтобы веса читались из flash один раз вместо двух. Сам model.bin объёмом 3,1 МБ отображается в адресное пространство через esp_partition_mmap, пик потребления SRAM держится около 250 КБ.
Обучение модели: CelebA и автоматические подписи
Датасет CelebA содержит 200 тысяч фотографий и 40 бинарных атрибутов: «улыбается», «очки», «борода» и другие. Текстовых описаний в нём нет, поэтому подписи автор генерирует сам, на лету и разные для каждого прохода по датасету. Признак упоминается с вероятностью p = 0.75, и модель учится работать с неполными описаниями вместо заученных формулировок.
Одна из синтетических подписей выглядит так: «молодая улыбающаяся девушка с чёлкой с волнистыми волосами». Подписи русскоязычные, и именно под них обучен текстовый энкодер со словарём из 45 токенов. Такой словарь покрывает ограниченный набор признаков, зато не требует ни мегабайтов эмбеддингов, ни токенизатора общего назначения.
Отладка: numpy-симулятор, повторяющий арифметику прошивки
Искать ошибку в квантовании прямо на микроконтроллере неудобно: привычных инструментов уровня промежуточных тензоров там нет, а любое расхождение выглядит как случайный шум на итоговой картинке. Автор собрал numpy-симулятор, который повторяет арифметику прошивки: те же масштабы int8, то же округление, тот же порядок операций. Расхождение между симулятором и прошивкой сразу указывает на конкретный слой или ассемблерную вставку, а не на «модель сломалась целиком».
Такой подход отделяет ошибки обучения от ошибок переноса на железо. Симулятор гоняется на большом компьютере, где промежуточные значения можно сравнить с эталоном на PyTorch, и только после сходимости код уходит на ESP32-S3.
Результаты: 5,4 секунды, 64×64 и честные ограничения
Итоговая производительность: лицо 64×64 за 5,4 секунды при 8 шагах rectified flow и за 2,9 секунды при 4 шагах. Промпт передаётся по USB, ESP32 отвечает логом шагов и картинкой, которую клиент рисует в терминале полублоками.
Качество нестабильное, и автор это не скрывает: примерно одна генерация из пяти выходит неудачной. Для демонстрации на микроконтроллере такая честность полезнее отобранных красивых примеров, поскольку сразу показывает реальный уровень технологии. Все цифры и примеры есть в публикации автора.
Ограничения проекта:
- домен только лица, другие сюжеты модель не умеет;
- разрешение 64×64: меньше - лицо уже не узнать, больше - не влезает по памяти;
- словарь из 45 токенов и русскоязычные подписи вместо универсального текстового энкодера;
- обучение на CelebA, а значит наследование смещений датасета;
- работа полностью локальная, без облака и без обращения к внешним сервисам.
Можно ли повторить проект и кому это нужно
Повторение требует трёх разных навыков: обучения маленьких диффузионных моделей, работы с ассемблером Xtensa и понимания встраиваемых систем. Кода здесь много, и он нетривиальный: ручной GEMM, numpy-симулятор, пайплайн генерации подписей. Железо при этом остаётся бюджетным: model.bin занимает 3,1 МБ, пик SRAM около 250 КБ.
Практическая ценность проекта в другом. Это разбор того, как бюджет памяти превращается в архитектурные решения: латент вместо пикселей, DiT вместо U-Net, rectified flow вместо DDPM, int8 и PIE вместо fp32 и GPU. Тот же набор приёмов пригодится при переносе любой модели на edge-устройства, где важны не только миллионы параметров, но и то, как они разложены по памяти.
Заменой Stable Diffusion это не станет, и автор таких заявлений не делает. Код и веса стоит искать у автора проекта (Nikuson): прямых ссылок на репозиторий в исходных материалах нет. Если хотите разобраться в теме на практике, начните с numpy-симулятора: он даёт площадку для проверки квантования и арифметики до того, как вы возьмёте в руки ассемблер и микроконтроллер.