Введение

Салют, Хабр! На связи Игорь Пасечник — технический лид направления XR RnD SberDevices. Сегодня я хочу рассказать про одно из наших направлений исследований — разработку генеративных моделей для 3D-контента. 

Современные методы генерации 2D-контента, такие, как 2D-диффузионные модели (Kandinsky 3.0, SDXL), уже достигли впечатляющих результатов и несколько лет являются неотъемлемой частью современности, генеративные видео модели также активно развиваются. Кульминацией развития таких подходов, вероятно, станет представленная не так давно модель Sora. Тем не менее большинство из этих моделей до сих пор испытывают проблемы при генерации консистентных 3D-сцен и объектов.

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

Учитывая это, мы решили свести свои исследования в области генеративного 3D к решению нескольких задач:

• Научиться генерировать консистентные multi-view изображения объектов по тексту или картинкам, используя 2D-генеративные модели. 

• Научиться генерировать 3D-ассеты для практического использования по тексту или картинкам.

• Научиться делать быструю и достаточно качественную 3D-реконструкцию.

• Научиться создавать консистентные визуализации объектов (3D-гифки).

Мы разработали семейство генеративных 3D-моделей на основе 2D-моделей Сбера, проведя работу с данными, выбрав лучшие компоненты из доступных на текущий момент архитектур, а также внеся дополнительные новшества, достигнув SOTA-результатов относительно open-source и проприетарных моделей.

Общая архитектура пайплайна (Text-to-3D)

На текущий момент существует несколько методов быстрой генерации 3D-контента по тексту:

1. С помощью обучения single-view реконструктора можно получать 3D-объект из любой модели генерации картинок (Kandinsky, Stable Diffusion и прочие). Данный реконструктор может быть feed-forward-моделью как в LRM без вариативности в генерации новых ракурсов, либо через диффузию, как DMV3D с возможностью обусловливания на текст напрямую, либо с обусловливанием на изображение. 

2. По аналогии можно обучить реконструктор sparse-view и использовать диффузионную модель генерации картинок, дообученную на генерацию sparse-view, для получения 3D-объектов. Свой sparse-view генератор изображений представили в ImageDream, MVDream, Instant3D, SPAD, CRM, Wonder3D. В реконструкторе sparse-view модели были обучены с разными рендерингами: volumetric rendering используется в трансформерной модели Instant3D, PF-LRM, гауссовский рендеринг используется в GRM, LGM, mesh-based рендеринг flexicubes используется в CRM.

3. В качестве генераторов dense-view используют диффузионные модели для видео, обученные генерировать облеты объектов в работах V3D, VFusion3D, SV3D. Для данных dense-view генераций уже необязательно использовать какие-то дополнительные реконструкционные модели.

4. Также долгое время в качестве текст-3D генераторов использовались модели на основе дистилляции текстового промпта через SDS loss. Данные методы работали по-объектно и занимали большое время на генерацию. Недавно было представлено несколько моделей на основе дистилляции ATT3D, AToM и LATTE3D, обученные на множестве промптов одновременно без использования 3D-данных в качестве референса.

5. Также есть подходы, которые учат текст-3D модели напрямую, без использования картиночных или видео генераторов. Один из таких подходов GVGEN учится на подсете из 36 тыс. объектов обжаверса, использует гауссовый рендеринг, для которого как раз и нужно иметь 3D-объект, а не его рендеринг.

Изучив все возможные варианты, мы пришли к выводу, что для нашей задачи оптимальным является использование двухстадийного подхода: 

• Модель первого этапа: дообученная 2D-диффузионная модель (SDXL/Kandinsky 3.0) для генерации консистентных multi-view картинок. Это позволяет сохранить широкий спектр генерируемых образов и визуальное качество изображений, унаследованных от базовой модели.

• Модель второго этапа: обучение универсальных моделей реконструкторов для генерации конечного 3D-представления, которые могут работать как на сгенерированных картинках из первого этапа, так и для реконструкции объектов по реальным фотографиям.

В рамках такого подхода можно выделить четыре относительно независимых компонента:

• модель первого этапа;

• реконструирующий генератор; 

• латентное 3D-представление;

• алгоритм дифференцируемого рендерера. 

