Инженерия, позволяющая запустить диффузию на смартфоне

Диффузионные модели создавались для GPU в дата-центрах. Как реальный движок диффузии уменьшают, чтобы он работал на смартфоне, и какие компромиссы приходится делать.

Что на самом деле делает диффузионная модель

Диффузионная модель создаёт изображение, начиная с чистого шума и убирая его шаг за шагом. На каждом шаге сеть выполняется один раз: она предсказывает немного более чистую версию картинки, и после достаточного числа шагов остаётся связное изображение. Именно этот цикл делает диффузию такой требовательной к вычислениям: для одного результата полная модель запускается десятки раз, а не один раз, как у классификатора.

Проблема размера

Эти модели проектировались для видеокарт с вентиляторами и отдельными кабелями питания. Полноценная диффузионная модель в полной точности может занимать несколько гигабайт, и при каждом шаге очистки от шума веса снова читаются из памяти. На смартфоне и объём хранилища, и пропускная способность памяти — ресурсы ограниченные, поэтому перенести десктопный конвейер целиком без изменений никогда не получилось бы.

Как её уменьшают

Главный инструмент — квантизация: веса модели хранятся в 8 или даже 4 битах вместо 16 или 32, что сокращает объём в разы при сохранении большей части её возможностей. Добавьте к этому нативный движок инференса, написанный для мобильных CPU и GPU, разумные разрешения по умолчанию и настроенное число шагов — и нагрузка дата-центра превращается в задачу, посильную для чипа в кармане.

Честные компромиссы

Уменьшенные квантизованные модели и меньшее число шагов означают, что мобильная диффузия немного теряет в резкости и в сложных формулировках промптов, которые под силу самым большим моделям. Генерация занимает от нескольких секунд до минуты, а не даёт почти мгновенный результат серверной фермы — нужно лишь немного терпения, никакой магии. Зато весь конвейер работает там, где возникла идея: без загрузки в сеть и без очередей.

Где всё работает нативно

OIG построен именно на этом: реальный движок диффузии работает нативно на смартфоне, а не прячется за облачным API. Задаёте промпт — получаете изображение, всё вычисляется на устройстве: тот же цикл очистки от шума, что и на десктопе, но адаптированный к мобильному железу. Для генерации изображений это разница между тем, чтобы владеть инструментом, и тем, чтобы лишь мельком взглянуть на него.