Сейчас технологии машинного обучения и нейронных сетей находят широкое применение в различных сферах, не исключая дизайн и ремонт помещений. Одной из таких технологий являются методы генеративных нейросетей, которые позволяют преобразовывать изображения, сохраняя основные элементы оригинала, но добавляя новые детали и стилистические изменения. Меня зовут Алексей Луговой, я работаю с Computer Vision в Самолете и уже обзорно рассказывал на Хабре, как мы применяем искусственный интеллект в строительстве. Сегодня же углубимся в задачу по созданию генеративной сети для создания проекта ремонта. Рассмотрим процесс использования диффузионных моделей на примере не самой обычной задачи — преобразования интерьера комнаты, начав с оригинального изображения голых стен и завершая итоговой генерацией в фирменном дизайнерском стиле.

Какая была задача

В сценарии существующего процесса в компании Самолёт клиент обращается к менеджеру по продажам с запросом на ремонт квартиры. Традиционный подход включает предложение клиенту описания работ, мудбордов и стилевых рекомендаций, после чего клиент вступает в процесс с дизайнерами, чтобы увидеть окончательное воплощение идей через определенное время. Вроде бы всё просто и понятно. Но во‑первых, это долго. Во‑вторых, не всегда команды дизайнеров свободны и процесс может растянуться на недели (!). В целом, да, всё сводится к тому, что это долго. А где много человеко‑часов, там ещё и много затрат. Мы решили, что можем сделать процесс не только быстрее, но и легче.

И разумеется, решили в контексте этой задачи обратиться к ИИ, потому что, как говорится, «всё лучше с искусственным интеллектом». Это позволит существенно сократить время и стоимость процесса. Вместо традиционного долгого и дорогостоящего процесса клиенту будет предложено воспользоваться возможностями генерации готового ремонта, который, опираясь на данные о предпочтениях клиента и визуализациях, создаст концепции и дизайн‑решения. Такой подход значительно ускорит и удешевит процесс, обеспечивая клиенту более быстрое и качественное выполнение задачи.

Как мы начали её решать и столкнулись с проблемами

Первоначальная задача заключалась в том, чтобы сделать ремонт в стиле «до и после», используя модель Stable Diffusion, предоставленную платформой Hugging Face (или любую другую open‑source). Мы искали все варианты, которые помогут нам решить задачу. Разумеется, смотрели в сторону DALL‑E и Midjourney, которые изначально хайпанули и приоткрыли ящик Пандоры для генеративных моделей. Но у Midjourney вообще нет API (а генерировать через Discord со всеми его ограничениями — это здорово для B2С, но не для B2B), а DALL‑E и иные подобные инструменты оказались недостаточно гибкими.

Поэтому выбор пал на Stable Diffusion. На момент начала разработки была только версия 1.5 (да, мы знаем про SD 2.0 и SDXL, SD 2.1, SDXL‑refiner и прочие), обученная на картинках 512×512 пикселей. И нам оказалось этого достаточно. Во‑первых это open‑source решение, а значит — гибкость использования. Во‑вторых, у SD хорошее SOTA качество, на момент разработки это одно из лучших решений. Ну и в‑третьих, у Stable Diffusion есть передовые технологии генерации (в отличие от GAN или Flow‑based моделей, которые качественно хуже подходят под эту задачу). Плюс, нам нужно было высокое качество и разнообразие. В общем, протестировали разные доступные варианты и выбрали наиболее подходящий под наши требования продукт.

Первые шаги

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

Было: 1024x767 пикселей, стало 512x512 и поехало всё, хотя сам концепт довольно симпатичный
Было: 1024x767 пикселей, стало 512x512 и поехало всё, хотя сам концепт довольно симпатичный
Код примера
from diffusers import AutoPipelineForImage2Image

CACHE_DIR = "."
DEVICE = "cuda:0"

pipe = AutoPipelineForImage2Image.from_pretrained(
    "runwayml/stable-diffusion-v1-5", 
    torch_dtype=torch.float16,
    variant="fp16",
    use_safetensors=True,
    cache_dir=CACHE_DIR
).to(DEVICE)