При создании 3D-объектов мы хотели дать модели всю необходимую информацию о сцене, исключив неопределенности в аспекте геометрии и текстуры. Для получения достаточной информации о генерируемом объекте требуется иметь несколько фотографий сцены с нескольких точек обзора с оптимальными параметрами расположения камер для перекрытия областей обзора. Мы пришли к выводу, что в контексте наших задач не имеет смысла использовать более четырёх ортогональных точек зрения на объект, поэтому мы дообучаем 2D-диффузионную модель для генерации четырех ортогональных ракурсов объекта. Это количество точек обзора является достаточным для реконструкции сцены и сохранения детализации текстур, текущее разрешение синтеза лучших моделей (1024 × 1024) позволяет разделить его на сетку 2 × 2, обеспечивая всё ещё приемлемое качество каждого отдельного вида объекта (512 × 512). Вторым этапом модели является реконструктор, который, используя четыре поступивших на вход изображения, создаёт промежуточное представление сцены в виде триплейнов — ортогональных латентных плоскостей. Наиболее ресурсоёмкой операцией при таком подходе является диффузионная модель. Применив дистилляцию к зафайнтюненой модели первого этапа, мы можем сократить общее время созданий 3D-ассета до 4 секунд.

Рассматривая данный подход, финальное качество работы модели будет определяться следующими факторам:

• Качеством модели первого этапа, уровнем её детализации; 

• Качеством работы реконструктора;

• Разрешающей способностью 3D-представления;

• Разрешающей способностью и экспрессивностью рендерера.

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

Рассмотрим поподробнее все этапы обучения, а также то, что нам показалось важным.

Первый этап. Text to consistent multi-view

Данные 

Так как размеры 3D-датасетов (GSO, ABO, Shapenet, Objaverse, OmniObject3D) не так велики, как используемые для обучения 2D-диффузионных моделей (LAION), то надо попытаться выжать из них всё. В первую очередь для подготовки данных к этому этапу обучения необходимо обработать с помощью пайплайна рендеринга используемые объекты с четырёх ортогональных направлений. Далее для каждой получившейся картинки получить текстовое описание через модель-разметчик. Оптимальным вариантом нам показалось использовать LLaVA-1.5 и синтезировать информацию из нескольких описаний в единое лаконичное описание сцены. Для более удобного отбора данных мы рассчитали значения CLIP и оценки эстетичности от предсказателя эстетики (обученного на классификаторе LAION-aesthetics). Также мы составили небольшую выборку 3D-моделей, разделили их на две категории «красивые» и «некрасивые», объединили свойства CLIP и обучили небольшой SVM-классификатор. Учитывая, что мы используем на данном этапе синтетические данные, ключевой особенностью, которая редко рассматривается в академической литературе, являются параметры рендеринга освещения материалов. Регулируя эти параметры, мы можем влиять на качество и особенности финальной модели, а также эстетического восприятия пользователем генерируемых объектов.

Обучение

Процесс дообучения модели аналогичен обучению основной модели: на вход модели подаются пары «текст-изображение», представленные в виде сетки из четырёх рендеров объекта в определённом порядке. Эти данные проходят через текстовый и визуальный энкодеры, результатом чего становятся текстовый эмбеддинг и эмбеддинг изображения. Затем представление изображения зашумляется в соответствии с конфигурацией расписания диффузии и передаётся на вход стандартной модели типа U-Net с условным входом в виде текстового эмбеддинга. На выходе получается прогноз модели об уровне зашумления латентного представления.

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

Благодаря тому, что мы переиспользовали диффузионную модель и дообучали её на небольшом датасете, она не растеряла способности генерировать объекты не из обучающей выборки, однако выучила новую способность — теперь она умеет видеть их с новой стороны. В процессе обучения происходит несколько процессов — смещение модели в сторону стиля, заданного синтетическим данными, деградация исходного разнообразия модели, снижение степени соответствия генерируемых изображений промпту, а также увеличение степени консистентности четырех проекций объектов между собой. При длительной оптимизации  модель сходится к максимальным 3D консистеным объектам, но при этом падает вариативность геометрии и текстур, поэтому локальный минимум с наилучшим балансом всех факторов достаточно узок и сильно зависит от гиперпараметров, а также оптимального набора данных. Для валидации и выбора оптимальной итерации можно использовать ошибку реконструкции модели второго этапа а также CLIP score.

Механизмы контроля генерации

