Вероятно, все уже слышали о Stable Diffusion — модели, способной создавать фотореалистичные изображения на основе текста. Благодаря библиотеке diffusers
от HuggingFace, использование этой модели очень просто.
Однако организация проекта и зависимостей для его запуска независимо от среды (будь то локально или в облаке), все еще может быть сложной задачей.
![](https://habrastorage.org/getpro/habr/upload_files/9b3/952/6db/9b39526db98ab12a3aeb64ad304a88cf.png)
В этой статье, я на простом примере расскажу о том, как решать эту проблему с помощью diffusers и dstack. Мы напишем скрипт для генерации изображений с помощью предобученной модели, взятой из удаленного репозитория, и покажем, как легко выполнять этот скрипт как локально, так и удаленно в облаке. Это ускоряет разработку и отладку локально, позволяя при необходимости переключаться в облако, запрашивая необходимые ресурсы on‑demand.
Библиотека
diffusers
Python предоставляет простой способ доступа к различным пред-обученным диффузионным моделям, опубликованным на 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.
![](https://habrastorage.org/getpro/habr/upload_files/ad2/f76/14a/ad2f7614ae0990dde8717fc27ece3fc2.png)
Генерация изображений по тексту
Давайте напишем скрипт, который генерирует изображения, используя предварительно обученную модель и заданный текст.
Вот пример файла 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
.
![](https://habrastorage.org/getpro/habr/upload_files/638/e7a/b1f/638e7ab1fb57818849907248e9775a4a.png)
Настройка 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.