prompt = "New modern style of livingroom, warm color palette, detailed, 8k"
img_gen = pipe(prompt, image=img_pil.resize((512, 512)), strenght=0.7, num_inference_steps=100)
img_gen.images[0]

Ссылка на GitHub

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

Стоит отметить, что на фоне нашей работы с v1.5, на рынок вышла новая модель SD XL. Только посмотрите, как менялся Stable Diffusion, за пару лет он прошел довольно большой путь, но когда мы начинали разработку, у нас этих опций не было.

  1. Stable Diffusion v1.1 — v1.5 (CompViz репозиторий, RunwayML репозиторий). Обучение проходило на подвыборке из LAION‑Aesthetics V2 в разрешении 512×512. Модель использует подход с ?‑предсказанием.

  2. Stable Diffusion v2.0 — v2.1 (StabilityAI репозиторий). Обучение осуществлялось на подвыборке LAION-5B в разрешении 768×768. Модель применяет подход с ?‑предсказанием. Она становится больше, обучается с нуля и использует более тяжелую версию OpenCLIP вместо OpenAI CLIP.

  3. Stable Diffusion XL (StabilityAI репозиторий, arXiv:2307.01 952). Модель имеет другую архитектуру и работает в две стадии: base и refiner. Включает в себя целых 2 CLIP text encoder»а. При обучении использовались дополнительные трюки — Micro‑Conditioning и SDEdit. Обучение проходило в разрешении 1024×1024. Модель снова стала применять подход с ?‑предсказанием.

  4. Stable Diffusion 3 (блогпост, arxiv.org:2403.03 206). Сейчас модель не находится в публичном доступе (хотя 12 июня уже выложили веса для исследований), но доступна по API и показывает потрясающие результаты. Из основных изменений: заметно изменена архитектура модели (используется DiT вместо Unet и 3 text encoder»а, в процессе инференса сам text condition обрабатывается вместе с латентами параллельно, обучение базируется на концепции Rectified Flow).

Смена модели на более новую

Этот момент можно интерпретировать как пример быстрого темпа развития технологий в данной области. Несмотря на то, что мы уже активно заняты исследованием и разработкой в v1.5, мы приняли решение внимательно изучить новую модель и оценить ее потенциал. Такой подход позволил нам оставаться на переднем крае инноваций и внедрять в нашу работу самые передовые технологические решения (про SOTA рассказывал выше). Для достижения лучшего качества мы решили использовать модель Stable Diffusion XL, которая представлена здесь.

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

Было: 1024x767, стало 1024x760 (из-за необходимости кратности на 8 — особенности просчета модели)
Было: 1024x767, стало 1024x760 (из-за необходимости кратности на 8 — особенности просчета модели)
Код примера
from diffusers import AutoPipelineForImage2Image

pipe = AutoPipelineForImage2Image.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", 
    torch_dtype=torch.float16,
    variant="fp16",
    use_safetensors=True,
    cache_dir=CACHE_DIR
).to(DEVICE)

prompt = "New modern style of livingroom, warm color palette, detailed, 8k"

img_gen = pipe(prompt, image=img_pil, strength=0.7, num_inference_steps=100)
img_gen.images[0]

Ссылка на GitHub

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

Использование SDXL и ControlNet для точной генерации 

Для достижения более точного сохранения габаритов и планировки комнаты, мы решили использовать исходную модель SDXL в сочетании с ControlNet. Это архитектура для глубокого обучения, которая позволяет управлять генерацией изображений на основе различных видов входных данных, таких как позы, контуры, или текст. Она добавляет к основной сети специальные слои, которые отвечают за обработку управляющих сигналов, обеспечивая точный контроль над конечным результатом. Основная идея заключается в том, чтобы сохранить возможности генерации изображений нейросетей, добавив при этом возможность точного контроля и настройки создаваемого контента. Если интересно, на Хабре уже была интересная статья про ControlNet.

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