Несмотря на успешность дообучения модели, часто требуется дополнительное управление моделью, чтобы гарантировать, что на итоговом изображении четыре синтезированных рендера будут разделены и каждый будет соответствовать проекции обособленного объекта, а не сцены с явно выраженным диффузионным фоном. Для этого приходится прибегать к различным хитростям. Например, можно использовать изображение, состоящее из 4 равномерно распределенных 2D-гауссиан для формирования стартовой точки при генерации итоговых изображений. После создания такого исходного изображения необходимо получить его латентное представление и добавить шум в соответствии с расписанием диффузионного процесса. Этим же методом можно контролировать размер объекта.

Второй этап. Реконструтор

Общий процесс обучения

В реконструируемом этапе общая идея модели, которую мы используем, выглядит следующим образом:

1. На вход мы принимаем 4 картинки объекта с разных ракурсов и с помощью реконструируемого генератора получаем из этих картинок triplane-представление, 3 feature-плоскости, на которые проецируется объект.

2. Затем мы семплируем точки в 3D-пространстве в зависимости от используемой модели рендеринга, получаем соответствующие этим точкам features из триплейнов и декодируем из них с помощью MLP в sdf/opacity/tex_features, которые потом будут использоваться в рендеринге.

3. Рендерим объект из полученного триплейна с четырёх случайных ракурсов, взятых из датасета, оптимизируем с лоссом LPIPS + mse между отрендеренными изображениями и изображениями из датасета, соответствующим этим ракурсам.

В рамках такой схемы мы рассмотрели два вида архитектур генераторов и несколько архитектур рендерера. 

Генератор триплейна на основе трансформера (instant3d)

Одной из первых работ, использующих такой подход к реконструкции, является Instant3d, она же на текущий момент демонстрирует одни из лучших результатов по синтетическим метрикам.

Однако базовая архитектура имеет ряд недостатков, не указанных в статье, а именно — ограниченную детализацию текстуры, блочные артефакты для высокочастотных регионов, низкое разрешение выходного триплейна. Модифицировав исходную архитектуру, включая расширение внутреннего представления генератора до 64 × 64, использовав DINOv2 энкодеры с регистрами, а также более сложный upsampler, мы смогли получить лучшие результаты и избавиться от артефактов, однако общая архитектура модели вносит ограничения в масштабируемость такого подхода, потому что увеличение количества параметров генератора не влияет в явном виде на детализацию текстур на объекте, а лишь на способность генератора фьюзить входные данные в целевое представление объекта.

Генератор на основе U-Net

Альтернативным вариантом для реконструкционного генератора является использование генератора на основе U-Net для предсказания триплейнов. Один из таких вариантов в частности используется в ряде подходов, включая CRM. Наилучшие результаты в рамках SBS в задаче text-to-3D из четырёх ортогональных картинок, полученных в результате применения первого этапа, получил именно такой подход. Четыре картинки каноническим образом могут быть отождествлены с рассматриваемой триплейновой репрезентацией. На наш взгляд предсказание более чем четырёх view на первом этапе является избыточным, поэтому для горизонтального триплейна используется обучаемый позиционный эмбеддинг. Тем не менее, на картинке видно, что модели удается выучить хорошую репрезентацию для этого триплейна. Основное преимущество данного генератора в возможности генерировать триплейны в разрешении равном входным картинкам (512x512), что позволяет повысить качество деталей при рендеринге, в сравнении с instant3d-генератором, максимальное разрешение триплейнов для которого не превышает 128 без использования апсемлера. Число каналов в используемом U-Net для каждого разрешения: [128, 256, 256, 512, 512, 1024], начиная с 512 и заканчивая 16, по 2 ResNet блока для каждого разрешения, что в результате дает около 300 млн обучаемых параметров. Также, начиная с разрешения 64, используется self-attention. 

Рендеринг

