Сергей

Ведущий инженер-программист Иностудио

Рассказываем о том, как внедряли новомодные AI-инструменты в проект. Как это повлияло на его популярность и что за этим последовало — читайте в статье.

Технические особенности проекта:

• Фреймворк — Laravel

• БД — PostgreSQL

• Кэш/очереди — Redis

• Архитектура — основной кластер DigitalOcean Kubernetes и графический кластер GKE

О самом приложении

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

Было решено интегрировать дополнительные инструменты для работы с видеоконтентом и продолжить работу над актуальностью ранее созданных. Но после анализа состояния проекта один из векторов дальнейшего развития заключался в работе c AI-решениями. О них и пойдёт речь в этой статье.

Утилита «Удаление и замена фона фотографий»

Технические особенности:

OpenCV, Pillow, Numpy, Onnxruntime

Остановимся на особенностях реализации. За основу мы взяли инструмент remBg. Он построен с использованием целого набора библиотек, указанных выше.

Алгоритм работы.

1. Открывается изначальное изображение:

…
	input = cv2.imread(input_path)
…

И происходит приведение к объекту Pillow:

…
	if isinstance(data, PILImage):
        return_type = ReturnType.PILLOW
        img = data
    elif isinstance(data, bytes):
        return_type = ReturnType.BYTES
        img = Image.open(io.BytesIO(data))
    elif isinstance(data, np.ndarray):
        return_type = ReturnType.NDARRAY
        img = Image.fromarray(data)
…

2. Прогрузка натренированной ранее модели на основе U2net в данном случае:

…
	session = new_session("u2net")
…

3. Получение маски исходного изображения на основе прогноза:

…
    ort_outs = self.inner_session.run(
            None,
            self.normalize(
                img, (0.485, 0.456, 0.406), (0.229, 0.224, 0.225), (320, 320)
            ),
        )

        pred = ort_outs[0][:, 0, :, :]

        ma = np.max(pred)
        mi = np.min(pred)

        pred = (pred - mi) / (ma - mi)
        pred = np.squeeze(pred)

        mask = Image.fromarray((pred * 255).astype("uint8"), mode="L")
        mask = mask.resize(img.size, Image.LANCZOS)
…

4. Форматирование полученной маски (сглаживание, обработка) с помощью OpenCV:

…
	mask = morphologyEx(mask, MORPH_OPEN, kernel)
	mask = GaussianBlur(mask, (5, 5), sigmaX=2, sigmaY=2, borderType=BORDER_DEFAULT)
	mask = np.where(mask < 127, 0, 255).astype(np.uint8) # convert again to binary
…

5. Наложение маски на исходное изображение:

…
	empty = Image.new("RGBA", (img.size), 0)
	cutout = Image.composite(img, empty, mask)
…

6. Сохранение нового изображения:

…
	cv2.imwrite(output_path, output)
…

Это «в двух словах» о принципе работы утилиты.

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

def add_image_as_background(bg_image_path: str, image: Image.Image) -> Image.Image:
    pre_processing = pre_process_file(bg_image_path)
    target_image = image.convert("RGBA")

................................................
................................................

   # Ресайз, если необходимо из расчётов
   target_image = target_image.resize((w_size, h_size), Image.ANTIALIAS)

................................................
................................................
        # Совмещение
        pre_final = Image.alpha_composite(bg_img, pre_final)

        return pre_final
.................................................

def add_background_color(color, image: Image.Image):
    new_image = Image.new("RGBA", image.size, color)
    new_image.paste(image, (0, 0), image)

    return new_image

В репозитории remBg можно выкачать уже обученные модели или обучить свою. Для уже обученных используется U2net и DUTS-TR dataset. Для более подробной информации есть описание. Для наших целей пока что достаточно обученных.

Скриншот утилиты:

В дальнейшем использовали эту реализацию для Profile Photo Maker — более продвинутого инструмента в виде конструктора.

О «внутренностях» продукта

Для каждого типа утилит используется своя очередь. 

При обработке, например джобы, происходит обновление статусов в БД, вытаскивается исходный файл из digitalOcean space. После этого задача отправляется на обработку через драйвер.

В качестве основы для драйвера взяли https://github.com/alchemy-fr/BinaryDriver Он отлично себя зарекомендовал в реализации PHP-FFMPEG.

Используется symfony/process для обработки с необходимой нам обёрткой в виде класса Abstract Binary.

Основной метод формирует CLI-команду в таком виде:

public function removeBg($input, $destination, $bgInput = null, string $color = null): RmbgTranscoder
{
   try {
       $rembgArg = [
           config('services.rembg'),
           '--input_path=' . $input,
           '--output_path=' . $destination,
       ];
       if (isset($color)) {
           $rembgArg[] = '--color=' . $color;
       }
       if (isset($bgInput)) {
           $rembgArg[] = '--bg_image_path=' . $bgInput;
       }
       $this->command($rembgArg);
   } catch (ExecutionFailureException $e) {
       throw new Exception('RemBg library was unable to remove background from image (command)', $e->getCode(), $e);
   } catch (Exception $e) {
       throw new Exception('RemBg library was unable to remove background from image (command) common exception',
           $e->getCode(), $e);
   }
   return $this;
}

