В данном посте я расскажу о том, как запускать модели 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

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

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. И много чего еще.

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