Как запустить Stable Diffusion с Core ML

Core ML — это формат моделей и библиотека машинного обучения, поддерживаемая фреймворками Apple. Если вы хотите запускать модели Stable Diffusion в своих приложениях для macOS или iOS/iPadOS, это руководство покажет, как преобразовать существующие контрольные точки PyTorch в формат Core ML и использовать их для инференса на Python или Swift.

Модели Core ML могут использовать все вычислительные движки, доступные в устройствах Apple: CPU, GPU и Apple Neural Engine (или ANE, тензорный ускоритель, доступный в Mac на Apple Silicon и современных iPhone/iPad). В зависимости от модели и устройства, на котором она выполняется, Core ML также может комбинировать вычислительные движки, поэтому, например, часть модели может выполняться на CPU, а другая — на GPU.

Вы также можете запускать кодовую базу diffusers на Python на Mac с Apple Silicon, используя акселератор mps, встроенный в PyTorch. Этот подход подробно описан в руководстве по mps, но он несовместим с нативными приложениями.

Чекпоинты Stable Diffusion Core ML

Веса Stable Diffusion (или контрольные точки) хранятся в формате PyTorch, поэтому перед использованием в нативных приложениях их необходимо преобразовать в формат Core ML.

К счастью, инженеры Apple разработали инструмент для конвертации на основе diffusers, который позволяет преобразовывать контрольные точки PyTorch в формат Core ML.

Прежде чем конвертировать модель, уделите время изучению Hugging Face Hub – скорее всего, интересующая вас модель уже доступна в формате Core ML:

Организация Apple включает версии Stable Diffusion 1.4, 1.5, 2.0 base и 2.1 base
coreml community включает в себя кастомные дообученные модели
используйте этот фильтр, чтобы вернуть все доступные контрольные точки Core ML

Если вы не можете найти интересующую вас модель, мы рекомендуем следовать инструкциям Apple по конвертации моделей в Core ML.

Выбор варианта Core ML

Модели Stable Diffusion можно преобразовать в различные варианты Core ML, предназначенные для разных целей:

Тип используемых блоков внимания. Операция внимания используется для «фокусировки» на взаимосвязях между различными областями в представлениях изображения и для понимания того, как связаны представления изображения и текста. Внимание требует значительных вычислительных ресурсов и памяти, поэтому существуют различные реализации, учитывающие аппаратные характеристики разных устройств. Для моделей Stable Diffusion в Core ML доступны два варианта внимания:
- split_einsum (предложенный Apple) оптимизирован для устройств ANE, которые доступны в современных iPhone, iPad и компьютерах серии M.
- «Оригинальное» внимание (базовая реализация, используемая в diffusers) совместимо только с CPU/GPU, но не с ANE. Запуск модели на CPU + GPU с использованием внимания original может оказаться быстрее, чем на ANE. Дополнительные сведения см. в этом бенчмарке производительности, а также в дополнительных замерах, предоставленных сообществом.
Поддерживаемый фреймворк для инференса.
- packages подходят для инференса на Python. Их можно использовать для тестирования сконвертированных моделей Core ML перед попыткой интеграции в нативные приложения или если вы хотите исследовать производительность Core ML, но не планируете поддерживать нативные приложения. Например, приложение с веб-интерфейсом может вполне использовать Python-бэкенд Core ML.
- Модели compiled необходимы для кода на Swift. Модели compiled в Hub разделяют веса большой модели UNet на несколько файлов для совместимости с устройствами iOS и iPadOS. Это соответствует параметру конвертации --chunk-unet. Если вы хотите поддерживать нативные приложения, вам нужно выбрать вариант compiled.

Официальные модели Stable Diffusion для Core ML включают следующие варианты, но варианты сообщества могут отличаться:

coreml-stable-diffusion-v1-4
├── README.md
├── original
│   ├── compiled
│   └── packages
└── split_einsum
    ├── compiled
    └── packages

Вы можете скачать и использовать нужный вам вариант, как показано ниже.

Вывод в Core ML на Python

Установите следующие библиотеки для выполнения инференса Core ML в Python:

pip install huggingface_hub
pip install git+https://github.com/apple/ml-stable-diffusion

Загрузка контрольных точек модели

Для выполнения инференса на Python используйте одну из версий, хранящихся в папках packages, так как версии из compiled совместимы только со Swift. Вы можете выбрать, использовать ли вам original или split_einsum attention.

Вот как скачать вариант original внимания из Hub в каталог models:

from huggingface_hub import snapshot_download
from pathlib import Path

repo_id = "apple/coreml-stable-diffusion-v1-4"
variant = "original/packages"

model_path = Path("./models") / (repo_id.split("/")[-1] + "_" + variant.replace("/", "_"))
snapshot_download(repo_id, allow_patterns=f"{variant}/*", local_dir=model_path, local_dir_use_symlinks=False)
print(f"Model downloaded at {model_path}")

Вывод

После загрузки снимка модели вы можете протестировать её с помощью Python-скрипта Apple.