Всё просто.

Утилита «Генератор изображений» 

Изначально работали со Stable Diffusion версии 1.4.

На huggingface довольно неплохо описан стандартный способ интеграции. Однако на нашей сборке GKE с T4 GPU возникли проблемы производительности, которые не удовлетворяли клиента.

Что такое Stable Diffusion

Stable Diffusion — это целый набор компонентов и моделей для преобразования текста в изображение. На первом этапе происходит преобразование текста в токены с помощью нейронки CLIPText. Далее текстовые токены и массив информации изображения (latent) отправляются в скрытое пространство нейросети Unet. Там происходит процесс диффузии с необходимым количеством итераций и последующая конвертация полученной после обработки информации в картинку.

Здесь, как мне кажется, отлично описали алгоритм.

Сейчас у нас запущена обработка на более современной модели.

В качестве метода обработки входного шума остановились на семплере DDIM (Denoising Diffusion Implicit Models) в 50 шагов. Наиболее оптимальные для нужд клиента результаты и время обработки. Upscaler не используем после получения результата — отдаём в 512x512.

Вернёмся к нашей проблеме.

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

Помимо этого установили xFormers для оптимизации памяти. Библиотека оптимизирует использование видеокарт Nvidia на архитектурах Pascal, Turing, Ampere (30XX), Lovelace (40XX) или Hopper (GH100).

В целом генерация вышла на комфортные 10 секунд. Однако выгрузка модели при «холодном» запуске занимает ещё около 10 секунд. То есть это время просто теряется, несмотря на то, что мы храним модель на ssd volume без предзагрузок. Был вариант написания «демона» для того, чтобы модель была всё время в памяти. Но пообщавшись с разработчиками, пришли к выводу не инициализировать веса с помощью device_map=«auto».

Это привело к значительному ускорению при инициализации. Рекомендации по поводу этого добавили в релиз и документацию, кстати ????

Реализация

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

Страница галереи выглядит так:

Обработка схожа с обработкой утилиты удаления фона. Генерация, как я сказал выше, происходит в кластере GKE на T4 GPU в 8 потоков.

Утилита стала очень популярна, GPU работают практически без перерывов. Клиент доволен, реклама запущена.

Утилита «Удаление объектов с фотографий»

За основу взята LaMa.

Backend. Сделали по аналогии с утилитами выше — драйвер обернули в очереди. Работает также в нашем выделенном кластере на GPU.

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

Скорость обработки замечательная, никаких проблем с оптимизацией не возникло.

Как это выглядит у пользователя:

Оригинал картинки:

После удаления:

UI представляет собой набор инструментов для выделения областей удаления — кисть и зумирование. Плюс реализована история изменений — пользователь может отменять действия.

Frontend написан на VueJS. Реализация инструментария с использованием KonvaJS.

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

Клиент в восторге, пользователи тоже. Сейчас рисуется новый дизайн для этой утилиты.

Набор утилит «Текстовые генераторы-помощники»

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

Используется OpenAI (GPT-3) с различными моделями и настройками. Для более популярных утилит используем более свежие модели.

Что такое GPT-3

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

Например, для утилиты генерации сочинений по заданной теме используется модель text-davinci-003. А для утилиты, которая исправляет ошибки в переданном тексте или утилите-генераторе списка FAQ, используется модель text-curie-001, которая характеризуется OpenAI как «очень способная, но быстрее и дешевле, чем Davinci».

Помимо модели настраиваем следующие показатели:

• temperature — чем ниже «температура», тем выше вероятность, что GPT-3 выберет слова с более высокой вероятностью появления; 

• frequency_penalty — «штраф за частоту» работает за счёт снижения шансов на то, что слово будет выбрано снова, если это слово уже было использовано;

• presence_penalty — «штраф за присутствие» учитывает не то, как часто слово использовалось, а только то, существует ли оно в тексте.

OpenAI не даёт неограниченное количество запросов сразу, даже несмотря на оплату. Лимит увеличивается постепенно в зависимости от популярности нашего аккаунта. 

Также OpenAI может в любой момент перестать обрабатывать поступающие запросы «просто так», аргументируя это большой нагрузкой в данный момент.

Был случай, когда в течение всего дня случалось 20% ошибок.

Для удобства добавили историю запросов.

Утилита «Конструктор статей для блога»

Наиболее интересная утилита из этой категории — это «Конструктор статей для блога». 

Алгоритм следующий.

1. Пользователь вводит набор ключевых слов, описывающих его будущую статью.

2. Мы подбираем 30 фраз в OpenAI, которые можно использовать в качестве подзаголовков. Дополнительно мы тянем несколько наборов заголовков из гугла по аналогичному запросу через API.

3. Пользователь выбирает необходимые подзаголовки.

4. По каждому подзаголовку генерируем текст статьи.

Полных скриншотов мы не можем предоставить, так как утилита только анонсирована.

Подведём итог

В результате внедрения разработанных интеграций трафик вырос в два раза и составляет на данный момент около 3,5 млн уникальных посещений в месяц. Каждый день обрабатывается около 50 000 задач от AI-утилит. Обращаю внимание, что речь именно об перечисленных утилитах. 

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