Если Stable Diffusion на PyTorch тормозит на видеокарте AMD или встроенной графике, конвертация модели в формат ONNX часто становится единственным способом получить приемлемую скорость генерации — через ONNX Runtime с провайдером DirectML нагрузка уходит на GPU даже без поддержки CUDA. Формат ONNX (Open Neural Network Exchange) — это открытый стандарт представления нейросетей, который позволяет запускать одну и ту же модель на разных движках вывода: CPU, DirectML, CUDA, TensorRT.
В этой статье разберём, как устроена ONNX-версия Stable Diffusion, какие инструменты используются для конвертации, как запустить инференс через ONNX Runtime и какие ошибки встречаются чаще всего. Материал ориентирован на пользователей Windows, но общие принципы применимы и к Linux.
Что такое ONNX и зачем он нужен для Stable Diffusion
ONNX — это промежуточное представление вычислительного графа нейросети. Модель, обученная в PyTorch, экспортируется в файлы .onnx, после чего её может выполнять любой совместимый рантайм. Для Stable Diffusion это особенно полезно, потому что оригинальная реализация жёстко привязана к PyTorch и CUDA.
Практических причин переходить на ONNX несколько. Во-первых, DirectML-провайдер позволяет задействовать видеокарты AMD и Intel, для которых поддержка в PyTorch-экосистеме ограничена. Во-вторых, ONNX Runtime умеет оптимизировать граф: выполнять слияние операций, квантование и другие преобразования, снижающие потребление памяти. В-третьих, ONNX-модели проще встраивать в десктопные приложения без тяжёлой зависимости от PyTorch.
Стоит понимать и ограничения. Конвертированный пайплайн обычно менее гибок: смена сэмплера, LoRA-адаптеров или нестандартных шагов генерации может потребовать повторной конвертации или дополнительного кода. Если вам нужна максимальная гибкость экспериментов, классический Automatic1111 или ComfyUI на PyTorch останутся удобнее.
Из чего состоит ONNX-версия пайплайна
Stable Diffusion — это не одна сеть, а связка из нескольких моделей, и в ONNX каждая конвертируется отдельно. Типичный набор файлов выглядит так:
- 📝 Text Encoder — преобразует текстовый промпт в эмбеддинги (модель CLIP).
- 🎨 UNet — основная диффузионная сеть, выполняющая итеративное удаление шума; самый тяжёлый компонент.
- 🖼️ VAE Decoder — декодирует латентное представление в итоговое изображение.
- 🔐 VAE Encoder — нужен для режимов img2img, кодирует исходную картинку в латенты.
- 🛡️ Safety Checker — опциональный фильтр контента, присутствует не во всех сборках.
Каждый компонент лежит в отдельной папке вместе со своим файлом .onnx и конфигурацией. Дополнительно рядом хранятся токенизатор и конфиг планировщика (scheduler) — они не являются нейросетями и в конвертации не нуждаются.
⚠️ Внимание: ONNX-модель привязана к разрешению и размеру батча, заданным при экспорте, если не использовались динамические оси. Картинка 512×512 и 768×768 — это, как правило, разные экспорты UNet. Проверяйте, под какое разрешение собрана конкретная сборка.
Способы получить ONNX-модель
Есть три рабочих пути: скачать готовую конвертированную модель, выполнить экспорт самостоятельно через Optimum от Hugging Face либо воспользоваться приложением, которое делает это автоматически при первом запуске. Выбор зависит от того, насколько стандартная модель вам нужна.
Самостоятельная конвертация через библиотеку Optimum выполняется одной командой. Потребуется установленный Python с пакетами optimum[onnxruntime] и diffusers:
optimum-cli export onnx --model stabilityai/stable-diffusion-2-1 --task stable-diffusion ./sd-onnxКоманда скачает веса с Hugging Face и разложит пайплайн по компонентам в указанную папку. Для кастомных чекпоинтов в формате .safetensors сначала потребуется конвертация в формат Diffusers — прямой экспорт из одиночного файла весов этой командой не поддерживается.
☑️ Подготовка к конвертации в ONNX
Запуск через ONNX Runtime: выбор провайдера
За выполнение ONNX-моделей отвечает ONNX Runtime, а конкретное устройство вычислений определяется выбранным execution provider. От правильного выбора зависит, попадёт ли нагрузка на GPU или вся генерация пойдёт через процессор.
| Провайдер | Устройство | Когда использовать |
|---|---|---|
| CUDA | NVIDIA GPU | Если PyTorch-вариант недоступен, но нужна скорость |
| DirectML | Любой GPU с DirectX 12 (AMD, Intel, NVIDIA) | Основной вариант для видеокарт AMD и Intel под Windows |
| CPU | Процессор | Резервный вариант; генерация заметно медленнее |
| TensorRT | NVIDIA GPU | Максимальная скорость, но требует отдельной сборки движка |
Для запуска удобнее всего использовать класс OnnxStableDiffusionPipeline из библиотеки diffusers — он принимает папку с ONNX-компонентами и параметр provider. Для DirectML устанавливается пакет onnxruntime-directml, для CUDA — onnxruntime-gpu. Ставить оба одновременно не нужно: они конфликтуют между собой.
Если после запуска скорость не отличается от CPU-режима, первым делом проверьте, какой провайдер реально активировался — рантайм при несовместимости молча откатывается на процессор. Диагностику даёт вывод доступных провайдеров через onnxruntime.get_available_providers().
Оптимизация скорости и памяти
Даже после конвертации генерация может оставаться медленной или упираться в нехватку видеопамяти. Несколько приёмов помогают выжать из ONNX-пайплайна больше.
- ⚡ Снижение числа шагов — планировщики вроде DPM-Solver или LCM дают приемлемый результат за меньшее количество итераций UNet.
- 🧮 Половинная точность (float16) — некоторые сборки экспортируются в fp16, что сокращает потребление памяти; поддержка зависит от провайдера и железа.
- 📉 Уменьшение разрешения — генерация 512×512 вместо 768×768 резко снижает нагрузку на UNet.
- 🔧 Готовые оптимизированные приложения — например, сборки на базе Olive от Microsoft выполняют дополнительную оптимизацию графа под конкретное железо.
⚠️ Внимание: квантование и агрессивная оптимизация графа могут заметно менять итоговое изображение — детали «плывут», цвета смещаются. Перед массовой генерацией сравните результат оптимизированной и исходной модели на одном и том же сиде и промпте.
Типичные ошибки и их решения
Чаще всего пользователи сталкиваются с ошибкой вида INVALID_ARGUMENT или несовпадением размерностей входа — это почти всегда означает, что разрешение или батч при запуске не совпадают с теми, что были зафиксированы при экспорте. Решение — использовать параметры, под которые собрана модель, или переконвертировать её с нужными значениями.
Вторая частая ситуация — DirectML не видит видеокарту. Проверьте, что установлен актуальный драйвер GPU и система поддерживает DirectX 12. Также убедитесь, что установлен именно пакет onnxruntime-directml, а не базовый onnxruntime: при неверном пакете провайдер DirectML в списке доступных просто не появится.
Третья группа проблем — качество результата. Если ONNX-версия выдаёт заметно более мыльные или артефактные картинки по сравнению с PyTorch-оригиналом, возможные причины — экспорт в fp16 на железе без надёжной поддержки половинной точности, несовпадение планировщика или ошибки при конвертации VAE. Для диагностики сгенерируйте изображение с идентичными промптом, сидом и числом шагов в обоих пайплайнах и сравните попиксельно.
Почему результаты ONNX и PyTorch могут отличаться
Даже без ошибок конвертации результаты редко совпадают бит-в-бит. Причины — другой порядок вычислений в оптимизированном графе, иные реализации операторов в рантайме и округление в fp16. Различия на уровне отдельных пикселей — норма; системная потеря детализации — уже повод проверять конвертацию.
Когда ONNX не нужен
Честно обозначим границы применимости. Если у вас видеокарта NVIDIA с достаточным объёмом памяти и всё работает через PyTorch, переход на ONNX редко даёт выигрыш, оправдывающий потерю гибкости. Экосистема LoRA, ControlNet и кастомных нод вокруг PyTorch-интерфейсов развита заметно шире.
ONNX имеет смысл, когда нужен запуск на GPU без CUDA, встраивание генерации в собственное приложение без тяжёлых зависимостей или максимально предсказуемое окружение для развёртывания. Оцените свою задачу до конвертации — обратный путь из ONNX в редактируемый чекпоинт невозможен.
Частые вопросы
Можно ли конвертировать любой чекпоинт Stable Diffusion в ONNX?
Технически да, если модель совместима с архитектурой пайплайна (SD 1.x, SD 2.x, SDXL). Одиночный файл .safetensors сначала преобразуется в формат Diffusers, затем экспортируется в ONNX. Модели с нестандартной архитектурой могут потребовать ручной доработки экспорта.
Работает ли ONNX-версия на видеокартах AMD?
Да, через провайдер DirectML под Windows — это один из главных сценариев использования ONNX для Stable Diffusion. Требуется актуальный драйвер GPU и пакет onnxruntime-directml. Производительность зависит от конкретной видеокарты.
Поддерживаются ли LoRA в ONNX-пайплайне?
Напрямую — в ограниченном виде. LoRA-веса обычно нужно объединить с базовой моделью до конвертации. Динамическая подгрузка адаптеров, как в Automatic1111, для ONNX не предусмотрена без дополнительных решений.
Почему генерация на ONNX идёт через CPU, хотя GPU есть?
ONNX Runtime молча откатывается на CPU, если выбранный провайдер недоступен. Проверьте список через onnxruntime.get_available_providers(), убедитесь, что установлен правильный пакет (onnxruntime-gpu или onnxruntime-directml) и обновлён драйвер видеокарты.
Можно ли менять разрешение генерации после конвертации?
Только если при экспорте были заданы динамические оси для размеров изображения. В большинстве готовых сборок разрешение зафиксировано, и для другого размера потребуется повторная конвертация модели.