На текущий момент существует множество подходов к дифференцируемому нейронному рендерингу: стандартный волюметрик, NeuS, gaussian splatting, DMTet, FlexiCubes и множество других. Проведя анализ и эксперименты, мы выбрали два рендерера для нашего пайплайна и двухэтапную схему обучения. На первом этапе мы обучаемся, используя волюметрик рендерер со стандартной экспоненциальной функцией активации — он хорошо изучен, стабильно оптимизируется и достигает высоких значений PSNR. На втором этапе мы дообучаем реконструктор в большем разрешении с рендером на основе DMTet и растеризацией, адаптируя модель к surface-based-геометрии и целевому бюджету по количеству вершин в меше, что позволяет нам впоследствии получать лучшую геометрию при экспорте 3D-ассетов. Перспективным также является использование gaussian splatting, который показывает хорошие результаты и метрики на задачах dense реконструкции, однако исходя из прошлых экспериментов мы пришли к выводу, что наивная адаптация такого подхода к генеративному сетапу требует отказа от нелинейных оптимизаций вроде клонирования и сплита гауссиан, что нивелирует преимущества метода. Также качество, демонстрируемое на задаче реконструкции, вероятно, не может быть обеспечено с текущими архитектурам и бюджету по параметрам генератора. 

Для волюметрик-рендерера мы используем достаточно стандартные параметры по аналогии с Instant3d и других подходов, а вот относительно DMTet стоит сказать пару слов более подробно.

DMTet рендеринг

Альтернативным вариантом дифференцируемого рендеринга является DMTet. Алгоритм DMTet состоит из деформируемой тетраэдральной сетки, содержащей функцию расстояния до объекта, signed distance function (sdf), которая в нашем случае аппроксимируется из триплейна на выходе генератора; и дифференцируемого алгоритма marching tetrahedra, который преобразует неявное представление sdf в явное представление поверхностной сетки. У DMTet существуют следующие преимущества перед volumetric рендерингом:

1. В процессе работы алгоритма получается явное представление 3D-модели в виде меша (набор вершин и граней). Такое представление можно без постпроцессинга использовать в 3D-сценах.

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

3. Отделение текстурной компоненты также открывает возможность разделения текстуры на параметры, такие как base color/albedo, specular, metallic (пример nvdiffrec). То есть, DMTet может выдавать текстуру для PhysicallyBased Rendering (PBR).

Переход от volumetric к DMTet

DMTet часто используется для получения мешей и текстур высокого качества в задачах реконструкции единичных объектов, например, nvdiffrec, Fantasia3d или отдельных категорий больших датасетов, например, GET3D. Особенность нашего эксперимента заключается в вариативности данных. Простая замена volumetric рендера на DMTet в нашей реконструкционной модели приводит к ухудшению сходимости модели, вплоть до обучения нулевых функций расстояния sdf. В нашем реконструкторе DMTet выступает финальным этапом тюнинга генератора, позволяющим улучшить геометрию, получить разложение текстуры на компоненты для PBR.

Технические особенности обучения DMTet

DMTet-алгоритм в нашем реконструкторе используется для улучшения качества меша и текстур в большем разрешении, поэтому мы используем достаточно большую тетраэдральную сетку 256 × 256. Триплейн на выходе генератора разделяется на два, а именно геометрический и текстурный. Их размеры тоже увеличиваются по сравнению с volumetric-реконструкцией до 256 для геометрического и 512 для текстурного. На текущем этапе мы не используем алгоритм полного разделения текстуры на параметры PBR, как в, например, Fantasia3D. Корректный PBR-рендеринг требует от модели предсказания не только текстурных свойств объекта, но и генерации окружения модели. Такой подход увеличивает размер модели и замедляет сходимость. В качестве альтернативы полного PBR-разложения мы используем предобученную модель света ENVIDR. Выучив, однако, геометрическую часть DMTet, текстурную можно заменить на корректное PBR-разложение при пообъектном тюнинге, где окружение модели статично и не требует дополнительного генератора. 

Обучение реконструирующих моделей

Данные

Для обучения мы используем объединение датасетов Objaverse 1.0 и OmniObject3D, а также Objaverse XL и MVimgNet. Каждый объект из OmniObject3D рендерился со 100 случайных камер на сфере радиуса 4, в качестве обучения использовались все объекты, без дополнительной фильтрации с нашей стороны. В датасете Objaverse, Objaverse XL много объектов низкого качества: не текстурированные, случайные расстановки геометрических объектов, плоскости, что потребовало дополнительной фильтрации датасета. Для этого мы посчитали aesthetic score четырёх противоположных рендерингов для каждого объекта и отфильтровали те объекты, что в среднем дают оценку меньше 3,5. Наилучшие результаты показало обучение в 2 этапа. На первом этапе модель обучается на большой базе с низким порогом отсечения по качеству, на втором этапе — используя только 30% самых лучших данных.

Ресурсы и скорость

