Запустите собственный сервер DALL-E Mini (Craiyon) на EC2

Если последние несколько месяцев вы были под камнем, DALL-E — это ML-модель от OpenAI, которая генерирует изображения на основе текстовых подсказок. DALL-E Mini (переименованная в Craiyon) Бориса Дайма и др. является менее мощной, но открытой версией DALL-E, а на сайте craiyon.com размещена версия, которую может попробовать каждый.

Однако если вы похожи на нас, вам захочется покопаться в коде и запустить модель самостоятельно. В этой статье мы сделаем это с помощью Meadowrun, библиотеки с открытым исходным кодом, которая позволяет легко запускать Python-код в облаке. В частности, для ML-моделей в недавнем выпуске мы добавили функцию запроса GPU-машин. Мы также будем подавать изображения, сгенерированные DALL-E Mini, в дополнительные модели обработки изображений (GLID-3-xl и SwinIR), чтобы улучшить качество генерируемых изображений. Попутно мы разберемся с проблемами, возникающими при запуске ML-моделей с открытым исходным кодом на EC2.

Запуск dalle-playground

В первой половине этой статьи мы покажем, как запустить saharmor/dalle-playground, который оборачивает код DALL-E Mini в HTTP API и предоставляет простую веб-страницу для генерации изображений через этот API.

dalle-playground предоставляет блокнот Jupyter, который можно запустить в Google Colab. Однако если вы собираетесь делать что-то большее, чем просто пробовать, вы столкнетесь с лимитом динамического использования в бесплатном уровне Colab. Вы можете перейти на Colab Pro ($9,99/месяц) или Colab Pro+ ($49,99/месяц), но мы получим эту функциональность за копейки, используя AWS напрямую!

Предварительные условия

Во-первых, вам понадобится учетная запись AWS. Если вы никогда раньше не использовали экземпляры GPU в AWS, вам, вероятно, придется увеличить квоты. Учетные записи AWS имеют квоты в каждом регионе, которые ограничивают количество CPU определенного типа экземпляров, которые вы можете запустить одновременно. Для экземпляров GPU существует 4 квоты:

• L-3819A6DF: «Все запросы на точечные экземпляры G и VT».
• L-7212CCBC: «Все запросы на точечные экземпляры P»
• L-DB2E81BA: «Работающие по требованию экземпляры G и VT».
• L-417A185B: «Запуск P экземпляров по требованию».

Все эти параметры установлены в 0 для новой учетной записи EC2, поэтому если вы попытаетесь выполнить приведенный ниже код, вы получите это сообщение от Meadowrun:

Unable to launch new g4dn.xlarge spot instances due to the L-3819A6DF
quota which is set to 0. This means you cannot have more than 0 CPUs
across all of your spot instances from the g, vt instance families.
This quota is currently met. Run `aws service-quotas
request-service-quota-increase --service-code ec2 --quota-code
L-3819A6DF --desired-value X` to set the quota to X, where X is
larger than the current quota. (Note that terminated instances
sometimes count against this limit:
https://stackoverflow.com/a/54538652/908704 Also, quota increases are
not granted immediately.)

Мы рекомендуем выполнить команду в этом сообщении или перейти по одной из ссылок в списке выше, чтобы запросить увеличение квоты, если вы решили попробовать (если вы используете ссылку, убедитесь, что вы находитесь в том же регионе, что и ваш AWS CLI, как указано в aws configure get region). Похоже, что у AWS есть человек, отвечающий за увеличение квоты, и, по нашему опыту, для получения разрешения на увеличение квоты может потребоваться день или два.

Во-вторых, нам понадобится локальная среда Python с Meadowrun, а затем мы установим Meadowrun в нашей учетной записи AWS. Вот пример с использованием pip в Linux:

$ python3 -m venv meadowrun-venv
$ source meadowrun-venv/bin/activate
$ pip install meadowrun
$ meadowrun-manage-ec2 install --allow-authorize-ips

Запуск DALL-E Mini

Теперь, когда мы разобрались с этим, легко запустить бэкенд dalle-playground!

import asyncio
import meadowrun

