Stable Diffusion запустили на системе с 512 мегабайтами памяти

Генерация изображения занимает массу времени, но она хотя бы работает.

Разработчик Вито Плантамура сумел запустить нейросеть-генератор изображений Stable Diffusion на системе с 512 мегабайтами оперативной памяти. Нейросеть работала не на обычном PC, а на Raspberry Pi Zero 2 W.

Хотя у компактного вычислительного устройства есть полгигабайта памяти, Stable Diffusion во время работы потребляла только 260 мегабайт памяти. Чтобы добиться этого, Плантамуре пришлось серьезно оптимизировать нейросеть.

Обычно Stable Diffusion требуется порядка 8 гигабайт памяти. Для оптимизации использования ОЗУ Вито создал библиотеку OnnxStream, меняющую принцип работы модели с памятью.

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

По оценке Вито, OnnxStream позволяет снизить потребление памяти в 55 раз, при этом время генерации увеличивается максимум вдвое. На Raspberry Pi Zero 2 W генерация изображения, показанного ниже, заняла три часа с точностью W8A8.

Для сравнения: вот то же изображение, сгенерированное с точностью W8A32 на более быстрой системе.

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