Генерация изображений на микроконтроллере: как ESP32-S3 запускает диффузионную модель
Habr AI · оригинал
Материал подготовлен автоматизированной редакционной системой. Факты можно сверить по указанному первоисточнику.
Разработчик реализовал генерацию портретов 64x64 пикселей на микроконтроллере ESP32-S3, используя квантованную диффузионную модель с трансформером. Проект демонстрирует возможность локального инференса нейросетей на встраиваемых устройствах с ограниченной памятью.
Демонстрация работы сложных архитектур (DiT) на микроконтроллере расширяет границы edge-AI, показывая, что генеративные модели могут работать без облака и мощных GPU при глубокой оптимизации.
Обычно запуск диффузионных моделей вроде Stable Diffusion или Flux требует видеокарт с 4–24 ГБ памяти. Однако инженеру удалось реализовать генерацию изображений на микроконтроллере ESP32-S3, который обладает всего 512 КБ оперативной памяти и двумя ядрами по 240 МГц. Результатом стала работающая система, способная за 5,4 секунды создать портрет размером 64x64 пикселя по текстовому описанию, полностью локально на чипе.
Ключевой проблемой было несоответствие ресурсов: веса стандартного текстового энкодера CLIP весят 250 МБ, что в 500 раз превышает всю оперативную память контроллера. Чтобы решить это, автор разработал собственную архитектуру под конкретное железо. Вместо универсальной генерации «всего чего угодно» модель специализируется на лицах, обучаясь на датасете CelebA с 200 тысячами фотографий. Для работы с текстом был создан крошечный трансформер на 300 тысяч параметров со словарем из 45 токенов, а для описаний использовался упрощенный русскоязычный стеммер, обрезающий слова до пяти букв.
Архитектура модели опирается на диффузионный трансформер (DiT), работающий в латентном пространстве 16x16x4. Это позволяет сократить объем данных в 12 раз по сравнению с работой с пикселями. Вместо классического U-Net используется набор матричных умножений, что идеально подходит для оптимизации под векторные инструкции процессора. Для ускорения обучения и инференса применен метод Rectified flow, который позволяет получить приемлемое качество всего за 8 шагов генерации вместо сотен.
Главным техническим достижением стала реализация ядра матричных умножений на ассемблере с использованием набора инструкций PIE (Processor Instruction Extensions). Инструкция EE.VMULAS.S8.ACCX позволяет перемножать 16 пар int8 чисел и добавлять результат в аккумулятор за один такт. Это обеспечило производительность в 1,5 миллиарда умножений в секунду, что в 78 раз быстрее наивной реализации на C. Также была оптимизирована работа с памятью: веса модели хранятся во flash-памяти и читаются напрямую через маппинг, а активации квантуются динамически для сохранения точности.
Особое внимание уделено экономии оперативной памяти. Свободных 288 КБ внутренней SRAM едва хватало для всех буферов. Инженер использовал единую арену памяти для разных фаз работы (энкодер, DiT, декодер), так как они не активны одновременно. Также были оптимизированы функции округления и экспоненты, заменив вызовы библиотечных функций на быстрые битовые манипуляции и полиномы.
Для отладки был создан numpy-симулятор, точно повторяющий арифметику прошивки, включая квантование и генерацию случайных чисел. Это позволило выявить тонкий баг с неверным шагом записи данных в буфере модуляции, который на обученной модели проявился бы лишь как незначительное искажение качества. В итоге средняя разница между квантованной моделью и эталоном составила всего 1,2 уровня яркости из 255, что практически незаметно глазу.
Проект доказывает, что при глубокой оптимизации под конкретное железо даже сложные генеративные модели могут работать на устройствах с крайне ограниченными ресурсами. Это открывает новые возможности для приватной генерации контента в умных устройствах без зависимости от облачных сервисов.