python -m python_coreml_stable_diffusion.pipeline --prompt "a photo of an astronaut riding a horse on mars" -i ./models/coreml-stable-diffusion-v1-4_original_packages/original/packages -o </path/to/output/image> --compute-unit CPU_AND_GPU --seed 93

Передайте путь к скачанной контрольной точке в скрипт с флагом -i. Флаг --compute-unit указывает на аппаратное обеспечение, которое разрешено использовать для инференса. Он должен быть одним из следующих вариантов: ALL, CPU_AND_GPU, CPU_ONLY, CPU_AND_NE. Также можно указать необязательный путь для вывода и seed для воспроизводимости.

Скрипт для инференса предполагает, что вы используете исходную версию модели Stable Diffusion, CompVis/stable-diffusion-v1-4. Если вы используете другую модель, вам необходимо указать её Hub id в командной строке инференса с помощью опции --model-version. Это работает как для уже поддерживаемых моделей, так и для собственных моделей, которые вы обучили или дообучили.

Например, если вы хотите использовать stable-diffusion-v1-5/stable-diffusion-v1-5:

python -m python_coreml_stable_diffusion.pipeline --prompt "a photo of an astronaut riding a horse on mars" --compute-unit ALL -o output --seed 93 -i models/coreml-stable-diffusion-v1-5_original_packages --model-version stable-diffusion-v1-5/stable-diffusion-v1-5

Вывод Core ML на Swift

Запуск инференса на Swift немного быстрее, чем на Python, поскольку модели уже скомпилированы в формате mlmodelc. Это заметно при запуске приложения, когда модель загружается, но не должно быть заметно, если после этого выполнить несколько генераций.

Загрузить

Чтобы выполнить инференс на Swift на вашем Mac, вам нужна одна из compiled версий контрольных точек. Мы рекомендуем скачать их локально с помощью Python-кода, похожего на предыдущий пример, но с одним из compiled вариантов:

from huggingface_hub import snapshot_download
from pathlib import Path

repo_id = "apple/coreml-stable-diffusion-v1-4"
variant = "original/compiled"

model_path = Path("./models") / (repo_id.split("/")[-1] + "_" + variant.replace("/", "_"))
snapshot_download(repo_id, allow_patterns=f"{variant}/*", local_dir=model_path, local_dir_use_symlinks=False)
print(f"Model downloaded at {model_path}")

Вывод

Для запуска инференса клонируйте репозиторий Apple:

git clone https://github.com/apple/ml-stable-diffusion
cd ml-stable-diffusion

Затем используйте командную утилиту Apple Swift Package Manager:

swift run StableDiffusionSample --resource-path models/coreml-stable-diffusion-v1-4_original_compiled --compute-units all "a photo of an astronaut riding a horse on mars"

В параметре --resource-path необходимо указать один из контрольных точек, загруженных на предыдущем шаге, поэтому убедитесь, что он содержит скомпилированные пакеты Core ML с расширением .mlmodelc. Значение параметра --compute-units должно быть одним из следующих: all, cpuOnly, cpuAndGPU, cpuAndNeuralEngine.

Для получения подробной информации обратитесь к инструкциям в репозитории Apple.

Поддерживаемые функции Diffusers

Модели Core ML и код для инференса не поддерживают многие функции, опции и возможности, доступные в 🧨 Diffusers. Вот некоторые из ограничений, о которых стоит помнить:

Модели Core ML подходят только для инференса. Их нельзя использовать для обучения или дообучения.
В Swift были перенесены только два планировщика: планировщик по умолчанию, используемый в Stable Diffusion, и DPMSolverMultistepScheduler, который мы перенесли в Swift из нашей реализации diffusers. Мы рекомендуем использовать DPMSolverMultistepScheduler, поскольку он обеспечивает то же качество примерно за половину шагов.
Негативные промпты, масштаб classifier-free guidance и задачи image-to-image доступны в коде инференса. Продвинутые функции, такие как depth guidance, ControlNet и latent upscalers, пока не поддерживаются.

Репозиторий Apple conversion and inference repo и наш собственный swift-coreml-diffusers предназначены в качестве демонстраций технологий, чтобы другие разработчики могли использовать их как основу для своих проектов.

Если вы считаете, что не хватает каких-либо функций, пожалуйста, не стесняйтесь создавать запрос на добавление функции или, что ещё лучше, PR с вкладом 🙂.

Нативное приложение Diffusers на Swift

Один из простых способов запустить Stable Diffusion на собственном оборудовании Apple — использовать наш репозиторий с открытым исходным кодом на Swift, основанный на diffusers и репозитории Apple для конвертации и инференса. Вы можете изучить код, скомпилировать его с помощью Xcode и адаптировать под собственные задачи. Для вашего удобства также доступно автономное приложение для Mac в App Store, так что вы можете поэкспериментировать с ним, не прибегая к работе с кодом или IDE. Если вы разработчик и пришли к выводу, что Core ML — лучшее решение для создания вашего приложения Stable Diffusion, то остальная часть этого руководства поможет вам начать работу над проектом. Мы с нетерпением ждём, что вы создадите 🙂.

Update on GitHub