Вероятно, все уже слышали о Stable Diffusion — модели, способной создавать фотореалистичные изображения на основе текста. Благодаря библиотеке diffusers от HuggingFace, использование этой модели очень просто.
Однако организация проекта и зависимостей для его запуска независимо от среды (будь то локально или в облаке), все еще может быть сложной задачей.
В этой статье, я на простом примере расскажу о том, как решать эту проблему с помощью diffusers и dstack. Мы напишем скрипт для генерации изображений с помощью предобученной модели, взятой из удаленного репозитория, и покажем, как легко выполнять этот скрипт как локально, так и удаленно в облаке. Это ускоряет разработку и отладку локально, позволяя при необходимости переключаться в облако, запрашивая необходимые ресурсы on‑demand.
Библиотека
diffusersPython предоставляет простой способ доступа к различным пред-обученным диффузионным моделям, опубликованным на Hugging Face, позволяя вам легко выполнять задачи инференса.
Инструмент
dstackпозволяет настраивать ML workflows (также иногда называемые пайплайнами) с помощью кода и запускать их локально или в облачном аккаунте, который вы настроили. Он позаботится о создании и удалении облачных ресурсов по мере выполнения.
Итак, давайте приступим.
Требования
Вот список библиотек Python, которые мы будем использовать:
diffusers
transformers
accelerate
scipy
ftfy
safetensorsПримечание: Используем библиотеку safetensors для хранения тензоров вместо
pickle, по рекоммендации Hugging Face, для большей безопасности и скорости.
Чтобы наши скрипты могли работать без проблем во всех средах, давайте включим их в файл stable_diffusion/requirements.txt.
Вы также можете установить эти библиотеки локально:
pip install -r stable_diffusion/requirements.txtУстановим также dstack CLI, поскольку мы будем использовать его локально:
pip install dstack --upgradeЗагрузка модели
В этом руководстве мы будем использовать модель runwayml/stable-diffusion-v1-5, обученную командой Runway ML. Однако на Hugging Face вы можете найти и ряд других моделей на выбор.
Давайте создадим следующий файл stable_diffusion/stable_diffusion.py:
import shutil
from diffusers import StableDiffusionPipeline
def main():
_, cache_folder = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5",
return_cached_folder=True)
shutil.copytree(cache_folder, "./models/runwayml/stable-diffusion-v1-5", dirs_exist_ok=True)Примечание: По умолчанию
diffusersзагружает модель в свою собственную директорию, управляемую с помощью симлинков. Посколькуdstackне поддерживает симлинки в артефактах, мы копируем файлы модели в локальную папку.
Чтобы запустить сценарий через dstack, сценарий должен быть определен как workflow через YAML-файл в .dstack/workflows.
Давайте определим следующий файл .dstack/workflows/stable_diffusion.yaml:
workflows:
- name: stable-diffusion
provider: bash
commands:
- pip install -r stable_diffusion/requirements.txt
- python stable_diffusion/stable_diffusion.py
artifacts:
- path: ./models
resources:
memory: 16GBТеперь workflow может быть запущен локально или в удаленной среде через dstack CLI.
Примечание: Прежде чем запускать workflow через
dstack, убедитесь, что в репозитории проекта настроена удаленная ветка Git, и вызовите командуdstack init, которая обеспечит доступdstackк репозиторию.
Вот как запустить workflow локально:
dstack run stable-diffusionКогда вы запустите его, dstack выполнит скрипт и сохранит папку models как артефакт. После этого вы можете повторно использовать артефакт в других workflows.
Присоединение интерактивной IDE
Иногда, прежде чем запустить workflows, вы можете захотеть выполнить код в интерактивном режиме, например, с помощью IDE или ноутбука.
Посмотрите на следующий пример:
workflows:
- name: code-stable
provider: code
deps:
- workflow: stable-diffusion
setup:
- pip install -r stable_diffusion/requirements.txt
resources:
memory: 16GBКак вы видите, code-stable ссылается на stable-diffusion как на зависимость. Запустите его.
Если вы это сделаете, в выводе вы увидите URL:
Он открывает VS Code, прикрепленный к вашему workflow, со всем настроенным: кодом, моделью и средой Python.
Генерация изображений по тексту
Давайте напишем скрипт, который генерирует изображения, используя предварительно обученную модель и заданный текст.
Вот пример файла stable_diffusion/prompt_stable.py:
import argparse
from pathlib import Path
import torch
from diffusers import StableDiffusionPipeline
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument("-P", "--prompt", action='append', required=True)
args = parser.parse_args()
pipe = StableDiffusionPipeline.from_pretrained("./models/runwayml/stable-diffusion-v1-5", local_files_only=True)
if torch.cuda.is_available():
pipe.to("cuda")
images = pipe(args.prompt).images
output = Path("./output")
output.mkdir(parents=True, exist_ok=True)
for i in range(len(images)):
images[i].save(output / f"{i}.png")Скрипт загружает модель из локальной папки ./models/runwayml/stable-diffusion-v1-5, генерирует изображения на основе заданного текста и сохраняет полученные изображения в локальной папке output.
Чтобы иметь возможность запускать его через dstack, давайте определим его в .dstack/workflows/stable_diffusion.yaml:
workflows:
- name: prompt-stable
provider: bash
deps:
- workflow: stable-diffusion
commands:
- pip install -r stable_diffusion/requirements.txt
- python stable_diffusion/prompt_stable.py ${{ run.args }}
artifacts:
- path: ./output
resources:
memory: 16GBКогда вы запустите этот workflow, dstack смонтирует артефакты stable‑diffusion в рабочую директорию. Таким образом, модель, которая была загружена ранее, окажется в локальной папке ./models/runwayml/stable-diffusion-v1-5.
Примечание: Команда
dstack runпозволяет передавать аргументы workflow через${{ run.args }}.
Давайте запустим workflow локально:
dstack run prompt-stable -P "cats in hats"Примечание: Выходные артефакты локального запуска хранятся в папке
~/.dstack/artifacts.
Настройка AWS в качестве remote
По умолчанию workflows в dstack запускаются локально. Однако у вас есть возможность настроить remote. Например, вы можете настроить свою учетную запись AWS как удаленную для запуска рабочих процессов.
Чтобы настроить remote, выполните следующую команду:
dstack configЭта команда предложит вам выбрать профиль AWS, регион AWS для выполнения workflows и S3 bucket для хранения артефактов.
AWS profile: default
AWS region: eu-west-1
S3 bucket: dstack-142421590066-eu-west-1
EC2 subnet: noneПримечание: В настоящее время
dstackподдерживает только AWS в качестве удаленного бэкенда. Поддержка GCP и Azure ожидается в одном из ближайших релизов.
Удаленный запуск
После настройки удаленного бэкенда вы можете использовать флаг --remote в команде dstack run для удаленного запуска workflows.
Давайте сначала запустим stable-diffusion:
dstack run stable-diffusion --remoteПримечание: Когда вы запускаете workflow удаленно,
dstackавтоматически создает ресурсы в настроенном облаке, сохраняет артефакты и освобождает их по завершении.
Когда вы запускаете workflow удаленно, вы можете настроить необходимые ресурсы для запуска рабочих процессов: либо через свойство resources в YAML, либо через аргументы команды dstack run, такие как --gpu, --gpu-name и т.д.
Давайте запустим prompt-stable удаленно и попросим использовать GPU:
dstack run prompt-stable --remote --gpu 1 -P "cats in hats"Примечание: По умолчанию
dstackвыбирает самую дешевую доступную машину, соответствующую требованиям к ресурсам. Например, в AWS, если вы запросите один GPU, он будет использоватьp2.xlargeинстанс с GPU NVIDIA Tesla K80.
Заключение
Если эта статья показалась вам интересной, вы можете углубиться в эту тему, изучив документацию по diffusers и dstack.
Исходный код можно найти на GitHub.
В одной из следующих статей мы углубимся не только в генерацию изображений, но и в файнтьюнинг Stable Diffusion.