async def run_dallemini():
    return await meadowrun.run_command(
        "python backend/app.py --port 8080 --model_version mini",
        meadowrun.AllocCloudInstance("EC2"),
        meadowrun.Resources(
            logical_cpu=1,
            memory_gb=16,
            max_eviction_rate=80,
            gpu_memory=4,
            flags="nvidia"
        ),
        meadowrun.Deployment.git_repo(
            "https://github.com/hrichardlee/dalle-playground",
            interpreter=meadowrun.PipRequirementsFile("backend/requirements.txt", "3.9")
        ),
        ports=8080
    )

asyncio.run(run_dallemini())

Краткое описание этого фрагмента:

• run_command указывает Meadowrun запустить python backend/app.py --port 8080 --model_version mini на экземпляре EC2. Это запускает бэкенд dalle-playground на порту 8080, используя мини-версию DALL-E Mini. Мини-версия в 27 раз меньше мега-версии DALL-E Mini, что делает ее менее мощной, но более простой в запуске.
• Следующие несколько строк сообщают Meadowrun, каковы требования для нашей работы: 1 процессор, 16 ГБ оперативной памяти, и мы не против точечных экземпляров с вероятностью выселения (прерывания) до 80%. Типы экземпляров, которые мы будем использовать, имеют тенденцию прерываться, поэтому, если это станет проблемой, мы можем изменить это значение на 0%, что скажет Meadowrun, что мы хотим получить экземпляр по требованию. Мы также просим графический процессор Nvidia с памятью GPU не менее 4 ГБ, что необходимо для мини-модели.
• Далее нам нужен код в репозитории https://github.com/hrichardlee/dalle-playground, и мы хотим создать окружение pip из файла backend/requirements.txt в этом репозитории. Мы почти смогли использовать репо saharmor/dalle-playground как есть, но нам пришлось сделать одно изменение, чтобы добавить пакет jax[cuda] в файл requirements.txt. Если вы еще не знакомы с jax, то jax — это библиотека машинного обучения от Google, примерно эквивалентная Tensorflow или PyTorch. Она сочетает в себе Autograd для автоматического дифференцирования и XLA (ускоренная линейная алгебра) для JIT-компиляции numpy-подобного кода для TPU от Google или API CUDA от Nvidia для GPU. Поддержка CUDA требует явного выбора опции [cuda] при установке пакета.
• Наконец, мы сообщаем Meadowrun, что хотим открыть порт 8080 на машине, на которой выполняется это задание, чтобы мы могли получить доступ к бэкенду с нашего текущего IP-адреса. Будьте осторожны с этим! dalle-playground не использует TLS, и это не очень хорошая идея — дать всем с вашего IP-адреса доступ к этому интерфейсу навсегда.

Для просмотра выбранных частей вывода этой команды:

Launched a new instance for the job:
ec2-3-138-184-193.us-east-2.compute.amazonaws.com: g4dn.xlarge (4.0
CPU, 16.0 GB, 1.0 GPU), spot ($0.1578/hr, 61.0% chance of
interruption), will run 1 workers

Здесь Meadowrun сообщает нам все, что нам нужно знать об экземпляре, который он запустил для этой работы, и сколько это будет стоить (всего 15 центов в час!).

Building python environment in container  eccac6...

Далее Meadowrun создает контейнер на основе содержимого файла requirements.txt, который мы указали. Это занимает некоторое время, но Meadowrun кэширует изображение в ECR для вас, поэтому это нужно сделать только один раз (пока ваш файл requirements.txt не изменится). Meadowrun также очищает изображение, если вы не используете его в течение некоторого времени.

--> Starting DALL-E Server. This might take up to two minutes.

Здесь мы добрались до кода в dalle-playground, которому нужно сделать несколько минут инициализации.

--> DALL-E Server is up and running!

И теперь мы готовы к работе!

Теперь нам нужно запустить front end на нашей локальной машине (если у вас нет npm, вам нужно будет установить node.js):

git clone https://github.com/saharmor/dalle-playground
cd dalle-playground/interface
npm start

Вы захотите создать URL бэкенда в отдельном редакторе, например, http://ec2-3-138-184-193.us-east-2.compute.amazonaws.com:8080 и скопировать/вставить его в webapp — введение его напрямую приводит к лишним запросам к частично полному URL, которые медленно завершаются.

