Как оптимизировать работу Stable Diffusion при помощи текстовой инверсии

Изображение, сгенерированное AI по промпту «photograph of a robot drawing in the wild, nature, jungle» («фотография робота, рисующего в природе, джунглях»)

22 августа 2022 года Stability.AI объявила о публичном релизе Stable Diffusion — мощной диффузионной модели text-to-image. Модель способна генерировать различные варианты изображений на основании текстового или графического ввода.

Стоит заметить, что «модель выпущена под лицензией Creative ML OpenRAIL-M. Лицензия допускает коммерческое и некоммерческое использование. Ответственность за этическое использование модели лежит на разработчиках. Это относится и к производным от неё моделям».

В этой статье я расскажу, как точно настраивать эмбеддинги для создания персонализированных изображений на основании произвольных стилей или объектов. Вместо переучивания модели мы можем представить собственный стиль в виде новых слов в пространстве эмбеддингов модели. В результате этого новое слово будет руководить созданием новых изображений интуитивно понятным образом.

Посмотрите на показанное ниже сравнение реального объекта и изображения, сгенерированного при помощи модели Stable Diffusion и точно настроенного эмбеддинга:

Давайте приступим к подготовке и установке.

Установка

Следует учесть, что этот туториал основан на форкнутой версии Stable Diffusion, в которую интегрирована текстовая инверсия для персонализации генерации изображений. В официальном репозитории рекомендуют использовать для установки Anaconda.

Для установки всех необходимых зависимостей изучите следующее руководство.

Если вы намереваетесь устанавливать все зависимости через pip install, то воспользуйтесь следующим файлом требований:

pip install -r requirements.txt

Если у вас возникнут проблемы с установкой, изучите следующее руководство по устранению проблем.

Убедитесь, что у вас есть все требуемые модели и чекпоинты в рабочей папке или в папке .cache.

Или же можно протестировать это при помощи следующего ноутбука Stable Diffusion AI на Google Colab. Следуйте инструкциям по запуску Stable Diffusion в изолированном окружении.

Подготовка данных

Первым делом в корневом каталоге (stable-diffusion) создайте папку training_data. Мы поместим в неё все наши обучающие изображения. Можно дать им любые имена, но они должны иметь следующие свойства:

• размер изображения: 512 x 512
• объекты на изображениях должны быть правильно повёрнуты
• для наилучших результатов используйте 3-5 изображений (если использовать слишком много изображений, модель может не сойтись)
• изображения должны содержать схожую контекстуальную информацию. Для точной настройки под стиль нужно использовать согласованные цветовые схемы и стиль графики. Для точной настройки под объекты на изображениях должен присутствовать одинаковый объект под разными углами (рекомендуется сохранять одно направление для объектов, передняя и задняя части которых различаются).

Для демонстрации я воспользуюсь своим брелоком в качестве данных обучения для точной настройки под объект.

Показанные выше фото сделаны камерой моего телефона, поэтому мне нужно было уменьшить их до размера 512 x 512. Внутри папки training_data я создал ещё две папки:

• raw — для сырых изображений, требующих изменения размера
• key — подпапка для всех обучающих изображений (512 x 512)

Для изменения размера обучающих изображений вы можете воспользоваться в качестве справки следующим скриптом:

from PIL import Image import glob  size = (512, 512)  # получаем все файлы в папке, проверяем, что все они являются изображениями files = glob.glob('./training_data/raw/*')  for fil in files:     # при необходимости реализуем здесь проверку типа файла          # получаем базовое имя, например "dragon.jpg" -> ("dragon", ".jpg")     basename = os.path.splitext(os.path.basename(fil))[0]      with Image.open(fil) as img:         # изменяем размер до 512 x 512         img = img.resize(size)                  # при необходимости поворачиваем изображение         # img = img.rotate(90)                  # сохраняем изменённое изображение, при необходимости измените способ ресэмплинга и папку для вывода         img.save(f"./training_data/key/{basename}.png", format="PNG", resample=Image.Resampling.NEAREST)

Конфигурация

Перейдите в папку <root>/configs/stable-diffusion/, в ней вы должны увидеть следующие файлы:

• v1-finetune.yaml
• v1-finetune_style.yaml
• v1-inference.yaml

Файл v1-finetune.yaml предназначен для точной настройки под объекты. Для точной настройки под стиль нужно использовать в качестве файла конфигурации v1-finetune_style.yaml.

Рекомендуется создать резервные копии файлов конфигурации на случай, если вы поломаете конфигурацию.

Стандартная конфигурация требует для обучения не менее 20 ГБ VRAM. Мы можем снизить требования к памяти, уменьшив размер батча и количество воркеров. Измените следующие строки:

data:   target: main.DataModuleFromConfig   params:     batch_size: 2     num_workers: 16 ... lightning:   callbacks:     image_logger:       target: main.ImageLogger       params:         batch_frequency: 500         max_images: 8

чтобы они выглядели так:

data:   target: main.DataModuleFromConfig   params:     batch_size: 1     num_workers: 8 ... lightning:   callbacks:     image_logger:       target: main.ImageLogger       params:         batch_frequency: 500         max_images: 1

• batch_size с 2 до 1
• num_workers с 16 до 8
• max_images с 8 до 1

Я протестировал новую конфигурацию, она работает на GeForce RTX 2080.

Для уменьшения объёма памяти можно изменить разрешение обучающего изображения в first_state_config. Рекомендую использовать 256 или 512, в зависимости от устройств GPU.

