Как запустить Text-to-Image на ESP32-S3: диффузионный трансформер на чипе за 500 рублей

Микроконтроллер ESP32-S3 способен генерировать реалистичные лица по текстовому описанию без подключения к интернету и облачным API. Инференс диффузионной модели занимает 5,4 секунды при разрешении 64×64 пикселя. Для запуска полного пайплайна на чипе с 512 КБ оперативной памяти потребовалось пересобрать архитектуру: заменить U-Net на трансформер DiT, отказаться от тяжелого текстового энкодера CLIP и сжать вычисления до векторных инструкций int8.

В этом руководстве мы разберем инженерные решения, которые позволили уместить генеративную диффузию в микроконтроллер ценой в несколько сотен рублей, обучить модель и запустить локальный Text-to-Image.

Зачем запускать диффузию на микроконтроллере

Полноразмерные диффузионные модели требуют колоссальных вычислительных мощностей и гигабайтов видеопамяти. Stable Diffusion запрашивает от 4 до 8 ГБ VRAM, а современный Flux требует минимум 12 ГБ, а лучше 24 ГБ оперативной памяти ускорителя. Вес одного только текстового энкодера CLIP составляет около 250 МБ, что в 500 раз превышает весь объем встроенной оперативной памяти стандартного микроконтроллера.

Автономный генеративный пайплайн на чипе решает задачу полной изоляции устройства от облака в умных гаджетах, игрушках или носимой электронике. Когда генерация происходит локально на микросхеме за 500 рублей, девайс не зависит от задержек сети, подписок на API и наличия Wi-Fi. Чтобы проект взлетел на кристалле без тензорных ядер, готовую модель приходится отбросить и проектировать архитектуру строго под аппаратные лимиты конкретного кремния.

Шаг 1. Сужаем домен и пишем стеммер на 45 токенов

Универсальная генерация любых объектов в лимитах микроконтроллера невозможна из-за взрывного роста весов. Выбор лица человека в качестве целевого домена дает возможность использовать датасет CelebA (200 тысяч фотографий с 40 бинарными атрибутами) и получить четкие контуры даже в низком разрешении 64×64 пикселя.

Текстовый энкодер CLIP мы заменяем собственной компактной моделью на 300 тысяч параметров со словарем всего из 45 слов. У CelebA нет готовых подписей, поэтому скрипт обучения синтезирует текстовые фразы на лету, случайно комбинируя доступные атрибуты пола, возраста, мимики и прически.

import random

def caption(a, rng=random):
    has = lambda n: a[I[n]] == 1
    male, young = has('Male'), has('Young')
    g = lambda m, f: m if male else f  # согласование по роду
    adj = []
    p = 0.75  # упоминаем признак с вероятностью, чтобы сеть видела неполные фразы

    if rng.random() < p:
        adj.append(g('молодой', 'молодая') if young else g('пожилой', 'пожилая'))
    if has('Smiling'):
        if rng.random() < p:
            adj.append(g('улыбающийся', 'улыбающаяся'))
        elif rng.random() < 0.3:
            adj.append(g('серьёзный', 'серьёзная'))
    return adj

Русский язык создает проблему вариативности окончаний: слова «улыбающийся», «улыбающаяся» и «улыбающимися» словарь считает разными сущностями. Запустить тяжелый стеммер или лемматизатор на микроконтроллере невозможно из-за дефицита памяти, поэтому токенизация решается грубой обрезкой слов до первых пяти символов.

import re

def tokenize(text: str) -> list[str]:
    # Обрезаем слово до 5 символов: "улыбающаяся" -> "улыба", "светлыми" -> "светл"
    return [w[:5] for w in re.findall(r'[а-яa-z]+', text.lower().replace('ё', 'е'))]

Каждое слово ужимается до корня, что сокращает словарь прошивки до 45 токенов, занимающих по 16 байт на запись.

Шаг 2. Собираем облегченный автоэнкодер

Генерировать пиксели напрямую слишком ресурсоемко, поэтому вычисления переносятся в латентное пространство. Автоэнкодер упаковывает изображение 64×64×3 в тензор 16×16×4. Это сокращает объем данных для диффузии ровно в 12 раз, а декодирование в исходные цвета выполняется всего один раз на финальном этапе генерации.

Энкодер используется только на этапе обучения на видеокарте и насчитывает 1,9 млн параметров. Декодер переносится в прошивку ESP32, поэтому его архитектура зачищена от тяжелых операций: исключена GroupNorm, которая медленно считается на CPU и плохо поддается int8-квантованию.

import torch.nn as nn

class Decoder(nn.Module):
    CH = (64, 48, 32)

    def __init__(self, z=4):
        super().__init__()
        a, b, c = self.CH
        self.net = nn.Sequential(
            nn.Conv2d(z, a, 3, padding=1),
            Res(a),
            nn.Upsample(scale_factor=2),
            nn.Conv2d(a, b, 3, padding=1),
            Res(b),
            nn.Upsample(scale_factor=2),
            nn.Conv2d(b, c, 3, padding=1),
            Res(c),
            nn.ReLU(),
            nn.Conv2d(c, 3, 3, padding=1)
        )