Время генерировать изображения!

Было довольно легко заставить это работать, но эта модель не делает того, о чем мы ее просим. Для первого набора изображений у нас есть фигура, похожая на Бэтмена, но он не молится, и я не уверен, что он находится в Гефсиманском саду. Что касается второго набора изображений, похоже, что мы получаем либо оливковое масло, либо планету, но мы не получим ни того, ни другого на одном изображении, не говоря уже о целой системе. Давайте посмотрим, сможет ли «мега» версия DALL-E Mini сделать что-то лучше.

Запуск DALL-E Mega

DALL-E Mega — это увеличенная версия DALL-E Mini, то есть она имеет ту же архитектуру, но больше параметров. Теоретически мы можем просто заменить --model_version mini на --model_version mega_full в предыдущем фрагменте и получить мега-версию. Однако, когда мы это делаем, код инициализации dalle-playground занимает около 45 минут.

Нам не нужно никакого реального профилирования, чтобы понять это — если вы просто убьете процесс после того, как он проработает некоторое время, трассировка стека ясно покажет, что виновником является функция from_pretrained, которая загружает предварительно обученные веса модели из Weights and Biases (aka wandb). Weights and Biases — это платформа MLOps, которая помогает отслеживать код, данные и анализы, которые используются для обучения и оценки ML-модели. Для целей данной статьи это место, куда мы переходим для загрузки весов предварительно обученной модели. Мы можем посмотреть спецификацию артефактов, которые мы загружаем из wandb, перейти к веб-обзору для мега-версии и увидеть, что основной файл, который нам нужен, занимает около 10 ГБ. Если мы войдем в экземпляр EC2, который Meadowrun создает для выполнения этой команды, и запустим iftop, то увидим, что мы получаем неспешные 35 Мбит/с от wandb.

Мы не хотим ждать 45 минут каждый раз, когда запускаем DALL-E Mega, и нам неприятно видеть, как мощная машина с GPU потребляет 35 Мбит/с из интернета, в то время как почти все ее ресурсы простаивают. Поэтому мы внесли некоторые изменения в dalle-playground, чтобы кэшировать артефакты в S3. cache_in_s3.py фактически вызывает wandb.Api().artifact("dalle-mini/dalle-mini/mega-1:latest").download(), а затем загружает артефакты в S3. Чтобы проследить за этим, сначала нужно создать ведро S3 и предоставить доступ к нему роли Meadowrun EC2:

aws s3 mb s3://meadowrun-dallemini
meadowrun-manage-ec2 grant-permission-to-s3-bucket meadowrun-dallemini

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

Затем мы можем использовать Meadowrun для запуска длительного задания по загрузке на более дешевой машине — обратите внимание, что мы запрашиваем только 2 ГБ памяти и никаких GPU для этого задания:

import asyncio
import meadowrun

async def cache_pretrained_model_in_s3():
    return await meadowrun.run_command(
        "python backend/cache_in_s3.py --model_version mega_full --s3_bucket meadowrun-dallemini --s3_bucket_region us-east-2",
        meadowrun.AllocCloudInstance("EC2"),
        meadowrun.Resources(1, 2, 80),
        meadowrun.Deployment.git_repo(
            "https://github.com/hrichardlee/dalle-playground",
            branch="s3cache",
            interpreter=meadowrun.PipRequirementsFile(
                "backend/requirements_for_caching.txt", "3.9"
            )
        )
    )

asyncio.run(cache_pretrained_model_in_s3())

Мы также изменили код модели, чтобы загружать файлы из S3 вместо wandb. Мы загружаем файлы в специальную папку /var/meadowrun/machine_cache, которая является общей для всех запущенных Meadowrun контейнеров на машине. Таким образом, если мы запустим один и тот же контейнер несколько раз на одной и той же машине, нам не придется заново скачивать эти файлы.

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

import asyncio
import meadowrun