Процесс генерации. Если не понимаете, в чём суть, отдалитесь от экрана или прищурьтесь.
Процесс генерации. Если не понимаете, в чём суть, отдалитесь от экрана или прищурьтесь.
Код примера
from diffusers import StableDiffusionControlNetPipeline, UniPCMultistepScheduler, ControlNetModel

controlnet = ControlNetModel.from_pretrained(
    "monster-labs/control_v1p_sd15_qrcode_monster",
    torch_dtype=torch.float16,
    cache_dir=CACHE_DIR
).to(DEVICE)

pipe = StableDiffusionControlNetPipeline.from_pretrained(
	"runwayml/stable-diffusion-v1-5", 
    torch_dtype=torch.float16,
    controlnet=controlnet,
    cache_dir=CACHE_DIR
).to(DEVICE)

pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)

prompt = ["award winning photography shirts, pants, socks, linens, sheets on clothesline in the sun high definition"]
negative_prompt = ["ugly. disfigured. low quality. blurry. nsfw. trees. green leaves. leaf. vegetation. paint. oil painting"] 

img_gen = pipe(
    prompt=prompt,
    negative_prompt=negative_prompt,
    image=img_pil,
    guidance_scale=9,
    controlnet_conditioning_scale=1.2,
    height=768,    # init image size
    width=768,
    generator=torch.Generator(DEVICE).manual_seed(22),
    num_inference_steps=100
)

img_gen.images[0]

Ссылка на GitHub

Необходимо учитывать многообразие различных моделей ControlNet. Например, вот в этом решении приведены различные примеры, однако они предназначены для версии v1.5 и не подходят для наших целей, связанных с использованием SD XL, т.к. количество доступных хороших вариантов вообще пока сильно ограничено ( вот ресурсы, если хотите поискать поглубже: HugginFace, Civitai).

Вообще на платформах много вариантов, но большинство из них — различные аниме‑вариации и ничего дельного для серьёзной работы нет. А нам важно выбрать те, которые наилучшим образом подходят для нашей задачи. Плюс проверить все особенности, потестировать и убедиться во всех нюансах. В данном случае, было решено воспользоваться моделями из официальных источников, предоставленных самими Diffusers, всё же оригиналы лучше форков. Это реимплементация оригинального кода, при этом максимально совместимая со всей остальной экосистемой HuggingFace. Среди них были выделены Canny и Depth (ниже о них подробно расскажу), так как они на момент выбора лучше всего соответствовали требованиям нашей задачи. Данное решение позволило нам продолжить работу без необходимости дополнительного поиска, обеспечивая при этом высокое качество и результативность работы.

Дополнительную информацию о контролнетах для SD XL можно найти на ресурсах:

Алгоритмы Canny и Depth: в поисках идеала

Как и сказано выше, мы выбрали надстройку ControlNet, использующий алгоритм Canny для выделения контуров исходного изображения. Алгоритм Кэнни (Canny edge detection) — это метод выделения границ (контуров) на изображениях, разработанный Джоном Кэнни в 1986 году. Пример использования данного подхода описан здесь.

Исходное изображение и выделение контуров с помощью Canny
Исходное изображение и выделение контуров с помощью Canny
Результат генерации с помощью Canny
Результат генерации с помощью Canny
Код примера
from diffusers import StableDiffusionXLControlNetPipeline, ControlNetModel, AutoencoderKL
from diffusers.utils import load_image
import numpy as np
import torch
import cv2
from PIL import Image


img_canny = np.array(img_pil).copy()
img_canny = cv2.Canny(img_canny, 100, 300)
img_canny = img_canny[:, :, None]
img_canny = np.concatenate([img_canny, img_canny, img_canny], axis=2)
img_canny = Image.fromarray(img_canny)


controlnet = ControlNetModel.from_pretrained(
    "diffusers/controlnet-canny-sdxl-1.0",
    torch_dtype=torch.float16,
    cache_dir=CACHE_DIR
).to(DEVICE)

pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    controlnet=controlnet,
    torch_dtype=torch.float16,
    cache_dir=CACHE_DIR
).to(DEVICE)


prompt = ["New modern style of livingroom, warm color palette, detailed, 8k"]
negative_prompt = ["low quality, bad quality, sketches"]

