Развертывание предварительно обученной стабильной диффузионной модели в AWS Lambda

Руководство по установке для развертывания ваших моделей генераторов искусственного интеллекта

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

Базовый код использовал Openvino для создания оптимизированной для ЦП версии Stable Diffusion.

Эта структура особенно полезна для сценариев использования Edge и Интернета вещей, но здесь мы будем использовать что-то совершенно другое — мы развернем действительно большую модель (~ 3 ГБ и успешно запустим ее на AWS Lambda).

Я не буду ничего учить про openvino, потому что я с ним не знаком. На самом деле, я просто использовал репозиторий с открытым исходным кодом в качестве основы.

Вместо этого я сосредоточусь на том, чтобы получить клей, чтобы развернуть его на AWS Lambda.

Описанный подход должен работать независимо от лежащего в его основе фреймворка (обнимание лица, PyTorch и т. д.), поэтому очень полезно знать его, если вы хотите выполнить вывод по бессерверной конечной точке HTTP.

Эта небольшая история делится на:

1. Краткая предыстория
2. Пошаговое руководство по развертыванию на AWS Lambda
3. Мои глупые ошибки, прежде чем исправить

Немного предыстории стабильной диффузии

Но прежде чем мы углубимся, немного контекста о том, с чем мы работаем. Вы, наверное, слышали об изображениях, созданных ИИ, и их недавнем буме со стабильной диффузией. Если нет, и если вам интересна эта тема, вы можете проверить этот отличный блог на эту тему от Hugging Face: https://huggingface.co/blog/stable_diffusion.

Вот образец изображения кота-астронавта, который я создал с помощью Stable Diffusion:

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

Поскольку модель Stable Diffusion является открытым исходным кодом, разные люди также работали над предложением оптимизированных альтернатив: оптимизация для процессоров MacBook M1, оптимизация для чипов Intel и т. д.

Как правило, время, которое требуется, во многом зависит от фактического оборудования. Я извлек следующее приближение для вычислений для разных типов процессоров (проверено локально с RTX 3090 и i9, и только читал о M1 в Интернете):

Как видите, при выполнении на i9 решение Openvino мучительно медленное по сравнению с альтернативами.

