31,18 ГиБ весов модели генерации изображений Krea 2 Turbo удалось уместить в 32 ГБ unified memory MacBook Pro. Без Python. Без облаков. Порт на Swift и MLX загружает компоненты поэтапно: сначала Qwen кодирует промпт, затем DiT выполняет диффузию, и только потом VAE декодирует латенты в картинку. Пик потребления памяти не превышает доступный лимит.
В процессе порта обнаружился критический баг с неправильными RoPE-позициями для пяти токенов. Старые тесты на cosine similarity показывали ~1.0 - почти полное совпадение. Визуально разница была незаметна. Но исправление ошибки изменило 99,32% пикселей итогового изображения. Этот случай - жёсткое напоминание: полагаться только на метрики близости эмбеддингов при валидации порта нельзя.
Одинаковый seed не гарантирует идентичность персонажа. Причина - недетерминированность в порядке загрузки и выгрузки компонентов между запусками. Реализованные Regional Prompts и Character Sheet работают в однопроходном режиме, давая контроль над композицией и внешностью персонажа без файнтюнинга. Разберём архитектуру, баг и практические уроки порта.
Архитектура порта: как 31,18 ГиБ весов уместились в 32 ГБ unified memory
Прямой запуск Krea 2 Turbo на потребительском Apple Silicon выглядит невозможным. 31,18 ГиБ весов при доступных 32 ГБ unified memory - запаса всего 0,82 ГБ на системные нужды, буферы и промежуточные тензоры. Python-стек с PyTorch и MPS бэкендом в эту память не влезает: накладные расходы на аллокаторы, фрагментация и копирования между CPU и GPU съедают гигабайты. Решение - нативный Swift-стек с MLX и поэтапная загрузка компонентов.
Почему Swift и MLX: отказ от Python ради unified memory
PyTorch на Apple Silicon работает через MPS-бэкенд, который эмулирует дискретную модель памяти. Тензоры копируются между пулами CPU и GPU, даже когда физически память едина. Фрагментация растёт. Пиковое потребление всегда выше суммы весов. Для модели на 31 ГиБ это означает немедленный out-of-memory.
MLX использует прямой доступ к unified memory. Нет копирований. Нет раздельных аллокаторов. Ленивые вычисления позволяют откладывать создание промежуточных тензоров до момента реальной необходимости, что срезает пики потребления. Swift даёт контроль над жизненным циклом объектов: компонент загрузился, отработал, веса выгружены - память освобождена детерминированно, без ожидания сборщика мусора.
SwiftUI в этом стеке решает единственную задачу - минимальный UI для ввода промпта и отображения результата. Никаких Electron-подобных прослоек, которые тоже жрут память. Всё нативное, всё в unified memory.
Поэтапная загрузка: Qwen → DiT → VAE
Схема загрузки трёх компонентов последовательно - ключевой архитектурный паттерн порта. Каждый компонент живёт в памяти ровно столько, сколько нужно для его задачи:
- Qwen (~8 ГБ) загружается первым. Кодирует текстовый промпт в эмбеддинги. Результат - компактный тензор, который остаётся в памяти.
- Qwen выгружается. Загружается DiT (~20 ГБ) - Diffusion Transformer. Выполняет процесс диффузии: от шума к латентному представлению изображения. Пик памяти здесь максимальный: веса DiT + латенты + эмбеддинги от Qwen.
- DiT выгружается. Загружается VAE (~3 ГБ). Декодирует латенты в итоговое изображение.
Суммарный пик памяти не превышает ~28-29 ГБ. Оставшиеся 3-4 ГБ уходят на системные нужды и фреймворк MLX. Веса могут подгружаться с SSD напрямую через memory-mapped файлы, но это замедляет инференс - для продакшен-сценариев лучше держать модель в RAM.
Этот подход применим к любой diffusion-модели с разделяемыми компонентами. Если текстовый энкодер, дифа и декодер - отдельные модули, их можно загружать последовательно. Ограничение - сумма пика весов самого тяжёлого компонента плюс промежуточные тензоры должна влезать в доступную unified memory. Для 32 ГБ машин потолок - модели с DiT до ~22 ГБ.
Критический баг RoPE: когда cosine similarity ~1.0 недостаточно
Баг обнаружился случайно. Порт выдавал визуально корректные изображения. Промежуточные эмбеддинги при сравнении с эталонным Python-раннером давали cosine similarity ~1.0 - расхождение в пятом-шестом знаке после запятой, которое обычно списывают на разницу в реализации операций. Но при попиксельном сравнении итоговых изображений выяснилось: различается 99,32% пикселей.
Как пять токенов сломали воспроизводимость
Проблема сидела в позиционном кодировании RoPE (Rotary Position Embedding). При токенизации промпта пять специальных токенов - <|startoftext|>, <|endoftext|> и три паддинг-токена - получали индексы позиций, смещённые на единицу относительно эталонной реализации. Ошибка в индексации цикла: сдвиг начинался не с нулевой позиции, а с первой.
Влияние на attention оказалось коварным. Модель «видела» неправильные расстояния между токенами. Для семантически значимых токенов ошибка в одну позицию на коротких последовательностях почти не меняла attention-веса - отсюда cosine similarity ~1.0. Но в процессе диффузии, где позиционная информация влияет на пространственное распределение признаков, накопленная ошибка давала совершенно другую картину латентов.
Связь с проблемой идентичности персонажа прямая. Даже при фиксированном seed, разные запуски порта могли давать разные результаты из-за недетерминированности в порядке загрузки компонентов. Пять токенов с неправильными позициями добавляли дополнительный источник вариативности, который не ловился стандартными тестами.
Уроки для тестирования генеративных моделей
Cosine similarity на эмбеддингах не работает как единственный критерий корректности порта. Промежуточные представления могут совпадать с точностью до пятого знака, а итоговое изображение - различаться почти полностью. Необходимо попиксельное сравнение эталонных и тестовых изображений.
Фиксация всех случайных факторов критична: seed, порядок загрузки модулей, типы данных (float32 vs bfloat16), даже версия MLX. Разница в реализации exp2 на GPU разных поколений Apple Silicon может давать расхождение в младших битах, которое каскадно расходится через 50 шагов диффузии.
Практический чек-лист для валидации порта diffusion-модели:
- Зафиксировать seed и все параметры генерации.
- Прогнать один и тот же промпт на эталонной и тестовой реализации.
- Сравнить попиксельно итоговые изображения (допустимое расхождение - 0 пикселей при одинаковых типах данных).
- Если расхождение есть - сравнивать латенты после каждого шага диффузии, чтобы локализовать шаг расхождения.
- Не полагаться на cosine similarity эмбеддингов как на достаточный тест.
Этот баг - не экзотика. При портировании моделей между фреймворками ошибки позиционного кодирования возникают регулярно. Разница в индексации (0-based vs 1-based), обработка специальных токенов, паддинг - типичные места слома.
Regional Prompts и Character Sheet в однопроходном режиме
Порт Krea 2 Turbo реализует две продвинутые функции контролируемой генерации без многократных прогонов модели. Обе работают за один проход диффузии.
Regional Prompts: контроль композиции без дообучения
Regional Prompts позволяют задать разные текстовые описания для разных областей изображения. Пример: «слева - рыжий кот на диване, справа - чёрная собака у окна». Области задаются координатами или масками. Механизм работает через модификацию attention-масок в DiT: токены левого промпта «видят» только левую часть латентного пространства, токены правого - только правую.
Взаимодействие с основным промптом настраивается весами. Основной промпт задаёт общую сцену и стиль, региональные - уточняют детали в конкретных зонах. Артефакты на стыках возможны, если маски имеют резкие границы. Размытие границ маски на 10-15% ширины региона сглаживает переходы.
Character Sheet: стабильность персонажа между генерациями
Character Sheet фиксирует внешность персонажа через эмбеддинги референсного изображения. Референс прогоняется через VAE-энкодер, полученные латенты конкатенируются с текстовыми эмбеддингами промпта. DiT получает сигнал «этот персонаж должен выглядеть вот так».
Ограничения метода: поза и освещение на референсе сильно влияют на результат. Если референс снят анфас при дневном свете, а промпт описывает персонажа в профиль при неоновом освещении, консистентность падает. Решение - несколько референсов с разных ракурсов, но это увеличивает объём эмбеддингов и пик памяти.
Обе функции реализованы в однопроходном режиме: не требуется генерировать изображение, потом маскировать, потом генерировать снова. Это критично для скорости - один проход диффузии на M2 Max занимает 8-12 секунд, каждый дополнительный проход удваивает время.
Практические выводы: когда стоит портировать модели на Swift/MLX
Порт на Swift/MLX оправдан в трёх случаях:
- Память - узкое место. Python-стек не влезает в доступную unified memory, а модель физически помещается при поэтапной загрузке. Выигрыш - возможность запуска на потребительском железе без облаков.
- Продакшен в Apple-экосистеме. Нативное приложение на SwiftUI с инференсом на MLX не требует серверной части, работает офлайн и использует все оптимизации чипа (ANE, AMX-блоки).
- Контроль над пайплайном. Поэтапная загрузка, детерминированное управление памятью, отсутствие фреймворк-специфичных сюрпризов вроде неявных копирований тензоров.
Минусы: меньше готовых инструментов. Приходится руками реализовывать токенизатор, attention-маски, позиционное кодирование. Отладка ошибок вроде бага с RoPE-позициями занимает дни. Привязка к Apple Silicon означает, что порт не запустится на NVIDIA-железе или в облаке.
Сравнение по скорости: на M2 Max с 32 ГБ unified memory порт Krea 2 Turbo генерирует изображение 1024x1024 за 8-12 секунд. Python-аналог с PyTorch и MPS на той же машине либо не запускается вовсе (out-of-memory), либо требует агрессивного оффлоада на SSD и выдаёт результат за 40-60 секунд. Выигрыш - в эффективном использовании unified memory, а не в сырой производительности вычислений.
Для разработчиков, которые рассматривают локальный инференс diffusion-моделей на Mac, порт Krea 2 Turbo - прецедент. Он доказывает: модели, которые «не влезают» в 32 ГБ по спецификации, могут работать на потребительском Apple Silicon. Цена - ручная реализация компонентов и тщательное тестирование каждого слоя.