img_gen = pipe(
    prompt, 
    negative_prompt=negative_prompt, 
    image=img_canny, 
    controlnet_conditioning_scale=0.6,
    guidance_scale=9,
    eta=0.0,
    generator=torch.Generator(DEVICE).manual_seed(2),
    num_inference_steps=100
)

img_gen.images[0]

Ссылка на GitHub

Однако результаты оказались неудачными. Контролирующие изображения, созданные алгоритмом Canny, подсвечивали мелкие и зачастую ненужные детали, что приводило к множеству артефактов и «галлюцинаций» модели. Вместо того, чтобы выделять основные элементы комнаты, алгоритм фокусировался на незначительных деталях, что негативно сказалось на итоговом изображении. Значит, продолжаем поиски.

Вообще для задачи определения границ оптимальным решением является использование модели MLSD (Machine Learning Line Detection System), доступна по ссылке. Существует множество примеров использования этой модели, которая по сути является детектором прямых линий. Она часто позиционируется для задач, связанных с «очерчиванием формы коробки» или других объектов, где ключевым является обнаружение прямых линий.

Пример работы MLSD
Пример работы MLSD

По сути, MLSD можно рассматривать как частный случай алгоритма Canny. Разница в том, что Canny — это аналитический алгоритм, быстро решающий задачу, а MLSD — это отдельная модель машинного обучения, что делает её вычислительно более затратной. Однако, на момент нашей разработки MLSD ещё не был доступен (или я его не нашел), да и сейчас есть только для версии 1.5 (а мы уже выбрали SD XL). Плюс этот подход требует значительной настройки для каждого изображения, что делает его мало подходящим для автоматизации. Поэтому мы продолжили исследовать другие методы и обнаружили интересные опции в официальной коллекции ControlNet от diffusers.

Пример работы карты глубины
Пример работы карты глубины

Особенно перспективным нам показалось использование карты глубины (Depth map). Как и MLSD, этот алгоритм должен подходить под нашу задачу: очерчивать квартиру, строить карту глубины перспективы. В общем, судя по описанию — это то, что нам было нужно. Значит, начинаем пробовать!

Для этой задачи мы решили использовать модель определения глубины Intel/dpt‑beit‑large-512, которая предоставляет более высокую детализацию, что идеально подходит для наших требований. Подробные сравнения моделей определения глубины можно найти здесь.

В анализе графиков проведён разрез между «качеством» и «скоростью», причем стоит отметить, что скорость не является определяющим фактором, поскольку модель функционирует в оффлайн режиме, не требуя непрерывного видеопотока или подобного. Из‑за этого было принято решение выбрать более ресурсоёмкую модель. Хотя могли бы быть рассмотрены иные интересные проекты для более глубокого исследования, например, Marigold, наше текущее решение оказалось достаточно удовлетворительным. Стоит также отметить, что в самом Diffuser имеются ссылки на указанные проекты, а также на модели глубины DPT.

Пример анализа алгоритмом Depth map
Пример анализа алгоритмом Depth map
Результат генерации
Результат генерации
Код примера
from transformers import DPTImageProcessor, DPTForDepthEstimation
from diffusers import ControlNetModel, StableDiffusionXLControlNetPipeline, AutoencoderKL


processor = DPTImageProcessor.from_pretrained("Intel/dpt-beit-large-512", cache_dir=CACHE_DIR)
model = DPTForDepthEstimation.from_pretrained("Intel/dpt-beit-large-512", cache_dir=CACHE_DIR)#.to(DEVICE)

# prepare image for the model
inputs = processor(images=img_pil, return_tensors="pt")

with torch.no_grad():
    outputs = model(**inputs)
    predicted_depth = outputs.predicted_depth

# interpolate to original size
prediction = torch.nn.functional.interpolate(
    predicted_depth.unsqueeze(1),
    size=img_pil.size[::-1],
    mode="bicubic",
    align_corners=False,
)

# visualize the prediction
output = prediction.squeeze().cpu().numpy()
formatted = (output * 255 / np.max(output)).astype("uint8")
img_depth = Image.fromarray(formatted)


