Обучение модели Stable Diffusion текстовой инверсии с помощью diffusers

Листая интернет на наличие интересных технологий в области нейронных сетей и различного искусства,я наткнулся на пост в Твиттере, в котором Suraj Patil объявил о возможности обучения модели Stable Diffusion текстовой инверсии используя всего 3-5 изображений.

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

Что такое текстовая инверсия?

Постараюсь рассказать это на более понятном языке.

Текстовая инверсия позволяет вам обучать модель новой «концепции», которую вы можете связать со «словом», без изменения весов, вместо этого точно настраивая векторы встраивания текста.

Так, например, я добавил кучу фотографий плюшевого пингвина. Затем я смог попросить модель, например, создать рисунок моего конкретного плюшевого пингвина . Или сгенерировать фотографию моего конкретного плюшевого пингвина , сидящего на вершине горы. Цель состоит в том, чтобы позволить вам импортировать конкретную идею/концепцию из фотографий/изображений, которые не были сгенерированы моделью, и научиться представлять их. Его также можно использовать для представления стилей. Итак, скажем, вы загрузили кучу рисунков определенного художника, затем вы можете попросить его сгенерировать изображения в этом конкретном стиле.

Использование текстовой инверсии

Сейчас я покажу как можно обучить Stable Diffusion текстовой инверсии. Этот код также можно найти в гитхаб репозитории или использовать этот Колаб.

До начала работы советую зарегистрироваться в HuggingFace и получить access token с настройкой доступа "write". Теперь можно приступать.

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

• Сначало необходимо установить необходимые библиотеки и зайти в HuggingFace. Для этого нам и нужен access token. Вход в HuggingFace позволит сохранить вашу обученную модель и при необходимости поделиться с ней в библиотеке контента Stable Diffusion

• В разделе "Settings for teaching your new concept (настройки обучения для нового концепта)" необходимо выбрать чекпоинт на котором и будет завязано все обучение. По умолчанию там находится чекпоинт "CompVis/stable-diffusion-v1-4". Вы можете указать любой другой

• Далее начинается работа с датасетом. Официальный колаб текстовой инверсии позволяет использовать прямые ссылки на изображения, но я для своего проекта с наследованием стиля художника Ильи Кувшинова использовал гугл диск на котором у меня хранится около 1000 изображений артов этого художника (использовал я только 30 изображений). Для этого вошёлел в свой аккаунт колаба и используя модуль shutil копировал папку со своего гугл диска:

from google.colab import drive
drive.mount('/content/drive')
import shutil
shutil.copytree('/content/drive/My Drive/ForSD/', '/content/kuvshinov')

• Далее пришлось слегка видоизменить код чтобы модель могла использовать мои изображения:

import glob

path = 'kuvshinov/'

def download_image(filename):
  return Image.open(path+filename).convert("RGB")

images = list(filter(None,[download_image(filename) for filename in os.listdir(path)]))
save_path = "./my_concept"
if not os.path.exists(save_path):
  os.mkdir(save_path)
[image.save(f"{save_path}/{i}.jpeg") for i, image in enumerate(images)]

• Теперь можно переходить к настройке модели. В разделе "what_to_teach" вы можете выбрать объект (позволяет обучить модель новым объектам) или стиль (как в моем случае стилю изображений художника). В разделе "placeholder_token" необходимо указать имя по которому в дальнейшем можно будет использовать на тренированную модель. У меня это <kuvshinov-style>. В разделе "initializer_token" необходимо указать 1 слово которое будет показывать о чем ваша модель. У меня это kuvshinov.

• Теперь можно начинать обучение модели. Для этого просто запустите ячейки находящиеся в этом разделе "Teach the model a new concept (fine-tuning with textual inversion)". Вы также можете посмотреть какие параметра там используются и если вам интересно даже отредактировать их. Иногда во время запуска ячеек этого раздела может возникать ошибка при которой код ругается на неправильный initializer_token. Чаще всего эту ошибку выдает одна и также ячейка поэтому вы можете закомментировать не нужный фрагмент кода как у меня. Название этой ячейки "Get token ids for our placeholder and initializer token. This code block will complain if initializer string is not a single token". После этого можете дальше запускать код

token_ids = tokenizer.encode(initializer_token, add_special_tokens=False)
# Check if initializer_token is a single token or a sequence of tokens
#if len(token_ids) > 1:
#    raise ValueError("The initializer token must be a single token.")

initializer_token_id = token_ids[0]
placeholder_token_id = tokenizer.convert_tokens_to_ids(placeholder_token)

• Обучение модели в среднем занимает +/- 3 часа. После этого вы можете протестировать полученную модель и опубликовать ее в библиотеке контента Stable Diffusion

• Для того чтобы протестировать и опубликовать модель необходимо зайти в раздел "Run the code with your newly trained model". В ячейке "Save your newly created concept to the library of concepts" вы можете указать название вашего концепта, а так же выбрать публиковать модель или нет. Если вы решите ее опубликовать то в строке hf_token_write необходимо указать токен доступа который вы указывали выше при входе в HuggingFace. Обязательно проверьте какой доступ у этого токена. Их всего 2: чтение и запись (write). Нам нужен токен записи.

• После этого вы можете запустить все остальные ячейки поочередно. В последней ячейке в строке prompt необходимо указать вашу текстовую подсказку на английском языке и указать ваш placeholder_token. На пример: a grafitti in a wall with a <kuvshinov-style> on it

Вопросы и ответы

1. Как много ресурсов требуется для обучения модели?

   Для обучения своей модели я использовал 30 изображений размером 1024х1024, но так же можно использовать и меньшее количество изображений. Разработчики данной технологии для своих тестов выбирали от 3 до 5 изображений

2. Зависит ли скорость обучения от размера датасета?

   Нет, не зависит. Я пробовал использовать как и 5, так и 30 изображений и в среднем скорость обучения оставалась той же

3. «Чтобы сохранить этот концепт для повторного использования, загрузите learned_embeds.binфайл или сохраните его в библиотеке концептов». Означает ли это, что я могу использовать это на моей локальной стабильной диффузии, которую я уже запускаю? Как мне это сделать, куда будет попадать этот файл .bin и как я скажу программе использовать его?

   Да, вы можете запустить вашу модель на локальном компьютере, но вы должны быть уверены что вам хватит видео памяти т.к в основе всего это лежит модель Stable Diffusion которая требует 8 ГБ видеопамяти. Чтобы посмотреть как это работает используйте этот колаб

4. Можно ли это использовать для добавления большего количества/новых вещей поп-культуры, чтобы иметь лучшие результаты для вещей, которым требуется больше данных в существующем наборе, а также для добавления новых вещей, которых нет в наборе данных, чтобы SD мог выводить данные по указанным вещам?

   Да. Можно

Спасибо что прочли данную статью. Это моя первая статья поэтому не судите строго. Если хотите почитать больше об этом то посмотрите гитхаб diffusers и notebook. Вы также можете задать некоторые вопросы мне, хоть я и не квалифицированный специалист и не могу точно рассказать об этом, но в чем нибудь да помогу.