Четыре года я занимаюсь разработкой различных спецэффектов для фото и видео в мобильных приложениях. Вроде бы это локальная и как бы несерьезная тема, но одну только плачущую маску в Snapchat посмотрели 9 млрд раз. Такие штуки пользуются бешеной популярностью и здорово повышают виральность мобильных приложений, но с каждым годом удивлять людей становится все сложнее.

В этой статье я разберу эволюцию видеоэффектов, поделюсь наблюдениями и раскрою пару инсайдов о том, как перенести стилизацию изображения из StableDiffusion на смартфоны.

Когда-то 3D-маски онлайн были прорывом

Чтобы разработать нечто новое, нужно хорошо знать историю вопроса, и понимать, какие эффекты больше привлекают пользователей прямо сейчас, а какие остались в прошлом.

История успеха белорусов из Masquerade Technologies уже подзабылась, но все еще будоражит воображение. В 2016 ребята одни из первых придумали, как накладывать маски на изображение с камеры смартфона в режиме реального времени и сорвали джекпот. Выпустили приложение MSQRD, продали технологию Цукербергу и уехали в Лондон. Тогда же взлетел и стартап Looksery, который почти сразу купил Snapchat за $300 млн.

У Masquerade был крутой опыт сфере 3D-рендеринга и графики и оригинальный подход к трекингу лиц в 3D-пространстве — универсальный, кроссплатформенный и нетребовательный к железу. А главное, это была свежая идея, которая зацепила пользователей. Не удивительно, что ей заинтересовались крупные игроки.

Сейчас подобные эффекты есть в приложении большинства соцсетей, к ним привыкли. Они еще могут влиять на продуктовые метрики, но дифференцирующим фактором для выбора приложения сами по себе уже не станут. То, что 8 лет назад было уникальной технологией, теперь — во множестве библиотек, которые позволяют любому приложению сделать все то же самое. Возможно, не в таком хорошем качестве, но просто, быстро и дешево. Теперь это есть у всех.

Классической стилизацией никого не удивишь, даешь нейросети

Следующим геймчейнджером стали нейросетевые фоторедакторы, например, Prisma. Она отличилась тем, что стилизовала изображения под произведения искусства с помощью real-time Style Transfer на нейросетях. Суть метода в обучении под каждый стиль сверточной нейронной сети с архитектурой вида энкодер-декодер и использованием функции потерь:

Схема обучения нейросети для real-time Style Transfer

Стилизация одного фото на базе этого метода занимает не более сотни миллисекунд и запускается на смартфоне. Но больше всех от внедрения нейросетей выиграл Snapchat со своими масками, которые стилизуют лицо с учетом дескрипторов. Если у тебя большой нос, то и у маски будет большой нос — сохраняется некоторая часть identity. На том, как это работает, стоит остановиться поподробнее.

Надеваем маску при помощи StyleGAN2

Наиболее известный подход к созданию современных нейросетевых масок связан с первыми моделями, которые научились хорошо генерировать лица — семейства GAN’ов — генеративно состязательных сетей или Generative Adversarial Networks.

Предыдущие методы на момент появления были хороши, но 3D-маски ограничены в применении и зачастую выглядят неестественно. Стайл-трансфер тоже имеет характерные паттерны, по которым человек может понять, что изображение искусственное. StyleGAN стал первым генеративным алгоритмом, который создавал лица, неотличимые от настоящих.

Архитектура StyleGAN и лица, сгенерированные архитектурой StyleGAN2, обученной на домене лиц. Одновременно обучаются две сети: генеративная, для создания лица, и дискриминативная, которая учится отличать результаты, сгенерированные первой моделью, от реальных картинок

Я работаю в компании ESN над видеоэффектами для социальных сетей Chips и Панч, и, конечно, мы сразу начали экспериментировать со StyleGAN2, предложенной Nvidia, и пытаться сделать маски с ее помощью. Расскажу, как это работает у нас.

Как это работает в реальном домене

StaleGAN оказался удивительно хорош в модификации атрибутов лица (Face Edit), например, пола, возраста или расы. Уже тогда StyleGAN2 Distillation, позволяла без помощи Midjorney и DALL·E 3 создавать разнообразные вариации известных персонажей.

Иллюстрация преобразования Гарри Поттера. Из статьи StyleGAN2-distilation 2020 и с помощью Midjorney 2023

