Текст в изображение на микроконтроллере: как ESP32-S3 запустил диффузионный трансформер
Диффузионные модели Stable Diffusion и Flux требуют мощных видеокарт и гигабайт памяти, что делает их недоступными для простых микроконтроллеров. Однако инженерам удалось адаптировать архитектуру DiT для ESP32-S3, создав систему, генерирующую изображения 64x64 пикселей непосредственно на чипе без доступа в интернет. Процесс занял 5,4 секунды и потребовал сложной оптимизации, включая использование ассемблера и квантование весов.
# Текст в изображение на микроконтроллере: как ESP32-S3 запустил диффузионный трансформер
Разработчики столкнулись с очевидной проблемой: для запуска современных генеративных моделей искусственного интеллекта требуются специализированные видеокарты с огромным объемом видеопамяти. Популярные модели, такие как Stable Diffusion, нуждаются в 4–8 ГБ, а более новые архитектуры типа Flux — в 12 или даже 24 ГБ. В то же время, отладочные платы на базе микроконтроллера ESP32-S3 стоят всего несколько сотен рублей и имеют лишь 512 КБ оперативной памяти, два процессорных ядра и отсутствие тензорных ускорителей.
Вопрос звучит провокационно: возможно ли прогнать диффузионную модель прямо на таком устройстве? Не отправляя запрос в облако, а выполняя вычисления локально? Ответ оказался положительным, хотя путь к нему был усыпан техническими сложностями.
Что уместить в 512 килобайт
Главным ограничением является объем памяти. Стандартные веси текстового энкодера CLIP занимают 250 МБ, что превышает свободную оперативную память чипа в 500 раз. Для решения этой проблемы команда разработчиков отказалась от готовых тяжелых моделей и начала строительство архитектуры с нуля, строго подгоняя её под железо.
Были приняты следующие ограничения: * Домен данных: Вместо универсального генератора была выбрана специализированная задача — создание лиц. Использовался датасет CelebA с 200 тысячами фотографий и набором из 40 атрибутов (возраст, пол, эмоции). * Разрешение: Выходная картинка ограничена разрешением 64x64 пикселя. Это минимальный размер для узнавания черт лица, но достаточно малый для размещения в памяти. * Латентная диффузия: Вмesto работы с пикселями, модель оперирует сжатым латентным пространством 16x16x4, что уменьшает объем данных в 12 раз. * Архитектура: Вместо стандартной U-Net использован диффузионный трансформер (DiT) на 64 токенах, сводящий всю сеть к одному типу матричных умножений.
Оптимизация алгоритмов и кода
Для реализации работы в рамках жестких ограничений разработчикам пришлось применить ряд агрессивных методов оптимизации.
Квантование и математика Все веса модели переведены в формат int8, а активации — в int8, с использованием 32-битного аккумулятора. Это позволяет эффективно использовать специализированные инструкции процессора ESP32-S3. Особое внимание уделено квантованию softmax: для его выполнения не требуется отдельный проход вычислений, что экономит ресурсы. Также реализован метод «грубого» стемминга русского языка, где слова обрезались до пяти символов, чтобы уместить словарь из всего 45 токенов.
Аппаратное ускорение Ключевую роль играют векторные инструкции PIE (Processor Instruction Extensions). С их помощью реализовано ядро для умножения матриц, способное перемножать 16 пар чисел int8 за одну инструкцию. Однако стандартный компилятор не умеет генерировать эти коды, поэтому разработчикам пришлось писать критические участки программы на ассемблере.
Работа с памятью Внутренняя SRAM имеет разгрузочную структуру и низкую пропускную способность. Чтобы избежать задержек при работе с PSRAM (медленной внешней памятью), выход данных оптимизирован для работы с внутренней памятью. Кроме того, использована техника «единой арены», когда различные части алгоритма (текстовый энкодер, сам генератор и декодер) работают в одном блоке памяти, переиспользуя его адресное пространство.
Ускорение вычислений Для повышения скорости были заменены стандартные функции математической библиотеки на аппаратные аналоги или таблицы. Например, функции экспоненты и GELU заменены на аппроксимации, а округление чисел выполнено через битовые операции, что в разы сокращает время выполнения.
Результаты работы
Финальная система демонстрирует впечатляющие показатели для устройства такого класса. Процесс генерации изображения занимает всего 5,4 секунды. Это время включает работу текстового энкодера, прохождение диффузионного процесса и декодирование латентного пространства в пиксели.
Сравнение качества изображений, сгенерированных на чипе (int8), и эталонной модели с использованием float32, показало минимальные различия. Разница в яркости составляет менее 1,2 уровня из 255, что практически невидимо глазу. Использование метода Rectified Flow позволяет получать качественные результаты всего за 4–8 шагов, а не за сотни итераций, необходимых классическим диффузионным моделям.
Этот эксперимент доказывает, что сложные алгоритмы искусственного интеллекта можно адаптировать даже для самых простых микроконтроллеров, при условии глубокого понимания архитектуры как программного обеспечения, так и целевого оборудования.