Как стать автором
Обновить

Запускаем StableDiffusionXl на Nvidia видеокартах с помощью python

Уровень сложностиСредний
Время на прочтение4 мин
Количество просмотров6.3K

В данном посте я расскажу о том, как запускать модели StableDiffusion, в том числе тысячи их производных с civitai.com

Работа SDXL
Работа SDXL

Предисловие

В интернете есть уже много гайдов, по запуску моделей. Почему еще одна статья?

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

Перед стартом, важно, чтобы у вас стоял на рабочей станции Docker выше версии 19.03, а также последние драйвера от Nvidia для вашей видеокарты. Cuda и другие штуки ставить не нужно.

Почему Docker?

Потому что мы будем использовать специальный образ от Nvidia для запуска питона, pytorch и jupyter, которые были скомпилированы и оптимизированы для данных видеокарт, с поддержкой бэкендов, разработанных Nvidia. Это позволит не только ускорить генерацию картинок, но и не засорять систему необходимыми драйверами и инструментами, по типу CUDA.

Начнем

Для начала, необходимо зайти на сайт с каталогом образов nvidia для pytorch.

Находим там самый последний образ и используем именно его. На текущий момент самый последний образ - это nvcr.io/nvidia/pytorch:24.05-py3. я буду использовать именно его.

Далее необходимо создать папку, из которой будет запускаться jupyter, в этой папке мы будем хранить модели, lora, и результаты наших генераций, если потребуется.

Создали? В моем случае, я создал папку jupyter на диске D. У вас может быть другой путь, в зависимости от ОС и воображения.

Переходим к запуску контейнера:

docker run --gpus all -it --name=pytorch -p 8888:8888 -v D:\jupyter:/jupyter-data --ipc=host nvcr.io/nvidia/pytorch:24.05-py3 jupyter notebook --ip 0.0.0.0 --notebook-dir="/jupyter-data"

Таким нехитрым способом, мы подняли контейнер с jupyter, включили там поддержку gpu и пробросили порт 8888 наружу, теперь мы можем цепляться к нему через IDE или просто открыть в браузере. Токен для доступа в jupyter можно скопировать из stdout вашего контейнера. Настоятельно рекомендую заменить его на пароль, при первом открытии. Но это не обязательно.

Ставим зависимости

Открываем веб версию Jupyter по адресу 127.0.0.1:8888, авторизуемся и видим что-то похожее на этот экран:

Стартовый экран Jupyter. Если вы только начали, то ваш список директорий будет пуст.
Стартовый экран Jupyter. Если вы только начали, то ваш список директорий будет пуст.

Создаем папку lora, через кнопку New. Или же назовите как хотите. Внутри этой папки создаем новый блокнот.

Блокнот ipykernel
Блокнот ipykernel

Откроем наш блокнот и установим зависимости

!pip install diffusers transformers accelerate safetensors ipywidgets compel peft --upgrade

После выполнения ячейки надо перезагрузить ядро.

Загружаем модель

from diffusers import (
    StableDiffusionXLPipeline,
    EulerAncestralDiscreteScheduler,
    AutoencoderKL
)
import torch

CHECKPOINT_NAME = 'stabilityai/stable-diffusion-xl-base-1.0'

torch.backends.cuda.matmul.allow_tf32 = True

# Load VAE component

vae = vae = AutoencoderKL.from_pretrained("stabilityai/sdxl-vae", torch_dtype=torch.float16)

# Configure the pipeline
pipe = StableDiffusionXLPipeline.from_pretrained(
    CHECKPOINT_NAME,
    vae=vae,
    torch_dtype=torch.float16,
)
pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(
    pipe.scheduler.config,
)
pipe.to('cuda')

# Отключение цензуры
pipe.safety_checker = None
pipe.requires_safety_checker = False

Исполняем ячейку и ждем, когда наша модель загрузится.

Hidden text

Не так давно, huggingface сломали torch компиляцию модели. Когда починят, можно будет добавить эту строчку для дополнительной оптимизации.

pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)

Готовим Compel

Compel - это библиотека, которая позволяет "перевзвесить" веса из вашего промпта, благодаря чему, можно засунуть большой текст в промпт вашей модели, игнорируя ограничение в 75 токенов у SD моделей.

Создаем и исполняем новую ячейку:

from compel import Compel, ReturnedEmbeddingsType

compel = Compel(tokenizer=[pipe.tokenizer, pipe.tokenizer_2], text_encoder=[pipe.text_encoder, pipe.text_encoder_2],
                returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
                requires_pooled=[False, True])

Генерируем картинку

В новой ячейке пишем простой код и исполняем его.

from IPython.display import display

pos = '' # позитивный промпт
neg = '' # негативный промпт

conditioning, pooled = compel(pos)
negative_embed, negative_pooled = compel(neg)
[conditioning, negative_embed] = compel.pad_conditioning_tensors_to_same_length([conditioning, negative_embed])

img = pipe(
    prompt_embeds=conditioning,
    pooled_prompt_embeds=pooled,
    negative_prompt_embeds=negative_embed,
    negative_pooled_prompt_embeds=negative_pooled,
    num_inference_steps=40,
    guidance_scale=7,
#     num_images_per_prompt=10 # количество картинок за генерацию
)

for image in img[0]:
    display(image)

Готово, теперь мы получили генерацию картинок с помощью StableDiffusionXL. Теперь можно подставлять самые разные модели в переменную CHECKPOINT_NAME и экспериментировать. Список моделей можно найти на https://huggingface.co/

Грузим модели с CivitAi

На данном ресурсе можно найти тысячи fine tune SD моделей, на любой вкус и цвет. Заменить готовые модели с huggingface на данные модели очень просто.

Checkpoints

Чтобы подгрузить готовую модель из чекпоинта, достаточно немного поменять нашу логику запуска модели, заменяем StableDiffusionXLPipeline.from_pretrained на StableDiffusionXLPipeline.from_single_file

CHECKPOINT_NAME = 'my_sdxl_model.safetensors'
pipe = StableDiffusionXLPipeline.from_single_file(
    CHECKPOINT_NAME,
    vae=vae,
    torch_dtype=torch.float16,
)

При этом CHECKPOINT_NAME заменяем на название файла, который вы скачали с civitai.

Lora

Чтобы подключить lora расширения, достаточно скачать ее и в новой ячейке подключить расширение к нашей модели

pipe.load_lora_weights('./', weight_name='lora_name.safetensors')

Что дальше?

Теперь, когда мы научились генерировать модели, можно приступать к экспериментам. Играться с моделями, lora, менять параметры num_inference_steps и guidance_scale, выбрать другой pipe.scheduler. И много чего еще.

На этом все. Надеюсь, данная статья поможет обойти вам те грабли, которые мне удалось собрать по пути.

Теги:
Хабы:
Всего голосов 5: ↑4 и ↓1+5
Комментарии5

Публикации

Истории

Работа

Data Scientist
71 вакансия
Python разработчик
101 вакансия

Ближайшие события

22 – 24 ноября
Хакатон «AgroCode Hack Genetics'24»
Онлайн
28 ноября
Конференция «TechRec: ITHR CAMPUS»
МоскваОнлайн
2 – 18 декабря
Yandex DataLens Festival 2024
МоскваОнлайн
11 – 13 декабря
Международная конференция по AI/ML «AI Journey»
МоскваОнлайн
25 – 26 апреля
IT-конференция Merge Tatarstan 2025
Казань