Это был простой, но эффективный метод. Вначале создавалось множество изображений, точнее, множество пар вида: изображение — латентный вектор. Затем применялся классификатор для разделения на мужской и женский пол. Из образцов, выбранных на основе уверенности классификации, формировались мужские и женские латентные вектора для преобразований между мужским и женским полом. Найдя такое преобразование в латентном пространстве, можно сгенерировать синтетический датасет с парами изображений мужчина->женщина и обучить для такого преобразования простой im2im (pix2pix) генератор.

Затем мы обратили внимание на генератор StyleGAN2 NADA, в котором объединены возможности StyleGAN2 и CLIP для адаптации домена генераторов изображений. Он позволяет создавать изображения, которые соответствуют заданным текстовым описаниям. В StyleGAN-NADA применяется специальная функция потерь, которая структурирована таким образом, чтобы оптимизировать выход генератора и согласовать с заданными текстовыми описаниями через модель CLIP.

Вдохновившись, мы изобрели собственный отчасти более эффективный подход и обучили специальную модель для модификации латентного пространства из одного класса в другой. Для ее обучения брали универсальный классификатор CLIP или специализированный классификатор (например, возраста), где при использовании CLIP заданный класс просто задается текстовым промтом (например, старый человек).

В качестве входных моделей с замороженными весами использовались генератор G (StyleGAN), классификатор С (CLIP) и входных латентных векторов W, а на выходе dW — изменение латентных векторов, которое пытается сохранить максимально контент и изменить класс изображения cls=C( G(W+dW) ), на заданный (например, возраст или пол).

В результате можно получить высококачественные пары изображений: одно из них является исходным, а второе — модифицированным. На модифицированном изображении лицо изменено минимально необходимым образом, что позволяет сохранить идентичность, но при этом вносит определенное свойство (например, возраст, пол, полноту или любое другое, которое придет вам в голову).

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

LEGO-стилизация и специализированное дообучение

На практике StyleGAN2 NADA работала неплохо, но для внедрения в продукт качества не хватало. Поэтому мы взяли в качестве примера человечков LEGO и взялись вносить изменения в базовую модель StyleGAN. Для этого используется техника специализированного дообучения.

G — генератор StyleGAN. Mapping — функция преобразования из шума в латентное пространство W. Mapping можно независимо применять к реальной части (слева) и к стилевой части справа. Синим отмечены замороженные веса. Справа — адаптер. Нижняя часть «пирамиды» отвечает за форму и dAffine – за обучение специализированного аффинного слоя. Аффинный слой адаптирует латентный вектор W под конкретный слой стеки и подает его на вход ModConv. Блок с dAffine+dModConv означает, что добавочные слои с более высоком разрешении обучаются оба блока

1. Берется обученная модель StyleGAN, которая умеет генерировать реалистичные изображения лиц, и фиксируется, т.е. веса замораживаются и не участвуют в обучении. При этом модель сохраняет возможность генерировать реалистичные лица.
2. К зафиксированной модели добавляются новые слои (так называемые адаптеры), которые будут поэтапно (послойно) модифицировать реалистичные изображения до стилизованных.
3. Смешивание стилей происходит благодаря коэффициенту, который стоит перед каждым новым слоем. Таким образом можно задавать уровень стилизации.
4. На выходе получается гибридная модель, которая может генерировать пары изображений — реальное лицо и его стилизованная LEGO-версия.

Атрибутивное выравнивание

Теперь нужно было обеспечить сходство стилизованного изображения с оригиналом.

Здесь есть концептуальная проблема. Стилизация может быть настолько радикальной, что становится сложно определить критерии похожести. Например, какие признаки должны сохраняться, чтобы человек узнал себя в виде фигурки LEGO? А орка или инопланетянина? Вопрос сложный, но эксперименты показали, что важнее всего сохранить ряд характерных атрибутов лица: очки, усы и бороду, форму лица и цвет волос. Тогда шанс на узнавание будет выше, даже при самых радикальных трансформациях

Существуют специализированные нейронные сети, которые находят такие атрибуты, но беда в том, что никто не догадался обучить их на человечках из детского конструктора. А вот zero-shot CLIP универсальный и справляется с задачей поиска атрибутов как на реальных лицах, так и на лицах в других доменах.

Интеграция сохранения атрибутов была внедрена в процесс обучения StyleGAN. Это привело к ухудшению FID — ключевой метрики для GAN, оценивающей, насколько распределение сгенерированных изображений отличается от исходного. Однако, на практике, при использовании небольшого коэффициента для функции потерь на атрибутах, получаемые пары изображений оказываются более качественными по сравнению с теми, что получаются без учета атрибутов.

Применяем атрибуты при построении пар