Мы провели стандартную инженерную оптимизацию, используемую при обучении больших моделей и неоднократно описанную в рекомендация pytorch и гайдах от крупных компаний. Однако в нашем случае был найден отдельный момент, заслуживающий внимания. Одним из самых узких мест в моделях, использующих сэмплирование из триплейна, является backward стандартной функции grid_sample. Мы написали модифицированную версию torch.nn.functional.grid_sample, которая считает градиенты только по триплейну, а также добавили несколько оптимизаций для ускорения интерполяции градиента в соседних элементах триплейна. Это позволило ускорить backward в 2 раза. 

В сумме все оптимизации позволили получить качество, соответствующее базовой модели instant3D, используя 32 карты A100 за 5 дней против 128 карт на 7 дней, как в базовом эксперименте.

Результаты реконструкции на данных из GSO

Оценка реконструируемой модели

Изучив текущие методы валидации подобных моделей, мы пришли к выводу, что наиболее распространенным вариантом является датасет GSO (Google Scanned Objects), состоящий из сканированных реальных объектов, однако этот датасет не имеет референсной методики измерения, и в разных работах используются разные опорные кадры и метод рендеринга референсных картинок, что не позволяет достоверно сравнивать подходы. Мы использовали параметры валидации из Instant3D — входные изображения брались с фиксированных позиций на elevation 20 и азимутом 0, 90, 180, 270. Метрики измерялись на разрешении 128. 

Метрики в разрешении 512.

Оценка и SBS

При выборе стратегии валидации для модели мы применяли как синтетические метрики, так и результаты side-by-side тестирования. Синтетические тесты и сравнение метрик позволяют оценить этапы, где можно достичь некоторого численного идеала, например, при реконструкции. Оценка визуальной привлекательности и консистентности генеративной модели с помощью набора чисел более сложна, и для этого мы проводили субъективные исследования с участием групп разметки. 

Мы использовали несколько банков текстовых описаний: в качестве базового использовали DF411 из статьи DreamFusion, который включает 415 сложных композиционных промптов (например: «Кенгуру сидит на лавочке и играет на аккордеоне»), часто используемых для валидации методов преобразования текста в 3D-модели. Однако она ограничена и не включает множество практических кейсов, в связи с этим мы составили собственный набор текстовых промптов, собранных из различных публичных источников, включая ручной сбор неудачных примеров, отобранные по разным критериями промпты 2D-бенчмарка JourneyDB, тестовую базу gpt3eval. 

Для оценки моделей и различных изменений мы предоставляли разметчикам два варианта видео виртуального облета объекта, сгенерированных двумя моделями. Затем они должны были ответить на несколько простых вопросов: «Какой объект вам больше нравится?» и «Какое видео лучше соответствует текстовому описанию?». Для более детальных UX-исследований им требовалось объяснить свое решение, а также указать, что им нравится или не нравится в отдельных видео.

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

В качестве бейзлайна мы выбрали одну из наиболее современных проприетарных моделей от Luma AI (Genie).

Примеры генераций

Использование предобученных 2D-диффузионных моделей позволяет использовать все стилевые модификаторы, хорошо изученные для 2D-кейсов, мы сохраняем все разнообразие исходной 2D-модели.

Генерации текущего поколения моделей дают достаточно интересные результаты, хотя в среднем худшее соответствие промпту, чем 2D-модель, и меньший success rate, учитывая необходимость генерации консистентных объектов.

Заключение

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

В текущий момент мы прорабатываем дополнительный функционал, включая генерацию PBR-текстур, hi-res-ассетов, оптимальной топологии меша, а также повышения качества и детализации моделей.

В ближайшее время мы собираемся предоставить доступ к подробному техническому репорту, метрикам, валидационным корзинам. А сейчас наши модели доступны в тестовом доступе на developers.sber.ru. Следите за новостями!

• Коллектив авторов: Михаил Мазуров, Наталья Соболева, Дарья Мусаткина, Фанзиль Набияров, Данил Ельцов, Марина Ярославцева, Игорь Пасечник, а также наша продуктовая команда Яна Чаруйская и Кирилл Субботин.

• Также рассказ от Михаила Мазурова с конференции GIGA R&D DAY.

По всем вопросам и предложениям, а также в поисках сотрудничества просьба связываться в телеграмме @ipasechnik@charuyskaya