first_stage_config:       target: ldm.models.autoencoder.AutoencoderKL       params:         embed_dim: 4         monitor: val/rec_loss         ddconfig:           double_z: true           z_channels: 4           resolution: 256

Кроме того, вам нужно будет изменить initializer_words как часть ввода для эмбеддинга. Он принимает список строк:

# точная настройка на изображение игрушки initializer_words: ["toy"] # точная настройка на стиль графики initializer_words: ["futuristic", "painting"]

Обучение

Закончив с этим, выполните следующую команду:

python main.py --base ./configs/stable-diffusion/v1-finetune.yaml \     -t \     --actual_resume ./models/ldm/stable-diffusion-v1/model.ckpt \     -n my_key \     --gpus 0, \     --data_root ./training_data/key/

Укажите --no-test в командной строке, чтобы игнорировать тестирование при точной настройке. Для изменения initializer_words можно использовать аргумент --init_word. Стоит заметить, что он работает только для одной строки. Если вы намереваетесь использовать несколько строк, то нужно изменить их в файле конфигурации.

На момент написания статьи обучение можно выполнять только на одном GPU и после числа при указании GPU нужно поставить запятую.

Если всё настроено правильно, начнётся обучение и будет создана папка logs. Внутри неё будет находиться папка для каждого проведённого вами обучения, и в ней хранятся следующие папки:

• checkpoints — содержит выходной эмбеддинг и чекпоинты обучения
• configs — файлы конфигурации в формате YAML
• images — пример текста промпта, изображение промпта и сгенерированные изображения
• testtube — файлы, связанные с тестированием

Обучение будет продолжаться бесконечно или в зависимости от числа max_steps, заданного в файле конфигурации:

trainer:   benchmark: True   max_steps: 4000

Для справки см. v1-finetune.yaml, так как эта конфигурация не включена в v1-finetune_style.yaml. Пропустите её, чтобы обучение выполнялось бесконечно.
Чтобы понять качество результатов работы, проверьте папку images. Изображения с префиксом sample_scaled_ представляют собой сгенерированные сэмплы и если обучение происходит правильно, они должны иметь схожие черты с данными обучения. Можно остановить обучение, когда оно достигнет примерно 5000-7000 шагов.

Внутри папки checkpoints находится довольно большое количество файлов:

• embeddings_gs-5699.pt — файл эмбеддинга на шаге 5699
• epoch=000010.ckpt — чекпоинт обучения на 10 эпохах
• last.ckpt — чекпоинт обучения последней эпохи
• embeddings.pt — файл эмбеддинга последнего шага

Файлы ckpt используются для продолжения обучения. Файлы pt — это файлы эмбеддингов, которые должны использоваться совместно с моделью Stable Diffusion. Просто скопируйте нужный файл эмбеддинга и поместите его в удобное место для получения логического вывода (inference).

Логический вывод

▍ Интерактивный терминал

Для использования в командной строке просто задайте флаг embedding_path при запуске ./scripts/invoke.py.

python ./scripts/invoke.py --embedding_path /path/to/embeddings.pt

Затем используйте в интерактивном терминале следующий промпт для text2image:

invokeai> "a photo of *"

Кроме того, можно выполнить и image2image, задав флаг init_img:

invokeai> "waterfall in the style of *" --init_img=./path/to/drawing.png --strength=0.7

▍ Упрощённый API

Также можно использовать модуль ldm.generate для создания логического вывода программным образом. Он имеет функцию prompt2image как единую точку входа для задач text2image и image2image.

Стоит заметить, что старая версия содержит следующие функции, которые устарели, но всё равно работают:

• txt2img
• img2img

Для получения подробной информации по входным аргументам изучите скрипт на Python.

Полный код логического вывода приведён в этом gist:

from ldm.generate import Generate  # конфигурация модели gr = Generate(     'stable-diffusion-1.4',     embedding_path='./models/embeddings.pt'  # измените путь к эмбеддингам )  # загрузка модели gr.load_model()  # инициализация переменной text = "a photo of *"  # логический вывод возвращает список кортежей results = gr.prompt2image(     prompt   = text,     outdir   = "./outputs/",     iterations=1,     steps=50 )  # сохраняем изображение в папку outputs for row in results:     im   = row[0]     seed = row[1]     im.save(f'./outputs/image-{seed}.png')

Результаты

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

Изображение, сгенерированное AI по промпту «watercolor painting of * with trees» («акварельный рисунок * с деревьями»)

Изображение, сгенерированное AI по промпту «a bagpack in the style of *» («рюкзак в стиле *»)

Изображение, сгенерированное AI по промпту «a shirt with * logo» («футболка с логотипом *»

Кроме того, я протестировал задачи image2image и получил следующие результаты:

Изображение справа сгенерировано по промпту «three * on a beach» («три * на пляже»)

Заключение

Выпуск Stable Diffusion предоставляет огромный потенциал и возможности разработчикам всего мира. Он совершенно точно сильно повлияет на большинство современных отраслей, изменив понимание того, чего можно достичь при помощи AI.

Всего лишь за одну неделю я увидел, как разработчики со всего мира совместно работают над созданием собственных опенсорсных проектов, от UI для веба до анимации.

Я надеюсь увидеть новые прекрасные инструменты и современные модели, созданные опенсорсными сообществами.

Ссылки

1. Stability.AI — Stable Diffusion public release
2. Github — Stable Diffusion in a dream bot style interface
3. Github — Textual Inversion