Весь декодер весит 179 тысяч параметров и состоит исключительно из сверток 3×3, блоков ReLU и nearest-апсемплинга. Модель обучается на Google Colab с GPU Nvidia T4 за 20 минут (20 тысяч шагов) с комбинацией функций потерь L1, LPIPS и небольшим добавлением KL-дивергенции.

Шаг 3. Строим DiT с оптимизацией adaLN-single

Классическая U-Net содержит ветвящиеся skip-соединения, требующие хранения промежуточных тензоров в памяти. Трансформер DiT (Diffusion Transformer) разбивает латент 16×16×4 на неперекрывающиеся патчи 2×2, превращая картинку в последовательность из 64 токенов по 16 признаков. Вся диффузия сводится к регулярным матричным умножениям одинакового размера.

Для решения этой проблемы заимствуется трюк adaLN-single из модели PixArt-α: вектор условия шага и текста проецируется один раз в начале сети, а блоки трансформера лишь масштабируют полученные коэффициенты.

Параметры трансформера DiT:
  Ширина слоя (D): 192
  Количество блоков: 6
  Головы внимания: 6 (по 32 размерности на голову)
  Размерность MLP: x3 (576)
  Всего параметров: 2.9 млн (включая текстовый энкодер)

Благодаря сквозному сведению математики к однотипным матричным операциям вычислительное ядро прошивки оптимизируется под одну функцию перемножения матриц.

Шаг 4. Оптимизируем диффузию под SIMD и память

Стандартный планировщик DDPM требует от 50 до 1000 итераций денойзинга, что неприменимо для микроконтроллера. Переход на Rectified Flow сокращает траекторию генерации всего до 8 шагов без потери качества лиц.

Алгоритм Classifier-Free Guidance (CFG) традиционно удваивает нагрузку, требуя два последовательных прогона трансформера на каждом шаге (с промптом и без). Чтобы исключить повторные чтения весов из медленной flash-памяти, условный и безусловный проходы склеиваются в один общий батч размером 2. Веса считываются из кэша один раз на два потока вычислений.

Вычисления переводятся в целочисленный формат: веса и активации квантуются в int8, а аккумуляция ведется в регистры int32. Расширение PIE (Processor Instruction Extension) в архитектуре Xtensa LX7 задействует 128-битные векторные регистры, выполняя 16 операций умножения int8 за один такт микроконтроллера.

Где ломается система

Главным узким местом ESP32-S3 при инференсе трансформеров становится пропускная способность шины памяти. Внешняя PSRAM и Flash подключены по quad SPI, поэтому частый случайный доступ к весам вызывает просадки частоты конвейера CPU.

При квантовании функций внимания softmax в int8 без запаса точности активации переполняют аккумуляторы int32. Это приводит к появлению контрастных артефактов и «выбитых» белых пикселей на сгенерированных лицах.

Ограниченный словарь из 45 токенов и пятибуквенный стемминг приводят к коллизиям понятий. Слова со сходными приставками могут активировать одни и те же нейроны, меняя атрибуты лица вопреки смыслу запроса.

Что попробовать дальше

Для развития автономной генерации на чипе стоит подключить компактный SPI-дисплей (например, ST7789 или ILI9341), настроив отрисовку картинки напрямую из буфера декодера без передачи по USB.

Второй вектор — добавление интерполяции в скрытом пространстве. Фиксируя шум и плавно смещая текстовый вектор между двумя промптами («улыбающийся» и «серьезный»), можно получить плавную анимацию мимики прямо на вычислительных ресурсах микроконтроллера.

FAQ

Можно ли сгенерировать пейзажи или технику вместо лиц?

Архитектура позволяет обучить модель на любом другом датасете с малым числом классов. Однако лица обладают выраженной геометрической симметрией, поэтому в разрешении 64×64 они выглядят узнаваемо, в то время как детализированные объекты со сложной сценой превратятся в кашу.

Подойдет ли обычный чип ESP32 или ESP8266?

Нет, проект требует именно чип ESP32-S3. В классическом ESP32 и ESP8266 отсутствуют 128-битные векторные инструкции PIE, а объема оперативной памяти недостаточно для хранения даже сжатых активаций латентного пространства.

Зачем нужен стеммер по первым пяти буквам?

Русский язык имеет развитую систему падежей и окончаний, из-за чего обычные словари разрастаются до десятков тысяч токенов. Обрезка до пяти символов сводит разные словоформы («молодой», «молодая», «молодому») к единому токену «молод», укладывая словарь в 45 записей.

Почему отказались от нормализации GroupNorm в декодере?

Слой GroupNorm требует вычисления среднего и дисперсии по группам каналов в реальном времени, что плохо распараллеливается на процессоре микроконтроллера и дает высокую погрешность при целочисленном int8-квантовании. Замена на стандартные свертки с ReLU сохранила производительность.

Источники

• Habr ML: Text‑to‑image на ESP32: диффузионный трансформер в микроконтроллере