Мы добились того, что наш модный StyleGAN одновременно генерировал довольно похожие пары реальное/модифицированное лицо. На следующем шаге мы сгенерировали несколько тысяч пар лиц и посчитали для них атрибуты в обоих доменах с помощью CLIP. Этот набор стал нашим кастомным словарем атрибутов.

Теперь при генерации нового изображения для реального домена мы могли вычислить его атрибуты, найти ближайшие K векторов aW в словаре для модифицированных изображений (атрибутов) и смещать латентный вектор W+aW в направлении aW, усредненный латентный вектор для К, ближайший латентный вектор в целевом домене из словаря.

Этот метод позволил сохранять идентичность лица с большей точностью.

Доделываем все остальное

Получается очень даже неплохо, но с помощью такой StyleGAN-сетки сложно сгенерировать что-то кроме лица человека или персонажа. Для генерации фона приходилось обучать и использовать другие сетки, попроще.

Чтобы получить определенный стиль, нужно найти изображения, вырезать лица, отфильтровать сеточками и на всякий случай закинуть для верификации данных в Толоку. Хорошо хоть, что это просто функция в скрипте.

Только после этого можно сгенерировать датасет, чтобы обучить на нем специализированную легковесную нейросеть — видеостилизатор. Да, мы выстроили автоматический пайплайн, но процесс разработки новых масок все равно остается сложным.

Пришествие диффузии и новые эксперименты

Недавно появилась новая впечатляющая технология — диффузионные модели. В качестве основного инструмента для многих стартапов здесь выступает полностью открытая StableDiffusion. Она дает возможность генерировать изображения с персонажами в заранее заданной позе, и зачастую нужный стиль уже есть в пользовательских моделях на huggingface или civitai. А если хочется сделать что-то более редкое — модель тюнится с помощью LoRA.

Сложно пройти мимо, и мы снова озадачились получением пар изображений. Логично было использовать метод img2img, который преобразует исходное изображение в нужный стиль, но оказалось, что это реализовано костыльно, так, будто бы мы продолжаем исходное изображение с некоторого этапа диффузионного процесса.

Допустим, наш диффузионный процесс состоит из 30 итераций. Вместо генерации всех 30-ти, берем изображение, накладываем на него шум, и подставляем в пайплайн, например, на место 10-й итерации и продолжаем процесс диффузии еще 20 итераций. Иногда это работает, но StableDiffusion становится не таким уж и stable

Что не так с этим подходом? С технической точки зрения хочется, чтобы алгоритм понимал семантику изображения и формировал из исходника сжатую карту признаков для воссоздания необходимого контента. Но как показывает практика — нельзя просто взять и превратить одно изображение в другое.

Применяем ControlNet

Чтобы упорядочить процесс генерации, можно использовать входное изображение-условие (контуры, глубина, поза и т. д.), которое затем используется в управлении процессом генерации изображений. Этот подход добавляет существующей диффузионной модели способность использовать дополнительные данные. Он реализован в рамках ControlNet.

Есть исходное изображение → к нему применяется простая операция извлечения контуров (или нейросетевой алгоритм построения глубины по этому изображению) → получается парное «изображение», содержащее специфическую часть исходника, и это накладывает ограничение на генерацию изображения

В процессе тренировки ControlNet научили восстанавливать целое изображение из контурного и других. Для каждого типа препроцессинга (контуры, позы, глубина и т. д.) существует метод этого препроцессинга и отдельная нейросеть ControlNet.

Архитектура ControlNet

Архитектурно ControllNet — дополнительная нейронная сеть, которая встраивается в архитектуру диффузионной модели, U-Net в Stable Diffusion. Она берет на вход вектор подсказки, представляющий желаемые свойства генерируемого изображения, обрабатывает его и выдает дополнительные карты признаков, которые подаются на вход декодера диффузионной модели вместе с шумом и картами признаков из ее энкодера.

Так, ControllNet динамически модулирует работу декодера в процессе генерации изображения, направляя его к более соответствующему результату согласно подсказке. Любопытно, что ControlNet работает именно с картинками — поза человека — это не набор признаков или карт признаков, а просто картинка с цветным «скелетом».

Дополняем ControlNet

Еще один инструмент для упорядочения выдачи StableDiffusion называется InstructPixPix. Здесь случайные изображения сначала генерируются моделью Stable Diffusion. Потом к ним применяются преобразования на базе текстовых инструкций, и создаются тренировочные пары до-после. И текст в них немного человечнее, чем в стандартных промтах для StableDiffusion. На этих данных обучаются дополнительные нейронные сети, которые затем комбинируются с исходной диффузионной моделью.