controlnet = ControlNetModel.from_pretrained(
    "diffusers/controlnet-depth-sdxl-1.0",
    variant="fp16",
    use_safetensors=True,
    torch_dtype=torch.float16,
    cache_dir=CACHE_DIR
).to(DEVICE)

pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    controlnet=controlnet,
    variant="fp16",
    use_safetensors=True,
    torch_dtype=torch.float16,
    cache_dir=CACHE_DIR
).to(DEVICE)


prompt = ["New modern style of livingroom, warm color palette, detailed, 8k"]
negative_prompt = ["low quality, bad quality, sketches"]

img_gen = pipe(
    prompt=prompt,
    negative_prompt=negative_prompt,
    image=img_depth,
    guidance_scale=13,
    num_images_per_prompt=1,
    num_inference_steps=100, 
    controlnet_conditioning_scale=0.5,
    generator = torch.Generator(DEVICE).manual_seed(8)
)

img_gen.images[0]

Ссылка на GitHub

Уже красиво, но надо ещё лучше!

Предметы на переднем плане

Проблема: Модель может интерпретировать какие‑то объекты, попавшие в кадр как объекты интерьера и пытаться заместить их новыми элементами, что часто приводит к появлению случайных искажений и неопознанных вещей на изображении. В том числе это работает и с водяными знаками, если их не удалить.

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

NB: На первой стадии также важно провести фильтрацию контента, включая запрещенный. Стоит учитывать возможное наличие NSFW.

Скрытые части на карте глубины

Проблема: Карта глубины может не учитывать важные объекты, которые находятся далеко или скрыты перекрывающими элементами. Это может привести к тому, что генеративная модель игнорирует необходимые объекты, такие как окна или двери, или наоборот, акцентирует внимание на незначительных объектах, например, целлофановом пакете или пальце, случайно попавшем в кадр (что, вероятно, лучше решить на шаге 1, если пакет небольшой и он на переднем плане).

Решение: Использование различных версий моделей определения глубины и проведение индивидуальной настройки параметров генерации может помочь решить эту проблему. Регуляризация с помощью промпт‑инжиниринга также будет полезна. Для этого можно тестировать и комбинировать разные модели, чтобы найти наилучший подход.

Изменение планировки

Проблема: Модель может изменять планировку помещения, добавляя или убирая окна, двери или другие элементы. Это особенно проблематично при работе с первичными и вторичными объектами недвижимости, так как они могут иметь разные особенности планировки и отделки. Например, стены могут быть голыми и рельефными на первичных объектах, что отразится на карте глубины и приведет к нежелательным результатам.

NB: Вообще проблема поехавшей планировки была первой и наиболее существенной. Общая планировка уже претерпела значительные улучшения благодаря предшествующим усилиям. Однако возникает новая проблема: модель способна определить наличие стены (нечто продольное и похожее на стену, например, шкаф), но не может определить её материал, что приводит к появлению окон или дверей (в том числе закрытых), где на самом деле должна быть сплошная стена без таких атрибутов или встроенный гардероб. Приведенный пример с полноэкранным окном иллюстрирует эту проблему: такие архитектурные элементы могут создать впечатление другой квартиры, не соответствующей запросам клиента.

Решение: Промпт‑инжиниринг с явным указанием нужных элементов может помочь минимизировать такие проблемы. Также полезна подстройка гиперпараметров модели и использование трюков с картой глубины, таких как «размывание» (в библиотеке CV2 есть множество решений) для усреднения деталей или комбинирование с алгоритмом Canny. Это позволит сгладить рельеф и уменьшить влияние мелких деталей на итоговое изображение.

Подстраивание расстановки мебели

Проблема: При работе с первичными объектами модель хорошо наполняет помещение мебелью. Однако, при работе со вторичными объектами, где мебель уже присутствует, модель может не учитывать это и пытаться разместить мебель там, где это не соответствует исходному расположению.

Решение: Настройка гиперпараметров модели может помочь справиться с этой проблемой. Нужно тщательно подбирать параметры, чтобы модель лучше учитывала существующую мебель и правильно интегрировала новые элементы в уже обставленное пространство.

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

