Запускаем StableDiffusionXl на Nvidia видеокартах с помощью python
В данном посте я расскажу о том, как запускать модели StableDiffusion, в том числе тысячи их производных с civitai.com
Предисловие
В интернете есть уже много гайдов, по запуску моделей. Почему еще одна статья?
Тут, я бы хотел описать специализированный запуск моделей, оптимизированный только под 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, авторизуемся и видим что-то похожее на этот экран:
Создаем папку lora, через кнопку New. Или же назовите как хотите. Внутри этой папки создаем новый блокнот.
Откроем наш блокнот и установим зависимости
!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
. И много чего еще.
На этом все. Надеюсь, данная статья поможет обойти вам те грабли, которые мне удалось собрать по пути.