Важная особенность InstructPixPix — использование механизма внимания (attention), подсетей, которые анализируют текст инструкции и выделяют релевантные для изображения фрагменты. Можно точнее определять, какие изменения применить к конкретным областям изображения в соответствии с инструкцией. Использование attention улучшает результаты: модель лучше фокусируется на значимых деталях при выполнении инструкций.

Изначально InstructPix2Pix обучалась как img2img. Это ровно та задача про построение пар (исходного и стилизованного изображений), которую мы должны решить для создания маски и преобразования внешнего вида пользователя или всей сцены. И важный аспект, что модель прямо обучалась накладывать стили. В отрыве от ControlNet получается создавать довольно сильно стилизованные изображения, при этом наследуя стабильность UnStable Diffusion.

Смешиваем инструменты

Если забросить InstructPix2Pix в котел за компанию с разными ControlNetами и StableDiffusion, а затем помешать, получается нечто интересное. Мы пробовали разные пайплайны, но самым удачным оказалось сочетание ControlNet c мягкими контурами, SOTA, DWPose, которая сохраняет нам пальцы, Mediape для выражения лица и InstructPix2Pix для усиления стиля. В текстовые запросы добавляем атрибуты лица: пол, цвет волос, наличие усов и бороды — это мы уже научились делать с помощью CLIP в середине статьи.

Коктейль готов! Получаются стабильные пары, на которых можно обучить видеостилизатор.

Диффузия на мобилках или цифровая дистилляция

Пока без видеостилизатора никак. StableDiffusion вполне справляется с генерацией видео, но делает это медленно, и стилизовать видеопоток в реальном времени на среднестатистическом смартфоне сложновато.

Чтобы в реальном времени преобразовывать видео с камеры телефона в нужный стиль — тот же LEGO, cyberpunk или еще что-нибудь модное, используется техника под названием дистилляции знаний. Это когда порождающая сеть (учитель) создает необходимый стиль в нужном качестве, а маленькая сеть (студент) послушно выучивает это соответствие. Такое возможно за счет сильного сужения пространства, с которым работает сеть. Все знания диффузии или StyleGan уже не нужны: стиль зафиксирован и контекст ограничен (вертикальные видео).

Но есть нюансы. Важно, чтобы видео было выровнено попиксельно, выход сетки был детерминированным (одному входу соответствует один выход) и, кроме изображения, хорошо бы кормить «студента» промежуточными данными. Эти условия соблюдаются, если сначала обучить сеть img2img/pix2pix с той же архитектурой, что и «студента», но размером побольше. Сети-«учитель» подобны, различается лишь так называемая ширина каждого слоя, грубо говоря, количество нейронов в слоях. Используя технику послойной передачи знаний, можно эффективно обучать легковесные сетки для стилизации видео.

Реализация нейросетевых масок в наших приложениях вдохновлена StyleGAN, ведь это одна из лучших GAN-архитектур ( архитектуры на базе StyleGAN до сих пор пытаются состязаться с дифузионными нейросетями, например StyleGAN-T), но это не совсем обычный StyleGAN, а нейросеть «гармошкой» (image to image):

Такая вот архитектура энкодер-декодер с «мостом». Энкодер кодирует изображение до определенного размера, допустим, сокращая исходное разрешение в 8 раз. В результате его работы получается карта признаков размером 64х64х256 (при разрешении 512х512 и заданном числе нейронов — 256). «Мост» перерабатывает признаки несколькими блоками ResNet-подобной сети, а энкодер восстанавливает исходное разрешение изображения.

Посмотрев внимательно статью и код StyleGAN, мы многое поняли.

Правильно применили методы инициализации, трюк с размытием, а еще корректно адаптировали функции потерь и методы регуляризации из StyleGAN2. И получили подход, который доставляет стиль в мобильный телефон. Разумеется, ценой бессонных ночей программирования мобильных устройств и шишек от iOS и Android. От Android больше, но об этом как-нибудь в следующий раз.

В IT даже в достаточно узких нишах каждый год происходит нечто новое и захватывающее. Работа позволила мне взглянуть на прогресс в области создания видеоэффектов для приложений на смартфонах.

Мы начинали с StуleGAN2 для редактирования лиц, но потом переключились на диффузионные модели для более крутых штук. И что нам это дало? Теперь алгоритмы стилизации крутятся на мобилках в реальном времени! Конечно, путь был не из легких: много кода, прочитанных статей, кофе и сражений с iOS и Android. Но оно того стоило. В этой области заложен огромный потенциал, который нам еще предстоит предстоит исследовать.