Универсальное решение: StableDiffusionXLControlNetImg2ImgPipeline 

Результатами Depth Map мы хоть и были довольны, но для работы с ним нужна была большая предварительная подготовка. Поэтому поиск не прекращался и спустя какое‑то время мы нашли, как нам сейчас кажется, универсальное решение — использование StableDiffusionXL ControlNetImg2ImgPipeline(да, название сложное, но это никак не влияет на работоспособность решения). Этот подход позволяет ещё немного улучшить качество генерации изображений к достаточному для нас, учитывая различные аспекты интерьера и сохраняя основные элементы планировки. То есть, он делает картинку ещё более красивой и решает почти все перечисленные выше нюансы. Подробнее об этом инструменте можно узнать по ссылке: StableDiffusionXLControlNetImg2ImgPipeline.

Пример работы через пайплайн
Пример работы через пайплайн
Код примера
from transformers import DPTImageProcessor, DPTForDepthEstimation
from diffusers import ControlNetModel, StableDiffusionXLControlNetImg2ImgPipeline


processor = DPTImageProcessor.from_pretrained("Intel/dpt-beit-large-512", cache_dir=CACHE_DIR)
model = DPTForDepthEstimation.from_pretrained("Intel/dpt-beit-large-512", cache_dir=CACHE_DIR)#.to(DEVICE)

# prepare image for the model
inputs = processor(images=img_pil, return_tensors="pt")

with torch.no_grad():
    outputs = model(**inputs)
    predicted_depth = outputs.predicted_depth

# interpolate to original size
prediction = torch.nn.functional.interpolate(
    predicted_depth.unsqueeze(1),
    size=img_pil.size[::-1],
    mode="bicubic",
    align_corners=False,
)

# visualize the prediction
output = prediction.squeeze().cpu().numpy()
formatted = (output * 255 / np.max(output)).astype("uint8")
formatted = cv2.bilateralFilter(formatted, 50, 90, 90)
img_depth = Image.fromarray(formatted)


controlnet = ControlNetModel.from_pretrained(
    'diffusers/controlnet-depth-sdxl-1.0-small',
    variant="fp16",
    use_safetensors=True,
    torch_dtype=torch.float16,
    cache_dir=CACHE_DIR
).to(DEVICE)

pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    controlnet=controlnet,
    variant="fp16",
    use_safetensors=True,
    torch_dtype=torch.float16,
    cache_dir=CACHE_DIR
).to(DEVICE)


prompt = ["New modern style of livingroom, warm color palette, detailed, 8k"]
negative_prompt = ["low quality, bad quality, sketches"]

img_gen = pipe(
    prompt=prompt,
    negative_prompt=negative_prompt,
    image=img_pil,
    control_image=img_depth,
    strength=0.9,
    eta=0.0,
    num_inference_steps=150,
    controlnet_conditioning_scale=0.4,
    guidance_scale=12,
    num_images_per_prompt=2,
    generator = torch.Generator(DEVICE).manual_seed(2)
)

img_gen.images[0]

Ссылка на GitHub

В работе с этим инструментом очень важно правильно подобрать параметры для генерации изображения, чтобы соблюсти баланс между возможностью модели фантазировать и удерживать её в заданных рамках, таких как стены квартиры, чтобы они не искажались. Пример таких параметров для SD XL:

  • Strength: Определяет, насколько далеко можно отходить от оригинала, от 0 (сам оригинал) до 1 (полностью другое изображение).

  • Eta: Уровень стохастичности (случайности) во время сэмплирования модели. Обычно значение 0 используется для большего контроля и качества.

  • Num_inference_steps: Количество шагов для генерации изображения. Чем больше шагов, тем лучше качество.

  • Controlnet_conditioning_scale: Уровень контроля в рамках контролирующего изображения ControlNet, где 0 — нет контроля, 1 — полный контроль.

  • Guidance_scale: Влияние текстовых подсказок на итоговое изображение. Чем больше значение, тем более строгое следование тексту, но увеличение значения до бесконечности может привести к галлюцинациям модели.

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