async def run_dallemega():
    return await meadowrun.run_command(
        "python backend/app.py --port 8080 --model_version mega_full --s3_bucket meadowrun-dallemini --s3_bucket_region us-east-2",
        meadowrun.AllocCloudInstance("EC2"),
        meadowrun.Resources(1, 32, 80, gpu_memory=12, flags="nvidia"),
        meadowrun.Deployment.git_repo(
            "https://github.com/hrichardlee/dalle-playground",
            branch="s3cache",
            interpreter=meadowrun.PipRequirementsFile("backend/requirements.txt", "3.9")
        ),
        ports=8080
    )

asyncio.run(run_dallemega())

Следует отметить несколько моментов в этом фрагменте:

• Мы просим Meadowrun использовать ветку s3cache нашего git-репо, которая включает изменения, позволяющие кэшировать/получать артефакты в S3.
• Мы увеличили требования до 32 ГБ основной памяти и 12 ГБ памяти GPU, что требуется более крупной модели.
• При первом запуске Meadowrun создает новый образ, поскольку мы добавили пакет boto3 для получения наших кэшированных файлов из S3.

И последнее замечание — install программы Meadowrun устанавливает AWS Lambda, которая периодически запускается и автоматически очищает ваши экземпляры, если вы давно не запускали задание. Для дополнительной безопасности вы также можете вручную очистить инстансы:

meadowrun-manage-ec2 clean

Вот что мы получаем:

Намного лучше! Что касается первого набора изображений, я не уверен, что Бэтмен молится на всех этих изображениях, но он определенно Бэтмен и он определенно в Гефсиманском саду. Для второго набора изображений у нас теперь есть тарелка, оливковое масло и уксус, и это определенно больше похоже на солнечную систему. Изображения еще не совсем соответствуют OpenAI’s DALL-E, но они заметно лучше! К сожалению, мы не можем сделать слишком много для улучшения перевода текста в изображение, кроме обучения нашей собственной модели с 12 миллиардами параметров, но мы попробуем добавить модель диффузии для улучшения более тонких деталей на изображениях. Мы также добавим модель для масштабирования изображений, поскольку сейчас они имеют разрешение всего 256×256 пикселей.

Построение конвейера генерации изображений

Во второй половине статьи мы будем использовать пакет meadowdata/meadowrun-dallemini-demo, который содержит блокнот для запуска нескольких моделей в качестве последовательных пакетных заданий для генерации изображений с помощью Meadowrun. Комбинация моделей вдохновлена jina-ai/dalle-flow.

• DALL-E Mini: Модель, на которой мы сосредоточились в первой половине этой статьи. Этот пост — хорошее руководство по тому, как устроен DALL-E 2 от OpenAI. Упрощенно говоря, DALL-E — это комбинация двух моделей. Первая модель обучается на изображениях и учится «сжимать» изображения в векторы, а затем «распаковывать» эти векторы обратно в исходные изображения. Вторая модель обучается на парах изображение/заголовок и учится превращать подписи в векторы изображений. После обучения мы можем ввести новые подписи во вторую модель, чтобы получить вектор изображения, а затем мы можем ввести этот вектор изображения в первую модель, чтобы получить новое изображение.
• GLID-3-xl: диффузионная модель. Диффузионные модели обучаются путем получения изображений, их размытия (оно же диффузия) и обучения модели на парах оригинальное/размытое изображение. Модель учится восстанавливать исходную, не размытую версию по размытой версии. Диффузионные модели можно использовать для различных задач, но в данном случае мы будем использовать GLID-3-xl, чтобы заполнить более мелкие детали на наших изображениях.
• SwinIR: модель для повышения качества изображений (также известная как восстановление изображений). Модели восстановления изображений обучаются путем получения изображений и их уменьшения. Модель учится создавать оригинальное изображение более высокого разрешения из уменьшенного изображения.

Для запуска этого конвейера, в дополнение к предварительным условиям из первой половины статьи, мы получим git-репо meadowrun-dallemini-demo и локальные зависимости, а затем запустим сервер Jupyter notebook:

git clone https://github.com/meadowdata/meadowrun-dallemini-demo
cd meadowrun-dallemini-demo
# assuming you are already in a virtualenv from before
pip install -r local_requirements.txt
jupyter notebook

