Оптимизация Stable Diffusion для iPhone, iPad и Mac: практический гайд по Core ML и квантованию

Stable Diffusion в базовой конфигурации занимает от 1.5 до 4 ГБ в FP16 и требует несколько секунд на генерацию одного изображения даже на мощных GPU. На iPhone с 6 ГБ общей памяти и пассивным охлаждением такой сценарий неработоспособен. Apple решила эту проблему через Core ML и coremltools 7.0: палетизация весов до 6 бит сжимает модель в 4-5 раз без заметной потери качества, а оптимизация внимания SPLIT_EINSUM_V2 дополнительно ускоряет инференс на Neural Engine. В этом гайде разбираем механику оптимизаций, приводим бенчмарки и показываем пошаговую конвертацию кастомной модели для iOS и macOS.

Практическая ценность подхода: Stable Diffusion запускается локально на устройствах, которые пользователь носит в кармане. Без облака, без подписки, без задержек сети. Для разработчиков это открывает сценарии офлайн-генерации в приложениях, приватной обработки изображений и встраивания диффузионных моделей в нативные продукты.

Ключевые инновации Apple: палетизация весов до 6 бит и другие оптимизации

Core ML поддерживает несколько методов сжатия моделей: линейное квантование, палетизацию и комбинированные схемы. Для Stable Diffusion наиболее эффективной оказалась палетизация весов, при которой значения группируются в кластеры, а каждый вес представляется индексом в кодовой книге. При 6 битах на вес кодовая книга содержит 64 кластера, что достаточно для сохранения распределения параметров UNet.

Палетизация весов: сжатие без потери качества

Механизм работает так: coremltools анализирует распределение весов каждого слоя, находит оптимальные центры кластеров методом k-means, затем заменяет исходные значения на индексы. При инференсе Core ML восстанавливает веса из кодовой книги на лету. В отличие от простого усечения битов, палетизация сохраняет структуру распределения, поэтому артефакты генерации минимальны даже при 6 битах.

Цифры: UNet модели Stable Diffusion 1.5 сжимается с 1.5 ГБ до примерно 200-300 МБ при 6-битной палетизации. Текстовый энкодер и VAE квантуются отдельно, обычно в 8 бит, поскольку они более чувствительны к потере точности. Суммарный размер модели уменьшается в 4-5 раз, что позволяет разместить её в памяти iPhone с 6 ГБ ОЗУ.

Для сравнения, 4-битная палетизация даёт дополнительное сжатие до 150-180 МБ, но появляются видимые артефакты на сложных текстурах и мелких деталях. 6 бит - это практический компромисс между размером и качеством, который Apple рекомендует как базовый.

Оптимизация внимания: SPLIT_EINSUM_V2

Слои внимания в UNet - основное узкое место при инференсе на Neural Engine. Операция einsum, используемая для вычисления матриц Query, Key и Value, плохо распараллеливается на специализированных ускорителях Apple. SPLIT_EINSUM_V2 разбивает вычисление на несколько более простых операций, которые Neural Engine выполняет параллельно.

Эффект особенно заметен на iPhone и iPad, где Neural Engine имеет больше вычислительных блоков, чем GPU. На Mac с мощным GPU ускорение от SPLIT_EINSUM_V2 скромнее, но всё равно измеримо: 10-15% на M1 Pro и до 20% на базовом M1. Опция включается при конвертации через coremltools и не требует изменений в архитектуре модели.

Результаты производительности: насколько быстрее и меньше

Бенчмарки на устройствах Apple показывают, что оптимизированная модель генерирует изображение 512x512 за 2-4 секунды на iPhone 15 Pro и за 1.5-2.5 секунды на iPad Pro с M2. MacBook Pro с M3 Max справляется за 0.8-1.2 секунды. Без оптимизаций эти же устройства либо не запускают модель вовсе из-за нехватки памяти, либо тратят 15-30 секунд на изображение.

Сравнение на разных устройствах Apple

На устройствах с Neural Engine ускорение достигает 3-5 раз по сравнению с запуском через PyTorch на CPU. Для Mac с дискретным GPU выигрыш меньше, но экономия памяти критична для многозадачных сценариев.

Если вас интересует оптимизация Stable Diffusion на CPU с использованием OpenVINO и Token Merging, в нашем материале по 8-битной квантизации и ToMe разобран альтернативный подход с ускорением в 5.1 раза.

Практическое руководство: конвертация и запуск Stable Diffusion с coremltools 7.0 и Xcode 15

Для конвертации нужны Python 3.10+, coremltools 7.0, Xcode 15 и модель Stable Diffusion в формате PyTorch. Процесс состоит из трёх этапов: подготовка модели, конвертация с оптимизациями, интеграция в приложение.

Шаг 1: Подготовка модели и установка coremltools

Установите coremltools и зависимости:

pip install coremltools==7.0 torch torchvision transformers

Загрузите модель Stable Diffusion с Hugging Face и экспортируйте её в формат, который принимает coremltools. Для кастомных моделей на базе SD 1.5 или SDXL используйте torchscript или ONNX как промежуточный формат. Проверьте, что все операции в модели поддерживаются coremltools: динамический контроль потока и нестандартные кастомные слои могут вызвать ошибки конвертации.

Шаг 2: Конвертация с квантованием и палетизацией

Ключевой шаг - вызов coremltools.convert с параметрами оптимизации:

import coremltools as ct

model = ct.convert(
    torchscript_model,
    compute_units=ct.ComputeUnit.ALL,
    minimum_deployment_target=ct.target.iOS17,
    convert_to="mlprogram",
    compute_precision=ct.precision.FLOAT16,
    quantization_mode="palettized",
    nbits=6,
    optimize_attention=True
)
model.save("stable_diffusion_6bit.mlpackage")

Параметр quantization_mode="palettized" включает палетизацию, nbits=6 задаёт битность. Для текстового энкодера и VAE используйте nbits=8. Параметр optimize_attention=True активирует SPLIT_EINSUM_V2. Диапазон квантования от 2 до 8 бит позволяет гибко настраивать компромисс между размером и качеством.

Шаг 3: Оптимизация внимания и финальная сборка

После сохранения .mlpackage добавьте модель в Xcode проект. В Info.plist укажите поддержку Neural Engine и минимальную версию iOS 17 или macOS 14. Для инференса используйте стандартный API Core ML:

let model = try MLModel(contentsOf: modelURL)
let prediction = try model.prediction(from: input)

Проверьте, что модель загружается на устройстве и использует Neural Engine. В Instruments отслеживайте пиковое потребление памяти: при 6-битной палетизации UNet должен занимать не более 300 МБ.

Для более глубокого понимания квантования на Apple Silicon рекомендуем разбор 4-битных методов в MLX, где сравниваются OptiQ, Unsloth и нативные схемы.

Ограничения и подводные камни

6-битная палетизация сохраняет качество генерации на уровне, близком к FP16, но 4-битная и ниже даёт заметные артефакты: размытие мелких деталей, цветовые искажения на градиентах, потерю текстур. Для моделей с тонкими стилевыми настройками агрессивное квантование может испортить результат.

Конвертация кастомных моделей не всегда проходит гладко. Модели с LoRA-адаптерами, кастомными слоями внимания или нестандартными операциями нормализации требуют ручных правок в torchscript-представлении. Некоторые операции, специфичные для PyTorch, не имеют аналогов в Core ML и вызывают ошибки на этапе конвертации.

Требования к памяти: для Stable Diffusion 1.5 с 6-битной палетизацией нужно минимум 6 ГБ ОЗУ на iPhone. На устройствах с 4 ГБ модель запускается, но при генерации изображений 512x512 возможны вылеты из-за нехватки памяти. Для SDXL требования выше: 8 ГБ и более.

Будущее: 4-битное квантование при обучении и дальнейшие перспективы

Пост-обучающая палетизация имеет предел: ниже 4 бит качество деградирует слишком сильно. Решение - квантование при обучении, когда модель изначально обучается с учётом низкой точности весов. Методы типа QLoRA позволяют дообучать Stable Diffusion в 4-битном представлении, сохраняя качество генерации на уровне 8-битных моделей.

Apple пока не поддерживает 4-битное квантование при обучении в coremltools, но направление активно развивается. Если Core ML добавит нативную поддержку моделей, обученных в 4 бит, размер UNet уменьшится до 100-150 МБ, а инференс ускорится ещё на 30-40%. Это сделает возможным запуск SDXL на iPhone с 6 ГБ памяти.

Параллельно развиваются методы спекулятивного декодирования для LLM на Apple Silicon. В статье про mlx-dspark и спекулятивное декодирование показано ускорение до 3 раз для Qwen3.8-27B, что указывает на общий тренд оптимизации инференса на устройствах Apple.

Заключение

Палетизация весов до 6 бит и SPLIT_EINSUM_V2 в coremltools 7.0 делают запуск Stable Diffusion на iPhone, iPad и Mac практичным. Модель сжимается с 1.5 ГБ до 280 МБ, генерация изображения 512x512 занимает 2-4 секунды на iPhone 15 Pro. Пошаговая конвертация занимает меньше часа при наличии подготовленной модели.

Попробуйте конвертировать свою кастомную модель по инструкции выше. Начните с 6-битной палетизации и optimize_attention=True, затем протестируйте качество на своих промптах. Если результат устраивает, попробуйте 4 бита для дополнительного сжатия. Документация Apple по coremltools содержит полный список параметров квантования и примеры для разных архитектур.

Подписаться на канал