Создание фирменного стиля и использование LoRA

Также важно отметить, что мы стремимся создать ремонт не абы какой, а в фирменном стиле. В перспективе у нас в Самолёте будет несколько стилей для создания дизайна ремонта. Технически мы пробовали задавать этот набор через промпты:

Пример положительного промпта

livingroom, Modern Interior Style, Clear lines and shapes, Mixed lighting, Functionality with high-quality finish, Minimalism using modern materials, Neutral color palette with matte surfaces, Open spaces with technological features, true-to-life textures, photorealistic, hyperrealistic, highly detailed, high definition, 8k resolution

Пример негативного промпта

Enclosed, isolated spaces, Classic furniture with heavy textiles and excessive decoration, Decorative wall patterns, moldings, cornices, Traditional materials, Bright, contrasting colors and patterns, low quality, sketch, cartoon, deformed, distortions, messy, exaggerated proportions, blurry, unfocused, surreal

Результаты оказались вполне приемлемыми, но мы решили попробовать дообучение модели. Рассмотрев несколько популярных методов, таких как DreamBooth и LoRA, остановились на LoRA, поскольку он требует меньше вычислительных ресурсов GPU, быстрее обучается и показывает хорошие результаты. Мы использовали мудборд фирменного стиля от наших дизайнеров для дообучения модели, и результат превзошел ожидания.

Источники стиля для LoRa
Источники стиля для LoRa
Итоговый результат
Итоговый результат

Дальнейшие шаги

Благодаря дообучению модели, помимо общей картины (цвета, фон, палитра красок и прочее), в фирменном стиле также появляются детали интерьера, такие как обивка дивана и стульев. В общем, мы получили то, что хотели, запрашивали и надеялись увидеть. И всё это гораздо быстрее, чем нарисовал бы дизайнер. Все решение мы обернули в платформу MLflow, разработали простенький сервис с API и интегрировали всё это в чат‑бот в Telegram для сотрудников.

Вот лишь некоторые результаты за короткое время работы проекта:

  • Более 80% положительных отзывов.

  • Всего выполнено свыше 15 000 генераций.

  • Среднее время инференса составляет 1 минуту, но на тестах новой версии уже достигаем 30 секунд в зависимости от входного изображения.

Также стоит отметить, что в боте реализована и другая задача — добавление реально существующей мебели из каталога (inpainting).

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

Заключение

Как вы могли наглядно изучить, использование пайплайна StableDiffusionXLControlNetImg2ImgPipeline в комбинации с различными моделями ControlNet и настройками позволяет значительно улучшить качество генерации изображений, сохраняя планировку и основные элементы интерьера. Адаптация стиля с помощью дообучения моделей также является эффективным решением для создания уникальных и качественных результатов. При этом, важно понимать, что модели и алгоритмы совершенствуются постоянно. И есть риски, что если вы сейчас начнете делать продукт на базе SDXL, то через полгода‑год появится какой‑нибудь SDXL 2 или SDXXL, и придётся всё переделывать, т.к. он «из коробки» будет работать лучше. В любом случае, процесс разработки ComputerVision‑решений сейчас больше напоминает исследование, тестирование и даже немножко творчество. Так что пробуйте, ошибайтесь, учитесь на нашем опыте и всё у вас получится!

Комментарии (10)


  1. Roland21
    02.07.2024 14:24
    +3

    Идея интересная.
    Правда что-то пропорции и размеры мебели не вяжутся с реальностью размера комнаты.


    1. Evengard
      02.07.2024 14:24
      +1

      Вот тоже показалось будто комната стала будто-бы "больше" чем она на фото до ремонта


    1. AlexeyLugovoy Автор
      02.07.2024 14:24
      +1

      Конечно, сейчас, генерация - это, скорее, про фантазию и вдохновение: помимо размеров еще куча артефактов и других мелочей, к сожалению.

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

      Чтож, будем стараться улучшать :)


  1. piton369
    02.07.2024 14:24
    +1

    Осталось, чтобы ещё чертежи генерировала.