Затем нам нужно будет открыть основной блокнот в Jupyter и отредактировать S3_BUCKET_NAME и S3_BUCKET_REGION, чтобы они соответствовали ведру, которое мы создали в первой половине этой статьи.

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

• Мы адаптировали код примера, который поставляется со всеми нашими моделями, для использования нашего кэша S3 и предоставили простые в использовании интерфейсы в dalle_wrapper.py, glid3xl_wrapper.py и swinir_wrapper.py.
• В model_requirements.txt мы ссылаемся на наши три модели непосредственно как на git-репозитории (поскольку они не доступны как пакеты в PyPI), но нам пришлось сделать несколько изменений, чтобы эти репозитории работали как пакеты pip. Pip ищет файл setup.py в git-репо, чтобы выяснить, какие файлы из репо нужно установить в окружение, а также каковы зависимости этого репо. GLID-3-xl и latent-diffusion (другая модель диффузии, от которой зависит GLID-3-xl) имели файлы setup.py, которые нужно было подправить, чтобы включить весь код, необходимый для запуска моделей. У SwinIR вообще не было файла setup.py, поэтому мы добавили его. Наконец, все эти файлы setup.py требовали дополнительных зависимостей, которые мы просто добавили в файл model_requirements.txt.
• Все эти модели довольно сложно запустить на чем-либо, кроме Linux, поэтому мы разделили файл local_requirements.txt и model_requirements.txt. Даже если вы работаете на Windows или Mac, у вас не должно возникнуть проблем с запуском этого блокнота — Meadowrun позаботится о создании «волосатой» среды модели на экземпляре EC2 под управлением Linux.

И еще несколько замечаний о Meadowrun:

• Поскольку мы запускаем эти модели как пакетные задания, а не как сервисы, Meadowrun будет использовать один экземпляр EC2 для их запуска.
• Если вы чувствуете себя амбициозным, вы даже можете использовать meadowrun.run_map для параллельного запуска этих моделей на нескольких машинах с GPU.

Давайте посмотрим на результаты!

• Ноутбук запрашивает текстовый запрос и просит DALL-E Mini сгенерировать 8 изображений:

• Мы выбираем одно из изображений, и GLID-3-xl создает 8 новых изображений на основе выбранного нами изображения.

• Наконец, мы выбираем одно из этих изображений, и SwinIR увеличивает его с 256×256 до 1024×1024 пикселей:

Неплохо, хотя на каждом этапе нам помогал человек!

Вот что генерирует OpenAI’s DALL-E на основе того же запроса:

Rúnar

@runarorama

Бэтмен молится в Гефсиманском саду

13:56 PM — 21 Jul 2022

А вот еще одно сравнение:

Даниэль Компьютерные изображения 💿

@djbaskin_images

Оливковое масло и уксус, вылитые на тарелку в форме Солнечной системы
🪄#dalle2 #dalle

12:42 PM — 27 Apr 2022

Все это подчеркивает, насколько впечатляющим является DALL-E от OpenAI. При этом с DALL-E Mini очень весело играть, он по-настоящему открыт и будет становиться только лучше по мере обучения.

Заключительные замечания

Этот пост демонстрирует, как использовать Meadowrun для выполнения GPU-вычислений, таких как ML-интерпретация в EC2. Meadowrun позаботится о таких деталях, как поиск самых дешевых доступных типов экземпляров GPU, а также убедится, что CUDA и Nvidia Container Runtime (ранее известный как Nvidia Docker) установлены в нужных местах.

Очень здорово, что мы можем указать Meadowrun на репозиторий вроде dalle-playground, сообщить ему, какие ресурсы ему нужны, и запустить его практически без лишних хлопот. Одна из самых неприятных вещей в программном обеспечении — заставить работать чужой код, и приятно видеть, что экосистема Python и ML достигла в этом отношении большого прогресса. Благодаря лучшим инструментам управления пакетами, инструментам MLOps, таким как Hugging Face и Weights and Biases, а также Meadowrun (если можно так выразиться) стало как никогда легко использовать чужие наработки.

Чтобы быть в курсе новостей о Meadowrun, присоединяйтесь к нам на Github или следите за нами в Twitter!