Альтернативой является использование версии ONNX, предоставленной HuggingFace (у которой аналогичное время вычислений после оптимизации с помощью onnx simplifier (https://github.com/daquexian/onnx-simplifier).

Примечание: для упрощения модели требуется около 27 ГБ оперативной памяти. Я подозреваю, что окончательные результаты были бы такими же, если бы я использовал версию ONNX.

Несмотря на медлительность вывода ЦП, интересно видеть, что он может выполняться в AWS Lambda, что означает, что его можно использовать для бесплатных пробных/демонстрационных версий из-за щедрого бесплатного уровня AWS Lambda. Я создаю что-то подобное, например, бесплатную игрушку, которую я опубликовал на https://app.openimagegenius.com.

Если вы хотите поэкспериментировать с ним, наберитесь терпения и подождите до пяти минут для создания изображения. Lambda использует только 12 шагов вывода для ускорения (после прогрева Lambda выполнение занимает около 60 секунд).

Исходный код этого примера можно найти здесь. Не стесняйтесь использовать его по своему усмотрению, не спрашивая моего разрешения (это лицензия MIT). Только, пожалуйста, имейте в виду лицензию моделей.

https://github.com/paolorechia/openimagegenius/tree/main/serverless/stable-diffusion-open-vino-engine

Хватит болтать. Давайте перейдем к решению.

Бессерверный клей

(Рабочая версия: Lambda на основе контейнера с EFS🎉)

Ручная часть

К сожалению, здесь много ручных шагов. Хотя большую часть этого можно было бы автоматизировать, я не думал, что для меня имеет смысл тратить так много времени. Кто-то, кто хорошо разбирается в CloudFormation или Terraform, вероятно, сможет автоматизировать большинство (если не все) этих шагов.

Если вы попытаетесь выполнить эти шаги и застрянете, не стесняйтесь связаться со мной, я буду рад помочь.

Начнем с VPC и EC2

1. Создайте VPC или используйте стандартное. В любом случае это хорошо.
2. Создайте инстанс EC2, подключенный к этому VPC, и желательно разверните его в общедоступной подсети. Вам нужно будет подключиться к экземпляру через SSH, что намного проще, если вы используете общедоступную подсеть.
3. Создайте новую подгруппу безопасности или измените ту, которую вы используете. Вам понадобятся порты 22 и 2049, открытые на входе.

Теперь EFS

1. Создайте ЭФС. Убедитесь, что он доступен в той же подсети, что и экземпляр EC2, и использует ту же группу безопасности, которую вы определили.
2. SSH в свой экземпляр EC2.
3. Следуйте руководству AWS о том, как смонтировать EFS в вашем EC2: https://docs.aws.amazon.com/efs/latest/ug/wt1-test.html
4. Таким образом, вы выполняете следующие команды (соответствующим образом измените параметры с фактической папкой монтирования и целевым DNS-сервером, извлеченным из консоли AWS / CLI — вы можете найти DNS на экране пользовательского интерфейса файловой системы EFS)
5. mkdir mnt-folder
6. sudo mount -t nfs -o nfsvers=4.1,rsize=1048576,wsize=1048576,hard,timeo=600,retrans=2,noresvport mount-target-DNS:/ ~/mnt-folder
7. Сохраните файлы моделей openvino в EFS. В моем случае я загрузил их вручную с помощью этого кода (https://github.com/bes-dev/stable_diffusion.openvino/blob/master/stable_diffusion_engine.py) и предварительно загрузил их в корзину S3. Затем в моем экземпляре EC2 я загрузил из корзины S3 в EFS

(Примечание: для этого может потребоваться назначить роль вашему EC2, как показано на снимке экрана ниже.)

После того, как вы правильно настроили роль, aws-cli должен работать, например, вы можете выполнять такие команды, как aws3 sync s3://your-bucket.

Затем создайте точку доступа для эластичной файловой системы.

Создайте точку доступа EFS для настроенной вами эластичной файловой системы.

Здесь важно обратить внимание на несколько вещей:

Разрешения пользователя файловой системы — если они слишком ограничены, вы получите PermissionError при доступе к файлам EFS из вашей Lambda. В моем случае эта EFS была посвящена этой лямбде, поэтому я не заботился о детализации и просто дал широко открытый доступ (позже я сделаю то же самое в файле serverless):

Кроме того, избегайте назначения / пути к корневому каталогу, который вы определяете для точки доступа, это может вызвать проблемы при монтировании. Кроме того, обязательно запишите значение, которое вы выбрали. Вам нужно будет использовать его внутри вашей лямбда-функции. Я лично использовал: /mnt/fs, как это было описано в другом руководстве.

Хорошо, мы почти закончили создавать ресурсы вручную.

Бессерверная платформа

Большую часть тяжелой работы с бессерверным шаблоном, касающуюся частей EFS, я собрал с https://medium.com/swlh/mount-your-aws-efs-volume-into-aws-lambda-with-the-serverless-framework. -470b1c6b1b2»d.

Вот полный шаблон. Вы в значительной степени должны заменить идентификаторы ресурсов на свои собственные.

service: stable-diffusion-open-vino
frameworkVersion: "3"
provider:
  name: aws
  runtime: python3.8
  stage: ${opt:stage}
  region: eu-central-1
  memorySize: 10240
  iam:
    role:
      statements:
        - Effect: Allow
          Action:
            - "elasticfilesystem:*"
          Resource:
            - "arn:aws:elasticfilesystem:${aws:region}:${aws:accountId}:file-system/${self:custom.fileSystemId}"
            - "arn:aws:elasticfilesystem:${aws:region}:${aws:accountId}:access-point/${self:custom.efsAccessPoint}"
functions:
  textToImg:
    url: true
    image:
      name: appimage
    timeout: 300
    environment:
      MNT_DIR: ${self:custom.LocalMountPath}
    vpc:
      securityGroupIds:
        - ${self:custom.securityGroup}
      subnetIds:
        - ${self:custom.subnetsId.subnet0}
custom:
  efsAccessPoint: YOUR_ACCESS_POINT_ID
  fileSystemId: YOUR_FS_ID
  LocalMountPath: /mnt/fs
  subnetsId:
    subnet0: YOUR_SUBNET_ID
  securityGroup: YOUR_SECURITY_GROUP
resources:
  extensions:
    TextToImgLambdaFunction:
      Properties:
        FileSystemConfigs:
          - Arn: "arn:aws:elasticfilesystem:${self:provider.region}:${aws:accountId}:access-point/${self:custom.efsAccessPoint}"
            LocalMountPath: "${self:custom.LocalMountPath}"

Несколько частей, о которых стоит упомянуть:

Размер памяти: по умолчанию у вас не будет доступа к 10 ГБ памяти. Вам нужно открыть тикет с AWS для поддержки этого варианта использования. Обратите внимание, что вы не найдете конкретного случая для этого запроса. Я попросил увеличить объем хранилища Lambda и объяснил, что мне нужно больше памяти. Потребовалось несколько дней, чтобы AWS принял его.

memorySize: 10240

URL-адрес функции: эта строка url: true позволяет общедоступному URL-адресу вызывать вашу функцию, в основном только для целей разработки/отладки.

Режим сборки Docker-контейнера

provider:
	...
  ecr:
    images:
      appimage:
        path: ./
...
functions:
  textToImg:
    url: true
    image:
      name: appimage
    timeout: 300

Бессерверная структура здесь многое делает для вас: одни только эти блоки:

1. создать частный репозиторий ECR
2. используйте локальный Dockerfile для создания контейнера
3. пометить изображение
4. отправить его в частный репозиторий ECR
5. создайте функцию Lambda, которая использует только что созданный образ докера

Как говорится, будьте готовы. Наша сборка/развертывание займет значительное время по сравнению с родной AWS Lambda.

Вот Dockerfile, который я использовал:

FROM python:3.9.9-bullseye
WORKDIR /src
RUN apt-get update && \\
    apt-get install -y \\
    libgl1 libglib2.0-0 \\
    g++ \\
    make \\
    cmake \\
    unzip \\
    libcurl4-openssl-dev
COPY requirements.txt /src/
 
RUN pip3 install -r requirements.txt --target /src/
COPY handler.py stable_diffusion_engine.py /src/
ENTRYPOINT [ "/usr/local/bin/python", "-m", "awslambdaric" ]
CMD [ "handler.handler" ]

Он устанавливает интерфейс времени выполнения AWS Lambda и зависимости, необходимые для выполнения стабильной диффузии (версия openvino).

Опять же, благодарность принадлежит оригинальному автору решения openvino: https://github.com/bes-dev/stable_diffusion.openvino.

Прежде чем создавать образ докера, вам необходимо адаптировать модуль stable_diffusion_engine и получить обработчик Lambda. Вы можете либо взять их из моего репозитория, либо адаптировать оригиналы из репозитория выше.

Основные изменения, необходимые для движка, касаются того, как загружать модели:

        self.tokenizer = CLIPTokenizer.from_pretrained(
             "/mnt/fs/models/clip")
				(...)
        self._text_encoder = self.core.read_model(
            "/mnt/fs/models/text_encoder/text_encoder.xml")
				(...)
        self._unet = self.core.read_model(
	    "/mnt/fs/models/unet/unet.xml")
				(...)
        self._vae_decoder = self.core.read_model(
            "/mnt/fs/models/vae_decoder/vae_decoder.xml")
				(...)
        self._vae_encoder = self.core.read_model(
            "/mnt/fs/models/vae_encoder/vae_encoder.xml")

Затем вы можете использовать модуль в моем обработчике (который является просто адаптацией из файла demo.py с https://github.com/bes-dev/stable_diffusion.openvino):

# -- coding: utf-8 --`
print("Starting container code...")
from dataclasses import dataclass
import numpy as np
import cv2
from diffusers import LMSDiscreteScheduler, PNDMScheduler
from stable_diffusion_engine import StableDiffusionEngine
import json
import os
@dataclass
class StableDiffusionArguments:
    prompt: str
    num_inference_steps: int
    guidance_scale: float
    models_dir: str
    seed: int = None
    init_image: str = None
    beta_start: float = 0.00085
    beta_end: float = 0.012
    beta_schedule: str = "scaled_linear"
    model: str = "bes-dev/stable-diffusion-v1-4-openvino"
    mask: str = None
    strength: float = 0.5
    eta: float = 0.0
    tokenizer: str = "openai/clip-vit-large-patch14"
def run_sd(args: StableDiffusionArguments):
    if args.seed is not None:
        np.random.seed(args.seed)
    if args.init_image is None:
        scheduler = LMSDiscreteScheduler(
            beta_start=args.beta_start,
            beta_end=args.beta_end,
            beta_schedule=args.beta_schedule,
            tensor_format="np",
        )
    else:
        scheduler = PNDMScheduler(
            beta_start=args.beta_start,
            beta_end=args.beta_end,
            beta_schedule=args.beta_schedule,
            skip_prk_steps=True,
            tensor_format="np",
        )
    engine = StableDiffusionEngine(
        model=args.model, scheduler=scheduler, tokenizer=args.tokenizer, models_dir=args.models_dir
    )
    image = engine(
        prompt=args.prompt,
        init_image=None if args.init_image is None else cv2.imread(
            args.init_image),
        mask=None if args.mask is None else cv2.imread(args.mask, 0),
        strength=args.strength,
        num_inference_steps=args.num_inference_steps,
        guidance_scale=args.guidance_scale,
        eta=args.eta,
    )
    is_success, im_buf_arr = cv2.imencode(".jpg", image)
    if not is_success:
        raise ValueError("Failed to encode image as JPG")
    byte_im = im_buf_arr.tobytes()
    return byte_im
def handler(event, context, models_dir=None):
    print("Getting into handler, event: ", event)
    print("Working dir at handler...", )
    current_dir = os.getcwd()
    print(current_dir)
    print(os.listdir(current_dir))
    print("Listing root")
    print(os.listdir("/"))
    # Get args
    # randomizer params
    body = json.loads(event.get("body"))
    prompt = body["prompt"]
    seed = body.get("seed")
    num_inference_steps: int = int(body.get("num_inference_steps", 32))
    guidance_scale: float = float(body.get("guidance_scale", 7.5))
    args = StableDiffusionArguments(
        prompt=prompt,
        seed=seed,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
        models_dir=models_dir
    )
    print("Parsed args:", args)
    image = run_sd(args)
    print("Image generated")
    body = json.dumps(
        {"message": "wow, no way", "image": image.decode("latin1")})
    return {"statusCode": 200, "body": body}

Тестирование развертывания

Когда вы закончите развертывание, бессерверная среда должна предоставить вам URL-адрес, который вы можете вызвать следующим образом:

curl -X POST \\
https://your_lambda_url_id.lambda-url.eu-central-1.on.aws/ \\
-H 'content-type: application/json' \\
-d '{"prompt": "tree"}'

Если все работает, вы должны увидеть в журналах Cloud Watch, что он генерирует следующее изображение:

Когда я тестировал, на основной цикл ушло почти три минуты, а на выполнение полной лямбды — 238 секунд (четыре минуты).

Завиток выше даст вам нечитаемую строку с изображением, закодированным в latin1. Если вы планируете фактически использовать свою Lambda, возможно, вам может понадобиться что-то вроде этого (я использовал это для локального тестирования своего контейнера, замените URL-адрес):

import requests
import json
headers = {"content-type": "application/json"}
url = "<http://localhost:9000/2015-03-31/functions/function/invocations>"
body = json.dumps({"prompt": "beautiful tree", "num_inference_steps": 1})
response = requests.post(url, json={"body": body}, headers=headers)
response.raise_for_status()
j = response.json()
body = json.loads(j["body"])
bytes_img = body["image"].encode("latin1")
with open("test_result.png", "w+b") as fp:
    fp.write(bytes_img)

Фу! Это было много шагов! Теперь вы можете развернуть свою модель стабильного распространения в AWS Lambda. Надеюсь, вам понравилось читать этот краткий урок. Я оставлю вас с некоторыми вещами, которые я пробовал — но не совсем получилось — так что, может быть, я смогу убедить вас не пробовать их.

История ошибок

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

Первая попытка: полная EFS, нативная AWS Lambda

Итак, я много раз читал о развертывании больших моделей на Lambda и использовании AWS Elastic File System. Так я и сделал.

Я настроил обычную AWS Lambda и подключил ее к EFS. Однако, когда я выполнил код, я столкнулся с ошибкой при импорте среды выполнения openvino: libm.so.6 not found.

После некоторых размышлений и исследований я узнал, что AWS Lambda работает на Amazon Linux и что мне, возможно, следует создавать зависимости моей библиотеки непосредственно внутри экземпляра EC2.

За исключением того, что когда я попробовал это, я обнаружил, что openvino не имеет версии среды выполнения 2022, доступной для Amazon Linux (https://pypi.org/project/openvino/). Угу, тупик.

Вторая попытка: режим полного контейнера

Через несколько дней, после обсуждения с другом, насколько я не скучаю по использованию Docker по сравнению с использованием бессерверных технологий, загорелась лампочка: а что, если бы я использовал образ докера для развертывания Lambda?

Оказывается, ограничение на размер образа контейнера составляет 10 ГБ, что довольно щедро (https://aws.amazon.com/blogs/aws/new-for-aws-lambda-container-image-support/) — что должен работать!

Ну не так быстро. Когда все, казалось, приблизилось к работе, я столкнулся с некоторыми проблемами.

[ERROR] RuntimeError: Model file /src/models/unet/unet.xml cannot be opened!
Traceback (most recent call last):
  File "/src/handler.py", line 129, in handler
    image = run_sd(args)
  File "/src/handler.py", line 76, in run_sd
    engine = StableDiffusionEngine(
  File "/src/stable_diffusion_engine.py", line 59, in __init__
    self._unet = self.core.read_model(unet_xml, unet_bin)

Хм? Я смотрел на эту ошибку пару часов, отлаживая свою среду, проверяя доступность файла и т. д. Поскольку, согласно справочнику API OpenVINO, core.read_model также может напрямую принимать двоичные данные, я немного изменил свой код и пытался заранее загрузить модели в словарь бинарных буферов.

models = {}
for model in ["text_encoder", "unet", "vae_decoder", "vae_encoder"]:
    with open(f"./models/{model}/{model}.xml", "r+b") as fp:
        models[f"{model}-xml"] = fp.read()
    with open(f"./models/{model}/{model}.bin", "r+b") as fp:
        models[f"{model}-bin"] = fp.read()

Кроме того, я все еще сталкивался с ошибками, но на этот раз они были более значимыми.

[ERROR] OSError: [Errno 30] Read-only file system: './models/text_encoder/text_encoder.xml'
Traceback (most recent call last):
  File "/usr/local/lib/python3.9/importlib/__init__.py", line 127, in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
  File "<frozen importlib._bootstrap>", line 1030, in _gcd_import
  File "<frozen importlib._bootstrap>", line 1007, in _find_and_load
  File "<frozen importlib._bootstrap>", line 986, in _find_and_load_unlocked
  File "<frozen importlib._bootstrap>", line 680, in _load_unlocked
  File "<frozen importlib._bootstrap_external>", line 850, in exec_module
  File "<frozen importlib._bootstrap>", line 228, in _call_with_frames_removed
  File "/src/handler.py", line 23, in <module>
    from stable_diffusion_engine import StableDiffusionEngine
  File "/src/stable_diffusion_engine.py", line 30, in <module>
    with open(f"./models/{model}/{model}.xml", "r+b") as fp:

Я перепроверил документацию по Python и понял, что r+b на самом деле означает «открыт для обновления (чтения и записи)». Возможно, файловая система доступна только для чтения. Давайте попробуем еще раз без него, используя вместо этого только rb:

[ERROR] PermissionError: [Errno 13] Permission denied: './models/text_encoder/text_encoder.xml'
Traceback (most recent call last):
  File "/usr/local/lib/python3.9/importlib/__init__.py", line 127, in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
  File "<frozen importlib._bootstrap>", line 1030, in _gcd_import
  File "<frozen importlib._bootstrap>", line 1007, in _find_and_load
  File "<frozen importlib._bootstrap>", line 986, in _find_and_load_unlocked
  File "<frozen importlib._bootstrap>", line 680, in _load_unlocked
  File "<frozen importlib._bootstrap_external>", line 850, in exec_module
  File "<frozen importlib._bootstrap>", line 228, in _call_with_frames_removed
  File "/src/handler.py", line 23, in <module>
    from stable_diffusion_engine import StableDiffusionEngine
  File "/src/stable_diffusion_engine.py", line 30, in <module>
    with open(f"./models/{model}/{model}.xml", "rb") as fp:

Хорошо, может быть, сначала нужно скопировать файлы в /tmp? Нет, это просто дало мне ту же ошибку. Я не мог в этом разобраться — тот же код отлично работал локально, и я даже тестировал его с помощью Эмулятора интерфейса Lambda Runtime. Это должно было быть что-то с окружающей средой.

AWS блокирует чтение двоичных файлов из пути к файлу образа контейнера из соображений безопасности. Я так и не понял, почему именно. Переход на гибридный подход, когда модели хранятся в EFS, а зависимости кода/библиотеки — в образе Docker, прошел гладко.

Ладно, на сегодня все!

Надеюсь, вам понравилось читать. Ваше здоровье!